Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion
Authors: Boyuan Chen
Diego Martí Monsó
Yilun Du
Max Simchowitz
Russ Tedrake
Vincent Sitzmann
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion
sequence modeling
decision making
planning
Primary area: diffusion_based_models
TL;DR: A novel way of training next-token prediction models to diffuse whole sequences at once enables long-horizon guidance, stable infinite rollout of continuous signals, and new planning & policy techniques.

Abstract: This paper presents Diffusion Forcing, a new training paradigm where a diffusion model is trained to denoise a set of tokens with independent per-token noise levels. We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens without fully diffusing past ones. Our approach is shown to combine the strengths of next-token prediction models, such as variable-length generation, with the strengths of full-sequence diffusion models, such as the ability to guide sampling to desirable trajectories. Our method offers a range of additional capabilities, such as (1) rolling-out sequences of continuous tokens, such as video, with lengths past the training horizon, where baselines diverge and (2) new sampling and guiding schemes that uniquely profit from Diffusion Forcing's variable-horizon and causal architecture,  and which lead to marked performance gains in decision-making and planning tasks. In addition to its empirical success, our method is proven to optimize a variational lower bound on the likelihoods of all subsequences of tokens drawn from the true joint distribution. Project website: https://boyuan.space/diffusion-forcing/

