Localized neural networks improve accuracy in complex equation models
Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks
Machine Learning
Summary
Many AI methods for solving complex equations struggle when the problems have sharp boundaries or mixed materials. The paper introduces a way to break down the problem into smaller parts using special networks that focus on local areas, blending their outputs smoothly. The authors show that this approach improves accuracy and helps the AI focus where details change quickly. This method works better for some types of problems and makes the AI’s reasoning easier to understand.
What this means in practice
- •For simulation engineers: Improve simulation accuracy for problems involving materials with sharp boundaries using locally focused neural operator models.
- •For geophysics modelers: Enhance porous media flow models by adapting neural networks to spatially heterogeneous subsurface properties.
Authors
Madison Cooley, Ramansh Sharma, Shandian Zhe, Robert M. Kirby, Varun Shankar
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.