[Paper Review] Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

Boyuan Chen, Diego Martí Monsó+4 more|arXiv (Cornell University)|Jul 1, 2024

Forensic and Genetic ResearchBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

This paper introduces Diffusion Forcing, a novel training paradigm that unifies next-token prediction with full-sequence diffusion by assigning independent noise levels to each token, enabling stable, variable-length generation and effective guidance in long-horizon sequence modeling. The method achieves state-of-the-art performance in video generation, planning, and time series forecasting by combining causal modeling with flexible, per-token denoising schedules.

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

Motivation & Objective

• To address the instability of autoregressive models in generating long sequences of continuous data, such as video, beyond the training horizon.
• To overcome the limitations of full-sequence diffusion models, which lack causality and cannot support variable-length generation or subsequence conditioning.
• To enable effective sequence guidance and planning by combining the flexibility of next-token prediction with the global optimization capabilities of diffusion models.
• To formalize a training objective that maximizes a variational lower bound on the likelihood of all sub-sequences from the true joint distribution.
• To demonstrate the method's effectiveness across diverse domains, including video generation, model-based planning, and time series forecasting.

Proposed method

• Diffusion Forcing trains a single causal next-token prediction model to denoise sequences where each token is independently noised with a unique noise level, enabling per-token denoising schedules.
• The model is trained to reconstruct any subset of tokens from a noisy sequence, effectively learning to 'unmask' tokens based on their individual noise levels.
• During inference, the model performs progressive denoising of a full sequence with arbitrary, per-token noise levels, enabling variable-length generation from the immediate next token onward.
• A novel sampling scheme, Monte Carlo Tree Guidance (MCTG), leverages the causal architecture and variable-horizon capabilities to guide sampling toward high-reward trajectories in planning tasks.
• The method uses a shared causal transformer architecture that conditions future token predictions on past tokens, preserving causality while enabling full-sequence denoising.
• The training objective optimizes a variational lower bound on the likelihood of all sub-sequences observed in the training data, ensuring consistency with the joint distribution.

Figure 2 : Method Overview. Diffusion Forcing trains causal sequence neural networks (such as an RNN or a masked transformer) to denoise flexible-length sequences where each frame of the sequence can have a different noise level. In contrast, next-token prediction models, common in language modeling

Experimental results

Research questions

• RQ1 Can a single model combine the variable-length generation capability of next-token prediction with the long-horizon guidance of full-sequence diffusion?
• RQ2 Does per-token noise scheduling enable stable generation of continuous sequences like video beyond the training horizon?
• RQ3 Can the causal architecture of next-token models be leveraged within a diffusion framework to enable new guidance schemes such as Monte Carlo Tree Guidance?
• RQ4 Is the proposed training objective theoretically justified by optimizing a lower bound on sub-sequence likelihoods?
• RQ5 How does Diffusion Forcing perform across diverse sequence modeling tasks, including video, planning, and time series forecasting?

Key findings

• Diffusion Forcing enables stable, long-horizon generation of continuous sequences such as video, with sequences extending thousands of tokens beyond the training horizon, where autoregressive baselines diverge.
• The method achieves state-of-the-art performance in video generation and time series forecasting, with CRPS scores significantly lower than baseline models across multiple GluonTS datasets.
• Monte Carlo Tree Guidance (MCTG) using Diffusion Forcing leads to marked improvements in decision-making and planning tasks, outperforming non-causal full-sequence diffusion baselines.
• The model generalizes compositionally to new trajectories and supports efficient tree search and online feedback control, combining the strengths of both autoregressive and diffusion models.
• Empirical results show that the training objective effectively maximizes a variational lower bound on sub-sequence likelihoods, validating the theoretical foundation of the approach.
• In visual imitation learning and model-based planning, Diffusion Forcing enables high-reward trajectory generation with improved sample efficiency and stability.