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

arXiv:2412.13074 · cs.LG · Submitted 2024-12-17 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Predicting Change, Not States".

Jane: The gist: Neural surrogates can gain accuracy and stability in finely-discretized regimes by predicting the temporal derivative instead of directly predicting the next state,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, looking at this paper, "Predicting Change, Not States: An Alternate Framework for Neural PDE Surrogates," it’s really about finding a different way to train these physics simulators.

Jane: The authors are suggesting that instead of directly predicting the next state u n+one we should train the AI to predict the temporal derivative d u / d t <ref:2412.13074#pg2>.

Lu: This framework offers a way to improve accuracy and flexibility because it lets you decouple the spatial dynamics from the time integration step, which is useful across different PDE setups <ref:2412.13074#pg3>.

Meng: It’s a practical trade-off where you might have a little extra inference cost with the ODE integrator, but you get better performance on certain problems <ref:2412.13074#pg3>.

Tom: The implication is that this isn't just about making one model better; it’s about creating a training and inference framework that is agnostic to the specific PDE or the surrogate architecture <ref:2412.13074#pg3>.

Jane: It suggests that for simulating complex physics, we might be better off focusing on learning how things change over time rather than trying to force the network to guess where they will end up <ref:2412.13074#pg1>.

Lu: They are opening the door for more research into this alternative training and inference structure across a wider variety of problems and models <ref:2412.13074#pg3>.

Tom: It gives us a solid alternative framework to explore when we're building these next-generation physics simulators <ref:2412.13074#pg1>.

Conclusion: Tom: So we’re wrapping up this look at "Predicting Change, Not States: An Alternate Framework for Neural PDE Surrogates." Basically, the paper argues that instead of training an AI to predict the exact next state in time, it's smarter to train it to predict the rate of change—the derivative.

Jane: That’s a big shift from just predicting u at time t+one. So what’s the core idea there?

Tom: The core idea is that you can decouple the spatial part of the problem from the time integration. Instead of forcing your neural surrogate to guess every single tiny step, it learns how fast things are changing right now.

Lu: And because we have this derivative prediction, we can then use a standard ODE integrator—like Runge-Kutta or Forward Euler—to evolve the solution forward in time. It’s like the AI does the "how fast" part and a separate tool handles the "where to next" part.

Meng: From an engineering standpoint, that sounds promising for stability. If we can decouple them, maybe we don't have to worry as much about making one massive neural network handle every single temporal discretization step at once.

Lalam: Based on the results, they found that this derivative approach actually led to lower errors compared to the direct state prediction method across different equations. It’s more stable for chaotic systems like the Kolmogorov Flow, too.

Jane: So for someone just listening, it means we can get better accuracy in those finely-discretized regimes where time steps are really tight, because the model isn't trying to be perfect at every single point in time simultaneously.

Tom: Exactly. The paper suggests this framework is agnostic to the specific PDE or even the surrogate architecture you're using right now. It’s a change in training objective that can actually make a difference when you’re simulating complex physics.

Lu: The authors are really interested in *why* it works and under what conditions it’s best. They show that while the integrator still introduces some numerical error, it’s usually much smaller than the error you get from using a bad model at all.

Meng: So, there's a trade-off where you introduce a small numerical integration cost for better accuracy on certain dynamics. That makes sense for deployment considerations.

Jane: It seems like the authors are really pushing this idea to see if it can replace the direct state prediction method in more real-world scenarios across different physics problems.

Tom: Right. So, we’ve seen how predicting change instead of states opens up a new way to think about training these AI surrogates. Next up, we're going to look at the specific training loss functions they used and how that actually looks in practice.

Anthony Zhou, Amir Barati Farimani

Department of Mechanical Engineering, Carnegie Mellon University · Machine Learning Department, Carnegie Mellon University

cs.LG

Submitted: 2024-12-17

Updated: 2025-04-16

Comments: 22 pages, 7 figures. For code see http://github.com/anthonyzhou-1/temporal_pdes

Journal ref: Computer Methods in Applied Mechanics and Engineering, Volume 441, 2025

DOI: 10.1016/j.cma.2025.117990

Code: https://github.com/anthonyzhou-1/temporal

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 80/100

The gist: The gist: Neural surrogates can gain accuracy and stability in finely-discretized regimes by predicting the temporal derivative instead of directly predicting the next state, which allows for

Key concepts

State Prediction Framework: The traditional method where a neural network is trained to directly predict the next state of the system, $u_{n+1} = F_ heta(u_n)$. While simple, this framework can be rigid and constrain models to specific temporal discretizations when predicting physics.
Temporal Derivative Prediction: The proposed alternative where the neural network is trained to predict the rate of change (the temporal derivative) of the system. This decouples spatial learning from time evolution, allowing an external ODE integrator to handle the actual time stepping, which improves flexibility.
ODE Integrator: A numerical method used during inference to evolve a solution forward in time based on the predicted temporal derivatives. Methods like Forward Euler or Runge-Kutta are employed. This allows for flexible time-stepping during inference, enabling arbitrary temporal resolutions.
Training Loss Function ($L_ heta$): The objective function used to train the neural model, which can be defined based on whether the model predicts the next state or the temporal derivative. Precomputing derivatives for training speeds up this process, although calculating them dynamically incurs a small overhead.

