Control-Augmented Diffusion for Autoregressive Data Assimilation

Prakhar Srivastava, Farrin Marouf Sofian, Francesco Immorlano, Stephan Mandt

Track: Track 1: Original Research/Position/Education/Attention Track

Keywords: data assimilation, autoregressive diffusion, stochastic optimal control, amortization

Abstract: Data assimilation (DA) in chaotic spatiotemporal systems, such as turbulent PDEs, is essential but computationally demanding, often requiring expensive adjoints, ensembles, or test-time optimization. We introduce an amortized framework that augments autoregressive diffusion models with learned feedback control. A pretrained diffusion model provides one-step forecasts, while a compact control network, trained offline, injects affine residuals into the DDIM denoising steps. These residuals gently nudge the sampler toward consistency with upcoming observations, preventing forecast drift during long observation gaps. At inference, assimilation reduces to a single forward rollout with on-the-fly corrections, avoiding optimization or ensembles. On chaotic Kolmogorov flow, our method yields improved long-horizon stability, substantial accuracy gains, and over $30\times$ faster runtime. To our knowledge, this is the first framework to integrate amortized assimilation directly into autoregressive diffusion models, opening a new direction for efficient learned control in high-dimensional PDE forecasting.

Submission Number: 230