Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Generative Latent Neural PDE Solver using Flow Matching

Zijie Li, Anthony Zhou, Amir Barati Farimani

TL;DR

This work tackles temporal instability and spectral bias in data-driven PDE solvers by introducing a generative latent PDE solver that operates in a mesh-reduced latent space via flow matching diffusion. An autoencoder maps irregular meshes to a structured latent grid, enabling efficient diffusion-based forecasting, while a Diffusion Transformer predicts latent velocity fields for multi-step evolution. Empirical results across 2D buoyancy flow, 3D MHD turbulence, and UAV airflow demonstrate improved long-term accuracy and stability over deterministic baselines, with ensembles further boosting robustness and spectral coherence. The approach offers a scalable, flexible surrogate framework for time-dependent PDEs on complex geometries and irregular discretizations.

Abstract

Autoregressive next-step prediction models have become the de-facto standard for building data-driven neural solvers to forecast time-dependent partial differential equations (PDEs). Denoise training that is closely related to diffusion probabilistic model has been shown to enhance the temporal stability of neural solvers, while its stochastic inference mechanism enables ensemble predictions and uncertainty quantification. In principle, such training involves sampling a series of discretized diffusion timesteps during both training and inference, inevitably increasing computational overhead. In addition, most diffusion models apply isotropic Gaussian noise on structured, uniform grids, limiting their adaptability to irregular domains. We propose a latent diffusion model for PDE simulation that embeds the PDE state in a lower-dimensional latent space, which significantly reduces computational costs. Our framework uses an autoencoder to map different types of meshes onto a unified structured latent grid, capturing complex geometries. By analyzing common diffusion paths, we propose to use a coarsely sampled noise schedule from flow matching for both training and testing. Numerical experiments show that the proposed model outperforms several deterministic baselines in both accuracy and long-term stability, highlighting the potential of diffusion-based approaches for robust data-driven PDE learning.

Generative Latent Neural PDE Solver using Flow Matching

TL;DR

This work tackles temporal instability and spectral bias in data-driven PDE solvers by introducing a generative latent PDE solver that operates in a mesh-reduced latent space via flow matching diffusion. An autoencoder maps irregular meshes to a structured latent grid, enabling efficient diffusion-based forecasting, while a Diffusion Transformer predicts latent velocity fields for multi-step evolution. Empirical results across 2D buoyancy flow, 3D MHD turbulence, and UAV airflow demonstrate improved long-term accuracy and stability over deterministic baselines, with ensembles further boosting robustness and spectral coherence. The approach offers a scalable, flexible surrogate framework for time-dependent PDEs on complex geometries and irregular discretizations.

Abstract

Autoregressive next-step prediction models have become the de-facto standard for building data-driven neural solvers to forecast time-dependent partial differential equations (PDEs). Denoise training that is closely related to diffusion probabilistic model has been shown to enhance the temporal stability of neural solvers, while its stochastic inference mechanism enables ensemble predictions and uncertainty quantification. In principle, such training involves sampling a series of discretized diffusion timesteps during both training and inference, inevitably increasing computational overhead. In addition, most diffusion models apply isotropic Gaussian noise on structured, uniform grids, limiting their adaptability to irregular domains. We propose a latent diffusion model for PDE simulation that embeds the PDE state in a lower-dimensional latent space, which significantly reduces computational costs. Our framework uses an autoencoder to map different types of meshes onto a unified structured latent grid, capturing complex geometries. By analyzing common diffusion paths, we propose to use a coarsely sampled noise schedule from flow matching for both training and testing. Numerical experiments show that the proposed model outperforms several deterministic baselines in both accuracy and long-term stability, highlighting the potential of diffusion-based approaches for robust data-driven PDE learning.

Paper Structure

This paper contains 18 sections, 12 equations, 6 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Neural PDE solver
  4. Neural solver on reduced-order representations
  5. Temporal stability of neural solver
  6. Diffusion probabilistic models in PDEs
  7. Methodology
  8. Diffusion probabilistic model
  9. Diffusion in a unified mesh-reduced space
  10. Autoencoder
  11. Diffusion Transformer
  12. Experiments and Discussion
  13. Problems
  14. 2D buoyancy-driven flow
  15. Magnetohydrodynamics compressible turbulence
  16. ...and 3 more sections

Figures (6)

  • Figure 1: Main architecture of the proposed latent neural solver based on flow matching. The input data is projected onto a unified latent space via encoder. In the latent space, the diffusion model iteratively transform the noise into the latent state of future timestep. The latent embedding can be projected back to the original mesh space via the decoder.
  • Figure 2: Spectrum averaged across a time-window centered at 10-th frame.
  • Figure 3: Spectrum averaged across a time-window centered at 50-th frame.
  • Figure 4: Spectrum averaged across a time-window centered at 90-th frame.
  • Figure 5: Visualization of the velocity field ($x$ component) in 3D MHD predicted by different models.
  • ...and 1 more figures