Test-time Generalization for Physics through Neural Operator Splitting

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 multi 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.

Physical systems like heat flow, fluid motion, and waves are described by equations that are slow to solve numerically. AI models called neural operators can predict these solutions almost instantly, but they usually only work on the situations they were trained on. In the context of a new starting condition, a different material property, or an unfamiliar mix of physics, their accuracy degrades dramatically. The common fix is to train on lots of varied physics and then retrain for each new problem, but that still needs examples of the new behavior, so it isn't really true zero-shot generalization. Instead, we propose a method that improves a pretrained model (DISCO) at the moment it's used, without changing any of its weights. Building on a system that stores a library of operators each expert in different dynamics, our method searches for ways to combine these experts to approximate unseen physics. On hard cases like extreme parameters and novel physics combinations, it achieves the best zero-shot results so far, and can even recover the hidden physical parameters. This points to test-time computation as a key route to flexible, compositional neural operators.