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

Independent noise vs AR-Diffusion

We thank the reviewer for highlighting the need for a more explicit discussion relative to AR Diffusion.

We would first like to clarify that the stabilization discussion in Appendix B.2 is orthogonal to AR-diffusion. The key to AR-Diffusion is training and sampling future tokens with linearly growing noise similar to the pyramid sampling of Diffusion Forcing. On the other hand, stabilization is a technique of assigning non-zero noise levels to context tokens (history) instead of future tokens. In fact, our stabilization experiment uses autoregressive diffusion, not pyramid noise like AR-Diffusion.

Furthermore, AR-diffusion style pyramid sampling is just one of the possible sampling schemes supported by diffusion forcing after training once. In online decision-making like our robot experiment, we can use autoregressive diffusion to achieve the best speed, use pyramid sampling to achieve better guidance, or many other schemes. In Figure 1 of rebuttal pdf, we further show that non-causal diffusion forcing can do frame interpolation by assigning noise level 0 to keyframes at arbitrary time steps, while AR-diffusion’s non-causal variant, Rolling Diffusion lacks this flexibility.

Finally, we found that Diffusion Forcing beats AR-Diffusion even using AR-Diffusion’s own sampling scheme as we show and analyze in our response to the next question. This shows that independent noise, just like standard techniques such as image pre-training, plays an important role in enhancing the underlying representation.

Experimental comparison to AR-Diffusion

We are happy to report that we have re-implemented and benchmarked with AR-Diffusion for video generation, as well as a concurrent work named Rolling Diffusion[1] which is basically non-causal AR-diffusion. As shown in Figure 4 of the rebuttal pdf, both AR-diffusion and Rolling-diffusion performed worse in FVD metric compared to diffusion forcing. Our main insight is twofold: 1. while both use linearly growing noise levels along the time axis, both seem to be sensitive with respect to the slope of this linear growth. We spent a fair amount of time tuning the two baselines, but since one has to use the same slope during both training and testing, it requires us to re-train many times to find one slope that’s reasonable for our data. 2. We observe that training with linearly growing noise has too much redundancy. The trajectories tend to have a higher signal-to-noise ratio which made the task too easy. We tried the boundary condition in Rolling Diffusion[1] and multiple tunings to mitigate this, but still failed to match Diffusion Forcing. In fact, table 2 of Rolling Diffusion [1] shows that linearly growing noise isn’t necessarily better than normal diffusion in a large video dataset, aligning with our observation. In addition, a fixed slope at training time creates more inconvenience for DDIM sampling as well due to rounding, creating practical pain for any user who wants to tune DDIM steps.

Reference Appendix B.6 in the main paper

Thank you for your suggestion - we will definitely reference Appendix B.6 in the main paper, and attempt to move it to the main paper entirely, space after incorporating all reviewer feedback permitting!

Typo

Thank you for catching the typo, we have fixed it!

Grammar and spelling in appendix

We have done a full pass over the appendix and polished grammar and spelling, thanks for pointing this out!

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.

We thank the reviewer for their in-depth review - we are particularly happy to be able to address some of the limitations of Diffuser that the reviewer had to contend with themselves in the past!

Clarification on Classifier Guidance Term

Sorry about the confusion. The reviewer is exactly correct in that future tokens are a differentiable function of past tokens via the posterior model, and that this enables us to perform guidance. We adopted the notation for simplicity and to avoid clutter/accommodate space restrictions, but we will absolutely clarify the meaning of that gradient in the text of the “Guidance” section in our subsequent revision. On the other hand, in our transformer implementation of diffusion forcing, any token can attend to the entire history, allowing us to reproduce the guided planning results as well.

The naming of Monte Carlo Tree Guidance is misleading

Thank you for your comment. Amusingly, we had an internal discussion among the authors on exactly this topic - we report that one of the parties feels highly vindicated by your comment. We will follow your advice and rename the method into Monte Carlo Guidance!

Consistency of action and state in high-dim control environments

We tried CDF on one high-dimensional control environment (hopper-medium-expert) without guidance to exclude confounding factors. Unfortunately, the observation still diverges from the ground truth eventually without feedback as no model is perfect even in deterministic environments. However, it’s clearly better than Diffuser just like our visualizations in the 2d maze experiment. In the generic RL setting, the environment can be stochastic, and therefore open-loop execution of a plan is expected to diverge from the plan even with an oracle model. Therefore, we believe a thorough evaluation of dynamics better falls in the result in Section 4.5, which shows CDF is competitive for high-dimensional time-series data.

Similarity of compositional generation environment with Janner et al. (2022).

Thank you for highlighting this relationship - we will absolutely add a reference to Janner et al. (2022) to Section 4.3!

Minor formatting problems

We apologize for these mistakes and have fixed all of them for the camera-ready!

Does the flexible horizon property derive from RNN

