DAISI: Data Assimilation with Inverse Sampling using Stochastic Interpolants

AI summary

Overview

• Research area: Machine learning for data assimilation (filtering), combining flow-based generative models (stochastic interpolants / diffusion-style models) with ensemble forecast methods for sequential state estimation.
• Technical level: Advanced. The paper assumes familiarity with stochastic differential equations, score-based/flow-based generative modeling, guidance-based conditional sampling, and classical ensemble data assimilation.
• Scope in one sentence: The paper introduces DAISI, a filtering algorithm that couples a stationary, pre-trained generative prior with any forecast ensemble through a novel "inverse sampling" step, then assimilates observations by guided conditional sampling, and evaluates it on Lorenz '63, Surface Quasi-Geostrophic dynamics, and the SEVIR radar dataset.

What This Paper Is About

Data assimilation (DA) estimates the hidden state of a dynamical system by combining imperfect model forecasts with sparse, noisy observations. Classical methods such as the Ensemble Kalman Filter (EnKF) and 4DVar rely on Gaussian approximations or MAP estimation, and particle filters suffer from the curse of dimensionality, so all struggle when observations are sparse, nonlinear, or noisy.

The goal of this paper is a scalable filtering method that can use an expressive, non-Gaussian learned prior without retraining it at every assimilation step, while still incorporating the dynamical information carried by an ensemble forecast.

Key Contributions

1. DAISI algorithm: A filtering algorithm built on flow-based generative models that uses a stationary, pre-trained generative prior (taken to be the invariant measure of the dynamical system) and conditions on new observations via guidance, avoiding retraining at each assimilation step.
2. Inverse sampling: A new step that runs the unconditional backward SDE from the forecast ensemble down to an intermediate time t_min, mapping forecasted states into the latent ("noise") space of the generative prior; these latents then initialize the guided forward SDE. This transfers dynamical information from the ensemble into the generative model without fine-tuning.
3. Theoretical analysis of the hyperparameters: An analysis of the DAISI analysis measure in the limits t_min → 1 and t_min → 0 and of the noise level epsilon, showing that each extreme matches only one component of the true filtering distribution and motivating an intermediate t_min and a small non-zero epsilon. An entropy-dissipation result of Bakry-Émery type is stated in Appendix E.
4. Modular, zero-shot design: The method is compatible with numerical and ML-based forecast and observation models, and supports any flow-based generative model and gradient-based guidance method. Code is released at https://github.com/Erik-Wikingsson/DAISI.

Main Findings

• Forecast information is essential: On Lorenz '63, a DAISI variant without inverse sampling (i.e., sampling directly from the guided invariant-measure posterior P_inf^y) performed substantially worse than the Bootstrap Particle Filter (BPF) on all metrics: RMSE 5.90 ± 1.09, CRPS 7.17 ± 1.38, SSR 1.40 ± 0.25, versus BPF's RMSE 1.18 ± 0.50, CRPS 1.46 ± 0.49, SSR 2.57 ± 0.56.
• Hyperparameter choice matters: With t_min ≈ 0 and epsilon = 0, DAISI underperformed even the no-inversion variant (RMSE 7.04 ± 0.90, CRPS 11.9 ± 1.55, SSR 0.06 ± 0.02) in the reported table.
• Tuning improves accuracy: Tuning t_min alone gave RMSE 2.44 ± 2.89, CRPS 3.58 ± 4.79, SSR 1.07 ± 1.09; jointly tuning t_min and epsilon gave the best DAISI results and the closest match to BPF: RMSE 2.03 ± 1.01, CRPS 2.61 ± 1.38, SSR 1.08 ± 0.41.
• epsilon maintains ensemble spread: Re-running with the tuned t_min but setting epsilon = 0 produced similar mean behaviour but narrower uncertainty bands (RMSE 2.05 ± 1.31, CRPS 3.19 ± 2.14, SSR 0.32 ± 0.09), showing that epsilon matters for spread and CRPS while t_min matters for accuracy.
• Toy-problem ablation supports the trade-off: In a 1D toy problem with epsilon = 0, t_min = 0.01 produced a distribution resembling P_inf^y, t_min = 0.6 matched hat{pi}_n, and the intermediate t_min = 0.3 best aligned with the filtering distribution, with a noticeably lower Maximum Mean Discrepancy (MMD) score.
• Comparison to SDEdit: In Appendix C.4.6, replacing the backward SDE with SDEdit is reported to give substantially worse results.
• Complexity: DAISI costs O(J(f(d_x) + T g(d_x, d_y))), compared with O(J(f(d_x) + (d_x + d_y)J + J^2)) for the ensemble transform Kalman filter (ETKF). Because the U-net drift evaluation scales linearly with input dimension, DAISI achieves comparable cost to ETKF when T ∼ J and scales linearly in ensemble size J rather than cubically. For small ensembles (e.g., O(10)), DAISI can be more expensive because T = O(10^2) Euler–Maruyama steps are needed. The forecast, inverse sampling, and guided sampling steps are all embarrassingly parallelizable.
• SQG and SEVIR experiments: DAISI is evaluated on SQG (64 × 64 grid, 100 time steps at 3-hour intervals, plus one 256 × 256 high-dimensional case) and on SEVIR, across noisy, sparse, multimodal, saturating, high-dimensional, and radar configurations, against LETKF, FlowDAS, EnSF, SDA (filtering and smoothing), and a DAISI-ML variant that replaces the numerical model with the learned FlowDAS model. The numeric CRPS values for these experiments are given in Table 3, but the table content was truncated in the material provided here, so the specific values are not reported in this summary.

