# Control-Augmented Autoregressive Diffusion for Data Assimilation

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 44 of its 63 references that are in the map.
- First posted on arXiv 8 Oct 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 44)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Denoising Diffusion Restoration Models](https://mlsanity.com/paper/2201.11793/) (2022)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion](https://mlsanity.com/paper/2506.08009/) (2025)
- [Directly Fine-Tuning Diffusion Models on Differentiable Rewards](https://mlsanity.com/paper/2309.17400/) (2023)
- [Diffusion Probabilistic Modeling for Video Generation](https://mlsanity.com/paper/2203.09481/) (2022)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [A General Framework for Inference-time Scaling and Steering of Diffusion Models](https://mlsanity.com/paper/2501.06848/) (2025)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [A Survey on Diffusion Models for Inverse Problems](https://mlsanity.com/paper/2410.00083/) (2024)
- [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://mlsanity.com/paper/2312.15796/) (2023)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling](https://mlsanity.com/paper/2309.15214/) (2023)
- [DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting](https://mlsanity.com/paper/2306.01984/) (2023)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review](https://mlsanity.com/paper/2407.13734/) (2024)
- [Score-based Data Assimilation](https://mlsanity.com/paper/2306.10574/) (2023)
- [Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review](https://mlsanity.com/paper/2501.09685/) (2025)
- [Lie Point Symmetry Data Augmentation for Neural PDE Solvers](https://mlsanity.com/paper/2202.07643/) (2022)
- [Training-free Linear Image Inverses via Flows](https://mlsanity.com/paper/2310.04432/) (2023)
- [DiffDA: a Diffusion Model for Weather-scale Data Assimilation](https://mlsanity.com/paper/2401.05932/) (2024)
- [On conditional diffusion models for PDE simulations](https://mlsanity.com/paper/2410.16415/) (2024)

## Built on by (most cited first, 1 of 1)

- [Iterative Refinement Diffusion for Super-Resolved Data Assimilation of Multiscale Physical Systems](https://mlsanity.com/paper/2608.14744/) (2026)

## Abstract (arXiv)

Despite advances in test-time scaling and diffusion finetuning, guidance for Auto-Regressive Diffusion Models (ARDMs) remains underexplored. We introduce an amortized framework that augments a pretrained ARDM with an offline-trained controller. By previewing future rollouts, the controller learns stepwise corrections that anticipate observations under a terminal-cost objective, yielding a reusable policy for guided generation. Motivated by a stochastic optimal control view of ARDM trajectories, our method injects small controls within each denoising sub-step while staying close to the pretrained dynamics. We study this approach for dataassimilation (DA) in chaotic spatiotemporal partial differential equations (PDEs), where existing methods are often computationally expensive and susceptible to forecast drift under sparse observations. At inference, DA becomes a feed-forward rollout with on-the-fly corrections, achieving an order-of-magnitude speedup over strong diffusion-based baselines. Across two canonical PDEs and a compact ECMWF Reanalysis v5 (ERA5) pilot spanning six observation regimes, our method consistently improves stability and accuracy over state-of-the-art alternatives, with similar improvements observed in a larger-scale GenCast study.