Source-linked AI summary
Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion
Boyuan Chen, Diego Marti Monso, Yilun Du, Max Simchowitz, Russ Tedrake, Vincent Sitzmann
TL;DR
Existing sequence models trade off variable-length, causal generation against full-sequence guidance and continuous-signal stability. Diffusion Forcing trains causal models to denoise tokens with independent noise levels, combining these capabilities; experiments report stable long rollouts and gains in sequential decision-making, while the current implementation remains limited in scale and model complexity.
Problem
Next-token models lack objective-guided sampling and can diverge on long continuous-data rollouts, while full-sequence diffusion restricts generation to fixed sequences despite supporting guidance.
Method
Diffusion Forcing trains a causal next-token or next-few-token model to denoise sequences with arbitrary independent per-token noise levels.
Results
Diffusion Forcing combines variable-length generation with guidance and supports stable continuous-token rollouts beyond the training horizon, with performance gains in sequential decision-making tasks.
Takeaways & Limitations
The method provides a sequence model that can function as both a policy and a planner while supporting flexible-horizon generation and guidance.
Takeaways & Limitations
The current causal implementation uses an RNN, and scaling to higher-resolution video, complex distributions, internet-scale datasets, and tasks is left for future work.
Abstract
from arXiv · showhide
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
1 Introduction
Diffusion Forcing addresses limitations of next-token prediction and full-sequence diffusion by combining variable-length generation, long-horizon guidance, and causal denoising with independent per-token noise levels.
- Next-token models support variable-length generation, variable history conditioning, tree search, and online feedback control.
- Teacher forcing provides no mechanism to guide sampling toward an objective, and continuous-data rollouts can diverge beyond the training horizon.Small frame-prediction errors accumulate during autoregressive video generation.
- Full-sequence diffusion enables guidance and excels at continuous signals, but its non-causal architecture restricts generation to fixed-length sequences.It applies identical noise levels across concatenated tokens and limits guidance and subsequence generation.
- Diffusion Forcing assigns each token an independent noise level and denoises arbitrary per-token schedules through a shared next-token or next-few-token model.The method treats noising as partial masking: zero noise leaves tokens unmasked, while complete noise fully masks them.
- Causal Diffusion Forcing generates variable-length continuous sequences while supporting stable long-horizon generation and guidance toward high-reward outcomes.Its causal architecture denoises all tokens simultaneously with independent noise levels during training and progressively denoises Gaussian-noise frames during sampling.
- The paper evaluates Diffusion Forcing across video generation, model-based planning, visual imitation learning, and time-series prediction, alongside formal likelihood guarantees.The contributions include a decision-making framework using the model as both policy and planner.
2 Related Work and Preliminaries
This section situates Diffusion Forcing among Bayesian filtering, diffusion models, next-token predictors, and diffusion-based sequence models, emphasizing their distinct modeling and guidance choices.
- The paper views sequence modeling as Bayesian filtering along time and diffusion along an uncertainty axis.Time uses subscript t, while noise level uses superscript k.
- Bayesian Filtering: Bayesian filtering recursively estimates latent states from incoming observations using prior, observation, and posterior models.
- Diffusion Models: Diffusion models add Gaussian noise through a Markov chain and learn a reverse process that reconstructs data from noise.The noise variance is controlled by a schedule {β_k ∈ (0, 1)}^K.
- Guidance of Diffusion Models: Classifier guidance biases diffusion sampling by modifying the noise-prediction gradient with the gradient of a desired-label log probability.Other guidance energies include least-squares objectives comparing model outputs with desirable ground truth.
- Next-Token Prediction Models: Next-token models predict x_t+1 from past frames, then autoregressively reuse predictions to sample a full sequence.They use squared-error loss for continuous data or cross-entropy for discrete data and do not accept multi-step guidance.
- Diffusion Sequence Models: Diffusion sequence models support controllable text, short-video synthesis with sliding-window rollout, and planning in offline reinforcement learning.
3 Method
Diffusion Forcing treats independent per-token noising as partial masking, enabling causal sequence models to denoise variable-length sequences while supporting flexible guidance. Its training objective has a variational likelihood interpretation, and its sampling paradigm supports stable long-horizon generation and sequential decision-making.
- 3.1 Noising as partial masking: Diffusion Forcing assigns each token an independently varying noise level, unifying temporal masking and diffusion-based masking.Zero noise leaves a token unmasked, while complete noise removes its information.
- 3.2 Diffusion Forcing: different noise levels for different tokens: Causal Diffusion Forcing uses a causal architecture to denoise all sequence tokens with independent noise levels while retaining variable-length sampling.Its implementation uses an RNN whose latent state depends only on past noisy tokens.
- 3.2 Diffusion Forcing: different noise levels for different tokens: The training objective optimizes a reweighted ELBO on expected likelihoods and, under appropriate conditions, lower bounds likelihoods for all noise-level sequences simultaneously.This includes conditional distributions obtained by setting prior-token noise levels to either zero or complete noise.
- 3.3 New Capabilities in Sequence Generation: Diffusion Forcing stabilizes autoregressive generation of continuous sequences beyond the training length by updating latents from slightly noisy previous tokens.The method targets divergence that arises when high-dimensional continuous sequences such as video are rolled out autoregressively.
- 3.3 New Capabilities in Sequence Generation: Future-token guidance gradients can propagate backward through causal dependencies, guiding less-diffused past tokens without fully diffusing them.This provides long-horizon guidance while respecting causality and significantly outperforms guided full-sequence diffusion in the reported planning comparison.
- 3.4 Diffusion Forcing for Flexible Sequential Decision Making: In sequential decision-making, Diffusion Forcing selects actions sequentially while flexibly changing the lookahead horizon and guiding plans with dense or sparse future rewards.Monte Carlo Guidance evaluates expected reward over future outcomes, and removing it degrades performance while leaving Diffusion Forcing competitive.
4 Experiments
Diffusion Forcing is evaluated across video prediction, planning, imitation learning, robustness, compositional sampling, and time-series forecasting. It remains stable beyond the training horizon, outperforms baselines across six planning environments, and supports memory-dependent robotic manipulation.
- Video Prediction: Diffusion Forcing stably rolls out video far beyond its training horizon, unlike the compared teacher-forced and causal full-sequence diffusion baselines.The comparison uses the same exact RNN architecture for all models.
- Planning: DF outperforms Diffuser and all baselines across all 6 environments in the offline RL benchmark.The benchmark evaluates 2D maze environments with sparse reward.
- Planning: Removing Monte Carlo Guidance degrades performance, although Diffusion Forcing remains competitive without it.Monte Carlo Guidance is used to maximize expected future rewards during decision-making.
- Planning: Diffusion Forcing supports flexible-horizon planning and compositional sampling by denoising different time steps on different schedules.This allows the far future to remain more uncertain than the near future and permits composition of subsequences observed during training.
- Robotics: DF achieves 80% success in a memory-dependent fruit-swapping task, while diffusion policy without memory fails.The task randomizes initial fruit positions, so the policy must remember the initial configuration.
- Robotics and Forecasting: DF can remain robust to noisy or missing observations and generate realistic robot-task video from a single initial frame.The paper also reports competitiveness with prior diffusion and transformer-based methods on multivariate time-series forecasting.
5 Discussion
The discussion presents Diffusion Forcing as a flexible sequence-modeling framework with empirical, decision-making, and theoretical advantages. Its theoretical guarantee is an ELBO-based surrogate under suitable expressivity and modeling conditions.
- Limitations: The paper does not investigate Diffusion Forcing’s scaling behavior on internet-scale datasets and tasks.Higher-resolution video or more complex distributions may require large transformer models.
- Discussion: Diffusion Forcing combines next-token models’ flexibility with full-sequence diffusion models’ long-horizon guidance.The method is also used as both a policy and a planner in decision-making.
- Theory: The training methods optimize a reweighted Evidence Lower Bound on the average log-likelihood of the data.For Gaussian diffusion, the general result specializes to the form used by the method.
- Theory: With a fully expressive network, Diffusion Forcing optimizes a valid surrogate for the likelihood of all sequences of noise levels simultaneously.The bound is expressed through nonnegative terms that correspond, after reweighting, to the training objective.
A.1 Main Results
The theoretical analysis derives an ELBO for general Markov forward processes and generative models, then specializes it to Gaussian diffusion. Under deterministic latents and sufficient expressivity, the resulting objective provides a tight and separable likelihood surrogate.
- Main Bound: The main theorem gives an evidence lower bound on the expected log-likelihood of partially noised sequences.The formulation does not require an explicit form for the forward process or generative model before specialization.
- Main Bound: The ELBO is tight when the variational approximation is exact, under the stated deterministic-latent conditions.The proof identifies equality for deterministic latents and exact variational approximation.
- General ELBO and Gaussian Specialization: The analysis first states an ELBO for general Markov forward processes and generative models, then specializes it to Gaussian diffusion.The Gaussian specialization recovers the paper’s diffusion loss.
- Gaussian Specialization: The result applies to Gaussian diffusion’s x-prediction formulation and can also be stated for ϵ-prediction and v-prediction.The prediction formalisms are described as affinely related.
- Optimization Consequence: With expressive enough latents and pθ, the objective separates across time steps, making its weighting terms immaterial to optimization.This supports dropping or rescaling weighting terms in the training objective.
A.1.2 Capturing all subsequences
Diffusion Forcing’s objective provides variational guarantees for noise-level sequences, including fixed schedules, and under expressivity can serve as a surrogate for likelihood maximization across possible noise sequences.
- Theorem A.3 extends the variational lower-bound result to arbitrary distributions over temporally organized noise levels.The noise sequence may be temporally correlated.
- The lower-bound equality conditions include deterministic latent transitions and a sharp variational approximation.
- For fixed noise sequences, the lower-bound terms are, up to reweighting, a subset of those optimized by the training objective.This applies when the noise distribution is concentrated on one sequence k1, . . . , kT.
- A fully expressive network can optimize all relevant terms simultaneously, making the objective a surrogate for maximizing likelihood over all possible noise sequences.
- Diffusion Forcing’s independent noise levels support causal uncertainty, autoregressive stabilization, and sampling-time guidance mechanisms unavailable to full-sequence diffusion.
B.8 Connection to other sequence training schemes
Diffusion Forcing unifies sequence training schemes by treating noise as masking, then uses independent noise levels to interpolate between autoregressive prediction, partial conditioning, and diffusion-based generation.
- Noise can be interpreted as masking: zero noise leaves a token unmasked, while complete noise fully masks it.
- Autoregressive: Diffusion Forcing can reproduce autoregressive sampling by fully diffusing future tokens before adding them to history.This sampling mode is applicable when the next action should be diffused quickly, such as imitation learning.
- Causal uncertainty: Independent noise levels enable causal uncertainty by keeping farther-future tokens more uncertain than nearer-future tokens.
- Stabilization: Diffusion Forcing stabilizes autoregressive rollout by treating generated images as noisy observations rather than clean ground truth.
- Implementations: The method extends across architectures and domains, with U-Net and GRU components for video and residual MLPs for non-spatial sequences.
D.5 Implementation Details of Sampling with Guidance
Guided Diffusion Forcing sampling uses a flexible noise schedule to denoise tokens sequentially while supporting variable horizons and reward-based planning.
- When a token must remain at one noise level, the implementation generally resamples through backward and forward diffusion for Monte Carlo guidance.For noise levels 0 and K outside Monte Carlo guidance, copying is used in the described experiments.
- Guidance can use rewards defined over future steps, including dense per-time-step trajectory rewards.
- Sampling initializes all tokens as white noise, then denoises them row-by-row and left-to-right according to a scheduling matrix K.
- Pyramid scheduling keeps far-future tokens at higher noise levels while progressively denoising nearer tokens.
D.8 Metrics for Maze Planning
The maze-planning evaluation reports episode reward but argues that this metric can misrepresent planning quality because the environment and dataset reward lingering near the goal.
- Episode reward is reported for Diffusion Forcing across maze-planning environments.
- Episode reward may favor slow plans that remain near the goal rather than agents that reach it quickly.
- The paper recommends time to first reach the goal as an alternative metric and reports that Diffusion Forcing excels at it.
- Time-series comparison: The time-series material defines CRPSsum as the forecasting metric, with lower values preferred, across six datasets.
E.2 Additional results in compositional generation
Diffusion Forcing can model subsequence distributions and compose sub-trajectories, with controllable memory determining whether generated plans preserve or recombine trajectory structure.
- The model can sample distributions over trajectory subsets and compose the resulting sub-trajectories into new trajectories.
- Full memory preserves the cross-shaped trajectory distribution, whereas memoryless short-horizon MPC can stitch sub-trajectories into a V-shaped plan.
- Diffusion Forcing models the joint distribution of all subsequences and can sample either full-horizon trajectories or Markovian dynamics.
E.3 Additional results in video prediction (wo/ cherry picking)
Without sliding windows, Diffusion Forcing generates video well beyond its training horizon on Minecraft and DMLab, while also producing consistent within-horizon videos without cherry-picking.
- Diffusion Forcing rolls out 180-frame videos without sliding windows after training on 72 Minecraft frames and 36 DMLab frames.The paper reports that this is 2x–5x the training length.
- The DMLab model trained on 36 frames generates 180-frame rollouts, with visualized samples described as non-cherry-picked.
- The Minecraft model trained on 72 frames generates 180-frame rollouts and has been tested for sequences as long as 2000 frames.
- Within the maximum training length, Diffusion Forcing generates consistent videos, and the reported results are not cherry-picked.
E.4 Additional results in planning
Additional planning results show that Diffusion Forcing supports flexible-horizon plans and can operate under visual observation corruption without provoking wild robot behavior.
- Diffusion Forcing generates shorter-horizon plans because its horizon is flexible, supporting causal planning and model predictive control.
- In the robot experiment, the target bag is introduced zero-shot as a visual distractor, and the typical failure is ignoring randomized object-location cues rather than wild actions.
F.1 Dataset for video diffusion
The paper evaluates video diffusion and related sequence modeling components on video, robotic, offline-RL, and time-series data, with dataset-specific preprocessing and covariates.
- Video diffusion: Video diffusion experiments use Minecraft and DMLab navigation datasets, with videos trimmed into fixed-length subsequences.Minecraft uses 128-pixel videos and 72-frame subsequences; DMLab uses 64-pixel videos and 48-frame subsequences.
- Video diffusion: The Minecraft dataset contains first-person random-walk videos in the swamp biome, while DMLab contains random walks in a 3D maze.
- Video diffusion: A dataset issue causes Minecraft agents to become stuck at tall obstacles, producing long runs of grey granite or brown dirt patterns.The authors attribute this problem to the dataset itself.
- Planning data: Offline-RL planning uses D4RL maze environments with state, action, and reward trajectories, while the reported maze setup controls two-dimensional robot acceleration.
- Time-series data: Time-series benchmarks use GluonTS datasets spanning different domains, dimensionalities, sampling frequencies, sequence lengths, and prediction lengths.
- Time-series data: Time-series models incorporate lagged inputs, learned embeddings, and handcrafted temporal covariates that are known throughout the inference interval.