PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent-MoE: Domain-Aware Mixture-of-Experts for PDEs with Multi-Regime Physics

Hanwen Wang, Paris Perdikaris

arXiv:2609.07814v1cs.LGphysics.comp-ph

TL;DR

PINNs can struggle when PDE physics varies across a domain, and the paper attributes this to long-range influence from translation-variant coordinate-network NTKs. It proposes Latent-MoE, combining domain-aware routing with a shared backbone; it matches baselines on homogeneous benchmarks and reports more than an order of magnitude lower relative L2 error on multi-stage time-variable physics.

  • Problem

    Translation-variant NTKs in standard coordinate networks let large-magnitude training coordinates disproportionately influence distant predictions, creating long-range coupling and gradient conflict when physics varies across a domain.

    Source: §1 · pp. 1–2

  • Method

    Latent-MoE interleaves domain-aware MoE blocks with fixed, compact-support routers in a shared backbone, using localization motivated by the banded NTK result.

    Source: §1 · p. 2, §4.2 · p. 5, §4.3 · p. 5

  • Results

    On multi-stage time-variable physics, Latent-MoE achieves more than an order of magnitude lower relative L2 error than baselines, while matching established baselines on homogeneous-physics benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported advantage is regime-specific: localization does not help on homogeneous benchmarks, while Latent-MoE preserves dynamics with reduced gradient conflict when physics changes across the domain.

    Source: §6 · p. 10

    Quantitative claims rest on three synthetic heterogeneous-physics benchmarks and homogeneous controls; evaluation on real engineering or geophysical problems remains a next step.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

PINNs embed PDE equations and available boundary or initial data in a neural-network training objective, but coordinate-network optimization pathologies and limited theory on domain partitioning leave open why partitioning helps. The paper links these problems to NTK structure and introduces Latent-MoE, reporting stronger results on multi-stage, time-variable physics than on homogeneous benchmarks.

1 Introduction

  • Prior domain-decomposition methods partition the spatio-temporal domain and assign subnetworks to regions, but why partitioning helps and which partition is best have received limited theoretical attention.The introduction frames the paper's NTK analysis as addressing that gap.Source: §1 · p. 2
  • On multi-stage time-variable physics, Latent-MoE achieves more than an order of magnitude lower relative L2 error than baselines, while matching established baselines on homogeneous-physics benchmarks.The reported multi-stage setting is one where global models collapse to spurious solutions and rigid decompositions fail to capture stage-specific dynamics.Source: §1 · p. 2
  • Standard ReLU coordinate networks have translation-variant NTKs, so large-magnitude training coordinates disproportionately influence distant predictions and can make gradients from different dynamics compete.The paper identifies this long-range coupling as a structural obstacle to PINN training when equations or forcing change across the domain.Source: §1 · p. 2
  • The paper proves that centered, compact-support domain-aware MoE routers yield a banded NTK with exponentially distance-decaying kernel-regression weights, localizing learning.Latent-MoE interleaves these experts within a shared backbone rather than partitioning the entire model.Source: §1 · p. 2

2 Preliminaries

The preliminaries describe PINNs as neural surrogates trained to satisfy PDE and boundary or initial conditions through residual losses. They also introduce the empirical NTK and its role in describing how training-data influence propagates to predictions.

2 Preliminaries

  • PINNs train a neural surrogate uθ to satisfy a PDE operator equation on the domain and a boundary-operator equation on the boundary or initial-data set.The setup specifies a PDE operator N, forcing f, boundary operator B, and data g on ∂Ω.Source: §2 · p. 2
  • The training objective minimizes a weighted sum of residual losses, with weights that may be fixed or adapted to stabilize multi-objective optimization.Higher-order automatic differentiation for the PDE residual makes the loss landscape stiff and motivates architectural and optimization methods.Source: §2 · p. 2
  • The empirical NTK is the Gram matrix of parameter-Jacobians; under appropriate scaling it converges to a deterministic kernel that remains constant during gradient-descent training.The associated kernel-regression weights characterize how training-data influence propagates across the input domain.Source: §2 · p. 2