Diffusion Forcing: Next-token Prediction Meets
Full-Sequence Diffusion
Boyuan Chen
Diego Martí Monsó Technical University of Munich diego martietum...
Yilun Du
MIT CSAIL
MIT CSAIL
boyuancemit.. edu
yilundu@mit, edu
Max Simchowitz MIT CSAIL
Russ Tedrake MIT CSAIL russt0mi edu
Vincent Sitzmann
MIT CSAIL
msimchow@mit.es edu
sitzmann@mit. edu
Abstract
This paper presents Diffusion Forcing, a new training paradigm where a diffusion model trained denoise a set of tokens with independent per-token noise levels. We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens without fully diffusing past ones. Our approach is shown to combine the strengths next-token prediction models, such as variable-length generation, with the strengths of full-sequence diffusion models, such the ability to guide sampling to desirable trajectories. Our method offers range of additional capabilities, such (1) rolling-out sequences of continuous tokens, such video, with lengths past the training horizon, where baselines diverge and (2) new sampling and guiding schemes that uniquely profit from Diffusion Forcing's variable-horizon and causal architecture, and which lead to marked performance gains decision-making and planning tasks. In addition its empirical success, our method proven optimize variational lower bound on the likelihoods of all subsequences of tokens drawn from the true joint distribution. Project website: httpa://boyuan, space/ eeffussionffoocee
Introduction
Probabilistic sequence modeling plays a crucial role in diverse machine learning applications including natural language processing 6 47). video prediction 69 and decision making [3] 22 Next-token prediction models in particular have a number of desirable properties. They enable the generation of
sequences with varying length [32 21 37] (generating only single token an "infinite" number
of tokens via auto-regressive sampling). can be conditioned on varying amounts of history [2] 37
support efficient tree search (70 23 25). and can be used for online feedback control 22
Current next-token prediction models are trained via teacher forcinn (64, where the model predicts the immediate next token based on ground truth history of previous tokens. This results in two
limitations: (1) there is no mechanism by which one can guide the sampling a sequence to minimize certain objective, and (2) current next-token models easily become unstable on continuous data. For example, when attempting to rugr-segressyvely generate a video (as opposed to text (6) or cector-quantiza latents (33) past the training horizon, slight emrors in frame-to hame-to-frame -o-mmrm predictions accumulate and the model diverges.
Work done as a visiting student at MIT.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Guidance
Tree Search
Compositionally
Causal Unceriainty
Flexible Horizon
Comp Comp
Velog Xala)
Teacher Forcing
Fall-Soq. Diffusion
Diffusion Forcing
Figure 1: Diffusion Forcing capabilities. Today, different applications such as language modeling [6], planning [36], or video generation [31 69 rely on either auto-regressive next-token prediction or full-sequence diffusion, according to their respective unique capabilities. The proposed Diffusion Forcing is novel sequence generative model that enjoys key strengths of both model types.
Full-sequence diffusion seemingly offers a solution. Commonly used in video generation and long-horizon planning, one directly models the joint distribution of fixed number of tokens by diffusing their concatenation [31 D. where the noise level is identical across all tokens. They
offer diffusion guidance [30 16 to guide sampling to a desirable sequence, invaluable in decision- making (planning) upplications [3 34). They further excel at generating continuous signals such as video [31]. However, full- sullene diffusion is universally parameterized via non-causal, unmasked architectures. In addition to restricting sampling to full sequences, as opposed to variable length generation, we show that this limits the possibilities for both guidance and subsequence generation (Figure [ Further, we demonstrate that a naive attempt at combining the best of both worlds by training a toke- prediction model for full-sequence diffusion leads to poor generations, intuitively because it does not model the fact that small uncertainty in an early token necessitates high uncertainty a later one.
In this paper, we introduce Diffusion Forcing (DF), a training and sampling paradigm where each token is associated with random, independent noise level, and where tokens can be denoised according to arbitrary, independent, per-token schedules through a shared noxt-or-next-ee--toke prediction model. Our approach is motivated by the observation that noising tokens is form of partial masking zero noise means a token is unmasked, and complete noise fully masks out token, Thus, DF forces the model to learn to "unmask" any collection of variably noised tokens (Figure ( Simultaneously, by parameterizing predictions as composition of next-token prediction models, our system can flexibly generate varying length sequences as well as compositionally generalize to new trajectories (Figure
We implement DF for sequence generation as Causal Diffuusion Forcing (CDF), in which future tokens depend on past ones via a causal architecture. We train the model to denoise all tokens of sequence at once, with an independent noise level per token. During sampling, CDF gradually denoises a sequence of Gaussian noise frames into clean samples where different frames may have different noise levels each denoising step. Like next-token prediction models, CDF can generate variable-lengee sequences; unlike next-token prediction, it does so stabily from the immediate next token thousands of tokens in the future even for continuous tokens. Moreover, like full-seqeence diffusion it accepts guidance towards high-reward generations. Synergistically leveraging causality, flexible horizon, and variable noise schedules, CDF enables a new capability, Monte Carlo Guidance (MCG), that dramatically improves the sampling of high-reward generations compared to non-causal full sequence diffusion models. Fig. [ overviews these capabilities.
In summary, our contributions are: (1) We propose Diffusion Forcing, a new probabilistic sequence model that has the flexibility of next-token prediction models while being able to perform long- horizon guidance like full-sequence diffusion models. (2) Taking advantage of Diffusion Forcing's ' unique capabilities, we introduce a novel decision-maide making framework that allows us to use Diffusion Forcing as simultaneously policy ((10) and as a planner ((36)). (3) We formally prove that, under appropriate conditions, optimizing our proposed training objective maximizes a lower bound on the likelihood of the joint distribution of all sub-sequences observed at training time. (4) We empirically evaluate CDF across diverse domains such as video generation, model-based planning, visual imitation learning, and time series prediction, and demonstrate CDF's unique capabilities,

---

Noise Masking
Diffusion Foreing
Teacher Foreing
Full-Seq. Diffusion
Noise
Training
wwpe
Time
Observation
Latent State
Sampling
Generation
Add Noise
Method Overview. Diffusion Foreing trains causal sequence neural networks (such as an RNN or a masked transformer) to denoise flexible-length sequences where each frame of the sequence can have a different noise level. In contrast, next-token prediction models, common in language modeling,
Figure 2:
trained to predict single next token from ground-truth truth sequence (teacher forcing (64)), and full-sequence diffusion, common in video generation, train non-causal architectures to denoise all frames in sequence once with the same noise level. Diffusion Forcing thus interleaves the time axis of the sequence and the noise axis of diffusion, unifying strengths both alternatives and enabling completely new capabilities (see Secs. 3.2 3.4
such as stabilizing long-rollout autoregressive video generation, composing sub-sequences those observed at training time with user-determined memory horizon, Monte Carlo Guidance, and more.
2 Related Work and Preliminaries
We discuss related work and preliminaries for application, sequence generative modeling: see Appendix D for further literature review.
Our method unifies two perspectives on sequence modeling: Bayesian filtering along the time axis, denoted by subscript and diffusion along an "uncertainty" (or noise level) axis denoted by superscript k. In the following,
denote observations
X and latent states
Z.
Bayesian Filtering. Given Hidden Markov Model (HMM) defined by latent states z, and obser- vations a Bayes filter is probabilistic method for estimating latent states recursively over time from incoming observations. prior model p(z++1 [2)) infers a belief over the next state given only the current state, and an observation model infers belief over the next observation given the current latent state p(x,) ).). When a new observation is made, posterior model p(z,41 1 [Z,, X++1) provides an updated estimation of the next latent state Zz+1. When trained end-to-end with neural networks 22 123, latent states are not an estimate of any physical quantity, but sufficiently expressive latent that summarizes past observations for predicting future observations (xr)sse in the sequence.
Diffusion Models. Diffusion models [6 28 have proven be highly expressive and reliable generative models. We review their essentials here. Let (xx) denote a data distribution of interest, q. We consider forward diffusion process that gradually adds Gaussian noise to
and let x data point over series of time steps. This process is modeled Markov chain, where the data at each step lo is noised incrementally:
q(x (x**|x*-1) -N(x;; -
(2.1)
where A is the normal distribution and By is the variance of the noise added at each step controlled by a schedule {3e e (0, 1) The process continues until the data is converted into pure noise at = The reverse process is also a Markov chain and attempts to recreate the original data from the noise with parameterized model po:
po (x xxxxxx -Nxx f1(x k), 7x
(2.2)
where the mean la model with a neural network, and where it is shown (9 that one can set the covariance to the identity scaled by fixed constant fy depending on k. Adopting the standard

---

exposition, we reparametrize the mean f in terms of noise prediction c = (vI -x)=xxx This leads (28] to the following least squares objective:
Vālu
[()) = xx, [lle* cØ(x*, R)]l3],
(2.3)
where x² varx + v1 cge and E ~ (O, One can then sample from this model via Langevin dynamics x² yae (xf 1-ay EØ x k) + Fwww 288
VI-ue
Guidance of Difffusion Models. Guidance [30 T6] allows biasing diffusion generation towards desirable predictions at sampling time. We focus on classifier guidance [[6]: given a classifier
cyyx** of some desired (e.g. class or success indicator), one modifies the Langevin sampling
[29] gradient e (x*,, k) to be cx(x,k) -v as V log ) This allows sampling from the joint distribution of x and class label y without the need to train conditional model. Other energies
such as least-squares objective comparing the model output to desirable ground truth have been
explored in applications such as decision making [16] (6).
Next-Token Prediction Models. Next-token prediction models are sequence models that predict the next frame X++1 given past frames X1:t- At training time, one feeds neural network with X1:t and minimizes | x|| for continuous data or a cross-entropy loss for discrete data [64]. At sampling time, one samples the next frame x following p(x(+1 Xx.e). If one treats X as X,-1, one can use the same model to predict Xt+22 and repeat until full sequence is sampled. Unlike full-sequence diffusion models, next-token models do not accept multi-step guidance, as prior frames must be fully determined to sample future frames.
Diffusion has been widely used in sequence modeling. 43 use full- sequence diffusion models to achieve controllable text generation via guidance, such as generating text following specified parts of speech. [1 trains full-sequence diffusion models synthesize short videos and uses a sliding window to roll out longer conditioned on previously generated frames. [36 uses full-sequence diffusion models as planners in offline reinforcement learning. This is achieved by
Diffusion Sequence Models.
training on a dataset of interaction trajectories with the environment and using classifier guidance sampling time to sample trajectories with high rewards towards chosen goal. [49] modifies auto-regressive models to denoise the next token conditioned on previous tokens. It trains with
teacher forcing (64] and samples next-token auto-regressively for time series data. Most similar to our work is AR-Diffusion [65]. which trains full-sequence text diffusion with causal architecture with linearly dependent noise level along the time axis. We provide detailed comparision between
this approach and ours in Appendix D
3 Method
3.1 Noising as partial masking
Recall that masking is the practice of occluding subset of data, such as patches of an image 26 or timesteps a sequence [SS 48 and training model to recover unmasked portions. Without loss of generality, we can view any collection of tokens, sequential or not, as an ordered set indexed by t. Training next-token prediction with teacher forcing can then be interpreted as masking each token x; at time t and making predictions from the past X1.t-1- Restricted to sequences, we refer to all these practices as masking along the time axis. We can also view full-sequence forward diffusion, i.e. gradually adding noise to the data x X1-7, masking along the noise axis. Indeed, after K steps of noising, x K iir is (approximately) pure white
form of partial masking, which refer as
noise without information about the original data.
We establish unified view along both axes of masking (see Fig. We denote X1.7 for a sequence of tokens, where the subscript indicates the time axis. As above, x," denotes at noise level k under the forward diffusion process (2.1 xf x the unnoised token, and x is white noise NO, I). Thus, (x4* DSSTT denotes a sequence of noisy observations where each token has different noise level ki, which can be seen the degree of partial masking applied to each token through noising.
3.2 Diffusion Forcing: different noise levels for different tokens
Diffusion Forcing (DF) is framework for training and sampling arbitrary sequence lengths of noisy SISIS, where critically, the noise level ke of each token can vary by time step. In this tokens (x &s)

---

Algorithm Diffusion Forcing Training
Algorithm 2 DF Sampling with Guidance
1: loop
Input: Model 0, scheduling matrix K, initial latent
2: Sample tajectory of observations (X1, XT).
Zo. guidance cost cf-).
3: for t 1,.... T....To do
2: Initialize X1, Xi...... XT ~ (0, x
4: Sample independent noise level ls.
3:
for
= M 1,..,, do
{0, (0.1...,,K) K}
4: for m l..... do
5:
x ForwardDiffuse(
2i po zt 12-1,, Xx, 1,)+
x/-1
V
6: t Km,,, w ~ ((,
Define ct VI-aN
xnew vaa (Xt
now X x(, VI-aN co co
1-ok
7:
Update 21 ~ po po 1,llz--1,,,, Ks).
J&W
8: Update Z: + 2,
Ser = ED 11c--1,00 Xx kit)
9: end for
9:
end for
10: X1:I AddiGuiddacccc Vx log c(x11)
10:
I. -MSELoss((E1,.... En] [er,.... (i......n, 6n))
11: end for
11: Backprop with and update 0
12:
Return X1.r.
end loop
paper, we focus on time series data, and thus instantiate Diffusion Foreing with causal architectures (where x depends only on past noisy tokens), which we call Causal Diffusion Forcing (CDF). For simplicity, focus minimal implementation with vanilla Recurrent Neural Network (RNN) []] Potential transformer implementation of Diffusion Forcing is also possible but we defer its discussion to Appendix I.I
The RNN with weights maintains latents z, capturing the influence of past tokens, and these evolve via dynamics Z₁ ~ (z; x,, kz)) with a recurrent layer. When an incoming noisy observation kt is made, the hidden state is updated a Markovian fashion 2, ~ Pa pa (Z, 22, x k2 L When k2 0, this is the posterior update in Bayes filtering: whereas when k, K (and x is pure noise and thus uninformative), this is equivalent to modeling the "prior distribution" pa [ze Zz-1) in Bayes filtering. Given latent 2z, an observation model po po(x) (x] (2.) predicts X,.
Training. The dynamics model po 12-1, x," Iee) and the observation model pe (x{ 2:) together form a RNN unit. Such unit has the same input-output behavior as a standard conditional diffusion model, using a conditioning variable 2z-1 and a noisy token x, as input to predict the noise-free X; = x and thus, indirectly, the noise c² via affine reparamctrization [29]. We can thus directly train (Causal) Diffusion Forcing with the conventional diffusion training objective. We parameterize the aforementioned unit in terms of noise prediction Eş(2, x kt). We then find parameters 0 by minimizing the loss
T
E
|le: co(zi-1, k.]l
(3.1)
Ri x kk)
where we sample ki.T uniformly from [K] XI.7 from our training data, and NN, o I) in accordance with the forward diffusion process (see Algorithm [or pseudocode). Importantly, the loss (3.1) captures essential elements of Bayesian filtering and conditional diffusion. In Appendix B. we further re-derive common techniques in diffusion model training for Diffusion Forcing, which proves extremely useful for video prediction experiments. In Appendix C.2, we discuss the need of sampling K1.T uniformly. Finally, prove the validity of this objective stated informally in the following Theorem 3.1 in Appendix(A)
Theorem 3.1 (Informal). The Diffiusion Forcing training procedure (Algorithm 7 optimizes reweighting of an Evidence Lower Bound (ELBO) on the expected ligliiklissods In (x ) (x())))))) where the expectation is averaged over noise levels kı.1 ~ [K]T and x noised according to the forward process. Moreover, under appropriate conditions, optimizing also maximizes a lower
bound on the likelihood for all sequences of noise levels, simultaneously.
We implement 2, po (2, 11 X kx) to be deterministic, with z, representing distribution over beliefs rather than sample from it. This allows training by backpropogating through the latent dynamics in Eq B.1)
5

---

We remark that special case of 'all sequences of noise levels' are those for which either k2 0 or K; thus, one can mask out any prior token and DF will learn to sample from the correct conditional distribution, modeling the distribution of all possible sub-sequences of the training set.
Sampling. Diffusion Forcing sampling is depicted in Algorithm and is defined by prescribing noise schedule on a 2D M x 'I' grid K E [K]м*Т; columns correspond to time step t and rows indexed by m determine noise-level. Km,r represents the desired noise level of the time-step t token for row m. To generate a whole sequence of length 7', initialize the tokens X1:T to be white noise, corresponding to noise level k K. We iterate down the grid row-by-row, denoising left- to-right across columns to the noise levels prescribed by K. By the last row m = O, the tokens are clean, i.e. their noise level Ko.t 0. Appendix discusses corner cases of this scheme; the hyperparameters (ag., 0g. .k) are set to their standard values [29). The matrix IC specifies how fast each token gets denoised at every step of sequence diffusion. Since Diffusion Forcing is trained
to denoise tokens of all sequences of noise levels, K. can be designed to flexibly achieve different behaviors without re-training the model.
3.3
New Capabilities in Sequence Generation
We now explain the new capabilities this flexible sampling paradigm has to offer.
Stable Auto-Rog,
Diffuse w/
Full Traj. Guidance
Condition Currupted Obs.
Rollout
Causal Uncertainty
Stabilizing autoregressive generation. For nigh-dimensional, continuous sequences such video, auto-regressive architectures are known to diverge, especially when sampling past the training horizon. In contrast, Diffusion Forcing can stably roll out long sequences even beyond the training sequence length by updating the latents using the previous latent associated with slightly "noisy tokens" for some small noise level < k < K. Our experiments (Sec. 4.1) illustrates the resulting marked improvements in long-horizen generation capabilities; App. C.4 provides further intuition.
seque
Keeping the future uncertain. Beginning from sequence of white noise tokens [x x x JT,
we may denoise the first token fully and the second token partially, yiclding [x]. x2 K/2 .x3
K/2
K then
[xf, xy. x3 k and finally denoising all tokens fully [x{. x², x Interpreting the noise level as uncertainty, this "zig-zag sampling scheme intuitively encodes the immediate future as more certain than the far future. Sec. 3.4 describes how this leads more effective sequence guidance.
Long-horizon Guidance. trajectory X1.T as in Sec. 2 Due to the dependency of future tokens on the past, guidance gradients from future tokens can propagate backwards in time. The unique advantage of Diffusion Forcing is that, because we can diffuse future tokens without fully diffusing the past, the gradient guides the sampling of past tokens, thereby achieving long-horizon guidance while respecting causality. We elaborate implementation details in Appendix C.3, As show in Section 4.2 planning in this manner significantly outperforms guided full-sequence diffusion models.
Line 10 of Algorithm one may add guidance to the partially diffused
3.4 Diffusion Forcing for Flexible Sequential Decision Making
The capabilities offered by Diffusion Foreing motivate our novel framework for sequential decision making (SDM), with key applications to robotics and autonomous agents. Consider Markov Decision Process defined by an environment with dynamics p(s(41|s, a,), observation p(or st) and reward psr, a.). The goal is to train policy π (aa/o1d)) such that the expected cumulative reward of a trajectory E re] is maximized. We assign tokens X; |ar, Ft, 0 A trajectory is a sequence X1.r, possibly of variable length; training is conducted as in Algorithm T At each step t of execution, past (noise-free) tokens X1:t-1 are summarized by a latent 2 Conditioned on this containing predicted
latent, we sample, via Algorithm 2 plan &e:+ with &, = [āe, f., 0z+1] actions, rewards and observations. is a look-ahead window, analogous to future predictions in model predictive control [20). After taking planned action â, the environment produces reward re

---

DMLab
Minecraft
draw multiple Teacher Forcing Causal Full-Seq. Diffusion Forcing
Inpul
Generations
496
500
996
1000
Impul
Generations
496 500
996 1000
Seq
Seq
Seq
Seg
Seq 2
Seq I
Seit
Seq 2
Ses
Seq
Figure 3: Video Generation. Among tested methods, Diffusion Foreing generations are uniquely temporally consistent and do not diverge when rolling well past the training horizon. Please see the project website for video results.
and next observation 0.+-1 yielding next token = [āc. Fr,, Oc+1] The latent is updated according to the posterior po (z, 2 Xx, 0). Our framework enables functionality as both policy and planner:
Flexible planning horizon. Diffusion Forcing (a) can be deployed on tasks of variable horizon, because each new action is selected sequentially, and (b) its lookahead window H can be shortened to lower latency (using Diffusion Foreing as a policy), or lengthened to perform long- morizoo planning (via guidance described below), without re-training or modifications of the architecture. Note that (a) is not possible for full-sequence diffusion models like Diffuser [36 with full-trajectory generation horizons, whereas diffusion policies [10 need fixed, small lookahead sizes, precluding (b).
Flexible reward guidance. As detailed in Appendix C.3 Diffusion Forcing can plan via guidance using any reward (in place of loge) specified over future steps: this includes dense per-time step rewards on the entire trajectory Σ Ft, dense rewards on future lookahead E Ft, and sparse rewards indicating goal completion -1/or gll2. Per-time step policies cannot take advantage of this latter, longer horizon guidance.
Monte Carlo Guidance (MCG), future uncertainty. Causal Diffusion Forcing allows influ- ence the generation of a token x by guidance on the whole distribution of future X++1.7. Instead of drawing single trajectory sample to calculate this guidance gradient,
samples of the future and average their guidance gradients. We call this Monte Carlo Guidance. In the spirit of so-called shooting methods like MPPI [63], x is then guided by the expected reward over the distribution of all future outcomes instead of one particular outcome. The effect of MCG is enhanced when combined with sampling schedules that keep the noise level of future tokens high when denoising immediate next tokens (e.g. the zig- aag-zagaan schedule described in Sec. 3.3), accounting for greater uncertainty farther into the future. Appendix C.5 further justifies the sigmificance of MCG, and why Diffusion Forcing uniquely takes advantage of it.
4 Experiments
We extensively evaluate Diffusion Forcing's merits generative sequence model across diverse applications video and time series prediction, planning, and imitation learning. Please find the dataset and reproducibility details in the Appendix, well video results on the project website
4.1 Video Prediction: Consistent, Stable Sequence Generation and Infinite Rollout.
We train convolutional RNN implementation of Causal Diffusion Forcing for video generative modeling on videos of Minccraft gameplay [68 and DMLab navigation [68). At sampling time,

---

Maze2d-medium-
Margee-
start
end
denoising sleps
denoising sleps
DAF (Ours)
Diffuser
Environment
MPPI
CQL
IQL
Diffuser*
Diffuser w/ diffused action
Ours wo/ MCG
Ours
Mazc2D
U-Maze
33.2
47.4
113.9 + 3.1
110.1 3.9
116.7 +: 2.0
Mazc2D
Medium
10.2
2.7
13.5:2.3
136.1 10.2
149.4 7.5
Maze2D
Large
5.1
12.5
58.6
123.0 + 6.4
6.3 =2.1
142.8 5.6
159.0 + 2.7
Single-tusk Average
16.2
7.7
47.0
119.5
8.7
129.67
141.7
Multi2D
U-Maze
41.2
24.8
128.9
# 1.8
32.811777
107.7 + 4,9
119.1 + 4.0
Multi2D
Medium
15.4
12.1
127.2
+ 3,4
22.042,77
145.6 1: 6.5
152.3 +: 9,9
Multi2D
Large
8.0
13.9
132.1
= 5.8
=1.7
129.8 + 1.5
167.1 12777
Multi-task Average
21.5
129.4
20.6
127.7
146.2
Table 1: Diffusion Forcing for Planning. (top) During sampling, Diffusion Foreing allows each time
step to be denoised on different noise schedules, enabling us to account for causal uncertainty during guided planning. Diffusion Forcing keeps the far future more uncertain than the near future while Diffuser (36] puts them at the same noise level during sampling. (bottom) Quantitatively, Diffusion Forcing achieves the highest average reward across runs. Diffuser fails dramatically when executing the actually generated actions, requiring hand-crafted PD controller (indicated by the asterisk) and ignoring generated actions.
perform auto-regressive rollout with stabilization proposed in Sec. We consider two baselines, both leveraging the same exact RNN architecture: next-frame diffusion baseline trained with teacher forcing (64) as well as a causal full-sequence diffusion model. Figure TigureBdisplle displays qualitative results of roll-outs generated by Diffusion Forcing and baselines starting from unseen frames for both datasets. While Diffusion Forcing succeeds at stably rolling out even far beyond its training horizon (e.g. 1000 frames), teacher forcing and full-sequence diffusion baselines diverge quickly. Further, within the training horizon, we observe that full-sequence diffusion suffers from frame-to-frame discontinuity where video sequences jump dramatically, while Diffusion Forcing roll-outs show ego-motion through consistent 3D environment. This highlights the ability of Diffusion Forcing to stabilize rollouts of high-dimensional sequences without compounding emror..
4.2 Diffusion Planning: MCG, Causal Uncertainty, Flexible Horizon Control.
Decision makinn uniquely benefits from Diffusion Forcing's capabilities. We evaluate our proposed decision-making framework in standard offline RL benchmark, D4RL [18]. Specifically, benchmark Diffusion Forcing on a set of 2D maze environments with sparse reward. An agent tasked with reaching designated goal position starting from random starting position. In Appendix F.5 we provide a detailed description of the environment. The benchmark provides a dataset of random walks through mazes (thus stochastic). We train one model per maze.
We benchmark the proposed decision-makie framework 3.4 with state-of- tfch---rt- offline RL methods and the recently introduced Diffuser (36). diffusion planning framework. See Fig. T for qualitative and quantitative results: DF outperforms Diffuser and all baselines across all 6 environments.
Benefit of Monte Carlo Guidance. The typical goal for an RL problem is to find actions that maximize the expected future rewards, which we achieve through MCG. Full-sequence diffusion models such as Diffuser do not support sampling maximize expected reward, as we formally derive Appendix C.5) To understand MCG's importance, we ablate it in Table Removing MCG guidance degrades performance, though Diffusion Forcing remains competitive even then.

---

Stateful Cases
Generated Video
Goal Stare Intern. State
Case
Case
Input
Predictions
120 180 240 300
360 420 480
Outs, Cam Wristeum
Figure 4: our real robot task, a robot arm is asked to swap the slots two fruits using a third slot. Since the fruits are input in random slots the beginning. cannot determine the steps from single observation without knowledge of the initial placement of the fruits. As illustrated in (a) and (b), the upper observation is the same but the desired outcome illustrated below vary-the the task thus requires remembering the initial configuration. In addition, as shown in (c), the same model that generates actions also synthesizes realistic video from just a single frame.
Benefit of Modeling Causality. Unlike pure generative modeling, sequential decision-making takes actions and receives feedback. Due to compounding uncertainty, the immediate next actions more important than those in the far future. Though Diffuser and subsequent models are trained to generate sequences of action-reward-state tuples [a,, Tr, oz], directly executing the actions will lead to trajectory that deviates significantly from the generated states. In other words, the generated states and actions not causally consistent with each other. To address this shortcoming, Diffuser's implementation ignores the generated actions and instead relies on a hand-crafted PD controller to infer actions from generated states. In Table we see that Diffuser's performance drops dramatically when directly executing generated actions. In contrast, Diffusion Forcing's 's raw action generations are self-consistent, outperforming even actions selected by combining Diffuser's state predictions with handcrafted PD controller.
Benefit of Flexible Horizon. Many RL tasks have fixed horizon, requiring the planning horizon to shrink as an agent makes progress in the task. Diffusion Forcing accomplishes this by design, while full-sequence models like Diffuser perform poorly even with tweaks, as explain in Appendix
4.3 Controllable Sequential Compositional Generation
We demonstrate that by only modifying the sampling scheme, can flexibly compose sub-sequences of sequences observed at training time. We consider a dataset of trajectories on 2D, square plane, where all trajectories start from one corner and end up in the opposite corner, forming a cross shape. As shown in Fig. [ when no compositional behavior is desired, one can let DF keep full memory, replicating the cross-shaped distribution. When one desires compositionality, one can let the model generate shorter plans without memory using MPC, leading to the stitching of the cross's sub- tobb-rrrectoriss forming a V-shaped trajectory. Due to limited space, we defer the result to Appendix E.2
4.4 Robotics: Long horizon imitation learning and robust visuomotor control
Finally, we illustrate that Diffusion Forcing (DF) opens up new opportunities in the visuomotor control of real-world robots. Imitation learning 110 a popular technique in robotic manipulation where one learns of memory often prevents imitation lcarning from accomplishing mong-oorionn tasks. DF not only alleviates this shortcoming but also provides way to make imitation learning robust.
observation-to-actiee
mapping from expert demonstrations, However, the lack
Imitation Learning with Memory. We collect dataset of videos and actions by teleoperating with a Franka robot. In the chosen task, one needs to swap the position of an apple and an orange, using a third slot. See Fig. for an illustration. The initial positions of the fruits are randomized such that there are two possible goal states. As illustrated in Fig, 4 when one fruit is in the third slot, the desired outcome cannot be inferred from the current observation-a policy must remember the initial configuration to determine which fruit to move. In contrast to common behavior cloning methods, DF naturally incorporates memory in latent state. We found that DF achieves 80% success rate while diffusion policy [10], ofate-of-the-ar imitation learning algorithm without memory, fails.

---

Robustness to missing or noisy observations. Because it incorporates principles from Bayes filtering, Diffusion Forcing can perform imitation learning while being robust noisy or missing observations. We demonstrate this by adding visual distractions and even fully occluding the camera during execution. DF allows us to easily indicate these observations as "noisy" by using which case DF relies heavily on its prior model to predict actions. Consequently, the success rate is only lowered by 4% to 76%. In contrast, a next-frame diffusion model baseline attains a success rate of 48%: it must treat perturbed observations ground truth and suffers out-of-distribution error.
O, in
Potential for pre-training with video. Finally, parallel generating actions, Fig. 4 illlustrates that Diffusion Forcing is capable of generating a video of the robot performing the task given only an initial frame, unifying diffusion policy/imitation learning and video generative modeling and paving the way to pre-traiee pre-training on unlabeled video.
4.5 Time Series Forecasting; Diffusion Forcing a Good General-purpose Sequence Model
In Appendix show that DF is competitive with prior diffusion 49 and transformer-based [50]
work on multivariate time series forecasting, following the experimental setup of [53]
Discussion
Limitations. Our current causal implementation based on an RNN. Applications to higher- resolution video or more complex distributions likely require large transformer models following instructions in Appendix C.T We not investigate the scaling behavior of Diffusion Foreing to internet-scale datasets and tasks.
Conclusion. In this paper, we introduced Diffusion Forcing, a new training paradigm where a model is trained to denoise sets of tokens with independent, per-token noise levels. Applied to time series data, we show how a next-token prediction model trained with Diffusion Forcing combines the benefits of both next-token models and full-sequence diffusion models. We introduced new sampling and guidance schemes that lead dramatic performance gains when applied tasks in sequential decision making. Future work may investigate the application of Diffusion Forcing to domains other than time series generative modeling, and scale up Diffusion Foreing to larger datasets.
10

---

References
[1] A. Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and Agrawal. Is conditional generative modeling all you need for decision-making arXiv preprint arXiv:2211. 2022.
[2] A. Alexandrov, K. Benidis, M. Bohlke- Schneider. V. Flunkert, J. Gasthaus, T. Januschowski, D. C. Maddix, S. Rangapuram, D. Salinas, J. Schulz, L. Stella, A. C. Türkmen, and Y. Wang. Gluonts: Probabilistic and neural time series modeling in python. Journal of Machine Learning Research, 21 21(116):1-6, 2020.
[3] C. Berner, G. Brockman, B. Chan, V. Cheung, P. Debiak, C. Dennison, D. Farhi, Q- Fis- cher, S. Hashme, C. Hesse, R. Józefowicz, S. Gray, C. Olsson, J. Pachocki, M. Petrov, H. P. de Oliveira Pinto, J. Raiman, T. Salimans, J. Schlatter, J. Schneider, S. Sidor, I. Sutskever, J. Tang, F. Wolski, and S. Zhang, Dota 2 with large scale deep reinforcement learning. CoRR, abs/ abs/1912.06680, 2019.
[4] A. Blattmann, T. Dockhorn, S. Kulal, D. Mendelevitch, M. Kilian, D. Lorenz, Y. Levi, Z. English, V. Voleti, A. Letts, V. Jampani, and R. Rombach. Stable video diffusion: Scaling latent video diffusion models to large datasets, 2023.
[5] A. Block, A. Jadbabaie, D. Pfrommer, M. Simchowitz, and R. Tedrake. Provable guarantees for generative behavior cloning: Bridging low-level stability and high-level behavior. In Thirtt--eventh Conference on Neural Information Processing Systems, 2023.
[6] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- Voss, G. Krueger, Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray,
B. Chess, Clark, C. Berner, S. McCandlish, Radford, I. Sutskever, and D. Amodei.
Language models are few-shot learners. CoRR, abs/2005. 2020.
[7] S. H. Chan. Tutorial diffusion models for imaging and vision. arXiv preprint 0rXiv:2403.18103 2024.
[8] M. Chen, A. Radford, R. Child, J. Wu, H. Jun, D. Luan, and Sutskever. Generative pretraining from pixels. In International conference machine learning, pages 1691-1703. PMLR, 2020.
[9] T. Chen. On the importance of noise scheduling for diffusion models, 2023.
[10] C. Chi, Z. Xu, S. Feng, E. Cousineau, Y. Du, B. Burchfiel, R. Tedrake, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion, 2024.
[11] K. Cho, B. van Merrienboer, Ç. Gülçehre, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. CoRR, abs/1406. 1078, 2014.
[12] J. Chung, K. Kastner, L. Dinh, K. Goel, A. C. Courville, and Y. Bengio. A recurrent latent variable model for sequential data. Advances in neural information processing systems, 28, 2015.
[13] J. Cohen, E. Rosenfeld, and Z. Kolter. Certified adversarial robustness via randomized smooth- ing. In international conference on machine learning. pages 1310-1320. PMLR, 2019.
[14] E. de Bézenac, S. S. Rangapuram, K. Benidis, M. Bohlke-Schneidee R. Kurle, L. Stella, H. Hasson, P. Gallinari, and T. Januschowski. Normalizing kalman filters for multivariate time series analysis. In Advances in Neural Information Processing Systems, volume 33, 2020.
[15] Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: pre-tiaining of deep bidirectional transformers for language understanding. CoRR, abs/ Abs//1810.04855, 2018.
[16] P. Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. CoRR, his/2105.00233. 2021.
[17] C. Feichtenhofer, Li, K. He, et Masked autoencoders as spatiotemporal learners. Advances in neural information processing systems, 35:35946-35953. 2022.

---

[18] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
[19] S. Gao, P. Zhou, M.-M. Cheng, and S. Yan. Masked diffusion transformer is a strong image synthesizer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 23164-23173, 2023.
[20] C. E. Garcia, D. M. Prett, and M. Morari. Model predictive control: Theory and practice survey, Automatica, 25(3)333--348 1989.
[21] F. Gers, J. Schmidhuber, and F. Cummins. Learning to forget: continual prediction with Istm. In 1999 Ninth International Conference on Artificial Neural Networks ICANN 99. (Conf. Publ. No. 470), volume 2, pages 850-855 vol.2, 1999.
[22] D. Hafner, T. P. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. CoRR, abs/1912.01603, 2019.
[23] D. Hafner, P. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and Davidson. Learning latent dynamics for planning from pixels. CoRR,
2018.
[24] Hang, S. Gu, C. Li, Bao, D. Chen, H. Hu, X. Geng, and Guo. Efficient diffusion training via min-snr weighting strategy, 2024.
[25] N. Hansen, Wang, and H. Su. Temporal difference learning for model predictive control, 2022.
[26] K. He, Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[27] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. CoRR, abs/1512.00385, 2015.
[28] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems (NeurIPS),
2020.
[29] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. CoRR, abs/2006. 11239, 2020.
[30] J. Ho and T. Salimans. Classifier-free diffusion guidance, 2022.
[31] J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet. Video diffusion models, 2022.
[32] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Comput., 18)5 nov 1997.
[33] A. Hu, I. Russell, H. Yeo, Z. Murez, G. Fedoseev, A. Kendall, J. Shotton, and G. Corrado. Gaia-l: A generative world model for autonomous driving. arXiv preprint arXiv:2309.1708) 2023.
[34] S. Huang, 7. Wang, P. Li, B. Jia, T. Liu, Y. Zhu, W. Liang, and S.-C. Zhu. Diffusion-based gen- eration, optimization, and planning in 3d scenes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16750-16761, 2023.
[35] R. Hyndman, A. B. Koehler, K. Ord, and R. D. Snyder. Forecasting with Exponential Smoothing: The State Space Approach. Springer Science & Business Media, 2008
[36] M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffuusion for flexible behavior synthesis. Proceedings of the International Conference on Machine Learning (ICML), 2022.
[37] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. CoRR, abs/2006. 16236, 2020.
12

---

[38] L. Ke, J. Wang, T. Bhattacharjee, B. Boots, and S. Srinivasa. Grasping with chopsticks: Combating covariate shift in model-free imitation learning for fine manipulation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 6185-6191. IEEE, 2021.
[39] D. Kingma and R. Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
[40] R. G. Krishnan, U. Shalit, and D. Sontag. Structured inference networks for nonlinear state space models. In AAAI, 2017.
[41] G. Lai, W. Chang, Y. Yang, and H. Liu. Modeling long- and short-term temporal patterns with deep neural networks. CoRR, abs/1703.0700 703.07015, 2017.
[42] M. Laskey, J. Lee, R. Fox, A. Dragan, and K. Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning, pages 143-156. PMLR, 2017.
[43] X. L. Li, J. Thickstun, I. Gulrajani, Liang, and T. B. Hashimoto. Diffusion-Im improves controllable text generation, 2022.
[44] H. Lütkepohl. New Introduction to Multiple Time Series Analysis. Springer Science & Business Media, 2005.
[45] E. Matheson and R. L. Winkler. Scoring rules for continuous probability distributions. Management Science, 22(10):0087-1096 1976.
[46] A. Nichol and P. Dhariwal. Improved denoising diffusion probabilistic models. CoRR, Abs/2102.09672,2000 2021. 2021.
[47] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derezynski, X. Du, M. Grella, K. Gv, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, Lin, K. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. Wind, S. Woćniak, Z. Zhang, Q. Zhou, J. Zhu, and R.-. Zhu. RWKV: Reinventing RNNs for the transformer era. In H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023, pages 14048- 14077, Singapore, Dec. 2023. Association for Computational Linguistics.
[48] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. CoRR, 0bs/1910.1068.. 2019.
[49] K. Rasul, C. Seward, 1. Schuster, and R. Vollgraf. Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, 2021.
[50] K. Rasul, A.-S. Sheikh, I. Schuster, U. M. Bergmann, and R. Vollgraf. Multivariate probabilistic time series forecasting via conditioned normalizing flows. In International Conference on Learning Representations, 2021.
[51] D. Ruhe, J. Heek, Salimans, and E. Hoogeboom. Rolling diffusion models. arXiv preprint arXiv:2402.09470, 2024.
[52] T. Salimans and Ho. Progressive distillation for fast sampling of diffusion models. CoRR. abs/2202.00512,, 2022.
[53] D. Salinas, M. dimensional multivariate forecasting with low-rank gaussian copula processes. In NeurIPS, 2019.
L. Callot, R. Medico. J. Gasthaus, and R. Medico. High-
[54] D. Salinas, V. Flunkert, J. Gasthaus, and T. Januschowkki. Deepar: Probabilistic forecasting with autoregressive recurrent networks. International Journal of Forecasting, 36(3):111--119, 2020.
13

---

[55] D. Salinas, V. Flunkert, J. Gasthaus, and T. Januschowski. Deepar: Probabilistic forecasting with autoregressive recurrent networks. International Journal of Forecasting. 36(3):1181-1191 2020.
[56] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the International Conference on Machine Learning (ICML), 2015.
[57] J. Song, C. Men,, and S. Ermon. Denoising diffusion implicit models. CoRR, abs/2010.02502, 2020.
[58] B. Tang and D. S. Matteson. Probabilistic transformer for time series analysis. In A. Beygelzimer, Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[59] H. Touvron, P. Bojanowski, M. Caron, M. Cord, A. El-Nouby, E. Grave, A. Joulin, G. Synnaeve, Verbeek, and H. Jégou. Resmlp: Feedforward networks for image classification with data- efficient training. CoRR, abs/2105.03400 2021.
[60] A. Van den Oord, N. Kalchbrenner, L. Espeholt, O. Vinyals, A. Graves, et al. Conditional image generation with pixelenn decoders. Advances in neural information processing systems, 29, 2016.
[61] R. der Weide. Go-garch: A multivariate generalized orthogonal garch model. Journal of Applied Econometrics, 17(5):549-564 2002.
[62] C. Wei, K. Mangalam, P.-Y. Huang, Y. Li, H. Fan, H. Xu, H. Wang, C. Xie, A. Yuille, and C. Feichtenhofer. Diffusion models as masked autoencoders. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16284-16294, 2023.
[63] G. Williams, A. Aldrich, and E. Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint 1rXiv:1509.01149 2015.
[64] R. J. Williams and D. Zipser. A Learning Algorithm for Continually Running Fully Recurrent Neural Networks. Neural Computation, 1(2):270-28)) 06 1989.
[65] T. Wu, Z. Fan, X. Liu, Y. Gong, Y. Shen, J. Jiao, H.-T. Zheng, J. Li, 7. Wei, J. Guo, N. Duan, and W. Chen. Ar-diffusion: Auto-regressive diffusion model for text generation, 2023.
[66] T. Wu, Z. Fan, X. Liu, H.-T. Zheng, Y. Gong, J. Jiao, J. Li, J. Guo. N. Duan, W. Chen, et al. Ar- diffusion: Auto-regressive diffusion model for text generation. Advances in Neural Information Processing Systems, 36:39957-39974 2023.
[67] T. Yan, H. Zhang, T. Zhou, Y. Zhan, and Y. Xia. Scoregrad: Multivariate probabilistic time series forecasting with continuous energy-base based generative models, 2021.
[68] W. Yan, D. Hafner, S. James, and P. Abbeel. Temporally consistent transformers for video generation, 2023.
[69] R. Yang, P. Srivastava, and S. Mandt. Diffusion probabilistic modeling for video generation. Entropy, 25(10): 2023.
[70] S. Yao, D. Yu, J. Zhao, I. Shafran, T. L. Griffiths, Y. Cao, and K. Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023.
[71] H. Yu, N. Rao, and I. S. Dhillon. Temporal regularized matrix factorization. CoRR, 00s/3359.0033,,,0 2015.
14

---

A Theoretical Justification
In this section, we provide theoretical justification for the train of Diffusion Forcing. The main contributions can be summarized as follows:
We show that our training methods optimize reweighting of the Evidence Lower Bound (ELBO) on the average log-likelihood of our data. We first establish this in full generality (Theorem and then specialize to the form of Gaussian diffusion (Corollary A.2). We show that the resulting terms decouple in such a fashion that, in the limit a fully expressive latent and model, makes the reweighting terms immaterial.
We show that the expected likelihood over any distribution over sequences of noise levels can be lower bounded by a sum over nonnegative terms which, when reweighted, correspond the terms optimized in the Diffusion Forcing training objective maximizes. Thus, for fully expressive network that can drive all terms to their minimal value, Diffusion Forcing optimizes a valid surrogate of the likelihood of all sequences of noise levels simultaneously.
We begin by stating an ELBO for general Markov forward processes g(-), and generative models (.(.), and then specialize to Gaussian diffusion, thereby recovering our loss. We denote our Markov forward process q(-) as
K
q(x g(x2:k
1:
x9) = II {x* 1
(A.1)
and parameterized probability model
po(((x]) ISSSK, 20)221)
(A.2)
We assume that pe satisfies the Markov property that
Po (z x(
11:-11 (x )=s0<)) po(zi, x (2-1)
(A.3)
that is, the latent codes z is a sufficient statistic for x given the history. We say that pe has deterministic latents if (z Z1:-11, (x )1<.<(, x"" is Dirac delta.
Remark 1. In order for pe to have deterministic latents and correspond to valid probability distri- bution, we need to view the latents 2, not as individual variables, but as a collection of variables ze (A1.a) indexed by E [T] and the history of noise levels k1.1 10,1..... KJ In this case, simply setting 24 (ki:t) = (ki:t, (x Sigage tautologically produces deterministic latents. The reason for indexing ze (k1:)) with ki.e then arises because, otherwise, (z ((x** Issese (x- )is.s)) would be ill-defined unless k2- kg for all < < and thus, Pe would not correspond probability measure. The exposition and theorem that follows allow z, (kize) to be indexed on past noise levels k1 but suppresses dependence on kıtt to avoid notational confusion.
joint
A.1 Main Results
We can now state our main theorem, which provides evidence lower bound (ELBO) the expected log-likelln eegliikei of partially-noised sequences (x/1 IISST, under uniformly sampled levels ke and x
obtained by noising according to q(.) as in (A.I). Notice that this formulation does not require an explicit for of q()) or pe. but will specialize to Gaussian diffusion in the following section.
Theorem A.1. Fix xi.r. Define the expectation over the forward process with random noise level kET as
E F
E key... unif [C] x xg(xx (x))
E
F].
(A.4)
forward
and the expectation over the latents under (- conditioned on kirr, (xf ISSeT
E
E
KI.T, (x+ )
(A.5)
P.ILT
J.ST
15

---

Then, as long as pe satisfies the Markov property,
E
[In poppo(()) )ISIST)]
2 C(xi.r
forward
K Σ j=2
E
Σ (=1
(x X,, 21 11 +
DXL (a(x}- xx), |l po(x) xx- z2-11)
forward P.PLT
K K+1
K+1
where C(xLer) is constant depending only on x(=r (the unnoised data). Moreover, if the latents deterministic (i.e. po(z: xx-1,x;") a Dirac distribution), then the inequality holds with inequality if and only if q(x) +11: x; px t:1 exact.
X Z1-1), i.e. the variational approximation is
The proof of the above theorem given in Appendix A.2 Remarkably, involves only two inequalities! The first holds with equality under deterministic latents and the second holds if and only if variational approximation is exact: x(x, ((x/++1:''' x po px(x (x/++ti x 2z-1). This tightness of the ELBO suggests that the expression in Theorem A. is a relatively strong surrogate objective for optimizing the likelihoods.
A.1.1 Specializing to Gaussian diffusion
We special Theorem A. to Gaussian diffusion. For now, focus the "x-prediction" formulation diffusion, which is the one used in our implementation. The "E-prediction" formalism, used throughout the main body of the text, can be derived similarly (see Section 2 of 7 for a clean exposition). The following theorem follows directly by apply standard likelihood and KL-divergence computations for the DDPM 28
to Theorem A.I
Corollary A.2. Let
xx,) = N(x;; /1 Bax ,B&II, q(x
(A.6)
and define an = (1- Bx), āx = =II aj. Suppose that we parameterize po(x xx 1.-1) = (Hol(x) ,2)-1,5), 09). where further.
(1 - aaj-) x āj Then, as long as po satisfies the Markov property, we obtained
(1- a1-aj)aj X (x),
of (1 - az) v 00--1
Ho(x/ 2.
1 āj
1- āj
E
[n po ((x; (ISST)))))))))))) + C(x. 2
E
c xx (x. 2i 11 xp|1
forward
forward p.
K
++
j=1
T
E
Σ
1(k 2 1} kickk, (x 21-1, - xp12
forward P.AL.T
[=1
(1-0,) Pa 21 (1=))"-))
where above, we define cj
Proof. The first inequality follows from the standard computations for the "x-prediction" formulation of Diffusion (see Section 2.7 of 7 and references therein). The second follows by replacing the sum
over j with an expectation over ke unif
{0, 1,.,...,AA KJ.
We make a couple of remarks:
As noted above, Corollary A.2 can also be stated for e-prediction, or the so-called "v- prediction" formalism, all are affinely related.
Define an idealized latent 2 consisting of all past tokens (x"*) as well as of their noise levels kt. This is sufficient statistic for Z;-1, and thus can always view
16

---

xx(x""
zc-1.ke) = xxx(x Z(-1. where zz-1 is just compressing 2 When ap-
plying the expectation of X1:7 ~ q to both sides of the bound in Corollary A.2 and taking infimum over possible function approximator &g. we obtain
inf E E
E po 2 forward P.E..T
12g (xxx 21,1, ki) x0'1² = inf E E
E
11x (x,' Zc-1)
x; I1
forward p.
Var, x
(xx )isese ki...... kool
This leads to a striking finding: with expressive enough latents and po, we can view the maximization of each term in Corollary separately across time steps. The absence of this coupling means that the weighting terms are immaterial to the optimization, and thus can be ignored.
Given the above remarks, we can optimize the ELBO by taking gradients through the objective specified by Corollary (A.2, and are free drop any weighting terms (or rescale them) as desired. Backpropagation through Epæur is straightforward due to deterministic latents. This justifies the correctness of our training objective (3.1) and protocol Algorithm
A.1.2 Capturing all subsequences
Theorem A.] stipulates that, up to reweighting, the Diffusion Forcing objective optimizes a valid ELBO on the expected log-likelihoods over uniformly sampled noise levels. The following theorem can be obtained by straightforward modification of the proof of Theorem A. generalizes this to arbitrary (possibly temporally correlated) sequences of noise.
Theorem A.3. Let 'D be an arbitrary distribution over [K]*, and define Pilj k1:-1) = Pro
j 11.-1]. Fix xi.r- Define the expectation over the forward process with random noise level ki:r as
E
E 11
E
(A.7)
forward,
and the expectation over the latent under pe(.) conditioned on ki (x ISSET as
E
E
KI. (x,, S1SIST
(A.8)
p.z.
=x lax. xxx 1.AST
Then, as long as Po satisfies the Markov property.
E forward,
[Inpo( S1stTT) C(xl.r) +
E forward, P.Z..T
E
Σ E
where
E. = P:(1 k1tt1) ln xt, 21 + Σ jPll Dr. (1(xt-1 x2 x!)
px
x 2
where C(X).T) is constant depending only on xi. (the noise-free data), and where the inequality an equality under the conditions that (a) (z 2
24-1, xX+) ) is Dirac distribution (deterministic x 2)-1), 1), i.e. the variational approximation is
latents), and (b) g(xt++ k:f sharp.
xx() = Po k +1:
In particular, in the Gaussian case of Corollary
A.2
we have
T Σ 1(k 2 13kIk,, (x 24-1, k.) x!112
E Iorward,
((x))(()))ss)) In ((x )IsIsT)] + C(x..r)
E
||
forward, P.Z1:T
The most salient case for us is the restriction of 'D to fixed sequences of noise ki...., ky (i.e. Dirac distributions on [K]T). In this case, Pí k1s-1) objective need not be a lower bound on Egorward, p [In (ox(("())))ee)"""""""" ) However, the terms in the lower bound are, up to reweighting, an subset of those terms optimized in the training objective. Thus, light the remarks following Corollary A.2, a fully expressive network can optimize all the terms in the loss simultaneously. We conclude that, for fully expressive neural network, optimizing the training objective (3.1 is a valid surrogate for maximizing the likelihood of all possible noise sequences.
for all but j k2 and thus our training
17

---

A.2 Proof of Theorem A.1
Define E<el] as shorthand for E begin with the following claim
unit Neo
[K]
E Extr
(e)) 1x))
E pl
We
Claim (Expanding the latents). The following lower bound holds:
E [I (o((() (ISEST)
E
E
E
In po m-))]: (A.9)
forward
unIf (o.I..... R) x-" ~xxxx (x.)
Moreover, this lower bound holds with equality is ~ p(z= 2 x5x) is a Dirac distribution (i.e. deterministic latents).
Proof. Let's fix a sequence ky.r. It holds that
pe (x( Ke (ISET)
p(x;' Zz (xx 22,)
=
x1 t=1
(x Ze m(-1)
(Markov Property)
t=1
p(zz
26- 1, x (x 22-1)
t=1
T II
E
po (2-1).
(Importance Sampling)
2. 22 lz 1xx
Thus, by Jensen's inequality,
T
T
po (ises)) 2
E
In po (x 21-1)
E
In px
Zz zi-1)
2, p(2.
where the inequality is and equality when pe (z, Eforward to both sides of the above display, and invoking the Markov property of the latents, we conclude that
2g-1; xx: Dirac distribution. By applying
T
E forward
IIn po ((x," )(sss))]
>
E
E
In po (x,
22-1)
forward P.
E
E
E
In po 12-1))]
kt
will
I0.
KJ x nxg (x?)
We now unpack the terms obtained from the preceding claim.
Claim (ELBO w.r.t. q). It holds that
[x 22-11) 2 Ci(xo +
po (x
Zz Z
In
(x ke xx))
~a(x
(x. Fi
F.+l:
x
where Ci (No, ke) is a constant depending only on Xo and Ke and where the inequality holds with k.+1.1
equality if and only if gxx,
xxx1 (
x 2(-1).
18

---

Proof. We have that
E
In po
m0-)]]
10000 1 (xc)
E
In px(xx
1 Jdxx dx +1k
x m(x, (xcp)
Po po b
z(-1)
x [ (x())
1:00
(x)
1.k bx
x((1
In po (x beik =z-1)
E te 1 (x..
(Jensen's inequality))
hi (xx (xp?
xx 1:
(x
ketkiK (xx,
E
In
Po (x
(x Bick
Be
(Markov property (()))
x
& - (xc))
g(x g(xftttt bo 1:k
x
In po (x, kilK 2
= Cx0o ki) +
E 2ell
Aell bef)
+1 x()
x5
I. R -oxx te (xy) Inn +1x px))
where the constant C1 (Xo, k.)
depends only on Xy and ky-
x
g(x xx)
To check the conditions for equality, note that if qxx/::-
1-.--117 1:7 xx(()) po (x xx zi-1): then
In po(xticck Zo z q(x++++11 (x((1)
E t xxt)
In po ze-1) +
E + (x
In po (x
22-1, x
Mit (x)
x
Since In(.) is strictly concave, E Px mglx ti (xtt) (xtt) (x( po(x titk (x-1,x(") = q(x, g(x/itlik ).
In [oppo 2e-1)]
- if and only if
Claim (Computing the expected ELBO).
(x kr:K 2(-1) In
E
(x?)
(x + 1:k
x(1
K-1
= Cc((x, kt) + 1(ke = 03 In px( x1 2,
1) +
Σ
16 kz DKL. (g(x) xt++ x() || po(x I xgt1 (4t112211) 22-1)).
where C2(Xo. kg) is some other constant depending on Xy and ky-
Proof. The proof invokes similar manipulations the standard ELBO derivation for diffusion, but with a few careful modifications to handle the fact that we only noise to level k2, As is standard, we require the identity
g(x/
x()
q(x/ xx/-1,x]] = -((x x((x²))
(A.10)
q(x; -1 xx])
19

---

Part 1: Expanding the likelihood ratios
Using the above identity, we obtain
In po p k 2(-1)
q(x; <tt: xx]
Zz-1) + In po (x Extet11 1
K
In po(x 11
x/x 21-1)
= =Imp(x
g(xtittl x()
q(x] xx= xx)
j j=
@ |nx 11-1) + In px(x te xxtett 10-1)
K
In po (x
1-
x 20-1)
+ In 2(4 j-1 x
Σ
+
c(xtrt2 xQ)
g(x}
x/xx((1)
g(xg xx?)
x=
(1) Inp(x K| (2-1) + In px(x t 1 ixfitt
; 21111
(xf(e)
K-1
Σ In px
xt | z--1) 1)
+ In
g(x² 1i-
xx()
g(x²
xx(")
x(x]
x xx²)
/=
1mp(x;
z1-1)
In po(x
x
i+ , zz-1)
+ In px xx/tt x
q(x;' (x
j= ((x}
Xx +1 x!)
= =In(o(x²'
xx xx)+1))(( +
p(x; 20-1) q(x,²
In po(x
x kc+1 z
K-1
po (x}
xx x1 z(-1)
In
In
xx
((x
Net )
q(x/ xx x9)
j=
where (7) uses A. (ii) invokes cancellation in the telescoping sum, and the final display follows from the computation
q(x+ kt x x
( (x(xt
{1 I |xptttt
kt = 0
(A.11)
ho 2 1
Observe that, because we don't parameterize p(xf 21-1), In in(ate (g(x)) xx (11-1110211) +
lap( (2-1) can be regarded as some constant C"(x"" , Thus,
x(x;" (x?")
In po x 11-1)
Cx x 1++1, xxx) + In px x t z
K-1
po(x)
xx i+1 21)
Σ j=ke+1
q(
t |xx)
g(x+ x
g(x}
xx+,,xx1
(A.12)
Part 2: Taking expecations. We can now simplify to taking expectations. Observe that
po (x/ x 22-1)
E
In
DKI. (u(x)
x i+1 x9) || po(x | xf+1 2)-)),,
x
(x} x²" x()
and similarly,
px xx z
[In (inpo xt, 21-1)
h =
E Ng(x
In
beck
(x))
(x
x
ki+ 1 [1,21]
DRL
(gxt)" x x xx) ll po px xx -1-2)))))))
hg 2 1.
Finally, E tox wg( bekk (xq)
C'(x. x H++1 xx{() is a constant Cz(ke,, xo) depending only on k2, Xg-
Thus, from A.12
In po Ke K
2
E
qx{ + x2)
(xp)
R-1
= =Ca(k, xo) + llke = 03 In (x x 2z 1)+ +
Σ
DKL. (x(x)) xx++1 x²) || po(x) x(t1121-1) 21-1))
K
-C((,, + 1lke = =) po 11) +
1fs 2 ke DEL (g(x)) |x++| xx] || po I xt+* , -2-11)
==1
20

---

Completing the proof of the ELBO. previous two claims, we have
We are now ready to complete the proof. By combining the
E
(x R1-1)]
mg((x) (xP]
K-1
Cx(x0, + 1[k= 03 In po x 21-) + Σ 1; (1(x) x ;x+) po(x) xxft1 1)-)
where Ca(xo, kit) Ci ki)) + C2(xo, hi)) and where again, the above is an equality when
q(x
xx+t) ((x
po
unif (x+++1177 xx 11) Taking an expectation over hs {0.1....: K).....K..... we
have
j ++1
E
[1kk = =0}l = K +1
E
102
(A.13)
K +1
(o.1.....KK
he " (0.1....... K}
and consequently,
E
E
In (x("")))) )ISSTT
unif ke (U.I.....
KJ x, ma(x 1el)SSITT
K-1
j ++
2 C1(x+) + K +
px x + Σ
DEL [g(x) xt xx) |I po(x .z-1)11))
K -
==1
Invoking Claim
D
E
In po )(ses)
forward
T
E unff K) x ~ (o()
[nppo(x [)-
T
K
(x,x) + K+1 loxo
x 21-1) +
Σ
j+1
DXL (9(x)) ,x,,x
|| po(x xxf+| 22-11)
K+1
j=1
We conclude by observing that ΣT= Eee [Ca(x]) is constant C(x4er)) and that
E <t
po (x
x 11] forward p.21.7
E
E [o
x x1, 2c-1)]
E
DKL
{x} xx+ xX) po xx/
x
21-1)]]
E
E
DEL (g(x))
x
x
||
po
xt =.
forward p.71.
since both terms only depend on k1:-11, (x )isese-1 and Z1 We conclude then that
E In px( )ISTT)) 2 C(x..T)
forward
T
K-1
E
Σ
K K+1 (x xt 2111 +
Σ j=1
j k DzC
(a(x) x3 x²) po(x) xxx sem1)))
forward P.EL.
needed. Lastly, recall that the above is an equality under the conditions that (a) (z ,,,1, is Dirac distribution, and (b) q(x,+++:T xx*)
k.)
po + x²,* 2(-1).
and we reindex j + j+ + 1-11 to ensure consistency with indexing in standard expositions the diffusion ELBBO.
B Additional Method Details
B.1 Fused SNR reweighting
SNR reweighting 24 is a widely used technique to accelerate the convergence of image diffusion models. In short, it reweighs the diffusion loss proportional to the signal-to-noist ratio (SNR) of
21

---

noisy x². In Diffusion Forcing, conditioning variable 2 can also contain a non- non-trivial amount of information about x,, in addition x For example, in deterministic markovian system, if x has its noise level 1 O, the posterior state 2c-1 contains all the information needed predict x regardless of the noise level of x
Therefore we re-derive SNR reweighting to reflect this change in Diffusion Forcing. We call this technique Fused SNR reweighting. We follow the intuition of original SNR reweighting to loosely define SNR in a sequence with independent levels of noises at different time steps. Denote S, as the normalized SNR reweighting factor for = following its normal derivation in diffusion models. For example, if one uses min snr strategy [24]. its reweighting factor will always fall between [o, C] which we divide by C to get St E [0, 1]. Define signal decay factor 0 < < 1, measuring what proportion of signal in x contribute to denoising x This is the simple exponential decay model of sequential information. Now, define cumulated SNR recursively as the running mean of S;: S, = - + (1 S, account for signals contributed by the entire noisy history to the denoising at time step t. The other factor that contributes to the denoising is S. of noisy observation k To combine them, we use a simplified model for independent events. Notice S, and S, always falls in range [0, 1]. and therefore can be reinterpreted as probabilities of having all the signal one needs to perfect denoise x Since the noise level at is independent of prior noise levels, we can view St and S 1 as probabilities of independent events and thus can composed to define joint probability S = 1 - (1 5,)(1 1,-1), and we use this Sy as our fused SNR reweighting factor for diffusion training.
R
24 derive the S. Our In our experiments, we choose to follow the min-SNR reweighting strategy Fused SNR reweighting proves extremely useful to accelerate the convergence of video prediction, while we didn't 't observe a boost on non-image domains so we didn't use it there.
B.2 Architecture
Video Diffusion We choose both the raw image x and latent state z to be 2D tensors with channel, width, and height. For simplicity. we use the same width and height for x and z. We then implement the transition model p(x.' 12:1) with a typical diffusion U-net [46]. We use the output of the U-net as the input to gated recurrent unit (GRU) and use 2 as the hidden state feed into GRU. The output of GRU is treated as 2, For observation model p(x, 2t). we use a 1-layer resnet (27] followed by a conv layer. We combine these two models to create an RNN layer, where the latent of a particular time step is 2,-1, input is x" and output is &. One can potentially obtain better results by training Diffusion Forcing with a causal transformer architecture. However, since RNN is more efficient for online decision-makin results.
also stick with it for video prediction and it already gives satisfying
We choose the number of channels in z to be 16 for DMlab and 32 for Minecraft. In total, our Minecraft model consists 36 million parameters and our DMlab model consists of 24 million parameters. We can potentially obtain a better Minecraft video prediction model with more parameters, but we defer that to future works to keep the training duration reasonable (< day). In maze planning, the number of total parameters is 4.33 million.
Von-Video Diffusion For non-spatial x that is not video nor images, we use residue MLPs instead of Unet as the backbone for the dynamics model. Residue MLP is basically the ResNet 27 equivalent for MLP. Similar to video prediction, we feed the output of resMLP into a GRU along with 2,- to get 2, Another ResMLP serves as the observation model.
B.3 Diffusion parameterization
In diffusion models, there are three equivalent prediction objectives, Xp, c (28), and v parameteriza- tion [52]. Different objectives lead to different reweighting of loss different noise levels, together with SNR reweighting. For example, c parameterization and parameterization are essential in generating pixel data that favors high-frequency details.
In our experiments, we use e parameterization for video prediction and found it essential to both convergence speed and quality.
22

---

We observe that Xo parameterization is strongly favorable in planning and imitation learning, likely because they don't 't favor an artificial emphasis on high-frequene, details. We observe the benefits of v-parameterization in time-series prediction.
B.4
Noise schedule
We use sigmoid noise schedule [9] for video prediction, linear noise schedule for maze planning, and cosine schedule for everything else.
B.5 Implementation Details of Sampling with Guidance
Corner case of sampling noise In our sampling algorithm, due to the flexibility of the scheduling matrix IC, there are corner cases when x is required to stay at its same noise level during a sampling step. The core question of this corner case is whether we should updatex, at all. One option is just copying over the old value. The other option is to run backward diffusion followed by forward diffusion back to its old noise level to resample under the diffusion process. While we conclude this can be an open question, we prefer the later approach, resampling, and use it in Monte Carlo Guidance to generate multiple samples. We note that even if one takes the first approach, the guidance gradient can still flow back in the time steps before as the dynamics model p(z:|x/) propagate the guidance gradient to 2:-1-
kt Z;-1) can still
Other than Monte Carlo Guidance, this corner case only happens when k = 0 or hS = K throughout experiments. That is, we chose our IC such that once any token gets diffused slightly, it will keep diffusing. In the case of k = K, keeping x the same noise level implies it will stay as white noise, and we don't I even need to sample another white noise. In case h; = 0, the time step is already completely diffused either approach should give the same result we just opt for copying over for simplicity.
Guidance for maze planning In maze planning, our main baseline Diffuer (36] discards the reward from the dataset and directly plans with the goal position and velocity. We adopt the same convention for Diffusion Forcing. One can perform guidance on goal position using likeliieddooo ||rr gll- but a flexible horizon model should not require users to manually specify T to reach its goal, instead want it try to reach the goal for any possible horizon. Therefore we use the reward model E. |Pr gll so any time step can be the final step to reach the goal. This objective is challenging due to the non-convex nature of 2D maze, but we found Diffusion Forcing can still reliably find plans without bumping into walls. However, we also observe that the agent tend to leave the goal location due to the nature of the provided dataset the goal location just one possible waypoint for the robot to pass through, and there are no trajectories that simply stay at the goal. We also tried this reward for guidance with Diffuser, but didn't work even with good amount of tuning.
B.6 Performance Optimization
Accelerating the diffuuson sampling of Diffusion Forcing is similar to that of normal diffusion models. We adopt DDIM 57 sampling for the diffusion of each token. While we use K 1000 steps of diffusion, we sample with only 100 DDIM for video prediction and 50 for non-video domains.
While Diffusion Forcing can be implemented with transformers, we use an RNN as the backbone for Diffusion Forcing experiments it's widely used in decision-making for its flexibility and efficiency in online decision-makinn systems. To further reduce training time and GPU memory usage, we use frame-stacking to stack multiple observed images as a single x. This is due to the fact that adjacent tokens can be very similar- e.g. recording the same motion at higher fps can lead to this. We deem that it's wasteful if we roll out the dynamics model multiple times to generate almost identical tokens. For video datasets, we manually examine how many time steps it takes to require a minimal level of prediction power instead of copying frames over. There is another reason why we use frame stacking many diffusion model techniques such as different noise schedules are designed to model x with correlated elements or redundancy. Low-dimensional systems may need drastically different hyperparameters when they lack the data redundancy these techniques are tested on. Frame stacking is thus also helpful for our non-image experiments we can start with canonical hyperparameters of diffusion models. We use frame stack of 4 for DMIab video prediction, 8 for Minecraft, and for maze planning.
23

---

At sampling time, we also have design choice to reduce compute usage, reflected in line 8 of Algorithm 2 In line 8, we directly assign z, Z,, instead of recalculating z, with posterior model p(zz|Zt-1, [2t-1,, X k 1). Since the model is trained to condition on estimated from arbitrary noisy history, we recognize that both are valid approaches. The reason why the choose line is twofold. First, it cuts the compute by half, avoiding computing posterior every step. Second, this happens to be what we want for stabilization i2 new already contains the information of the clean x under our simplified observation model, and happens be estimated with k k2, noise level higher than that of x This happens to implement the behavior we want for stabilization.
B.7 Sampling schedule for causal uncertainty
Inference is depicted in Algorithm 2 and Figure 2 In Equation B.I we illustrate specific instantiation of the K. matrix we used for causal planning. For simplicity, we denote the case where a latent 2g given and aim to generate X1:H+1-
K K
K
K
K
K
K
K
K 2
K
-
K
K
pyramid
2
3
(B.1)
O
1
2
H
0000
0
Diffusion Forcing begins by sampling our sequences as white noise with noise level K. It then denoises along each sequentially through frames t = 1.... T, updating the latent (Line 5 Algorithm 2}. and then partially applying the backward process to noise level ic Kmm, dictated by the scheduling matrix IC (Line 6-7 of Algorithm D We call a IC like this pyramid scheduling, as the tokens in the far future kept higher noise level than future.
1, ................................ M of K. in decreasing order. It does by proceeding
B.8 Metrics for Maze Planning
We report the episode reward of Diffusson Forcing for different maze planning environments in Table However, we found that the episode reward isn't necessarily good metric: Intuitively, maze planning should reward smart agents that can find the fastest route to the goal, not a slow-walking agent that goes there the end of the episode. The dataset never contains data the behavior of staying at the goal, so agents are supposed to walk away after reaching the goal with sequence planning methods. Diffuser may had unfair advantage of just generating slow plans, which happens to let the agent stay in the neighborhood of the goal for more steps and get a very high reward as result. This metric seems exploit flaws in the environment design good design would involve penalty of longer time taken to reach the goal. Therefore, in future works based on our paper, we encourage alternative metries like the time takes reach the goal for the first time, which Diffusion Forcing excels at.
B.9 Implementation Details of Timeseries Regression
We follow the implementation of pytorch-ts, where the validation set is random subset of the training set with the same number sequences as the test set. We use carly stopping when validation crps-sum hasn't t increased for 6 epochs. We leverage the same architecture mlp and grus) well as a batch size of 32.
B.10 Compute Resources
All of our experiments use /p16 mixed precision training. Time series, maze planning, composition- ally, and visual imitation experiments can trained with single 20807i with 11GB of We tune the batch size such that we fully use the memory of GPUs. This translates to batch size of
memory.
24

---

2048 for maze planning and compositional experiments, and 32 for visual imitation learning, While we use early stopping on the validation set for time series experiments. we did not carefully search for the minimal number of training steps required, though the model usually converges between 50k to 100k steps. The above environments thus usually take 8 hours to train although there is without doubt significant potential for speed up.
Video prediction is GPU intensive. We use 8 A100 GPUs for both video prediction datasets. We train for 50K steps with a batch size of 8 x 16. It usually takes 12 hours to converge at 40K steps of training (occasional validation time also included).
C
Additional Intuitions and Explainations
C.1 Extension to transformer backbone
While this paper focuses on a causal implementation of Diffusion Forcing with RNNs, it's easy to adopt Diffusion Forcing with modern architectures like transformers. One can simply modify transformer- based sequence diffusion model to train with independent noise levels across tokens and follow the techniques listed in Section [B.] A strict implementation of causal Diffusion Forcing would involve a causal attention mask on the transformer. However, Diffusion Forcing's 's fractional masking can do something more interesting: Consider the scenario that we use transformer without a causal mask. We can still implement causality by controlling noise. By labeling the future as full white noise, there is information leaked into the past tokens. By labeling future tokens free of noise, we make the model completely non-causal. By labeling the future tokens as noisy, slight amount of information about the future is provided for the prediction of past tokens. This effectively states that one only needs a non-causal architecture, but controlling fractional noise of the future, to achieve partial or complete causality. These extensions are beyond the scope of this paper, but we already verified their effectiveness and thus provide them as intuitions for future works.
C.2
The need for independent noise levels
When training Diffusion Forcing, we choose to sample per-token noise level following i.i.d uniform distribution from [1, k]. One may wonder about the necessity of this choice. Here we discuss the unique abilities of independent noise and the compute overhead added by it.
The use of independent noise confers a number of special capabilities in our model, including stabilization of autoregressive rollout 3.33 modeling causal uncertainty (33] and removing the need for expensive reconstruction guidance when conditioning on context [C.6, None of these capabilities can be achieved by full-sequence diffusion. AR-diffusion (66) and Rolling Diffusion [51] can only achieve the first and third onc. There are more sampling-time applications such as flexible frame interpolation. Finally, we also saw the practical benefits of using independent noise in hyperparameter tuning. One can simply try different sampling schemes to figure out the most effective one for their applications. All these capabilities only require training the model once with Diffusion Foreing. In contrast, any tuning of the sampling scheme would require re- re-training the model for AR AR-diffusion and Rolling Diffusion.
On the other hand, we didn't observe much computing overhead when comparing Diffusion Forcing to full-sequence diffusion, as soon as one closely follows our training techniques like B. The empirical evidence is based on our experiments with an experimental transformer implementation of Diffusion Forcing and is thus not fully consistent with the main paper. However, we present the high-level descriptions below for readers interested in more insights: The complexity added by independent noise levels is in the temporal dimension. Therefore, we first adopt a standard technique for video diffusion models- image pre- mre-trainngg, to abstract away the complexity of the image pixels themselves. Then the complexity left is temporal prediction only. We then take the pre-trained image-only model and continue training it on video data. It turns out the sampling result of Diffusion Forcing with fewer training steps in this second stage already better than that of full-sequence diffusion at convergence. We speculate that the better result is due to the same data-augmentation effect described prior works [39]. This shows that the overhead added by independent noise is werr-nerranted when considering the overall training compute (including image mraitriin
25

---

C.3 Guidance as planning
As stated in Section 2 one can use the gradient of the logarithmic of a classifier log e(ylx;) ) to guide the sampling process diffusion model towards samples with a desired attribute y. For example, y can refer to the indicator of success event. However, we can consider the logarithmic of a more general energy function c(x+). This has the interpretation Pr(y]x²), where Pr(y = 1 xt] ex(x) Some popular candidate energies include
=
Σ r
(x
x
(C.1)
corresponding to cost-to-go; we can obtain unbiased estimates of this gradient by using cumulative reward c(xf) Eess (x, ke ). We can also use goal distance - -|| gll as terminal
reward. We provide details about the guidance function deployed in the maze2d planning experiment in Appendix B.3
C.4 Noising and stabilizing long-horizon generations
Here, we explain detail how we use noising to stabilize long- generation. At each time t, during the denoising, we maintain a latent z,-1 from the previous time step. with 0 ksmall < K corresponding to some small amount of noise. We then do next token diffusion to diffuse the token x; across noise levels x x² K-1 ,...,x? x (corresponding to Algorithm 2 with horizon = initial latent 2²_1. and noise schedule Kmm,1 = m); this process also produces latents K zi z K-1 .....,2; zi associated with each noise level. From these, use the latent z,
Asmall to repeat the process. This noised latent be interpreted as implementation of conditioning on x xffmalll in an autoregressive process. In a potential transformer implementation of Diffusion Forcing as we discussed in Appendix C.I one can instead run a forward diffusion fully diffused token to achieve stabilization.
is widely appreciated that adding noise data ameliorates long-term compounding error in behavior cloning applications [88 42). and even induces robustness to -equeequlttaal adversarial attacks [L3 In autoregressive video generation, the noised x Aamall is in-distribution for training, because Diffusion Forcing trains from noisy past observation in its training objective. Hence, this method can be interpreted as a special case of the DART algorithm for behavior cloning [4),, where the imitiator (in case, video generator) is given actions (in our case, next video frames) from noisy observations (in our case, noised previous frames). Somewhat more precisely, because we use both tokens at training time train Diffusion Forcing, and using slightly noised tokens for autoregression at test time, our approach inherits the theoretical guarantee of the HINT algorithm (5).
C.5 Why Monte Carlo Guidance relies on Diffusion Forcing
Monte Carlo Guidance provides substantial variance reduction in our estimate of cost-to-go guidance (C.I). This technique crucially relies on the ability to roll out future tokens from current ones to use these sample rollouts to get Monte Carlo estimates for gradients. This is not feasible with full-sequence diffusion, because this requires denoising all tokens in tandem; thus, for a given fixed noise level, there is obvious source of randomness use for the Monte Carlo estimate. It may be possible to achieve variable horizon via the trick proposed in the following subsection to simulate future rollouts, but our knowledge, this approach is nonstandard.
C.6 Does the replacement technique lead to flexible horizons in full-sequence diffusion?
naive way to obtain flexible horizon generation in full- sequence diffusion is via the "replacement trick": consider a full sequence model trained to diffuse X1.r, which we partition into X1:-11, X Having diffused tokens X we can attempt to denoise tokens of the form [ x where we fix x = X1:-1 to be the previously generated token, and only have score gradients update the remaining x One clear disadvantage of this method is inefficiency one still needs to diffuse the whole sequence even when there is one step left / = T 1. What's more, [31 points out that this approach of conditioning, named "conditioning by replacement", is both mathematically unprincipled
26

---

Sample multiple futures compute guidance gradient for x allowing planning with expected reward
Corrupted History:
Classifier guidance on 3y can propagate through the trajectory because near future hasn't been fully diffused
Set (0, K) so model cautiously treats x ass noisy observation.
Far future:
History: model trusts x as observation to condition
Loop k from K... while staying bigger than the near future. Model diffuses x conditioned on partially diffused near future.
Near future:
Loop from K...0
model diffuses x prediction.
Figure 5:
Diffusion Forcing is trained on independent level of noises different timesteps. As result, we can control the noise level l to achieve different effects on conditioning and prediction.
and can lead inconsistency in the generated sequence. The best fix proposed by [T] incorporates an additional gradient term with respect Xx.r at every diffusion step; this is still an incomplete fix and suffers the computation cost of an backward propagation for every sampling step.
C.7 Further connection to Bayesian filtering
The core idea of Diffusion Foreing can be interpreted as using diffusion to construct an interpolation between prior distribution and posterior distribution of Bayes filter. Consider the hybrid distribution
(/ xf). When = 0, this hybrid distribution becomes the posterior p( p(z: Zt-1. Xt). On the other hand, when & K, the hybrid distribution becomes p(z||zt-1,1) n) for
(0, 1). Since the independent Gaussian noise term n contains no information about z. this exactly the prior distribution p( 12-1). By varying & between K and 0, the same neural network parameterize everything between prior and posterior.
C.8 Connection to other sequence training schemes
Noise as masking provides a unified view of different sequence training schemes. The following exposition uses length 3 sequence as example: We always start with fully masked sequence [xX x K1 with the goal of denoising it "clean sequence" of zero noise. [x]. x², x(] Assume all diffusions are sampled with 3-step DDIM.
Autoregressive. In teacher forcing, one trains a model to predict the next token conditioned on prior observations. One can train next-token diffusion models with teacher forcing such [4]] feed neural network with past observations as well as current observation and ask it to predict clean current observation. A typical training pair can have the input of [x]. xy. $1 and target of [xf. x0, T
27

---

At sampling time, one fully diffuses the next token before adding the diffused observation to history perform an autoregressive rollout. The diffusion process would thus look like
xx x2 x3 1
X [x] K//2 X2 K
[x] X2 X
[x] X2 K//2 xx3 x
[x] X3 1T
x x2, K|/2T
[x] x2, x
Notably, Diffusion Forcing can also perform this sampling scheme at sampling time for applications like imitation learning, when one wants diffuse the next action as fast as possible.
Full Sequence Diffusion. Full sequence diffusion models accept a noisy sequence and denoises level-by-level
K T [x X2 X3
[x] K112 K112 ,x3 X2
K//2
[x]. x2: X3]
Notably, Diffusion Forcing can also perform this sampling scheme at sampling time.
Diffusion Forcing with causal uncertainty As shown in Figure 2 to model causal uncertainty, Diffusion Forcing keeps the far future more uncertain than the near future by having larger noise level k, at any time of diffusion. An example pattern looks like this:
X X2 x3 1T
K//2 [x1 X2 X3
T
[x] X2 K//2 xx3 T
x 0 K//2TT X3
[x] x2, TI1
Notable, (65] the first one to propose such a linear uncertainty sampling scheme for causal diffusion models, although Diffusion Forcing provides generalization of such scheme in combination with other abilities.
Diffusion Forcing with stablization Previously we introduced the autoregressive sampling scheme that Diffusion Forcing can also do. However, such a scheme can accumulate single-step errors ground truth observation. Diffusion Forcing addresses this problem
because it treats predicted by telling the model that generated images should be treated as noisy ground truth, as shown in 2
It first fully diffuses the first token,
x X JT
x K//2 K KIT
[ x x51
Then, it feed the diffused xf into the model but tell of a slightly higher noise level, as x to diffuse X2.
[ x X2 X
K//2
K
[x] x2 X3
Then, feeds the diffused x6 into the model but tells is of higher noise level, x!.
K//2 [x x X3
[x] x x3]
01 T
28

---

Frecktion windwe
Nozian prodicion
52% prodiction anceral
Feature
Feoture
Feature
Feature
Feature 5
Peoture
Peature
Peature
100
WWV
123
Faature 9
Feature 10
Faature 11
Fenture 12
300
100x
1019
122
Feature 13
Feature 14
Feature 15
Feature 16
1/4
173
Figure 6: Prediction intervals of Diffusion Foreing for the first prediction window of the test set in the Electricity time series dataset. Only the first 16 features out of 370 are plotted.
D Extended Related Work
Reconstructing masked tokens. Masked Autoencoders for images 226 and videos 17 are
popular method for representation learning in pixel space. They have been extended perform diffusion to generate masked patches conditioned on unmasked ones 62
19
Casting Image Generation as Sequence Generation. of non-sequential data, such as images, can be fruitfully cast as sequence generative modeling.
660 show that generative modeling
Non-Diffusion Probabilistic Sequence Models. 112 parameterize token-to-token transitions via variational auto-encoder. This makes them probabilistic, but does not directly maximize the joint probability of sequences, but rather, enables sampling from the distribution of single-step transitions.
Sequence Diffusion with Varying Noise Levels. Most similar to our work is AR-Diffusion [65 which similarly aims to train next-token prediction models for sequence diffusion. Key differences are that AR-Diffusion proposes noise level that is linearly dependent on the position of each word in the sequence, while our critical contribution is to have each noise level be independent, as this uniquely enables our proposed sampling schemes, such as stabilizing auto-regressive generation and conditioning corrupted observations. Further, AR-Diffusion only explores language modeling and does not explore guidance, while we investigate Diffusion Forcing as a broadly applicable sequence generative model with particular applications to sequential decision-making In particular, we introduce Monte-Carlo Guidance as a novel guidance mechanism. Another closely related work Rolling Diffusion [51]. which proposes to diffuse sequence with near future more certain and far future more uncertain, resembling the causally uncertain sampling scheme of Diffusion Forcing. Like DR-fiffussion, Rolling Diffusion's 's training noise levels are linearly dependent the positions of tokens and must use the exact same noise level scheme at sampling time. It, therefore, shares the aforementioned limitations of AR-Diffusion well.
E
Additional Experiment Results
E.1
Multivariate Probabilistic Time Series Forecasting
To illustrate Diffusion Forcing's 's new training objective does not degrade it as a generic sequence model, we evaluate Diffusion Forcing high-dimensional and long-horizon sequence prediction
29

---

tasks in time series prediction. We adopt multiple time series datasets with real-world applications from GluonTS [ and evaluate Diffusson Forcing with strong baselines with standard metries in this domain. In this section, we mainly focus on the results and analysis. For detailed description of datasets and the metric, we refer the reader to Appendix F.4
Problem Formulation Let X {xx)/=1 be a sequence (multivariate time series) of D-dimensional observations x E RRD of some underlying dynamical process, sampled in discrete time steps t E (1..... T) where 7 E N. In the problem setting of probabilistic time series forecasting, the sequence X = {X p is split into two subsequences at time step ig E N with 1 < to S T: the (also called history or evidence) of length 40 - 1, and the prediction window X, {x1 ) of length T to + 1 (also known as the prediction horizon). Then, the task is to model the conditional joint probability distribution
context window X == {x fo-1
T (Xx1,:: X1:to-1) = II q(x, x1=-1) 1=
(E.1)
over the samples in the prediction window. If we know the distribution in (E.I). we can sample forecast prediction sequences given some initial context from the evidence sequence. However, most time- -epeeddenddnt data generation processes in nature have complex dynamics and no tractable formulation of (Xx,,:T' X1:to-1). Instead, we construct a statistical model that approximates the generative process in E and estimates quantiles via Monte Carlo sampling of simulated trajectories. In this way, confidence levels or uncertainty measures can be calculated, and point forecasts can be produced as the mean or median trajectory 35
Table 2: Results for time series forecasting. We report the test set CRPS.um (the lower, the better) of comparable methods on six time series datasets. We measure the mean and standard deviation of our method from five runs trained with different seeds.
Method
Exchange
Solar
Electricity
Traffic
Taxi
Wikipodia
VES 35
0.005 + 0.000
0.900 + 0.003
0.880 + 0.004
0,350 +: 0.002
VAR 44
0.005 0.000
0,830 + 0.006
0.039 0.001
0.290 0.001
VAR-Lasso 44
0.012 + 0.000
0,510 + 0.006
0.025 + 0.000
0.150 - 0.002
3.100 1: 0.004
GARCH
0.023 1: 0.000
0,880 0.002
0.190 0.001
0.370 0.001
DeepAR [1
0.336 + 0.014
0.023 - 0.001
0.055 0.003
0,127 + 0.042
LSTM-Copula 33
0.007 = 0.000
0,319 + 0.011
0.064 = 0.008
0.103 + 0.006
0.326 0.007
0,241 0.033
GP-Copula. 53
0.007 + 0.000
0.337 + 0.024
0.025 = 0.002
0,078 + 0.002
0.208 0.183
0,086 +: 0.004
KVAE [40
0.014 = 0.002
0,340 + 0.025
0.051 = 0.019
0.100 - 0.005
0.095 +: 0.012
[4]
0.320 + 0.020
0.002
forner MAF 50
0.005 1 0.003
0,301 1: 0.014
0,021 1 0.000
0.056 0.001
0,179 1 0.002
0.063 1: 0.003
TimeGrad [0
0.006 1 0.001
0,287 1: 0.020
0.021 0.001
0.044 0.006
0.114 0.020
0.049
ScoreGrad sub-VP SDE (1]
0.006 = 0.001
0.256 + 0.015
0.019 = 0.001
0.041 + 0.004
0,101 0.004
0.043 + 0.002
Ours
0.003 - 0.001
0,289 + 0.002
0.023 = 0.001
0.040 + 0.004
0.075 0.002
0.085 0.007
Results. We evaluate the effectiveness of Diffusion Forcing as a sequence model on the canonical task of multivariate time series forecasting by following the experiment setup of [53] 50 49 58 67 Concretely, we benchmark Diffusion Forcing on the datasets Solar, Electricity, Traffic, Taxi, and Wikipedia. These datasets have different dimensionality. domains, and sampling frequencies, and capture seasonal patterns of different lengths. The features of each dataset are detailed in Table We access the datasets from GluonTS (2), and set the context and prediction windows the same length for each dataset. Additionally, we employ the same covariates as [49]. We evaluate the performance of the model quantitatively by estimating the Summed Continuous Ranked Probability Score CRPS..mm via quantiles. As a metric, CRPS measures how well a forecast distribution matches the ground truth distribution. We provide detailed descriptions of the metric in Appendix F.4 We benchmark with other diffusion-basee methods in time series forecastings, such as TimeGrad [49] and the mansformer-basel MrAn [50]. In particular, the main baseline of interest, TimeGrad (49), is a next-token diffusion sequence model trained with teacher forcing. We track the CRPS metric on the validation set and use early stopping when the metric has not improved for 6 consecutive epochs, while all epochs are fixed to 100 batches across datasets. We then measure the CRPS on the test set at the end of the training, which we report Table 2 We use the exact same architecture and hyperparameters for all time series datasets and experiments. Diffusion Forcing outperforms all prior methods except for 67 with which Diffusion Forcing is overall tied, except for the Wikipedia dataset, on which Diffusion Forcing takes fourth place. Note that time series is not the core application
30

---

of Diffusion Foreing, and that we merely seek to demonstrate that the Diffusion Forcing objective is applicable diverse domains with no apparent trade-off in performance over baseline objectives.
E.2 Additional results in compositional generation
Since Diffusion Foreing models the joint distribution of any subset sequence, we can leverage this unique property to achieve compositional behavior i.e., Diffusion Forcing sample from the distribution subsets of the trajectory and compose these sub-trajectories into new trajectories.
In particular, we show that we can also have flexible control over how compositional Diffusion Forcing is. As shown 7 consider a dataset of trajectories on a 2D, square plane, where trajectories start from one corner and end up in the opposite corner, forming a cross shape. When no compositional behavior is desired, one can let the models replicate the cross-shaped distribution by allowing full memory of the HMM model. When one desires compositional such generating a V-shaped trajectory, which stitches two sub-trajectories together, one can let the model generate shorter plans with no-memory context using MPC. (Add figures).
Dataset
(b) W/ memory
W/o memory
Figure 7: Given a dataset of trajectories Diffusion Forcing models the joint distribution of all subsequences of arbitrary length. At sampling time, we can sample from the trajectory distribution by sampling Diffusion Forcing with full horizon (b) or recover Markovian dynamics by disregarding
previous states
E.3
Additional results in video prediction (wo/ cherry picking)
Infinite Rollout without sliding window Diffusion Forcing can rollout longer than maximum train- ing horizonwithout sliding window. That is, we run Diffusion Forcing's ' RNN continuously without ever reinitializing Zy. This is surprising effect observed from the rollout stabilization property of Diffusion Forcing. In Figure 8 [0] we use Diffusion Forcing to generate video sequences of length 180 and visualize subsampled sequences. Notably, Diffusion Forcing used in these visualizations is trained with maximum length of 72 frames for Minecraft and 36 frames for DMLab, illustrating rollout 2x-5x times longer than it's trained without sliding window. In addition, rolling these models out for 2000 frames and without seeing the model blowing up on both datasets. There are occasional cases where the Minecraft agent gets stuck and the entire screen the "dirt" block, but this is more of a dataset issue E.5 and the agent is able to recover after it turns around.
also tried
Consistency We also present additional results where only generate within maximum training length. As shown in figure 12 Diffusion Forcing can generate consistent videos. Results are not cherry-picked.
31

---

Figure Visualization shows Diffusion Foreing trained on 72 frames is able to rollout 180 frames on Minecraft dataset without sliding window. The visualization shows non-cherry-pice subsampling of these 180 frames, although Diffusion Forcing can roll out much longer (such as 2000 frames) on this dataset.
E.4 Additional results in planning
We provide some additional visualizations of causal planning in [5] We also present additional visualization of Diffusion Forcing performing model predictive control in action. As shown in figure 14 Diffusion Forcing generate plans of shorter horizons since it's flexible horizon.
32

---

Figure 9: Diffusion Foreing trained on 72 frames able to rollout 180 frames on Minecraft dataset without sliding window. The visualization shows con-cherry-pickee subsampling of these 180 frames, although Diffusion Forcing can roll out much longer (such as 2000 frames) on this dataset. The first few frames marked in red are the ground truth images of the dataset used for conditioning.
E.5 Real robot experiment setup
In Figure 16 we visualize our robot experiment setup with corruption on observation. The dataset is collected when the target bag isn't present, while with such bag in the zero-shot for the imitation learning experiment with observation corruption. The typical failure mode is when the robot no longer reacts to the visual clues of the randomized location of objects. We didn't observe the robot act wildly due to visual distractors.
F
Additional details about datasets
F.1 Dataset for video diffusion
We adopt the video prediction dataset Minecraft and DMlab used by TECO 68
Minecraft Navigation The Minecraft navigation dataset consists of first-person-vied videos of random walks in the Minecraft 'swamp' biome. The agent walks via a technique called 'sprint jump' which allows jump across blocks without getting stuck block obstacles. The agent walks straight most of the time, with small chances of turning left or right. The height and width of the video 128 pixels and we trim long videos to subsequences of 72 frames. The dataset comes with paired action data but we discard them to bring more stochasticity to the prediction task. Due to limited compute, we only train on about 10% of the total subsequences.
One problem we noticed about the dataset when the agent runs into obstacles with a height of blocks or more. In this case, the agent will get stuck and the entire video sequence will consist of grey granite patterns or brown dirty patterns. This leads to a huge amount of frames with these patterns, making video models predict meaningless frames. Yet, we deem this problem this dataset itself.
DMLab Navigation Deepmind Lab navigation dataset consists of random walks in 3D maze environment. For DMLab, the resolution is 64 pixels and we use subsequences of 48 frames. We also disregard the provided actions due training.
We note that the VQ-VAE latent that stable video diffusion diffuses also only 128 128 indicating Diffusion Forcing has the potential to scale up to higher resolution images with pre-trained
33

---

Figure 10: Visualization shows Diffusion Forcing trained on 36 frames able to rollout frames on DMLab dataset without sliding window. The visualization shows non-cherry -iickdd subsampling of these 180 frames, although Diffusion Forcing can roll out almost infinitely on this dataset. The first few frames marked in red are the ground truth images of the dataset used for conditioning.
image encoder and decoders. Due to the sheer size of the datasets, only use about 10% of the total data sequences for training due to limited computing, as we observe that doing so already allows us to make good generations from initial frames from the test set.
E.2 Dataset for planning
is a standard offline RL benchmark featuring a wide range of reinforcement learning environments. Each environment is associated with a provided dataset of offline interactions with the environment featuring state, action, and reward trajectories.
D4RL 18
we choose the 3 maze environments they are challenging long-horizon, multi- modal, sparse reward problems uniquely suited for visualization and evaluating planning algorithms. The IDs for the 3 used environments are "maze2d-medium-vi "maze2d-large-vl" "maze2d-umar
Like Diffuer (36]
34

---

D
Figure 11: Visualization shows Diffusion Forcing trained on 36 frames able to rollout frames on DMLab dataset without sliding window. The visualization shows non-cherry -icked subsampling of these 180 frames, although Diffusion Forcing can roll out almost infinitely on this dataset. The first few frames marked in red are the ground truth images of the dataset used for conditioning.
v1". each environment, one controls the acceleration of a robot to walk it towards a goal. The observation space is 4 dimensional, fcaturing 2D location and velocity. The action space is 2D acceleration. The agent always receives a random start location and the goal is to reach a fixed goal position for each maze. The agent receives reward of 1 if it is within circle of radius 0.5 centered at the goal state, and otherwise.
The offline RL dataset for the maze environments consists of random walks in the maze. Specifically, the authors first designate intersections and turn in the maze as waypoints and code an agent to navigate between waypoints with some randomization. As result, the random walks are generated way that the path is collision-free with the walls. The random walks introduce stochasticity to the dataset, as trajectories in the dataset are never towards specific goal.
35

---

Additional non-cherry-picked video prediction results on DMLab dataset, generated within maximum training length. The first few frames marked in red are the ground truth images of the dataset used for conditioning.
Figure 12:
There are few choices adopted from our main baseline Diffuser 36 dataset and plan with goals only. We also evaluate a multi-gout variant of each environment (labeled as "multi" Table D. where the goal randomized just like the starting position.
we disregard the reward the
F.3 Dataset for robot learning
We choose long horizon robotic manipulation task as described in Section 4.4 Consider tabletop
with three slots where we can place objects. One places an apple at slot or slot B randomly, and then places an orange at the other slot between A and B. A robot is challenged to swap the position of two fruits using the third slot C. That is, can only move a fruit to an empty slot at time. For example, when the apple is at slot A and the orange is at slot B, it may move the apple to slot C, leaving slot A empty. Then move the orange to slot A and finally move the apple from slot C to slot B. In figure 4 we illustrate the non-mane mor-kvakov property of the task: When the apple is at slot B and the orange is at slot C, one cannot tell what the immediate action is without knowing the initial positions of objects.
We put stickers on the table indicating circular region occupied by any slot. Each eircular region is designed to be about double the diameter of fruit. To make sure the task requires visual feedback, we also randomize the location of fruit inside the slot. We collected 150 expert demonstrations of Franka robot performing the task using VR teleoperation and impedance control. Among them, each initial slot configuration makes up half of the dataset. We record videos from two camera views, one from hand camera and one in the front capturing all three slots. Each demonstration
36

---

Additional non-cherry cheer--icceed video prediction results on the Minecraft dataset, generated within maximum training length. The first few frames marked in red are the ground truth images of the dataset used for conditioning.
Figure 13:
also comes with dof actions of the robot hand. During the data collection, since successful demonstration will swap the position of two objects, its end configuration will naturally serve as the starting configuration of the other randomized location, which we leverage
time.
Each demonstration comprises 500 - 600 frames and actions. We train Diffusion Forcing on entire sequence. However, since adjacent frames are visually close, pad and downsample the vidcos to 40 frames where each frame is bundled with 15 actions.
F.4 Dataset for time series
Table 3: Characteristics of the GluonTS datasets used to benchmark Diffusion Forcing in the domain of time series forecasting.
Dataset
Dimension
Domain
Frequency
Steps
Prediction length
Exchange
R²
BUSINESS DAY
Solar
137
R
HOUR
7,009
24
Electricity
370
R²
HOUR
5,833
24
Traffic
(0.1)
Taxi
1,214
N
30-MIN
1,488
24
Wikipedia
2,000
N
DAY
792
We set of time series datasets accessible via Gluon'TS [2), which adopted from prior works like [71] (5). These datasets capture real-world data of high-dimensional dynamics like monetary exchange rates the electricity grid. In Table these datasets, such as the dimensionality, the domains, the sampling frequency, the length of the multivariate sequence in the training set, and the prediction length. We access the datasets in Table 3 via GluonTS and wrap the data processing functions implemented in GluonTS in our own dataloaders. Each dataset consists of one long multivariate sequence, which is the training split, and set of short sequences that make up the test split. We construct validation set of the same cardinality as the held-out set as a randomly sampled subset of subsequences from the training set. All splits are normalized by the mean and the standard deviation of the features in the training split.
provide summary of the features of
Covariates Often, statistical models that approximate (E.1 benefit from manually curated features as additional input to the observations. A sequence covariates C {ct)r= can be constructed to help the model recognize seasonal patterns and other temporal dependencies. We follow the
implementation in [50] to construct the covariate sequence as a function of the frequency of each
dataset in Table 3 As such, our covariates are composed of lagged inputs, well as learned
embeddings and handcrafted temporal features that encode information such as the hour of the day or the day of the month, depending on the sampling rate of the particular time series that is being
37

---

Figure 14: Example MPC planning for maze medium environment. Blue indicated trajectories actually executed already. Red is the plan.
Figure 15: Example plans generated for maze medium (above) and maze large (below) environments.
modeled. Therefore, covariates are known for the entire interval [1, even inference. We can easily incorporate covariates into the probabilistic framework
g(xtoiT X1:to-1.1.17
II 9(xx X1:te-1, C1:T).
(F.1)
The benefit obtained from covariates highly dependent on the characteristics of both the dataset and the model used, well as the feature engineering practices followed.
38

---

We randomly throw a target bag on the table as strong visual distractor. Diffusion Forcing can be prompted to treat observation as corrupted rather than ground truth.
Figure 16:
Metric The Continuous Ranked Probability Score (CRPS) 45 is a scoring function that measures how well the forecast distribution matches the ground truth distribution:
CRPS(F,z)
(P(=)-114 S 2))² dz.
where F(2) is the univariate cumulative distribution function (CDF) over the predicted value, æ is ground truth observation, and I. {x =} is the indicator function that is one if S and otherwise. By summing the D-dimensional time series along the feature dimension for simulated samples (resulting in F Fsum (t)) and ground truth data (as Σ, xz,e).
we can report the CRPS
CRPS EE-Le(tn.I))
CRPS
(aam (1),
Σ J.e
the average the prediction window. The lower the CRPS value, the better the predicted distribution match the data distribution.
First, manually the time series along the feature dimension and estimate the CDF Fsum (t) via 19 quantile levels each time step & from 100 sampled trajectories. We then use the implementation in Gluon's 22] to compute the CRPS, which report as CRPS in Table While aggregate the data manually, we verify that the numerical error relative to the GluonTS implementation remains orders of magnitude below the precision threshold of the reported metric.
39

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction supported by experimental results and are contextualized with respect to competing methods.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in the final section of the paper (
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "LLimitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
40

---

Answer: [Yes]
Justification: Derivations of relevant expressions provided in the supplement. Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide reproducibility details for each experiment in the supplement.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
41

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Code has been released publicly.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/pulee wttpcccpppppccceeeeeee.ee..c.lct http://uuiiieeeeeeebbabieieeee hat////oooSemssseeeeeee oossiiccyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Experiment details are provided the paper and supplement. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: We report all quantitative results in terms of the mean and standard deviation over several runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
42

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All experiments were performed on the same device, the details of which are described in the supplement.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: To the best of our knowledge, we conform to the Code of Ethics. At this time we do not see our method providing a straightforward avenue for abuse by bad actors.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Due the low-level nature of method, do not see it directly facilitating any negative social impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
43

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All data used is cited in accordance with the provided licenses.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, thcoe com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
44

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
45


Reviews:
[{"review_id": "87XcPcyG1o", "reviewer": "Reviewer_a6Fo", "summary": "This paper introduces Diffusion Forcing, a novel training paradigm for sequential generative modeling using diffusion models. Diffusion Forcing learns from sequential tokens with varying independent noise levels, enabling more flexible sampling strategies and general capabilities such as guidance. The experimental results demonstrate that Diffusion Forcing outperforms existing methods, including full sequence diffusion and teacher forcing, across various tasks.", "questions": "1. **Consistency between training and sampling algorithms.**  \n   In Algorithms 1 and 2, there appear to be inconsistencies between the training and sampling algorithms. Can the authors provide an intuitive explanation for these inconsistencies? Specifically:\n   - During training, the predicted noise $\\hat{\\epsilon}\\_t$ is calculated using the latent from the previous step $z\\_{t-1}$, whereas during sampling, $\\hat{\\epsilon}_t$ is calculated using the latent from the current step $z_t^{\\text{new}}$.\n   - Similarly, during training, $\\hat{\\epsilon}\\_t = \\epsilon\\_\\theta(z, x\\_t\\^{k\\_t}, k\\_t)$ uses the same noise level $k\\_t$ of the noisy observation $x\\_t\\^{k\\_t}$, but during sampling, $x\\_t$ has a noise level $\\mathcal{K}\\_{m+1,t}$ instead of $k = \\mathcal{K}_{m,t}$.\n\n2. **Stabilizing auto-regressive generation.**  \n   The authors propose conditioning on the (latent of) slightly noisy previous tokens with a noise level $0 < k \\ll K$ to stabilize the auto-regressive generation. How were the values of $k$ chosen in the experiments? Could the authors provide ablation studies on the impact of using this trick?\n\n**I promise to raise the score once all the weaknesses/questions are solved.**", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed Diffusion Forcing method is general and flexible, making it applicable to various tasks.\n2. The paper provides a comprehensive discussion on the capabilities of Diffusion Forcing.\n3. The experiments are well-designed and effectively demonstrate the proposed method's effectiveness.", "weaknesses": "1. **Writing clarity and organization.**  \n   The writing style impacts readability, making the paper challenging to follow. It would benefit from a clearer organization. The paper primarily covers three points: (a) the proposed Diffusion Forcing (DF) method with independent noise levels and its theoretical analysis, (b) the capabilities of DF, including flexible sampling strategies, and (c) experimental results on various tasks. However, the current structure does not clearly present these points, particularly the DF method. Separating the design of DF and the intuitive explanation from the Bayesian filtering perspective, and listing the resulting capabilities in a separate section, would enhance clarity.\n\n2. **Clarity of figures.**  \n   The figures are not well-explained and are difficult to understand without referring to the text. For instance, Figure 1 omits latent states in the sampling process for both Diffusion Forcing and Teacher Forcing, which is confusing.\n\n3. **Minor issues and typos.**  \n   - Line 97: missing a \")\"\n   - Line 139: \"nevel\" should be \"level\"\n   - Line 186: \"$x^{K/2_3}$\" should be \"$x^{K/2}_3$\"\n   - Line 178, 184, etc.: paragraph titles are inconsistently formatted\n   - Line 522: missing a \"(\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful feedback! We respond to your comments below:\n\n> Writing clarity and style & typos\n\nSorry about the confusion. It’s true that we opted out of an independent “Method” section to introduce the intuitions first. With the extra 1 page for the camera-ready version, we promise to structure the paragraphs more clearly. In particular, following your advice, we will make (1) one “methods” section that discusses DF in one subsection as well as its intuition obtained from Bayesian Filtering in another subsection, (2) one “Capabilities” section discussing novel capabilities, and (3) one ” Results” section.  Thank you for your time on the list of typos, we will make sure to fix them in camera-ready.\n\n> Figure clarity\n\nThank you for the constructive advice. We will make sure to caption the figures with more details with the 1 extra page limit of the camera-ready version so readers can easily understand them without referring to the main text. For the sampling scheme in Figure 2 (Assuming the figure index in your comment is a typo), we were trying to illustrate a scheme that illustrates both the transformer and RNN as mentioned in the caption. Therefore, we will likely instead remove latent from the training subfigure to make sure it’s consistent with the sampling figure and add a new figure in the appendix to illustrate RNN variant with latent state. We provide a preview of modification in **Figure 5 of the rebuttal pdf**.\n\n> Consistency between training and sampling algorithms.\n\nThank you for pointing this out. We will definitely fix the algorithm box carefully: 1. $z_{t-1}$ in line 8 of the training algorithm should be changed to $z_{t}$ 2. In line 7 of the sampling algorithm, the $k$ should indeed be corrected to be $\\mathcal{K}_{m+1,t}$. We will fix line 5 to reflect this. 3. Outside the two errors you pointed out, some $\\alpha$ are accidentally indexed via $t$ or $k$, not $k_t$. \n\n> Details about stabilization and ablating $k$\n\nIn our implementation, we choose $k=20$ for $K=1000$ but this value is largely dependent on the difficulty of the task. In **Figure 3(a)** of the rebuttal pdf, we present the result of the requested ablation of this stabilization value. We choose a slightly harder setting for this ablation, reducing the frame rate by half and using lower numbers of DDIM steps to make compounding error more obvious. One can observe that the FVD metric improves as we gradually increase the stabilization level from 0, reaching optimal around 100 before FVD starts to rise. We found stabilization to be extremely important for datasets with high stochastic like the BAIR robot pushing dataset shown in **Figure 3(b) of the rebuttal pdf**. We cannot achieve stable autoregressive rollout even with very high DDIM steps. Stabilization is a fundamental technique to get this dataset to work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qF6W9kNF0j", "reviewer": "Reviewer_tz82", "summary": "This paper proposes to augment autoregressive models with diffusion. Specifically, rather than generating every token in one shot (one neural network evaluation), the paper proposes to gradually denoise the tokens following an autoregressive order. That is, every token is given a different noise level (lower for former tokens and higher for latter ones), and the tokens are jointly denoised to generate better samples. Compared to pure autoregressive prediction, diffusion forcing allows the model to refine the samples through the diffusion process. Compared to diffusion models, the proposed model is capable of variable-length generation and extrapolation.\n\nThe authors also demonstrate additional potential generation tasks that can be done by diffusion-forcing models such as guided autoregressive sampling. \n\nEmpirical results demonstrate that diffusion forcing performs well on video prediction and various planning tasks.", "questions": "How easy or difficult can diffusion forcing be applied to non-autoregressive generation? Although diffusion forcing improves the performance in the autoregressive generation regime, some tasks (e.g., constrained text generation) require awareness of future tokens to generate the current ones. I wonder if diffusion forcing can be extended to this regime.", "limitations": "As discussed by the authors, one of the main limitations is that diffusion forcing is only tested with RNN base models but not other autoregressive models such as autoregressive transformers.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper proposes an interesting combination of autoregressive models and diffusion models and demonstrates that the combination of both outperforms both individual models in terms of performance. Further, the diffusion-forcing paradigm offers many more applications that are otherwise impossible. For example, while doing variable-length generation, the model can leverage classifier-based/-free conditions. This provides much better flexibility to inference-demanding tasks such as planning and control.\n\nThe authors propose a training objective of diffusion forcing models based on noise prediction. The objective is proved to be a reweighted version of the evidence lower bound and thus is sound.\n\nDiffusion forcing achieves much better performance compared to autoregressive models and diffusion models in long-horizon generation tasks.", "weaknesses": "A more detailed discussion of the noise schedule is desired to better understand the effectiveness of diffusion forcing. Is it necessary to use different noise schedules in different tasks to achieve good performance? Further, can we train the model with various/arbitrary noise schedules and at evaluation time find a good schedule? If any of these is possible it will greatly reduce the training complexity and extend diffusion forcing to more applications.\n\nTheorem 3.1 states that the proposed objective is equivalent to a reweighting of the evidence lower bound. However, it is unclear how the noise schedule biases the reweighting since a very badly balanced ELBO can render the training process unstable.\n\nHow diffusion forcing balances efficiency and performance. In the extreme case where only one denoising step per token is allowed, diffusion forcing reduces to autoregressive generation. How much performance gain can we expect if we allow for more computation time?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful feedback. We respond to your comments below:\n\n> Can we train with arbitrary noise schedules and find a good schedule at evaluation time? \n\nYes - at the core of Diffusion Forcing lies exactly the idea that by training with **arbitrary, random** noise schedules, we can experiment with arbitrary noise schedules at evaluation time to find the optimal schedule for each task! In practice, we found this to be very useful as users can try different noise schemes at sampling time for different applications after training once while the slightest tuning would require retraining for various baselines. The proof of the ELBO extends to arbitrary noise schedules as well. \n\n> Is it necessary to use different noise schedules in different tasks?\n\nIn short, yes. For example, we found that pyramid sampling with causal diffusion forcing (CDF) is best for MCTG planning while autoregressive diffusion offers the best consistency for causal video. When conditioning on ground truth context frames, CDF works best without stabilization; In contrast, when context is previously generated, CDF works best with stabilization as shown in **Figure 3 of rebuttal pdf**. Furthermore, even if we commit to one specific schedule such as pyramid sampling, it is critical to tune the rate with which uncertainty grows with time via the noise schedule under pyramid sampling. DF allows us to do this easily after training once, while AR-Diffusion [2] (suggested by 35x2 and 2PA2) makes this difficult as this rate is a training hyperparameter and fixed at test time. \n\nWe will use the extra page from the camera ready to include this discussion of noise schedules!\n\n> Reduce the training complexity\n\nSurprisingly, we found that random, independent noise schedules at training time do not significantly add to the training complexity. Due to limited space, please refer to the experiment in general response’s “Improved Model Performance” section which shows the overhead of independent noise is insignificant when considering the overall training compute. \n\nWe speculate that random, independent noise levels at training time can be seen as a form of data augmentation[3], leading to a bigger “bang for the buck” of training data. Further, they might in expectation improve the quality of gradient flow backward in time. \n\n> Balance between performance and computation time\n\nThe suggested sampling scheme reduces to auto-regressive diffusion with ddim_steps of 1, which mainly captures the low-frequency part of the data only. If we have more compute budget, we can sample auto-regressively with >1 ddim_steps and get much higher quality samples. As shown in **Figure 2 of rebuttal pdf**, the main video metric fvd decreases as the sampling time increases. A similar trend follows for maze planning. \n\n> reweighted ELBO can be poorly balanced\n\n\nIndeed, the formal version of Theorem 3.1 in the Appendix, Theorem A.1 demonstrates the weighting explicitly: that the ELBO on the expected log-likehoods, over the randomness of the noise schedule, corresponds to an objective which reweights the contribution of terms from denoising steps k by their step number k. Thus, the reweighting factor is by at most a factor of the total number of denoising steps, K. Note that this weighting by k is independent of the noise schedule (that is, of the variance of the noise at each step k), and depends only on the position k of the noising step (as well as our decision to select noise levels uniformly and independently). We will add this as a remark in the revision. \n\nThough we could choose to up weight gradients from earlier steps by more (as Thm A.1 makes the weights explicit), not reweighting the gradients according to the ELBO corresponds to effectively weighting gradients at smaller denoising steps k more, which may be desirable as these steps capture the “higher resolution” features of the model. \nLastly, as noted in the remarks following Corollary A.2, our ELBO derivation is sufficiently tight that, given a sufficiently representative neural network, the optima of all (non-trivial) reweightings of our ELBO are also optima of the desired expected likelihood. Hence, the weighting does not meaningfully bias the training objective. We do acknowledge that the weighting may be salient for the dynamics of optimization, but empirically, we find that our decision not to compensate for the step-wise weighting of the ELBO leads to reliable training. We will clarify the remarks following Corollary A.2 to emphasize this point further. \n\n> Main result is RNN\n\nWe’ve reimplemented diffusion forcing with both transformer and 3D-Unet architectures after the draft submission. Please see the general response for details due to limited space.\n\n> Apply diffusion forcing to non-autoregressive generation\n\nIn the maze planning experiments, we already show how guidance can be used to achieve certain future outcomes. In fact, reconstruction guidance discussed in Appendix B.4 has been widely used to achieve consistency with context through guidance.  Alternatively, one simply uses a non-causal neural network architecture, like our implementation of 3D-Unet, and then - as discussed in the paper - varies the per-token noise level at training time. One can further mask out padded entries by marking them as full noise. We are actively pursuing this direction in follow-up work: Initial results in **Figure 1 of rebuttal pdf** showcase frame interpolation result of Diffusion Forcing without the expensive reconstruction guidance. However, a detailed discussion of this setting would exceed the capacity of the present paper, though we will discuss this direction in the discussion section.\n\n[1] Michael Janner, Planning with Diffusion for Flexible Behavior Synthesis, 2022\n\n[2] Tong Wu et al. “Ar-diffusion: Auto-regressive diffusion model for text generation”, Neurips 2023\n\n[3] Diederik P. Kingma, “Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation”, Neurips 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "khcIDlwvi4", "reviewer": "Reviewer_35x2", "summary": "The authors introduce Diffusion Forcing (DF), a method for diffusion of sequential data where the noise level at each token can be different (“independent”). The authors show that DF provides more flexible steerability properties and more stable rollouts compared to full-sequence diffusion and teacher forcing. Experimentally, these enable stable video prediction along several timesteps, improved performance on planning tasks, and more robust robotic visuomotor control, relative to the respective relevant baselines.\n\nThe following is a more detailed summary.\n\n### Overview of the method\n\nDiffusion Forcing (DF) denoises a sequence of noisy tokens $x^k\\_1, \\cdots, x^k\\_T$ at noise level $k$, starting at $k=K$ (maximum noise) and finishing at $k = 0$. A sequence of hidden states $z\\_1, \\cdots, z\\_T$ is also maintained throughout. Importantly, different tokens can be denoised by different amounts at each denoising step.\n\nThe architecture has two main components: \nan encoder $p\\_\\theta(z\\_t | z\\_{t-1}, x^{k\\_t}\\_t, k\\_t)$ mapping the previous hidden state $z\\_{t-1}$, the current noisy token $x^{k\\_t}\\_t$ and the noise level $k$ to the new value of the current hidden state $z\\_t$;\na denoiser $\\epsilon\\_\\theta(z\\_t, x^{k\\_t}\\_t, k\\_t)$, which is used to denoise $x^{k\\_t}\\_t$.\n\nAt training time, the noise levels $(k\\_t)\\_{1 \\leq t \\leq T}$ are sampled independently, and the encoder and denoiser are trained jointly using the usual diffusion loss on the output of $\\epsilon\\_\\theta$.\n\nAt inference time, the tokens are initialized with independent Gaussians. They are then denoised by first computing hidden states from left to right (via an RNN, in this case) using $p\\_\\theta$, and then by updating the values of the tokens using their current values and the hidden states.\n\nThe authors provide an ELBO interpretation for their loss function in the appendix.\n\n### Features of Diffusion Forcing \n\n- The authors highlight the following features of DF:\n- It supports classifier guidance, like ordinary diffusion;\n- It allows for keeping the noise level higher for future tokens. This makes intuitive sense in an auto-regressive setting, where future tokens depend on past tokens.\n- It supports a flexible planning horizon, as tokens are denoised sequentially.\n- It supports a more flexible form of classifier guidance (or reward guidance): past tokens can be guided by rewards that depend on future tokens, due to DF’s autoregressive architecture.\n\nWhen doing reward guidance, the authors propose drawing many samples of possible future trajectories, and averaging their rewards, rather than using a single sample as in ordinary classifier guidance. They term this approach Monte Carlo Tree Guidance (MCTG).\n\n### Overview of experimental findings\n\n- The authors evaluate Diffusion Forcing on video prediction, planning and robotics tasks. Their findings can be summarized as:\n- In video prediction (datasets: Minecraft gameplay and DMLab), DF provides more stable rollouts than full-sequence diffusion and teacher forcing. In particular, DF’s rollouts do not diverge as the number of tokens increases.\n- In planning (environment: Maze2d from D4RL), DF produces more more consistent trajectories, and executing the generated actions indeed produces a trajectory similar to that given by the generated states. - In addition, DF with MCTG significantly outperforms Diffuser on Maze2d environments.\n- In robotics, DF is robust to missing or noisy observations and can perform imitation learning with memory (as it maintains a hidden state, rather than directly mapping observations to actions).\nIn the appendix, the authors provide additional experiments on compositionality and time series prediction.", "questions": "1. Does the flexible planning horizon (line 209) of Diffusion Forcing not derive from the choice of an RNN as the architecture, rather than e.g. a UNet? Would implementing existing methods such as Diffuser (Janner et al. 2022) not allow for a similar property?\n\n1. In the paragraph “Benefit of Modeling Causality”, the authors highlight that states and actions produced by DF are causally related, which does not hold in practice for Diffuser. Do the authors claim this is due to DF explicitly incorporating temporal structure into its architecture? Could it not also be due to the use of an observation model $p\\_\\theta(x^0\\_t|z\\_t)$ to predict the noise-free token $x^0\\_t$ from a hidden state $z\\_t$?\n\n1. At first sight and in its current form, the method seems tailored to the use of an RNN architecture, rather than a Transformer. For example, the denoiser is applied token-wise, with the information from previous tokens affecting the current token only via the hidden states $z\\_t$. How would the method have to be adapted, if at all, to work with transformers, in case one wants to scale up Diffusion Forcing?\n\n1. Janner et al. (2022) showcase in Section 5.4 show how to apply Diffuser with a variable planning budget, and study how the resulting performance varies with the planning budget. Can Diffusion Forcing also be run with a variable planning budget, through warm-starting (as for Diffuser) or otherwise? If so, it would strengthen the paper if the authors described how, and included a similar budget vs. performance analysis, especially in planning and robotics tasks.", "limitations": "The authors address relevant limitations in Section 5 and social impacts in the checklist.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The authors propose an original and performant method combining strengths of diffusion (steerability, robustness to noise, high-quality gradual sample generation) and auto-regressive sequence modelling (flexible horizons, temporal causality, memory in the case of RNNs).\n\n1. In addition, the authors provide a theoretical justification of their loss function in terms of an evidence lower bound (ELBO).\n\n1. The paper is written clearly, providing a clear motivation for the authors’ approach, contextualizing DF relative to existing work (especially Diffuser, AR-Diffusion and Diffusion Policy), and highlighting the main contributions of the method conceptually and experimentally.\n\n1. Trajectory inconsistency is a major limitation of Diffuser, which I have contended with in my own research. Mitigating this limitation is an important enabler of bringing the strengths of diffusion to bear in sequential decision making.\n\n1. Monte Carlo Tree Guidance can be seen as maximizing an empirical estimate of the expected future reward. From a policy optimization perspective, this seems more principled than doing gradient ascent on the realized cumulative reward of a given trajectory, as is done in full-sequence diffusion (e.g. Diffuser). As the authors explain in Appendix B.3, this technique relies on the architecture of DF to be effective.\n\n1. The results on video prediction, available in an anonymized project website provided in the abstract, are particularly impressive in terms of stability and 3D consistency. This, together with results on planning and robotics, indicates DF might contribute to advances in diffusion world models; a research area of established relevance that has received significant attention recently.", "weaknesses": "1. Clarification on classifier guidance term $\\nabla\\_x \\log c (x^{\\textrm{new}}\\_{1:H})$: If this term is to be understood as the gradient of $x \\mapsto \\log c(x)$ evaluated at $x^{\\textrm{new}}\\_{1:H}$, then the gradients of $c$ on future tokens would not flow to previous tokens, as the inputs $x$ are “frozen” before being fed into $\\log c$. It seems that what the authors mean to say is that future tokens are treated as a differentiable function of  past tokens when computing the gradients. It would strengthen the exposition if the authors either clarify this point in the paper, or update the notation to avoid confusion, as the current notation might lead the reader to believe that the gradients from future tokens do not flow into past ones.\n\n1. The naming of Monte Carlo Tree Guidance seems to misleadingly suggest a similarity with Monte Carlo Tree Search (MCTS). However, the method consists of sampling several future trajectories independently and averaging their guidance gradients, which seems quite divorced from MCTS, which involves actual search on a tree of states and actions and backpropagation of rewards through this tree. As such, I believe naming the technique Monte Carlo Guidance would be more appropriate.\n\n1. High-dimensional control evaluation: Janner et al. (2022) evaluate Diffuser on high-dimensional control locomotion tasks from D4RL. It would be interesting to see an evaluation of Diffusion Forcing in this setting, in particular regarding the consistency between states and actions. I recall from my own experience that executing longer plans from Diffuser in these locomotion environments in an open-loop fashion (i.e. no re-planning) led to trajectories diverging from the generated states, as noted by the authors. It would be interesting to see whether this is addressed by Diffusion Forcing on these higher-dimensional environments.\n\n1. The compositional generation environment referenced in Section 4.3 is very similar (if not identical) to the one used by Janner et al. (2022) in Figure 1b of their paper. I believe it is likely worth mentioning this in Section 4.3.\n\n1. Minor formatting problems\n    \n    1. Line 186: $x^{K/2\\_3}$ -> $x^{K/2}\\_3$\n        \n    1. Table 1 caption: “Diffusion Forcingkeeps” -> “Diffusion Forcing keeps”; “Diffusion Forcingachieves” -> “Diffusion Forcing achieves”\n        \n    1. Line 495: “in full abstraction” -> “in full generality”\n    Line 503: “likelihood for likelihood of all” -> “likelihood for all”\n        \n    1. Equation A.3: superscript $k\\_2$ on the LHS should be $k\\_s$\n    Line 516: revise the bracketing of the expression involving $p\\_\\theta$.\n        \n    1. Line 522: “under uniform levels” -> “under uniformly sampled levels”\n        \n    1. Line 524: “in the sequel” -> “in the following section”\n    Equation A.5: $s \\leq T$ -> $1 \\leq s \\leq T$\n        \n    1. Line 592: specify range for $s$ on the first expectation\n        \n    1. Line 598: correct superscripts $t\\_k$ to $k\\_t$\n        \n    1. Line 608: revise bracketing of the numerator inside the $\\ln$\n        \n    1. Line 616: In the last and penultimate lines, replace $\\frac{\\ln p(...)}{q(...)}$ by $\\ln \\frac{p(...)}{q(...)}$\n        \n    1. Line 628: “we” -> “we have”\n        \n    1. Line 631: correct superscript of $x\\_t$ on the second line\n        \n    1. Line 634: expression with $p\\_\\theta$ broken between lines\n        \n    1. Line 635: capitalize Dirac\n        \n    1. Equation B.1: include \\left[ \\right] in the brackets\n        \n    1. Line 664: “we are” -> “we use”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their in-depth review - we are particularly happy to be able to address some of the limitations of Diffuser that the reviewer had to contend with themselves in the past!\n\n> Clarification on Classifier Guidance Term\n\nSorry about the confusion. The reviewer is exactly correct in that future tokens are a differentiable function of past tokens via the posterior model, and that this enables us to perform guidance. We adopted the notation for simplicity and to avoid clutter/accommodate space restrictions, but we will absolutely clarify the meaning of that gradient in the text of the “Guidance” section in our subsequent revision. On the other hand, in our transformer implementation of diffusion forcing, any token can attend to the entire history, allowing us to reproduce the guided planning results as well.\n\n> The naming of Monte Carlo Tree Guidance is misleading\n\nThank you for your comment. Amusingly, we had an internal discussion among the authors on exactly this topic - we report that one of the parties feels highly vindicated by your comment. We will follow your advice and rename the method into Monte Carlo Guidance!\n\n> Consistency of action and state in high-dim control environments\n\nWe tried CDF on one high-dimensional control environment (hopper-medium-expert) without guidance to exclude confounding factors. Unfortunately, the observation still diverges from the ground truth eventually without feedback as no model is perfect even in deterministic environments. However, it’s clearly better than Diffuser just like our visualizations in the 2d maze experiment. In the generic RL setting, the environment can be stochastic, and therefore open-loop execution of a plan is expected to diverge from the plan even with an oracle model. Therefore, we believe a thorough evaluation of dynamics better falls in the result in Section 4.5, which shows CDF is competitive for high-dimensional time-series data.\n\n> Similarity of compositional generation environment with Janner et al. (2022). \n\nThank you for highlighting this relationship - we will absolutely add a reference to Janner et al. (2022) to Section 4.3!\n\n> Minor formatting problems\n\nWe apologize for these mistakes and have fixed all of them for the camera-ready!\n\n> Does the flexible horizon property derive from RNN\n\nFlexible horizon planning doesn’t depend on RNN. We’ve reimplemented diffusion forcing with both transformer and 3D-Unet (with temporal attention), showing flexible horizon planning is fully reproducible with attention instead of RNN. For causal attention, this is obvious. For non-causal variants, diffusion forcing allows us to mark padded future entries as pure noise to achieve a flexible horizon. The only ability that depends on RNN is infinite rolling **without** sliding window because state space models have more invariance, although transformer can roll out longer **with** sliding window. \n\n> Could combining Diffuser with RNNs also lead flexible planning horizon?\n\nDiffuser’s maze planning result is strictly fixed horizon. While Diffuser made an argument about flexible horizon via 1D convolution, its maze planning results critically depend on the replacement technique we mentioned in Appendix B.4. That is, Diffuser fixed the last token to be at the goal throughout the diffusion process to generate a fixed length plan of a carefully chosen length. If the objective is to reach the goal asap, this objective would fail, as one doesn’t know which token shall be replaced as the goal. We implement this on Diffuser and cannot reproduce their numbers. We suspect that the mere reason why Diffuser got high scores is that their fixed length objective coincidentally encourages very slow plans that made the agent stay near the goal for longer, an out-of-distribution behavior that doesn’t exist in the dataset. \n\n> Source of causality in DF vs. Diffuser: architecture or observation model\n\nIt’s most likely the architecture itself. Our transformer implementation is direct empirical evidence that a separate observation model is the key. It also follows theoretically: consider an alternative formulation of $x_t=[o_t, a_t, r_t]$ like diffuser instead of $x_t=[o_{t+1}, a_t, r_t]$ used in our paper. Even though the hidden state and observation model are untouched, causality is mathematically broken - at time step t, one wants to diffuse action $a_t$ given current $o_t$, but the process of diffusing $a_t$ samples another $o_t$, which might be different from the existing observation. Such weird inconsistency suggests the temporal structure is critical.\n\n> scale up Diffusion Forcing for transformers \n\nBetween the initial draft and rebuttal, we’ve already reimplemented diffusion forcing with architectures like Transformers or 3D-Unet and reproduced results. We also find Diffusion Forcing to work well in latent diffusion settings for higher-resolution videos (**Figure 1 of rebuttal pdf**). \n\n\n> Can Diffusion Forcing, like the Diffuser, be run with a variable planning budget, through warm-starting or otherwise?\n\nUnlike Diffuser[1], diffusion forcing implements DDIM sampling which already achieves good speed without Diffuser’s warm-starting technique. Therefore, we provide an alternative planning budget analysis here in terms of frequency of replanning. In **Figure 2 of rebuttal pdf**, we show the performance when we replan 1x, 0.5x, 0.25x, 0.125x of episode length and for 50 steps. We found that diffusion forcing’s performance worsens as we replan less frequently. We’d like to clarify that this is partially due to the dataset itself - the maze dataset never contains the behavior of staying at the goal once reached, and the agent always walks away from the goal even after reaching it. Therefore, diffusion forcing is supposed to follow this suboptimal behavior if there is less replanning. \n\n[1] Michael Janner, Planning with Diffusion for Flexible Behavior Synthesis, 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "USXEJoBrUA", "reviewer": "Reviewer_2PA2", "summary": "This work presents Diffusion Forcing, a new framework for probabilistic sequence modeling that combines diffusion models with Bayesian filtering. This framework builds on state of the art approaches to sequence modeling using diffusion models, but has several novel contributions.\nFirst, it allows the model to use *independent* noise levels per element in the sequence, which is a key factor for the stability of autoregressive generation and conditional, guided generation.\nSecond, this work casts the proposed method for sequential decision making, by defining a new guidance technique that allows the generation of the next toke by guidance on the full distribution of future tokens.\nThird, the authors go at great length in demonstrating empirically that their proposed framework is general and can be applied beyond text generation, as opposed to related work.\nDiffusion forcing relies on simple ideas: noising is understood as (partial) masking, and it is cast for sequential data, giving rise to a causal variant of diffusion forcing. In practical terms, we have a dynamical system modeled with a simple RNN, in which hidden states follow the Markovian principle: the next hidden state depends on the previous hidden state and a current observation. Previous, next and current are to be intended as indexes in the sequence. Observations are obtained by running a diffusion model with independent noise levels per sequence index, and noisy observations can be used to transition to the next hidden state. A connection with Bayesian filtering is made clear in the paper. Then, we end up with an observation model (for state transitions) and a denoising model (for the diffusion of the observations).\nThe authors provide a sound development of the training procedure and objective, by showing that their training algorithm optimizes a weighted ELBO on the expected log-likelihood.", "questions": "Q.1: could you please provide a clear summary of why *independent* noise levels are key for your method, and substantiate the difference with respect to AR-Diffusion [62]? I have read Appendix B.2 and Appendix C, where you attempt at clarifying, but I think the benefits for stability, and conditioning on corrupted observations is not spelled out sufficiently\n\nQ.2: is there a way to compare your work to AR-Diffusion that would not require a substantial re-factoring of their code, such that it can be applied to one (e.g. video) use case in your experiments? Another way to go would be to modify your CDF method and use linearly dependent noise levels, to ablate on the requirement for independent noise levels\n\nMinor (no need to answer):\n* typos: I could spot one typo in line 186: $[x_1^0, x_2^0, x_3^{K/2}]$\n* please check the proofs in the appendix as there are some typos that slipped there, as well as the text in the appendix that has several grammar problems, missing verbs and the like\n\n\n====== POST REBUTTAL MESSAGE ======\n\nThank you for the rebuttal. I have raised my score.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This work presents substantial improvement over the literature on the joint application of diffusion models and autoregressive models\n* The proposed methodology is technically sound, and well supported by intuition, formal proofs and a wide range of experiments\n* The experimental section expands over the literature by focusing on several domains including long-range video generation, planning, compositional generation and multivariate time series forecasting", "weaknesses": "* The intuition of the effects of noising on long-horizon generation (appendix B.2) is very similar to the ideas described in a related work AR-Diffusion [62]. This does not highlight the contribution of *independent* noise levels per sequence index\n* Experiments do not compare (at least to the best of my understanding) Causal Diffusion Forcing to AR-Diffusion, which would be the natural competitor. Nevertheless, I understand that this would require considerable amount of adaptation work, since AR-Diffusion tackles language modeling mainly\n* I liked Appendix B.6, but it is not referenced in the main, and I think this would be more helpful than the figures in sec 3.1", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Independent noise vs AR-Diffusion\n\nWe thank the reviewer for highlighting the need for a more explicit discussion relative to AR Diffusion. \n\nWe would first like to clarify that the stabilization discussion in Appendix B.2 is orthogonal to AR-diffusion. The key to AR-Diffusion is training and sampling **future** tokens with linearly growing noise similar to the pyramid sampling of Diffusion Forcing. On the other hand, stabilization is a technique of assigning non-zero noise levels to **context** tokens (history) instead of future tokens. In fact, our stabilization experiment uses autoregressive diffusion, not pyramid noise like AR-Diffusion. \n\nFurthermore, AR-diffusion style pyramid sampling is just one of the possible sampling schemes supported by diffusion forcing **after training once**. In online decision-making like our robot experiment, we can use autoregressive diffusion to achieve the best speed, use pyramid sampling to achieve better guidance, or many other schemes. In **Figure 1 of rebuttal pdf**, we further show that non-causal diffusion forcing can do frame interpolation by assigning noise level 0 to keyframes at arbitrary time steps, while AR-diffusion’s non-causal variant, Rolling Diffusion lacks this flexibility. \n\nFinally, we found that Diffusion Forcing beats AR-Diffusion even using AR-Diffusion’s own sampling scheme as we show and analyze in our response to the next question. This shows that independent noise, just like standard techniques such as image pre-training, plays an important role in enhancing the underlying representation. \n\n> Experimental comparison to AR-Diffusion\n\nWe are happy to report that we have re-implemented and benchmarked with AR-Diffusion for video generation, as well as a concurrent work named Rolling Diffusion[1] which is basically non-causal AR-diffusion. As shown in **Figure 4 of the rebuttal pdf**, both AR-diffusion and Rolling-diffusion performed worse in FVD metric compared to diffusion forcing. Our main insight is twofold: 1. while both use linearly growing noise levels along the time axis, both seem to be sensitive with respect to the slope of this linear growth. We spent a fair amount of time tuning the two baselines, but since one has to use the same slope during both training and testing, it requires us to re-train many times to find one slope that’s reasonable for our data. 2. We observe that training with linearly growing noise has too much redundancy. The trajectories tend to have a higher signal-to-noise ratio which made the task too easy. We tried the boundary condition in Rolling Diffusion[1] and multiple tunings to mitigate this, but still failed to match Diffusion Forcing. In fact, table 2 of Rolling Diffusion [1] shows that linearly growing noise isn’t necessarily better than normal diffusion in a large video dataset, aligning with our observation. In addition, a fixed slope at training time creates more inconvenience for DDIM sampling as well due to rounding, creating practical pain for any user who wants to tune DDIM steps. \n\n\n> Reference Appendix B.6 in the main paper\n\nThank you for your suggestion - we will definitely reference Appendix B.6 in the main paper, and attempt to move it to the main paper entirely, space after incorporating all reviewer feedback permitting!\n\n> Typo\n\nThank you for catching the typo, we have fixed it!\n\n> Grammar and spelling in appendix\n\nWe have done a full pass over the appendix and polished grammar and spelling, thanks for pointing this out!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## General Response
We thank the reviewers for their comments and suggestions. We are pleased that the reviewers find our paper original & interesting (Reviewers 35x2,tz82), general & flexible (Reviewers a6Fo, tz82), that it has great performance (Reviewers a6Fo, tz82,35x2) with substantial improvement over prior methods (Reviewers 2PA2) backed by well-designed experiments (Reviewers a6Fo,2PA2) and theoretical justification (Reviewers tz82,35x2,2PA2).

The outstanding concerns center around comparison to prior methods, scaling up beyond RNNs (Reviewers tz82,35x2), the need for independent noise (Reviewers tz82, 2PA2), and more ablations to help understanding (Reviewers a6Fo, tz82,35x2,2PA2). We present further discussion, supported by additional experimental results, to address these concerns.

> Comparison to new baselines and variants

We reimplemented AR-diffusion[1] and its non-causal variant, Rolling-diffusion[2] with multiple rounds of tuning. As we detailed in Appendix C, AR-Diffusion uses linearly growing noise during training and sampling. Their sampling scheme closely resembles Diffusion Forcing’s pyramid sampling so we also present a result of diffusion forcing but using AR-diffusion’s own pyramid sampling scheme. **In Table 4(a) of the rebuttal pdf**, we find that Diffusion Forcing has a clear advantage over all baselines, causal or non-causal. Furthermore, the fact that Diffusion Forcing can beat new baselines with their own sampling schemes further highlights the importance of independent noise, as we detail below. 

>  The need for independent noise

We’ve included additional investigative experiments that demonstrate the importance of independent noise, both for **improved model performance** as well as **numerous model capabilities**. We summarize them here and provide detailed discussions in the individual responses. 

### Improved Model Performance
- In **Table 4(a) of rebuttal pdf** We trained with Diffusion Forcing but using AR- and Rolling-diffusion’s sampling schemes for inference: this yields improvements over those baselines, suggesting the benefit lies in the training via independent noise. 
- We present experimental evidence that, when adopting a standard technique for video diffusion models - image pre-training - the added complexity of independent noise is well warranted: we pretrain Diffusion Forcing and all baselines with images with no temporal structure, before training on video data. While this pretraining improves the performance of all baselines, Diffusion Forcing still (1) performs the best, and (2) its performance in earlier training iterations (20k) is superior to that of baselines at convergence. This shows that the overhead of independent noise is justified when considering the overall training compute (equivalent to 100k video training steps)  
- We see a strong practical benefit of independent noise when tuning hyperparameters such as uncertainty rate in pyramid sample, which is painfully inconvenient in AR-Diffusion since it requires retraining. 

### Model Capabilities
- As noted in the main text, the use of independent noise confers a number of additional capabilities in our model, including stabilization of autoregressive rollout (Sec. 4.1), modeling causal uncertainty (Sec. 4.2), and removing the need for expensive reconstruction guidance when conditioning on context  (Appendix B.4). None of these capabilities can be achieved by full-sequence diffusion and AR-diffusion can only achieve the first and third one. To demonstrate yet another capability, we present preliminary results on using Diffusion Forcing for frame interpolation without reconstruction guidance in video prediction **Figure 1 of rebuttal pdf**. 

> Further ablations

We’ve added the following ablations. 
1. An ablation of stabilization level k for our ability to stabilize autoregressive generation. As shown in **Figure 3(a) of the rebuttal pdf**, we found that video metrics monotonically improve as one goes from no stabilization to a certain level and monotonically worsen as we further increase the value.  We further present a qualitative visualization in **Figure 3(b) of the rebuttal pdf** for a more stochastic dataset and found stabilization is an indispensable technique to prevent blowing up.
2. Multiple ablations of performance vs computing budget across domains. In **Figure 2 of the rebuttal pdf**, we found that diffusion forcing can trade off speed for better quality by varying DDIM sampling steps for video and replan frequency for planning.  

> Scaling up beyond RNNs

It is straightforward to adapt Diffusion Forcing to transformers, and no changes to the proposed framework have to be made. One simply uses any architecture - causal or non-causal, and then - as discussed in the paper - varies the per-token noise level at training time.

After the initial draft, we’ve further reimplemented Diffusion Forcing with alternative architectures such as the transformer and 3D-UNet with temporal attention. We found that our approach similarly works. In addition, we’ve integrated diffusion forcing with modern techniques such as latent diffusion, allowing it to scale up to much longer video sequences (300) or those with a much higher resolution (512x512). In **Figure 1 of the rebuttal pdf**, we present non-cherry-picked samples of Diffusion Forcing on RE10K, a hard dataset that’s high-resolution, with latent diffusion and temporal attention. We’ve successfully reproduced maze planning results using transformers as well. These results show positive signs of scaling up diffusion forcing. However, we find that a detailed quantitative evaluation of these new variants would exceed the capacity of the present paper, though we will discuss this direction and interesting implications in the discussion section.



[1] Tong Wu et al. “Ar-diffusion: Auto-regressive diffusion model for text generation”. Neurips 2023.

[2]​​ David Ruhe et al. “Rolling Diffusion Models”. Arxiv 2024

Meta-review: The paper presents a new algorithm, called Diffusion Forcing, that attempts
to achieve the strengths of both purely autoregressive generation and
full-sequence generation.  The algorithm is well motivated by the
limitations of those approaches.  The results of the algorithm are
impressive both in terms of numerical performance on control benchmarks and
in the stability of longer range video prediction.

The authors were able to address some concerns about comparisons and
ablations during the rebuttal period.
Decision: Accept (poster)