Control-Augmented Autoregressive Diffusion for Data Assimilation
This addresses the problem of computationally prohibitive and drift-prone data assimilation methods for chaotic PDEs, representing a novel method for a known bottleneck.
The paper tackles data assimilation for chaotic spatiotemporal PDEs by augmenting pretrained autoregressive diffusion models with a lightweight controller network, reducing inference to a single forward rollout and outperforming four state-of-the-art baselines in stability, accuracy, and physical fidelity across two PDEs and six observation regimes.
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 pretrained ARDMs with a lightweight controller network, trained offline by previewing future ARDM rollouts and learning stepwise controls that anticipate upcoming observations under a terminal cost objective. We evaluate this framework in the context of data assimilation (DA) for chaotic spatiotemporal partial differential equations (PDEs), a setting where existing methods are often computationally prohibitive and prone to forecast drift under sparse observations. Our approach reduces DA inference to a single forward rollout with on-the-fly corrections, avoiding expensive adjoint computations and/or optimizations during inference. We demonstrate that our method consistently outperforms four state-of-the-art baselines in stability, accuracy, and physical fidelity across two canonical PDEs and six observation regimes. We will release code and checkpoints publicly.