Methodology in Plain English

The starting point is a generative model trained to produce samples from the long-run (invariant) behaviour of the dynamical system — a "background measure" that plays the role a static climatological covariance plays in hybrid ensemble-variational DA, but is far more expressive than a Gaussian.

Filtering then alternates two steps:

1. Forecast: Push each ensemble member forward one time step with the dynamics model (numerical or ML-based) to get forecast particles.
2. Analysis: (a) Run the generative model's backward SDE from each forecast particle down to an intermediate time t_min, turning each forecast into a latent noise-space vector that encodes the forecast information. (b) Run the guided forward SDE from those latents up to time 1, using the observation to steer the sampling. The output particles approximate the filtering distribution while still carrying information from the forecast ensemble.

Two knobs control the behaviour. The stopping time t_min sets how much the analysis resembles a pure conditional sample from the learned prior (t_min → 1 keeps more forecast, t_min → 0 keeps more prior) — extremes match only part of the true filtering distribution, so an intermediate value is best. The noise level epsilon controls resampling: at epsilon = 0 the particles are deterministic and can collapse onto a single state, while small non-zero epsilon injects diversity but still keeps the particles close to the forecasts (as epsilon → ∞, all ensemble information is lost).

Why This Matters

• Research impact: The work shows how stationary, pre-trained flow-based generative priors can be brought into sequential online filtering without retraining or fine-tuning, bridging offline inverse-problem techniques with operational DA, and provides an analysis of the bias trade-off introduced by the inverse-sampling scheme.
• Weather forecasting: The paper motivates state estimation with chaos, high dimensionality, and nonlinear noisy measurements, citing early warning systems and renewable energy management as downstream users.
• Geophysical and turbulent flows: Evaluated on Surface Quasi-Geostrophic dynamics, a standard benchmark for turbulent geophysical flows.
• Real-world radar data: Evaluated on SEVIR, a real-world radar dataset.
• Industrial relevance: Because DAISI is zero-shot compatible with numerical and ML-based forecast and observation models, it can be dropped on top of existing forecasting pipelines (including learned forecasters like FlowDAS) without a bespoke adjoint model or per-step retraining — useful for companies running weather, energy, or robotics forecasting stacks. The released code (https://github.com/Erik-Wikingsson/DAISI) lowers the barrier to adoption.

Future Directions

• Closing the residual gap to exact filters: DAISI's tuned results approach but do not match BPF on Lorenz '63; understanding and reducing the remaining bias in the analysis measure is an open question.
• Guarantees for the invariant measure assumption: The paper notes that existence and ergodicity of the invariant measure P_inf are not generally guaranteed, though reasonable for many geophysical systems — establishing conditions would broaden applicability.
• Hyperparameter selection at scale: t_min and epsilon were tuned on a short held-out validation trajectory to minimize CRPS; a principled, scalable selection rule (especially in high dimensions) is a natural next step.
• Extending beyond Gaussian observation operators: The SQG and SEVIR experiments stay within Gaussian observations because the guidance methods used have been applied almost exclusively in that setting, though the authors state this is not a limitation of DAISI itself.

Target Audience

Researchers and practitioners in machine learning for scientific computing, data assimilation, and geophysical/atmospheric modeling; graduate students and engineers working with diffusion or flow-based generative models applied to inverse and sequential problems; and industry teams building forecasting systems who want a modular, retraining-free way to combine learned priors with existing numerical or ML forecast models. Readers should be comfortable with SDEs and score-based generative modeling.

Authors’ abstract

Data assimilation (DA) is a cornerstone of scientific and engineering applications, combining model forecasts with sparse and noisy observations to estimate latent system states. Classical high-dimensional DA methods, such as the ensemble Kalman filter, rely on Gaussian approximations that are violated for complex dynamics or observation operators. To address this limitation, we introduce DAISI, a scalable filtering algorithm built on flow-based generative models that enables flexible probabilistic inference using data-driven priors. The core idea is to use a stationary, pre-trained generative prior that first incorporates forecast information through a novel inverse-sampling step, before assimilating observations via guidance-based conditional sampling. This allows us to leverage any forecasting model as part of the DA pipeline without having to retrain or fine-tune the generative prior at each assimilation step. Experiments on challenging nonlinear systems show that DAISI achieves accurate filtering results in regimes with sparse, noisy, and nonlinear observations where traditional methods struggle. The code for DAISI is available at https://github.com/Erik-Wikingsson/DAISI.