Flexible horizon planning doesn’t depend on RNN. We’ve reimplemented diffusion forcing with both transformer and 3D-Unet (with temporal attention), showing flexible horizon planning is fully reproducible with attention instead of RNN. For causal attention, this is obvious. For non-causal variants, diffusion forcing allows us to mark padded future entries as pure noise to achieve a flexible horizon. The only ability that depends on RNN is infinite rolling without sliding window because state space models have more invariance, although transformer can roll out longer with sliding window.

Could combining Diffuser with RNNs also lead flexible planning horizon?

Diffuser’s maze planning result is strictly fixed horizon. While Diffuser made an argument about flexible horizon via 1D convolution, its maze planning results critically depend on the replacement technique we mentioned in Appendix B.4. That is, Diffuser fixed the last token to be at the goal throughout the diffusion process to generate a fixed length plan of a carefully chosen length. If the objective is to reach the goal asap, this objective would fail, as one doesn’t know which token shall be replaced as the goal. We implement this on Diffuser and cannot reproduce their numbers. We suspect that the mere reason why Diffuser got high scores is that their fixed length objective coincidentally encourages very slow plans that made the agent stay near the goal for longer, an out-of-distribution behavior that doesn’t exist in the dataset.

Source of causality in DF vs. Diffuser: architecture or observation model

It’s most likely the architecture itself. Our transformer implementation is direct empirical evidence that a separate observation model is the key. It also follows theoretically: consider an alternative formulation of $x_t=[o_t, a_t, r_t]$ like diffuser instead of $x_t=[o_{t+1}, a_t, r_t]$ used in our paper. Even though the hidden state and observation model are untouched, causality is mathematically broken - at time step t, one wants to diffuse action $a_t$ given current $o_t$, but the process of diffusing $a_t$ samples another $o_t$, which might be different from the existing observation. Such weird inconsistency suggests the temporal structure is critical.

scale up Diffusion Forcing for transformers

Between the initial draft and rebuttal, we’ve already reimplemented diffusion forcing with architectures like Transformers or 3D-Unet and reproduced results. We also find Diffusion Forcing to work well in latent diffusion settings for higher-resolution videos (Figure 1 of rebuttal pdf).

Can Diffusion Forcing, like the Diffuser, be run with a variable planning budget, through warm-starting or otherwise?

Unlike Diffuser[1], diffusion forcing implements DDIM sampling which already achieves good speed without Diffuser’s warm-starting technique. Therefore, we provide an alternative planning budget analysis here in terms of frequency of replanning. In Figure 2 of rebuttal pdf, we show the performance when we replan 1x, 0.5x, 0.25x, 0.125x of episode length and for 50 steps. We found that diffusion forcing’s performance worsens as we replan less frequently. We’d like to clarify that this is partially due to the dataset itself - the maze dataset never contains the behavior of staying at the goal once reached, and the agent always walks away from the goal even after reaching it. Therefore, diffusion forcing is supposed to follow this suboptimal behavior if there is less replanning.

[1] Michael Janner, Planning with Diffusion for Flexible Behavior Synthesis, 2022

Thank you for your helpful feedback. We respond to your comments below:

Can we train with arbitrary noise schedules and find a good schedule at evaluation time?

Yes - at the core of Diffusion Forcing lies exactly the idea that by training with arbitrary, random noise schedules, we can experiment with arbitrary noise schedules at evaluation time to find the optimal schedule for each task! In practice, we found this to be very useful as users can try different noise schemes at sampling time for different applications after training once while the slightest tuning would require retraining for various baselines. The proof of the ELBO extends to arbitrary noise schedules as well.

Is it necessary to use different noise schedules in different tasks?

In short, yes. For example, we found that pyramid sampling with causal diffusion forcing (CDF) is best for MCTG planning while autoregressive diffusion offers the best consistency for causal video. When conditioning on ground truth context frames, CDF works best without stabilization; In contrast, when context is previously generated, CDF works best with stabilization as shown in Figure 3 of rebuttal pdf. Furthermore, even if we commit to one specific schedule such as pyramid sampling, it is critical to tune the rate with which uncertainty grows with time via the noise schedule under pyramid sampling. DF allows us to do this easily after training once, while AR-Diffusion [2] (suggested by 35x2 and 2PA2) makes this difficult as this rate is a training hyperparameter and fixed at test time.

We will use the extra page from the camera ready to include this discussion of noise schedules!

Reduce the training complexity

Surprisingly, we found that random, independent noise schedules at training time do not significantly add to the training complexity. Due to limited space, please refer to the experiment in general response’s “Improved Model Performance” section which shows the overhead of independent noise is insignificant when considering the overall training compute.

We speculate that random, independent noise levels at training time can be seen as a form of data augmentation[3], leading to a bigger “bang for the buck” of training data. Further, they might in expectation improve the quality of gradient flow backward in time.

Balance between performance and computation time

The suggested sampling scheme reduces to auto-regressive diffusion with ddim_steps of 1, which mainly captures the low-frequency part of the data only. If we have more compute budget, we can sample auto-regressively with >1 ddim_steps and get much higher quality samples. As shown in Figure 2 of rebuttal pdf, the main video metric fvd decreases as the sampling time increases. A similar trend follows for maze planning.

