Physics aware AI improves solving complex changing problems in math models
Latent-MoE: Domain-Aware Mixture-of-Experts for PDEs with Multi-Regime Physics
Machine Learning
Summary
Physics-informed neural networks often have trouble solving equations when the rules change in different parts of the problem. The authors found that standard networks mistakenly let faraway parts strongly influence each other, causing confusion during training. They introduce Latent-MoE, a method that divides the problem into experts focused on local regions while still sharing some common information. This approach keeps learning focused where it matters and shares knowledge across the whole problem, leading to much better performance on difficult cases where physics changes over time. It also reduces conflicts that make training hard.
Physics-informed neural networksPartial differential equationsNeural tangent kernelMixture-of-expertsDomain decompositionGradient conflictMulti-regime physicsPhysics simulationMachine learningLocalization
Authors
Hanwen Wang, Paris Perdikaris
Abstract
Physics-informed neural networks (PINNs) struggle on PDEs whose governing physics varies across the domain. We trace this to a structural property of standard coordinate networks: their neural tangent kernel (NTK) is translation-variant and lets training points of large coordinate magnitude disproportionately influence predictions elsewhere, producing long-range coupling and gradient conflict during training. We show analytically and empirically that mixture-of-experts (MoE) architectures with centered, compact-support routers yield a uniformly banded NTK whose kernel-regression weights decay exponentially with distance, localizing the learning. Building on this, we propose \emph{Latent-MoE}, which interleaves domain-aware MoE blocks within a shared backbone. Unlike FB-PINNs or X-PINNs, which rigidly partition both the domain and the parameters so that the parameters on different subdomains are updated independently, Latent-MoE is designed to preserve the localization benefit of domain-aware routing while allowing capacity to flow across regions through the shared backbone. On standard homogeneous-physics benchmarks Latent-MoE is competitive with established baselines; on benchmarks with multi-stage time-variable physics, where global models and rigid domain decompositions both fall into spurious solutions, it improves over them by more than an order of magnitude, with markedly reduced gradient conflict during training.