3 Related Work

Related work spans rigid domain-decomposition PINNs, mixture-of-experts architectures, NTK analyses, and complementary training strategies. Latent-MoE combines interleaved experts and a shared backbone with fixed routers informed by PDE geometry.

3 Related Work

  • X-PINNs enforce interface agreement, while FB-PINNs and multilevel FB-PINNs [13] blend independently parameterized experts; AB-PINNs adapt partitions to the residual landscape.These end-to-end decompositions localize capacity but prevent its reallocation across regions with differing physics difficulty.Source: §3 · p. 3
  • Modern mixture-of-experts designs interleave sparse expert layers in a shared backbone with learnable, input-agnostic routers [21]; Latent-MoE instead uses fixed routers informed by PDE geometry.The paper characterizes domain-decomposition PINNs as a special case with a single output MoE layer and coordinate-based fixed routing.Source: §3 · p. 3
  • Prior NTK studies analyze PINN multi-loss gradient pathologies and Fourier-feature spectral bias, alongside gradient alignment [25] and spurious solutions under stiff physics.The paper situates its work within NTK-based analysis of learning dynamics in wide networks [22] [18].Source: §3 · p. 3
  • Complementary PINN strategies include PirateNet [26], causal training, sinusoidal activations [9], and Fourier features [10]; Latent-MoE uses Fourier features and causal weighting in its loss.These methods are described as largely orthogonal to changing domain structure.Source: §3 · p. 3

4 Method

The method uses NTK analysis to motivate localized learning with centered, compact-support routing, then embeds domain-aware experts within a shared backbone. Latent-MoE combines subdomain-bound experts with shared capacity across regions.

4.1 NTK analysis: long-range coupling and its localization

  • Standard ReLU coordinate networks couple distant training points, while centered compact-support experts yield banded NTKs and exponentially decaying kernel-regression weights.Empirically, centering bounds effective neighborhood widths uniformly across the domain, whereas MoE alone does not; smaller α-effective widths indicate more localized learning.Source: §4.1 · pp. 3–4

4.2 From localization to architecture

  • Latent-MoE keeps experts subdomain-bound while its shared backbone lets representational capacity flow across regions, unlike end-to-end domain partitioning.Rigid partitions fix per-subdomain capacity, which must be sized for the worst-case region; the authors hypothesize that capacity-limited early experts can propagate errors across boundaries.Source: §4.2 · pp. 4–5

4.3 The Latent-MoE architecture

  • Latent-MoE alternates shared residual and sparse domain-aware MoE blocks; fixed compact-support routers activate covering experts, and centered shared encoding outperforms independent or uncentered alternatives.The encoder centers each expert’s input on its subdomain centroid; independent random Fourier projections can fragment the latent space and disrupt consistency for the shared backbone.Source: §4.3 · pp. 5–6, Figure 2 · p. 5

5 Experiments

The experiments compare Latent-MoE with parameter-matched baselines on homogeneous and heterogeneous PDE benchmarks, including multi-stage time-variable physics and a spatially layered wave equation. They assess accuracy, robustness to transitions and router alignment, and gradient conflict.

5 Experiments

  • On homogeneous benchmarks, all architectures perform comparably; on multi-stage time-variable benchmarks, Latent-MoE beats every baseline by more than an order of magnitude, while domain-decomposition models lead on layered wave.The comparisons use median relative L2 errors; models have comparable parameter counts within each benchmark.Source: §5 · p. 6, Table 1 · p. 7

5.1 Advection–diffusion with time-variable coefficients

  • 1.0 × 10^-2 median relative L2 error: Latent-MoE beats every baseline by more than an order of magnitude on advection–diffusion and tracks the solution through all three stages.Table 1 reports medians across five random seeds; ResNet collapses after the diffusion coefficient grows.Source: Figure 3 · p. 7, §5.1 · p. 8

5.2 Damped wave equation with staged forcing

  • ResNet, PirateNet, and FB-PINNs all reach near-trivial solutions with relative L2 errors above 0.7; Latent-MoE recovers the damped-wave resonant frequency signature.The benchmark adds staged forcing to a time-variable operator.Source: Figure 3 · p. 7, §5.2 · p. 8