Terminology

Neural Surrogates — 4× in this paper · explained in 3 episodes
Accuracy — 1× in this paper · explained in 1 episode

Summary

The gist: Neural surrogates can gain accuracy and stability in finely-discretized regimes by predicting the temporal derivative instead of directly predicting the next state, which allows for flexible time-stepping during inference

Alternate Framework for Neural Surrogates

The paper evaluates an alternate framework where 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. This approach contrasts with the dominant state prediction framework where models are trained to directly predict un+1 = Fθ(un). The core idea is that predicting the temporal derivative allows models to not be constrained to a specific temporal discretization, enabling flexible time-stepping during inference or training on higher-resolution PDE data.

Benefits of Derivative Prediction

Derivative prediction offers several advantages over state prediction, primarily by changing the training objective and inference procedure. The main benefit is that it allows the use of an ODE integrator to evolve the solution based on the predicted temporal derivative, which can be quickly queried at inference time. This effectively decouples the prediction problem, where the model learns spatial dynamics while an ODE integrator updates the solution in time. Furthermore, derivative prediction can be more flexible during inference than state prediction because it can generate solutions of arbitrary temporal resolutions by modifying ∆t in the ODE integrator during inference.

Training and Inference Details

The training loss function Lθ is defined based on whether the model predicts the next state or the temporal derivative, as shown in equation (2). The temporal derivatives for a given dataset can be precomputed and cached for training speed, although calculating them on the fly during training incurs a small overhead. During inference, an ODE integrator is used to evolve the solution using rules such as Forward Euler or Runge-Kutta methods. For explicit multistep methods like Adams-Bashforth, previous derivative estimates can be cached during inference to take more accurate forward steps without extra computational cost.

Performance Comparison and Results

Experimental results show that derivative prediction can achieve lower errors than state prediction across equations and model choices. For chaotic systems such as the KS equation or Kolmogorov Flow, models using derivative prediction tend to remain accurate for longer. In general, the paper finds that derivative prediction can be more stable than state prediction because it simply adds a change to the previous state rather than constructing a new state.

Limitations and Future Directions

The main limitation of derivative prediction is that it reintroduces numerical error and discretization constraints back into the neural surrogate framework. Specifically, even with perfect estimates of the instantaneous derivative, numerical integrators can still accumulate error over time with large timesteps ∆t or stiff PDE dynamics. For steady-state problems or boundary value problems, derivative prediction would not work in its current formulation. The paper concludes that the integrator can introduce numerical error, however it is very small with respect to model error, even when considering numerically unstable step sizes.

Conclusion

The evaluation suggests that learning the temporal derivative can be beneficial due to decomposing the solution step into a spatial and temporal update. The framework offers useful data and inference flexibility, making it a viable alternative for neural PDE surrogates. This work provides an alternate training/inference framework that is agnostic to the PDE or surrogate architecture, as well as investigating an alternate training/inference framework. The results suggest that there can be benefits in learning the temporal derivative, rather than directly predicting the next state, and in practice, implementing this through different ODE integrators can further affect accuracy. The paper hopes that future work can apply derivative prediction across a wider variety problems and model architectures, and improve performance at coarser step sizes. The data availability is provided at https://huggingface.co/datasets/ayz2/temporal pdes. The code to reproduce this work is provided at https://github.com/anthonyzhou-1/temporal pdes.

--- Page 1 ---

The gist: Neural surrogates can gain accuracy and stability in finely-discretized regimes by predicting the temporal derivative instead of directly predicting the next state

Benefits of Derivative Prediction

Improvements for AI systems

The system can predict future solutions by predicting temporal derivatives rather than directly predicting the next state, which can improve neural surrogate accuracies and allows models to be not constrained to a specific temporal discretization.
During inference, the model can be used with an ODE integrator, such as Forward Euler or Runge-Kutta, where the ODE integrator evolves the solution based on the predicted temporal derivative, allowing for flexible time-stepping during inference or training on higher-resolution PDE data.
The system can utilize higher-order integrators during inference to achieve greater accuracy and stability, as shown in Figure 4 where RK4 improves accuracy for chaotic systems like the KS equation.
The model can be trained on the full-resolution data for derivative prediction, which can avoid discarding a considerable amount of existing data, potentially leading to highly accurate estimates of the instantaneous derivative.
The system can adaptively change the step size during inference, taking larger steps where the solution changes slowly and smaller steps where the solution changes quickly, based on its predicted temporal derivatives.

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.