# COSMOS: Soft Mechanism Mixtures with Verifiable Routing for Long-Horizon PDE Forecasting

arXiv: https://arxiv.org/abs/2610.04427 · HTML: https://mlsanity.com/paper/2610.04427/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 15 references that are in the map.
- First posted on arXiv 3 Oct 2026.

## Builds on (most cited first, 13 of 13)

- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [PDE-Net: Learning PDEs from Data](https://mlsanity.com/paper/1710.09668/) (2017)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Disentangling Physical Dynamics from Unknown Factors for Unsupervised Video Prediction](https://mlsanity.com/paper/2003.01460/) (2020)
- [Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting](https://mlsanity.com/paper/2010.04456/) (2020)
- [Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2510.25803/) (2025)
- [CompNO: A Novel Foundation Model approach for solving Partial Differential Equations](https://mlsanity.com/paper/2601.07384/) (2026)
- [NESTOR: A Nested MOE-based Neural Operator for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2602.22059/) (2026)
- [Towards Generalizable PDE Dynamics Forecasting via Physics-Guided Invariant Learning](https://mlsanity.com/paper/2509.24332/) (2025)

## Abstract (arXiv)

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× 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.