5.3 Layered wave equation

  • Latent-MoE does not require exact alignment with the four physical layers, and a plateau ratio of 1 outperforms the native ratio of 1.5 for both MoE models.The benchmark also tests shifted expert subdomains and overlap-ratio changes.Source: §5.3 · p. 8

5.4 Robustness to stage-transition placement

  • Across shifted transition times, Latent-MoE remains in the 10^-3–10^-2 relative L2 error range while ResNet stays above 0.3.The temporal subdomains remain fixed as the first transition time varies from 0.1 to 0.7.Source: §5.4 · pp. 8–9, Table 3 · p. 9

5.5 Gradient conflict and the localization hypothesis

  • Latent-MoE maintains lower gradient conflict than dense baselines on both benchmarks; replacing centered expert encoders with a global encoder raises conflict toward the dense-baseline range.FB-PINNs can show mechanically inflated conflict because disjoint expert parameters give per-chunk Jacobians nearly disjoint support.Source: Figure 4 · p. 9, §5.5 · p. 10

6 Discussion

Latent-MoE combines domain-aware experts with a shared backbone, pairing localized learning with capacity shared across regions. Its NTK analysis motivates this design, and its benefits are strongest when physics varies across space or time.

6 Discussion

  • Latent-MoE interleaves domain-aware MoE blocks within a shared backbone, unlike rigidly partitioned approaches whose subdomain parameters update independently.The shared backbone preserves flexibility to allocate capacity across regions.Source: §6 · p. 10
  • Latent-MoE matches established baselines on homogeneous-physics benchmarks, while its advantage emerges when governing equations or forcing vary across the domain.Global baselines collapse to spurious solutions in the heterogeneous settings discussed; reduced gradient conflict accompanies Latent-MoE's preserved dynamics.Source: §6 · p. 10
  • The evaluation uses three synthetic heterogeneous-physics benchmarks, and validation on real engineering or geophysical problems remains future work.The analysis also assumes infinite width and squared-error loss; a finite-width characterization under residual loss remains an open gap.Source: §6 · p. 10
  • The paper identifies well-defined physics that varies substantially across time or space as Latent-MoE's natural target, including layered media, reactive flows, and time-variable forcing.The broader outlook argues that routing informed by physical structure can repurpose expert architectures for scientific computing.Source: §6 · p. 10

A.1 Preliminaries

  • The standard coordinate-network NTK is translation-variant, with scale growing quadratically with input magnitude; centered compact-support experts instead yield a banded NTK and exponentially decaying regression weights.For small bias and input-dominated conditions, the linearized NTK grows with input magnitude, potentially inducing spurious temporal correlations.Source: §A.1 · pp. 17–18

A.2 Proof of Theorem 1

Compact-support routers make the composite MoE NTK banded, which supports exponentially localized kernel-regression influence. Centering also produces more regular learning behavior in the visualized comparisons.

A.2 Proof of Theorem 1

  • If training points are farther apart than router width w, the composite MoE NTK entry is zero, making the matrix wΔx-banded.The result uses evenly spaced compact-support routers; bounded positive-definite banded matrices have inverses with exponentially decaying entries.Source: §A.2 · p. 19

A.3 2D effective neighborhood width

  • The temporal α-effective neighborhood width is defined for a 2D space-time model partitioned only along time.The input is represented as x = [t, s], and the regression-weight matrix has shape Nt × Ns.Source: §A.3 · p. 19

A.4 NTK diagnosis

  • Centering the experts yields more localized, regular NTK and regression-weight structures than non-centered routing, with increasing expert counts progressively uniformizing learning behavior.The comparison includes both NTKs and absolute kernel-regression weights across expert counts.Source: §A.4 · p. 19, Figure 6 · p. 20, Figure 5 · p. 20

B Architecture Details

The architecture uses overlapping bump routers for soft domain partitioning and periodic or Fourier-based coordinate features for its encoder.

