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_cc_2024_Conference
Published: 2024-05-13 00:52:38

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/

Main paper: 
 introduction: 
 Probabilistic sequence modeling plays a crucial role in diverse machine learning applications including natural language processing [6, 47] , video prediction [31, 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 a single token or an "infinite" number of tokens via auto-regressive sampling), can be conditioned on varying amounts of history [21, 37] , support efficient tree search [70, 23, 25] , and can be used for online feedback control [22, 3] .Current next-token prediction models are trained via teacher forcing [64] , where the model predicts the immediate next token based on a ground truth history of previous tokens. This results in two limitations: (1) there is no mechanism by which one can guide the sampling of a sequence to minimize a certain objective, and (2) current next-token models easily become unstable on continuous data. For example, when attempting to auto-regressively generate a video (as opposed to text [6] or vector-quantized latents [33] ) past the training horizon, slight errors in frame-to-frame predictions accumulate and the model diverges.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 a 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 a fixed number of tokens by diffusing their concatenation [31, 1] , where the noise level is identical across all tokens. They offer diffusion guidance [30, 16] to guide sampling to a desirable sequence, invaluable in decisionmaking (planning) applications [36, 34] . They further excel at generating continuous signals such as video [31] . However, full-sequence 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 1 ). Further, we demonstrate that a naive attempt at combining the best of both worlds by training a next-token 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 in a later one.In this paper, we introduce Diffusion Forcing (DF), a training and sampling paradigm where each token is associated with a random, independent noise level, and where tokens can be denoised according to arbitrary, independent, per-token schedules through a shared next-or-next-few-token prediction model. Our approach is motivated by the observation that noising tokens is a form of partial masking-zero noise means a token is unmasked, and complete noise fully masks out a token. Thus, DF forces the model to learn to "unmask" any collection of variably noised tokens (Figure 2 ). Simultaneously, by parameterizing predictions as a composition of next-token prediction models, our system can flexibly generate varying length sequences as well as compositionally generalize to new trajectories (Figure 1 ). We implement DF for sequence generation as Causal Diffusion Forcing (CDF), in which future tokens depend on past ones via a causal architecture. We train the model to denoise all tokens of a 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 at each denoising step. Like next-token prediction models, CDF can generate variable-length sequences; unlike next-token prediction, it does so stabily from the immediate next token to thousands of tokens in the future -even for continuous tokens. Moreover, like full-sequence 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. 1 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 longhorizon guidance like full-sequence diffusion models. (2) Taking advantage of Diffusion Forcing's unique capabilities, we introduce a novel decision-making framework that allows us to use Diffusion Forcing as simultaneously a 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, In contrast, next-token prediction models, common in language modeling, are trained to predict a single next token from a ground-truth sequence (teacher forcing [64] ), and full-sequence diffusion, common in video generation, train non-causal architectures to denoise all frames in a sequence at once with the same noise level. Diffusion Forcing thus interleaves the time axis of the sequence and the noise axis of diffusion, unifying strengths of both alternatives and enabling completely new capabilities (see Secs. 3.2,3.4).such as stabilizing long-rollout autoregressive video generation, composing sub-sequences of those observed at training time with user-determined memory horizon, Monte Carlo Guidance, and more. 
related work and preliminaries: 
 We discuss related work and preliminaries for our core 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 t, and diffusion along an "uncertainty" (or noise level) axis denoted by superscript k. In the following, we denote observations as x ∈ X and latent states as z ∈ Z.Bayesian Filtering. Given a Hidden Markov Model (HMM) defined by latent states z t and observations x t , a Bayes filter is a probabilistic method for estimating latent states recursively over time from incoming observations. A prior model p(z t+1 |z t ) infers a belief over the next state given only the current state, and an observation model infers a belief over the next observation given the current latent state p(x t |z t ). When a new observation is made, a posterior model p(z t+1 |z t , x t+1 ) provides an updated estimation of the next latent state z t+1 . When trained end-to-end with neural networks [22, 23] , latent states are not an estimate of any physical quantity, but a sufficiently expressive latent that summarizes past observations for predicting future observations (x t ′ ) t ′ >t in the sequence.Diffusion Models. Diffusion models [56, 28] have proven to be highly expressive and reliable generative models. We review their essentials here. Let q(x) denote a data distribution of interest, and let x 0 ≡ x ∼ q. We consider a forward diffusion process that gradually adds Gaussian noise to a data point over a series of time steps. This process is modeled as a Markov chain, where the data at each step k is noised incrementally:q(x k |x k-1 ) = N (x k ; 1 -β k x k-1 , β k I) (2.1)where N is the normal distribution and β k is the variance of the noise added at each step controlled by a schedule {β k ∈ (0, 1)} K k=1 . The process continues until the data is converted into pure noise at x K . The reverse process is also a Markov chain and attempts to recreate the original data from the noise with a parameterized model p θ :EQUATIONwhere the mean µ is a model with a neural network, and where it is shown [29] that one can set the covariance to the identity scaled by a fixed constant γ k depending on k. Adopting the standard exposition, we reparametrize the mean µ in terms of noise prediction ϵ= ( √ 1 -ᾱt ) -1 x kt t - √ ᾱt µ.This leads [28] to the following least squares objective:EQUATIONwhere [28] . Guidance of Diffusion Models. Guidance [30, 16] allows biasing diffusion generation towards desirable predictions at sampling time. We focus on classifier guidance [16] : given a classifier c(y|x k ) of some desired y (e.g. class or success indicator), one modifies the Langevin sampling [29] x k = √ ᾱt x 0 + √ 1 -ᾱt ϵ k and ϵ k ∼ N (0, I) . One can then sample from this model via Langevin dynamics x k-1 ← 1 √ α k (x k t -1-α k √ 1-ᾱk ϵ θ (x k t , k) + σ k w)gradient ϵ θ (x k , k) to be ϵ θ (x k , k) - √ 1 -ᾱk ∇ x k log c(y|x k). This allows sampling from the joint distribution of x and class label y without the need to train a conditional model. Other energies such as a least-squares objective comparing the model output to a desirable ground truth have been explored in applications such as decision making [16, 36] .Next-Token Prediction Models. Next-token prediction models are sequence models that predict the next frame x t+1 given past frames x 1:t . At training time, one feeds a neural network with x 1:t and minimizes ||x -x|| 2 for continuous data or a cross-entropy loss for discrete data [64] . At sampling time, one samples the next frame xt+1 following p(x t+1 |x 1:t ). If one treats xt+1 as x t+1 , one can use the same model to predict x t+2 and repeat until a 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 Sequence Models. Diffusion has been widely used in sequence modeling. [43] use fullsequence diffusion models to achieve controllable text generation via guidance, such as generating text following specified parts of speech. [31] trains full-sequence diffusion models to 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 training on a dataset of interaction trajectories with the environment and using classifier guidance at sampling time to sample trajectories with high rewards towards a 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 a causal architecture with linearly dependent noise level along the time axis. We provide a detailed comparision between this approach and ours in Appendix D.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 (5) .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 are encouraged to create a separate "Limitations" 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
