Ensemble and Mixture-of-Experts DeepONets For Operator Learning

## Abstract

Abstract 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 ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT 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.

## 1 Introduction

In recent years, machine learning (ML) has been applied with great success to problems in science and engineering. Notably, ML architectures have been leveraged to learn *operators*, which are function-to-function maps. In many of these applications, ML-based operators, often called *neural operators*, have been utilized to learn solution maps to partial differential equations (PDEs). This area of research, known as operator learning, has shown immense potential and practical applicability to a variety of real-world problems such as weather/climate modeling (Bora et al., ; Pathak et al., ), earthquake modeling (Haghighat et al., ), material science (Gupta & Brandstetter, ; Oommen et al., ), and shape optimization (Shukla et al., ). Some popular neural operators that have emerged are deep operator networks (DeepONets) (Lu et al., ), Fourier neural operators (FNOs) (Li et al., ), and graph neural operators (GNOs) (Li et al., ). DeepONets have also been extended to incorporate discretization invariance (Zhang et al., ), more general mappings (Jin et al., ), and multiscale modeling (Howard et al., ). In this work, we focus on the DeepONet architecture due to its ability to separate the function spaces involved in operator learning; for completeness, we discuss one possible extension to the FNO in Appendix B.

At a high level, operator learning consists of learning a map from an input function to an output function. The DeepONet architecture is an inner product between a *trunk network* that is a function of the output function domain, and a *branch network* that learns to combine elements of the trunk using transformations of the input function. In fact, one can view the trunk as a set of learned, nonlinear, data-dependent basis functions. This perspective was first leveraged to replace the trunk with a set of basis functions learned from a proper orthogonal decomposition (POD) of the training data corresponding to the output functions; the resulting POD-DeepONet achieved state-of-the-art accuracy on a variety of operator learning problems (Lu et al., ). More recently, this idea was further generalized by extracting a basis from the trunk as a postprocessing step (Lee & Shin, ); this approach proved to be highly successful in learning challenging operators (Peyvan et al., ).

In this work, we present the ensemble DeepONet, a DeepONet architecture that explicitly enables enriching a trunk network with multiple distinct trunk networks; however, this enriched/augmented trunk uses a single branch that learns how to combine multiple trunks in such a way as to minimize the DeepONet loss function. The ensemble DeepONet essentially provides a natural framework for *basis function enrichment* of a standard (vanilla) DeepONet trunk. We also introduce a novel partition-of-unity (PoU) mixture-of-experts (MoE) trunk, the PoU-MoE trunk, that produces smooth blends of spatially-localized, overlapping, distinct trunks. The use of compactly-supported blending functions allows the PoU formulation to have a strong inductive bias towards spatial locality. Acknowledging that such an inductive bias is not always appropriate for learning inherently global operators, we simply introduce this PoU-MoE trunk into our ensemble DeepONet as an ensemble member alongside other global bases such as the POD trunk.

Our results show that the ensemble DeepONet, especially the POD-PoU ensemble, shows 2-4x accuracy improvements over vanilla-DeepONets with single branches and up to 2x accuracy improvements over the POD DeepONet (also with a single branch) in challenging 2D and 3D problems where the output function space of the operator has functions with sharp spatial gradients. In Section 4, we summarize the relative strengths of five different ensemble formulations, each carefully selected to answer a specific scientific question about the effectiveness of ensemble DeepONets. We conclude that the strength of ensemble DeepONets lie not merely in overparametrization but rather in the ability to incorporate spatially local information into the basis functions.

### 1.1 Related work

Basis enrichment has been widely used in the field of scientific computing in the extended finite element method (XFEM) (McQuien et al., ; Belytschko & Black, ; Ballard et al., ), modern radial basis function (RBF) methods (Flyer et al., ; Bayona et al., ; Shankar & Fogelson, ; Shankar et al., ), and others (Cai et al., ). In operator learning, basis enrichment (labeled “feature expansion”) with trigonometric functions was leveraged to enhance accuracy in DeepONets and FNOs (Lu et al., ). The ensemble DeepONet generalizes these prior results by providing a natural framework to bring data-dependent, locality-aware, basis function enrichment into operator learning. PoU approximation also has a rich history in scientific computing (Melenk & Babuvska, ; Larsson et al., ; Shcherbakov & Larsson, ; Heryudono et al., ; Safdari-Vaighani et al., ; Shankar & Wright, ), and has recently found use in ML applications (Han et al., ; Cavoretto et al., ; Trask et al., ). In (Trask et al., ), which targeted (probabilistic) regression applications, the authors used trainable partition functions that were effectively black-box ML classifiers with polynomial approximation on each partition. In Han et al. () (which also targeted regression), the authors used compactly-supported kernels as weight functions (like in this work), but used kernel-based regressors on each partition. Our PoU-MoE formulation generalizes both these works by using neural networks on each partition and further generalizes the technique to operator learning. In general, ensemble learning and MoE have a rich history, and we provide a more in-depth overview in Appendix A. The ensemble and PoU-MoE DeepONets introduced here extend this body of work to deterministic operator learning and PDE applications.

Broader Impacts: To the best of the authors’ knowledge, there are no negative societal impacts of our work including potential malicious or unintended uses, environmental impact, security, or privacy concerns.

Limitations: Ensemble DeepONets, especially when using PoU-MoE trunks, contain 2-3x as many trainable trunk network parameters as a vanilla-DeepONet and consequently require more time to train (see Section 3.4 for runtime results and discussion); however, in future work, we plan to ameliorate this issue with a novel parallelization strategy for the PoU-MoE trunk. Further, due to limited time, we used a single branch network that outputs to $\mathbb{R}^{p}$ for all our results (an unstacked branch) rather than using $p$ branch networks that each output to $\mathbb{R}$ (a stacked branch) from Lu et al. (). This choice may result in lowered accuracy for all methods (not just ours), but certainly resulted in fewer parameters. However, our results extend straightforwardly to stacked branches also.

## 2 Ensemble DeepONets

Figure: Figure 1: An ensemble DeepONet containing a POD trunk and a PoU-MoE trunk.
Refer to caption: extracted/5894272/figures/diagram.png

In this section, we first discuss the operator learning problem, then present the ensemble DeepONet architecture for learning these operators. We also present the novel PoU-MoE trunk and a modification the POD trunk from the POD-DeepONet, both for use within the ensemble DeepONet.

### 2.1 Operator learning with DeepONets

Let $\mathcal{U}\left(\Omega_{u};\mathbb{R}^{d_{u}}\right)$ and $\mathcal{V}\left(\Omega_{v};\mathbb{R}^{d_{v}}\right)$ be two separable Banach spaces of functions taking values in $\Omega_{u}\subset\mathbb{R}^{d_{u}}$ and $\Omega_{v}\subset\mathbb{R}^{d_{v}}$, respectively. Further, let $\mathcal{G}:\mathcal{U}\rightarrow\mathcal{V}$ be a general (nonlinear) operator. The operator learning problem involves approximating $\mathcal{G}:\mathcal{U}\to\mathcal{V}$ with a parametrized operator $\hat{\mathcal{G}}:\mathcal{U}\times\Theta\to\mathcal{V}$ from a finite number of function pairs $\{(u_{i},v_{i})\}$, $i=1,\ldots,N$ where $u_{i}\in\mathcal{U}$ are typically called *input functions*, and $v_{i}\in\mathcal{V}$ are called *output functions*, *i.e.*, $v_{i}=\mathcal{G}(u_{i})$. The parameters $\Theta$ are chosen to minimize $\|\mathcal{G}-\hat{\mathcal{G}}\|$ in some norm.

