Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks

Madison Cooley; Ramansh Sharma; Shandian Zhe; Robert M. Kirby; Varun Shankar

cs.LGOct 3, 2026

Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks

Authors: Madison Cooley, Ramansh Sharma, Shandian Zhe, Robert M. Kirby, Varun Shankar

Organizations: Scientific Computing Institute, University of Utah, Salt Lake City, UT · Kahlert School of Computing, University of Utah, Salt Lake City, UT

Abstract

Operator learning methods such as DeepONets and FNOs often struggle with PDE families featuring sharp interfaces, heterogeneous coefficients, and localized multiscale structures. We introduce a partition-of-unity (POU) mixture-of-experts framework for localized operator learning, in which geometry-aware gating networks produce smooth spatial partitions which blend the contributions of local expert networks. Our main contribution is HiRefPOU, a residual-style hierarchical POU architecture for DeepONets that organizes localized representations through nested parent-child partitions while preserving global continuity. We also show that the same POU principle can be incorporated into Fourier Neural Operators to introduce spatial adaptivity without modifying the underlying spectral layers. On heterogeneous Darcy and reaction-diffusion benchmarks, HiRefPOU achieves substantially lower error than global DeepONet and static POU-MoE baselines, while the broader operator-learning experiments show that the benefits of localization depend on the PDE structure and the chosen neural-operator backbone. The learned partitions are interpretable and align with interfaces and regions of rapid solution variation. These results show that explicit geometric localization can improve both accuracy and interpretability in neural operator learning.

Figures & tables

Panel of Figure 1 — Figure 1: Residual-style hierarchical architecture used in HiRefPOU. The input $x$ is embedded and passed through a global backbone to obtain $u_{\mathrm{g}}(x)$ . A coarse Mixture-of-Experts (MoE) layer predicts $u_{c}(x)$ . A second MoE layer predicts the fine- level contributions $u_{f}(x)$ through nested parent–child partitioning. The final output is the sum of all three components, $u(x)=u_{\mathrm{g}}(x)+u_{c}(x)+u_{f}(x)$ .

Term	Meaning in this work
Adaptive	The parameters defining the POU gates, including centers, scales, neural gate parameters, and mixing coefficients, are learned jointly from the training objective rather than prescribed a priori. In operator-learning experiments, the gates depend only on evaluation coordinates and are shared across input samples; they are therefore not instance-conditioned.
Hierarchical	The POU weights are organized through nested parent–child relationships. A child partition is spatially weighted by its parent partition, producing a nested POU structure across levels.
Residual-style	Successive POU levels contribute additively to the model representation. The term describes the additive architecture and does not imply stage-wise training against detached residual targets.
Coarse / fine level	Labels identifying relative levels of the spatial partition hierarchy. They do not denote prescribed low- and high-frequency components or an orthogonal multiresolution decomposition.
Localization	Spatially varying POU weights assign different relative contributions to local expert representations across the evaluation domain.
Expert specialization	An expert may contribute more strongly in regions where its learned POU weight is large. Specialization is an empirical property of the learned representation rather than an imposed functional decomposition.

Table 2: Terminology used throughout this work. These definitions describe the specific architectural meaning of each term in HiRefPOU and are not intended as general definitions for all hierarchical or mixture-of-experts models.

Panel of Figure 2 — Figure 2: A POU-Mixture-of-Experts (MoE) layer predicts $u_{p}(x)$ .

Panel of Figure 3 — Figure 3: Residual-style hierarchical architecture used in HiRefPOU-DeepONet. The input $x$ is embedded and passed through a global backbone to obtain $t_{\mathrm{g}}(x)$ . A coarse Mixture-of-Experts (MoE) layer predicts $t_{c}(x)$ . A second MoE layer further produces fine-level contributions $t_{f}(x)$ through nested parent–child partitioning. The final output is the sum of all three components, $t(x)=t_{\mathrm{g}}(x)+t_{c}(x)+t_{f}(x)$ .

Figure 4: Example input–output pair from the lid-driven cavity flow experiment. (a) One-dimensional lid velocity profile

u(x)

prescribed along the top boundary of the cavity and used as the model input. (b) Corresponding two-dimensional cavity flow solution for a representative test instance.

