FlowDAS: A Stochastic Interpolant-based Framework for Data Assimilation

TL;DR

FlowDAS introduces a stochastic interpolant-based framework for data assimilation that learns state transition dynamics directly from data, enabling more accurate and physically consistent forecasts in complex systems.

Contribution

It proposes a novel step-by-step stochastic interpolant approach for data assimilation, improving modeling of unknown dynamics and measurement consistency.

Findings

Outperforms traditional model-driven methods in accuracy.

Achieves better physical plausibility in forecasts.

Effective in complex, real-world scenarios like weather prediction.

Abstract

Data assimilation (DA) integrates observations with a dynamical model to estimate states of PDE-governed systems. Model-driven methods (e.g., Kalman, particle) presuppose full knowledge of the true dynamics, which is not always satisfied in practice, while purely data-driven solvers learn a deterministic mapping between observations and states and therefore miss the intrinsic stochasticity of real processes. Recently, score-based diffusion models learn a global diffusion prior and provide a good modeling of the stochastic dynamics, showing new potential for DA. However, their all-at-once generation rather than step-by-step transition limits their performance when dealing with highly complex stochastic processes and lacks physical interpretability. To tackle these drawbacks, we introduce FlowDAS, a generative DA framework that uses stochastic interpolants to directly learn state transition dynamics and achieve step-by-step transition to better model the real dynamics. We also improve the framework by combining the observation, better suiting the DA settings. Directly learning the underlying dynamics from collected data removes restrictive dynamical assumptions, and conditioning on observations at each interpolation step yields stable, measurement-consistent forecasts. Experiments on Lorenz-63, Navier-Stokes super-resolution/sparse-observation scenarios, and large-scale weather forecasting -- where dynamics are partly or wholly unknown -- show that FlowDAS surpasses model-driven methods, neural operators, and score-based baselines in accuracy and physical plausibility.