In practice, the problem must be discretized. First, one puts samples the input and output functions at a finite set of function sample locations $X\in\Omega_{u}$ and $Y\in\Omega_{v}$, respectively; also let $N_{x}=|X|$ and $N_{y}=|Y|$. One then requires that $\|v_{i}(y)-\hat{\mathcal{G}}(u_{i})(y)\|^{2}_{2}$ is minimized over $(u_{i},v_{i})$, $i=1,\ldots,N$, where $u_{i}$ are sampled at $x\in X$ and $v_{i}$ at $y\in Y$. The vanilla-DeepONet is one particular parametrization of $\hat{\mathcal{G}}(u)(y)$ as $\hat{\mathcal{G}}(u)(y)=\left\langle\boldsymbol{\tau}(y),\boldsymbol{\beta}(u)
\right\rangle+b_{0}$ where $\left\langle,\right\rangle$ is the $p$-dimensional inner product, $\boldsymbol{\beta}:\mathbb{R}^{N_{x}}\times\Theta_{\boldsymbol{\beta}}\to
\mathbb{R}^{p}$ is the *branch* (neural) network, $\boldsymbol{\tau}:\mathbb{R}^{d_{v}}\times\Theta_{\boldsymbol{\tau}}\to\mathbb
{R}^{p}$ is the trunk network, and $b_{0}$ is a trainable bias parameter; $p$ is a hyperparameter that partly controls the expressivity of $\hat{\mathcal{G}}(u)(y)$. $\Theta_{\boldsymbol{\beta}}$ and $\Theta_{\boldsymbol{\tau}}$ are the trainable parameters in the branch and trunk, respectively.

### 2.2 Mathematical formulation

We now present the new ensemble DeepONet formulation; an example is illustrated in Figure 1. Without loss of generality, assume that we are given three distinct trunk networks $\boldsymbol{\tau}_{1}(y;\theta_{\boldsymbol{\tau}_{1}})$,$\boldsymbol{\tau}_{2}(y;\theta_{\boldsymbol{\tau}_{2}})$, and $\boldsymbol{\tau}_{3}(y;\theta_{\boldsymbol{\tau}_{3}})$, where $y$ corresponds to the domain of the output function $v(y)$. Assume further that $\boldsymbol{\tau}_{j}:\mathbb{R}^{d}\times\Theta_{\boldsymbol{\tau}_{j}}\to
\mathbb{R}^{p_{j}}$, $j=1,2,3$. Then, given a single branch network $\hat{\boldsymbol{\beta}}(u;\theta_{b})$, the ensemble DeepONet is given in vector form by:

$$ $\displaystyle\hat{\mathcal{G}}(u,y)$ $\displaystyle=\left\langle[\boldsymbol{\tau}_{1}(y;\theta_{\boldsymbol{\tau}_{ 1}}),\boldsymbol{\tau}_{2}(y;\theta_{\boldsymbol{\tau}_{2}}),\boldsymbol{\tau} _{3}(y;\theta_{\boldsymbol{\tau}_{3}})],\hat{\boldsymbol{\beta}}(u;\theta_{b}) \right\rangle+b_{0}=\left\langle\hat{\boldsymbol{\tau}},\hat{\boldsymbol{\beta }}(u;\theta_{b})\right\rangle+b_{0}.$ (1) $$

Here, $\hat{\boldsymbol{\tau}}:\mathbb{R}^{d_{v}}\times\Theta_{\boldsymbol{\tau}_{1}}
\times\Theta_{\boldsymbol{\tau}_{2}}\times\Theta_{\boldsymbol{\tau}_{3}}\to
\mathbb{R}^{p_{1}+p_{2}+p_{3}}$ is the *ensemble trunk*. Clearly, the individual trunks simply “stack” column-wise to form the ensemble trunk $\hat{\boldsymbol{\tau}}$; in Appendix C, we discuss other suboptimal attempts to form an ensemble trunk. The ensemble trunk now consists of $p_{1}+p_{2}+p_{3}$ (potentially trainable) basis functions, necessitating that the branch $\hat{\boldsymbol{\beta}}:\mathbb{R}^{N_{x}}\times\Theta_{\hat{\boldsymbol{
\beta}}}\to\mathbb{R}^{p_{1}+p_{2}+p_{3}}$.

##### A universal approximation theorem

###### Theorem 1 .

Let $\mathcal{G}:\mathcal{U}\to\mathcal{V}$ be a continuous operator. Define $\hat{\mathcal{G}}$ as $\hat{\mathcal{G}}(u,y)=\left\langle\hat{\boldsymbol{\tau}}(y;\theta_{
\boldsymbol{\tau}_{1}};\theta_{\boldsymbol{\tau}_{2}};\theta_{\boldsymbol{\tau
}_{3}}),\hat{\boldsymbol{\beta}}(u;\theta_{b})\right\rangle+b_{0}$, where $\hat{\boldsymbol{\beta}}:\mathbb{R}^{N_{x}}\times\Theta_{\hat{\boldsymbol{
\beta}}}\to\mathbb{R}^{p_{1}+p_{2}+p_{3}}$ is a branch network embedding the input function $u$, $b_{0}$ is the bias, and $\hat{\boldsymbol{\tau}}:\mathbb{R}^{d_{v}}\times\Theta_{\hat{\boldsymbol{\tau}
}_{1}}\times\Theta_{\hat{\boldsymbol{\tau}}_{2}}\times\Theta_{\hat{\boldsymbol
{\tau}}_{3}}\to\mathbb{R}^{p_{1}+p_{2}+p_{3}}$ is an *ensemble* trunk network. Then $\hat{\mathcal{G}}$ can approximate $\mathcal{G}$ globally to any desired accuracy, *i.e.*,

$$ $\displaystyle\mathcal{G}(u)(y)-\hat{\mathcal{G}}(u)(y)\|_{\mathcal{V}}\leq\epsilon,$ (2) $$

where $\epsilon>0$ can be made arbitrarily small.

###### Proof.

This automatically follows from the (generalized) universal approximation theorem (Lu et al., ) which holds for arbitrary branches and trunks.
∎

###### Theorem 1 .

###### Proof.

#### 2.2.1 The PoU-MoE trunk

Figure: Figure 2: Enriched bases on the 2D reaction-diffusion problem 3.2. The solutions exhibit sharp gradients (left); the PoU-MoE trunk has learned spatially-localized basis functions (middle); the POD trunk has learned a global basis function (right).
Refer to caption: extracted/5894272/figures/diffrec/diffrec_pod_pu_bases.png

We now present the PoU-MoE trunk architecture, which leverages partition-of-unity approximation. We begin by partitioning $\Omega_{v}$ into $P$ overlapping circular/spherical patches $\Omega_{k}$, $k=1,\ldots,P$, with each patch having its own radius $\rho_{k}$ and containing a set of sample locations $Y_{k}$; of course, $\bigcup\limits_{k=1}^{P}Y_{k}=Y$. The key idea behind the PoU-MoE trunk is to employ a separate trunk network on each patch $\Omega_{k}$ and then blend (and train) these trunks appropriately to yield a single trunk network on $\Omega$. Each $\boldsymbol{\tau}_{k}$ is trained at data on $Y_{k}$, but may also be influenced by spatial neighbors. The PoU-MoE trunk $\boldsymbol{\tau}_{{}_{\text{PU}}}(x)$ is given as follows:

$$ $\displaystyle\boldsymbol{\tau}_{{}_{\text{PU}}}(y;\theta_{\boldsymbol{\tau}_{{ }_{\text{PU}}}})$ $\displaystyle=\sum_{k=1}^{P}w_{k}(y)\boldsymbol{\tau}_{k}(y;\theta_{ \boldsymbol{\tau}_{k}}),$ (3) $$

where $\theta_{\boldsymbol{\tau}_{k}}$, $k=1,\ldots,P$ are the trainable parameters for each trunk. In this work, we choose the weight functions $w_{k}$ to be (scaled and shifted) compactly-supported, positive-definite kernels $\psi_{k}:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ that are $\mathbb{C}^{2}\left(\mathbb{R}^{d}\right)$. More specifically, on the patch $\Omega_{k}$, we select $\psi_{k}$ to be the $\mathbb{C}^{2}\left(\mathbb{R}^{3}\right)$ Wendland kernel (Wendland, ; ; Fasshauer, ; Fasshauer & McCourt, ), which is a *radial* kernel given by

$$ $\displaystyle\psi_{k}(y,y^{c})=\psi_{k}\left(\frac{\|y-y^{c}_{k}\|}{\rho_{k}} \right)=\psi_{k}(r)$ $\displaystyle=\begin{cases}(1-r)^{4}(4r+1),&\text{if $r\leq 1$}\\ 0,&\text{ if $r>1$}\end{cases},$ (4) $$

where $y^{c}_{k}$ is the center of the $k$-th patch. The weight functions are then given by

$$ $\displaystyle w_{k}(y)$ $\displaystyle=\frac{\psi_{k}(y)}{\sum_{j}\psi_{j}(y)},\;k,j=1,\ldots,P,$ (5) $$

