Guided Autoregressive Diffusion Models with Applications to PDE Simulation

Federico Bergamin, Cristiana Diaconu, Aliaksandra Shysheya, Paris Perdikaris, José Miguel Hernández-Lobato, Richard E. Turner, Emile Mathieu

Keywords: neural PDE solver, PDE, partial differential equation, forecasting, data-assimilation, diffusion, denoising, autoregressive, neural surrogate, reconstruction guidance

TL;DR: We develop a joint diffusion model trained on short trajectory segments and conditioned a posteriori allowing for autoregressive generation of accurate and stable predictions of long PDE trajectories.

Abstract: Solving partial differential equations (PDEs) is of crucial importance in science and engineering. Yet numerical solvers necessitate high space-time resolution which in turn leads to heavy computational cost. Often applications require solving the same PDE many times, only changing initial conditions or parameters. In this setting, data-driven machine learning methods have shown great promise, a principle advantage being the ability to simultaneously train at coarse resolutions and produce fast PDE solutions. In this work we introduce the Guided AutoRegressive Diffusion model (GuARD), which is trained over short segments from PDE trajectories and a posteriori sampled by conditioning over (1) some initial state to tackle forecasting and/or over (2) some sparse space-time observations for data assimilation purposes. We empirically demonstrate the ability of such a sampling procedure to generate accurate predictions of long PDE trajectories.

Submission Number: 60