PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Anupam Rawat, Manikandan Padmanaban, Jagabondhu Hazra

arXiv:2610.04427v1cs.LGcs.AI

TL;DR

Neural operators commonly fit one map per training distribution, while real systems can combine physical mechanisms and standard benchmarks lack trajectory-level mixture labels for checking mechanism claims. COSMOS combines process-biased specialists with dense soft routing and introduces a labeled compositional benchmark. Seed audits do not support a reliable accuracy-win claim over tuned FNO, and the paper limits its claim to compositional accuracy rather than mechanism identity.

  • Problem

    Neural operators often fit one map per training distribution, and equation-family benchmarks lack trajectory-level mixture labels needed to test whether routing identifies active mechanisms.

    Source: §1 · p. 2

  • Method

    COSMOS combines process-biased specialists with dense continuous gating and introduces an operator-splitting corpus with known mixture weights for routing diagnostics.

    Source: §1 · p. 2

  • Results

    The 11-seed audit overturns the locked 3-seed diffusion–reaction win, and dense soft routing has 2.2× lower rollout error than hard top-1 routing at identical fusion.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    The labeled benchmark shows that COSMOS supports compositional accuracy claims, not mechanism identity claims about its specialists.

    Source: §1 · p. 2

    Three-seed comparisons are underpowered for training-run variability, and the median seed in the 11-seed diffusion–reaction audit trails FNO.

    Source: §7 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

COSMOS addresses the mismatch between single-distribution neural operators and systems with concurrent physical mechanisms, while also making mechanism claims testable with known trajectory-level mixture labels. Its benchmark and audits support caution about small-seed accuracy comparisons and limit its claim to compositional accuracy rather than mechanism identity.

Motivation

  • A generalist PDE surrogate must handle multiple mechanisms and their mixtures, unlike operators fitted to one training distribution or sparse top-K mixtures that switch experts off.The introduction frames simultaneous processes such as advection and diffusion as a blend that should keep specialists active.Source: §1 · p. 2

Verification gap

  • Standard equation-family benchmarks lack trajectory-level process-mixture labels, leaving claims that routing discovers mechanisms unfalsifiable.COSMOS adds an operator-splitting corpus that records each trajectory's generating mixture weights and enables generator-truth substitution and a falsification test.Source: §1 · p. 2

Contributions

  • COSMOS combines dense continuous gating over process-biased specialists with concatenation-plus-MLP fusion, a shared latent library, and a zero-initialized full-resolution residual.The paper also introduces a labeled compositional corpus for controlled routing diagnostics.Source: §1 · p. 2

Evidence and scope

  • The 11-seed audit overturns the locked 3-seed diffusion–reaction win, so the introduction does not establish a reliable accuracy advantage over tuned FNO.The paper describes its evidence as PDEBench comparisons, routing controls, a soft-versus-hard comparison, and explicit negative results.Source: §1 · p. 2
  • The study is small-scale and architecture-first, uses separately trained synthetic and PDEBench models, and excludes stochastic forcing and calibrated uncertainty.COSMOS has about 3.4M parameters versus FNO's about 1.2M, and external tables are grid-only.Source: §1 · p. 2

2 Related Work

Related work spans grid-based neural operators, physics-informed decompositions, and sparse mixture-of-experts operators. COSMOS instead uses dense soft routing and evaluates mechanism claims against known mixture labels.

Neural operators

  • FNO is the standard grid baseline, while GINO and Transolver extend operator learning beyond regular grids.Other cited approaches emphasize local updates and rollout-aware training, or map variable-resolution inputs onto fixed latent grids.Source: §2 · p. 3

Physics and compositional operators

  • PDE-Net, APHYNITY, and PhyDNet decompose dynamics by term or prior, whereas COSMOS routes a library of process operators with learned continuous weights.COSMOS evaluates against generator-known mixture labels on a controlled compositional benchmark, using labels for evaluation rather than training.Source: §2 · p. 3

Mixture and compositional operators

  • COSMOS differs from cited MoE operators through dense soft gating rather than top-K sparsity, concat-plus-MLP fusion rather than weighted sums, and process-biased rather than homogeneous specialists.Its small-scale MoE-POT and NESTOR comparisons are matched from-scratch tests outside those methods' pretraining regime.Source: §2 · p. 3

Verification and interpretability

  • A labeled generator makes expert-specialization claims checkable, and the paper reports a case where the mechanistic interpretation does not survive testing.A companion study finds gated multi-stream encoding outperforms flat concatenation for PDE-family identification; COSMOS adopts the same principle.Source: §2 · p. 3

