Diffusion-Based Language Models

Autoregressive models generate left to right. Diffusion models generate everywhere at once, like a toddler with finger paint.

Spectra, Simultaneously Denoising AI Agent

1. From Continuous to Discrete Diffusion

A Quick Review: Diffusion in Images

In image diffusion models, the forward process gradually adds Gaussian noise to an image over many steps until it becomes pure noise. The reverse process learns to denoise step by step, recovering the original image. At generation time, you start from random noise and iteratively denoise to produce a new image. This works beautifully because pixel values are continuous, and Gaussian noise is a natural perturbation in continuous space.

Text, however, is discrete: each position holds a token from a finite vocabulary (e.g., 50,000 words). You cannot add Gaussian noise to the word "cat" in any meaningful sense. This fundamental difference requires a completely different formulation of the diffusion process.

Discrete Diffusion for Text

Instead of adding continuous noise, discrete diffusion corrupts tokens. The most common approach replaces tokens with a special [MASK] token (absorbing diffusion) or with random tokens from the vocabulary (uniform diffusion). Over many forward steps, a clean sentence becomes a sequence of masked or random tokens. The reverse process learns to recover the original tokens from this corrupted state.

Figure 5.4.2: Discrete diffusion for text. The forward process masks tokens; the reverse process recovers them. Multiple tokens can be unmasked per step.

2. Key Models and Approaches

Comparison of Approaches

3. Parallel Generation: The Speed Advantage

The most exciting property of diffusion language models is parallel token generation. In autoregressive models, generating N tokens requires N sequential forward passes. Each pass depends on the output of the previous one, creating an inherent sequential bottleneck that limits throughput regardless of hardware.

Diffusion models sidestep this entirely. At each denoising step, the model predicts all positions simultaneously. A 1,000-token output might require only 20 to 50 denoising steps, where each step processes the entire sequence in parallel. On modern GPU hardware with massive parallel processing capability, this can translate to order-of-magnitude latency reductions for long outputs. Code Fragment 5.4.1 below puts this into practice.

# Conceptual comparison of generation steps
import numpy as np

def compare_generation_steps(sequence_lengths, diffusion_steps=30):
 """Compare sequential steps needed for AR vs diffusion generation."""
 print(f"{' Length':>8s} | {'AR Steps':>10s} | {'Diffusion Steps':>16s} | {'Speedup':>8s}")
 print("-" * 52)
 for length in sequence_lengths:
 ar_steps = length # one forward pass per token
 diff_steps = diffusion_steps # fixed number of denoising steps
 speedup = ar_steps / diff_steps
 print(f"{length:>8d} | {ar_steps:>10d} | {diff_steps:>16d} | {speedup:>7.1f}x")

compare_generation_steps([50, 200, 500, 1000, 4000])

The speedup grows linearly with output length because autoregressive models scale linearly while diffusion models use a fixed number of steps. For a 4,000-token output, the theoretical speedup is over 100x in sequential steps. Of course, each diffusion step processes the entire sequence (which is more expensive per step than a single autoregressive step), so the wall-clock speedup is smaller, typically 3x to 10x for long sequences. Nevertheless, this represents a fundamental architectural advantage. Code Fragment 5.4.2 below puts this into practice.

# Simplified discrete diffusion process (conceptual)
import torch