reweighted ELBO can be poorly balanced

Indeed, the formal version of Theorem 3.1 in the Appendix, Theorem A.1 demonstrates the weighting explicitly: that the ELBO on the expected log-likehoods, over the randomness of the noise schedule, corresponds to an objective which reweights the contribution of terms from denoising steps k by their step number k. Thus, the reweighting factor is by at most a factor of the total number of denoising steps, K. Note that this weighting by k is independent of the noise schedule (that is, of the variance of the noise at each step k), and depends only on the position k of the noising step (as well as our decision to select noise levels uniformly and independently). We will add this as a remark in the revision.

Though we could choose to up weight gradients from earlier steps by more (as Thm A.1 makes the weights explicit), not reweighting the gradients according to the ELBO corresponds to effectively weighting gradients at smaller denoising steps k more, which may be desirable as these steps capture the “higher resolution” features of the model. Lastly, as noted in the remarks following Corollary A.2, our ELBO derivation is sufficiently tight that, given a sufficiently representative neural network, the optima of all (non-trivial) reweightings of our ELBO are also optima of the desired expected likelihood. Hence, the weighting does not meaningfully bias the training objective. We do acknowledge that the weighting may be salient for the dynamics of optimization, but empirically, we find that our decision not to compensate for the step-wise weighting of the ELBO leads to reliable training. We will clarify the remarks following Corollary A.2 to emphasize this point further.

Main result is RNN

We’ve reimplemented diffusion forcing with both transformer and 3D-Unet architectures after the draft submission. Please see the general response for details due to limited space.

Apply diffusion forcing to non-autoregressive generation

In the maze planning experiments, we already show how guidance can be used to achieve certain future outcomes. In fact, reconstruction guidance discussed in Appendix B.4 has been widely used to achieve consistency with context through guidance. Alternatively, one simply uses a non-causal neural network architecture, like our implementation of 3D-Unet, and then - as discussed in the paper - varies the per-token noise level at training time. One can further mask out padded entries by marking them as full noise. We are actively pursuing this direction in follow-up work: Initial results in Figure 1 of rebuttal pdf showcase frame interpolation result of Diffusion Forcing without the expensive reconstruction guidance. However, a detailed discussion of this setting would exceed the capacity of the present paper, though we will discuss this direction in the discussion section.

[1] Michael Janner, Planning with Diffusion for Flexible Behavior Synthesis, 2022

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

[3] Diederik P. Kingma, “Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation”, Neurips 2023

Thank you for your helpful feedback! We respond to your comments below:

Writing clarity and style & typos

Sorry about the confusion. It’s true that we opted out of an independent “Method” section to introduce the intuitions first. With the extra 1 page for the camera-ready version, we promise to structure the paragraphs more clearly. In particular, following your advice, we will make (1) one “methods” section that discusses DF in one subsection as well as its intuition obtained from Bayesian Filtering in another subsection, (2) one “Capabilities” section discussing novel capabilities, and (3) one ” Results” section. Thank you for your time on the list of typos, we will make sure to fix them in camera-ready.

Figure clarity

Thank you for the constructive advice. We will make sure to caption the figures with more details with the 1 extra page limit of the camera-ready version so readers can easily understand them without referring to the main text. For the sampling scheme in Figure 2 (Assuming the figure index in your comment is a typo), we were trying to illustrate a scheme that illustrates both the transformer and RNN as mentioned in the caption. Therefore, we will likely instead remove latent from the training subfigure to make sure it’s consistent with the sampling figure and add a new figure in the appendix to illustrate RNN variant with latent state. We provide a preview of modification in Figure 5 of the rebuttal pdf.

Consistency between training and sampling algorithms.

Thank you for pointing this out. We will definitely fix the algorithm box carefully: 1. $z_{t-1}$ in line 8 of the training algorithm should be changed to $z_{t}$ 2. In line 7 of the sampling algorithm, the $k$ should indeed be corrected to be $\mathcal{K}_{m+1,t}$. We will fix line 5 to reflect this. 3. Outside the two errors you pointed out, some $\alpha$ are accidentally indexed via $t$ or $k$, not $k_t$.

Details about stabilization and ablating $k$

In our implementation, we choose $k=20$ for $K=1000$ but this value is largely dependent on the difficulty of the task. In Figure 3(a) of the rebuttal pdf, we present the result of the requested ablation of this stabilization value. We choose a slightly harder setting for this ablation, reducing the frame rate by half and using lower numbers of DDIM steps to make compounding error more obvious. One can observe that the FVD metric improves as we gradually increase the stabilization level from 0, reaching optimal around 100 before FVD starts to rise. We found stabilization to be extremely important for datasets with high stochastic like the BAIR robot pushing dataset shown in Figure 3(b) of the rebuttal pdf. We cannot achieve stable autoregressive rollout even with very high DDIM steps. Stabilization is a fundamental technique to get this dataset to work.