Control-Augmented Autoregressive Diffusion for Data Assimilation

📄 arXiv: 2510.06637v1 📥 PDF

作者: Prakhar Srivastava, Farrin Marouf Sofian, Francesco Immorlano, Kushagra Pandey, Stephan Mandt

分类: cs.LG, cs.AI, cs.CV

发布日期: 2025-10-08

💡 一句话要点

提出控制增强自回归扩散模型，用于解决数据同化中预测漂移问题。

🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)

关键词: 自回归扩散模型 数据同化 控制增强 偏微分方程 预测漂移

📋 核心要点

  1. 现有数据同化方法在处理混沌偏微分方程时，计算成本高昂，且在稀疏观测下易产生预测漂移。
  2. 提出一种控制增强自回归扩散模型，通过轻量级控制器网络预测未来观测，实现即时校正。
  3. 实验表明，该方法在稳定性、准确性和物理保真度方面优于现有方法，且避免了昂贵的计算。

📝 摘要（中文）

尽管扩散模型在测试时缩放和微调方面取得了进展，但自回归扩散模型（ARDM）中的引导仍然未被充分探索。我们引入了一种摊销框架，通过一个轻量级控制器网络来增强预训练的ARDM。该控制器网络通过预览未来的ARDM展开并学习逐步控制来预测即将到来的观测，从而在终端成本目标下进行离线训练。我们在混沌时空偏微分方程（PDE）的数据同化（DA）背景下评估了该框架，在这种情况下，现有方法通常在计算上是禁止的，并且在稀疏观测下容易出现预测漂移。我们的方法将DA推理简化为单次前向展开和即时校正，避免了推理期间昂贵的伴随计算和/或优化。我们证明了我们的方法在两个典型的PDE和六个观测方案中，在稳定性、准确性和物理保真度方面始终优于四种最先进的基线方法。我们将公开发布代码和检查点。

🔬 方法详解

问题定义：论文旨在解决数据同化（DA）问题，特别是在混沌时空偏微分方程（PDE）的背景下。现有的DA方法，如卡尔曼滤波及其变体，在处理高维、非线性系统时计算成本非常高昂。此外，在观测数据稀疏的情况下，这些方法容易出现预测漂移，导致长期预测结果与真实情况偏差较大。

核心思路：论文的核心思路是利用自回归扩散模型（ARDM）生成高质量的预测，并引入一个轻量级的控制器网络来对ARDM的生成过程进行引导。该控制器通过观察ARDM的中间状态和未来的观测数据，学习如何逐步调整ARDM的生成过程，从而实现对预测结果的校正。这种方法避免了传统DA方法中昂贵的伴随计算和优化过程。

技术框架：整体框架包含两个主要部分：预训练的ARDM和一个轻量级的控制器网络。ARDM负责生成初始的预测轨迹，控制器网络则负责根据观测数据对ARDM的生成过程进行实时校正。控制器网络通过预览ARDM未来的展开轨迹，学习在每一步采取何种控制动作，以最小化终端成本。整个过程可以看作是一个强化学习问题，控制器网络充当策略网络，ARDM充当环境。

关键创新：最重要的技术创新在于将控制理论与自回归扩散模型相结合，提出了一种控制增强的扩散模型。与传统的DA方法相比，该方法避免了昂贵的伴随计算和优化过程，实现了高效的数据同化。此外，通过离线训练控制器网络，可以在推理阶段实现快速的即时校正。

关键设计：控制器网络是一个轻量级的神经网络，输入包括ARDM的当前状态和未来的观测数据，输出是控制信号，用于调整ARDM的生成过程。损失函数采用终端成本的形式，鼓励控制器网络在整个预测过程中逐步逼近真实状态。具体的网络结构和参数设置需要根据具体的PDE问题进行调整。

📊 实验亮点

实验结果表明，该方法在两个典型的偏微分方程（PDE）和六个观测方案中，在稳定性、准确性和物理保真度方面始终优于四种最先进的基线方法。具体而言，该方法能够显著减少预测漂移，提高长期预测的准确性，并且在计算效率方面具有显著优势。

🎯 应用场景

该研究成果可应用于气象预报、海洋环境监测、金融风险管理等领域。通过更准确地同化观测数据，可以提高预测的准确性和可靠性，为决策提供更可靠的依据。未来，该方法有望扩展到其他类型的动态系统，并与其他数据同化技术相结合，进一步提升数据同化的性能。

📄 摘要（原文）

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.