B.1 Domain-aware Routers

  • Overlapping bump-bridge functions softly partition the domain, allowing neighboring experts to blend smoothly.The router construction uses C1 bridges, while C2 bridges make both first and second derivatives vanish at bridge endpoints and are used for second-order-in-time benchmarks.Source: Figure 7 · p. 21, §B.1 · p. 21

B.2 Encoder

  • For periodic spatial domains, the encoder represents space with sine and cosine features at selected frequencies and appends a temporal feature.The periodic embedding is followed by a Gaussian random projection and Fourier embedding; Burgers with Dirichlet boundaries omits the periodic embedding.Source: §B.2 · p. 21
  • The random Fourier-feature scales are 4 for both space and time, except for KdV, where both are 8.The projection matrix is learnable, followed by a trainable linear layer mapping the embedding to backbone width.Source: §B.2 · p. 21

C Experimental Details

The appendix documents shared experiment settings, optimization procedures, and the constructed heterogeneous benchmarks. It also describes how baseline errors and model statistics are visualized and reported.

C Experimental Details

  • Networks have width 128 and four residual blocks; in Latent-MoE, two blocks are MoE, while corresponding non-MoE blocks are widened to keep total parameter counts comparable, aside from small bias-term differences.Experiments use JAX and Flax in Python 3.11, and all experiments use an unofficial JAX implementation of SOAP.Source: C Experimental Details · p. 22

C.1 Huber loss function; C.2 Causality Enforcement; C.3 Loss balancing

  • The training procedures use Huber loss, dynamic causal weighting for non-standard benchmarks, and gradient-based balancing of PDE, initial, and boundary losses.Eight temporal chunks are used; standard JAXPI benchmarks retain unnormalized causal weighting.Source: §C.1 · p. 23, §C.2 · p. 23

C.4 Gradient Conflict Index

  • The gradient-conflict index is presented as a practical proxy for negative correlation that tracks training dynamics and loss evolution.Its NTK interpretation is asymptotic in network width and sampled points per chunk.Source: §C.4 · p. 24

C.5 Advection-Diffusion

  • The advection-diffusion benchmark combines oscillating, drifting advection speed with a diffusion coefficient taking three stage-specific values joined by smooth ramps.The ramps occur at t = 0.5 and t = 4/3.Source: §C.5 · p. 24, Figure 8 · p. 25

C.6 Damped Wave Equation with forcing

  • The damped-wave benchmark injects forcing during its middle stage, after damping, with free evolution before damping and after forcing ends.The force ramps up linearly over the first third of its stage, then remains constant to the stage boundary.Source: §C.6 · p. 24, Figure 9 · p. 25

C.7 Layered Wave Equation

  • The layered-wave benchmark uses four periodic spatial layers with distinct wave speeds and a layer-dependent initial frequency, with zero initial velocity.Its four spatial experts are tiled uniformly, and router boundaries can be shifted to test sensitivity to misalignment; wrapping routers preserve periodicity.Source: §C.7 · pp. 24–25, Figure 11 · p. 26, Figure 10 · p. 26

C.8 Baseline model visualizations

  • Appendix visualizations show ground truths and median prediction errors for FB-PINNs and PirateNet across the three heterogeneous benchmarks.The prediction-error figure arranges benchmarks by rows: advection-diffusion, damped wave, then layered wave.Source: §C.8 · p. 26, Figure 12 · p. 26, Figure 13 · p. 27

C.9 Detailed model statistics

  • Table 5 reports relative L2 error mean, standard deviation, and median across five seeds for six benchmarks.The median is included to mitigate outlier influence, and the best performance for each benchmark is highlighted.Source: §C.9 · pp. 26–27, Table 5 · p. 28

D Ablation study

The ablations examine encoder design, expert overlap, and the relative causality coefficient using a greedy one-hyperparameter-at-a-time process. The reported best choices favor a centered, shared encoder, a plateau ratio of 1, and benchmark-specific causality coefficients.

D Ablation study

  • The ablation process greedily changes one hyperparameter at a time, carries each round’s best choice forward, and uses one fixed seed per setting.Ablations cover both advection-diffusion and damped-wave benchmarks.Source: D Ablation study · p. 29