noising as partial masking: 
 Recall that masking is the practice of occluding a subset of data, such as patches of an image [26] or timesteps in a sequence [15, 48] , and training a 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 t at time t and making predictions from the past x 1: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 0 1:T ≡ x 1:T , as a form of partial masking, which we refer to as masking along the noise axis. Indeed, after K steps of noising, x K 1:T is (approximately) pure white noise without information about the original data.We establish a unified view along both axes of masking (see Fig. 2 ). We denote x 1:T for a sequence of tokens, where the subscript indicates the time axis. As above, x kt t denotes x t at noise level k t under the forward diffusion process (2.1); x 0 t = x is the unnoised token, and x K t is white noise N (0, I). Thus, (x kt t ) 1≤t≤T denotes a sequence of noisy observations where each token has a different noise level k t , which can be seen as the degree of partial masking applied to each token through noising. k ← Km,t, w ∼ N (0, I). 7:x new t ← 1 √ α k (xt -1-α k √ 1-ᾱk ϵ θ (z new t , xt, k)) + σ k w 8:Update zt ← z new t . 9: end for 10: x1:H ←AddGuidance(x new 1:H , ∇x log c(x new 1:H )) 11: end for 12: Return x1:T . paper, we focus on time series data, and thus instantiate Diffusion Forcing with causal architectures (where x kt t depends only on past noisy tokens), which we call Causal Diffusion Forcing (CDF). For simplicity, we focus on a minimal implementation with a vanilla Recurrent Neural Network (RNN) [11] . Potential transformer implementation of Diffusion Forcing is also possible but we defer its discussion to Appendix C.1.The RNN with weights θ maintains latents z t capturing the influence of past tokens, and these evolve via dynamics z t ∼ p θ (z t |z t-1 , x kt t , k t ) with a recurrent layer. When an incoming noisy observation x kt t is made, the hidden state is updated in a Markovian fashion z t ∼ p θ (z t |z t-1 , x kt t , k t )2 . When k t = 0, this is the posterior update in Bayes filtering; whereas when k t = K (and x K t is pure noise and thus uninformative), this is equivalent to modeling the "prior distribution" p θ (z t | z t-1 ) in Bayes filtering. Given latent z t , an observation model p θ (x 0 t |z t ) predicts x t . Training. The dynamics model p θ (z t |z t-1 , x kt t , k t ) and the observation model p θ (x 0 t |z t ) together form a RNN unit. Such unit has the same input-output behavior as a standard conditional diffusion model, using a conditioning variable z t-1 and a noisy token x kt t as input to predict the noise-free x t = x 0 t and thus, indirectly, the noise ϵ kt via affine reparametrization [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 ϵ θ (z t-1 , x kt t , k t ). We then find parameters θ by minimizing the lossEQUATIONwhere we sample k 1:T uniformly from [K] T , x 1:T from our training data, and ϵ t ∼ N (0, σ 2 kt I) in accordance with the forward diffusion process (see Algorithm 1 for pseudocode). Importantly, the loss (3.1) captures essential elements of Bayesian filtering and conditional diffusion. In Appendix B.1, 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 k 1:T uniformly. Finally, we prove the validity of this objective stated informally in the following Theorem 3.1 in Appendix A. Theorem 3.1 (Informal). The Diffusion Forcing training procedure (Algorithm 1) optimizes a reweighting of an Evidence Lower Bound (ELBO) on the expected log-likelihoods ln p θ ((x kt t ) 1≤t≤T ), where the expectation is averaged over noise levels k 1:T ∼ [K] T and x kt t noised according to the forward process. Moreover, under appropriate conditions, optimizing (3.1) also maximizes a lower bound on the likelihood for all sequences of noise levels, simultaneously.We remark that a special case of 'all sequences of noise levels' are those for which either k t = 0 or k t = 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 2 and is defined by prescribing a noise schedule on a 2D M × T grid K ∈ [K] M ×T ; columns correspond to time step t and rows indexed by m determine noise-level. K m,t represents the desired noise level of the time-step t token for row m. To generate a whole sequence of length T , initialize the tokens x 1:T to be white noise, corresponding to noise level k = K. We iterate down the grid row-by-row, denoising leftto-right across columns to the noise levels prescribed by K. By the last row m = 0, the tokens are clean, i.e. their noise level is K 0,t ≡ 0. Appendix B.5 discusses corner cases of this scheme; the hyperparameters (α k , ᾱk , σ k ) are set to their standard values [29] . The matrix K 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.We now explain the new capabilities this flexible sampling paradigm has to offer. Stabilizing autoregressive generation. For high-dimensional, continuous sequences such as 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 0 < k ≪ K. Our experiments (Sec. 4.1) illustrates the resulting marked improvements in long-horizon generation capabilities; App. C.4 provides further intuition.Keeping the future uncertain. Beginning from a sequence of white noise tokens [x K 1 , x K 2 , x K 3 ] ⊤ , we may denoise the first token fully and the second token partially, yielding[x 0 1 , x K/2 2 , x K 3 ] ⊤ , then [x 0 1 , x 0 2 , x K/2 3] ⊤ , and finally denoising all tokens fully to [x 0 1 , x 0 2 , x 0 3 ] ⊤ . 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 to more effective sequence guidance.Long-horizon Guidance. In Line 10 of Algorithm 2, one may add guidance to the partially diffused trajectory x 1: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 on implementation details in Appendix C.3. As we show in Section 4.2, planning in this manner significantly outperforms guided full-sequence diffusion models.The capabilities offered by Diffusion Forcing motivate our novel framework for sequential decision making (SDM), with key applications to robotics and autonomous agents. 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 Forcing as a policy), or lengthened to perform long-horizon 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). Monte Carlo Guidance (MCG), future uncertainty. Causal Diffusion Forcing allows us to influence the generation of a token x k t by guidance on the whole distribution of future x t+1:T . Instead of drawing a single trajectory sample to calculate this guidance gradient, we can draw multiple 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 k t 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-zag schedule described in Sec. 3.3), accounting for greater uncertainty farther into the future. Appendix C.5 further justifies the significance of MCG, and why Diffusion Forcing uniquely takes advantage of it.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 as 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 or 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We extensively evaluate Diffusion Forcing's merits as a generative sequence model across diverse applications in video and time series prediction, planning, and imitation learning. Please find the dataset and reproducibility details in the Appendix, as well as video results on the project website.We train a convolutional RNN implementation of Causal Diffusion Forcing for video generative modeling on videos of Minecraft gameplay [68] and DMLab navigation [68] . At sampling time, weMaze2d-large-v1 perform auto-regressive rollout with stabilization proposed in Sec. 3.3. We consider two baselines, both leveraging the same exact RNN architecture: a next-frame diffusion baseline trained with teacher forcing [64] as well as a causal full-sequence diffusion model. Figure 3 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 a consistent 3D environment. This highlights the ability of Diffusion Forcing to stabilize rollouts of high-dimensional sequences without compounding errors.Planning: MCG, Causal Uncertainty, Flexible Horizon Control.Decision-making uniquely benefits from Diffusion Forcing's capabilities. We evaluate our proposed decision-making framework in a standard offline RL benchmark, D4RL [18] . Specifically, we benchmark Diffusion Forcing on a set of 2D maze environments with sparse reward. An agent is tasked with reaching a designated goal position starting from a random starting position. In Appendix E.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-making framework 3.4 with state-of-the-art offline RL methods and the recently introduced Diffuser [36] , a diffusion planning framework. See Fig. 1 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 to maximize expected reward, as we formally derive in Appendix C.5. To understand MCG's importance, we ablate it in Table 1 . Removing MCG guidance degrades our performance, though Diffusion Forcing remains competitive even then. Figure 4 : In our real robot task, a robot arm is asked to swap the slots of two fruits using a third slot. Since the fruits are input in random slots at the beginning, one cannot determine the next steps from a 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 can vary-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 are more important than those in the far future. Though Diffuser and subsequent models are trained to generate sequences of action-reward-state tuples [a t , r t , o t ], directly executing the actions will lead to a trajectory that deviates significantly from the generated states. In other words, the generated states and actions are 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 1 , we see that Diffuser's performance drops dramatically when directly executing generated actions. In contrast, Diffusion Forcing's raw action generations are self-consistent, outperforming even actions selected by combining Diffuser's state predictions with a handcrafted PD controller.Benefit of Flexible Horizon. Many RL tasks have a 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 we explain in Appendix C.6.We demonstrate that by only modifying the sampling scheme, we can flexibly compose sub-sequences of sequences observed at training time. We consider a dataset of trajectories on a 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. 1 , 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-trajectories, forming a V-shaped trajectory. Due to limited space, we defer the result to Appendix E.2.Finally, we illustrate that Diffusion Forcing (DF) opens up new opportunities in the visuomotor control of real-world robots. Imitation learning [10] is a popular technique in robotic manipulation where one learns an observation-to-action mapping from expert demonstrations. However, the lack of memory often prevents imitation learning from accomplishing long-horizon tasks. DF not only alleviates this shortcoming but also provides a way to make imitation learning robust.Imitation Learning with Memory. We collect a 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. 4 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 its latent state. We found that DF achieves 80% success rate while diffusion policy [10] , a state-of-the-art 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 to 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 k > 0, in 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 as ground truth and suffers out-of-distribution error.Potential for pre-training with video. Finally, in parallel to generating actions, Fig. 4 illustrates 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-training on unlabeled video.In Appendix E, we 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 is based on an RNN. Applications to higherresolution video or more complex distributions likely require large transformer models following instructions in Appendix C.1. We do not investigate the scaling behavior of Diffusion Forcing to internet-scale datasets and tasks.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 to dramatic performance gains when applied to 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 Forcing to larger datasets.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 a reweighting of the Evidence Lower Bound (ELBO) on the average log-likelihood of our data. We first establish this in full generality (Theorem A.1), 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 of 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 to the terms optimized in the Diffusion Forcing training objective maximizes. Thus, for a 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 q(•), and generative models p θ (•), and then specialize to Gaussian diffusion, thereby recovering our loss. We denote our Markov forward process q(•) asq(x 1:K | x 0 ) = K k=1 q(x k | x k-1 ), (A.1)and a parameterized probability modelp θ (((x k t ) 1≤k≤K , z t ) t≥1 ) (A.2)We assume that p θ satisfies the Markov property thatp θ (z t , x kt t | z 1:t-1 , (x ks s ) 1≤s<t ) = p θ (z t , x kt | z t-1 ) (A.3)that is, the latent codes z t-1 is a sufficient statistic for x kt given the history. We say that p θ has deterministic latents if p θ (z t | z 1:t-1 , (x ks s ) 1≤s<t , x kt t ) is a Dirac delta. Remark 1. In order for p θ to have deterministic latents and correspond to a valid probability distribution, we need to view the latents z t not as individual variables, but as a collection of variables z t (k 1:t ) indexed by t ∈ [T ] and the history of noise levels k 1:t ∈ {0, 1, . . . , K} t . In this case, simply setting z t (k 1:t ) = (k 1:t , (x ks s ) 1≤s≤t tautologically produces deterministic latents. The reason for indexing z t (k 1:t ) with k 1:t then arises because, otherwise, p θ (z t | ((x ks s ) 1≤s≤t , (xk ′ s s ) 1≤s≤t ) would be ill-defined unless k s = k ′s for all 1 ≤ s ≤ t, and thus, p θ would not correspond to a joint probability measure. The exposition and theorem that follows allow z t (k 1:t ) to be indexed on past noise levels k 1:t but suppresses dependence on k 1:t to avoid notational confusion.We can now state our main theorem, which provides an evidence lower bound (ELBO) on the expected log-likelihood of partially-noised sequences (x kt t ) 1≤t≤T , under uniformly sampled levels k t and x kt t obtained by noising according to q(•) as in (A.1). Notice that this formulation does not require an explicit for of q(•) or p θ , but we will specialize to Gaussian diffusion in the following section.Theorem A.1. Fix x 0 1:T . Define the expectation over the forward process with random noise levelk 1:T as E forward [•] := E k1,...,k T unif ∼ [K] E x ks s ∼q(x ks s |x 0 s ),1≤s≤T [•], (A.4)and the expectation over the latents under p θ (•) conditioned on k 1:T , (x kt s ) 1≤t≤T asE p,z 1:T [•] := E zs∼p(zs|zs-1,x ks s ),s≤T • | k 1:T , (x kt t ) 1≤t≤T (A.5)Then, as long as p θ satisfies the Markov property,E forward [ln p θ ((x kt t ) 1≤t≤T )] ≥ C(x 0 1:T ) + E forward E p,z 1:T   T t=1   1 K + 1 ln p θ (x 0 t | x 1 t , z t-1 ) + K j=2 j K + 1 D KL q(x j-1 t | x j t , x 0 t ) ∥ p θ (x j t | x j-1 t , z t-1 )     ,where C(x 0 1:T ) is a constant depending only on x 0 1:T (the unnoised data). Moreover, if the latents are deterministic (i.e. p θ (z t | z t-1 , x kt t ) is a Dirac distribution), then the inequality holds with inequality if and only if q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1 ), i.e. the variational approximation is exact.The proof of the above theorem is given in Appendix A.2. Remarkably, it involves only two inequalities! The first holds with equality under deterministic latents and the second holds if and only if variational approximation is exact:q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1). This tightness of the ELBO suggests that the expression in Theorem A.1 is a relatively strong surrogate objective for optimizing the likelihoods.We now special Theorem A.1 to Gaussian diffusion. For now, we focus on the "x-prediction" formulation of diffusion, which is the one used in our implementation. The "ϵ-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, 7] to Theorem A.1.Corollary A.2. Let q(x k+1 | x k t ) = N (x k ; 1 -β k x k-1 , β k I), (A.6)and defineα k = (1 -β k ), ᾱk = k j=1 α j . Suppose that we parameterize p θ (x j t | x j+1 t , z t-1 ) = N (µ θ (x j+1 t , z t-1 , j), σ 2 j ), where further,µ θ (x j t , z t-1 , j) = (1 -ᾱj-1 ) √ α j 1 -ᾱj x j t + (1 -α j ) √ ᾱj-1 1 -ᾱj xθ (x j t , z t-1 , j), σ 2 j := (1 -α j )(1 - √ ᾱj-1 )1 -ᾱj .Then, as long as p θ satisfies the Markov property, we obtainedE forward [ln p θ ((x kt t ) 1≤t≤T )] + C(x 0 1:T ) ≥ E forward E p,z 1:T   T t=1 j K + 1 K j=1 c j ∥x 0 θ (x j t , z t-1 , j) -x 0 t ∥ 2   = E forward E p,z 1:T T t=1 1{k t ≥ 1} • k t c kt ∥x 0 θ (x kt t , z t-1 , k t ) -x 0 t ∥ 2 ,where above, we definec j = (1-αj ) 2 ᾱj-1 2σ 2 (1-ᾱj ) 2 .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 k t unif ∼ {0, 1, . . . , K}.We make a couple of remarks:• As noted above, Corollary A.2 can also be stated for ϵ-prediction, or the so-called "vprediction" formalism, as all are affinely related.• Define an idealized latent zt-1 consisting of all past tokens (x kt t ) as well as of their noise levels k t . This is a sufficient statistic for z t-1 , and thus we can always viewx0 θ (x kt t , z t-1 , k t ) = x0 θ (x kt t , zt-1 , k t ), where z t-1 is just compressing zt-1 . When applying the expectation of x 1:T ∼ q to both sides of the bound in Corollary A.2, and taking an infimum over possible function approximator x0 θ , we obtain infp θ E q E forward E p,z 1:T ∥x 0 θ (x kt t , z t-1 , k t ) -x 0 t ∥ 2 = inf p θ E q E forward E p,z 1:T ∥x 0 θ (x kt t , zt-1 ) -x 0 t ∥ 2 = Var q [x 0 t | (x ks s ) 1≤s≤t , k 1 , . . . , k t ]. This leads to a striking finding: with expressive enough latents and p θ , we can view the maximization of each term in Corollary A.2 separately across time steps. The absence of this coupling means that the weighting terms are immaterial to the optimization, and thus can be ignored. Ξ t :=   P t (1 | k 1:t-1 ) ln p θ (x 0 t | x 1 t , z t-1 ) + K j=2 jP t (j | k 1:t-1 )D KL q(x j-1 t | x j t , x 0 t ) ∥ p θ (x j t | x j-1 t , z t-1 )   ,where C(x 0 1:T ) is a constant depending only on x 0 1:T (the noise-free data), and where the inequality is an equality under the conditions that (a)p θ (z t | z t-1 , x kt t ) is a Dirac distribution (deterministic latents), and (b) q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1), i.e. the variational approximation is sharp.In particular, in the Gaussian case of Corollary A.2, we haveE forward,D [ln p θ ((x kt t ) 1≤t≤T )] + C(x 0 1:T ) ≥ E forward,D E p,z 1:T T t=1 1{k t ≥ 1}k t c kt ∥x 0 θ (x kt t , z t-1 , k t ) -x 0 t ∥ 2 ,The most salient case for us is the restriction of D to fixed sequences of noise k 1 , . . . , k T (i.e. Dirac distributions on [K] T ). In this case, P t (j | k 1:t-1 ) = 0 for all but j = k t , and thus our training objective need not be a lower bound on E forward,D [ln p θ ((x kt t ) 1≤t≤T )]. However, the terms in the lower bound are, up to reweighting, an subset of those terms optimized in the training objective. Thus, in light of the remarks following Corollary A.2, a fully expressive network can optimize all the terms in the loss simultaneously. We conclude that, for a fully expressive neural network, optimizing the training objective (3.1) is a valid surrogate for maximizing the likelihood of all possible noise sequences. where the inequality is and equality when p θ (z s | z s-1 , x ks s ) is a Dirac distribution. By applying E forward to both sides of the above display, and invoking the Markov property of the latents, we conclude thatEQUATIONE forward [ln p θ ((x kt t ) 1≤t≤T )] ≥ E forward E p,z 1:T T t=1 ln p θ (x kt t | z t-1 ) = T t=1 E <t E kt unif ∼ {0,1,...,K} E x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt t | z t-1 ).We now unpack the terms obtained from the preceding claim.Claim 2 (ELBO w.r.t. q). It holds thatE x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt t | z t-1 ) ≥ C 1 (x 0 , k t ) + E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ).where C 1 (x 0 , k t ) is a constant depending only on x 0 and k t , and where the inequality holds with equality if and only if q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1 ).Proof. We have thatE x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt t | z t-1 ) = E x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt:K t | z t-1 )dx kt+1:K t = E x k t t ∼q(x k t t |x 0 t ) ln E x k t +1:K t ∼q(x k t +1:K t |x k t t ) p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x kt t ) ≥ E x k t t ∼q(x k t t |x 0 t ) E x k t +1:K t ∼q(x k t +1:K t |x k t t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x kt t ) ((Jensen's inequality)) = E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x kt t ) (Markov property of q(•)) = C 1 (x 0 , k t ) + E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ),where the constant C 1 (x 0 , k t ) = E x k t :K t ∼q(x k t :K t |x 0 t ) ln q(x k t +1:K t |x 0 t ) q(x k t +1:K t |x k t t )depends only on x 0 and k t .To check the conditions for equality, note that if q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1 ), then E x k t +1:K t ∼q(x k t +1:K t |x k t t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x kt t ) = ln p θ (x kt t | z t-1 ) + E x k t +1:K t ∼q(x k t +1:K t |x k t t ) ln p θ (x kt+1:K t | z t-1 , x kt t ) Since ln(•) is strictly concave, E x k t +1:K t ∼q(x k t +1:K t |x k t t ) ln p θ (x kt t | z t-1 ) = 0 if and only if p θ (x kt+1:K t | z t-1 , x kt t ) = q(x kt+1:K t | x kt t ).Claim 3 (Computing the expected ELBO).E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ) = C 3 (x 0 , k t ) + 1{k t = 0} ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 1{j ≥ k t }D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) ,where C 2 (x 0 , k t ) is some other constant depending on x 0 and k t .Proof. The proof invokes similar manipulations to the standard ELBO derivation for diffusion, but with a few careful modifications to handle the fact that we only noise to level k t . As is standard, we require the identityq(x j t | x j-1 t , x 0 t ) = q(x j-1 t | x j t , x 0 t ) • q(x j t | x 0 t ) q(x j-1 t | x 0 t ) . (A.10)Part 1: Expanding the likelihood ratios . Using the above identity, we obtainln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ) = ln p(x K t | z t-1 ) + ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt+1 t | x 0 t ) + K j=kt+2 ln p θ (x j-1 t | x j t , z t-1 ) q(x j t | x j-1 t , x 0 t ) (i) = ln p(x K t | z t-1 ) + ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt+1 t | x 0 t ) + K j=kt+2 ln p θ (x j-1 t | x j t , z t-1 ) q(x j-1 t | x j t , x kt t ) + ln q(x j-1 t | x 0 t ) q(x j t | x 0 t ) (ii) = ln p(x K t | z t-1 ) + ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt+1 t | x 0 t ) + ln q(x kt+1 t | x kt t ) q(x K t | x kt t ) + K-1 j=kt+1 ln p θ (x j t | x j+1 t , z t-1 ) q(x j t | x j+1 t , x 0 t ) = ln p(x K t | z t-1 ) q(x K t | x kt t ) + ln p θ (x kt t | x kt+1 t , z t-1 ) + K-1 j=kt+1 ln p θ (x j t | x j+1 t , z t-1 ) q(x j t | x j+1 t , x 0 t ) = ln q(x kt t | x kt+1 t ) 1{kt≥1} + ln p(x K t | z t-1 ) q(x K t | x kt t ) + ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt t | x kt+1 t ) 1{kt≥1} + K-1 j=kt+1 ln p θ (x j t | x j+1 t , z t-1 ) q(x j t | x j+1 t , x 0 t ), where (i) uses A.10, (ii) invokes a cancellation in the telescoping sum, and the final display follows from the computationq(x kt t | x kt+1 t ) 1{kt≥1} = 1 k t = 0 q(x kt t | x kt+1 t ) k t ≥ 1 . (A.11)Observe that, because we don't parameterize p(xK t | z t-1 ), ln q(x kt t | x kt+1 t ) 1{kt≥1} + ln p(x K t |zt-1) q(x K t |x k t t )can be regarded as some constant C ′ (x kt t , x kt+1 t , x K t ). Thus,ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ) = C ′ (x kt t , x kt+1 t , x K t ) + ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt t | x kt+1 t ) 1{kt≥1} + K-1 j=kt+1 ln p θ (x j t | x j+1 t , z t-1 ) q(x j t | x j+1 t , x 0 t ) (A.12)Part 2: Taking expecations. We can now simplify to taking expectations. Observe thatE x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x j t | x j+1 t , z t-1 ) q(x j t | x j+1 t , x 0 t ) = D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) ,and similarly,E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt t | x kt+1 t , z t-1 ) q(x kt t | x kt+1 t ) 1{kt≥1} = ln p θ (x 0 t | x 1 t , z t-1 ) k t = 0 D KL q(x kt t | x kt+1 t , x 0 t ) ∥ p θ (x kt t | x j+1 t , z t-1 ) k t ≥ 1.Finally,E x k t :K t ∼q(x k t :K t |x 0 t ) C ′ (x kt t , x kt+1 t , x K t ) is a constant C 2 (k t ,x 0 ) depending only on k t , x 0 . Thus, from (A.12)E x k t :K t ∼q(x k t :K t |x 0 t ) ln p θ (x kt:K t | z t-1 ) q(x kt+1:K t | x 0 t ) = C 2 (k t , x 0 ) + 1{k t = 0} ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=max{1,kt} D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) = C 2 (k t , x 0 ) + 1{k t = 0} ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 1{j ≥ k t }D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) .Completing the proof of the ELBO. We are now ready to complete the proof. By combining the previous two claims, we haveE x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt t | z t-1 ) ≥ C 3 (x 0 , k t ) + 1{k t = 0} ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 1{j ≥ k t }D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) ,where C 3 (x 0 , k t ) = C 1 (x 0 , k t ) + C 2 (x 0 , k t ) and where again, the above is an equality when q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1 ). Taking an expectation over k t unif ∼ {0, 1, . . . , K}, we haveE kt unif ∼ {0,1,...,K} [1{k t = 0}] = 1 K + 1 , E kt unif ∼ {0,1,...,K} 1{j ≥ k t } = j + 1 K + 1 . (A.13)and consequently,E kt unif ∼ {0,1,...,K} E x k t t ∼q(x k t t |x 0 t ),1≤t≤T ln p θ ((x kt t ) 1≤t≤T ) ≥ C 4 (x 0 t ) + 1 K + 1 ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 j + 1 K + 1 D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) Invoking Claim 1, E forward [ln p θ ((x kt t ) 1≤t≤T )] ≥ T t=1 E <t E kt unif ∼ {0,1,...,K} E x k t t ∼q(x k t t |x 0 t ) ln p θ (x kt t | z t-1 ) = T t=1 E <t   C 4 (x 0 t ) + 1 K + 1 ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 j + 1 K + 1 D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 )  We conclude by observing that T t=1 E <t C 4 (x 0 t ) is a constant C(x 0 1:T ), and thatE <t ln p θ (x 0 t | x 1 t , z t-1 ) = E forward E p,z 1:T ln p θ (x 0 t | x 1 t , z t-1 ) E <t D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) = E forward E p,z 1:T D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 ) ,since both terms only depend on k 1:t-1 , (x ks s ) 1≤s≤t-1 and z 1:t-1 . We conclude then thatE forward [ln p θ ((x kt t ) 1≤t≤T )] ≥ C(x 0 1:T ) + E forward E p,z 1:T   T t=1   1 K + 1 ln p θ (x 0 t | x 1 t , z t-1 ) + K-1 j=1 j + 1 K + 1 D KL q(x j t | x j+1 t , x 0 t ) ∥ p θ (x j t | x j+1 t , z t-1 )     ,as needed. Lastly, we recall that the above is an equality under the conditions that (a) p θ (z t | z t-1 , x kt t ) is a Dirac distribution, and (b) q(x kt+1:T t | x kt t ) ≡ p θ (x kt+1:T t | x kt t , z t-1), and we reindex j ← j + 1 to ensure consistency with indexing in standard expositions of the diffusion ELBO.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-noise ratio (SNR) of noisy x k . In Diffusion Forcing, conditioning variable z t-1 can also contain a non-trivial amount of information about x t , in addition to x kt t . For example, in a deterministic markovian system, if x kt-1 t-1 has its noise level k t-1 = 0, the posterior state z t-1 contains all the information needed to predict x 0 t regardless of the noise level of x kt t . 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 t as the normalized SNR reweighting factor for x kt t following its normal derivation in diffusion models. For example, if one uses min snr strategy [24] , its reweighting factor will always fall between [0, C] which we divide by C to get S t ∈ [0, 1]. Define signal decay factor 0 < γ < 1, measuring what proportion of signal in x kt-1 t-1 contribute to denoising x kt t . This is the simple exponential decay model of sequential information. Now, define cumulated SNR recursively as the running mean of S t : St = γ St-1 + (1 -γ)S t to 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 t of noisy observation x kt t . To combine them, we use a simplified model for independent events. Notice S t and St always falls in range [0, 1], and therefore can be reinterpreted as probabilities of having all the signal one needs to perfect denoise x kt t . Since the noise level at t is independent of prior noise levels, we can view S t and St-1 as probabilities of independent events and thus can composed to define a joint probabilityS ′ t = 1 -(1 -S t )(1 -St-1), and we use this S ′ t as our fused SNR reweighting factor for diffusion training.In our experiments, we choose to follow the min-SNR reweighting strategy [24] to derive the S. Our Fused SNR reweighting proves extremely useful to accelerate the convergence of video prediction, while we didn't observe a boost on non-image domains so we didn't use it there.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 kt t |z t-1 ) with a typical diffusion U-net [46] . We use the output of the U-net as the input to a gated recurrent unit (GRU) and use z t-1 as the hidden state feed into a GRU. The output of GRU is treated as z t . For observation model p(x t |z t ), 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 z t-1 , input is x kt t and output is x. One can potentially obtain better results by training Diffusion Forcing with a causal transformer architecture. However, since RNN is more efficient for online decision-making, we also stick with it for video prediction and it already gives us satisfying results.We choose the number of channels in z to be 16 for DMlab and 32 for Minecraft. In total, our Minecraft model consists of 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 (< 1 day). In maze planning, the number of total parameters is 4.33 million.Non-Video Diffusion For non-spatial x that is not video nor images, we use residue MLPs [59] 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 z t-1 to get z t . Another ResMLP serves as the observation model.In diffusion models, there are three equivalent prediction objectives, x 0 , ϵ [28] , and v parameterization [52] . Different objectives lead to different reweighting of loss at different noise levels, together with SNR reweighting. For example, ϵ parameterization and v parameterization are essential in generating pixel data that favors high-frequency details.In our experiments, we use v parameterization for video prediction and found it essential to both convergence speed and quality.We observe that x 0 parameterization is strongly favorable in planning and imitation learning, likely because they don't favor an artificial emphasis on high-frequency details. We observe the benefits of v-parameterization in time-series prediction.We use sigmoid noise schedule [9] for video prediction, linear noise schedule for maze planning, and cosine schedule for everything else.Corner case of sampling noise In our sampling algorithm, due to the flexibility of the scheduling matrix K, there are corner cases when x kt t 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 kt t at all. One option is just copying over the old value. The other option is to run a backward diffusion followed by a 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 t as the dynamics model p(z t |x kt t , z t-1 ) can still propagate the guidance gradient to z t-1 .Other than Monte Carlo Guidance, this corner case only happens when k t = 0 or k t = K throughout our experiments. That is, we chose our K such that once any token gets diffused slightly, it will keep diffusing. In the case of k t = K, keeping x kt t at the same noise level implies it will stay as white noise, and we don't even need to sample another white noise. In case k t = 0, the time step is already completely diffused either approach should give us the same result so 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 log-likelihood ||p T -g||, but a flexible horizon model should not require users to manually specify a T to reach its goal, instead we want it to try to reach the goal for any possible horizon. Therefore we use the reward model t ||p T -g|| 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 is 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 it didn't work even with a good amount of tuning.Accelerating the diffusion 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-making 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 so we can start with canonical hyperparameters of diffusion models. We use a frame stack of 4 for DMlab video prediction, 8 for Minecraft, and 10 for maze planning.At sampling time, we also have a design choice to reduce compute usage, as reflected in line 8 of Algorithm 2. In line 8, we directly assign z new t to z t , instead of recalculating z t with posterior model p(z t |z t-1 , x new t , k -1). Since the model is trained to condition on z t estimated from arbitrary noisy history, we recognize that both are valid approaches. The reason why the choose line 8 is twofold. First, it cuts the compute by half, avoiding computing posterior every step. Second, this happens to be what we want for stabilization -z new t already contains the information of the clean x new t under our simplified observation model, and happens to be estimated with k = k t , a noise level higher than that of x new t . This happens to implement the behavior we want for stabilization.Inference is depicted in Algorithm 2 and Figure 2 . In Equation B.1, we illustrate a specific instantiation of the K matrix we used for causal planning. For simplicity, we denote the case where a latent z 0 is given and aim to generate x 1:H+1 .K pyramid =                K K K ... K K -1 K K ... K K -2 K -1 K ... K . . . . . . . . . . . . . . . 1 2 3 ... H 0 1 2 ... H -1 . . . . . . . . . . . . . . . 0 0 0 ... 1 0 0 0 ... 0                (B.1)Diffusion Forcing begins by sampling our sequences as white noise with noise level K. It then denoises along each row m = 1, . . . , M of K in decreasing order. It does so by proceeding sequentially through frames t = 1, . . . , T , updating the latent (Line 5 of Algorithm 2), and then partially applying the backward process to noise level k = K m,t dictated by the scheduling matrix K (Line 6-7 of Algorithm 2). We call a K like this pyramid scheduling, as the tokens in the far future are kept at higher noise level than near future.We report the episode reward of Diffusion Forcing for different maze planning environments in Table 1 . However, we found that the episode reward isn't necessarily a 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 at the end of the episode. The dataset never contains data on 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 an 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 a result. This metric seems to exploit flaws in the environment design -a good design would involve a penalty of longer time taken to reach the goal. Therefore, in future works based on our paper, we encourage alternative metrics like the time it takes to reach the goal for the first time, which Diffusion Forcing excels at.We follow the implementation of pytorch-ts, where the validation set is a random subset of the training set with the same number of sequences as the test set. We use early stopping when validation crps-sum hasn't increased for 6 epochs. We leverage the same architecture (1 mlp and 4 grus) as well as a batch size of 32.All of our experiments use f p16 mixed precision training. Time series, maze planning, compositionally, and visual imitation experiments can be trained with a single 2080T i with 11GB of memory. We tune the batch size such that we fully use the memory of GPUs. This translates to a batch size of 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 4 -8 hours to train although there is without doubt a 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 × 16. It usually takes 12 hours to converge at 40K steps of training (occasional validation time also included).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 a transformer-based sequence diffusion model to train with independent noise levels across tokens and follow the techniques listed in Section B.1. A strict implementation of causal Diffusion Forcing would involve a causal attention mask on the transformer. However, Diffusion Forcing's fractional masking can do something more interesting: Consider the scenario that we use a transformer without a causal mask. We can still implement causality by controlling noise. By labeling the future as full white noise, there is no information leaked into the past tokens. By labeling future tokens as free of noise, we make the model completely non-causal. By labeling the future tokens as noisy, a 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.When training Diffusion Forcing, we choose to sample per-token noise level following i.i.d uniform distribution from [1, 2...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.3, modeling causal uncertainty 3.3, 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 one. 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 Forcing. In contrast, any tuning of the sampling scheme would require re-training the model for 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.1. 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-training, 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 is 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 in prior works [39] . This shows that the overhead added by independent noise is well-warranted when considering the overall training compute (including image pre-training).As stated in Section 2, one can use the gradient of the logarithmic of a classifier log c(y|x k t ) to guide the sampling process of diffusion model towards samples with a desired attribute y. For example, y can refer to the indicator of a success event. However, we can consider the logarithmic of a more general energy function c(x k t ). This has the interpretation as Pr y|x k t , where Pr y = 1 | x k t = e c(x k t ) . Some popular candidate energies includec(x k t ) = E t ′ >t r′(x k t ′ t ′ ) | x k t , (C.1)corresponding to a cost-to-go; we can obtain unbiased estimates of this gradient by using cumulative reward c(x k t ) = t ′ >t r′(x k t ′ t ′ ).We can also use goal distance c = -∥x k T T -g∥ 2 as a terminal reward. We provide details about the guidance function deployed in the maze2d planning experiment in Appendix B.5.Here, we explain in detail how we use noising to stabilize long-horizon generation. At each time t, during the denoising, we maintain a latent z k small t-1 from the previous time step, with 0 < k small ≪ K corresponding to some small amount of noise. We then do next token diffusion to diffuse the token x t across noise levels x K t , x K-1 t , . . . , x 0 t (corresponding to Algorithm 2 with horizon T = 1, initial latent z k t-1 , and noise schedule K m,1 = m); this process also produces latents z K t , z K-1 t , . . . , z 0 t associated with each noise level. From these, we use the latent z k small t to repeat the process. This noised latent can be interpreted as an implementation of conditioning on x k small t in an autoregressive process. In a potential transformer implementation of Diffusion Forcing as we discussed in Appendix C.1, one can instead run a forward diffusion on a fully diffused token to achieve stabilization.It is widely appreciated that adding noise to data ameliorates long-term compounding error in behavior cloning applications [38, 42] , and even induces robustness to non-sequential adversarial attacks [13] . In autoregressive video generation, the noised x k small t 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 [42] , where the imitiator (in our 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 to train Diffusion Forcing, and using slightly noised tokens for autoregression at test time, our approach inherits the theoretical guarantee of the HINT algorithm [5] .Monte Carlo Guidance provides substantial variance reduction in our estimate of cost-to-go guidance (C.1). 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 no obvious source of randomness to 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 to our knowledge, this approach is nonstandard.A naive way to obtain flexible horizon generation in full-sequence diffusion is via the "replacement trick": consider a full sequence model trained to diffuse x 1:T , which we partition into x 1:t-1 , x t:T ]. Having diffused tokens x 1:t-1 , we can attempt to denoise tokens of the form [x k 1:t-1 , x k t:T ], where we fix xk 1:t-1 = x 1:t-1 to be the previously generated token, and only have score gradients update the remaining x k t:T . One clear disadvantage of this method is inefficiency -one still needs to diffuse the whole sequence even when there is one step left at t = T -1. What's more, [31] points out that this approach of conditioning, named "conditioning by replacement", is both mathematically unprincipled Reconstructing masked tokens. Masked Autoencoders for images [26] and videos [17] are a popular method for representation learning in pixel space. They have been extended to perform diffusion to generate masked patches conditioned on unmasked ones [62, 19] .Casting Image Generation as Sequence Generation. [60, 8] show that even generative modeling of non-sequential data, such as images, can be fruitfully cast as sequence generative modeling.Non-Diffusion Probabilistic Sequence Models. [12] parameterize token-to-token transitions via a 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 a 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 on 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 is Rolling Diffusion [51] , which proposes to diffuse a sequence with near future more certain and far future more uncertain, resembling the causally uncertain sampling scheme of Diffusion Forcing. Like AR-Diffussion, Rolling Diffusion's training noise levels are linearly dependent on 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 as well.To illustrate Diffusion Forcing's new training objective does not degrade it as a generic sequence model, we evaluate Diffusion Forcing on high-dimensional and long-horizon sequence prediction tasks in time series prediction. We adopt multiple time series datasets with real-world applications from GluonTS [2] and evaluate Diffusion Forcing with strong baselines with standard metrics in this domain. In this section, we mainly focus on the results and analysis. For a detailed description of datasets and the metric, we refer the reader to Appendix F.4.Problem Formulation Let X = {x t } T t=1 be a sequence (multivariate time series) of D-dimensional observations x t ∈ R D of some underlying dynamical process, sampled in discrete time steps t ∈ {1, . . . , T }, where T ∈ N. In the problem setting of probabilistic time series forecasting, the sequence X = {X c , X p } is split into two subsequences at time step t 0 ∈ N with 1 < t 0 ≤ T : the context window X c := {x t } t0-1 t=1 (also called history or evidence) of length t 0 -1, and the prediction window X p := {x t } T t=t0 of length T -t 0 + 1 (also known as the prediction horizon). Then, the task is to model the conditional joint probability distribution q(x t0:T | x 1:t0-1 ) :=T t=t0 q(x t | x 1:t-1 ) (E.1)over the samples in the prediction window. If we know the distribution in (E.1), we can sample forecast prediction sequences given some initial context from the evidence sequence. However, most time-dependent data generation processes in nature have complex dynamics and no tractable formulation of q(x t0:T | x 1:t0-1 ). Instead, we construct a statistical model that approximates the generative process in (E.1) 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 sum (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 Wikipedia 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 ± 0.004 GARCH [61] 0.023 ± 0.000 0.880 ± 0.002 0.190 ± 0.001 0.370 ± 0.001 --DeepAR [54] -0.336 ± 0.014 0.023 ± 0.001 0.055 ± 0.003 -0.127 ± 0.042 LSTM-Copula [53] 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 NKF [14] -0.320 ± 0.020 0.016 ± 0.001 0.100 ± 0.002 -0.071 ± 0.002 Transformer-MAF [50] 0.005 ± 0.003 0.301 ± 0.014 0.021 ± 0.000 0.056 ± 0.001 0.179 ± 0.002 0.063 ± 0.003 TimeGrad [49] 0.006 ± 0.001 0.287 ± 0.020 0.021 ± 0.001 0.044 ± 0.006 0.114 ± 0.020 0.049 ± 0.002 ScoreGrad sub-VP SDE [67] 0 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 3 . We access the datasets from GluonTS [2] , and set the context and prediction windows to 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 sum via quantiles. As a metric, CRPS sum 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-based methods in time series forecastings, such as TimeGrad [49] and the transformer-based Transformer-MAF [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 sum 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 sum on the test set at the end of the training, which we report in 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 of Diffusion Forcing, and that we merely seek to demonstrate that the Diffusion Forcing objective is applicable to diverse domains with no apparent trade-off in performance over baseline objectives.Since Diffusion Forcing models the joint distribution of any subset of a sequence, we can leverage this unique property to achieve compositional behavior -i.e., Diffusion Forcing can sample from the distribution of 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 in 7, consider a dataset of trajectories on a 2D, square plane, where all 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 as 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). Infinite Rollout without sliding window Diffusion Forcing can rollout longer than maximum training horizonwithout sliding window. That is, we run Diffusion Forcing's RNN continuously without ever reinitializing z 0 . This is a surprising effect we observed from the rollout stabilization property of Diffusion Forcing. In Figure 8 , 10, 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 a maximum length of 72 frames for Minecraft and 36 frames for DMLab, illustrating it can rollout 2x-5x times longer than it's trained on without sliding window. In addition, we also tried 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 is the "dirt" block, but this is more of a dataset issue E.5 and the agent is able to recover after it turns around.Consistency We also present additional results where we only generate within our maximum training length. As shown in figure 13 12, Diffusion Forcing can generate consistent videos. Results are not cherry-picked. We provide some additional visualizations of causal planning in 15. We also present additional visualization of Diffusion Forcing performing model predictive control in action. As shown in figure 14 , Diffusion Forcing can generate plans of shorter horizons since it's flexible horizon. The first few frames marked in red are the ground truth images of the dataset used for conditioning.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 we test with such a bag in the scene 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.We adopt the video prediction dataset Minecraft and DMlab used by TECO [68] .The Minecraft navigation dataset consists of first-person-view videos of random walks in the Minecraft 'swamp' biome. The agent walks via a technique called 'sprint jump' which allows it to jump across blocks without getting stuck at 1 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 is 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 is when the agent runs into obstacles with a height of 2 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 as a problem of this dataset itself.DMLab Navigation Deepmind Lab navigation dataset consists of random walks in a 3D maze environment. For DMLab, the resolution is 64 pixels and we use subsequences of 48 frames. We also disregard the provided actions due to training.We note that the VQ-VAE latent that stable video diffusion [4] diffuses is also only 128 × 128 × 3, indicating Diffusion Forcing has the potential to scale up to higher resolution images with pre-trained image encoder and decoders. Due to the sheer size of the datasets, we 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.D4RL [18] 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.Like Diffuer [36] , we choose the 3 maze environments as they are challenging long-horizon, multimodal, sparse reward problems uniquely suited for visualization and evaluating planning algorithms. The IDs for the 3 used environments are "maze2d-medium-v1", "maze2d-large-v1", "maze2d-umaze- v1". In each environment, one controls the acceleration of a robot to walk it towards a goal. The observation space is 4 dimensional, featuring 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 a reward of 1 if it is within a circle of radius 0.5 centered at the goal state, and 0 otherwise.The offline RL dataset for the maze environments consists of random walks in the maze. Specifically, the authors first designate all intersections and turn in the maze as waypoints and code an agent to navigate between waypoints with some randomization. As a result, the random walks are generated in a 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 a specific goal. There are a few choices adopted from our main baseline Diffuser [36] : we disregard the reward in the dataset and plan with goals only. We also evaluate a multi-goal variant of each environment (labeled as "multi" in Table 1 ), where the goal is randomized just like the starting position.We choose a long horizon robotic manipulation task as described in Section 4.4: Consider a tabletop with three slots where we can place objects. One places an apple at slot A 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, it can only move a fruit to an empty slot at a 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-markovian 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 a circular region occupied by any slot. Each circular region is designed to be about double the diameter of a fruit. To make sure the task requires visual feedback, we also randomize the location of a fruit inside the slot. We collected 150 expert demonstrations of a 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 a hand camera and one in the front capturing all three slots. Each demonstration also comes with 6 dof actions of the robot hand. During the data collection, since one 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 to save time.Each demonstration comprises 500 -600 frames and actions. We train Diffusion Forcing on the entire sequence. However, since adjacent frames are visually close, we pad and downsample the videos to 40 frames where each frame is bundled with 15 actions. We use a set of time series datasets accessible via GluonTS [2] , which are adopted from prior works like [71, 41, 55] . These datasets capture real-world data of high-dimensional dynamics like monetary exchange rates or the electricity grid. In Table 3 , we provide a summary of the features of 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 a set of short sequences that make up the test split. We construct a validation set of the same cardinality as the held-out test 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.Covariates Often, statistical models that approximate (E.1) benefit from manually curated features as additional input to the observations. A sequence of covariates C = {c t } T t=1 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, as 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 (F.1)The benefit obtained from covariates is highly dependent on the characteristics of both the dataset and the model used, as well as the feature engineering practices followed. 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 can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