3 Method

COSMOS models dynamics as mixtures of process operators inferred from input history, using process-biased specialists, dense routing, and concatenation fusion on a shared latent grid. A zero-initialized full-resolution residual restores detail, and training uses an unrolled multi-step objective with routing regularization.

3 Method

  • The method targets dynamics expressible as combinations of transport, diffusion, wave, and reaction operators, with deterministic datasets setting the unmodeled-physics term to zero.The model infers mixture structure from encoded input history without access to generator weights.Source: §3 · pp. 3–4

3.1 Canonical latent and process-biased specialists

  • Four specialists share a 16×16 latent grid and use distinct architectural biases: displacement warping, learned spectral filters, dispersion phase, or pointwise reaction processing.These biases shape the hypothesis space but do not enforce semantic mechanism identity after training.Source: §3.1 · p. 4

3.2 Dense routing and concatenation fusion

  • A pooled latent summary produces continuous routing weights; every specialist remains active, and weighted specialist features are concatenated for MLP fusion.Concatenation can represent interactions beyond a convex combination; the reported weighted-sum error is 0.176 versus 0.137 for concatenation, while hard top-1 costs 0.301.Source: §3.2 · p. 4

3.3 Full-resolution residual

  • A native-resolution correction addresses sharp detail lost when a 64×64 field is compressed to a 16×16 latent, and zero initialization starts optimization at the latent-only update.The residual branch can then learn corrections beyond that initial update.Source: §3.3 · p. 4

3.4 Objective

  • Training minimizes an unrolled multi-step loss with mildly increasing step weights, while route regularization penalizes entropy and load imbalance to discourage premature specialist collapse.The loss is applied over multiple forecast steps.Source: §3.4 · p. 4

4 An Instrument for Mechanism Claims

The section introduces an operator-splitting corpus with known per-trajectory process mixtures, enabling tests of compositional generalization and mechanism claims. It also describes the external benchmark comparison and the limits of interpreting routing without generator labels.

4 An Instrument for Mechanism Claims

  • The operator-splitting corpus records known process coefficients for each trajectory, providing a generator-truth substitution and a test of whether routing recovers the mixture.The corpus includes pure and mixed mechanism families; trajectories are also resampled to mesh and point-cloud layouts for a transfer check.Source: §4 · pp. 4–5
  • Table 1 reports PDEBench roll@20, where lower is better, comparing COSMOS means over 3 seeds with FNO means over 5 seeds at a validation-selected learning rate.The models differ in parameter count, so the table is not a strict capacity match.Source: Table 1 · p. 5
  • The 11-seed audit reports DR2D median 0.548, worse than FNO, and NS2D median 0.469, matching the locked estimate.The supplied passages do not specify the metric label for these medians.Source: §4 · p. 5
  • The primary split trains on all four pure families and tests on unseen three-mechanism mixtures, while the harder split tests advection–diffusion after training only on pure transport and diffusion.The primary split probes compositional generalization within represented mechanisms; the harder split is an out-of-regime probe.Source: §4 · p. 5
  • Generator weights w⋆ are normalized relative process coefficients used for evaluation, not as a training signal; replacing learned weights with w⋆ measures gate dependence.The substitution does not assume w⋆ is optimal for the trained specialists, and the weights cannot be recovered post hoc for PDEBench.Source: §4 · p. 5

5 Results

Under the locked 3-seed protocol, COSMOS leads on diffusion–reaction, is roughly level on Navier–Stokes, and trails on shallow water; an 11-seed audit removes a reliable diffusion–reaction win. Controls show that learned dense gating and the specialist mixture contribute to stable-regime rollouts, while seed audits expose rollout instability.

