cs.LGOct 3, 2026

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

Authors: Anupam Rawat, Manikandan Padmanaban, Jagabondhu Hazra

Organizations: Indian Institute of Technology Bombay · IBM Research

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-KK 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⋆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×2.2\times lower error than hard top-1 routing at identical fusion; using generator weights w⋆w^\star at inference further lowers rollout error to 0.0340.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.

Figures & tables

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. U-HNO: A U-shaped Hybrid Neural Operator with Sparse-Point Adaptive Routing for Non-stationary PDE Dynamics

    May 13, 2026Yingzhe Ma, Xiao Yang, Yuxin Xie +2Fourier Neural OperatorPDE Surrogate Modeling

  2. Green-Routed Neural Operators: Physics Determines Where the Network Reads

    Oct 4, 2026Chenhao Si, Ming YanPDE Surrogate ModelingPDE Operator Learning

  3. Learning Physical Operators using Neural Operators

    Feb 26, 2026Vignesh Gopakumar, Ander Gray, Dan Giles +5PDE Surrogate ModelingPhysics-Informed Neural Operators