safeguards: 
 Question: Does the paper describe 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)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 a best faith effort. 12. Licenses for existing assets Question: 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? 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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 the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.We implement zt = p θ (zt|zt-1, x k t t , kt) to be deterministic, with zt representing a distribution over beliefs rather than a sample from it. This allows training by backpropogating through the latent dynamics in Eq.(3.1).quantile levels at each time step t from 100 sampled trajectories. We then use the implementation in GluonTs[2] to compute the CRPS, which we report as CRPS sum in Table2. While we 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 1. **Writing clarity and organization.**  
   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.

2. **Clarity of figures.**  
   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.

3. **Minor issues and typos.**  
   - Line 97: missing a ")"
   - Line 139: "nevel" should be "level"
   - Line 186: "$x^{K/2_3}$" should be "$x^{K/2}_3$"
   - Line 178, 184, etc.: paragraph titles are inconsistently formatted
   - Line 522: missing a "("
questions: questions: 1. **Consistency between training and sampling algorithms.**  
   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:
   - 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}}$.
   - 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}$.

2. **Stabilizing auto-regressive generation.**  
   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?

**I promise to raise the score once all the weaknesses/questions are solved.**
limitations: limitations: Yes.
paper_summary: paper_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.

The authors also demonstrate additional potential generation tasks that can be done by diffusion-forcing models such as guided autoregressive sampling. 

