Generative Latent Neural PDE Solver using Flow Matching

Autoregressive next-step prediction models have become standard for building data-driven neural solvers to predict time-dependent partial differential equations (PDEs). The use of diffusion models has been shown to enhance the temporal stability of neural solvers, while its stochastic inference mechanism enables ensemble predictions and uncertainty quantification. However, a key drawback of diffusion models is the need to sample a series of discretized timesteps during both training and inference, which increases computational overhead. In addition, most diffusion models operate on structured, uniform grids, limiting their adaptability to irregular domains. To address these shortcomings, we propose a latent flow matching (FM) model for PDE simulation that embeds the PDE state in a lower-dimensional latent space, which reduces computational costs. In addition, we design an autoencoder to map different meshes onto a unified, structured latent grid, which allows predictions on complex geometries. Furthermore, we show that FM can result in faster and more accurate predictions than diffusion-based models, even with a coarser noise schedule. Numerical experiments show that the proposed model outperforms several deterministic and probabilistic baselines in both accuracy and long-term stability, highlighting the potential of FM-based approaches for data-driven PDE learning.