which automatically satisfy $\sum_{k}w_{k}(y)=1$. Each trunk $\boldsymbol{\tau}_{k}$ can be viewed as an “expert” on its own patch $\Omega_{k}$, thus leading to a *spatial* MoE formulation via the PoU formalism. Both training and evaluation of $\boldsymbol{\tau}_{{}_{\text{PU}}}$ can proceed locally in that each location $y$ lies in only a few patches; our implementation leverages this fact for efficiency. Further, since the weight functions $w_{k}(y)$ are each compactly-supported on their own patches $\Omega_{k}$, $\boldsymbol{\tau}_{{}_{\text{PU}}}$ can be viewed as *sparse* in its constituent spatial experts $\boldsymbol{\tau}_{k}$. Nevertheless, by ensuring that neighboring patches overlap sufficiently, we ensure that $\boldsymbol{\tau}_{{}_{\text{PU}}}$ still constitutes a global set of basis functions. For simplicity, we use the same $p$ value within each local trunk $\boldsymbol{\tau}_{k}$. Figure 2 (middle) shows one of the learned PoU-MoE basis functions in the POD-PoU ensemble; the learned basis function exhibits strong spatial locality corresponding to partitions. In Appendix G, we present more evidence for this spatial localization in the PoU-MoE basis functions.

##### Partitioning:

We placed the patch centers in a bounding box around $\Omega$, place a Cartesian grid in that box, then simply select $P$ of the grid points to use as centers. In this case, the uniform radius $\rho$ is determined as (Larsson et al., ) $\rho=(1+\delta)0.5H\sqrt{d}$ where $\delta$ is a free parameter to describe the overlap between patches and $H$ is the side length of the bounding box. However, as a demonstration, we also used variable radii $\rho_{k}$ in Section 3.1. In this work, we placed patches by using spatial gradients of a vanilla-DeepONet as our guidance, attempting to balance covering the whole domain with resolving these gradients; see Appendix E.3 for a more in-depth discussion on partitioning strategies.

##### A universal approximation theorem

###### Theorem 2 .

Let $\mathcal{G}:\mathcal{U}\rightarrow\mathcal{V}$ be a continuous operator. Define $\mathcal{G}^{\dagger}$ as $\mathcal{G}^{\dagger}(u)(y)=\left\langle\boldsymbol{\beta}(u;\theta_{b}),\sum
\limits_{j=1}^{P}w_{j}(y)\boldsymbol{\tau}_{j}(y;\theta_{\boldsymbol{\tau}_{j}
})\right\rangle+b_{0}$, where $\boldsymbol{\beta}:\mathbb{R}^{N_{x}}\times\Theta_{\boldsymbol{\beta}}\to
\mathbb{R}^{p}$ is a branch network embedding the input function $u$, $\boldsymbol{\tau}_{j}:\mathbb{R}^{d_{v}}\times\Theta_{\boldsymbol{\tau}_{j}}
\to\mathbb{R}^{p}$ are trunk networks, $b_{0}$ is a bias, and $w_{j}:\mathbb{R}^{d_{v}}\to\mathbb{R}$ are compactly-supported, positive-definite weight functions that satisfy the partition of unity condition $\sum_{j}w_{j}(y)=1,j=1,\ldots,P$. Then $\mathcal{G}^{\dagger}$ can approximate $\mathcal{G}$ globally to any desired accuracy, *i.e.*,

$$ $\displaystyle\mathcal{G}(u)(y)-\mathcal{G}^{\dagger}(u)(y)\|_{\mathcal{V}}\leq\epsilon,$ (6) $$

where $\epsilon>0$ can be made arbitrarily small.

###### Proof.

See Appendix D for the proof. The high level idea is to use the fact that the (generalized) universal approximation theorem (Chen & Chen, ; Lu et al., ) already holds for each local trunk on a patch, then use the partition of unity property to effectively blend that result over all patches to obtain a global estimate.
∎

###### Theorem 2 .

###### Proof.

#### 2.2.2 The POD trunk

The POD trunk is a modified version of the trunk used in the POD-DeepONet (Chatterjee, ) of the output function data. First, we remind the reader of the POD procedure. Recalling that $\{v_{i}(y)\}_{i=1}^{N}$ are the output functions, first define the matrix $V_{ij}=\frac{1}{\sigma_{i}}(v_{i}(y_{j})-\mu_{i})$, where $\mu_{i}$ is the spatial mean of the $i$-th function and $\sigma_{i}$ is its spatial standard deviation. Define the matrix $T=\frac{1}{N}VV^{T}$, and let $\Phi$ be the matrix of eigenvectors of $T$ ordered from the smallest eigenvalue to the largest. Then, the POD-DeepONet involves selecting the first $p$ columns of $\Phi$ to be the trunk of a DeepONet so that $G_{{}_{\text{POD}}}(u,y)=\sum\limits_{i=1}^{p}\boldsymbol{\beta}_{i}(u)\phi_{i
}(y)+\phi_{0}(y)$, where $\phi_{0}(y)$ is the mean function of $v(y)$ computed from the training dataset, and $\phi_{i}(y)$ are the columns of $\Phi$ as explained above. In this work, we use a POD trunk that includes the mean function $\phi_{0}$ in the set of basis functions. We label this the “Modified-POD” trunk in our experiments; this “Modified-POD” trunk $\boldsymbol{\tau}_{{}_{\text{POD}}}$ is given by

$$ $\displaystyle\boldsymbol{\tau}_{{}_{\text{POD}}}(y)=\begin{bmatrix}\phi_{0}(y) &\phi_{1}(y)&\ldots&\phi_{p-1}(y)\end{bmatrix},$ (7) $$

Consistent with the POD-DeepONet philosophy, no activation function is needed and the POD trunk has no trainable parameters. Figure 2 (right) shows one of the learned POD basis functions in the POD-PoU ensemble.

#### 2.2.3 Other Neural Operators

While we restricted our attention to DeepONets in this work, the ensemble idea naturally extends to other neural operator architectures. In Appendix B, we briefly discuss our ideas on creating ensembles of global and local basis functions within the FNO.

## 3 Results

We present results of our comparison of the new ensemble DeepONet (with and without a PoU-MoE trunk) against vanilla and POD DeepONets. We considered different ensemble combinations of the vanilla, POD, and PoU-MoE trunks. Each of the following ensembles attempted to address a specific scientific question:

- 1.
Vanilla-POD: Does adding POD modes to a vanilla trunk enhance expressivity over using either trunk in isolation?
- 2.
Vanilla-PoU: Does spatial locality introduced by the PoU-MoE trunk aid a DeepONet?
- 3.
POD-PoU: Does having both POD global modes and PoU-MoE local expertise enhance expressivity over simply using a vanilla trunk?
- 4.
Vanilla-POD-PoU: If the answer above is affirmative, then does adding a vanilla trunk (representing extra trainable parameters) to a POD-PoU ensemble help further enhance expressivity?
- 5.
$(P+1)$-Vanilla: Is spatial localization truly important or is simple overparametrization all that is needed? We use $P+1$ vanilla trunks in this model, where $P$ is the number of PoU-MoE patches. This ensemble thus contains as many trunks as the vanilla-PoU or POD-PoU ensembles, but all basis functions are purely *global* in this setting.

The answers to these questions are shown in Table 4 and summarized in Section 4. In a nutshell, spatial localization is indeed important, as is using a mix of global and localized basis functions; simple overparametrization is insufficient to attain state-of-the-art accuracy. We now describe our experimental setup, and both the standard and novel benchmark test results that led us to this conclusion.

##### Important DeepONet details.

In all cases, for parsimony in the number of training parameters, we used a single branch (the unstacked DeepONet) that outputs to $\mathbb{R}^{p}$ rather than $p$ branches. We found that output normalization did not help significantly in this case. We scaled all our POD architecture outputs by $\frac{1}{p}$ (standalone or in ensembles), as advocated in Lu et al. ().

##### Experiment design.

In the remainder of this section, we establish the performance of ensemble DeepONets on benchmarks such as a 2D lid-driven cavity flow problem (Section 3.1) and a 2D Darcy flow problem on a triangle (Appendix F.1), both common in the literature (Lu et al., ; Batlle et al., ). However, we also wished to develop challenging new spacetime PDE benchmarks where the PDE solutions (output functions) possessed steep gradients, while the input functions were well-behaved. To this end, we present results for both a 2D reaction-diffusion problem (Section 3.2) and a 3D reaction-diffusion problem with sharply (spatially) varying diffusion coefficients (Section 3.3). In both cases, we constructed spatially discontinuous reaction terms that resulted in PDE solutions (output functions) with steep gradients. Such PDE solutions abound in scientific applications. We note at the outset that the ensemble DeepONet with the PoU-MoE trunk performed best when the solutions had steep spatial gradients. Results on the Darcy problem show that the ensemble approaches tested here were not as effective on that problem.

