Adaptive SDE Interpolants for Calibrated Probabilistic PDE Forecasting

Jorge Mifsut Benet, Armand Kassaï Koupaï, Ramon Daniel Regueiro-Espino, Nicolas Baskiotis, Patrick Gallinari

Keywords: uncertainty quantification, stochastic interpolants, pde, partial differential equations

Abstract: Neural surrogate models for PDEs increasingly rely on generative methods to produce predictive ensembles, but existing approaches typically employ fixed noising schedules, leading to poorly calibrated uncertainty, especially under autoregressive rollouts and limited inference budgets. We propose a framework for calibrated probabilistic PDE forecasting that learns an adaptive latent stochastic transport between consecutive time steps. We model the transition between latent VAE posteriors using a learned Itô SDE with state-dependent drift and diffusion, trained via closed-form conditional supervision derived from a Gaussian stochastic interpolant, and regularized with an energy distance to improve distributional fidelity. This formulation enables simulation-free training and allows the diffusion term to adapt to the local dynamical regime. Across several PDE benchmarks, our method consistently improves calibration in short-to-intermediate rollouts compared to fixed-schedule generative baselines, while maintaining comparable mean accuracy under matched compute-constrained inference budgets.

Journal Opt In: No, I do not wish to participate

Submission Number: 58