Control-Augmented Autoregressive Diffusion for Data Assimilation

2025-10-08

AI Summary

Addressing the challenge of data assimilation for chaotic spatiotemporal partial differential equations (PDEs) under sparse observations—where existing methods suffer from high computational cost and forecast drift—this paper proposes Control-Enhanced Autoregressive Diffusion (CEDiff). CEDiff introduces a lightweight controller network trained end-to-end offline using a terminal-cost-based stepwise control policy, reducing assimilation inference to a single forward pass with dynamic correction, thereby avoiding adjoint computations and iterative optimization. A future-trajectory preview mechanism enables the controller to predict per-step regulation signals, ensuring efficient and physically consistent integration of observational information. Evaluated on two canonical chaotic PDEs across six sparse observation configurations, CEDiff consistently outperforms four state-of-the-art baselines, achieving superior accuracy, robustness, and dynamical fidelity.

Technology Category

• Intelligent Robots: State Estimation
• Computer Vision: Diffusion Models for Vision
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems

Application Category

• Responsible Web: Data and user privacy-enhancing technologies for the Web
• Web Mining and Content Analysis: Content-based information diffusion
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Enhancing autoregressive diffusion models with control guidance
• Improving data assimilation for chaotic spatiotemporal PDEs
• Reducing computational costs while maintaining physical fidelity

Innovation

• Lightweight controller network augments pretrained ARDMs
• Offline training anticipates observations using terminal cost
• Single forward rollout with corrections replaces adjoint computations