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

TMLR Paper8807 Authors

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 and route computation to specialized local experts.
Our main contribution is HiRefPOU, a hierarchical residual POU architecture for DeepONets that enables coarse-to-fine refinement 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.

Submission Type: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Yinpeng_Dong2

Submission Number: 8807