Empirical results demonstrate that diffusion forcing performs well on video prediction and various planning tasks.
strength_weakness: strength_weakness: 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.

Theorem 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.

How 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?
questions: 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: 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.
paper_summary: paper_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.

The following is a more detailed summary.

### Overview of the method

Diffusion 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.

The architecture has two main components: 
an 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$;
a denoiser $\epsilon\_\theta(z\_t, x^{k\_t}\_t, k\_t)$, which is used to denoise $x^{k\_t}\_t$.

At 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$.

At 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.

The authors provide an ELBO interpretation for their loss function in the appendix.

### Features of Diffusion Forcing 

- The authors highlight the following features of DF:
- It supports classifier guidance, like ordinary diffusion;
- 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.
- It supports a flexible planning horizon, as tokens are denoised sequentially.
- 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.

When 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).

### Overview of experimental findings

- The authors evaluate Diffusion Forcing on video prediction, planning and robotics tasks. Their findings can be summarized as:
- 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.
- 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.
- 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).
In the appendix, the authors provide additional experiments on compositionality and time series prediction.
strength_weakness: strength_weakness: 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.

1. 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.

1. 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.

1. 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.

1. Minor formatting problems
    
    1. Line 186: $x^{K/2\_3}$ -> $x^{K/2}\_3$
        
    1. Table 1 caption: “Diffusion Forcingkeeps” -> “Diffusion Forcing keeps”; “Diffusion Forcingachieves” -> “Diffusion Forcing achieves”
        
    1. Line 495: “in full abstraction” -> “in full generality”
    Line 503: “likelihood for likelihood of all” -> “likelihood for all”
        
    1. Equation A.3: superscript $k\_2$ on the LHS should be $k\_s$
    Line 516: revise the bracketing of the expression involving $p\_\theta$.
        
    1. Line 522: “under uniform levels” -> “under uniformly sampled levels”
        
    1. Line 524: “in the sequel” -> “in the following section”
    Equation A.5: $s \leq T$ -> $1 \leq s \leq T$
        
    1. Line 592: specify range for $s$ on the first expectation
        
    1. Line 598: correct superscripts $t\_k$ to $k\_t$
        
    1. Line 608: revise bracketing of the numerator inside the $\ln$
        
    1. Line 616: In the last and penultimate lines, replace $\frac{\ln p(...)}{q(...)}$ by $\ln \frac{p(...)}{q(...)}$
        
    1. Line 628: “we” -> “we have”
        
    1. Line 631: correct superscript of $x\_t$ on the second line
        
    1. Line 634: expression with $p\_\theta$ broken between lines
        
    1. Line 635: capitalize Dirac
        
    1. Equation B.1: include \left[ \right] in the brackets
        
    1. Line 664: “we are” -> “we use”