##### Error calculations.

For all problems, we compared the vanilla- and POD-DeepONets with the five different ensemble architectures described at the top of Section 3. We also compared these ensembles against a DeepONet with the modified POD trunk from Section 2.2.2 (labeled Modified-POD). For all experiments, we first computed the relative $l_{2}$ error for each test function, $e_{\ell_{2}}=\frac{\|\tilde{\underline{u}}-\underline{u}\|_{2}}{\|\underline{u
}\|_{2}}$ where $\underline{u}$ was the true solution vector and $\tilde{\underline{u}}$ was the DeepONet prediction vector; we then computed the mean over those relative $\ell_{2}$ errors. For vector-valued functions, we first computed pointwise magnitudes of the vectors, then repeated the same process. We also report a squared error (MSE) between the DeepONet prediction and the true solution averaged over $N$ *functions* $e_{\text{mse}}(y)=\frac{1}{N}\left(\tilde{u}(y)-u(y)\right)^{2}.$

##### Notation.

In the following text, we denote the space and time domains with $\Omega$ and $T$ respectively; the spatial domain boundary is denoted by $\partial\Omega$. A single spatial point is denoted by $y$, which can either be a point $(y_{1},y_{2})$ in $\mathbb{R}^{2}$ or a point $(y_{1},y_{2},y_{3})$ in $\mathbb{R}^{3}$.

##### Setup.

We trained all models for 150,000 epochs on an NVIDIA GTX 4080 GPU. All results were calculated over five random seeds. We annealed the learning rates with an inverse-time decay schedule. We used the Adam optimizer (Kingma & Ba, ) for training on the Darcy flow and the cavity flow problems, and the AdamW optimizer (Loshchilov & Hutter, ) on the 2D and 3D reaction-diffusion problems. Other DeepONet hyperparameters and the network architectures are listed in Appendix E.

**Table 1: Relative $l_{2}$ errors (as percentage) on the test dataset for the 2D Darcy flow, cavity flow, and reaction-diffusion, and the 3D reaction-diffusion problems. RD stands for reaction-diffusion.**
|  | Darcy flow | Cavity flow | 2D RD | 3D RD |
| --- | --- | --- | --- | --- |
| Vanilla | $0.857\pm 0.08$ | $5.53\pm 1.05$ | $0.144\pm 0.01$ | $0.127\pm 0.03$ |
| POD | $0.297\pm 0.01$ | $7.94\pm 2e-5$ | $5.06\pm 8e-7$ | $9.40\pm 8$ |
| Modified-POD | $0.300\pm 0.04$ | $7.93\pm 2e-5$ | $0.131\pm 4e-5$ | $0.155\pm 4e-5$ |
| (Vanilla, POD) | $0.227\pm 0.03$ | $0.310\pm 0.03$ | $0.0751\pm 4e-5$ | $5.24\pm 10.4$ |
| $(P+1)$-Vanilla | $1.19\pm 0.06$ | $2.17\pm 0.3$ | $0.0644\pm 0.02$ | $5.25\pm 10.3$ |
| (Vanilla, PoU) | $0.976\pm 0.03$ | $1.06\pm 0.05$ | $0.0946\pm 0.03$ | $5.25\pm 10.3$ |
| (POD, PoU) | $0.204\pm 0.02$ | $\mathbf{0.204\pm 0.01}$ | $\mathbf{0.0539\pm 4e-5}$ | $\mathbf{0.0576\pm 0.05}$ |
| (Vanilla, POD, PoU) | $\mathbf{0.187\pm 0.02}$ | $0.229\pm 0.01$ | $0.0666\pm 8e-5$ | $5.22\pm 10.4$ |

### 3.1 2D Lid-driven Cavity Flow

Figure: Figure 3: The 2D lid-driven cavity flow problem. We show in (A) an example input function; in (B) an example output function component; in (C) the four patches used for the PoU-MoE trunk; in (D), (E), and (F) the spatial mean squared error (MSE) for the vanilla, ensemble vanilla-POD-PoU, and ensemble POD-PoU DeepONets respectively.
Refer to caption: extracted/5894272/figures/cavity_steady/cavity_flow_final_fig.png

The 2D lid-driven cavity flow problem involves solving for fluid flow in a container whose lid moves tangentially along the top boundary. This can be described by the incompressible Navier-Stokes equations (with boundary conditions),

$$ $\displaystyle\frac{\partial\mathbf{u}}{\partial t}+(\mathbf{u}\cdot\nabla) \mathbf{u}$ $\displaystyle=-\nabla\mathbf{p}+\nu\Delta\mathbf{u},\;\;\nabla\cdot\mathbf{u}= 0,\;y\in\Omega,\;t\in T,$ (8) $\displaystyle\mathbf{u}$ $\displaystyle=\mathbf{u}_{b},$ (9) $$

where $\mathbf{u}=(u(y),v(y))$ is the velocity field, $p$ is the pressure field, $\nu$ is the kinematic viscosity, and $\mathbf{u}_{b}=\left(u_{b},v_{b}\right)$ is the Dirichlet boundary condition. We focused on the steady state problem and used the dataset specified in Lu et al. (, Section 5.7, Case A). We set $\Omega=[0,1]^{2}$ and learned the operator $\mathcal{G}:\mathbf{u}_{b}\rightarrow\mathbf{u}$. The steady state boundary condition is defined as,

$$ $\displaystyle u_{b}=U\left(1-\frac{\cosh\left(r(x-\frac{1}{2})\right)}{\cosh \left(\frac{r}{2}\right)}\right),\;\;v_{b}=0,$ (10) $$

where $r=10$. The other boundary velocities were set to zero. As described in Lu et al. (), the equations were then solved using a lattice Boltzmann method (LBM) to generate 100 training and 10 test input and output function pairs. All function pairs were generated over a range of Reynolds numbers in the range $[100,2080]$ (with $U$ and $\nu$ chosen appropriately), with no overlap between the training and test dataset. Figure 3 shows the four patches used to partition the domain.

We report the relative $\ell_{2}$ errors (as percentage) on the test dataset in Table 1. The vanilla-, modified POD-, and POD-DeepONets had the highest errors (in increasing order). The POD-PoU ensemble was the most accurate model by about an order of magnitude over the vanilla-DeepONet, and almost two orders of magnitude over the POD variants. While all ensembles outperformed the standalone DeepONets, the ensembles possessing POD modes appeared to do best in general. Further, adding a PoU-MoE trunk to the ensemble seemed to aid accuracy in general, but especially when POD modes were present. The spatial MSE figures in Figure 3 reflect the same trends.

### 3.2 A 2D Reaction-Diffusion Problem

Figure: Figure 4: The 2D reaction-diffusion problem. We show in (A) an example input function; in (B) an example output function; in (C) the six patches used for the PoU-MoE trunk; in (D), (E), and (F) the spatial mean squared error (MSE) for the vanilla, ensemble $(P+1)$-vanilla, and ensemble POD-PoU DeepONets respectively.
Refer to caption: extracted/5894272/figures/diffrec/diffrec_final_fig.png

Next, we present experimental results on a 2D reaction-diffusion problem. This equation governs the behavior of a chemical whose concentration is $c(y,t)$, and is given (along with boundary conditions) below:

$$ $\displaystyle\frac{\partial c}{\partial t}=k_{\text{on}}\left(R-c\right)c_{ \text{amb}}-k_{\text{off}}\;c+\nu\Delta c,\;y\in\Omega,\;t\in T,$ (11) $$

with the boundary condition $\nu\frac{\partial c}{\partial n}=0$ on $\partial\Omega$. The first r.h.s term is a binding reaction term modulated by $k_{\text{on}}$ and the second term an unbinding term modulated by $k_{\text{off}}$. $c_{\text{amb}}(y,t)=1+\cos(2\pi y_{1})\cos(2\pi y_{2}))\exp(-\pi t)$ is a background source of chemical available for reaction, $\nu=0.1$ is the diffusion coefficient, $R=2$ is a throttling term, and $n(y)$ is the unit outward normal vector on the boundary. In our experiments, we used $\Omega=\;[0,2]^{2}$ and $T=[0,0.5]$. We set the initial condition as a spatial constant $c(y,0)\sim\mathcal{U}(0,1)$. More importantly, $k_{\text{on}}$ and $k_{\text{off}}$ are discontinuous and given by

