Predicting Change, Not States: An Alternate Framework for Neural PDE Surrogates

TL;DR

This paper proposes an alternative neural PDE surrogate framework that predicts temporal derivatives and uses ODE integration, improving accuracy, stability, and flexibility over traditional black-box methods.

Contribution

It introduces a neural PDE surrogate approach that predicts derivatives and employs numerical integration, enhancing performance and adaptability across various PDEs and discretizations.

Findings

Improved accuracy and stability in finely-discretized regimes.

Enhanced flexibility with variable time-stepping.

Broad applicability across different PDE models.

Abstract

Neural surrogates for partial differential equations (PDEs) have become popular due to their potential to quickly simulate physics. With a few exceptions, neural surrogates generally treat the forward evolution of time-dependent PDEs as a black box by directly predicting the next state. While this is a natural and easy framework for applying neural surrogates, it can be an over-simplified and rigid framework for predicting physics. In this work, we evaluate an alternate framework in which neural solvers predict the temporal derivative and an ODE integrator forwards the solution in time, which has little overhead and is broadly applicable across model architectures and PDEs. We find that by simply changing the training target and introducing numerical integration during inference, neural surrogates can gain accuracy and stability in finely-discretized regimes. Predicting temporal derivatives also allows models to not be constrained to a specific temporal discretization, allowing for flexible time-stepping during inference or training on higher-resolution PDE data. Lastly, we investigate why this framework can be beneficial and in what situations does it work well.