questions: 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?

1. 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$?

1. 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?

1. 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: limitations: The authors address relevant limitations in Section 5 and social impacts in the checklist.
paper_summary: paper_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.
First, 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.
Second, 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.
Third, 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.
Diffusion 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).
The 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.
strength_weakness: strength_weakness: * 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
* 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
* 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
questions: 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

Q.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

Minor (no need to answer):
* typos: I could spot one typo in line 186: $[x_1^0, x_2^0, x_3^{K/2}]$
* 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


====== POST REBUTTAL MESSAGE ======

Thank you for the rebuttal. I have raised my score.
limitations: limitations: Yes

Comments: title: Thank you to the reviewers
Thank you to the reviewers
Thank you to the reviewers
Thank you to the reviewers
Thank you to the authors
Thank you for the rebuttal
comment: Dear Reviewer, thank you for your insightful feedback! We will carefully revise our paper to incorporate your suggestions and the latest results from our discussion.
Dear Reviewer, thank you for your insightful feedback! We will carefully revise our paper to incorporate your suggestions and the latest results from our discussion.
Dear Reviewer, thank you for your insightful feedback! We will carefully revise our paper to incorporate your suggestions and the latest results from our discussion.
Dear Reviewer, thank you for your insightful feedback! We will carefully revise our paper to incorporate your suggestions and the latest results from our discussion.
Thanks to the authors for taking my suggestions into consideration. I consider the achievements of this paper to be impressive and highly relevant for the line of work at the intersection of control and generative modeling. Hence, I maintain my position of strongly recommending acceptance.
Thank you for your response. I highly appreciate the contribution of your work, including novel frameworks, capabilities and extensive experiments, and I expect the writing could be refined in further versions. I raise my score to 6.
I thank the reviewer for their detailed response and for adding more experiments and ablations, which strengthen the paper. Therefore, I will maintain my positive rating.
Dear Authors,
thank you for the rebuttal, which answered all my questions clearly. Also, thank you for the additional experiments provided, as well as the thorough discussions with the other reviewers.

For all these reasons, I will raise my score.