Abstract
Neural operators have shown promise in learning solution maps of partial differential equations, but they often struggle to generalize when test inputs lie outside the training distribution — novel initial conditions, unseen coefficients, or unseen physics. Prior work addresses this by large-scale multi-physics pretraining followed by fine-tuning; that still requires examples from the new dynamics, falling short of true zero-shot generalization.
We propose enhancing generalization at test time, without modifying pretrained weights. Building on DISCO, which provides a dictionary of neural operators trained across different dynamics, we introduce a neural-operator-splitting strategy that, at test time, searches over compositions of training operators to approximate unseen dynamics. On challenging out-of-distribution tasks — parameter extrapolation and novel combinations of physics phenomena — our approach achieves state-of-the-art zero-shot results while recovering the underlying PDE parameters.
Method · no retraining, no fine-tuning, no labels
Search compositions of frozen operators at test time
We begin with a DISCO-pretrained backbone. A transformer hypernetwork ψα maps each training trajectory ui1:L — a history of L snapshots — to a small U-Net operator fθi, where θi = ψα(ui1:L). After pretraining we freeze everything and treat the encoded set {f1, …, fN} as a fixed dictionary of differential operators — one per training environment.
At test time we observe only a short trajectory — the same L ≈ 16 snapshots — of a new, possibly out-of-distribution dynamics. From it we search for a subset S ⊆ {f1, …, fN} whose sum best fits the observed evolution. The fitting objective is the averaged NRMSE between the next-step prediction obtained via operator splitting and the observed test data. We compare two search strategies: uniform random sampling of subsets, and beam search that greedily grows compositions while keeping the top-B candidates per depth.
Once a composition is chosen, the next step ut → ut+1 is realised through Strang operator splitting: each fi is integrated for a fractional timestep in sequence, producing a single combined integrator. Composition lengths up to M = 5 suffice across all benchmarks.
Throughout the figures and videos, DISCO original (a.k.a. direct) is the base DISCO prediction with no composition, and DISCO beam is our composition found by beam search.
The core idea · learn in isolation, combine at test time
Operators learned from physics in isolation, composed for unseen physics
During pretraining the model is trained on many different physics, each in isolation, and DISCO learns a physics operator from each trajectory (or environment of trajectories). At test time our method meets a new trajectory governed by unseen physics — without ever observing its parameters — and approximates it by searching over and combining operators from the learned dictionary. No retraining, no labels.
↓ search & combine operators at test time ↓
Table 1 results · seven OOD combinations
Composition of unseen physics
We evaluate on seven held-out physics combinations drawn from four PDE datasets — advection–diffusion, the combined equation (nonlinear advection, diffusion and dispersion), Gray–Scott reaction–diffusion, and 2D Navier–Stokes — reporting NRMSE against the baselines below.
α nonlinear advectionβ diffusionγ dispersionν viscosity (2D)
Reported NRMSE · lower is better · bold marks the best per column
| Method | Adv + Diff | NonlinAdv + Diff | NonlinAdv + Disp | Diff + Disp | All three | React + Diff | Euler + Diff |
|---|---|---|---|---|---|---|---|
| MPP | 0.270 | 0.050 | 0.105 | 0.091 | 0.128 | 0.191 | 0.273 |
| FNO | 0.318 | 0.031 | 0.165 | 0.038 | 0.129 | 0.224 | 0.241 |
| Zebra | 0.893 | 0.022* | 0.241 | 0.069 | 0.193 | 0.127 | 0.198 |
| GEPS | 0.039 | 0.039 | 0.249 | 0.229 | 0.265 | 0.128 | 0.786 |
| DISCO (Original) | 0.170 | 0.085 | 0.100 | 0.120 | 0.164 | 0.245 | 0.572 |
| Ours (Uniform) | 0.043 | 0.068 | 0.103 | 0.043 | 0.075 | 0.089 | 0.209 |
| Ours (Beam) | 0.015 | 0.056 | 0.049 | 0.007 | 0.036 | 0.089 | 0.066 |
Each column is a held-out physics combination. Methods see each phenomenon individually during pretraining; at test time multiple phenomena act simultaneously. *Zebra wins one column (nonlinear advection + diffusion); we win the remaining six.
Composition over time · watch beam search add operators
Deep dive: the best composition of one, two, then three operators
Navier–Stokes is the Euler equation (nonlinear advection) plus viscous diffusion. The model saw Euler and diffusion as separate training dynamics, but never their viscous combination — so this 2D flow is genuinely out of distribution, and the hardest case. Below, what beam search finds at each level of complexity — the best single operator, the best pair, the best triple — each from the same search over the frozen dictionary.
In every clip the top row is the ground truth. Each row below is the best composition of that many operators: as more operators are allowed, the prediction (left) tracks the truth more closely and the residual |pred − GT| (right) darkens toward zero. The banner names the operators in each composition.
The examples span viscosities from ν = 3×10⁻⁴ (nearly inviscid, turbulent) to 1×10⁻² (smooth, diffusion-dominated) — a range no single pretrained operator covers, yet a short composition does.
Notice how few operators are needed: even the turbulent, low-viscosity flows are matched by a composition of just two or three. The search keeps reusing a small set of building blocks across the whole range — a sign it is recovering physically meaningful operators rather than overfitting the handful of observed snapshots.
How to read it top row = ground truth · each row below = best composition of that many operators · right column = residual
2D Navier–Stokes
Table 2 results · coefficients pushed outside the training range
Parameter extrapolation
Two datasets, with test trajectories generated from coefficients set outside the training range: advection–diffusion (extrapolating the advection speed c and the diffusion D) and the combined equation (extrapolating nonlinear advection α and dispersion γ). The same test-time search still approximates the dynamics where the baselines break down.
c advection speedD diffusionα nonlinear advectionγ dispersion
Reported NRMSE · lower is better · bold marks the best per column
| Method | Adv–Diff c | Adv–Diff D | Combined α | Combined γ |
|---|---|---|---|---|
| MPP | 0.588 | 0.409 | 0.134 | 0.369 |
| FNO | 0.492 | 0.166 | 0.166 | 0.317 |
| Zebra | 1.070 | 1.579 | 0.128 | 0.448 |
| GEPS | 0.848 | 0.267 | 0.020 | 0.782 |
| DISCO | 0.768 | 0.159 | 0.088 | 1.007 |
| Ours (Uniform) | 0.113 | 0.055 | 0.027 | 0.070 |
| Ours (Beam) | 0.052 | 0.002 | 0.016 | 0.022 |
Advection–diffusion: train on small c and D, test on values roughly 3× larger. Combined-equation: extrapolate α and γ beyond the training distribution. Beam search consistently recovers OOD parameters that all baselines fail on — most strikingly on Adv-Diff D, where the error drops by nearly two orders of magnitude.
Builds on · related work
Literature
This work builds directly on DISCO, which learns the dictionary of neural operators we compose at test time. Below: DISCO, followed by the neural-operator baselines compared in the tables above.
DISCO: learning to DISCover an evolution Operator for multi-physics-agnostic predictionfoundation
Multiple Physics Pretraining for Physical Surrogate Models
Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
GEPS: Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning
Fourier Neural Operator for Parametric Partial Differential Equations