Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

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.

Strengths

* This work presents substantial improvement over the literature on the joint application of diffusion models and autoregressive models * The proposed methodology is technically sound, and well supported by intuition, formal proofs and a wide range of experiments * The experimental section expands over the literature by focusing on several domains including long-range video generation, planning, compositional generation and multivariate time series forecasting

Weaknesses

* The intuition of the effects of noising on long-horizon generation (appendix B.2) is very similar to the ideas described in a related work AR-Diffusion [62]. This does not highlight the contribution of *independent* noise levels per sequence index * 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

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

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.

Strengths

1. The authors propose an original and performant method combining strengths of diffusion (steerability, robustness to noise, high-quality gradual sample generation) and auto-regressive sequence modelling (flexible horizons, temporal causality, memory in the case of RNNs). 1. In addition, the authors provide a theoretical justification of their loss function in terms of an evidence lower bound (ELBO). 1. The paper is written clearly, providing a clear motivation for the authors’ approach, contextualizing DF relative to existing work (especially Diffuser, AR-Diffusion and Diffusion Policy), and highlighting the main contributions of the method conceptually and experimentally. 1. Trajectory inconsistency is a major limitation of Diffuser, which I have contended with in my own research. Mitigating this limitation is an important enabler of bringing the strengths of diffusion to bear in sequential decision making. 1. Monte Carlo Tree Guidance can be seen as maximizing an empirical estimate of the expected future reward. From a policy optimization perspective, this seems more principled than doing gradient ascent on the realized cumulative reward of a given trajectory, as is done in full-sequence diffusion (e.g. Diffuser). As the authors explain in Appendix B.3, this technique relies on the architecture of DF to be effective. 1. The results on video prediction, available in an anonymized project website provided in the abstract, are particularly impressive in terms of stability and 3D consistency. This, together with results on planning and robotics, indicates DF might contribute to advances in diffusion world models; a research area of established relevance that has received significant attention recently.

Weaknesses

1. Clarification on classifier guidance term $\nabla\_x \log c (x^{\textrm{new}}\_{1:H})$: If this term is to be understood as the gradient of $x \mapsto \log c(x)$ evaluated at $x^{\textrm{new}}\_{1:H}$, then the gradients of $c$ on future tokens would not flow to previous tokens, as the inputs $x$ are “frozen” before being fed into $\log c$. It seems that what the authors mean to say is that future tokens are treated as a differentiable function of past tokens when computing the gradients. It would strengthen the exposition if the authors either clarify this point in the paper, or update the notation to avoid confusion, as the current notation might lead the reader to believe that the gradients from future tokens do not flow into past ones. 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

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.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors address relevant limitations in Section 5 and social impacts in the checklist.

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.

Strengths

This paper proposes an interesting combination of autoregressive models and diffusion models and demonstrates that the combination of both outperforms both individual models in terms of performance. Further, the diffusion-forcing paradigm offers many more applications that are otherwise impossible. For example, while doing variable-length generation, the model can leverage classifier-based/-free conditions. This provides much better flexibility to inference-demanding tasks such as planning and control. The authors propose a training objective of diffusion forcing models based on noise prediction. The objective is proved to be a reweighted version of the evidence lower bound and thus is sound. Diffusion forcing achieves much better performance compared to autoregressive models and diffusion models in long-horizon generation tasks.

Weaknesses

A more detailed discussion of the noise schedule is desired to better understand the effectiveness of diffusion forcing. Is it necessary to use different noise schedules in different tasks to achieve good performance? Further, can we train the model with various/arbitrary noise schedules and at evaluation time find a good schedule? If any of these is possible it will greatly reduce the training complexity and extend diffusion forcing to more applications. 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

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

1. The proposed Diffusion Forcing method is general and flexible, making it applicable to various tasks. 2. The paper provides a comprehensive discussion on the capabilities of Diffusion Forcing. 3. The experiments are well-designed and effectively demonstrate the proposed method's effectiveness.

Weaknesses

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

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

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

Thank you for the rebuttal

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.

Thank you to the authors

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.