Figure 5: Comparison of single-level (1L) and hierarchical two-level (2L) HiRefPOU DeepONet models for the cavity flow problem. Both models employ six total partitions initialized at identical center locations. The left and center panels show the mean absolute error (MAE) fields for the 1L ( $e_{1L}$ ) and 2L ( $e_{2L}$ ) architectures, respectively. The rightmost panel displays the local error difference $\delta e=e_{1L}-e_{2L}$ , where positive values indicate regions in which the hierarchical model improves upon the single-level model.

Figure 6: Learned partition-of-unity (POU) gating functions for the cavity flow problem using the single-level (1L) and hierarchical two-level (2L) HiRefPOU-DeepONet models. (A) Single-level (1L) model: spatially varying coarse gating functions associated with individual experts, together with the induced hard partition obtained via the pointwise ArgMax. The ArgMax maps are shown only to visualize the induced partition structure; model predictions use the continuous POU weights and dense evaluation of experts. (B) Hierarchical two-level (2L) model: learned coarse-level gating functions (top row) and their corresponding fine-level child partitions (bottom row). Rightmost panels show the induced hard partitions at each level. Colors indicate gating weight magnitude, with warmer colors corresponding to larger POU gating weights. The hierarchical model subdivides the learned coarse partitions into localized fine-level child regions while preserving the nested parent–child structure.

Figure 7: Example input–output pair for the two-dimensional reaction-diffusion problem. (a) Spatial input field $c(y,0)$ provided as input to the operator-learning model. (b) Corresponding solution field $c(y,T)$ at the final time.

Figure 8: Comparison of single-level (1L) and hierarchical two-level (2L) HiRefPOUDeepONet models. Both models employ six total partitions with identical initial center locations. The left and center panels show the mean absolute error (MAE) fields for the 1L ( $e_{1L}$ ) and 2L ( $e_{2L}$ ) architectures, respectively. The rightmost panel displays the local error difference $\delta e=e_{1L}-e_{2L}$ , where positive values indicate regions in which the hierarchical model improves upon the single-level model.

Figure 9: Comparison of single-level (1L) and hierarchical two-level (2L) HiRefPOU DeepONet learned partitions. Each panel shows the learned POU gating functions corresponding to different experts within the network. The top row illustrates the partitions produced by the single-level (1L) model, while the middle and bottom rows correspond to the hierarchical two-level (2L) model, where coarse partitions are further subdivided into locally refined subregions. The hierarchical model learns localized second-level child partitions near the coefficient discontinuity while preserving coherent coarse-level partitions across the domain.

Figure 10: Example input–output pair for the three-dimensional reaction-diffusion problem. (Left) Spatial input field provided as input to the operator-learning model. (Right) Corresponding solution field at the final time.

Figure 11: Comparison of single-level (1L) and hierarchical two-level (2L) HiRefPOU DeepONet learned partitions. Each panel shows the learned POU gating functions corresponding to different experts within the network. The top row illustrates the partitions produced by the single-level (1L) model, while the middle and bottom rows correspond to the hierarchical two-level (2L) model, where coarse partitions are further subdivided into locally refined subregions. The hierarchical model learns localized second-level child partitions aligned with regions of strong diffusion variation while preserving coherent coarse-level partitions across the ball domain.

Panel of Figure 12 — Figure 12: (1D Burgers Problem.) Learned partition-of-unity (POU) gates for HiRefPOU-DeepONet models. (Top left) Continuous coarse-level gating functions $\{\phi_{j}(x)\}_{j=1}^{3}$ learned by the HiRefPOU-1L model. Colored markers indicate the learned center locations associated with each partition. (Top right) Induced hard partition obtained via $\arg\max_{j}\phi_{j}(x)$ , illustrating the discrete domain decomposition implied by the learned soft POU representation. (Bottom left) Continuous coarse- and fine-level gating functions $\{\phi^{c}_{j}(x)\}_{j=1}^{1}$ and $\{\phi^{f}_{k}(x)\}_{k=1}^{2}$ learned by the HiRefPOU-2L model. (Bottom right) Induced hard fine-level partition obtained via $\arg\max_{k}\phi^{f}_{k}(x)$ .

