Ensemble and Mixture-of-Experts DeepONets For Operator Learning

Introduction

We present a novel deep operator network (DeepONet) architecture for operator learning, the ensemble DeepONet, that allows for enriching the trunk network of a single DeepONet with multiple distinct trunk networks. This trunk enrichment allows for greater expressivity and generalization capabilities over a range of operator learning problems. We also present a spatial mixture-of-experts (MoE) DeepONet trunk network architecture that utilizes a partition-of-unity (PoU) approximation to promote spatial locality and model sparsity in the operator learning problem. We first prove that both the ensemble and PoU-MoE DeepONets are universal approximators. We then demonstrate that ensemble DeepONets containing a trunk ensemble of a standard trunk, the PoU-MoE trunk, and/or a proper orthogonal decomposition (POD) trunk can achieve 2-4x lower relative $l_2$ errors than standard DeepONets and POD-DeepONets on both standard and challenging new operator learning problems involving partial differential equations (PDEs) in two and three dimensions. Our new PoU-MoE formulation provides a natural way to incorporate spatial locality and model sparsity into any neural network architecture, while our new ensemble DeepONet provides a powerful and general framework for incorporating basis enrichment in scientific machine learning architectures for operator learning.

Requirements

For running the code, the following Python libraries are required:

1. numpy
2. matplotlib
3. json
4. scipy
5. torch
6. jax
7. haiku

We provide a requirements.txt file that can be used to install the libraries with pip:

>> pip install -r requirements.txt

Dataset

Please refer to the paper for the datasets used.

DeepONet code

We provide the Python code for the four problems (Darcy flow, lid-driven cavity flow, 2D reaction-diffusion, and 3D reaction-diffusion with variable coefficient diffusion) in this paper in the src/ folder. The code files are

1. darcy_flow_deeponet.py
2. cavity_flow_deeponet.py
3. diffrec_deeponet.py
4. diffrec_3d_deeponet.py

These files run the vanilla, POD, modified-POD, and ensemble DeepONets, and also contain plotting code. All the DeepONet models are defined in the jax_networks.py file. The utils.py file contains helper functions for forming the POD basis and implements the Partition-of-Unity (PoU) method. The train_utils.py file contains a generic Jax training function for DeepONets.

Hyperparameters: One can change the learning rate, activation function, and other DeepONet hyperparameters within the code files. The code also automatically saves the results in a descriptive folder name for each problem separately.

Citation

This repository is part of the following paper. Please cite the following paper if you would like to cite this repository and find its contents useful:

@article{SharmaShankar2025ensembledeeponet,
  title={Ensemble and Mixture-of-Experts DeepONets For Operator Learning},
  author={Sharma, Ramansh and Shankar, Varun},
  journal={Transactions on Machine Learning Research},
  arxiv = {https://arxiv.org/abs/2405.11907},
  url = {https://openreview.net/forum?id=MGdydNfWzQ},
  year={2025}
}