Summary
Simulating physical processes like heat spreading, waves, or fluid flow often involves solving complex equations. The authors propose a new method called COSMOS that mixes multiple specialized neural modules softly instead of choosing a single one, better reflecting how real-world processes blend together. This approach improves prediction stability and accuracy over time for certain physical systems, though the improvements vary by problem. The method also includes tests to verify how well the model understands the underlying physics, revealing that while it predicts well, the internal experts do not always align exactly with specific physical mechanisms.
What this means in practice
- •For computational physicists: Improve long-term predictions of complex PDE systems by blending multiple neural experts continuously instead of selecting one exclusively.
- •For climate model developers: Enhance modular forecasting components for physical processes like diffusion, wave propagation, and transport within climate simulations.
Abstract
Neural operators offer an efficient alternative to classical PDE solvers, but most learn a monolithic map per equation family and discretization. Real systems are compositional: transport, diffusion, wave, and reaction processes can act simultaneously. Existing mixture-of-experts operators typically use sparse top-$K$ routing, although concurrent physics is naturally a blend rather than a discrete choice. We propose COSMOS (Cooperative Operator Specialists with Mechanism-level Operator Soft-routing), a soft mechanism-mixture neural operator. Four process-biased specialists remain active at every step and are continuously mixed by a learned gate, with their features fused by a small network. Specialists share a coarse latent grid, while a zero-initialized full-resolution residual restores detail lost through the bottleneck. We also introduce an operator-splitting compositional benchmark with known mixture weights $w^\star$ per trajectory. Against family-tuned FNO under 20-step rollouts, an initial 3-seed evaluation suggested gains on diffusion--reaction, parity on Navier--Stokes, and weaker shallow-water performance. An 11-seed audit showed that the diffusion--reaction gain was unstable, motivating caution in small-seed rollout comparisons and precluding a reliable accuracy-win claim on these families. Ablations show that uniform routing or removing the specialist mixture substantially degrades stable-regime rollouts. On the labeled benchmark, dense soft routing yields $2.2\times$ lower error than hard top-1 routing at identical fusion; using generator weights $w^\star$ at inference further lowers rollout error to $0.034$, diagnosing limitations of the learned gate. However, routing labels do not align with the specialists' intended mechanisms: COSMOS supports compositional accuracy, not mechanism identity.