Control-Augmented Autoregressive Diffusion for Data Assimilation

• URL: http://arxiv.org/abs/2510.06637v1
• Date: Wed, 08 Oct 2025 04:37:32 GMT
• Title: Control-Augmented Autoregressive Diffusion for Data Assimilation
• Authors: Prakhar Srivastava, Farrin Marouf Sofian, Francesco Immorlano, Kushagra Pandey, Stephan Mandt,
• Abstract summary: We introduce an amortized framework that augments pretrained ARDMs with a lightweight controller.<n>We evaluate this framework in the context of data assimilation (DA) for chaotic partial differential equations (PDEs)<n>Our approach reduces DA inference to a single forward rollout with on-the-fly corrections, avoiding expensive adjoint computations and/or optimizations during inference.
• Score: 17.305296093966803
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Out-of-distribution transfer of PDE foundation models to material dynamics under extreme loading [86.6550968435969]
  Most PDE foundation models are pretrained and fine-tuned on fluid-centric benchmarks.<n>We benchmark out-of-distribution transfer on two discontinuity-dominated regimes in which shocks, evolving interfaces, and fracture produce highly non-smooth fields.<n>We evaluate two open-source PDE foundation models, POSEIDON and MORPH, and compare fine-tuning from pretrained weights against training from scratch across training-set sizes to quantify sample efficiency under distribution shift.
  arXiv  Detail & Related papers  (2026-03-04T18:19:35Z)
• Data-driven learning of feedback maps for explicit robust predictive control: an approximation theoretic view [15.111522780173777]
  We establish an algorithm to learn feedback maps from data for a class of robust model predictive control (MPC) problems.<n>We employ a couple of approximation schemes that furnish tight approximations within preassigned uniform error bounds on the admissible state space to learn the unknown feedback policy.
  arXiv  Detail & Related papers  (2025-10-15T13:14:14Z)
• ResAD: Normalized Residual Trajectory Modeling for End-to-End Autonomous Driving [64.42138266293202]
  ResAD is a Normalized Residual Trajectory Modeling framework.<n>It reframes the learning task to predict the residual deviation from an inertial reference.<n>On the NAVSIM benchmark, ResAD achieves a state-of-the-art PDMS of 88.6 using a vanilla diffusion policy.
  arXiv  Detail & Related papers  (2025-10-09T17:59:36Z)
• Revisiting Multivariate Time Series Forecasting with Missing Values [65.30332997607141]
  Missing values are common in real-world time series.<n>Current approaches have developed an imputation-then-prediction framework that uses imputation modules to fill in missing values, followed by forecasting on the imputed data.<n>This framework overlooks a critical issue: there is no ground truth for the missing values, making the imputation process susceptible to errors that can degrade prediction accuracy.<n>We introduce Consistency-Regularized Information Bottleneck (CRIB), a novel framework built on the Information Bottleneck principle.
  arXiv  Detail & Related papers  (2025-09-27T20:57:48Z)
• WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training [64.0932926819307]
  We present Warmup-Stable and Merge (WSM), a framework that establishes a formal connection between learning rate decay and model merging.<n>WSM provides a unified theoretical foundation for emulating various decay strategies.<n>Our framework consistently outperforms the widely-adopted Warmup-Stable-Decay (WSD) approach across multiple benchmarks.
  arXiv  Detail & Related papers  (2025-07-23T16:02:06Z)
• Training-Free Stein Diffusion Guidance: Posterior Correction for Sampling Beyond High-Density Regions [46.59494117137471]
  Training free diffusion guidance provides a flexible way to leverage off-the-shelf classifiers without additional training.<n>We introduce Stein Diffusion Guidance (SDG), a novel training-free framework grounded in a surrogate SOC objective.<n>Experiments on molecular low-density sampling tasks suggest that SDG consistently surpasses standard training-free guidance methods.
  arXiv  Detail & Related papers  (2025-07-07T21:14:27Z)
• Elucidated Rolling Diffusion Models for Probabilistic Weather Forecasting [52.6508222408558]
  We introduce Elucidated Rolling Diffusion Models (ERDM)<n>ERDM is the first framework to unify a rolling forecast structure with the principled, performant design of Elucidated Diffusion Models (EDM)<n>On 2D Navier-Stokes simulations and ERA5 global weather forecasting at 1.5circ resolution, ERDM consistently outperforms key diffusion-based baselines.
  arXiv  Detail & Related papers  (2025-06-24T21:44:31Z)
• Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence [19.484676783876306]
  Diffusion models have emerged as powerful tools for generative modeling.<n>We propose a control framework for fine-tuning diffusion models.<n>We show that PI-FT achieves global convergence at a linear rate.
  arXiv  Detail & Related papers  (2024-12-24T04:55:46Z)
• Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control [25.219524290912048]
  We formulate offline Reinforcement Learning as a two-stage optimization problem. First, we pretrain expressive generative policies on reward-free behavior datasets, then fine-tune these policies to align with task-specific annotations like Q-values. This strategy allows us to leverage abundant and diverse behavior data to enhance generalization and enable rapid adaptation to downstream tasks using minimal annotations.
  arXiv  Detail & Related papers  (2024-07-12T06:32:36Z)
• Amortizing intractable inference in diffusion models for vision, language, and control [89.65631572949702]
  This paper studies amortized sampling of the posterior over data, $mathbfxsim prm post(mathbfx)propto p(mathbfx)r(mathbfx)$, in a model that consists of a diffusion generative model prior $p(mathbfx)$ and a black-box constraint or function $r(mathbfx)$.<n>We prove the correctness of a data-free learning objective, relative trajectory balance, for training a diffusion model that samples from
  arXiv  Detail & Related papers  (2024-05-31T16:18:46Z)
• Data Assimilation in Chaotic Systems Using Deep Reinforcement Learning [0.5999777817331317]
  Data assimilation plays a pivotal role in diverse applications, ranging from climate predictions and weather forecasts to trajectory planning for autonomous vehicles. Recent advancements have seen the emergence of deep learning approaches in this domain, primarily within a supervised learning framework. In this study, we introduce a novel DA strategy that utilizes reinforcement learning (RL) to apply state corrections using full or partial observations of the state variables.
  arXiv  Detail & Related papers  (2024-01-01T06:53:36Z)
• Re-Evaluating LiDAR Scene Flow for Autonomous Driving [80.37947791534985]
  Popular benchmarks for self-supervised LiDAR scene flow have unrealistic rates of dynamic motion, unrealistic correspondences, and unrealistic sampling patterns. We evaluate a suite of top methods on a suite of real-world datasets. We show that despite the emphasis placed on learning, most performance gains are caused by pre- and post-processing steps.
  arXiv  Detail & Related papers  (2023-04-04T22:45:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.