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

In plain words

Usually a model that writes a sequence one step at a time (like a language model predicting the next word) and a model that starts from noise and cleans up a whole sequence at once (like image diffusion) are trained in different ways. This paper trains one model that can do both. Every item in the sequence gets its own amount of noise during training. Some items are almost clean, some are mostly noise, and the model learns to clean up whichever ones it is given.

Because of this, generation can be arranged as a schedule. The near future can be cleaned up before the far future, so the model keeps the far future uncertain. The past can stay clean while new items are generated, which lets sequences like video keep going well past the lengths the model was trained on. The same flexibility lets you steer the sampling toward a goal, for example a robot reaching a point in a maze, and plan with a window that can be short or long.

TL;DR

Diffusion Forcing trains a denoiser in which each token of a sequence has an independent noise level, and it uses a causal recurrent model so each token only sees the past. Sampling follows a per-token noise schedule, which enables stable autoregressive video rollouts, guidance over future tokens, and flexible-horizon planning. The paper reports the best average reward on its Maze2D planning benchmarks against the offline-RL and diffusion-planning baselines it compares to, and competitive time-series forecasting. All numbers are author-reported.

Problem

Next-token prediction models are good at variable-length generation but can diverge when they roll out high-dimensional continuous sequences such as video beyond their training horizon §4.1. Full-sequence diffusion models generate consistent samples and can be guided toward desired trajectories, but they fix the sequence length and must denoise all tokens together, so they cannot take actions or update a plan step by step §2 §3.4. The paper argues both properties are wanted and proposes a single model that has both §1 §3.1.

Method

Notation. A sequence is x_{1:T}. Each token x_t at noise level k_t is written x_t^{k_t}. Level 0 is clean and level K is pure white noise §3.1. The paper uses K = 1000 diffusion steps app. D.

Noising as partial masking §3.1. Next-token training with teacher forcing masks a token and predicts it from the past. Full-sequence diffusion can be read as partial masking of every token at once. Diffusion Forcing combines the two views by giving each token its own partial mask, which is its own noise level.

Recurrent denoiser §3.2. A recurrent network with latent z_t holds a summary of the past. At each step the latent is updated from the previous latent, the noisy token and that token's level, z_t ~ p(z_t | z_{t-1}, x_t^{k_t}, k_t). The paper implements this update as deterministic, so z_t is a belief state rather than a sample. An observation head predicts the clean token from z_t. The paper's training is a standard conditional diffusion loss on the noise epsilon_t, applied to each token and summed over the sequence §3.2.

Training §3.2 alg. training. Sample a clean sequence. Sample every token's noise level independently and uniformly from the K levels. Noise each token at its level. Run the recurrent unit over the sequence and sum a noise-prediction loss over all tokens. The paper's informal theorem says this objective optimises a reweighted evidence lower bound on the log-likelihood of the noised sequences. Under suitable conditions it also bounds the likelihood of every subsequence of noise levels at once §3.2 app. A.

# Training step  [§3.2] [alg. training]
# clean_tokens: (T, B, D)
def training_step(clean_tokens):
    # Sample an independent noise level for every token  [§3.2]
    noise_levels = sample_uniform_levels(K, shape=(T, B))  # (T, B)
    # Draw Gaussian noise for every token  [§3.2]
    eps = normal(0, 1, shape=clean_tokens.shape)  # (T, B, D)
    # Forward diffusion, each token at its own level  [§3.2]
    noisy_tokens = forward_diffuse(clean_tokens, noise_levels, eps)
    latent = initial_latent(B)  # z_0  [§3.2]
    loss = 0.0
    for t in range(T):
        # Predict the noise from the previous latent and the noisy token [§3.2]
        eps_hat = observe(latent, noisy_tokens[t], noise_levels[t])
        loss = loss + squared_error(eps[t], eps_hat)
        # Update the latent with the noisy token and its level  [§3.2]
        latent = dynamics(latent, noisy_tokens[t], noise_levels[t])
    return loss

Sampling with a schedule §3.2 alg. sampling. Sampling is defined by a scheduling matrix K with M rows and T columns. Column t is the time step and row m gives the target noise level for that token. Start with every token at level K. Walk down the rows from the top. In each row, step through the tokens from left to right, update each token's latent, and denoise the token to its target level in K. After the last row, every token is clean. The paper uses standard DDIM-style updates at each level §3.2 app. D.