Figure 13: (1D Burgers Problem.) Operator learning predictions at final time $t=1$ . (Top row) Reference solution $u(x,1)$ , HiRefPOU-1L prediction, and HiRefPOU-2L prediction, evaluated over the full test set. Each heatmap displays solution values as a function of spatial coordinate $x\in[0,1]$ (horizontal axis) and test sample index (vertical axis). (Bottom row) Point-wise absolute error fields $\lvert u_{\mathrm{ref}}-u_{\mathrm{pred}}\rvert$ for the 1L and 2L models, along with their difference. Positive values indicate locations where the 2L model achieves lower error than the 1L model.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Benchmark	Variant	Experts by level	Parent–child allocation	Leaf partitions
Cavity	1L	$N_{1}=4$	–	4
	2L	$(N_{1},N_{2})=(2,2)$	$L_{1}\!\to\!L_{2}:$ $[0,2]$	3
	3L	$(N_{1},N_{2},N_{3})=(2,2,4)$	$L_{1}\!\to\!L_{2}:\ [0,2]$ , $L_{2}\!\to\!L_{3}:\ [2,2]$	4
2D RD	1L	$N_{1}=6$	–	6
	2L	$(N_{1},N_{2})=(2,4)$	$L_{1}\!\to\!L_{2}:$ $[2,2]$	4
	3L	$(N_{1},N_{2},N_{3})=(2,4,8)$	$L_{1}\!\to\!L_{2}:\ [2,2]$ , $L_{2}\!\to\!L_{3}:\ [2]_{\times 4}$	8

Appendix

Table 11: HiRefPOU hierarchy configurations used for each benchmark.

N_{\ell}

denotes the number of experts at level

\ell

. Parent–child allocations list the number of children assigned to each parent at successive levels. The notation

[2]_{\times k}

denotes a length-

k

list in which every parent has two children.

Benchmark	POD rank	Basis samples	Centering / normalization
Cavity	$6$ $(3+3)$	100 training outputs	Pointwise training mean/std; separate basis per component
2D reaction–diffusion	20	1,000 training outputs	Pointwise training mean/std
3D reaction–diffusion	20	1,000 training outputs	Pointwise training mean/std
1D Burgers	20	1,000 training outputs	Pointwise training mean/std
Full-Field Burgers	20	1,000 training outputs	Pointwise training mean/std
Kuramoto–Sivashinsky	20	1,000 training outputs	Global training mean/std, followed by pointwise training mean/std

Appendix

Table 14: POD basis construction used in the DeepONet experiments. All POD statistics and basis vectors are computed exclusively from the training outputs. For cavity flow, separate rank-3 bases are constructed for the two velocity components, giving a total POD rank of 6. The POD ranks were fixed a priori and were not selected using test-set performance.

Figure 14: (Top) the true solution and fluxes. (Bottom) Predicted solutions, fluxes, and absolute errors for the Darcy flow benchmark using HiRefPOU-1L, APINN, and PirateNet. Left bottom figures show the predictions while the right bottom figures show the absolute errors.

Model	Main / expert width	Main / expert depth	Gate width/depth	POU experts by level	Trainable parameters
APINN	32	4	$32/3$	2	11,940
PirateNet	32	3	–	–	11,684
HiRefPOU-1L	32	3	$16/1$	2	11,935
HiRefPOU-2L	32	3	$16/1$	$2+3$	22,162

Appendix

Table 19: Architecture specifications for the Darcy PINN experiments. Width and depth refer to the primary network for PirateNet and to the expert networks for the POU-based models. Gate specifications are reported as hidden width/depth. The HiRefPOU-2L entry lists coarse and fine experts, respectively. Parameter counts include all trainable network, embedding, expert, and gating parameters.

Gate	Test rel. $\ell_{2}$	Std.	Test rel. $\ell_{\infty}$	Std.	Epoch time
RBF gate	$0.177134$	$0.004654$	$0.457744$	$0.009059$	$0.000542$
Box gate	$0.176758$	$0.004868$	$0.457380$	$0.007711$	$0.000606$
Hybrid gate	$\mathbf{0.172893}$	$0.001394$	$\mathbf{0.454045}$	$0.005357$	$0.000629$

Appendix

Table 20: Standard 1D Burgers benchmark: gate ablation for the HiRefPOU-1L DeepONet with