def simulate_diffusion_generation(seq_length, vocab_size, num_steps=20):
 """
 Simulate the discrete diffusion generation process.
 At each step, the model predicts all masked positions simultaneously.
 We unmask a fraction of positions per step.
 """
 MASK_TOKEN = vocab_size # special mask token ID

 # Start fully masked
 sequence = torch.full((seq_length,), MASK_TOKEN)
 tokens_per_step = max(1, seq_length // num_steps)

 print(f"Generating {seq_length} tokens in {num_steps} parallel steps\n")

 for step in range(num_steps):
 # Find masked positions
 masked_positions = (sequence == MASK_TOKEN).nonzero(as_tuple=True)[0]
 if len(masked_positions) == 0:
 break

 # "Model prediction": in reality, a neural network predicts all positions
 # Here we simulate by filling with random tokens
 n_to_unmask = min(tokens_per_step, len(masked_positions))
 # Unmask positions with highest "confidence" (simulated)
 chosen = masked_positions[torch.randperm(len(masked_positions))[:n_to_unmask]]
 sequence[chosen] = torch.randint(0, vocab_size, (n_to_unmask,))

 n_remaining = (sequence == MASK_TOKEN).sum().item()
 pct_done = 100 * (1 - n_remaining / seq_length)
 print(f"Step {step+1:2d}: unmasked {n_to_unmask:3d} tokens | {pct_done:5.1f}% complete | {n_remaining:3d} masked")

 return sequence

result = simulate_diffusion_generation(seq_length=100, vocab_size=50000, num_steps=10)

4. Gemini Diffusion

The Gemini Diffusion approach represents the first serious industrial investment in diffusion-based text generation at frontier scale. Key reported characteristics include:

• Adaptive denoising schedules: The number of denoising steps is not fixed but varies based on the complexity of what is being generated. Factual recall might need only a few steps, while multi-step reasoning might need dozens.
• Confidence-ordered unmasking: Tokens the model is most confident about are unmasked first, allowing later steps to condition on the high-confidence "skeleton" of the output.
• Hybrid autoregressive-diffusion: Some implementations use autoregressive generation for the first few tokens (establishing topic and tone) and then switch to diffusion for the bulk of the output.

5. Advantages and Limitations

Figure 5.4.3: Comparing the strengths and weaknesses of autoregressive and diffusion-based text generation paradigms.

The Quality Gap for Reasoning

The most significant current limitation of diffusion language models is their performance on tasks requiring complex multi-step reasoning. Autoregressive models naturally "think step by step" because each token is generated in sequence, allowing each step to build on the previous. Diffusion models generate all positions in parallel, which makes it harder to enforce logical dependencies between distant parts of the output. For tasks like mathematical proofs, code generation, and chain-of-thought reasoning, autoregressive models still hold a significant advantage.

6. TraceRL: Reinforcement Learning for Diffusion LLMs

The TraceRL approach works by treating the entire denoising trajectory as a sequence of decisions:

1. Generate a complete denoising trajectory: from fully masked to final output, recording the decisions at each step
2. Score the final output using a reward model (the same kind used in autoregressive RLHF)
3. Attribute the reward back to each denoising step using a credit assignment mechanism inspired by policy gradient methods
4. Update the model to increase the probability of denoising trajectories that led to high-reward outputs

This is conceptually similar to how REINFORCE or PPO work in autoregressive models, but adapted for the parallel, iterative structure of diffusion. The "trace" in TraceRL refers to the recorded sequence of denoising steps, which serves as the equivalent of a token-by-token trajectory in autoregressive generation. Code Fragment 5.4.3 below puts this into practice.

# Conceptual pseudocode for TraceRL training loop
# (simplified for illustrative purposes)

def trace_rl_training_step(diffusion_model, reward_model, prompt):
 """One step of TraceRL training (conceptual)."""

 # 1. Generate a denoising trajectory
 trajectory = []
 x_t = initialize_fully_masked(prompt)

 for t in reversed(range(T)):
 # Model predicts which tokens to unmask and what they should be
 predictions, log_probs = diffusion_model.denoise_step(x_t, t)
 trajectory.append({"step": t, "log_probs": log_probs, "state": x_t})
 x_t = apply_predictions(x_t, predictions)

 # 2. Score the final output
 final_text = x_t
 reward = reward_model(prompt, final_text)

 # 3. Compute policy gradient with reward attribution
 loss = 0
 for step_info in trajectory:
 # Attribute reward to each step (with discount)
 step_reward = attribute_reward(reward, step_info["step"], T)
 loss -= step_reward * step_info["log_probs"].mean()

 # 4. Update model
 loss.backward()
 optimizer.step()

 return reward

# This enables RLHF-style training for diffusion language models,
# previously an unsolved problem.
print("TraceRL: enables reward-based training for diffusion LLMs")
print("Key contribution: credit assignment across parallel denoising steps")

7. The Road Ahead

Diffusion-based language models are at a similar stage to where autoregressive transformers were around 2018 to 2019: the fundamental ideas are in place, early results are promising, but enormous scaling and engineering work remains. Several open questions will determine whether diffusion models become a mainstream alternative to autoregressive generation:

• Scaling laws: Do diffusion language models follow similar scaling laws to autoregressive models? Early evidence suggests they do, but with different constants.
• Alignment: Can we develop efficient techniques for aligning diffusion models with human preferences? TraceRL is a first step, but the field needs the equivalent of DPO, constitutional AI, and other alignment breakthroughs.
• Hybrid architectures: Will the future be purely diffusion, purely autoregressive, or a hybrid? Several teams are exploring models that use autoregressive generation for planning and diffusion for execution.
• Multimodal diffusion: Unified diffusion models that generate text, images, code, and structured data simultaneously could be a natural extension of this paradigm.

pip install transformers torch matplotlib numpy

import torch
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()

def get_logits(input_ids):
 """Get next-token logits from the model."""
 with torch.no_grad():
 outputs = model(input_ids)
 return outputs.logits[0, -1, :] # logits for the last position

def greedy_decode(prompt, max_tokens=30):
 """Greedy decoding: always pick the most likely next token."""
 input_ids = tokenizer.encode(prompt, return_tensors="pt")
 generated = list(input_ids[0])

 for _ in range(max_tokens):
 logits = get_logits(torch.tensor([generated]))
 next_token = torch.argmax(logits).item()
 generated.append(next_token)
 if next_token == tokenizer.eos_token_id:
 break

 return tokenizer.decode(generated)

prompt = "The future of artificial intelligence is"
print("GREEDY:")
print(greedy_decode(prompt))

def sample_with_strategy(prompt, max_tokens=30, temperature=1.0,
 top_k=0, top_p=1.0):
 """Decode with configurable temperature, top-k, and top-p."""
 input_ids = tokenizer.encode(prompt, return_tensors="pt")
 generated = list(input_ids[0])

 for _ in range(max_tokens):
 logits = get_logits(torch.tensor([generated]))

 # Apply temperature
 logits = logits / temperature

 # Apply top-k filtering
 if top_k > 0:
 top_k_values, _ = torch.topk(logits, top_k)
 threshold = top_k_values[-1]
 logits[logits < threshold] = float("-inf")

 # Apply top-p (nucleus) filtering
 if top_p < 1.0:
 sorted_logits, sorted_indices = torch.sort(logits, descending=True)
 cumulative_probs = torch.cumsum(
 torch.softmax(sorted_logits, dim=-1), dim=-1)
 # Remove tokens with cumulative probability above the threshold
 sorted_mask = (cumulative_probs
 - torch.softmax(sorted_logits, dim=-1) >= top_p)
 sorted_logits[sorted_mask] = float("-inf")
 # Scatter back to original indexing
 logits = torch.zeros_like(logits).scatter(
 0, sorted_indices, sorted_logits)

 # Sample from the filtered distribution
 probs = torch.softmax(logits, dim=-1)
 next_token = torch.multinomial(probs, num_samples=1).item()
 generated.append(next_token)
 if next_token == tokenizer.eos_token_id:
 break

 return tokenizer.decode(generated)

prompt = "The future of artificial intelligence is"
print("TEMPERATURE 0.3 (focused):")
print(sample_with_strategy(prompt, temperature=0.3))
print("\nTEMPERATURE 1.5 (creative):")
print(sample_with_strategy(prompt, temperature=1.5))
print("\nTOP-K 10:")
print(sample_with_strategy(prompt, top_k=10))
print("\nTOP-P 0.9 (nucleus):")
print(sample_with_strategy(prompt, top_p=0.9))

import matplotlib.pyplot as plt

input_ids = tokenizer.encode(prompt, return_tensors="pt")
logits = get_logits(input_ids)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

temperatures = [0.3, 1.0, 2.0]
for ax, temp in zip(axes, temperatures):
 probs = torch.softmax(logits / temp, dim=-1).numpy()
 top_indices = np.argsort(probs)[-15:][::-1]
 top_probs = probs[top_indices]
 top_tokens = [tokenizer.decode([i]).strip() for i in top_indices]

 ax.barh(range(len(top_tokens)), top_probs, color="#3498db")
 ax.set_yticks(range(len(top_tokens)))
 ax.set_yticklabels(top_tokens, fontsize=8)
 ax.set_xlabel("Probability")
 ax.set_title(f"Temperature = {temp}")
 ax.invert_yaxis()

plt.suptitle("How Temperature Reshapes Token Probabilities",
 fontweight="bold")
plt.tight_layout()
plt.savefig("decoding_temperature.png", dpi=150)
plt.show()

from transformers import set_seed

input_ids = tokenizer.encode(prompt, return_tensors="pt")
set_seed(42)

strategies = {
 "Greedy": dict(max_new_tokens=30, do_sample=False),
 "Top-k (k=10)": dict(max_new_tokens=30, do_sample=True, top_k=10),
 "Top-p (p=0.9)": dict(max_new_tokens=30, do_sample=True, top_p=0.9),
 "Temp 0.3": dict(max_new_tokens=30, do_sample=True, temperature=0.3),
 "Beam (n=4)": dict(max_new_tokens=30, num_beams=4, do_sample=False),
}

print(f"Prompt: '{prompt}'\n")
for name, kwargs in strategies.items():
 set_seed(42)
 output = model.generate(input_ids, **kwargs,
 pad_token_id=tokenizer.eos_token_id)
 text = tokenizer.decode(output[0], skip_special_tokens=True)
 print(f"{name:15s}: {text}")

print("\nFrom scratch: ~60 lines per strategy.")
print("Hugging Face: 1 line per strategy with generate().")