DRIFT OR DYNAMICS: RESIDUAL CORRECTION OF LONG-HORIZON NEURAL PDE ROLLOUTS

Learned simulators offer a way to make expensive physical simulation much faster, with potential applications across science and engineering, from fluid and material dynamics to weather and climate modeling. But using a learned model to predict far into the future remains difficult: each prediction becomes the input to the next, so small errors can accumulate until the simulation no longer resembles the true system. Existing approaches improve long-horizon behavior by training models to make longer sequences of predictions, repeatedly refining their outputs, or adding mechanisms that pull predictions back toward physically or statistically plausible states. These strategies can extend the useful prediction horizon, but they leave an important question unanswered: what should we do when a prediction moves away from the states seen during training? Pulling it back seems natural—but a physical system may be genuinely evolving into a state the model has never seen before. We introduce a geometry-aware correction that pulls back only the part of a prediction outside the principal directions of the training data, leaving the in-subspace component untouched. In controlled experiments, this targeted correction reduces long-horizon error by up to 90%. Importantly, we find that the geometry of the correction matters, and when it removes meaningful dynamics, error can instead increase by up to 130×. Our results show that long-horizon stability is not simply a matter of keeping predictions close to the training data. The critical question is whether a departure from the data is an error or a part of the dynamics we are trying to predict.