# Sampling with a noise schedule  [§3.2] [alg. sampling]
# schedule: (M, T) integer grid of noise levels, row 0 is clean
def sample_sequence(schedule, guidance_cost=None):
    # Start every token from pure noise, at level K  [§3.2]
    tokens = normal(0, 1, shape=(T, B, D))  # (T, B, D)
    latent = initial_latent(B)  # z_0  [§3.2]
    for row in range(M - 1, -1, -1):
        for t in range(T):
            # Update the latent using the level from the row above  [§3.2]
            new_latent = dynamics(
                latent, tokens[t], schedule[row + 1, t]
            )
            # Denoise token t to the target level for this row  [§3.2]
            tokens[t] = denoise_step(
                tokens[t], new_latent, schedule[row, t]
            )
            latent = new_latent
        if guidance_cost is not None:
            # Add guidance from the partially denoised trajectory  [§3.3]
            tokens = add_guidance(tokens, guidance_cost)
    return tokens

New sampling capabilities §3.3.

• Stabilised autoregressive generation. Condition on past tokens that are slightly noisy (a small level 0 < k << K) rather than fully clean. The paper reports this keeps long rollouts stable, and the paper's appendix explains why §3.3 app. B.
• Zig-zag schedule. Denoise the first token fully, the second partly, the third not at all, and so on. Near tokens are then more certain than far ones §3.3.
• Long-horizon guidance. Guidance gradients on future tokens propagate back to past tokens, because the model is causal and the past is only partly noised §3.3 app. D.

Planning and decision making §3.4. Tokens are triplets of action, reward and next observation. At each step the model plans a window of future tokens conditioned on the latent of the past, executes the first action, and updates the latent with the new token §3.4. The window can be shortened to act as a policy, or lengthened to plan, without retraining §3.4. Guidance can use dense rewards, rewards over a future window, or a sparse goal term §3.4 app. D.

Monte Carlo Guidance (MCG) §3.4. Instead of guiding with one sampled future, draw several samples of the future tokens and average their guidance gradients. The paper argues this works best with a schedule that keeps far tokens noisier §3.4 app. B. The paper does not give the number of samples used.

# Monte Carlo Guidance for one token  [§3.4]
# future_samples: list of rollouts of the tokens after t
def monte_carlo_guidance(token_t, future_samples, cost):
    # Guidance gradient for each sampled future  [§3.4]
    grads = [grad_wrt(token_t, cost(f)) for f in future_samples]
    # Average the gradients over the samples  [§3.4]
    return mean(grads)

Hyperparameters

Values are from the paper unless the Where column names the code. "not stated" means neither the paper nor the code (as read here) gives the value.

Results

All author-reported.

Planning, D4RL Maze2D and multi-goal Multi2D §4.2. Average episode reward. Diffuser* is run with the hand-crafted PD controller the paper uses for it. "Diffuser w/ diffused action" executes the actions the diffusion model generates. The paper notes that the reward metric may favour slow plans that stay near the goal app. D, so compare these numbers with care.

Source: the planning table in §4.2 (typeset inside a figure environment in the source; its number in the compiled paper was not checked). The paper says Diffusion Forcing outperforms Diffuser and all baselines on all six environments §4.2.

Video §4.1. Qualitative only, with no table of metrics in the main text. The paper reports that its rollouts stay stable up to 1000 frames, far beyond the training horizon, while teacher-forcing and full-sequence baselines diverge §4.1 (video figure in the source; numbering not checked).

Robotics §4.4. Franka robot, apple-and-orange swap task with a memory requirement. Diffusion Forcing reaches 80% success. A diffusion policy without memory fails §4.4. With visual distractions or a fully occluded camera, success falls to 76% for Diffusion Forcing and 48% for a next-frame diffusion baseline §4.4.

Compositional generation §4.3 app. E. Qualitative: switching between full memory and short MPC-style plans recovers either the cross-shaped or the V-shaped trajectory distribution §4.3.

Time series, CRPS_sum, lower is better §4.5 app. E. Mean and standard deviation over five seeds for the paper's method.

The paper says Diffusion Forcing is overall tied with ScoreGrad, and places fourth on Wikipedia §4.5. Dataset details are in app. F.

Paper vs code

The official repository is buoyancy99/diffusion-forcing, linked from the project page. Its main branch is a reimplementation with temporal attention (its README calls it the latest version). The README says the RNN code the paper used is on a separate paper branch. I could not read that branch, so the citations below are to main at code_commit. Statements about the RNN implementation are unverified.