$$ $\displaystyle k_{\text{on}}=\begin{cases}2,&y_{1}\leq 1.0,\\ 0,&\text{otherwise}\end{cases},\;\;k_{\text{off}}=\begin{cases}0.2,&y_{1}\leq 1 .0,\\ 0,&\text{otherwise}\end{cases},$ (12) $$

where $y_{1}$ is the horizontal direction. This discontinuity induces a sharp solution gradient at $y_{1}=1.0$ (see Figure 4 (B)). Our goal was to learn the solution operator $G:c(y,0)\to c(y,0.5)$. We solved the PDE numerically at $N_{y}=2207$ collocation points using a fourth-order accurate RBF-FD method (Shankar & Fogelson, ; Shankar et al., ); using this solver, we generated 1000 training and 200 test input and output function pairs. We sampled the random spatially-constant input on a regular spatial grid for the branch input. We used six patches for the PoU trunks as shown in Figure 4.

The third column of Table 1 shows that the POD-PoU ensemble achieved the lowest error, with an error reduction of almost 3x over the standalone DeepONets. The $(P+1)$-vanilla ensemble also performed reasonably well, with a greater than 2x error reduction over the same; this indicates that overparametrization indeed helped on this test case. However, the relatively higher errors of the vanilla-PoU ensemble (compared to the best results) indicate that POD modes are possibly vital to fully realizing the benefits of the PoU-MoE trunk. Once again, the spatial MSE plots in Figure 4 corroborate the relative errors.

### 3.3 3D Reaction-Variable-Coefficient-Diffusion

Figure: Figure 5: The 3D reaction-diffusion problem. We show in (A) an example input function; in (B) an example output function; in (C) the eight patches used for the PoU-MoE trunk; in (D), (E), and (F) the spatial mean squared error (MSE) for the vanilla, modified POD, and ensemble POD-PoU DeepONets respectively.
Refer to caption: extracted/5894272/figures/diffrec_3d/diffrec_3d_final_fig.png

Finally, we present results on a 3D reaction-diffusion problem with *variable-coefficient diffusion*. We used a similar setup to the 2D case but significantly also allow the diffusion coefficient to vary spatially via a function $K(y)$, $y\in\mathbb{R}^{3}$. The PDE and boundary conditions are given by

$$ $\displaystyle\frac{\partial c}{\partial t}$ $\displaystyle=k_{\text{on}}\left(R-c\right)c_{\text{amb}}-k_{\text{off}}\;c+ \nabla\cdot\left(K(y)\nabla c\right),\;y\in\Omega,\;t\in T,$ (13) $$

with $K(y)\frac{\partial c}{\partial n}=0$ on $\partial\Omega$. Here, $\Omega$ was the unit ball, *i.e.*, the interior of the unit sphere $\mathbb{S}^{2}$, and $T=[0,0.5]$. We set the $k_{\text{on}}$ and $k_{\text{off}}$ coefficients to the same values as in 2D in $y_{1}\leq 0$, and to zero in the $y_{1}>0$ half of the domain. We set $c_{\text{amb}}=(1+\cos(2\pi y_{1})\cos(2\pi y_{2})\sin(2\pi y_{3}))e^{(-\pi t)}$. All other model parameters were kept the same. $K(y)$ was chosen to have steep gradients, here defined as

$$ $\displaystyle K(y)$ $\displaystyle=B+\frac{C}{\tanh(A)}\left((A-3)\tanh(8x-5)-(A-15)\tanh(8x+5)+A \tanh(A)\right),$ (14) $$

where $A=9$, $B=0.0215$, and $C=0.005$. Once again, we learned the operator $\mathcal{G}:c(y,0)\rightarrow c(y,0.5)$. We again used the same RBF-FD solver to generate 1000 training and 200 test input/output function pairs (albeit at 4325 collocation points in 3D). We used eight spatial patches for the PoU trunks as shown in Figure 5. The last column in Table 1 shows that most of the ensemble DeepONets did poorly, as did the POD-DeepONet. However, the POD-PoU ensemble achieved almost a 2x reduction in error over the vanilla-DeepONet.

### 3.4 Runtime Comparison

The ensemble DeepONet architectures all have more trainable parameters than the vanilla and POD DeepOnets. This leads to higher training and inference times. We report the average time per training epoch and inference time on the test dataset in Tables 2 and 3 respectively. The training times were larger in ensemble DeepONets with more trunk networks, considerably so when the PoU-MoE trunks were used (an order of magnitude increase in training time on the 3D reaction-diffusion problem). The inference times showed a similar trend, although much less pronounced (only half an order of magnitude slowdown in the 3D problem). These slowdowns are because our current PoU-MoE implementation contains a serial loop over the patches in the forward pass, leading to slower back-propagation over its parameters. In future work, we plan to address this with a novel parallelization strategy; we believe this will speed up the ensemble architectures with the PoU-MoE trunk considerably. It is also important to note that despite this increased cost, Table 1 shows that the POD-PoU ensemble is more than 2x as accurate as the vanilla-DeepONet; the POD-DeepONet (and other ensembles) have errors that are two orders of magnitude worse!

**Table 2: Average time per training epoch in seconds. RD stands for reaction-diffusion.**
|  | Darcy flow | Cavity flow | 2D RD | 3D RD |
| --- | --- | --- | --- | --- |
| Vanilla | $8.93e-4$ | $3.99e-4$ | $2.97e-4$ | $2.10e-4$ |
| POD | $5.19e-4$ | $2.46e-4$ | $2.06e-4$ | $1.22e-4$ |
| Modified-POD | $6.86e-4$ | $2.49e-4$ | $2.08e-4$ | $1.22e-4$ |
| (Vanilla, POD) | $9.80e-4$ | $3.92e-4$ | $3.03e-4$ | $2.32e-4$ |
| $(P+1)$-Vanilla | $1.10e-3$ | $8.51e-4$ | $7.27e-4$ | $9.45e-4$ |
| Vanilla-PoU | $8.67e-4$ | $9.52e-4$ | $1.03e-3$ | $1.39e-3$ |
| POD-PoU | $6.74e-4$ | $8.21e-4$ | $9.24e-4$ | $1.28e-3$ |
| Vanilla-POD-PoU | $8.55e-4$ | $9.48e-4$ | $1.05e-3$ | $1.43e-3$ |

**Table 3: Inference time on the test dataset in seconds. RD stands for reaction-diffusion.**
|  | Darcy flow | Cavity flow | 2D RD | 3D RD |
| --- | --- | --- | --- | --- |
| Vanilla | $1.66e-4$ | $1.39e-4$ | $1.32e-4$ | $7.20e-5$ |
| POD | $1.57e-4$ | $1.12e-4$ | $1.12e-4$ | $6.42e-5$ |
| Modified-POD | $1.34e-4$ | $1.08e-4$ | $9.94e-5$ | $6.62e-5$ |
| (Vanilla, POD) | $1.69e-4$ | $1.33e-4$ | $1.20e-4$ | $7.76e-5$ |
| $(P+1)$-Vanilla | $2.08e-4$ | $2.12e-4$ | $1.71e-4$ | $1.48e-4$ |
| Vanilla-PoU | $1.91e-4$ | $2.42e-4$ | $2.21e-4$ | $2.37e-4$ |
| POD-PoU | $1.63e-4$ | $1.94e-4$ | $1.96e-4$ | $2.30e-4$ |
| Vanilla-POD-PoU | $2.00e-4$ | $2.18e-4$ | $2.28e-4$ | $2.41e-4$ |

## 4 Conclusions and Future Work

**Table 4: Effectiveness of different trunk choices. The yes/no refers to whether the strategy beats a vanilla-DeepONet. The bolded results are the best strategy for each experiment. RD stands for reaction-diffusion.**
| Trunk Choices | Darcy flow | Cavity flow | 2D RD | 3D RD |
| --- | --- | --- | --- | --- |
| Only POD global modes | Yes | No | No | No |
| Only modified POD global modes | Yes | No | No | No |
| \hdashlineAdding POD global modes | Yes | Yes | Yes | No |
| Adding spatial locality | No | Yes | Yes | No |
| Only POD global modes + spatial locality | Yes | Yes | Yes | Yes |
| Only POD global modes + spatial locality + mild overparametrization | Yes | Yes | Yes | No |
| Adding excessive overparametrization | No | Yes | Yes | No |

