Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks
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
| 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. |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Benchmark | Variant | Experts by level | Parent–child allocation | Leaf partitions |
| Cavity | 1L | – | 4 | |
| 2L | 3 | |||
| 3L | , | 4 | ||
| 2D RD | 1L | – | 6 | |
| 2L | 4 | |||
| 3L | , | 8 |
| Benchmark | POD rank | Basis samples | Centering / normalization |
| Cavity | 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 |
| Model | Main / expert width | Main / expert depth | Gate width/depth | POU experts by level | Trainable parameters |
| APINN | 32 | 4 | 2 | 11,940 | |
| PirateNet | 32 | 3 | – | – | 11,684 |
| HiRefPOU-1L | 32 | 3 | 2 | 11,935 | |
| HiRefPOU-2L | 32 | 3 | 22,162 |
| Gate | Test rel. | Std. | Test rel. | Std. | Epoch time |
| RBF gate | |||||
| Box gate | |||||
| Hybrid gate |
| Configuration | Test rel. | Std. | Test rel. | Std. | Epoch time |
| Hybrid gate, | |||||
| Hybrid gate, | |||||
| Hybrid gate, | |||||
| Hybrid gate, |
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