• Paper says the planning objective is written on the noise epsilon in the training loss §3.2. The planning config sets objective: pred_x0 df_planning.yaml, and the loss is an MSE on the clean signal weighted by clipped SNR diffusion.py. Appendix D says planning uses x_0 prediction app. D, so the appendix and the code agree but the main text's equation does not match them.
• Paper says the noise levels are sampled uniformly from [K] §3.2. The code samples integers in [0, timesteps) df_base.py. Frames that are masked out are set to the top level timesteps - 1 df_base.py.
• Paper does not describe a special first-frame noise rule. The planning code draws the first frame's level from the lower quarter of the range with probability 0.5 df_planning.py.
• Paper uses a recurrent backbone for the dynamics and a residual MLP for non-video data app. D. The main planning config uses a transformer with 12 layers, width 128 and 4 heads df_planning.yaml.
• Paper's sampling uses the schedule matrix K §3.2. The code's pyramid schedule builds a matrix with sampling_timesteps + (horizon - 1) * uncertainty_scale + 1 rows, so the number of rows depends on the horizon df_base.py.
• Paper uses the reweighting from min-SNR for the base loss app. D. The code's default use_snr is False df_base.yaml, with fused SNR selectable by config df_base.yaml. The paper states fused SNR is used for video only app. D.
• Paper says DDIM sampling with 50 steps for non-video app. D. The base config uses ddim_sampling_eta: 1.0 df_base.yaml, while the planning config uses 0.0 df_planning.yaml. The paper does not state eta.
• Sampling line 8 of the algorithm, as the paper describes it, can either recompute the posterior latent or reuse the sampled one. Appendix D says the authors reuse it, to halve compute and to get the stabilisation behaviour app. D. I did not check which option the main sampler takes.

No discrepancies were found in the planning and base diffusion settings that both the paper and the code state (noise schedule, frame stack, diffusion steps and sampling steps all match).

Pitfalls

• Noise level range. Paper says k_t is drawn from [K]. Code draws from [0, timesteps) and treats the top level as pure noise. An implementer must decide whether level 0 is included, and whether the top level is used for training or only for masking.
• Unstated values an implementer must choose: learning rate, optimiser, weight decay, warm-up, RNN width and depth, guidance scale, the stabilisation level k, and the MCG sample count. The paper gives none of them, and the values the code uses for some of these are listed in the table above.
• Stabilisation level. The paper uses "slightly noisy" tokens with 0 < k << K but gives no number §3.3. The planning code uses 10 out of 1000 df_planning.yaml. Treat this as a guess, not a paper value.
• Copy versus resample when the level stays the same. Paper says the corner case where a token keeps its level can either copy the old value or resample. The authors prefer resampling, and use it for MCG. For k=0 and k=K, copying is equivalent in practice app. D.
• Sampling cost. Denoising runs over every row of the schedule, and each row runs the recurrent update for every token. A long horizon with a pyramid schedule gives many rows df_base.py. The paper's compute claims are about training only app. D.
• Reward metric. The paper warns that episode reward on Maze2D can reward a slow agent that stays near the goal, and it suggests time-to-goal as an alternative app. D. Do not compare reward numbers with other work without checking the metric.
• Guidance on the maze. The maze reward is the sum over time steps of the distance to the goal, so any step can count as the final step app. D. The agent tends to leave the goal, because the training data never stays there app. D.
• Time-series metric. CRPS_sum is computed by summing over features and estimating quantiles from 100 sampled trajectories app. F. A different number of samples or quantiles gives different numbers.
• Scale. Video experiments use a 36M-parameter model, trained on 10% of the dataset for compute reasons app. D app. F. The paper does not study scaling beyond these sizes §5.

Lineage

• Builds on: masked modelling and teacher forcing §3.1; denoising diffusion (DDPM and DDIM) §2 §3.2 app. D; Bayesian filtering §2.
• Compared against: TimeGrad, Transformer-MAF, ScoreGrad, and other time-series models §4.5; Diffuser for planning §4.2; teacher-forced next-frame diffusion and full-sequence diffusion for video §4.1; diffusion policy for robot imitation §4.4.
• Related work the paper names: AR-Diffusion and Rolling Diffusion, which also use per-token noise, but with noise levels tied to token position app. C.
• Superseded by: not known from this wiki's content. The project's own README points to a follow-up repository, Diffusion Forcing v2 (kwsong0113/diffusion-forcing-transformer).