Control-Augmented Autoregressive Diffusion for Data Assimilation
Extracted source text · arXiv 2510.06637
\documentclass{article}
\usepackage[preprint]{neurips_2025}
\input{math_commands.tex}
\usepackage{hyperref}
\hypersetup{
colorlinks=true,
citecolor=magenta,
urlcolor=blue,
linkcolor=black
}
\usepackage{url}
\usepackage{booktabs}
\usepackage{amsfonts}
\usepackage{nicefrac}
\usepackage{microtype}
\usepackage[table]{xcolor}
\usepackage{amsmath}
\usepackage{mathtools}
\usepackage{tabularx}
\usepackage{booktabs}
\usepackage{amssymb}
\usepackage{algorithm}
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{stfloats}
\usepackage{algpseudocode}
\algrenewcommand\algorithmicrequire{\textbf{Input:}}
\algrenewcommand\algorithmicensure{\textbf{Output:}}
\usepackage{siunitx,booktabs,tabularx}
\sisetup{
detect-weight=true,
detect-inline-weight=math,
table-number-alignment = center,
table-space-text-post = \,,
}
\newcolumntype{Y}{>{\centering\arraybackslash}X}
\definecolor{lightgray}{gray}{0.9}
\newcommand{\placeholder}[1]{\color{gray}\fbox{\ttfamily #1}}
\newcommand{\XX}[1]{\textcolor{gray}{\textsf{#1}}}
\title{Control-Augmented Autoregressive\\Diffusion for Data Assimilation}
\author{
\textbf{Prakhar Srivastava} \quad
\textbf{Farrin Marouf Sofian} \quad
\textbf{Francesco Immorlano}\\
\textbf{Kushagra Pandey} \quad
\textbf{Stephan Mandt} \\
University of California, Irvine \\
\texttt{\{prakhs2, fmaroufs, fimmorla, pandeyk1, mandt\}@uci.edu}
}
\begin{document}
\maketitle
\begin{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 \emph{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.
\end{abstract}
%Input tex files
\input{tex/introduction}
\input{tex/background}
\input{tex/method}
\input{tex/related}
\input{tex/experiments}
\input{tex/conclusion}
\bibliographystyle{neurips_2025}
\bibliography{neurips_2025}
\appendix
\input{tex/appendix}
\end{document}
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
\newcommand{\kp}{\textcolor{red}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
\newcommand{\kl}[2]{D_{\mathrm{KL}}\!\left(#1 ~ \| ~ #2\right)}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\newcommand{\sm}[1]{{\color{blue} SM: #1}}
\section{Introduction}
Recently, much progress has been made on inference-time scaling \citep{uehara2025inferencetimealignmentdiffusionmodels, chung2023diffusion} and finetuning \citep{xu2023imagereward, clark2023directly, uehara2024understandingreinforcementlearningbasedfinetuning, domingo-enrich2025adjoint} in diffusion models. However, while Auto-Regressive Diffusion Models (ARDMs) \citep{Ge2022, yang2023diffusion, yu2023language, huang2025selfforcing} form the backbone of many scientific and engineering applications \citep{pathak2024kilometerscaleconvectionallowingmodel, Mardani2025}, exploring different strategies for finetuning in the context of ARDMs remains in its infancy.
In this work, we build on recent formulations of diffusion model guidance as stochastic control \citep{pandey2025variational, rout2025rbmodulation} and propose a finetuning framework for ARDMs that embeds a learned control mechanism directly into the generative dynamics. More specifically, a pretrained autoregressive diffusion model provides the backbone transitions, while a \emph{controller} network---finetuned in a separate stage---injects affine \emph{controls} into each denoising step, guiding the autoregressive dynamics to satisfy a terminal cost. We explore our proposed framework in the context of data assimilation (DA) in chaotic spatiotemporal PDEs. In DA, the goal is to forecast high-dimensional dynamics---like PDE states or global atmospheric weather states---while taking into account observational data, e.g., from sensors, weather stations, or satellites. Without such continuous adaptation, a forecast quickly diverges from the ground truth trajectory---even under almost perfectly specified initial conditions---because of the chaotic nature of the underlying dynamical system \citep{kalnay_2002_atmospheric, carrassi2018data, evensen2022data}.
\begin{figure*}[!htbp]
\centering
\begin{subfigure}[t]{0.68\textwidth}
\centering
\includegraphics[width=\linewidth]{figures/overview_autoreg.png}
\caption{Method overview}
\label{fig:overview:a}
\end{subfigure}\hfill
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\linewidth]{figures/spacetime_grid_1d.png}
\caption{Qualitative results}
\label{fig:overview:b}
\end{subfigure}
\caption{
% (a) \textbf{Overview.} We augment a pretrained autoregressive diffusion model (ARDM) with an offline-trained \emph{controller} that injects stepwise controls during denoising, guided by previewed future observations within a fixed window.
(a) \textbf{Overview.} A pretrained ARDM forms an autoregressive chain $\rvx_{t_0}\!\to\!\rvx_{t_0+1}\!\to\!\dots\!\to\!\rvx_{t_0+\Lambda}$. At each timestep, denoising sub-steps refine noise to data while a \emph{controller} injects affine controls conditioned on the preview of sparse future observations $\rvy_\tau$ within the active $\Lambda$ window, enabling small anticipatory nudges that steer forecasts toward consistency with $\rvy_\tau$ while preserving causality (no information beyond the preview window is accessed).
(b) \textbf{Stability.} Our method yields superior long-horizon rollouts on the 1D Kuramoto–Sivashinsky PDE (horizon 640) under sparse spatiotemporal observations (green dots). Darker colors indicate lower forecast error.
}
\label{fig:overview}
\end{figure*}
DA counters this by incorporating sparse, noisy observations, and producing improved analyses that extend predictability \citep{wang2025phyda, carrassi2018data}. Classical schemes, such as ensemble Kalman filters or variational approaches \citep{le1986variational, courtier1998ecmwf, tr2006accounting}, perform posterior inference, typically under Gaussian assumptions. They either optimize over sliding assimilation windows or update ensemble sequentially, demanding costly adjoint or ensemble computations \citep{wang2021physics}. However, when the underlying predictive model is a conditional diffusion model \citep{dhariwal2021diffusion}, one can leverage the rich methodology of guidance \citep{chung2023diffusion, uehara2025inferencetimealignmentdiffusionmodels} to explicitly account for the lossy measurement process and naturally cast DA as a sequential inverse problem. Recent works apply diffusion for Bayesian DA, either by guiding sampling at test time \citep{rozet2023score, qu2024deep, manshausen2025generative} or by conditioning training directly on observations \citep{huang2024diffda}. While promising, these approaches remain limited: inference-only guidance only corrects at observation times, allowing forecasts to drift in between ; and naïve conditional training destabilizes learning. To address these issues, we make the following contributions.
\begin{itemize}
\item We propose a diffusion-based DA framework that introduces a learned control mechanism into the generative dynamics of a pretrained ARDM (see Fig.~\ref{fig:overview} for an overview)
\item We train the controller in an offline fashion on synthetic assimilation scenarios. During inference, our system performs causal, feed-forward rollouts avoiding any further computationally-intensive optimization, thus enabling accurate and stable assimilation in chaotic PDE forecasting.
\item We demonstrate that our method outperforms four state-of-the-art baselines in terms of stability and accuracy across both short- and long-horizon rollouts on two canonical PDE datasets under six observation regimes, and further show that it more closely adheres to domain-standard physical diagnostics.
\end{itemize}
\section{Background}
\subsection{Autoregressive Diffusion Models}
Given a \emph{noising kernel} $p(\rvx^{(s)}|\rvx^{(0)})=\gN(\mu_s\rvx^{(0)}, \sigma_s^2\mI_d)$ which converts data to noise, diffusion models \citep{sohl2015deep, ho2020denoising} learn a corresponding \emph{denoising process},
\begin{equation}
\label{eq:unguided-diffusion}
q(\rvx^{(0:S)}) = q(\rvx^{(S)})\prod_{s=0}^{S-1} q(\rvx^{(s)}|\rvx^{(s+1)}),
\end{equation}
which transforms noise to data. The reverse diffusion posterior is specified as $q(\rvx^{(s)}|\rvx^{(s+1)}) = \gN(\vmu_\theta(\rvx^{(s+1)}, s+1), \sigma_{s+1}^2\mI_d)$ where $\vmu_\vtheta(.,.)$ is learned via denoising score matching \citep{hyvarinen2005estimation,vincent2011connection,song2019generative}.
In this paper, we consider autoregressive diffusion models in the \emph{time} domain. While upper indices referred to the noising/denoising steps $s$, the additional time indices are now denoted by \emph{lower} indices $t$. To this end, we consider an autoregressive process,
\begin{equation}
\label{eq:Qtheta-finite}
\gQ_\theta \;=\; p_0(\rvx_{0})\;\prod_{t \ge 0}^{} q_\theta(\rvx_{t+1}\mid \rvx_{t}),
\end{equation}
with $p_0(\rvx_0)$ being an initial distribution. Each transition distribution $q_\theta(\rvx_{t+1}\mid \rvx_t)$ is defined by a conditional diffusion model. More formally,
\begin{equation}
q_\theta(\rvx_{t+1}\!\mid \rvx_t)
=\!
\int \!\Bigg[
\prod_{s=0}^{S-1}
q_\theta\big(\rvz_{t+1}^{(s)} \,\big|\, \rvz_{t+1}^{(s+1)};\, \rvx_t\big)
\Bigg]\,
p\big(\rvz_{t+1}^{(S)}\big)\;
d \rvz_{t+1}^{(1:S)}.
\label{eq:unguided_ardm}
\end{equation}
The latents $\rvz_{t+1}^{(s)}$ represent intermediate noisy states and $p\big(\rvz_{t+1}^{(S)}\big)$ represents the prior distribution for sampling initial noise. Thus, at each autoregressive timestep $t$, the next state $\rvx_{t+1}\equiv\rvz^{(0)}_{t+1}$ is obtained by simulating $S$ diffusion sampling steps conditioned on $\rvx_t$.
\subsection{Guidance in Diffusion models}
Given a pretrained diffusion model, it is often desirable to guide the diffusion process conditioned on input $\rvy$.
The conditional score to be estimated can be usually decomposed as the sum of the unconditional score from the pretrained diffusion model and the score of the noisy likelihood distribution $p(\rvy|\rvx_t)$.
In classifier guidance \citep{dhariwal2021diffusion}, the noisy likelihood distribution can be modeled by training a classifier. However, it is also common to estimate this likelihood via $p(\rvy|\rvx_t)=\int p\left(\rvx^{(0)}|\rvx^{(t)}\right)p\left(\rvy|\rvx^{(0)}\right) d\rvx^{(0)}$. For example, Diffusion Posterior Sampling (DPS) \citep{chung2023diffusion} approximates the diffusion posterior as, $p\left(\rvx^{(0)}|\rvx^{(t)}\right) = \delta\left(\E\left[\rvx^{(0)}|\rvx^{(t)}\right]\right)$. More expressive approximations of the diffusion posterior were explored in \citep{song2022pseudoinverse, pandey2024fast} and presented in more detail in \citep{daras2024survey}
\section{Control Augmented Data Assimilation (CADA)}
\label{sec:method}
\vspace{-0.3em}
\paragraph{Notation.}
We denote physical time indices as $t \in \mathbb{N}$; diffusion denoising sub-steps $s \in \{S{-}1,\dots,0\}$; state space $\rvx_t$;
observation arrival indices $\mathcal{T}\subseteq\mathbb{N}$ with observations $\rvy_{\gT}\triangleq\{\rvy_\tau\}_{\tau\in\gT}$.
\subsection{Problem: Chaotic Forecasting with Delayed, Sparse Observations}
\label{subsec:problem}
Forecasting chaotic systems with ARDMs is difficult because trajectories diverge exponentially from the underlying dynamics. Data assimilation (DA) mitigates this instability by periodically steering forecasts back toward trajectories consistent with partial ground-truth observations. In practice, however, these observations arrive only at sparse timesteps, leaving long gaps in the autoregressive rollout without direct correction. Relying solely on past observations therefore allows forecasts to drift substantially before the next arrival. To address this, we assume access to all future observations within a fixed window $\Lambda$, analogous to finite assimilation windows in DA (e.g., fixed-lag smoothing in 4D-Var/EnKS or offline conditional DA in \citet{shysheya2024conditional}). Each observation $\rvy_\tau$ is further accompanied by operator-specific metadata, such as masks used in linear masking.
\subsection{Unguided and guided autoregressive dynamics}
\begin{figure}[t]
\centering
\includegraphics[width=0.95\linewidth]{./figures/snapshot_grid_2d_compact_flipped.png}
\caption{
\textbf{Our method improves rollout stability and reconstruction consistency over long horizons.} Qualitative results on 2D Kolmogorov flow (horizon 180) show snapshots at representative timesteps, including regions with and without observations. While most baselines drift around $t \approx 35$, our approach preserves sharp reconstructions. Joint-AR avoids catastrophic divergence but loses fine-scale structures at later steps, increasingly so for longer horizons (see Tab.~\ref{tab:rmsd_clean},\ref{tab:hct_clean}).
}
\label{fig:compare}
\end{figure}
Given the DA setting in Section \ref{subsec:problem}, our main idea is to augment the pretrained ARDM dynamics with a \emph{controller} network. Intuitively, the \emph{controller} applies anticipatory corrections to the unguided autoregressive dynamics, essentially guiding the predictions at any given timestep toward consistency with the next observation arrival, while staying close to the unguided ARDM dynamics. Consequently, we refer to the resulting dynamics as \emph{guided ARDM}. This aligns with existing approaches for finetuning diffusion models \citep{uehara2024understandingreinforcementlearningbasedfinetuning}, which tilt the learned distribution to maximize a reward model while regularizing to stay close to the unguided diffusion dynamics.
With this intuition in place, we now define the unguided and guided ARDM dynamics.
\paragraph{Unguided Autoregressive Dynamics.}
We refer to Eq. \ref{eq:unguided_ardm} as the unguided ARDM dynamics.
\paragraph{Guided Autoregressive Dynamics.}
Next, at any timestep $t$, we augment the unguided one-step transition kernel $q_\theta\big(\rvz_{t+1}^{(s)} \,\big|\, \rvz_{t+1}^{(s+1)};\, \rvx_t\big)$, with \emph{controls} $\rvu_{t+1}^{(s)}$ and define,
\begin{equation}
\label{eq:controlled-kernel}
p(\rvx_{t+1}\!\mid \rvx_t;\,\mU_{t+1})
=\!
\int \!\Bigg[\prod_{s=0}^{S-1}
p\big(\rvz_{t+1}^{(s)} \,\big|\, \rvz_{t+1}^{(s+1)};\,\rvu_{t+1}^{(s)},\rvx_t\big)\Bigg]\,
p\big(\rvz_{t+1}^{(S)}\big)\;
d \rvz_{t+1}^{(1:S)}.
\end{equation}
where for brevity, we denote $\mU_{t+1} = (\rvu_{t+1}^{(0)},\dots,\rvu_{t+1}^{(S-1)})$. Furthermore, following \citet{pandey2025variational}, we parameterize the guided diffusion posterior in Eq. \ref{eq:controlled-kernel} as,
\begin{equation}
\label{eq:controlled-substep}
p\big(\rvz_{t+1}^{(s)} \,\big|\, \rvz_{t+1}^{(s+1)};\,\rvu_{t+1}^{(s)},\rvx_t\big)
\;\triangleq\;
q_\theta\Big(\rvz_{t+1}^{(s)} \,\Big|\, \rvz_{t+1}^{(s+1)} + \gamma \rvu_{t+1}^{(s)};\;\rvx_t\Big).
\end{equation}
where $\gamma \in \R$ is a scalar hyperparameter of our method which perturbs the noisy state $\rvz_{t+1}^{(s+1)}$ in the direction of the controls $\rvu_{t+1}^{(s)}$. Therefore, we define \emph{guided ARDM} dynamics as,
\begin{equation}
\label{eq:controlled-path}
\gP
\;=\;
p_0(\rvx_0)\;\prod_{t\ge 0}
p\big(\rvx_{t+1}\,\big|\,\rvx_t;\,\mU_{t+1}\big).
\end{equation}
Next, we discuss learning the controls using upcoming observations $\rvy_t$ within a fixed time horizon.
\subsection{Learning the controls}
We begin by connecting incoming future observations to our optimization objective, and then describe how we parameterize the controls $\mU_{t+1}$.
\paragraph{Observations as arrival-time costs.}
The main intuition behind our approach is to tilt the unguided ARDM dynamics towards a reward which ensures consistency between the between the forecast state $\rvx_\tau$ and the corresponding upcoming observation $\rvy_\tau$ at arrival time $\tau$ (defined via $\Phi$). More formally, we want to sample from the following tilted distribution,
%
\begin{equation}
\gP^* \propto \gQ_\theta\exp{(-\frac{1}{\beta}\sum_\tau \Phi(\rvx_\tau; \rvy_\tau))}
\label{eq:tilt_distribution}
\end{equation}
which corresponds to solving the following optimization problem (proof in App.~\ref{app:proof-tilt}),
\begin{equation}
\label{eq:arrival-cost}
\gC(\rvx) \;\triangleq\; \sum_{\tau\in\gT} \E_{\rvx_\tau\sim \gP}\Big[\Phi(\rvx_\tau;\,\rvy_\tau)\Big] + \beta\kl{\gP}{\gQ_\theta},
\end{equation}
where $\gP$ denotes the guided ARDM process defined in Eq. \ref{eq:controlled-path}. The first term penalizes the mismatch between the forecast state and the corresponding observation while the second term penalizes the deviation of the guided from the unguided ARDM dynamics and acts as a regularizer.
Intuitively, the optimization objective in Eq.~\ref{eq:arrival-cost} \emph{tilts} the base unguided ARDM dynamics towards minimizing the aggregate cost $\Phi(\rvx_t; \rvy_t)$. This cumulative cost provides the assimilation signal for learning the controls $\mU_{t+1}$.
In our experiments, $\Phi$ is instantiated as a linear masking or downsampling operator with a least-squares penalty (see App.~\ref{app:obs-ops}). Moreover, we impose differentiability constraints on $\Phi(\cdot\,;\rvy_t)$, to enable gradient based learning of the controls.
\paragraph{Why Direct Tilting is Intractable for guided ARDM?}
The tilted posterior in Eq. \ref{eq:tilt_distribution} is in principle the optimal distribution, but computing its normalization constant $Z_\beta$ or sampling from it exactly is intractable. Moreover, importance sampling from $\gQ_\theta$ would require an intractable number of full autoregressive rollouts of the diffusion backbone to obtain reliable estimates of the importance weights which in chaotic regimes, would likely concentrate on a vanishing fraction of trajectories. Therefore, we instead define a tractable parametric family of distributions $\gP_\psi$, by injecting amortized controls into $\gQ_\theta$. This retains the pretrained diffusion prior while enabling lightweight corrections without direct sampling from the tilted posterior. Therefore, given upcoming observations $\rvy_\tau$, previous context $\rvx_{t-1}$ and the current diffusion noisy state $\rvz_t^{(s)}$, we amortize the controller (with learnable parameters $\psi$) as follows,
\begin{equation}
\rvu_{t}^{(s)} = \vu_\psi\!\big(\rvz_{t}^{(s)},\;\rvx_{t-1},\;\rvy_\tau,\;s,\;\tau-t\big),
\end{equation}
%
\paragraph{Why amortize the controls?}
While a natural strategy is to optimize controls per trajectory at test time,
in an autoregressive setting
the state at an arrival index $\tau\in\mathcal{T}$ depends on the entire preceding trajectory, so evaluating $\Phi(\rvx_\tau;\rvy_\tau)$ requires an explicit rollout through all intermediate denoising steps. As a result, test-time optimization per control would entail repeated full rollouts, which is computationally infeasible. We therefore \emph{amortize} control selection: a lightweight policy $\vu_\psi$ is trained offline on short preview rollouts, enabling single forward-pass control application at test time. This amortization avoids costly trajectory-level optimization while retaining the pretrained ARDM backbone for stability and expressivity.
\subsection{Practical Instantiation: Preview Windows as Fixed-Lag Assimilation}
\label{subsec:preview}
Chaotic rollouts often drift because feedback from sparse measurements arrives only intermittently. While the standard fixed-lag smoother---which waits until an observation is received and then attempts to retroactively adjust the preceding states---is effective in some settings, such backward corrections can arrive too late in fast, chaotic systems. In contrast we use a bounded \emph{preview} of imminent observations: during each rollout step, the controller is allowed to “peek” a short distance $\Lambda$ ahead into the observation schedule and make small, anticipatory nudges that steer the forecast toward the next arrival. This construction resembles fixed-lag smoothing in spirit (a limited horizon of lookahead), but differs in mechanism: rather than re-estimating past states, we integrate the preview directly into the autoregressive diffusion sampler so that the forward trajectory is nudged on the fly. Thus causality is preserved—no information beyond the preview horizon is used—and the control acts prospectively rather than retrospectively.
\paragraph{Anchored windows, prospective guidance.}
Training proceeds on short, anchored windows of length $\Lambda$. We pick a start $t_0$, roll out the guided dynamics (Eq.~\ref{eq:controlled-path}) for $\Lambda$ steps, and restrict attention to the arrivals that actually occur inside this window:
$
\mathcal{A}_{t_0,\Lambda} \;\triangleq\; \gT \cap [\,t_0{+}1,\,t_0{+}\Lambda\,].$
At each step $t\in\{t_0,\dots,t_0{+}\Lambda{-}1\}$, $\mathcal{P}_\psi$ receives only a compact summary of the imminent arrivals within the current window (the nearest upcoming measurement and its time-to-arrival, see App.~\ref{app:selector}), denoted $\omega_{t\mid t_0}$. The policy then emits per-substep controls using this preview,
$
\rvu_{t+1}^{(s)} \;=\; \vu_\psi\!\big(\rvz^{(s)}_{t+1},\rvx_t,\;\omega_{t\mid t_0},\;s\big),
$
and these are injected into the diffusion sub-steps via Eq.~\ref{eq:controlled-substep}.
\paragraph{Training on what is previewed.}
Consequently, the loss in Eq. \ref{eq:arrival-cost} is also localized to the preview window. Only states that coincide with arrivals inside $[t_0{+}1,\,t_0{+}\Lambda]$ contribute to the loss term. We set $\beta$ from Eq.~\ref{eq:arrival-cost} to zero and roll out our trajectory from $\gP^\psi$.
\begin{equation}
\label{eq:window-loss}
\gC(\rvx) \;=\; \sum_{t\in \mathcal{A}_{t_0,\Lambda}} \E_{\rvx_t \sim \gP^\psi}\Big[\Phi_t(\rvx_t;\,\rvy_t)\Big].
\end{equation}
We present the corresponding training algorithm in Alg.~\ref{alg:training} (App.~\ref{app:algo}).
\paragraph{Inference by sliding previews.}
At test time, we advance autoregressively and form previews of size $\Lambda$. In practice, we implement this procedure using DDIM sampling \citep{song2020ddim}, but the preview mechanism is sampler-agnostic and applies equally well to alternative diffusion samplers. The procedure functions analogously to fixed-lag smoothing in classical DA—improving stability by exploiting short previews of incoming observations—but remains causal: no information beyond the preview horizon is accessed. Alg.~\ref{alg:sampling-chunked} (App.~\ref{app:algo}) describes this inference procedure.
\section{Related Work}
\paragraph{Guidance in Diffusion Models.} Some existing works on guidance in diffusion models rely on explicit approximations of the score of the noisy likelihood score by approximating the diffusion posterior $p\left(\rvx^{(0)}|\rvx^{(t)}\right)$ \citep{chung2022diffusion, song2023pseudoinverseguided, kawar2022denoising, pandey2024fast, pokle2024trainingfree}. While this can result in accurate guidance and faster sampling, a large proportion of these methods are limited to linear inverse problems, which further limits their application. More recent works \citep{pandey2025variational, rout2025rbmodulation} alleviate some of these problems by formulating guidance as optimal control. Our method directly builds on top of \citet{pandey2025variational} by amortizing the controls in a separate finetuning stage and extending their framework to autoregressive diffusion models. There has also been recent work in finetuning diffusion models \citep{clark2023directly, fan2023reinforcement, domingo-enrich2025adjoint} which is complimentary to our proposed framework.
\paragraph{Data assimilation}
Several recent works utilize diffusion models for DA. \citet{rozet2023score} train score-based diffusion models on short trajectory segments to generate full long trajectories during inference. Their framework has been also applied in \citet{qu2024deep} and \citet{manshausen2025generative}. However, these approaches rely on inference-time-only guidance, lacking trajectory consistency mechanisms that avoid error accumulation during observational gaps.
Autoregressive methods \citep{shysheya2024conditional, gao2024bayesian} and DiffDA \citep{huang2024diffda} improve stability but remain computationally expensive due to inference-time optimization. Latent space-based approaches \citep{foroumandi2025harnessing, fan2025physically} reduce dimensionality but introduce reconstruction biases and latent-physical decoupling errors. We address these limitations by integrating learned feedback controls directly into autoregressive diffusion denoising, enabling inference as a single forward rollout with robust long-horizon stability and substantial computational efficiency.
\section{Experiments}
\label{sec:exp}
We evaluate Control–Augmented Data Assimilation (CADA) on two chaotic PDE benchmarks: 2D Kolmogorov flow and 1D Kuramoto–Sivashinsky (KS). To highlight challenges in delayed and sparse observations, we test six regimes combining spatial downsampling and temporal masking. Baselines include four strong diffusion-based DA methods (joint- and conditional-score) plus two CADA ablations. Our study addresses three questions: \emph{(i) Accuracy under delayed preview}—does preview-aware control reduce drift with infrequent, lagged data? \emph{(ii) Stability across horizons}—do improvements hold for long rollouts where chaos dominates? \emph{(iii) Mechanism}—are gains driven by amortization or brute-force search? We next detail datasets, training, metrics, baselines, and provide comprehensive quantitative and qualitative results.
\begin{figure*}[!t]
\centering
\begin{subfigure}[t]{0.5\textwidth}
\centering
\includegraphics[width=\linewidth]{figures/tv1d_lineplot.png}
\caption{Total Variation for Kuramoto-Sivashinsky}
\label{fig:tv}
\end{subfigure}\hfill
\begin{subfigure}[t]{0.5\textwidth}
\centering
\includegraphics[width=\linewidth]{figures/kolmogorov_dissipation_rate.png}
\caption{Dissipation Rate for Kolmogorov Flow}
\label{fig:dr}
\end{subfigure}
\caption{
\textbf{Our method better preserves physics-aware diagnostics under sparse observations.}
(a) Total variation (TV) error for the 1D Kuramoto–Sivashinsky system over a 640-step rollout; lower TV error indicates more faithful preservation of spatial oscillations. (b) Dissipation rate error for 2D Kolmogorov flow over a 60-step rollout; accurate dissipation reflects correct energy cascade to small scales.
Both figures correspond to the \textsc{MR4} observation regime (details in Sec.~\ref{sec:exp}).
}
\label{fig:physics}
\end{figure*}
\paragraph{Dataset}
To assess the effectiveness of CADA under chaotic dynamics, we consider two canonical PDE benchmarks: the 1D Kuramoto–Sivashinsky (KS) equation and the 2D Kolmogorov flow. Both systems exhibit nonlinear instabilities and long-range correlations, making them challenging testbeds for data assimilation with sparse observations \citep{du2021evolutional,lippe2023pde}.
The KS equation is a fourth-order nonlinear PDE modeling flame front instabilities and solidification dynamics, with dynamics
$
\partial_\tau u + u \,\partial_x u + \partial^2_x u + \nu \,\partial^4_x u = 0,
$
where $\nu>0$ denotes the viscosity. We solve it on a periodic domain with 256 spatial points and a fixed time step $\Delta \tau = 0.2$. Training trajectories span $140 \Delta \tau$, while validation and test trajectories extend to $640 \Delta \tau$. The generation strategy and data splits follow \citet{shysheya2024conditional, brandstetter2022lie}.
Kolmogorov flow is a 2D variant of the incompressible Navier–Stokes equations, describing the dynamics of an incompressible fluid:
$
\partial_\tau u + u \cdot \nabla u - \nu \nabla^2 u + \tfrac{1}{\rho}\nabla p - f = 0, \nabla \cdot u = 0,
$
with velocity field $u$, viscosity $\nu$, density $\rho$, pressure $p$, and $f$ an external forcing term. Trajectories span 64 states for training and 180 states for test and validation, each represented on a $64 \times 64$ grid, with $\Delta \tau = 0.2$. Data generation and splits follow \citet{shysheya2024conditional, rozet2023score}. Evaluation is carried out on the scalar vorticity field $\Omega = \partial_x u_y - \partial_y u_x$, which captures rotational structures.
\paragraph{Observation Regimes.}
We evaluate six assimilation settings for each dataset, designed to probe robustness under varying degrees of spatial sparsity and temporal delay. The first set of regimes applies \emph{spatial downsampling} with factors $\{2,4,8\}$ (denoted \textsc{DS-2/4/8}), where observations are available at every simulator step within the preview window $\Lambda$, but only on coarsened spatial grids. The second set applies \emph{regular strided masking} with the same factors (denoted \textsc{MS-2/4/8}), where observations are reported every fourth simulator step within the window (temporal stride of 4) and further subsampled spatially by the given factor. In all cases, operator metadata (e.g., masks) are carried with each observation, as detailed in App.~\ref{app:obs-ops}.
\paragraph{Experimental Setup.}
We train a separate \emph{controller} network for each observation regime, using a common pretrained ARDM (distinct for each dataset). The ARDM backbone in CADA employs DDIM sampling with $S{=}3$ denoising sub-steps per simulator transition. The preview horizon is set to $\Lambda{=}16$ for Kolmogorov flow and $\Lambda{=}54$ for Kuramoto–Sivashinsky, with the active observation selected by the nearest-arrival rule within the anchored window (see App.~\ref{app:selector}). To strengthen the assimilation signal, we evaluate the arrival cost not only at observation indices but also at their intermediate denoising sub-steps, using Tweedie estimates of the forecast state at each sub-step. We present architectural details of the ARDM and controller network in App.~\ref{app:ardm_details} and App.~\ref{app:ctrl-net}, respectively. Forecasts are evaluated under both short and long horizons to separate near-term correction from long-term stability: 60 and 180 steps for Kolmogorov flow, and 140 and 640 steps for Kuramoto–Sivashinsky.
\paragraph{Evaluation Metrics.}
We evaluate forecasts using both trajectory-based and physics-aware metrics. First, we report root-mean-square error (RMSE), a time-averaged pointwise measure of forecast fidelity. Second, we compute the high correlation time (HCT), defined as the last index $\ell_{\max}$ for which the Pearson correlation $\rho(\ell)$ between forecast and ground truth remains above a fixed threshold $\phi$ (we use $\phi=0.9$): $\ell_{\max} \;=\; \max \{\ell : \rho(\ell) \ge \phi \}, \quad t_{\max} = \ell_{\max}\,\Delta t.$
To assess physical fidelity beyond pointwise errors, we include two domain-specific diagnostics.
For the 1D Kuramoto--Sivashinsky system, where $z(\xi)$ denotes the state field over space $\xi$, we report the
\emph{total variation} (TV),
$
\mathrm{TV}(z) \;=\; \int \big|\partial_\xi z(\xi)\big|\, d\xi,
$
which quantifies spatial oscillations and the sharpness of evolving patterns. For the 2D Kolmogorov flow, $z(\xi,\eta)$ denotes the streamfunction,
with $(\xi,\eta)$ the spatial coordinates. We measure the \emph{dissipation rate},
$
\varepsilon \;=\; \nu \int \|\nabla z(\xi,\eta)\|_2^2 \, d\xi\,d\eta.
$
This canonical diagnostic quantifies the rate at which kinetic energy is dissipated at small scales.
Together, RMSE and HCT evaluate assimilated forecast accuracy and temporal coherence, while TV and dissipation provide assessments of structural fidelity in chaotic PDE dynamics.
\paragraph{Baselines}
Score-based Data Assimilation \citep{rozet2023score} learns local joint scores over short segments under a $k$-order Markov factorization and reconstructs a full-trajectory score to sample \emph{all-at-once} (AAO). This provides a principled joint generative model over trajectories, but is memory-intensive for long sequences and in practice requires sequential evaluation of local windows; AAO often benefits from additional corrector steps to enforce start–end consistency.
Using the same local joint score, \citet{shysheya2024conditional} adopt an autoregressive rollout that generates $P$ future states conditioned on the past $C$ (with $P{+}C{=}2k{+}1$), iterating along the horizon. Relative to AAO, the AR factorization increases the effective Markov order seen at each step and empirically yields more stable long-range rollouts at the cost of additional neural function evaluations. Both of these joint-score based methods are referred to as \textit{Joint AAO} and \textit{Joint AR} in this text.
We also evaluate two conditioning architectures from \citet{shysheya2024conditional}: \emph{Plain Amortized}, which fixes the number of conditioning frames $C$ during training, and \emph{Universal Amortized}, which samples $C\sim\mathrm{U}({0,\dots,2k})$ and uses masking to admit variable $(C,P)$ at test time while keeping a fixed input window. Both can be combined with reconstruction guidance for partial observations, following the conditional score formulations.
\paragraph{Our ablations.}
(i) \textit{TTO} (test-time optimization): a non-amortized variant that updates controls per step via inner optimization under a terminal-cost objective, requiring explicit rollouts to evaluate arrival costs; this parallels \citet{pandey2025variational} by extending it to autoregressive temporal setting. (ii) \textit{BoN} (Best-of-$n$): a selection baseline that samples $n{=}16$ independent latent seeds and returns the trajectory with the lowest terminal cost—representing a simple inference-time selection heuristic used in contemporary alignment/scaling studies \citep{singhal2025general,gao2023scaling}.
\paragraph{Qualitative and Quantitative Analysis.} Tab.~\ref{tab:rmsd_clean} reports quantitative comparisons between CADA, state-of-the-art baselines, and ablations. CADA consistently outperforms all baselines across metrics. In particular, it maintains stability over long horizons, as reflected in nearly constant RMSE from short to long rollouts. High Correlation Time (HCT) (Tab.~\ref{tab:hct_clean}, see App.~\ref{app:hct}) corroborates this finding, with correlations remaining strong throughout full trajectories. By contrast, other ARDM-based models such as Plain and Universal Amortized exhibit significant long-horizon degradation and even struggle in the short-horizon regime. This instability is undesirable for operational DA deployments, where stable forecasts are often required over fixed-lag windows resembling our experimental setup.
Not using amortization (TTO) leads to a noticeable drop in performance—mild in short rollouts but increasingly severe over long horizons. This shows that simply tilting the distribution at arrival times is insufficient; amortizing over observations and preview windows is essential for stability. The Best-of-$n$ (BoN) heuristic, which relies on sampling multiple trajectories and selecting the lowest-cost one, performs significantly worse. While BoN can improve sample quality in image generation, it fails here because DA requires consistent step-by-step correction rather than ex-post trajectory selection, underscoring the importance of amortized control for sequential inverse problems.
Fig.~\ref{fig:overview:b} and \ref{fig:compare} illustrate performance under masked observation regimes for long-horizon rollouts on both Kuramoto (1D) and Kolmogorov (2D) benchmarks. These settings are especially challenging, as observations are sparse in both time and space. On Kuramoto, autoregressive diffusion baselines diverge after roughly 200 steps, whereas our method remains stable for the full 640-step horizon. Joint score–based methods also maintain stability but exhibit strong visual artifacts. On Kolmogorov, ARDM baselines such as Plain and Universal Amortized begin diverging shortly after step 35. Joint AR shows greater stability, but in later timesteps its predictions lose fine-scale detail. In contrast, our method preserves sharper, more physically consistent structures throughout.
Fig.~\ref{fig:tv} tracks the evolution of total variation in the Kuramoto system, a proxy for capturing fine-scale spatiotemporal variability. All baselines either underestimate or overestimate fine scale features. In contrast, \textsc{CADA} tracks the ground truth closely across the entire horizon, with small variance and almost no drift, highlighting its robustness under chaotic dynamics. Fig.~\ref{fig:dr} reports dissipation rate in Kolmogorov flow, a canonical diagnostic of energy transfer in turbulence. Joint-score methods underestimate dissipation, and amortized ARDMs overshoot, leading to unphysical rollouts. \textsc{CADA}, by contrast, aligns closely with the ground truth, preserving the correct energy balance. These results highlight that control augmentation not only reduces forecast drift but also enforces domain-relevant physical invariants.
\begin{table}[t]
\centering
\tiny
\caption{
\textbf{Our method outperforms baselines (RMSE $\downarrow$) across six obervation regimes}. Results on Kolmogorov flow (60/180 steps) and Kuramoto–Sivashinsky (140/640 steps) under short- and long-horizon rollouts show CADA consistently superior. Ablations confirm that removing amortization (TTO) or relying on heuristic selection (BoN) substantially degrades performance. Observation regimes (Sec.~\ref{sec:exp}) include downsampled (\textsc{DS}, every step observed) and masked (\textsc{MS}, observations every fourth step). Refer to Tab.~\ref{tab:hct_clean} in App.\ref{app:hct} for HCT $\uparrow$ metric.}
\label{tab:rmsd_clean}
\setlength{\tabcolsep}{4pt}
\begin{tabularx}{\linewidth}{l *{12}{S[table-format=1.3]}}
\toprule
& \multicolumn{2}{c}{\textbf{DS-2}} & \multicolumn{2}{c}{\textbf{DS-4}} &
\multicolumn{2}{c}{\textbf{DS-8}} & \multicolumn{2}{c}{\textbf{MS-2}} &
\multicolumn{2}{c}{\textbf{MS-4}} & \multicolumn{2}{c}{\textbf{MS-8}} \\
\cmidrule(lr){2-3}\cmidrule(lr){4-5}\cmidrule(lr){6-7}\cmidrule(lr){8-9}\cmidrule(lr){10-11}\cmidrule(lr){12-13}
& {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} \\
\midrule
\multicolumn{13}{c}{\textbf{Kolmogorov}} \\
\midrule
\textbf{CADA (ours)} & \bfseries 0.016 & \bfseries 0.016 & \bfseries 0.020 & \bfseries 0.020 & \bfseries 0.138 & 0.351 & \bfseries 0.017 & \bfseries 0.017 & \bfseries 0.024 & \bfseries 0.024 & \bfseries 0.060 & \bfseries 0.286 \\
Joint AAO & 0.041 & 0.045 & 0.210 & 0.189 & 0.380 & 0.244 & 0.141 & 0.171 & 0.358 & 0.559 & 0.465 & 0.523 \\
Joint AR & 0.038 & 0.031 & 0.185 & 0.115 & 0.366 & \bfseries 0.218 & 0.046 & 0.129 & 0.152 & 0.261 & 0.404 & 0.574 \\
Plain Amortized & 0.109 & 0.814 & 0.229 & 1.033 & 0.712 & 1.276 & 0.245 & 0.454 & 0.302 & 0.477 & 0.316 & 0.479 \\
Universal Amortized & 0.295 & 1.398 & 1.061 & 1.566 & 1.612 & 1.766 & 0.186 & 0.397 & 0.323 & 0.469 & 0.351 & 0.483 \\
\arrayrulecolor{lightgray}\hline\arrayrulecolor{black}
TTO & 0.040 & 0.243 & 0.027 & 0.115 & 0.156 & 0.401 & 0.078 & 0.298 & 0.113 & 0.357 & 0.215 & 0.433 \\
BoN & 0.258 & 0.420 & 0.264 & 0.440 & 0.299 & 0.442 & 0.265 & 0.447 & 0.266 & 0.454 & 0.306 & 0.488 \\
\midrule
\multicolumn{13}{c}{\textbf{Kuramoto--Sivashinsky}} \\
\midrule
\textbf{CADA (ours)} & \bfseries 0.006 & \bfseries 0.006 & \bfseries 0.006 & \bfseries 0.006 & \bfseries 0.009 & \bfseries 0.009 & \bfseries 0.006 & \bfseries 0.006 & \bfseries 0.011 & \bfseries 0.011 & \bfseries 0.011 & \bfseries 0.096 \\
Joint AAO & 0.017 & 0.017 & 0.091 & 0.092 & 0.417 & 0.424 & 0.045 & 0.038 & 0.210 & 0.195 & 0.614 & 0.599 \\
Joint AR & 0.018 & 0.018 & 0.091 & 0.093 & 0.413 & 0.428 & 0.026 & 0.009 & 0.041 & 0.032 & 0.134 & 0.136 \\
Plain Amortized & 0.041 & 9.787 & 0.146 & 10.73 & 1.859 & 11.81 & 0.034 & 1.163 & 0.036 & 1.165 & 0.039 & 1.211 \\
Universal Amortized & 0.043 & 5.574 & 0.146 & 6.210 & 2.096 & 6.947 & 0.041 & 1.098 & 0.044 & 1.197 & 0.048 & 1.239 \\
\arrayrulecolor{lightgray}\hline\arrayrulecolor{black}
TTO & 0.016 & 8.288 & 0.009 & 0.053 & 0.418 & 0.634 & 0.016 & 0.298 & 0.114 & 0.363 & 0.081 & 0.580 \\
BoN & 0.046 & 1.257 & 0.046 & 1.498 & 0.048 & 3.122 & 0.045 & 1.987 & 0.046 & 1.644 & 0.049 & 2.128 \\
\bottomrule
\end{tabularx}
\end{table}
\section{Conclusion}
We introduced Control-Augmented Data Assimilation (CADA), a finetuning framework that integrates a learned control mechanism into pretrained autoregressive diffusion models. By amortizing anticipatory corrections through preview windows, CADA transforms data assimilation into a feed-forward process that is both computationally efficient and stable across long horizons.
On two canonical chaotic PDE benchmarks, CADA consistently outperforms state-of-the-art diffusion-based DA methods, yielding more accurate forecasts, improved long-term stability, and closer adherence to domain-standard physical diagnostics. Our experiments show that amortization is key: test-time-only optimization or naive trajectory selection cannot match the robustness achieved by offline-trained control policies.
Beyond PDEs, our framework suggests a general recipe for embedding control into generative dynamics. This perspective opens avenues for extending diffusion models to other sequential inverse problems where observations are delayed, sparse, or noisy—from atmospheric science and climate modeling to robotics and scientific simulation. Future work will explore adaptive preview horizons, integration with real-world observational data, and scaling to higher-dimensional systems.
\section*{Ethics Statement}
Our work focuses on amortized guidance for autoregressive diffusion models. While we present these ideas in the context of data assimilation in PDEs, our setup could be adapted for malicious terminal costs which can lead to potential misuse. Therefore, responsible deployment, monitoring, and safeguards are critical to balance performance gains with societal risks.
\section*{Reproducibility Statement}
We include proofs for all theoretical results introduced in the main text in Appendix \ref{app:proof-tilt}. We include further experimental and implementation details (including model architectures and other hyperparameter choices) in Appendix \ref{app:ctrl-net}. Our code will be made available by the time of publication.
\section*{Acknowledgments}
We thank Felix Draxler for additional discussions and feedback. This project was funded through support from the Chan Zuckerberg Initiative. Additionally, Stephan Mandt acknowledges funding from the National Science Foundation (NSF) through an NSF CAREER Award IIS-2047418, IIS-2007719, the NSF LEAP Center, and the Hasso Plattner Research Center at UCI. Parts of this research were supported by the Intelligence Advanced Research Projects Activity (IARPA) via Department of Interior/ Interior Business Center (DOI/IBC) contract number 140D0423C0075. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright annotation thereon. Disclaimer: The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of IARPA, DOI/IBC, or the U.S. Government.
\section{Proof of the Tilted Distribution}
\label{app:proof-tilt}
We show that the optimization problem in Eq.~\ref{eq:arrival-cost} admits the tilted distribution in Eq.~\ref{eq:tilt_distribution} as its optimal solution.
\paragraph{Setup.}
Recall the objective
\[
\gC(\rvx) \;=\; \sum_{t\in\gT} \E_{\rvx_t \sim \gP}\Big[\Phi_t(\rvx_t;\,\rvy_t)\Big] + \beta\,\KL(\gP \,\|\, \gQ_\theta),
\]
where $\gP$ is the guided ARDM distribution over trajectories, $\gQ_\theta$ the unguided distribution, and $\Phi_t$ an arrival-time cost.
\paragraph{Variational form.}
Expanding the KL divergence,
\[
\KL(\gP \,\|\, \gQ_\theta) \;=\; \E_{\gP}\!\left[ \log \frac{\gP}{\gQ_\theta} \right].
\]
Thus the objective reads
\[
\gC(\gP) \;=\; \E_{\gP}\!\Bigg[ \sum_{t\in\gT}\Phi_t(\rvx_t;\,\rvy_t) + \beta \log \frac{\gP}{\gQ_\theta}\Bigg].
\]
\paragraph{Lagrangian minimization.}
Consider minimizing $\gC(\gP)$ over distributions $\gP$ subject to normalization $\int \gP = 1$. The corresponding Lagrangian is
\[
\mathcal{L}(\gP,\lambda) \;=\; \E_{\gP}\!\Bigg[ \sum_{t\in\gT}\Phi_t(\rvx_t;\,\rvy_t) + \beta \log \frac{\gP}{\gQ_\theta}\Bigg] + \lambda \Big(\int \gP - 1\Big).
\]
\paragraph{Stationary point.}
Taking the functional derivative w.r.t. $\gP$ gives
\[
\frac{\delta \mathcal{L}}{\delta \gP} \;=\; \sum_{t\in\gT}\Phi_t(\rvx_t;\,\rvy_t) + \beta \Big(1 + \log \tfrac{\gP}{\gQ_\theta}\Big) + \lambda.
\]
Setting this derivative to zero yields
\[
\log \gP \;=\; \log \gQ_\theta - \tfrac{1}{\beta}\sum_{t\in\gT}\Phi_t(\rvx_t;\,\rvy_t) - \tfrac{\lambda+ \beta}{\beta}.
\]
\paragraph{Closed form.}
Exponentiating both sides gives
\[
\gP^*(\rvx_{0:T}) \;\propto\; \gQ_\theta(\rvx_{0:T}) \exp\!\Bigg(-\tfrac{1}{\beta}\sum_{t\in\gT}\Phi_t(\rvx_t;\,\rvy_t)\Bigg),
\]
which is exactly the tilted distribution in Eq.~\ref{eq:tilt_distribution}.
\section{Observation Operators}
\label{app:obs-ops}
Our experiments employ linear observation operators that map the full state $\rvx$ to observed signals $\rvy_k$.
\paragraph{Masked observations.}
For temporally strided or spatially sparse measurements, we define a (possibly time-varying) binary mask $\mM \in \{0,1\}^{1\times D}$ broadcast across channels, with
\[
A_{\mathrm{mask}}(\rvx) \;=\; \mM \odot \rvx,
\qquad
\Phi^{\mathrm{mask}}(\rvx;\rvy) \;=\; \frac{\|\mM \odot (\rvx - \rvy)\|_2^2}{\|\mM\|_1}.
\]
\paragraph{Downsampled observations.}
For coarse-resolution sensing, we apply average pooling $P_f$ over non-overlapping $f{\times}f$ blocks followed by nearest-neighbor upsampling $U_f$:
\[
A_{\downarrow f}(\rvx) \;=\; U_f(P_f \rvx),
\qquad
\Phi^{\mathrm{ds}}(\rvx;\rvy) \;=\; \|U_f(P_f \rvx) - U_f(P_f \rvy)\|_2^2.
\]
These operators produce the observed signals $\rvy_\tau$ that define the arrival-time costs in Eq.~\ref{eq:arrival-cost}. While we restrict to masking and downsampling here, any differentiable operator $\Phi$ could be incorporated within our framework without modification.
\section{Active Observation Selector}
\label{app:selector}
We maintain a preview buffer containing all observations from $\mathcal{T}$ that lie within a fixed lookahead horizon $\Lambda$ from index $t_0$. Each entry in the buffer is represented as a triplet $(\vy_j, \mM_j, \Delta_j)$, where
(i) $j \in \mathcal{T}$ is the physical time index of the observation;
(ii) $\vy_j$ is the observed signal, lifted to full resolution when necessary;
(iii) $\mM_j$ is an auxiliary mask (see App.~\ref{app:obs-ops}), while for other operators $\mM_j$ may be ignored or replaced with suitable metadata; and
(iv) $\Delta_j = j - t + 1$ is the lead time relative to the current forecast step $t$.
At each forecast step $t$, the active preview is defined as the nearest available future observation within the lookahead window
\[
\mathcal{W}_{t\mid t_0} \;\triangleq\; \{\, j \in \mathcal{T} : t{+}1 \le j \le t_0{+}\Lambda \,\}.
\]
The chosen preview is then
\[
\omega_{t\mid t_0} = (\vy_{j^\star}, \mM_{j^\star}, \Delta_{t,j^\star}),
\qquad
j^\star = \arg\min_{j \in \mathcal{W}_{t\mid t_0}} \{\Delta_{t,j} : \Delta_{t,j} \ge 0\}.
\]
In words, at each step the selector activates the nearest previewed observation within the anchored preview wiindow, along with its associated metadata and lead time.
\section{Training and Sampling Algorithm}
\label{app:algo}
\begin{algorithm}[H]
\caption{Preview-aware controlled DDIM one-step ($\rvx_t \!\to\! \rvx_{t+1}$)}
\label{alg:controlled-step}
\begin{algorithmic}[1]
\Require current state $\rvx_t$; preview $\omega_{t\mid t_0}=\{(\rvy_j, \mM_j, \Delta_{t,j}) : j \in \mathcal W_{t\mid t_0}\}$ (see App.~\ref{app:selector}); pretrained ARDM kernels $q_\theta$; control policy $\vu_\psi(\cdot)$; step $\gamma>0$
\Ensure next state $\rvx_{t+1}$ and (if applicable) arrival-time cost $\ell_{t+1}$
\State Sample parent latent $\rvz_{t+1}^{(S)} \sim p_S$
\For{$s = S-1, S-2, \dots, 0$} \Comment{DDIM sub-steps}
\State $\rvu_{t+1}^{(s)} \gets \vu_\psi\!\big(\rvx_t,\;\omega_{t\mid t_0},\;s\big)$
\State $\tilde \rvz \gets \rvz_{t+1}^{(s+1)} + \gamma\,\rvu_{t+1}^{(s)}$
\State $\rvz_{t+1}^{(s)} \sim q_\theta\!\big(\,\cdot\,\big|\, \tilde \rvz;\;\rvx_t\big)$
\EndFor
\State $\rvx_{t+1} \gets \rvz_{t+1}^{(0)}$
\If{$t{+}1 \in \mathcal T$} \Comment{Eq.~\ref{eq:window-loss}}
\State $\ell_{t+1} \gets \Phi_{t+1}(\rvx_{t+1};\,\rvy_{t+1})$
\Else
\State $\ell_{t+1} \gets 0$
\EndIf
\State \Return $\rvx_{t+1},\;\ell_{t+1}$
\end{algorithmic}
\end{algorithm}
\begin{algorithm}[H]
\caption{Training the \emph{controller} network}
\label{alg:training}
\begin{algorithmic}[1]
\Require pretrained ARDM kernel $q_\theta$; stream $\{\rvy_j\}_{j\in\mathcal T}$; preview horizon $\Lambda$; strength $\beta>0$; optimizer for $\psi$; \emph{controller} network $\vu_\psi(\cdot)$; step $\gamma>0$
\Ensure trained parameters $\psi$
\Repeat
\State Sample rollout start $t_0$ and initial $\rvx_{t_0}\sim p_0$
\State $\mathcal A_{t_0,\Lambda} \gets \mathcal T \cap [t_0{+}1,\,t_0{+}\Lambda]$, \quad $\widehat{\mathcal C} \gets 0$
\For{$t = t_0, t_0{+}1, \dots, t_0{+}\Lambda-1$} \Comment{anchored preview; see App.~\ref{app:selector}}
\State $\mathcal W_{t\mid t_0} \gets \{\, j\in\mathcal T : t{+}1 \le j \le t_0{+}\Lambda \,\}$
\State Optionally specialize to nearest arrival $j^\star{=}\min(\mathcal W_{t\mid t_0})$ if $\mathcal W_{t\mid t_0}\!\neq\!\varnothing$
\State Build $\omega_{t\mid t_0}$ from $\mathcal W_{t\mid t_0}$ (or $\{j^\star\}$)
\State $(\rvx_{t+1},\ell_{t+1}) \gets \textsc{ControlledStep}(\rvx_t,\omega_{t\mid t_0},q_\theta,\vu_\psi,\gamma)$ \Comment{Alg.~\ref{alg:controlled-step}}
\State $\rvx_t \gets \rvx_{t+1}$; \quad $\widehat{\mathcal C} \gets \widehat{\mathcal C} + \ell_{t+1}\,\mathbf{1}\{t{+}1\in\mathcal A_{t_0,\Lambda}\}$
\EndFor
\State $\widehat{\mathcal C} \gets \widehat{\mathcal C}\Big/\max\{|\mathcal A_{t_0,\Lambda}|,\,1\}$ \Comment{arrival normalization}
\State $\mathcal L(\psi) \gets \widehat{\mathcal C}$
\State Update $\psi$ by descending $\nabla_\psi \mathcal L(\psi)$
\Until{convergence}
\end{algorithmic}
\end{algorithm}
\begin{algorithm}[H]
\caption{Preview-aware forecasting with $\Lambda$-chunk anchoring}
\label{alg:sampling-chunked}
\begin{algorithmic}[1]
\Require pretrained ARDM kernel $q_\theta$ (Eq.~\eqref{eq:Qtheta-finite}); trained $\vu_\psi$; initial $\rvx_{t_0}$; forecast horizon $H$; preview horizon $\Lambda$; stream $\{\rvy_j\}_{j\in\mathcal T}$; step $\gamma>0$
\Ensure forecast $\rvx_{1:H}=(\rvx_{t_0+1},\dots,\rvx_{t_0+H})$
\State $C \gets \left\lceil H/\Lambda \right\rceil$
\For{$c = 0,1,\dots,C{-}1$} \Comment{chunk index}
\State $t_0^{(c)} \gets t_0 + c\,\Lambda$
\State $\Lambda_c \gets \min\{\Lambda,\; H - c\,\Lambda\}$ \Comment{last chunk may be shorter}
\For{$t = t_0^{(c)},\; \dots,\; t_0^{(c)}+\Lambda_c-1$}
\State $\mathcal W_{t\mid t_0^{(c)}} \gets \{\, j\in\mathcal T : t{+}1 \le j \le t_0^{(c)}{+}\Lambda_c \,\}$
\State Build $\omega_{t\mid t_0^{(c)}}$ from $\mathcal W_{t\mid t_0^{(c)}}$ \hfill (nearest-arrival option allowed)
\State $\rvx_{t+1} \gets \textsc{ControlledStep}(\rvx_t,\omega_{t\mid t_0^{(c)}},\{q_\theta^{(s)}\},\vu_\psi,\gamma).\textsc{state}$
\EndFor
\State \emph{// autoregressive handoff: last state becomes next chunk's initial condition}
\State $\rvx_{t_0^{(c+1)}} \gets \rvx_{t_0^{(c)}+\Lambda_c}$ \Comment{only if $c{+}1<C$}
\EndFor
\State \Return $\rvx_{1:H}$
\end{algorithmic}
\end{algorithm}
\section{Pretraining ARDM}
\label{app:ardm_details}
\paragraph{Implementation details.}
We adapt the 1D and 2D diffusion implementations from \texttt{lucidrains/denoising-diffusion-pytorch}\footnote{\url{https://github.com/lucidrains/denoising-diffusion-pytorch}} into an autoregressive diffusion model (ARDM) tailored for PDE forecasting. Each ARDM transition corresponds to one-step forecasting via a DDIM sampler with $S{=}3$ denoising steps, $v$-parameterization, and a sigmoid schedule for $\alpha$.
The backbone is a residual U-Net with multi-resolution attention and learned sinusoidal time embeddings:
\begin{verbatim}
dim = 64,
dim_mults = (1, 2, 4, 8),
learned_sinusoidal_dim = 128
\end{verbatim}
Attention layers are applied at intermediate and coarse resolutions, while residual blocks follow the standard Conv–Norm–SiLU design.
\paragraph{Training configuration.}
Models are trained with mixed precision (\textsc{fp16}) and exponential moving average (EMA). The configuration is:
\begin{verbatim}
train_batch_size = 32
train_lr = 3.2e-4
train_num_steps = 1000000
gradient_accumulate_every = 1
ema_decay = 0.995
ema_every = 10
\end{verbatim}
\section{Control Network Architecture}
\label{app:ctrl-net}
\paragraph{Overview.}
The \emph{controller} network $\vu_\psi$ produces controls $\rvu_t^{(s)}$ used at each denoising sub-step. In the rollout, we write $\rvu_{t+1}^{(s)} = \vu_\psi(\cdot)$ for brevity; here we provide a detailed description of the architecture and conditioning.
\paragraph{Inputs.}
At each sub-step, the network receives five spatial tensors concatenated along channels:
(i) the current latent $\rvz_{t+1}^{(s)}$,
(ii) the previous state $\rvx_t$,
(iii) the preview observation $\rvy_t^\star$,
(iv) the auxiliary mask $\mM_t^\star$, and
(v) the previous control $\rvu_{\mathrm{prev}}$.
In addition, it conditions on scalar metadata: the preview lag $\Delta_t^\star$, the local frame index $\tau$ within the preview window, and the current $\log\mathrm{SNR}(s)$ from the DDIM schedule.
\paragraph{Backbone encoder.}
The concatenated inputs are passed through a shallow convolutional encoder with two $3{\times}3$ layers and group normalization. A downsample/upsample path provides limited multi-scale context: features are reduced to half resolution, then upsampled and fused back with the original resolution. A $1{\times}1$ fusion convolution followed by group normalization yields the encoded representation.
\paragraph{FiLM conditioning.}
Each scalar input is normalized to $[0,1]$ and embedded via a two-layer MLP of dimension $\texttt{hid}$. The three embeddings (lag, frame index, SNR) are concatenated and mapped to $(\gamma,\beta)$ coefficients through a linear layer. These coefficients modulate the encoded features in a FiLM style,
$
\text{feat} \;\mapsto\; \text{feat} \cdot (1+\gamma) + \beta.
$
\paragraph{Residual head.}
A $3{\times}3$ convolutional head outputs the control increment $\Delta_\psi$. This is added to a normalized copy of the previous control $\rvu_{\mathrm{prev}}$, producing
$
\rvu_t^{(s)} = \text{GroupNorm}(\rvu_{\mathrm{prev}}) + \Delta_\psi.
$
At the first denoising step, $u_{\mathrm{prev}}$ is set to zero.
\paragraph{Configuration.}
In our experiments we instantiate the control network as
$
\texttt{hid=768},
$
with group normalization (8 groups), hidden dimension $\texttt{hid}$ for the encoder, and FiLM embeddings of dimension $\texttt{hid}$. The architecture is lightweight relative to the ARDM UNet (App.~\ref{app:ardm_details}) but sufficiently expressive to incorporate preview information into the denoising dynamics.
\paragraph{Training details.}
Gradients flow only into $\psi$ (the UNet $\theta$ is frozen).
We use gradient checkpointing at each UNet call and detach $\rvu_{\mathrm{prev}}$ within a frame to avoid deep denoising-step recurrences; memory scales with the number of checkpoints.
\section{Additional Results}
\label{app:hct}
Tab.~\ref{tab:hct_clean} reports the HCT metric for all our experiments.
\begin{table}[h]
\centering
\tiny
\caption{Quantitative comparison (HCT $\uparrow$) between our method and competitive baselines across six observation regimes (see Sec.~\ref{sec:exp}). We evaluate both 1D and 2D PDE benchmarks—Kolmogorov (60/180 steps) and Kuramoto–Sivashinsky (140/640 steps)—under short and long horizons. Our method consistently outperforms alternatives. In ablations, removing amortization (TTO) or using simple heuristic selection (BoN) leads to significant degradation.}
\label{tab:hct_clean}
\setlength{\tabcolsep}{4pt}
\begin{tabularx}{\linewidth}{l *{12}{S[table-format=3.0]}}
\toprule
& \multicolumn{2}{c}{\textbf{DS-2}} & \multicolumn{2}{c}{\textbf{DS-4}} &
\multicolumn{2}{c}{\textbf{DS-8}} & \multicolumn{2}{c}{\textbf{MS-2}} &
\multicolumn{2}{c}{\textbf{MS-4}} & \multicolumn{2}{c}{\textbf{MS-8}} \\
\cmidrule(lr){2-3}\cmidrule(lr){4-5}\cmidrule(lr){6-7}\cmidrule(lr){8-9}\cmidrule(lr){10-11}\cmidrule(lr){12-13}
& {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} & {short} & {long} \\
\midrule
\multicolumn{13}{c}{\textbf{Kolmogorov}} \\
\midrule
\textbf{CADA (ours)} & \bfseries 60 & \bfseries 180 & \bfseries 60 & \bfseries 180 & \bfseries 50 & 50 & \bfseries 60 & \bfseries 180 & \bfseries 60 & \bfseries 180 & \bfseries 60 & \bfseries 70 \\
Joint AAO & 60 & 180 & 60 & 180 & 50 & 180 & 60 & 173 & 25 & 25 & 10 & 18 \\
Joint AR & 60 & 180 & 60 & 180 & 50 & \bfseries 180 & 60 & 177 & 60 & 62 & 13 & 13 \\
Plain Amortized & 60 & 64 & 41 & 45 & 13 & 17 & 38 & 42 & 32 & 33 & 31 & 32 \\
Universal Amortized & 37 & 13 & 17 & 17 & 8 & 8 & 50 & 50 & 32 & 32 & 28 & 28 \\
TTO & 60 & 73 & 60 & 180 & 37 & 40 & 60 & 60 & 48 & 40 & 27 & 25 \\
BoN & 40 & 40 & 32 & 26 & 21 & 26 & 40 & 26 & 32 & 26 & 32 & 27 \\
\midrule
\multicolumn{13}{c}{\textbf{Kuramoto--Sivashinsky}} \\
\midrule
\textbf{CADA (ours)} & \bfseries 140 & \bfseries 640 & \bfseries 140 & \bfseries 640 & \bfseries 140 & \bfseries 640 & \bfseries 140 & \bfseries 640 & \bfseries 140 & \bfseries 640 & \bfseries 140 & \bfseries 640 \\
Joint AAO & 140 & 640 & 139 & 640 & 139 & 640 & 140 & 640 & 140 & 640 & 129 & 633 \\
Joint AR & 139 & 639 & 139 & 639 & 139 & 639 & 139 & 639 & 139 & 639 & 139 & 639 \\
Plain Amortized & 140 & 211 & 140 & 147 & 76 & 56 & 140 & 253 & 140 & 264 & 140 & 242 \\
Universal Amortized & 140 & 262 & 140 & 168 & 62 & 62 & 140 & 286 & 140 & 217 & 140 & 274 \\
TTO & 140 & 218 & 140 & 640 & 139 & 633 & 140 & 638 & 140 & 637 & 140 & 634 \\
BoN & 140 & 260 & 140 & 250 & 140 & 240 & 140 & 243 & 140 & 250 & 140 & 246 \\
\bottomrule
\end{tabularx}
\end{table}
\begin{thebibliography}{50}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Brandstetter et~al.(2022)Brandstetter, Welling, and Worrall]{brandstetter2022lie}
Johannes Brandstetter, Max Welling, and Daniel~E Worrall.
\newblock Lie point symmetry data augmentation for neural pde solvers.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 2241--2256. PMLR, 2022.
\bibitem[Carrassi et~al.(2018)Carrassi, Bocquet, Bertino, and Evensen]{carrassi2018data}
Alberto Carrassi, Marc Bocquet, Laurent Bertino, and Geir Evensen.
\newblock Data assimilation in the geosciences: An overview of methods, issues, and perspectives.
\newblock \emph{WIREs Climate Change}, 9\penalty0 (5):\penalty0 e535, 2018.
\newblock \doi{https://doi.org/10.1002/wcc.535}.
\newblock URL \url{https://wires.onlinelibrary.wiley.com/doi/abs/10.1002/wcc.535}.
\bibitem[Chung et~al.(2022)Chung, Kim, Mccann, Klasky, and Ye]{chung2022diffusion}
Hyungjin Chung, Jeongsol Kim, Michael~Thompson Mccann, Marc~Louis Klasky, and Jong~Chul Ye.
\newblock Diffusion posterior sampling for general noisy inverse problems.
\newblock In \emph{The Eleventh International Conference on Learning Representations}, 2022.
\bibitem[Chung et~al.(2023)Chung, Kim, Mccann, Klasky, and Ye]{chung2023diffusion}
Hyungjin Chung, Jeongsol Kim, Michael~Thompson Mccann, Marc~Louis Klasky, and Jong~Chul Ye.
\newblock Diffusion posterior sampling for general noisy inverse problems.
\newblock In \emph{The Eleventh International Conference on Learning Representations}, 2023.
\newblock URL \url{https://openreview.net/forum?id=OnD9zGAGT0k}.
\bibitem[Clark et~al.(2023)Clark, Vicol, Swersky, and Fleet]{clark2023directly}
Kevin Clark, Paul Vicol, Kevin Swersky, and David~J Fleet.
\newblock Directly fine-tuning diffusion models on differentiable rewards.
\newblock \emph{arXiv preprint arXiv:2309.17400}, 2023.
\bibitem[Courtier et~al.(1998)Courtier, Andersson, Heckley, Vasiljevic, Hamrud, Hollingsworth, Rabier, Fisher, and Pailleux]{courtier1998ecmwf}
Philippe Courtier, E~Andersson, W~Heckley, D~Vasiljevic, M~Hamrud, A~Hollingsworth, F~Rabier, M~Fisher, and J~Pailleux.
\newblock The ecmwf implementation of three-dimensional variational assimilation (3d-var). i: Formulation.
\newblock \emph{Quarterly Journal of the Royal Meteorological Society}, 124\penalty0 (550):\penalty0 1783--1807, 1998.
\bibitem[Daras et~al.(2024)Daras, Chung, Lai, Mitsufuji, Ye, Milanfar, Dimakis, and Delbracio]{daras2024survey}
Giannis Daras, Hyungjin Chung, Chieh-Hsin Lai, Yuki Mitsufuji, Jong~Chul Ye, Peyman Milanfar, Alexandros~G Dimakis, and Mauricio Delbracio.
\newblock A survey on diffusion models for inverse problems.
\newblock \emph{arXiv preprint arXiv:2410.00083}, 2024.
\bibitem[Dhariwal \& Nichol(2021)Dhariwal and Nichol]{dhariwal2021diffusion}
Prafulla Dhariwal and Alexander~Quinn Nichol.
\newblock Diffusion models beat {GAN}s on image synthesis.
\newblock In A.~Beygelzimer, Y.~Dauphin, P.~Liang, and J.~Wortman Vaughan (eds.), \emph{Advances in Neural Information Processing Systems}, 2021.
\newblock URL \url{https://openreview.net/forum?id=AAWuCvzaVt}.
\bibitem[Domingo-Enrich et~al.(2025)Domingo-Enrich, Drozdzal, Karrer, and Chen]{domingo-enrich2025adjoint}
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T.~Q. Chen.
\newblock Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.
\newblock In \emph{The Thirteenth International Conference on Learning Representations}, 2025.
\newblock URL \url{https://openreview.net/forum?id=xQBRrtQM8u}.
\bibitem[Du \& Zaki(2021)Du and Zaki]{du2021evolutional}
Yifan Du and Tamer~A Zaki.
\newblock Evolutional deep neural network.
\newblock \emph{Physical Review E}, 104\penalty0 (4):\penalty0 045303, 2021.
\bibitem[Evensen et~al.(2022)Evensen, Vossepoel, and Van~Leeuwen]{evensen2022data}
Geir Evensen, Femke~C Vossepoel, and Peter~Jan Van~Leeuwen.
\newblock \emph{Data assimilation fundamentals: A unified formulation of the state and parameter estimation problem}.
\newblock Springer Nature, 2022.
\newblock \doi{https://doi.org/10.1007/978-3-030-96709-3}.
\bibitem[Fan et~al.()Fan, Bai, Fei, Xiao, Chen, Liu, Qu, Ling, and Gentine]{fan2025physically}
Hang Fan, Lei Bai, Ben Fei, Yi~Xiao, Kun Chen, Yubao Liu, Yongquan Qu, Fenghua Ling, and Pierre Gentine.
\newblock Physically consistent global atmospheric data assimilation with machine learning in latent space.
\newblock URL \url{http://arxiv.org/abs/2502.02884}.
\bibitem[Fan et~al.(2023)Fan, Watkins, Du, Liu, Ryu, Boutilier, Abbeel, Ghavamzadeh, Lee, and Lee]{fan2023reinforcement}
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee.
\newblock Reinforcement learning for fine-tuning text-to-image diffusion models.
\newblock In \emph{Thirty-seventh Conference on Neural Information Processing Systems}, 2023.
\newblock URL \url{https://openreview.net/forum?id=8OTPepXzeh}.
\bibitem[Foroumandi \& Moradkhani(2025)Foroumandi and Moradkhani]{foroumandi2025harnessing}
Ehsan Foroumandi and Hamid Moradkhani.
\newblock Harnessing generative deep learning for enhanced ensemble data assimilation.
\newblock \emph{Water Resources Research}, 61\penalty0 (7):\penalty0 e2025WR040078, 2025.
\newblock \doi{https://doi.org/10.1029/2025WR040078}.
\newblock URL \url{https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2025WR040078}.
\newblock e2025WR040078 2025WR040078.
\bibitem[Gao et~al.()Gao, Han, Fan, Sun, Liu, Duan, and Wang]{gao2024bayesian}
Han Gao, Xu~Han, Xiantao Fan, Luning Sun, Li-Ping Liu, Lian Duan, and Jian-Xun Wang.
\newblock Bayesian conditional diffusion models for versatile spatiotemporal turbulence generation.
\newblock 427:\penalty0 117023.
\newblock ISSN 0045-7825.
\newblock \doi{10.1016/j.cma.2024.117023}.
\newblock URL \url{https://www.sciencedirect.com/science/article/pii/S0045782524002792}.
\bibitem[Gao et~al.(2023)Gao, Schulman, and Hilton]{gao2023scaling}
Leo Gao, John Schulman, and Jacob Hilton.
\newblock Scaling laws for reward model overoptimization.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 10835--10866. PMLR, 2023.
\bibitem[Ge et~al.(2022)Ge, Hayes, Yang, Yin, Pang, Jacobs, Huang, and Parikh]{Ge2022}
Songwei Ge, Thomas Hayes, Harry Yang, Xi~Yin, Guan Pang, David Jacobs, Jia-Bin Huang, and Devi Parikh.
\newblock \emph{Long Video Generation with Time-Agnostic VQGAN and Time-Sensitive Transformer}, pp.\ 102–118.
\newblock Springer Nature Switzerland, 2022.
\newblock ISBN 9783031197901.
\newblock \doi{10.1007/978-3-031-19790-1_7}.
\newblock URL \url{http://dx.doi.org/10.1007/978-3-031-19790-1_7}.
\bibitem[Ho et~al.(2020)Ho, Jain, and Abbeel]{ho2020denoising}
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
\newblock Denoising diffusion probabilistic models.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 6840--6851, 2020.
\bibitem[Huang et~al.(2024)Huang, Gianinazzi, Yu, Dueben, and Hoefler]{huang2024diffda}
Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter~D Dueben, and Torsten Hoefler.
\newblock Diffda: a diffusion model for weather-scale data assimilation.
\newblock \emph{arXiv preprint arXiv:2401.05932}, 2024.
\bibitem[Huang et~al.(2025)Huang, Li, He, Zhou, and Shechtman]{huang2025selfforcing}
Xun Huang, Zhengqi Li, Guande He, Mingyuan Zhou, and Eli Shechtman.
\newblock Self forcing: Bridging the train-test gap in autoregressive video diffusion.
\newblock \emph{arXiv preprint arXiv:2506.08009}, 2025.
\bibitem[Hyv{\"a}rinen \& Dayan(2005)Hyv{\"a}rinen and Dayan]{hyvarinen2005estimation}
Aapo Hyv{\"a}rinen and Peter Dayan.
\newblock Estimation of non-normalized statistical models by score matching.
\newblock \emph{Journal of Machine Learning Research}, 6\penalty0 (4), 2005.
\bibitem[Kalnay(2002)]{kalnay_2002_atmospheric}
Eugenia Kalnay.
\newblock \emph{Atmospheric Modeling, Data Assimilation and Predictability}.
\newblock Cambridge University Press, 2002.
\newblock \doi{https://doi.org/10.1017/CBO9780511802270}.
\bibitem[Kawar et~al.(2022)Kawar, Elad, Ermon, and Song]{kawar2022denoising}
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song.
\newblock Denoising diffusion restoration models.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 23593--23606, 2022.
\bibitem[Le~Dimet \& Talagrand(1986)Le~Dimet and Talagrand]{le1986variational}
Fran{\c{c}}ois-Xavier Le~Dimet and Olivier Talagrand.
\newblock Variational algorithms for analysis and assimilation of meteorological observations: theoretical aspects.
\newblock \emph{Tellus A: Dynamic Meteorology and Oceanography}, 38\penalty0 (2):\penalty0 97--110, 1986.
\bibitem[Lippe et~al.(2023)Lippe, Veeling, Perdikaris, Turner, and Brandstetter]{lippe2023pde}
Phillip Lippe, Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter.
\newblock Pde-refiner: Achieving accurate long rollouts with neural pde solvers.
\newblock \emph{Advances in Neural Information Processing Systems}, 36:\penalty0 67398--67433, 2023.
\bibitem[Manshausen et~al.(2025)Manshausen, Cohen, Harrington, Pathak, Pritchard, Garg, Mardani, Kashinath, Byrne, and Brenowitz]{manshausen2025generative}
Peter Manshausen, Yair Cohen, Peter Harrington, Jaideep Pathak, Mike Pritchard, Piyush Garg, Morteza Mardani, Karthik Kashinath, Simon Byrne, and Noah Brenowitz.
\newblock Generative data assimilation of sparse weather station observations at kilometer scales, 2025.
\newblock URL \url{https://arxiv.org/abs/2406.16947}.
\bibitem[Mardani et~al.(2025)Mardani, Brenowitz, Cohen, Pathak, Chen, Liu, Vahdat, Nabian, Ge, Subramaniam, Kashinath, Kautz, and Pritchard]{Mardani2025}
Morteza Mardani, Noah Brenowitz, Yair Cohen, Jaideep Pathak, Chieh-Yu Chen, Cheng-Chin Liu, Arash Vahdat, Mohammad~Amin Nabian, Tao Ge, Akshay Subramaniam, Karthik Kashinath, Jan Kautz, and Mike Pritchard.
\newblock Residual corrective diffusion modeling for km-scale atmospheric downscaling.
\newblock \emph{Communications Earth \&; Environment}, 6\penalty0 (1), February 2025.
\newblock ISSN 2662-4435.
\newblock \doi{10.1038/s43247-025-02042-5}.
\newblock URL \url{http://dx.doi.org/10.1038/s43247-025-02042-5}.
\bibitem[Pandey et~al.(2024)Pandey, Yang, and Mandt]{pandey2024fast}
Kushagra Pandey, Ruihan Yang, and Stephan Mandt.
\newblock Fast samplers for inverse problems in iterative refinement models.
\newblock In \emph{The Thirty-eighth Annual Conference on Neural Information Processing Systems}, 2024.
\newblock URL \url{https://openreview.net/forum?id=qxS4IvtLdD}.
\bibitem[Pandey et~al.(2025)Pandey, Sofian, Draxler, Karaletsos, and Mandt]{pandey2025variational}
Kushagra Pandey, Farrin~Marouf Sofian, Felix Draxler, Theofanis Karaletsos, and Stephan Mandt.
\newblock Variational control for guidance in diffusion models.
\newblock \emph{arXiv preprint arXiv:2502.03686}, 2025.
\bibitem[Pathak et~al.(2024)Pathak, Cohen, Garg, Harrington, Brenowitz, Durran, Mardani, Vahdat, Xu, Kashinath, and Pritchard]{pathak2024kilometerscaleconvectionallowingmodel}
Jaideep Pathak, Yair Cohen, Piyush Garg, Peter Harrington, Noah Brenowitz, Dale Durran, Morteza Mardani, Arash Vahdat, Shaoming Xu, Karthik Kashinath, and Michael Pritchard.
\newblock Kilometer-scale convection allowing model emulation using generative diffusion modeling, 2024.
\newblock URL \url{https://arxiv.org/abs/2408.10958}.
\bibitem[Pokle et~al.(2024)Pokle, Muckley, Chen, and Karrer]{pokle2024trainingfree}
Ashwini Pokle, Matthew~J. Muckley, Ricky T.~Q. Chen, and Brian Karrer.
\newblock Training-free linear image inverses via flows.
\newblock \emph{Transactions on Machine Learning Research}, 2024.
\newblock ISSN 2835-8856.
\newblock URL \url{https://openreview.net/forum?id=PLIt3a4yTm}.
\bibitem[Qu et~al.(2024)Qu, Nathaniel, Li, and Gentine]{qu2024deep}
Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine.
\newblock Deep generative data assimilation in multimodal setting.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops}, pp.\ 449--459, June 2024.
\bibitem[Rout et~al.(2025)Rout, Chen, Ruiz, Kumar, Caramanis, Shakkottai, and Chu]{rout2025rbmodulation}
Litu Rout, Yujia Chen, Nataniel Ruiz, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu.
\newblock {RB}-modulation: Training-free stylization using reference-based modulation.
\newblock In \emph{The Thirteenth International Conference on Learning Representations}, 2025.
\newblock URL \url{https://openreview.net/forum?id=bnINPG5A32}.
\bibitem[Rozet \& Louppe(2023)Rozet and Louppe]{rozet2023score}
Fran{\c{c}}ois Rozet and Gilles Louppe.
\newblock Score-based data assimilation.
\newblock \emph{Advances in Neural Information Processing Systems}, 36:\penalty0 40521--40541, 2023.
\bibitem[Shysheya et~al.(2024)Shysheya, Diaconu, Bergamin, Perdikaris, Hern{\'a}ndez-Lobato, Turner, and Mathieu]{shysheya2024conditional}
Aliaksandra Shysheya, Cristiana Diaconu, Federico Bergamin, Paris Perdikaris, Jos{\'e}~Miguel Hern{\'a}ndez-Lobato, Richard Turner, and Emile Mathieu.
\newblock On conditional diffusion models for pde simulations.
\newblock \emph{Advances in Neural Information Processing Systems}, 37:\penalty0 23246--23300, 2024.
\bibitem[Singhal et~al.(2025)Singhal, Horvitz, Teehan, Ren, Yu, McKeown, and Ranganath]{singhal2025general}
Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath.
\newblock A general framework for inference-time scaling and steering of diffusion models.
\newblock \emph{arXiv preprint arXiv:2501.06848}, 2025.
\bibitem[Sohl-Dickstein et~al.(2015)Sohl-Dickstein, Weiss, Maheswaranathan, and Ganguli]{sohl2015deep}
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
\newblock Deep unsupervised learning using nonequilibrium thermodynamics.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 2256--2265. PMLR, 2015.
\bibitem[Song et~al.(2020)Song, Meng, and Ermon]{song2020ddim}
Jiaming Song, Chenlin Meng, and Stefano Ermon.
\newblock Denoising diffusion implicit models.
\newblock \emph{arXiv preprint arXiv:2010.02502}, 2020.
\bibitem[Song et~al.(2023{\natexlab{a}})Song, Vahdat, Mardani, and Kautz]{song2022pseudoinverse}
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz.
\newblock Pseudoinverse-guided diffusion models for inverse problems.
\newblock In \emph{International Conference on Learning Representations}, 2023{\natexlab{a}}.
\newblock URL \url{https://openreview.net/forum?id=9_gsMA8MRKQ}.
\bibitem[Song et~al.(2023{\natexlab{b}})Song, Vahdat, Mardani, and Kautz]{song2023pseudoinverseguided}
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz.
\newblock Pseudoinverse-guided diffusion models for inverse problems.
\newblock In \emph{International Conference on Learning Representations}, 2023{\natexlab{b}}.
\newblock URL \url{https://openreview.net/forum?id=9_gsMA8MRKQ}.
\bibitem[Song \& Ermon(2019)Song and Ermon]{song2019generative}
Yang Song and Stefano Ermon.
\newblock Generative modeling by estimating gradients of the data distribution.
\newblock \emph{Advances in neural information processing systems}, 32, 2019.
\bibitem[Tr'emolet(2006)]{tr2006accounting}
Yannick Tr'emolet.
\newblock Accounting for an imperfect model in 4d-var.
\newblock \emph{Quarterly Journal of the Royal Meteorological Society: A journal of the atmospheric sciences, applied meteorology and physical oceanography}, 132\penalty0 (621):\penalty0 2483--2504, 2006.
\bibitem[Uehara et~al.(2024)Uehara, Zhao, Biancalani, and Levine]{uehara2024understandingreinforcementlearningbasedfinetuning}
Masatoshi Uehara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine.
\newblock Understanding reinforcement learning-based fine-tuning of diffusion models: A tutorial and review, 2024.
\newblock URL \url{https://arxiv.org/abs/2407.13734}.
\bibitem[Uehara et~al.(2025)Uehara, Zhao, Wang, Li, Regev, Levine, and Biancalani]{uehara2025inferencetimealignmentdiffusionmodels}
Masatoshi Uehara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani.
\newblock Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review, 2025.
\newblock URL \url{https://arxiv.org/abs/2501.09685}.
\bibitem[Vincent(2011)]{vincent2011connection}
Pascal Vincent.
\newblock A connection between score matching and denoising autoencoders.
\newblock \emph{Neural computation}, 23\penalty0 (7):\penalty0 1661--1674, 2011.
\bibitem[Wang et~al.(2025)Wang, Han, Fan, Zhang, and Liu]{wang2025phyda}
Hao Wang, Jindong Han, Wei Fan, Weijia Zhang, and Hao Liu.
\newblock Phyda: Physics-guided diffusion models for data assimilation in atmospheric systems.
\newblock \emph{arXiv preprint arXiv:2505.12882}, 2025.
\bibitem[Wang \& Yu(2021)Wang and Yu]{wang2021physics}
Rui Wang and Rose Yu.
\newblock Physics-guided deep learning for dynamical systems: A survey.
\newblock \emph{arXiv preprint arXiv:2107.01272}, 2021.
\bibitem[Xu et~al.(2023)Xu, Liu, Wu, Tong, Li, Ding, Tang, and Dong]{xu2023imagereward}
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong.
\newblock Imagereward: Learning and evaluating human preferences for text-to-image generation.
\newblock \emph{Advances in Neural Information Processing Systems}, 36:\penalty0 15903--15935, 2023.
\bibitem[Yang et~al.(2023)Yang, Srivastava, and Mandt]{yang2023diffusion}
Ruihan Yang, Prakhar Srivastava, and Stephan Mandt.
\newblock Diffusion probabilistic modeling for video generation.
\newblock \emph{Entropy}, 25\penalty0 (10):\penalty0 1469, 2023.
\bibitem[Yu et~al.(2023)Yu, Lezama, Gundavarapu, Versari, Sohn, Minnen, Cheng, Birodkar, Gupta, Gu, et~al.]{yu2023language}
Lijun Yu, Jos{\'e} Lezama, Nitesh~B Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Vighnesh Birodkar, Agrim Gupta, Xiuye Gu, et~al.
\newblock Language model beats diffusion--tokenizer is key to visual generation.
\newblock \emph{arXiv preprint arXiv:2310.05737}, 2023.
\end{thebibliography}