Control-Augmented Autoregressive Diffusion for Data Assimilation

Despite recent advances in test-time scaling and finetuning of diffusion models, guidance in Auto-Regressive Diffusion Models (ARDMs) remains underexplored. We introduce an amortized framework that augments a pretrained ARDM with a lightweight controller network, trained offline by previewing future rollouts to output stepwise controls that anticipate upcoming observations under a terminal-cost objective. Our approach is motivated by viewing guided generation as an entropy-regularized stochastic optimal control problem over ARDM trajectories: we learn a reusable policy that injects small control corrections inside each denoising sub-step while remaining anchored to the pretrained dynamics. We evaluate this framework in the context of data assimilation (DA) for chaotic spatiotemporal partial differential equations (PDEs), where existing methods can be computationally prohibitive and prone to forecast drift under sparse observations. At inference, DA reduces to a single causal forward rollout with on-the-fly corrections, requiring neither adjoint computations nor gradient-based optimization, and yields an order-of-magnitude speedup over strong diffusion-based DA baselines. Across two canonical PDEs and six observation regimes, our method consistently improves stability, accuracy, and physics-aware fidelity over state-of-the-art baselines. We will release code and checkpoints publicly.

💡 Research Summary

This paper addresses the challenge of guiding pretrained autoregressive diffusion models (ARDMs) for long‑horizon data assimilation (DA) in chaotic spatiotemporal systems. While diffusion models have recently benefited from test‑time scaling and fine‑tuning techniques, guidance mechanisms for ARDMs remain underexplored, especially in settings where small one‑step errors quickly amplify. The authors propose Control‑Augmented Data Assimilation (CAD A), an amortized framework that augments a frozen ARDM with a lightweight controller network. The controller injects additive corrections u(s)ₜ₊₁ into each denoising sub‑step of the diffusion process, steering the trajectory toward upcoming observations while staying close to the original generative dynamics.

The problem is formalized as an entropy‑regularized stochastic optimal control task. The target posterior distribution is P* ∝ Q·exp(−∑ₜ Φ(xₜ; yₜ)/β), where Q denotes the pretrained ARDM prior and Φ measures observation‑state mismatch (e.g., squared error after a degradation operator A). Direct sampling from P* is intractable, so the authors minimize the variational objective C(P)=∑ₜ Eₓ∼P