Test-time Generalization for Physics through Neural Operator Splitting

TL;DR

This paper tackles the lack of zero-shot generalization in neural operators for PDEs under test-time distribution shifts. It introduces a test-time neural operator splitting approach that composes pretrained operators from a DISCO dictionary, selecting the best combination via search and evolving the system with neural operator splitting, all without updating weights. Across parameter extrapolation and multi-physics composition tasks, the method achieves state-of-the-art zero-shot generalization and enables PDE parameter identification from test-time observations. The results highlight test-time computation as a powerful avenue for building flexible, compositional, and generalizable neural operators with practical impact for complex spatiotemporal dynamics.

Abstract

Neural operators have shown promise in learning solution maps of partial differential equations (PDEs), but they often struggle to generalize when test inputs lie outside the training distribution, such as novel initial conditions, unseen PDE coefficients or unseen physics. Prior works address this limitation with large-scale multiple physics pretraining followed by fine-tuning, but this still requires examples from the new dynamics, falling short of true zero-shot generalization. In this work, we propose a method to enhance generalization at test time, i.e., 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 including parameter extrapolation and novel combinations of physics phenomena, our approach achieves state-of-the-art zero-shot generalization results, while being able to recover the underlying PDE parameters. These results underscore test-time computation as a key avenue for building flexible, compositional, and generalizable neural operators.

Figures (17)

• Figure 1: Test-time generalization through neural operator splitting. During pretraining (left), DISCO learns a dictionary of operators for different physics, such as reaction dynamics (green) and diffusion+kill dynamics (red), with a hypernetwork generating corresponding operator weights $\theta_i, \theta_j$. At test time (right), on OOD dynamics such as reaction-diffusion, our method searches over combinations of these operators to approximate the new dynamics (e.g., $f_{OOD} \approx f_{i_1} + f_{i_2}$), and evolves $u^t \rightarrow u^{t+1}$ using neural operator splitting via sequential operator applications.
• Figure 2: Test-time generalization on Gray--Scott equations. Our neural operator search correctly predicts an unseen, non-trivial dynamics (compare first and second rows), which differs substantially from the pure reaction (third row) or pure diffusion (fourth row) seen during training, demonstrating that our method, based on combining simple operators, can capture complex phenomena.
• Figure 3: Test-time scaling laws. (Left) Performance of our test-time method with uniform search on an out-of-distribution (OOD) advection–diffusion dynamics. We report fitting and prediction error, measured as the mean NRMSE over 34 rollout steps, as a function of the number of uniform-search trials. (Right) Mean absolute error (MAE) for PDE parameter identification versus the number of trials, showing that the selected operators recover meaningful and accurate physical parameters.
• Figure 4: Zero-shot generalization to unseen advection-diffusion PDEs. Contour maps report the NRMSE (averaged over 34 rollout steps; lower is better) over the grid of advection (x-axis) and diffusion coefficients (y-axis). (a) DISCO predictions obtained with the operator produced by the encoder and hypernetwork. (b) Our test-time method: beam search selects operators from the dictionary to minimize the fitting error; the selected operators are then composed and unrolled via neural operator splitting.
• Figure 5: Fitting error versus computational cost (FLOPs) for uniform and beam search strategies across three tasks: diffusion extrapolation, advection extrapolation, and combined advection–diffusion. Beam search proceeds sequentially over operator complexity, first selecting the best single operators, and subsequently higher-order compositions. At intermediate FLOP budgets, beam search achieves substantially lower error than uniform search, while uniform search exhibits a smooth, approximately power-law decay with increasing computation.