We presented the ensemble DeepONet, a method of enriching a DeepONet trunk with arbitrary trunks. We also developed the PoU-MoE trunk to aid in spatial locality. Our results demonstrated significant accuracy improvements over standalone DeepONets on several challenging operator learning problems, including a particularly challenging 3D problem in the unit ball. One of the goals of this work was to provide insight into choices for ensemble trunk members. Thus, we considered different combinations of three very specific choices: a vanilla-DeepONet trunk (vanilla trunk), the POD trunk, and the new PoU-MoE trunk. Our results (summarized in Table 4) make clear that while different ensemble strategies beat the vanilla-DeepONet in different circumstances, only the POD-PoU ensemble consistently beats the vanilla-DeepONet across all problems. Simple overparametrization ($(P+1)$-Vanilla DeepONet) is not enough and sometimes deteriorates accuracy; a judicial combination of local and global basis functions is vital. Further, adding the PoU-MoE trunk aids expressivity in every problem that involves steep spatial gradients in either the input or output functions. Finally, it appears that the full benefits of the PoU-MoE trunk are mainly achieved when the POD trunk is also used in the ensemble.

Given the generality of our work, there are numerous possible extensions along the lines of problem-dependent choices for the ensemble members. The PoU-MoE trunk merits further investigation. It is plausible that adding adaptivity to the PoU weight functions could improve its accuracy further, as could a spatially hierarchical formulation. Our work also paves the way for the use of other non-neural network basis functions within the ensemble DeepONet.

## Appendix A Ensemble Learning and Mixture-of-Experts (MoE)

The key idea behind ensemble learning is to combine a diverse set of learnable features from individual models into a single model (Polikar, ; Dong et al., ; Zhou, ; Dasarathy & Sheela, ). This technique has been used for both supervised and unsupervised feature selection in a variety of applications (Saeys et al., ; Li et al., ; Elghazel & Aussem, ; Abeel et al., ; Van Landeghem et al., ; Awada et al., ; Dittman et al., ; Piao et al., ). Guan et al. () draws an important distinction between using ensemble learning for feature selection and using feature selection for ensemble learning (where the former category is known to overcome the problem of local minima in machine learning). It is generally known that these methods are more stable than single base learners (Tuv, ; Guan et al., ). While ensemble methods have traditionally been studied with a statistical perspective (Krogh & Sollich, ; Chipman et al., ), we focus more on its feature selection capability, i.e. our work falls in the category of ensemble learning for feature selection. We use ensemble learning specifically to aggregate the global and local spatial features learned by the vanilla, POD, and PoU-MoE trunks into a single ensemble trunk.

MoE, first introduced in (Jacobs et al., ; Jordan & Jacobs, ), is a method in which an “expert” model focuses on learning from a subset of the training dataset. These models can be support vector machines (Lima et al., ; ; Collobert et al., ), Gaussian Processes (Ng & Deisenroth, ; Yuan & Neubauer, ; Gadd et al., ), and neural networks. Similar to ensemble learning, the MoE idea has also proven to be very successful in diverse ML applications (Masoudnia & Ebrahimpour, ; Yuksel et al., ; Chen et al., ). Most recently, MoE has also been used in physics-informed learning; Chalapathi et al. () uses MoE across the spatial domain with non-overlapping patches to decompose global physical hard constraints into multiple local constraints. Our PoU-MoE trunk uses a similar methodology where it has individual expert trunk networks on each patch in the domain, albeit with overlapping patches. However, instead of learning physical constraints, we let our experts learn spatially local basis functions (see Appendix G for further discussion on this spatial locality).

## Appendix B Speculation on an ensemble FNO

Here, we show one possible technique for incorporating the PoU-MoE localized bases into the FNO architecture, *i.e.*, we show how to create an ensemble FNO. FNOs consist of a *lifting* operator that lifts the input functions to multiple channels, a *projection* operator that undoes the lift, and intermediate layers (Fourier layers) consisting of kernel-based integral operators discretized by the fast Fourier transform (FFT); these integral operators are also typically augmented by pointwise convolution operations. Let $f_{t}$ denote the intermediate function at the $t^{th}$ Fourier layer. Then, the output $f_{t+1}$ of this layer (and the input to the next layer) is given by

$$ $\displaystyle f_{t+1}(y)$ $\displaystyle=\sigma\left(\int_{\Omega}\mathcal{K}(x,y)f_{t}(x)\;dx\;+\;Wf_{t} (y)\right),\;x\in\Omega,$ (15) $$

where $\sigma$ is an activation function applied pointwise, $\mathcal{K}$ is a matrix-valued kernel learned in Fourier space via the FFT, and $W$ is the aforementioned pointwise convolution (Li et al., ). Since FNOs use the FFT to compute the integral operator in (15), this effectively constitutes a projection of $f_{t}(x)$ onto a set of *global* Fourier modes (trigonometric polynomials or complex exponentials). One possible method for creating an ensemble FNO would involve modifying (15) to incorporate a set of localized basis functions using the PoU-MoE formulation as follows:

$$ $\displaystyle f_{t+1}(y)$ $\displaystyle=\sigma\left(\underbrace{\int_{\Omega}\mathcal{K}(x,y)f_{t}(x)\; dx}_{\text{Global basis}}\;+\underbrace{\sum_{k=1}^{P}w_{k}(y)\int_{\Omega_{k} }\mathcal{K}(x,y)\left.f_{t}(x)\right|_{\Omega_{k}}\;dx}_{\text{Localized basis}}\;+\;Wf_{t}(y)\right),$ (16) $$

where $P$ is the number of spatial patches (all of which are hypercubes). The PoU-MoE formulation now combines a set of *localized* integrals on each patch, each of which when computed by an FFT would constitute a projection of $f_{t}$ (restricted to $\Omega_{k}$) onto a local Fourier basis. This loosely resembles the Chebyshev polynomial PoU approximation introduced by Aiton & Driscoll ().

It is worth mentioning that this is one of many ways to combine different basis functions in FNOs. Another way is to introduce a set of local basis functions at the final projection operator that maps to the output function. The projection operator’s final layer can be enlarged to weight the additional basis functions, closely resembling how the branch weights the ensemble trunk in ensemble DeepONets. Similar extensions are possible for the GNO and even kernel/GP-based operator learning techniques.

## Appendix C Suboptimal ensemble trunk architectures

We document here our experience with other ensemble trunk architectures. We primarily made the following two other attempts:

A residual ensemble: Our first attempt was to combine the different trunk outputs using weighted residual connections with trainable weights, then activate the resulting output, then pass that activated output to a dense layer. For instance, given two trunks $\boldsymbol{\tau}_{1}$ and $\boldsymbol{\tau}_{2}$, this residual ensemble trunk would be given by

$$ $\displaystyle\hat{\boldsymbol{\tau}}_{\text{res}}=W\sigma\left(\tanh(w_{1}) \boldsymbol{\tau}_{1}+\tanh(w_{2})\boldsymbol{\tau}_{2}\right)+b,$ (17) $$

where $\sigma$ was some nonlinear activation, $W$ was some matrix of weights, and $b$ a bias. We also attempted using the sigmoid instead of the tanh. The major drawback of this architecture was that the output dimensions of the individual trunks had to match, *i.e.*, $p_{1}$ = $p_{2}$ to add the results (otherwise, some form of padding would be needed). We found that this architecture indeed outperformed the vanilla-DeepONet in some of our test cases, but required greater fine tuning of the output dimension $p$. In addition, we found that this residual ensemble failed to match the accuracy of our final ensemble architecture.
An activated ensemble: Our second attempt resembled our final architecture, but had an extra activation function and weights and biases. This activated ensemble trunk would be given by

$$ $\displaystyle\hat{\boldsymbol{\tau}}_{\text{act}}=W\sigma\left([\tau_{1},\tau_ {2}]\right)+b.$ (18) $$

