Test-time Generalization for Physics through Neural Operator Splitting

TL;DR

This paper introduces a test-time neural operator splitting method that composes trained operators to achieve zero-shot generalization for PDE solutions, outperforming prior approaches on out-of-distribution tasks.

Contribution

It proposes a novel test-time composition strategy for neural operators that enhances zero-shot generalization without retraining or fine-tuning.

Findings

State-of-the-art zero-shot generalization on out-of-distribution PDE tasks

Ability to recover underlying PDE parameters

Effective composition of trained operators for unseen physics

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.