D.1 Encoder

  • The centered, weight-shared MoE encoder performs best on both benchmarks, outperforming centered independent and uncentered encoders.The ablations use a regular encoder based on periodic embedding followed by random Fourier features.Source: §D.1 · p. 29, Table 6 · p. 29

D.2 Bump function plateau ratio

  • A plateau ratio of 1 is best, setting the plateau width equal to the bridge width and controlling overlap between neighboring experts.A ratio of 0 gives pure bridges with maximal overlap; as the ratio grows toward infinity, the partition becomes fully disjoint.Source: §D.2 · p. 29, Table 7 · p. 30

D.3 Relative causality coefficient ϵ

  • The best relative causality coefficient is 4 for advection-diffusion and 2 for damped wave, with performance not very sensitive within a reasonable range.Source: §D.3 · p. 30, Table 8 · p. 30

E Scaling experiments

Scaling experiments compare parameter-matched Latent-MoE and ResNet models on advection-diffusion and damped-wave benchmarks. Latent-MoE benefits clearly from adding experts, while widening ResNet yields only marginal improvement.

E Scaling experiments

  • Latent-MoE benefits clearly from scaling the number of experts, whereas widening parameter-matched ResNet feedforward layers yields only marginal improvement.Figure 14 reports scaling diagnostics for advection-diffusion and damped wave using the same protocol as Section 5; each point is the median of three seeds.Source: E Scaling experiments · p. 30, Figure 14 · p. 30

F Variable stage transition experiments

Variable stage transitions challenge all compared architectures, but Latent-MoE’s errors remain more localized around transition boundaries. Its error is also more consistent across evaluation horizons that extend into the multi-stage regime.

F Variable stage transition experiments

  • Latent-MoE remains more robust to sharp stage transitions, with errors more localized around transition boundaries than the baselines.The windowed relative L2-error heatmaps compare ResNet, PirateNet, FB-PINNs, and Latent-MoE across both benchmarks; each window has size Δt = 1/16.Source: F Variable stage transition experiments · pp. 30–31, Figure 15 · p. 31
  • Latent-MoE maintains a more consistent median relative L2 error across tmax, while ResNet and PirateNet errors rise sharply as horizons enter the multi-stage regime.ResNet and PirateNet can be competitive or better at short horizons before the first stage transition.Source: Figure 16 · p. 31

G Variable time horizon experiments

Across variable time horizons, Latent-MoE’s error remains more consistent as evaluation extends into multi-stage dynamics, unlike the sharp error increases of ResNet and PirateNet. The section also reports seed sensitivity, computational costs, reproducibility limits, and broader-impact cautions.

G Variable time horizon experiments

  • Latent-MoE maintains a substantially more consistent error profile across tmax, while ResNet and PirateNet errors rise rapidly once evaluation extends into the multi-stage regime.At short horizons before the first transition, ResNet and PirateNet can be competitive or even achieve lower errors.Source: G Variable time horizon experiments · pp. 31–32

H Sensitivity against random seeds

I Computational cost

  • On damped wave, Latent-MoE is more computationally efficient than FB-PINNs, and inference can use at most two experts per point on non-overlapping subdomains.Figure 19 visualizes compute-cost curves; Table 9 reports training and inference times at different widths.Source: I Computational cost · p. 33, Figure 19 · p. 33, Table 9 · p. 33

J Reproducibility Statement

  • Exact reproducibility is not guaranteed because GPU operations are nondeterministic, an effect amplified by the stiff PINN loss landscape.Outcomes may vary slightly across hardware accelerators, software versions, and operating systems.Source: J Reproducibility Statement · p. 34

K Declaration of LLM usage

  • The authors used publicly available LLMs for coding, writing, and proofreading assistance, audited outputs, and ran experiments manually without exposing results to the LLMs.Source: K Declaration of LLM usage · p. 34

L Broader Impact

  • The authors caution that apparent fidelity does not replace validation against reference solvers or experimental data, and advise independent verification before safety-critical deployment.Source: L Broader Impact · p. 34
Loading 2609.07814v1…