M=3

adaptive partitions. We compare RBF, box, and hybrid gate formulations. Reported values are mean

\pm

standard deviation over random seeds. Epoch time denotes the average wall-clock time per training epoch.

Configuration	Test rel. $\ell_{2}$	Std.	Test rel. $\ell_{\infty}$	Std.	Epoch time
Hybrid gate, $M=2$	$\mathbf{0.172622}$	$0.002183$	$\mathbf{0.452593}$	$0.002469$	$0.000547$
Hybrid gate, $M=3$	$0.176576$	$0.001287$	$0.459151$	$0.006000$	$0.000629$
Hybrid gate, $M=4$	$0.174892$	$0.002286$	$0.458236$	$0.001768$	$0.000692$
Hybrid gate, $M=6$	$0.173398$	$0.002347$	$0.453612$	$0.006004$	$0.000833$

Appendix

Table 21: Standard 1D Burgers benchmark: sensitivity of the HiRefPOU-1L DeepONet to the number of adaptive partitions

M

when using the hybrid gate. Reported values are mean

\pm

standard deviation over random seeds. Epoch time denotes the average wall-clock time per training epoch.

PDE Operator Learning

Partial differential equation (PDE) operator learning builds models that map inputs such as coefficients and initial conditions to solution fields.

See all 225 papers

Linear-Nonlinear Fusion Neural Operator for Partial Differential Equations
Mar 25, 2026Heng Wu, Junjie Wang, Benzhuo LuPDE Operator Learning Neural Operators
SKLMS, ICMSEC, NCMIS, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China · School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China
Smooth Piecewise Cutting for Neural Operator to Handle Discontinuities and Sharp Transitions
May 19, 2026Ha Dang, Sebastian Schmidt, Juergen HesserPDE Operator Learning Deep Operator Networks
Mannheim Institute for Intelligent Systems in Medicine, Heidelberg University · Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University · Heidelberg Institute for Theoretical Studies (HITS), Heidelberg University +2
Fractional is Better: Learnable Derivative Orders in Neural Operator Learning
Oct 14, 2026Fares B. Mehouachi, Saif Eddin JabariPDE Operator Learning Neural Operators
1New York University Abu Dhabi, Abu Dhabi, UAE · 2New York University, New York, NY, USA

Mar 25, 2026J/K move · Enter open · S save

Linear-Nonlinear Fusion Neural Operator for Partial Differential Equations

Heng Wu, Junjie Wang, Benzhuo Lu

SKLMS, ICMSEC, NCMIS, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China · School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China

PDE Operator Learning Neural Operators PDE Surrogate Modeling

Open paper AskPDF arXiv

Neural operator learning directly constructs the mapping relationship from the equation parameter space to the solution space, enabling efficient direct inference in practical applications without the need for repeated solution of partial differential equations (PDEs) -- an advantage that is difficult to achieve with traditional numerical methods. In this work, we find that explicitly decoupling linear and nonlinear effects within such operator mappings leads to improved learning efficiency. This yields a novel network structure, namely the Linear-Nonlinear Fusion Neural Operator (LNF-NO), which models operator mappings via the multiplicative fusion of a linear component and a nonlinear component, thus achieving a lightweight and interpretable representation. This linear-nonlinear decoupling enables efficient capture of complex solution features at the operator level while maintaining stability and generality. LNF-NO naturally supports multiple functional inputs and is applicable to both regular grids and irregular geometries. Across a diverse suite of PDE operator-learning benchmarks, including nonlinear Poisson-Boltzmann equations and multi-physics coupled systems, LNF-NO is typically substantially faster to train than several representative neural operator baselines, while achieving comparable or improved accuracy across most tested cases. On the tested 3D Poisson-Boltzmann case, LNF-NO achieves strong accuracy while requiring substantially less training time than the three-dimensional Fourier Neural Operator and Transolver baselines.

Abstract

PDE Operator Learning

Explore similar work

Linear-Nonlinear Fusion Neural Operator for Partial Differential Equations

Smooth Piecewise Cutting for Neural Operator to Handle Discontinuities and Sharp Transitions

Fractional is Better: Learnable Derivative Orders in Neural Operator Learning