5 Results

  • The shared PDEBench comparison measures roll@20 on 64×64 grids, using 3 COSMOS seeds and FNO’s best validation-selected rate across five test seeds.FNO received a broader learning-rate search than COSMOS; the models have approximately 3.4M and 1.2M parameters, respectively.Source: §5 · p. 5
  • On diffusion–reaction trajectories at 16×16 latent resolution, roughly 95% of initial-condition spectral energy lies above the latent Nyquist; before residual enablement, latent-only COSMOS sits essentially at persistence, repeating the last input frame.Source: §5 · p. 5
  • At 3 seeds, COSMOS leads on DR, is roughly level on NS2D, and trails on SWE; the 11-seed audit rules out a reliable DR2D win, while the NS2D result stands.The DR2D median is 0.548 versus FNO’s tuned mean of 0.407; the NS2D 11-seed median is 0.469.Source: §5 · pp. 5–6
  • Fixing the gate to uniform weights throughout training raises DR error from 0.241 to 2.54, beyond persistence.This is a from-scratch retrain, so the specialists also learn under the fixed, non-input-dependent weighting.Source: §5 · p. 6
  • Residual-only COSMOS reaches 0.295 on DR at seed 0, versus 0.241 for the full model, showing the mixture improves on the residual alone.The residual-only three-seed mean is 0.516±0.472; the passage notes that one divergent seed inflates it.Source: §5 · p. 6
  • On DR, process-biased specialists diverge on 4 of 11 seeds versus 1 of 11 generic-block seeds; errors grow sharply by step 20 after looking similar at one step.A divergence is defined as rollout error above 1.0; the immediate failure is attributed to learned specialist and fusion dynamics, not the precise inference-time gate.Source: §5 · p. 6
  • At matched small scale, COSMOS reaches 0.28 on DR and 0.137 on synthetic S2, versus NESTOR’s 1.12 and 0.452 and MoE-POT’s 2.21 and 1.074.These are contact probes, not evidence against either method in its intended multi-PDE-pretrained regime.Source: §5 · p. 6

6 Routing Analysis on the Labeled Benchmark

On the labeled compositional benchmark, dense soft routing outperforms hard top-1 routing at matched fusion, and substituting generator weights substantially reduces rollout error. Routing diagnostics reveal limits to mechanism identification, especially on external PDEBench data.

Dense soft routing versus hard top-1

  • 2.2× lower rollout error than hard top-1 routing at identical fusion supports dense soft routing over sparse selection.Soft-flat routing scores 0.137 versus 0.301 for hard top-1.Source: §6 · p. 6, Table 2 · p. 7

Elaboration does not help

  • Hierarchical routing, cross-expert attention, shared experts, and gated constraints did not improve the chain in either fusion mode.The compact configuration is reported instead.Source: §6 · p. 7

Where the error actually lives

  • Roll@20 falls from 0.137 to 0.034 when generator weights w⋆ replace the learned router, with specialists unchanged.This is a routing diagnostic, not a guaranteed optimum.Source: §6 · p. 7, Table 2 · p. 7

Routing diagnostics

  • Learned weights correlate with w⋆ at 0.67–0.69 across seeds, and the largest weight matches the labeled specialist for three of four pure mechanisms.Pure-wave trajectories instead route to transport, with mean routed weights of 0.82–0.93; the pooled router summary misses the distinction in temporal evolution.Source: §6 · p. 7, Figure 2 · p. 8

External routing

  • On external PDEBench diffusion–reaction, routing mass favors transport and wave rather than diffusion and reaction, so mechanism identity is not supported externally.The external learned gate is near-uniform; replacing its outputs with exact uniform weights changes rollout by under 0.01%.Source: §6 · pp. 7–8

Gate behavior

  • Forcing a trained gate to uniform weights changes external diffusion–reaction rollout by under 0.01%, while removing the training regularizer sharpens routing to 42% entropy and nearly quadruples error in a single-seed diagnostic.The learned gate’s entropy is 98% of maximum over 5000 external predictions.Source: §6 · p. 7

Constructed siblings

  • Reusing trained weights on constructed mesh and point-cloud siblings raises roll@20 from 0.137 to 0.175 and 0.180, respectively.These results support a shared latent library across constructed layouts only.Source: §6 · p. 8

Scaling

  • Preliminary joint pretraining and per-family fine-tuning recovered rollout at or below the specialist solo on five 2D families and transferred to five 1D families.These runs lack the seed-count audit and are reported as preliminary.Source: §6 · p. 8

7 Limitations

The paper limits its claims because an 11-seed audit overturns the diffusion–reaction win, out-of-regime mixtures trail FNO, and external results do not establish mechanism identity. Its external evaluations also have architectural, training, and data-layout constraints.

Mechanism identity and scope

  • PDEBench diffusion–reaction routing does not align with intended specialist physics, so the paper claims no external mechanism interpretability.COSMOS also has more parameters than FNO, uses multi-step training, and reports external tables on regular grids only; whether rollout-trained FNO narrows the gap remains open.Source: §7 · p. 8

