Paper deep dive

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

Anupam Rawat, Manikandan Padmanaban, Jagabondhu Hazra

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 43

Intelligence

Status: not_run | Model: - | Prompt: - | Confidence: 0%

Entities (0)

No extracted entities yet.

Relation Signals (0)

No relation signals yet.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

40,033 characters extracted from source content.

Expand or collapse full text

AI for Stochastic Dynamics COSMOS: Soft Mechanism Mixtures with Verifiable Routing for Long-Horizon PDE Forecasting Anupam Rawat Affiliation: Indian Institute of Technology Bombay Email: anupam.rawat@iitb.ac.in Manikandan Padmanaban Affiliation: IBM Research Email: manipadm@in.ibm.com Jagabondhu Hazra Affiliation: IBM Research Email: jahazra1@in.ibm.com Abstract Neural operators are an efficient alternative to classical PDE solvers, but most learn one monolithic map per equation family and per discretization. Real systems are more diverse: transport, diffusion, wave, and reaction processes often act together. Existing mixture-of-experts operators follow the sparse top-K routing template, which selects a few experts and switches the rest off. When several processes act at once, the update is a blend, not a discrete choice, so sparse selection throws away part of the physics. Public benchmarks such as PDEBench also evaluate equation families separately and carry no mechanism-level mixture labels, so compositional routing cannot be tested on them directly. 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; a learned gate mixes them with continuous weights, and a small network fuses their concatenated features. Specialists share one coarse latent grid, and a parallel zero-initialized full-resolution residual restores sharp detail the bottleneck discards. We also introduce a compositional benchmark generated by operator splitting, with known mixture weights w⋆w for every trajectory. On three PDEBench families, we compare against the Fourier neural operator tuned per family (three learning rates, five seeds each, best rate chosen on held-out validation) under a 20-step rollout protocol (normalized RMSE; lower is better): a locked 3-seed run is ahead on diffusion–reaction (0.2410.241 vs. 0.4070.407), roughly level on Navier–Stokes (0.4930.493 vs. 0.4690.469), and behind on shallow water (0.0910.091 vs. 0.0610.061). Because a 3-seed mean can hide instability, we re-ran the same recipe at 11 seeds: the Navier–Stokes result holds, but the diffusion–reaction win does not – that recipe diverges in a large minority of seeds, and its 11-seed median trails FNO. We therefore do not claim a reliable accuracy win on any of the three families; the audit itself, and what it implies about trusting small-seed rollout comparisons, is as much this paper’s result as the architecture is. Retraining from scratch with the gate fixed to uniform weights, or removing the specialist mixture while keeping the residual, both degrade rollout substantially in the stable regime. On the labeled benchmark, dense soft routing has 2.2×2.2× lower error than hard top-1 at identical fusion, and substituting the generator coefficients w⋆w at inference lowers rollout to 0.0340.034. That gap is a diagnostic of the learned gate, not a claim that w⋆w is optimal for the trained specialists. The labels also falsify our own interpretability claim: external routing does not align with the specialists’ intended physics, so we claim accuracy, not mechanism identity. We do not claim zero-shot multi-family coverage, 3D, or unstructured meshes. 1 Introduction Fast PDE surrogates matter wherever numerical simulation is the bottleneck, from weather and subsurface flow to engineering design. Learned operators such as the Fourier neural operator (FNO) (Li et al., 2021), GINO (Li et al., 2023), and Transolver (Wu et al., 2024) make inference orders of magnitude cheaper than classical solvers, but they share a structural assumption: one map is fitted to one training distribution. A generalist PDE surrogate must instead serve many mechanisms and their mixtures; COSMOS (Cooperative Operator Specialists with Mechanism-level Operator Soft-routing) is designed as a step in that direction. Mixture-of-experts operators are the natural response to physical diversity (Wang et al., 2025; Sun et al., 2026; Li et al., 2025; Hmida et al., 2026), and most adopt the sparse template from large-scale language modeling: a gate selects a top-K expert and switches the rest off. This is a mismatch of kind, not of how many experts to keep. A state that is partly advective and partly diffusive needs both processes at once; a hard gate must drop one of them, and that choice is least reliable exactly where the two processes overlap. The gate should stay dense: keep every specialist active and let continuous weights carry the blend. Section 6 measures the cost: at identical fusion, hard top-1 routing has 2.2×2.2× higher rollout error than soft, against the sparse template most MoE operators take from language modeling. Externally, against a properly tuned FNO, a locked 3-seed run is roughly level on one of three PDEBench families and ahead on a second, but a follow-up seed audit shows the second win is not reliable (Section 5). A verification problem sits underneath the modeling problem. MoE operators are routinely described as if experts specialize and routing discovers structure. These are claims about what a model has learned, and on standard benchmarks they are unfalsifiable: PDEBench (Takamoto et al., 2022) and comparable corpora are organized by equation family and carry no trajectory-level record of which processes were active in what proportion. A model can be called mechanism-aware with no experiment able to show otherwise. This paper treats that gap as the methodological contribution rather than a caveat. We construct a corpus by operator splitting in which every trajectory records the mixture w⋆w that generated it, supplying two things no public benchmark can: a generator-truth substitution that measures what changes when the learned gate is replaced by the known coefficients, and a falsification test for the mechanistic reading of the router. Applied to our own model, the test shows that external routing does not align with the specialists’ intended physics; we claim accuracy rather than mechanism identity, and take the result as a working demonstration of the instrument. Contributions. 1. Method. COSMOS: dense continuous gating over process-biased specialists, concat++MLP fusion, a shared latent library, and a zero-initialized full-resolution residual (Section 3). 2. Instrument. A compositional corpus with ground-truth mixture labels, enabling a generator-truth routing diagnostic and a controlled routing ablation that public benchmarks cannot support (Section 4). 3. Evidence, controls, and a self-audit. Shared-harness PDEBench results against a per-family tuned FNO, three ablative controls, a soft-vs-hard comparison at fixed fusion, and explicit negative results, including an 11-seed audit that overturns the locked 3-seed DR win (Sections 5–7). Scope. A small-scale, architecture-first study: COSMOS ∼ 3.4M parameters vs. FNO ∼ 1.2M, a 16×1616×16 latent, grid-only external tables; synthetic and PDEBench models are trained separately and never averaged. All experiments use deterministic solvers; we do not evaluate stochastic forcing or calibrated uncertainty. 2 Related Work Neural operators. FNO (Li et al., 2021) is the standard grid baseline and our primary comparison. GINO (Li et al., 2023) and Transolver (Wu et al., 2024) extend operator learning beyond regular grids; message-passing solvers (Brandstetter et al., 2022) emphasize local updates and rollout-aware training, whose unrolled objective we adopt. As in Perceiver IO (Jaegle et al., 2022), we map variable-resolution inputs onto a fixed latent grid. A separate line decomposes dynamics by term or prior: PDE-Net learns per-term differential operators (Long et al., 2018), APHYNITY splits physics from a learned residual (Yin et al., 2021), PhyDNet uses a two-stream physics++residual design (Le Guen and Thome, 2020). COSMOS differs in routing a library of process operators with learned continuous weights, evaluated against generator-known mixture labels on a controlled compositional benchmark. The labels are used for evaluation, not as a training signal. Mixture and compositional operators. CompNO (Hmida et al., 2026) assembles pretrained process blocks; iMOOE (Li et al., 2025) uses an invariance-aligned mixture of operator experts; MoE-POT (Wang et al., 2025) and NESTOR (Sun et al., 2026) scale sparse MoE routing for multi-PDE pretraining. COSMOS differs on three axes: dense soft gating rather than top-K sparsity (a return to the classical mixture formulation of Jacobs et al. (1991); top-K sparsity exists to save compute at hundreds of experts, and is the wrong default at small M), concat++MLP fusion rather than a weighted sum, and process-biased specialists rather than homogeneous blocks. Our comparisons against MoE-POT and NESTOR are matched from-scratch contact tests at small scale, outside the pretraining regime those methods target. Verification and interpretability of learned operators. The claim that an expert has specialized is an interpretability claim, and plausible attributions often fail under controlled test. Our contribution is narrow: a labeled generator makes the claim checkable, and we report a case where it does not survive. A companion study (Rawat et al., 2026) finds gated multi-stream encoding outperforms flat concatenation for PDE-family identification; we adopt the same principle. 3 Method Figure 1: COSMOS forward pass. Frames are encoded onto a shared 16×1616×16 latent, processed in parallel by four process-biased specialists, mixed by a dense router (no top-K), fused by concatenation++MLP, and decoded. A parallel full-resolution residual is zero-initialized, so training begins exactly at the latent-only update and the branch learns only what the coarse latent cannot represent. (Figure notation: Zop=z^tZ^op= z^t, Δlocal=Rρ _local=R_ρ; the decoder predicts the change from utu_t.) We target dynamics expressible as a combination of process operators, ∂u∂t=∑m=1Mαm(x,t)ℒm[u]+ε,M=4, ∂ u∂ t= _m=1^M _m(x,t)\,L_m[u]+ , M=4, (1) with transport, diffusion, wave, and reaction in the library. The term ε collects unmodeled physics, discretization error, and any stochastic forcing. Every dataset here comes from a deterministic solver, so we set ε=0 =0; we do not evaluate noise or calibrated uncertainty. The surrogate learns ut+1=θ(ut,…,ut−k+1)u^t+1=G_θ(u^t,…,u^t-k+1) and is applied autoregressively; mixture structure must be inferred from the encoded input history, without access to w⋆w . 3.1 Canonical latent and process-biased specialists The encoder EϕE_φ maps k past frames at native resolution to zt∈ℝd×Hc×Wcz^t ^d× H_c× W_c with Hc=Wc=16H_c=W_c=16, d=128d=128. The specialist library is defined once on this fixed grid, so only EϕE_φ and the decoder DωD_ω change for a different input layout. Each specialist fmf_m carries a lightweight architectural bias: transport predicts a displacement field and warps the latent by it; diffusion applies a learned per-wavenumber low-pass filter; wave applies a spectral filter with a learned dispersion phase (a bounded per-wavenumber complex rotation); reaction is a pointwise MLP with no spatial coupling. These biases shape the hypothesis space; they do not enforce semantic mechanism identity after training, which is precisely the gap Section 6 measures. 3.2 Dense routing and concatenation fusion A router reads a pooled summary of ztz^t and emits continuous weights w=softmax⁡(gψ(mean-pool(zt))/τ)∈ΔM−1,w=softmax\! (g_ψ(mean -pool(z^t))/τ )∈ ^M-1, (2) Every specialist is evaluated at every step; there is no top-K selection. Rather than forming ∑mwmfm(zt) _mw_mf_m(z^t), we concatenate weighted specialist features and fuse, z^t=MLPξ([w1f1(zt)‖⋯‖wMfM(zt)]), z^t=MLP_ξ\! ( [\,w_1f_1(z^t)\,\|·s\|\,w_Mf_M(z^t)\, ] ), (3) which can represent interactions between specialist outputs, not only their convex combination. Both choices are ablated in Section 6: the weighted sum costs 0.1760.176 against concatenation’s 0.1370.137, and hard top-1 costs 0.3010.301. 3.3 Full-resolution residual Compressing a 64×6464×64 field onto a 16×1616×16 latent discards sharp detail. We add a native-resolution correction, ut+1=ut+Dω(z^t)+Rρ(ut,…,ut−k+1),u^t+1=u^t+D_ω( z^t)+R_ρ(u^t,…,u^t-k+1), (4) with the final layer of RρR_ρ zero-initialized, so optimization begins exactly at the latent-only update ut+Dω(z^t)u^t+D_ω( z^t) and the mixture is undisturbed at step 00. 3.4 Objective We minimize an unrolled multi-step loss with mildly increasing step weights λk _k, ℒ=∑j=1Kλj‖u^t+j−ut+j‖2‖ut+j‖2+βℒroute,L= _j=1^K _j \| u^\,t+j-u^t+j \|_2 \|u^t+j \|_2+β\,L_route, (5) where ℒrouteL_route adds entropy and load-balancing regularization on w, discouraging premature collapse onto a single specialist. 4 An Instrument for Mechanism Claims External benchmarks establish whether a design transfers. They cannot establish why it works, because they carry no mechanism labels. We therefore generate a corpus by operator splitting, integrating each process operator in alternation with known coefficients, so the generator records exactly which processes acted. Construction. Each trajectory carries w⋆=(wT,wD,wW,wR)w =(w_T,w_D,w_W,w_R) summing to one, derived from the generator coefficients. Four pure families activate a single mechanism; five mixed families compose two or three: advection–diffusion (transport++diffusion; “advection” and “transport” are used interchangeably), reaction–diffusion, damped wave (wave++diffusion), advection–reaction–diffusion, and damped-wave–reaction (wave++diffusion++reaction). After integration on a 33×3333×33 grid, each trajectory is resampled onto an unstructured mesh and a point cloud (n=1024n=1024), sharing w⋆w and coefficients and differing only in coordinate layout. These are constructed siblings, not independent unstructured simulations, and we use them only as a transfer check. Splits. The primary split S2 trains on all four pure families (grid++mesh) and tests on unseen three-mechanism mixtures, a compositional generalization test within the mechanisms the generator represents. The harder split S1 trains on pure transport and diffusion only and tests on adv-diff, a deliberately out-of-regime probe. What the labels buy. w⋆w is never used as a training signal for any model reported here; it is the normalized relative strengths of the generator’s per-process coefficients, and serves evaluation only. It supplies two things that make a mechanism claim testable. A generator-truth substitution: replacing the learned weights by w⋆w at inference, specialists unchanged, measures how much performance depends on the gate. The resulting gap is a diagnostic of the learned router; it does not assume that w⋆w is the optimal routing vector for the trained specialists. A falsification test: correlating learned w against w⋆w asks whether the router recovers the generator’s mixture; externally, where w⋆w does not exist, we can only ask the weaker qualitative question of whether routing mass aligns with each family’s known physics (Section 6). w⋆w exists because we generated the data and cannot be recovered post hoc for PDEBench, so the corpus licenses a mechanism claim only where the generating mixture is known. 5 Results Protocol. All headline numbers are roll@20: RMSE after 20 autoregressive steps, divided by the norm of the true field, so 00 is perfect and the trivial repeat-last-frame baseline (persistence) sits at 2.512.51 on shallow water. One-step error is a secondary diagnostic: an order of magnitude smaller than rollout error and closely clustered across methods, so it does not distinguish them. COSMOS (∼ 3.4M, d=128d=128) and FNO (∼ 1.2M) are evaluated on 64×6464×64 PDEBench grids under one shared rollout harness. COSMOS uses three seeds, the setting under which its recipe was locked. For FNO, since a single learning rate does not fit every family well, we trained it at three rates and five seeds per rate, picked the best rate on held-out validation, and report only that rate’s five test seeds – a broader hyperparameter search than COSMOS received, so the comparison is not driven by an obviously suboptimal FNO learning rate. The latent bottleneck is real, and it motivates the residual. Before the residual is enabled, latent-only COSMOS on diffusion–reaction (DR) sits essentially at persistence (the trivial baseline that repeats the last input frame). Roughly 95%95\% of the initial-condition spectral energy on these trajectories lies above the 16×1616×16 latent Nyquist, so the coarse library cannot carry the sharp modes at all. This is a statement about these initial conditions at this latent resolution, not a general claim about latent operators. It is, however, what Equation (4) exists to fix, and every number below includes that path. Table 1: PDEBench roll@20 under the shared evaluation harness, lower is better. COSMOS: mean± over 3 seeds (∼ 3.4M parameters). FNO: mean± over 5 seeds at its best learning rate, chosen per family on a held-out validation split (∼ 1.2M parameters). This is not a strict capacity match; three controls reported in the text address different sources of the observed differences. Family COSMOS FNO (Li et al., 2021) Ratio Diffusion–reaction (DR) 0.241±0.061†0.241±0.061 0.407±0.2110.407±0.211 0.59×†0.59× Navier–Stokes (NS2D) 0.493±0.078†0.493±0.078 0.469±0.0240.469±0.024 1.05×†1.05× Shallow water (SWE) 0.091±0.0300.091±0.030 0.061±0.0290.061±0.029 1.49×1.49× † Locked 3-seed protocol estimate, not the recipe’s typical seed: §5’s 11-seed audit gives DR2D median 0.5480.548 (worse than FNO) and NS2D median 0.4690.469 (matches). Headline. Table 1 reports the comparison, with FNO tuned as described above rather than run at one fixed rate: the single rate we tried first was a poor fit for two of three families, leaving FNO 22–13×13× worse than its own best rate on DR and NS2D, while already close to best on SWE. At the locked 3-seed protocol, COSMOS leads on DR, is roughly level on NS2D, and trails on SWE. The table footnote gives the reason we do not stop there: §5’s 11-seed audit finds DR’s outcome is bimodal and its median seed loses to FNO, so we do not claim a reliable DR2D win; NS2D’s locked number matches its 11-seed median almost exactly, so that result stands. We report the tuned baseline, weaker for COSMOS on two of three families, because it is the fair comparison, and we report the audit alongside the locked table because the locked number alone would overstate DR2D. Control 1: is the gate doing anything? Training the full recipe from scratch with the gate fixed to uniform weights throughout, rather than learned, raises DR error from 0.2410.241 to 2.542.54, a collapse past persistence. This is a from-scratch retrain, not a swap on an already-trained model: fixing the gate before training changes what the specialists themselves learn to do: every specialist then receives the same fixed weight rather than an input-dependent training signal. A learnable dense gate is therefore part of the training recipe, not a post-hoc mixer; §6 returns to what the learned gate does once training is finished. Control 2: is the residual doing all the work? Because FNO carries no comparable native-resolution branch, we isolate the residual’s contribution directly. A residual-only model that removes the mixture entirely but keeps the full-resolution path reaches 0.2950.295 on DR (seed0; three-seed mean 0.516±0.4720.516±0.472, inflated by one divergent seed) against the full model’s 0.2410.241. The mixture therefore improves on the residual alone, and on the well-behaved seeds the residual-only model still sits below FNO’s tuned mean of 0.4070.407. Both ingredients contribute; neither explains the result by itself. Control 3: are the process-biased specialists themselves reliable? A 3-seed mean can look stable even when individual seeds are not (§5), so we trained the same locked recipe at 11 random seeds on DR and on NS2D, together with a control in which the four process-biased specialists are replaced by four generic blocks of equal size and no physics bias. The generic control uses 11 seeds on DR and the eight completed seeds on NS2D. We call a rollout error above 1.01.0, worse than persistence, a divergence. On DR, the process-biased specialists diverge on 44 of 1111 seeds. The worst of these is 34.8×34.8× higher error than the good seeds. The generic blocks diverge on only 11 of 1111 seeds, and that one case is barely over the threshold, not a collapse. On NS2D the pattern reverses: the process-biased specialists diverge on 22 of 1111 seeds, while the generic blocks diverge on 33 of 88. In every seed, good or bad, the one-step error looks the same: normal and small. The difference only shows up in the rollout. Error tracks normally through roughly step 1010, then grows sharply by step 2020. So the failure sits in the autoregressive loop, not in training. It is also not caused by the routing weights: forcing the learned gate to exactly uniform weights at inference, on both a stable and a divergent seed, changes the result by under 0.01%0.01\%. The immediate failure is therefore not the precise inference-time gate; it sits in the learned specialist and fusion dynamics. That does not rule out a role for routing during training. We read this as two separate findings, not one. Process-biased specialists carry a real robustness cost on DR specifically, and the 11-seed median (0.5480.548) trails FNO’s tuned mean (0.4070.407): the locked 3-seed win does not hold at typical seeds. NS2D, on the other hand, is rollout-fragile for every method we tested there, FNO included, independent of architecture, and its 11-seed median (0.4690.469) matches the locked estimate, so that result stands. Contact with MoE operators. At matched ∼ 4M from-scratch capacity, COSMOS reaches 0.280.28 on DR and 0.1370.137 on synthetic S2, against NESTOR (1.121.12, 0.4520.452) and MoE-POT (2.212.21, 1.0741.074); a time-boxed fine-tune of a pretrained MoE-POT Tiny checkpoint on DR reached 0.9280.928. These are contact probes at matched small scale, not evidence against either method in its intended multi-PDE-pretrained regime. 6 Routing Analysis on the Labeled Benchmark Table 2: Compositional S2 holdout, grid, roll@20 (3 seeds). All COSMOS rows use identical concat++MLP fusion, so the routing comparison is clean. Generator-truth substitutes w⋆w at inference while leaving the trained specialists unchanged: a diagnostic, not a baseline or a guaranteed optimum. “—” = not recorded under this protocol. Method one-step roll@20 Generator-truth w⋆w (diagnostic) 0.0021±0.00030.0021±0.0003 0.034±0.0130.034±0.013 COSMOS dense soft routing (primary) 0.0031±0.00010.0031±0.0001 0.137±0.0650.137±0.065 weighted-sum fusion 0.0032±0.00010.0032±0.0001 0.176±0.0230.176±0.023 soft hierarchical 0.0032±0.00010.0032±0.0001 0.220±0.0570.220±0.057 uniform weights 0.0032±0.00010.0032±0.0001 0.283±0.0560.283±0.056 hard top-1 — 0.301±0.0060.301±0.006 FNO (Li et al., 2021) 0.0034±0.00010.0034±0.0001 0.326±0.0380.326±0.038 GINO (Li et al., 2023) (2.72.7M) — 0.369±0.0210.369±0.021 Persistence 0.00350.0035 0.3700.370 NESTOR (Sun et al., 2026) 0.0046±0.00050.0046±0.0005 0.452±0.2330.452±0.233 Dense monolith (matched) 0.0236±0.01090.0236±0.0109 0.572±0.1620.572±0.162 MoE-POT (Wang et al., 2025) 0.0065±0.00140.0065±0.0014 1.074±0.5921.074±0.592 Dense soft routing versus hard top-1. Table 2 holds fusion fixed and varies only the gate, which makes the comparison attributable. Soft-flat routing (0.1370.137) improves on hierarchical routing (0.2200.220), uniform weights (0.2830.283), and hard top-1 (0.3010.301). The last comparison is the one that carries the argument of Section 1: at matched fusion, dense continuous routing is 2.2×2.2× better than the sparse template MoE operators inherit from language modeling. Paired per-trajectory Wilcoxon tests rank the primary configuration above every learned baseline at p<10−40p<10^-40 (Appendix C); those tests are conditional on the trained checkpoints and do not replace seed-level uncertainty. Elaboration does not help. Hierarchical routing, cross-expert attention, shared experts, and gated constraints were all implemented; none improved the chain in either fusion mode (Appendix B). We report the compact configuration that survives, not the one we expected. Where the error actually lives. Replacing the learned router by w⋆w lowers roll@20 from 0.1370.137 to 0.0340.034, about 4×4×, with the specialists left unchanged. Routing is therefore a large, recoverable source of error in this controlled setting. That does not make w⋆w the optimal routing vector for the trained specialists; it is a diagnostic that public benchmarks cannot supply. Routing diagnostics. On the compositional corpus, learned weights correlate with w⋆w at 0.670.67–0.690.69 across seeds, and mean routed weight lands on the correspondingly labeled specialist for three of the four mechanisms (transport 1.001.00, diffusion 0.950.95–0.970.97, reaction 1.001.00 across seeds; Figure 2). The fourth is instructive: pure-wave trajectories route to the transport specialist. Linear and nonlinear probes did not support pooling loss, router capacity, or data diversity as sufficient explanations: transport and wave are both spectrum-preserving operators, so a single frame of either is statistically a band-limited field (encoded separability caps at ≈0.76≈0.76, vs. 0.960.96–1.001.00 for every other pair). The distinction lives in temporal evolution (first-order directional vs. second-order isotropic), which the pooled router summary does not see; diffusion and reaction transform the spectrum and are identifiable from one frame. Mixture accuracy is unaffected, since soft fusion keeps both operators active – itself an argument for dense routing over a hard choice. On external PDEBench diffusion–reaction, routing mass concentrates on the transport and wave specialists rather than diffusion and reaction. We therefore do not claim mechanism identity on external data, and read the specialist names as design-time inductive biases throughout. What the trained gate actually does at inference. Control 1 shows a gate fixed to uniform from the start of training costs accuracy – a claim about training, not about the converged, learned gate, so we checked the latter directly. Over 50005000 external DR predictions, the learned gate’s per-sample entropy sits at 98%98\% of its maximum: on almost every input it outputs close to equal weight on all four specialists. Forcing a trained checkpoint’s weights to exactly uniform at inference, no retraining, changes the result by under 0.01%0.01\% on both a well-behaved and a divergent Control 3 seed. Turning off the training-time regularizer that discourages a sharp gate has the opposite effect: routing sharpens (entropy falls to 42%42\% of maximum) and rollout error nearly quadruples. That last observation is a single seed, so we treat it as diagnostic rather than conclusive. Put together: the regularizer appears to keep the gate from collapsing early, which would starve the other specialists of gradient. For the external DR checkpoints we inspected, the precise converged weights matter little at inference: replacing their near-uniform outputs by exact uniforms changes rollout by under 0.01%0.01\%. Control 1 still stands — a learnable gate matters to the training recipe — but the external DR evidence does not support reading the converged gate as selecting a physically correct mixture on each input. Figure 2: Mean routed weight per specialist on pure-mechanism trajectories, three seeds. The largest mean weight aligns with the correspondingly labeled specialist for three of four mechanisms; pure wave routes to transport (0.820.82–0.930.93), a spectral ambiguity discussed in Section 6. Constructed siblings. Reusing the same trained weights on mesh and point-cloud siblings degrades roll@20 from 0.1370.137 to 0.1750.175 (1.27×1.27×) and 0.1800.180 (1.31×1.31×), supporting a single latent library across layouts on constructed siblings only – not evidence on independent unstructured solvers. Scaling, preliminary. A capacity-matched joint pretrain over five 2D families, followed by a short per-family fine-tune, recovers rollout at or below the specialist solo on all five, and the same recipe transfers to five 1D families (Appendix A). These runs lack the seed-count audit of §5, so we report them as preliminary only. 7 Limitations Mechanism identity. The routing diagnostics of Section 6 hold on the labeled corpus, not externally: on PDEBench DR, routing mass does not align with the specialists’ intended physics. We claim no interpretability on external data; specialist names are design-time biases. Robustness, and this paper’s clearest limitation. §5’s 11-seed audit overturns the DR2D win, not just caveats it: the median seed trails FNO, and the process-biased specialists diverge far more often than a matched generic-expert control on DR. We could not identify which specialist causes it, or a fix, in the time available; reporting this over the locked 3-seed table matters more than a result that looks cleaner than the method is. Three-seed comparisons are underpowered for training-run variability; Appendix C reports paired trajectory tests separately and does not mix them with seed-level statistics. Out-of-regime mixtures. On the harder S1 split (train on two mechanisms, test on their unseen combination), COSMOS trails FNO: soft mixtures do not extrapolate beyond the trained regime. Baseline matching and scale. COSMOS carries more parameters than FNO and a native-resolution branch FNO lacks (the residual-only control bounds this rather than eliminating it), and trains multi-step unrolled where FNO trains one-step; whether a rollout-trained FNO narrows the gap is open. Increasing the external model’s width also degrades DR rollout, unexplained so far; external tables are regular grids only, mesh/point-cloud results are constructed siblings not independent simulations, single-seed 3D results trail specialized baselines, and MoE-operator comparisons are from-scratch at matched capacity, outside those methods’ pretraining regime. 8 Conclusion COSMOS mixes process-biased specialists with a dense continuous gate, concat++MLP fusion, and a zero-initialized full-resolution residual. A locked 3-seed run is ahead of tuned FNO on DR and level on NS2D; the 11-seed audit overturns the DR result at typical seeds (§5) and leaves the NS2D estimate unchanged. A labeled generator and a larger seed sweep each falsified one of our own claims. Closing the generator-truth–learned routing gap and isolating DR’s unstable specialist are next. References Brandstetter et al. (2022) J. Brandstetter, D. Worrall, and M. Welling Message passing neural PDE solvers. In International Conference on Learning Representations (ICLR), Cited by: §2. Hmida et al. (2026) H. Hmida, H. Joly, and Y. Mesri CompNO: a novel foundation model approach for solving partial differential equations. Applied Sciences 16 (2), p. 972. External Links: Document Cited by: §1, §2. Jacobs et al. (1991) R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton Adaptive mixtures of local experts. Neural Computation 3 (1), p. 79–87. Cited by: §2. Jaegle et al. (2022) A. Jaegle, S. Borgeaud, J. Alayrac, C. Doersch, C. Ionescu, D. Ding, S. Koppula, D. Zoran, A. Brock, E. Shelhamer, O. Hénaff, M. Botvinick, A. Zisserman, O. Vinyals, and J. Carreira Perceiver IO: a general architecture for structured inputs and outputs. In International Conference on Learning Representations (ICLR), Cited by: §2. Le Guen and Thome (2020) V. Le Guen and N. Thome Disentangling physical dynamics from unknown factors for unsupervised video prediction. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Cited by: §2. Li et al. (2025) S. Li, Y. Chen, Y. Guo, M. Huang, and H. Xiong Towards generalizable PDE dynamics forecasting via physics-guided invariant learning. arXiv preprint arXiv:2509.24332. Cited by: §1, §2. Li et al. (2021) Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations (ICLR), Cited by: §1, §2, Table 1, Table 2. Li et al. (2023) Z. Li, N. Kovachki, C. Choy, B. Li, J. Kossaifi, S. Otta, M. A. Nabian, M. Stadler, C. Hundt, K. Azizzadenesheli, and A. Anandkumar Geometry-informed neural operator for large-scale 3d pdes. Advances in Neural Information Processing Systems (NeurIPS). Cited by: §1, §2, Table 2. Long et al. (2018) Z. Long, Y. Lu, X. Ma, and B. Dong PDE-net: learning PDEs from data. In International Conference on Machine Learning (ICML), Cited by: §2. Rawat et al. (2026) A. Rawat, M. Padmanaban, and J. Hazra Spectral probe: routing signals for mixture-of-experts PDE operators. Technical report IBM Research. Note: Companion technical report Cited by: §2. Sun et al. (2026) D. Sun, X. Zhou, X. Wang, H. Si, W. Lyu, J. Tang, and B. Luo NESTOR: a nested MOE-based neural operator for large-scale PDE pre-training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 6147–6156. Cited by: §1, §2, Table 2. Takamoto et al. (2022) M. Takamoto, T. Praditia, R. Leiteritz, D. MacKinlay, F. Alesiani, D. Pflüger, and M. Niepert PDEBench: an extensive benchmark for scientific machine learning. In Advances in Neural Information Processing Systems (NeurIPS), Datasets and Benchmarks Track, Cited by: §1. Wang et al. (2025) H. Wang, H. Xin, J. Wang, X. Yang, F. Zha, H. Dong, and Y. Jiang Mixture-of-experts operator transformer for large-scale PDE pre-training. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 2510.25803 Cited by: §1, §2, Table 2. Wu et al. (2024) H. Wu, H. Luo, H. Wang, J. Wang, and M. Long Transolver: a fast transformer solver for pdes on general geometries. In International Conference on Machine Learning (ICML), Cited by: §1, §2. Yin et al. (2021) Y. Yin, V. Le Guen, J. Dona, E. de Bézenac, I. Ayed, N. Thome, and P. Gallinari Augmenting physical models with deep networks for complex dynamics forecasting. In International Conference on Learning Representations (ICLR), Cited by: §2. Appendix A Scaling the Recipe Five 2D families with joint pretraining. Beyond per-family specialists, we train one capacity-matched joint backbone (d=256d=256, Km=8K_m=8, ∼ 16.2M) on five PDEBench 2D families and apply a short per-family fine-tune. Zero-shot joint transfer is not sufficient (SWE in particular degrades sharply), but after fine-tuning, rollout recovers at or below the specialist solo on all five families (Table 3). We do not claim a single frozen checkpoint for all of PDEBench, and the capacity here differs from Table 1, so rows are not comparable across the two. Table 3: Joint pretrain then short per-family fine-tune, roll@20. Ratio is fine-tuned/solo (acceptance gate ≤1.10×≤ 1.10×; 5/5 pass). Family Solo Joint ++ FT vs. solo SWE 0.01600.0160 0.00140.0014 0.09×0.09× DR2D 0.13540.1354 0.12030.1203 0.89×0.89× NS2D 0.46730.4673 0.37880.3788 0.81×0.81× CFD2D 0.16720.1672 0.08870.0887 0.53×0.53× Darcy 0.89040.8904 0.90900.9090 1.02×1.02× Five 1D families. The same recipe with a 1D strip residual also covers five PDEBench 1D families under a matched solo protocol (per-family nMSE 0.00010.0001–0.0760.076; see the 1D detail in Appendix E), showing the residual mechanism is not specific to the 2D setting. Appendix B Full Cooperation Ablation The complete a–h chain over both fusion modes. Under weighted-sum fusion: dense soft routing 0.1760.176, hierarchical 0.2760.276, cross-expert attention 0.2850.285, shared experts 0.3260.326, gated constraints 0.2860.286, hard top-1 0.3080.308, generator-truth 0.0390.039. Under concatenation fusion: dense soft routing 0.1370.137, hierarchical 0.2200.220, uniform 0.2830.283, hard top-1 0.3010.301, generator-truth 0.0340.034. No cooperation mechanism beyond dense soft mixing improved the chain in either fusion mode. Concatenation improves on weighted sum for both dense soft and hierarchical routing, which is why the primary configuration uses it. Appendix C Paired Significance Tests Paired per-trajectory Wilcoxon tests on S2 roll@20 (n=384n=384 held-out trajectories) rank the primary configuration above every learned baseline at p<10−40p<10^-40. These tests are conditional on the trained models and measure paired consistency across trajectories; they are not a substitute for seed-level variation. The only non-significant pair in the full matrix is GINO versus persistence (p=0.16p=0.16). Seed-level n=3n=3 statistics are reported separately and never combined with these. Appendix D Compositional Corpus Examples Figure 3: Example pure-mechanism training trajectories. Each row is one of the four generator mechanisms; columns are time snapshots. w⋆w is known by construction for every pure and mixed trajectory because the generator records the operator-splitting coefficients. Appendix E Experimental Details Five 1D families. Per-family nMSE under the matched solo protocol (3 seeds): advection 0.041±0.0130.041±0.013, Burgers 0.038±0.0140.038±0.014, diffusion–reaction 0.044±0.0110.044±0.011, compressible flow 0.076±0.0030.076±0.003, sorption 0.0001±0.0000.0001±0.000 (near-floor under this subsample; not over-read). FNO baseline. 12 retained Fourier modes per dimension, width 32, 4 spectral layers, trained per family with the shared next-step objective and evaluated under the same rollout harness as COSMOS. Test populations. Paired per-trajectory tests on the compositional S2 holdout use 384 held-out trajectories; each external PDEBench family is evaluated on 100 held-out trajectories per seed. Compositional corpus. On the order of 10310^3 trajectories per family across four pure and five mixed families, each emitted as grid (33×3333×33), mesh, and point-cloud siblings; splits and generator settings ship with the corpus. PDEBench sources. Shallow water uses the 2D radial dam-break subset; Navier–Stokes uses the incompressible velocity fields; diffusion–reaction uses the 2D two-channel subset. Up to 400 trajectories per family are adapted to 64×6464×64 grids for training and evaluation. Hardware. Training and evaluation ran on single NVIDIA H200-class GPUs. Appendix F Protocol Seams Three groups of numbers appear in this paper and are never averaged across: (i) external per-family COSMOS (∼ 3.4M, d=128d=128) versus FNO (∼ 1.2M) on 64×6464×64 PDEBench grids, three seeds under the shared evaluation harness; (i) compositional synthetic (33×3333×33, ∼ 4M) with mechanism labels; (i) joint pretrain++fine-tune (d=256d=256, Km=8K_m=8, ∼ 16.2M). Every table is annotated with its group. All headline numbers are drawn from frozen claim records rather than transcribed from logs.

Instant research discovery