This architecture allowed for different $p$ dimensions (columns) in $\tau_{1}$ and $\tau_{2}$. However, we found that this architecture did not perform well when the POD trunk was one of the constituents of the ensemble; this is likely because it is suboptimal to activate a POD trunk, which is already a data-dependent basis. There would also be no point in moving the activation function onto the other ensemble trunk constituents, since these are always activated if they are not POD trunks. Finally, though $W$ and $b$ allowed for a trainable combination rather than simple stacking, they did not offer greater expressivity over simply allowing a wider branch to combine these different trunks. We found that this architecture also underperformed our final reported architecture.

## Appendix D Proof of Universal Approximation Theorem for the PoU-MoE DeepONet

We have

$$ $\displaystyle\|\mathcal{G}(u)(y)-\mathcal{G}^{\dagger}(u)(y)\|_{\mathcal{V}}$ $\displaystyle=\left\|\mathcal{G}(u)(y)-\left\langle\boldsymbol{\beta}(u;\theta _{b}),\sum\limits_{j=1}^{P}w_{j}(y)\boldsymbol{\tau}_{j}(y;\theta_{\boldsymbol {\tau}_{j}})\right\rangle-b_{0}\right\|_{\mathcal{V}},$ $\displaystyle=\left\|\underbrace{\left(\sum\limits_{j=1}^{P}w_{j}(y)\right)}_{ =1}\mathcal{G}(u)(y)-\left\langle\boldsymbol{\beta}(u;\theta_{b}),\sum\limits_ {j=1}^{P}w_{j}(y)\boldsymbol{\tau}_{j}(y;\theta_{\boldsymbol{\tau}_{j}})\right \rangle\right.$ $\displaystyle\hskip 113.81102pt\left.-\underbrace{\left(\sum\limits_{j=1}^{P}w _{j}(y)\right)}_{=1}b_{0}\right\|_{\mathcal{V}},$ $\displaystyle=\left\|\sum\limits_{j=1}^{P}w_{j}(y)\left(G(u)(y)-\left\langle \boldsymbol{\beta}(u;\theta_{b}),\boldsymbol{\tau}_{j}(y;\theta_{\boldsymbol{ \tau}_{j}})\right\rangle-b_{0}\right)\right\|_{\mathcal{V}},$ $\displaystyle\leq\sum\limits_{j=1}^{P}w_{j}(y)\|\mathcal{G}(u)(y)-\left\langle \boldsymbol{\beta}(u;\theta_{b}),\boldsymbol{\tau}_{j}(y;\theta_{\boldsymbol{ \tau}_{j}})\right\rangle-b_{0}\|_{\mathcal{V}}.$ $$

Given a branch network $\boldsymbol{\beta}$ that can approximate functionals to arbitrary accuracy, the (generalized) universal approximation theorem for operators automatically implies that (Chen & Chen, ; Lu et al., ) a trunk network $\boldsymbol{\tau}_{j}$ (given sufficient capacity and proper training) can approximate the restriction of $\mathcal{G}$ to the support of $w_{i}(\mathbf{y})$ such that:

$$ $\|\mathcal{G}(u)(y)-\left\langle\boldsymbol{\beta}(u;\theta_{b}),\boldsymbol{ \tau}_{j}(y;\theta_{\boldsymbol{\tau}_{j}})\right\rangle-b_{0}\|_{\mathcal{V}} \leq\epsilon_{j},$ $$

for all $y$ in the support of $w_{j}$ and any $\epsilon_{j}>0$.
Setting $\epsilon_{j}=\epsilon$, $j=1,\ldots,P$, we obtain:

$$ $\displaystyle\|\mathcal{G}(u)(y)-\mathcal{G}^{\dagger}(u)(y)\|_{\mathcal{V}}$ $\displaystyle\leq\epsilon\underbrace{\sum\limits_{j=1}^{P}w_{i}(y)}_{=1},$ $\displaystyle\implies\|\mathcal{G}(u)(y)-\mathcal{G}^{\dagger}(u)(y)\|_{ \mathcal{V}}$ $\displaystyle\leq\epsilon.$ $$

where $\epsilon>0$ can be made arbitrarily small. This completes the proof.

## Appendix E Hyperparameters

### E.1 Network architecture

In this section, we describe the architecture details of branch and trunk networks. The architecture type, size, and activation functions are listed in Table 5. The CNN architecture consists of two five-filter convolutional layers with 64 and 128 channels respectively, followed by a linear layer with 128 nodes. Following Lu et al. (), the last layer in the branch network does not use an activation function, while the last layer in the trunk does. The individual PoU-MoE trunks in the ensemble models also use the same architecture as the vanilla trunk. We use the unstacked DeepONet with bias everywhere (except the POD-DeepONet which does not use a bias).

**Table 5: DeepONet network architectures across all models and problems. The CNN architecture is described in Appendix E.1.**
|  | Branch | Trunk | Activation function |
| --- | --- | --- | --- |
| Darcy flow | 3 layers, 128 nodes | 3 layers, 64 nodes | Leaky-ReLU |
| 2D Reaction-Diffusion | CNN | 3 layers, 128 nodes | ReLU |
| Cavity flow | CNN | [128, 128, 128, 100] | tanh |
| 3D Reaction-Diffusion | 3 layers, 128 nodes | 3 layers, 128 nodes | ReLU |

### E.2 Output dimension p 𝑝 p italic_p

We list the relevant DeepONet hyperparameters we use below. The $p$ ($p_{{}_{\text{POD}}}$ for POD) values are listed in Table 6 for all the DeepONets.

**Table 6: $p$ ($p_{{}_{\text{POD}}}$ for POD) values for the various DeepONet models. For $(P+1)$-vanilla DeepONet, the total number of basis functions is shown below. RD stands for reaction-diffusion.**
|  | Darcy flow | Cavity flow | 2D RD | 3D RD |
| --- | --- | --- | --- | --- |
| Vanilla | $100$ | $100$ | $100$ | $100$ |
| POD | $20$ | $6$ | $20$ | $20$ |
| Modified-POD | $20$ | $6$ | $20$ | $20$ |
| (Vanilla, POD) | $(100,20)$ | $(100,6)$ | $(100,20)$ | $(100,20)$ |
| $(P+1)$-Vanilla | $400$ | $500$ | $700$ | $900$ |
| Vanilla-PoU | $(100,100)$ | $(100,100)$ | $(100,100)$ | $(100,100)$ |
| POD-PoU | $(20,100)$ | $(6,100)$ | $(20,100)$ | $(20,100)$ |
| Vanilla-POD-PoU | $(100,20,100)$ | $(100,6,100)$ | $(100,20,100)$ | $(100,20,100)$ |

### E.3 Partitioning

The PoU-MoE trunk has certain hyperparameters that must be chosen. In our experiments, to maximize accuracy, we chose the patch size and the number of patches that produced the smallest possible patches and the smallest number of patches, while simultaneously seeking that the domain was covered and ensuring that the patches did not extend too far outside the domain boundary. Coincidentally, this strategy coincided with placing individual patches over regions of high spatial error in the vanilla-DeepONet solution (effectively, patches over "features of interest"). In the reaction-diffusion examples, even though we used uniform patch radii, we ensured that the patches did not overlap horizontally over line of the discontinuity. This choice combined with the use of ReLU activation ensured that we resolved that discontinuity better than vanilla-DeepONet; we believe this is one of the unique strengths of the PoU-MoE approach. Currently, we use the same trunk architectures on each patch as in the vanilla-DeepONet. In future work, adaptive patch selection strategies (such as making the patch centers and radii trainable or enforcing soft constraints on them as part of the loss function) can be used to automate determining patch placement and patch size. Furthermore, the patch shape can be changed depending on the problem domain; elongated/ellipsoidal patches can be used in narrower regions where spherical patches are not well suited.

## Appendix F Additional Results

We present additional results and figures in this section related to the problems in Section 3.

### F.1 2D Darcy flow

Figure: Figure 6: The 2D Darcy flow problem. (A) and (B) show example input and output functions respectively. (C) shows the three patches used for the PoU-MoE trunk. (D), (E), and (F) show the spatial mean squared error (MSE) for the vanilla, ensemble vanilla-POD-PoU, and ensemble POD-PoU DeepONets respectively.
Refer to caption: extracted/5894272/figures/Darcy_triangular/darcy_final_fig.png

The 2D Darcy flow problem models fluid flow within a porous media. The flow’s pressure field $u(y)$ and the boundary condition are given by