Robustness

  • The 11-seed audit overturns the diffusion–reaction win: the median seed trails FNO, and process-biased specialists diverge more often than a matched generic-expert control.Three-seed comparisons are underpowered for training-run variability, and the cause and a fix remain unidentified.Source: §7 · p. 8

Out-of-regime mixtures

  • On the harder S1 split, COSMOS trails FNO when trained on two mechanisms and tested on their unseen combination.The result bounds the claim to mixtures within the trained regime.Source: §7 · p. 8

8 Conclusion

The conclusion presents COSMOS as a dense-gated mixture of process-biased specialists, while emphasizing that its empirical claims are constrained by an audit and a labeled-generator diagnostic.

Conclusion

  • COSMOS combines process-biased specialists through a dense continuous gate, concat+MLP fusion, and a zero-initialized full-resolution residual.Source: §8 · p. 8
  • The 11-seed audit overturns the locked three-seed diffusion–reaction result at typical seeds, while leaving the Navier–Stokes estimate unchanged.Source: §8 · p. 8
  • The labeled generator and larger seed sweep falsified claims, and the next steps are closing the generator-truth–learned routing gap and isolating unstable diffusion–reaction behavior.Source: §8 · p. 8

A Scaling the Recipe

Joint pretraining followed by a short per-family fine-tune recovered rollouts at or below specialist-solo levels across five 2D families, and the recipe also covered five 1D families.

A Scaling the Recipe

  • The 2D result does not establish one frozen checkpoint for all PDEBench families, and its capacity differs from Table 1, so those rows are not comparable.Source: A Scaling the Recipe · p. 10
  • 5/5 2D families passed the roll@20 acceptance gate after joint pretraining and short per-family fine-tuning, with rollout at or below specialist-solo levels.The Table 3 gate is fine-tuned/solo ≤1.10×.Source: Table 3 · p. 10, A Scaling the Recipe · p. 10
  • Five 1D families achieved per-family nMSE values from 0.0001 to 0.076 under the matched solo protocol, extending the residual recipe beyond 2D.The sorption result is described as near-floor under this subsample and should not be over-read.Source: A Scaling the Recipe · p. 10

B Full Cooperation Ablation

Dense soft routing led the tested cooperation ablations under both fusion modes, and paired trajectory tests ranked the primary configuration above every learned baseline.

B Full Cooperation Ablation

  • 0.137 rollout error for dense soft routing with concatenation, versus 0.301 for hard top-1; no tested cooperation mechanism beyond dense soft mixing improved the chain.With weighted-sum fusion, dense soft routing scored 0.176 and hard top-1 scored 0.308; generator-truth routing scored 0.034 with concatenation and 0.039 with weighted-sum fusion.Source: B Full Cooperation Ablation · p. 10

C Paired Significance Tests

  • Paired tests on 384 held-out S2 trajectories ranked the primary configuration above every learned baseline on roll@20 at p < 10^-40.These tests measure paired consistency across trajectories, not seed-level variation.Source: C Paired Significance Tests · p. 10

E Experimental Details

The experimental details distinguish compositional and external evaluations, their test populations, and three reporting groups; the compositional corpus also provides known mechanism weights.

E Experimental Details

  • The FNO baseline used 12 retained Fourier modes per dimension, width 32, and four spectral layers, with per-family training and the shared next-step objective.It used the same rollout harness as COSMOS.Source: E Experimental Details · p. 10
  • External PDEBench families were evaluated on 100 held-out trajectories per seed, while paired tests on the compositional S2 holdout used 384 held-out trajectories.Source: E Experimental Details · p. 10
  • The compositional corpus contains on the order of 10^3 trajectories per family across four pure and five mixed families, with grid, mesh, and point-cloud siblings.The grid representation is 33×33.Source: E Experimental Details · p. 10
  • Figure 3 shows four pure-mechanism rows across time snapshots; generator-recorded operator-splitting coefficients provide known w⋆ for pure and mixed trajectories.Source: Figure 3 · p. 11
  • External PDEBench families used up to 400 trajectories adapted to 64×64 grids; shallow water, Navier–Stokes, and diffusion–reaction used specified subsets.Training and evaluation ran on single NVIDIA H200-class GPUs.Source: E Experimental Details · p. 11

F Protocol Seams

  • Results belong to three separate groups: external per-family COSMOS versus FNO, compositional synthetic data, and joint pretrain plus fine-tune; their numbers are never averaged across groups.Source: F Protocol Seams · p. 11
Loading 2610.04427v1…