# PnP-DA: Towards Principled Plug-and-Play Integration of Variational Data Assimilation and Generative Models

arXiv: https://arxiv.org/abs/2508.00325 · HTML: https://mlsanity.com/paper/2508.00325/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 47 references that are in the map.
- First posted on arXiv 1 Aug 2025.

## Builds on (most cited first, 20 of 20)

- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [A Variational Perspective on Solving Inverse Problems with Diffusion Models](https://mlsanity.com/paper/2305.04391/) (2023)
- [Plug-and-Play Methods Provably Converge with Properly Trained Denoisers](https://mlsanity.com/paper/1905.05406/) (2019)
- [Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors](https://mlsanity.com/paper/2405.18782/) (2024)
- [PnP-Flow: Plug-and-Play Image Restoration with Flow Matching](https://mlsanity.com/paper/2410.02423/) (2024)
- [Machine learning with data assimilation and uncertainty quantification for dynamical systems: a review](https://mlsanity.com/paper/2303.10462/) (2023)
- [DiffDA: a Diffusion Model for Weather-scale Data Assimilation](https://mlsanity.com/paper/2401.05932/) (2024)
- [An Ensemble Score Filter for Tracking High-Dimensional Nonlinear Dynamical Systems](https://mlsanity.com/paper/2309.00983/) (2023)
- [Gradient-Free Generation for Hard-Constrained Systems](https://mlsanity.com/paper/2412.01786/) (2024)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [Deep Generative Data Assimilation in Multimodal Setting](https://mlsanity.com/paper/2404.06665/) (2024)
- [Generative Data Assimilation of Sparse Weather Station Observations at Kilometer Scales](https://mlsanity.com/paper/2406.16947/) (2024)
- [Fuxi-DA: A Generalized Deep Learning Data Assimilation Framework for Assimilating Satellite Observations](https://mlsanity.com/paper/2404.08522/) (2024)
- [Online model error correction with neural networks in the incremental 4D-Var framework](https://mlsanity.com/paper/2210.13817/) (2022)
- [4D-Var using Hessian approximation and backpropagation applied to automatically-differentiable numerical and machine learning models](https://mlsanity.com/paper/2408.02767/) (2024)
- [Strictly Constrained Generative Modeling via Split Augmented Langevin Sampling](https://mlsanity.com/paper/2505.18017/) (2025)
- [Joint Parameter and Parameterization Inference with Uncertainty Quantification through Differentiable Programming](https://mlsanity.com/paper/2403.02215/) (2024)
- [Ensemble Kalman filter in latent space using a variational autoencoder pair](https://mlsanity.com/paper/2502.12987/) (2025)
- [Neural Incremental Data Assimilation](https://mlsanity.com/paper/2406.15076/) (2024)
- [VAE-Var: Variational-Autoencoder-Enhanced Variational Assimilation](https://mlsanity.com/paper/2405.13711/) (2024)

## Abstract (arXiv)

Earth system modeling presents a fundamental challenge in scientific computing: capturing complex, multiscale nonlinear dynamics in computationally efficient models while minimizing forecast errors caused by necessary simplifications. Even the most powerful AI- or physics-based forecast system suffer from gradual error accumulation. Data assimilation (DA) aims to mitigate these errors by optimally blending (noisy) observations with prior model forecasts, but conventional variational methods often assume Gaussian error statistics that fail to capture the true, non-Gaussian behavior of chaotic dynamical systems. We propose PnP-DA, a Plug-and-Play algorithm that alternates (1) a lightweight, gradient-based analysis update (using a Mahalanobis-distance misfit on new observations) with (2) a single forward pass through a pretrained generative prior conditioned on the background forecast via a conditional Wasserstein coupling. This strategy relaxes restrictive statistical assumptions and leverages rich historical data without requiring an explicit regularization functional, and it also avoids the need to backpropagate gradients through the complex neural network that encodes the prior during assimilation cycles. Experiments on standard chaotic testbeds demonstrate that this strategy consistently reduces forecast errors across a range of observation sparsities and noise levels, outperforming classical variational methods.