FlowDA: Accurate, Low-Latency Weather Data Assimilation via Flow Matching

Data assimilation (DA) is a fundamental component of modern weather prediction, yet it remains a major computational bottleneck in machine learning (ML)-based forecasting pipelines due to reliance on traditional variational methods. Recent generative ML-based DA methods offer a promising alternative but typically require many sampling steps and suffer from error accumulation under long-horizon auto-regressive rollouts with cycling assimilation. We propose FlowDA, a low-latency weather-scale generative DA framework based on flow matching. FlowDA conditions on observations through a SetConv-based embedding and fine-tunes the Aurora foundation model to deliver accurate, efficient, and robust analyses. Experiments across observation rates decreasing from $3.9%$ to $0.1%$ demonstrate superior performance of FlowDA over strong baselines with similar tunable-parameter size. FlowDA further shows robustness to observational noise and stable performance in long-horizon auto-regressive cycling DA. Overall, FlowDA points to an efficient and scalable direction for data-driven DA.

💡 Research Summary

This paper introduces FlowDA, a novel generative data assimilation (DA) framework for weather prediction that leverages flow matching to address the computational bottleneck of traditional variational DA methods. Data assimilation is the critical process of combining numerical model forecasts (background fields) with real-world observations to produce an optimal initial condition (analysis) for weather prediction. While machine learning-based weather forecasting has dramatically reduced prediction costs, DA remains a slow, optimization-heavy component. Existing generative ML-DA approaches, often based on diffusion models, require hundreds of sampling steps, leading to high latency and error accumulation in long-horizon cycling assimilation.

FlowDA tackles these issues through three key innovations. First, it employs a flow matching generative framework. Instead of iteratively denoising over many steps like diffusion models, flow matching learns a continuous probability path from the background distribution to the analysis distribution, governed by an ordinary differential equation (ODE). This allows high-quality sample generation with far fewer integration steps (e.g., 32), enabling low-latency inference.

Second, it uses a SetConv-based module to handle sparse, irregular observational data. Real observations from satellites and stations are unstructured. FlowDA’s SetConv layer acts as an inverse observation operator, mapping these discrete measurements into a continuous gridded field and an associated observation density map. This decouples the optimal use of observational information from the core assimilation dynamics, enhancing modularity and efficiency.

Third, FlowDA efficiently adapts a large-scale Earth-system foundation model, Aurora, for the DA task. Aurora, pre-trained on massive geophysical datasets, possesses inherent knowledge of atmospheric dynamics and structure. The authors fine-tune it in two parameter-efficient ways: “FlowDA-LoRA” adds lightweight Low-Rank Adaptation (LoRA) layers to the 1.3B-parameter base model (37M tunable parameters), and “FlowDA-Full” fully fine-tunes the 113M-parameter Aurora-Small checkpoint. This transfers the foundation model’s powerful representations to DA at a low cost.

During inference, FlowDA takes a background forecast and sparse observations. The SetConv layer processes the observations into a gridded field. The fine-tuned Aurora model then predicts a velocity field that guides the transformation from the background state to the analysis state via ODE integration.

Experiments were conducted using the ERA5 reanalysis dataset at 0.25° resolution, covering 69 atmospheric variables. Key findings demonstrate FlowDA’s superiority: 1) Accuracy: It outperformed strong baselines like DiffDA (diffusion-based) and VAE-Var across a wide range of observational coverage, from 3.9% down to an extremely sparse 0.1%. 2) Efficiency: Its few-step sampling provides a significant latency advantage over multi-step diffusion models. 3) Robustness: It maintained stable performance when observational noise was introduced. 4) Stability: It exhibited consistent performance in long-horizon auto-regressive cycling DA experiments, crucial for operational pipelines.

In summary, FlowDA presents a significant advancement in data-driven weather data assimilation by combining the fast sampling of flow matching, effective handling of sparse observations via SetConv, and efficient knowledge transfer from a foundation model. It points toward an accurate, low-latency, and scalable future for initializing machine learning-based weather forecasts directly from raw observations.