$$ $\displaystyle-\nabla\cdot\left(K(y)\ \nabla u(y)\right)$ $\displaystyle=f(y),\;y\in\Omega,$ (19) $\displaystyle u(y)$ $\displaystyle\sim\mathcal{GP}\left(0,\mathcal{K}(y_{1},y_{1}^{\prime})\right),$ (20) $$

where $K(y)$ is the permeability field, and $f(y)$ is the forcing term. The Dirichlet boundary condition was sampled from a zero-mean Gaussian process with a Gaussian kernel as the covariance function; the kernel length scale was $\sigma=0.2$. As in Lu et al. (), we learned the operator $\mathcal{G}:u(y)|_{\partial\Omega}\rightarrow u(y)|_{\Omega}$. We used the dataset provided in Lu et al. () which contains 1900 training and 100 test input and output function pairs. $\Omega$ was a triangular domain (shown in Figure 6). The permeability field and the forcing term were set to $K(y)=0.1$ and $f(y)=-1$. Example input and output functions, and the three patches for PoU trunks are shown in Figure 6. The partitioning always ensures that the regions with high spatial gradients are captured completely or near-completely by a patch.

We report the relative $\ell_{2}$ errors (as percentages) on the test dataset for the all the models in Table 1. The vanilla-POD-PoU ensemble was the most accurate model with a 4.5x error reduction over the vanilla-DeepONet and a 1.5x reduction over our POD-DeepOnet. The POD-PoU ensemble was second best with a 3.7x error reduction over the vanilla-DeepONet and a 1.5x reduction over the POD-DeepONet. The highly overparametrized $(P+1)$-vanilla model was less accurate than the standalone DeepONets. On this problem, overparametrization appeared to help only when spatial localization was also present; the biggest impact appeared to be from having both the right global and local information. The MSE errors as shown in Figure 6 corroborate these findings.

## Appendix G Evidence for Spatial Localization of the PoU-MoE Basis

Figure: Figure 7: Vanilla-DeepONet and PoU (from POD-PoU ensemble DeepONet) basis functions for the largest branch modes on the 2D reaction-diffusion problem.
Refer to caption: extracted/5894272/figures/diffrec/diffrec_basis_vanilla_pu.png

Here, we present further evidence showing that the PoU-MoE trunk learns spatially local features. In Figure 7, we show basis functions from the vanilla-DeepONet trunk and the PoU-MoE trunk of the POD-PoU ensemble DeepONet. Unlike the basis functions shown in Figure 2, these correspond to the largest branch coefficients in the respective models, i.e., the most “important” basis functions. Clearly, the PoU basis has a significantly higher spatial variation than the vanilla basis. We believe that this learned spatial locality helps the ensemble DeepONets with the PoU-MoE trunk achieve superior accuracy on problems with strong local features (such as those tested in this work).

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: See above.
soundness: 3
strengths: 1. Integrating the mixture-of-experts paradigm into operator learning enhances the model's capacity for effectively learning operators.
2. Building on the mixture-of-experts approach, the authors introduce a partition-of-utility strategy to encourage spatial locality and promote model sparsity.
3. A MoE model enhances the model's capacity by incorporating a diverse set of basis functions, which enables improved approximation accuracy across various operator learning tasks.
4. Comprehensive ablation study regarding key design factors of the MoE framework.
summary: This paper applies classical mixture-of-expert paradigm to learn mathematical operators. By incorporating different experts, and therefore enhancing basis representation, the resulting network has stronger learning power.
weaknesses: While the paper claims to present a novel framework combining classical expert neural networks, it largely repackages existing concepts rather than introducing groundbreaking ideas. The mixture-of-experts (MoE) paradigm is a well-established approach in machine learning, although innovative when applied in scientific ML, does not fundamentally transform the field. The combination of these elements lacks a compelling new problem formulation or a significant shift in methodology. The quality of the paper is undermined by several factors. There is a lack of critical evaluation of the scenarios where the QE-MoE approach may fail or perform suboptimally, especially concerning the boundary learning. While distributing data spatially to different expert models could potential relief the learning complexity compared to learning the data globally, it also introduces significant complexities depending on number of mixtures and data partitioning strategy. The experimental results do not convincingly demonstrate a substantial improvement over state-of-the-art methods but a comparison over several of its MoE variants.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Can you clearly explain what each of these baselines are in Table 1?
2. The choice of testing datasets are not so conventional, how does it perform on e.g. NS from Li et. 2020?
3. Why is (P + 1)-Vanilla MUCH worse than Vanilla (ours)? To my understanding, you simply add a layer to (P + 1)? Then with residual connection, it should be at least similar to Vanilla? Is there no residuals in your network?
soundness: 3
strengths: - The partition-of-unity mixture-of-experts (PoU-MoE) trunk introduces a novel approach that enhances spatial locality and promotes model sparsity.

- As reported, the ensemble DeepONets, particularly the POD-PoU variant, achieve substantial error reductions (2-4x) compared to standard DeepONets.

- Universal approximation capabilities are analyzed.
summary: This paper introduces an innovative enhancement for DeepONets, integrating multiple trunk networks to boost expressivity and generalization. It proposes a partition-of-unity mixture-of-experts (PoU-MoE) trunk structure to promote spatial locality, offering a refined approach to operator learning. Theoretical guarantees and extensive experiments across various PDE problems reveal that ensemble DeepONets, particularly the POD-PoU variant, achieve error reductions of 2-4x compared to standard DeepONets. This work offers valuable insights into effective trunk configurations and highlights a promising direction for advancing operator learning, albeit with increased computational requirements.
weaknesses: - The presentation could be improved. For instance, a clear description and detailed experimental setup for each baseline in Table 1 should be provided.
- The scope of this work appears somewhat limited, as it primarily focuses on testing enrichment strategies for basis functions within the specific context of operator learning. Applying these strategies to other popular frameworks is not trivial. Although the authors suggest that these methods might extend to FNO, sufficient details and evidence to substantiate this claim are lacking (Appendix B is not convincing enough).

- No comparison with other popular frameworks e.g. CNO [1], FNO, variants of FNO, etc.. Although this is also partially due to the fact that the scope of this work is limited to DeepONet. 

- While DeepONet is among the most well-known neural operators with clearly identifiable basis functions (the trunks), exploring additional works could help assess the generalizability of the proposed frameworks, such as [2] (treating the INRs/SIREN as producing basis functions as in DeepONet) .

[1] Convolutional Neural Operators for robust and accurate learning of PDEs; Bogdan Raonić, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, Emmanuel de Bézenac; 2023

[2] Operator Learning with Neural Fields: Tackling PDEs on General Geometries; Louis Serrano, Lise Le Boudec, Armand Kassaï Koupaï, Thomas X Wang, Yuan Yin, Jean-Noël Vittaut, Patrick Gallinari; 2024

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. How sensitive is the method's performance to the number of partitions? It would be valuable to see an ablation study on this hyperparameter.

2. Given that POD bases are computed from discretized output functions, how does the method handle evaluation at arbitrary points y not present in the training discretization?

3. How well does the method generalize to learning mappings from initial conditions to full spatiotemporal solutions, rather than just single-time predictions?
soundness: 3
strengths: 1. The paper presents its technical content in a clear, organized manner.

2. The integration of partition-of-unity principles into the DeepONet framework represents an interesting approach to incorporating spatial locality.

3. The work combines theoretical analysis (universal approximation theorems) with systematic empirical validation.

4. The proposed method demonstrates consistent performance improvements over standard DeepONets across multiple PDE examples.
summary: This paper presents a spatial mixture-of-experts (MoE) DeepONet architecture that utilizes partition-of-unity (PoU) approximation to combine expert networks across overlapping spatial patches. By integrating this localized approach with proper orthogonal decomposition (POD), the authors achieve 2-4x lower errors compared to standard DeepONets on several PDE problems. Besides, the authors also provide theoretical guarantees  for the proposed architecture through universal approximation theorems.
weaknesses: 1.  The experimental comparisons are primarily focused on DeepONet variants, omitting comparisons with other popular neural operators like FNO, which would provide broader context for the method's effectiveness.

2. The time-dependent PDE examples are restricted to single-step predictions (from one time point to another), leaving open questions about the method's capability to learn full temporal trajectories when time coordinates are included in the trunk network inputs.

3. The dependence on predefined partitions may limit the method's flexibility and generalizability, particularly for problems where optimal partition locations are not known a priori.