What is Diffusion?

Forward Process

Gradually add noise to data over multiple timesteps

Reverse Process

Learn to denoise, going from noise back to data

Efficiency Tricks

• Can jump directly to any noise level \(t\) without going step-by-step!
  • Cumulative signal retention: Let \(\alpha_t = 1 - \beta_t\) and \(\bar{\alpha}_t = \prod_{s=1}^t \alpha_s\)
• We can propagate gradients through the reparameterization trick:
  • \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon\)
  • where \(\epsilon \sim \mathcal{N}(0, I)\) is standard Gaussian noise
  • \(\bar{\alpha}_t\) controls the signal-to-noise ratio at time \(t\)

Sampling Strategies

Stochastic Sampling:

\(x_{t-1} = \text{denoise}(x_t) + \sigma_t \epsilon\) where \(\epsilon \sim \mathcal{N}(0, I)\)

• More diverse outputs
• Explores different possible generations

Deterministic Search (Mode-seeking):

\(x_{t-1} = \text{denoise}(x_t)\)

• Faster generation (can skip steps)
• Less diverse outputs

The Discrete Problem

Images are Continuous:

• Images are continuous-valued (pixels in [0, 1])
• Size is fixed (e.g., 32x32, 256x256)

Language is Discrete:

• Tokens from a discrete vocabulary
• Variable-length sequences

Diffusion-LM: The Approach

Process:

1. Embed: Map discrete tokens to continuous vectors \(\text{emb}(w) = E_w\) where \(E \in \mathbb{R}^{V \times d}\) is the embedding matrix, \(V\) is vocabulary size, \(d\) is embedding dimension

2. Diffuse: Add noise in embedding space \(x_t = \sqrt{\bar{\alpha}_t} \text{emb}(w) + \sqrt{1-\bar{\alpha}_t} \epsilon\)

3. Denoise: Learn to predict clean embeddings \(\epsilon_\theta(x_t, t) \approx \epsilon\)

4. Round: Map back to discrete tokens \(w' = \arg\max_w \cos(E_w, \hat{x}_0)\) where \(E_w\) is the embedding for word \(w\)

Challenge: Rounding is non-differentiable

Diffusion-LM: The Clamping Trick

Problem: Even with \(x_0\) prediction, rounding errors accumulate during sampling

Solution - Clamping at Decode Time:

Standard sampling from predicted \(x_0\):

With clamping:

where \(\text{Clamp}(x) = \arg\max_{w \in V} \cos(E_w, x)\) (nearest embedding by cosine similarity)

Effect: Forces intermediate predictions to commit to valid word embeddings, making subsequent predictions more precise

Diffusion-LM: Results

• Higher perplexity (~36 vs. ~23 for GPT-2) but better controllability
• Controllability: Outperforms PPLM and FUDGE on 6 control tasks (syntax, semantics, infilling)
• Robustness: More resilient to local failures than autoregressive models (no exposure bias)

Weaknesses:

• Slower decoding (2000 diffusion steps vs. sequence-length AR steps)
• Slower training convergence than AR models

MDLM: Forward Process

Forward Process: Gradually mask tokens

At time \(t\), randomly mask each token with probability \(m_t\):

$q(x_t \mid x_0) = \prod_i \begin{cases} \text{[MASK]} & \text{with probability } m_t \ x_{0,i} & \text{with probability } 1-m_t \end{cases}$

where \(m_t\) is the masking rate at timestep \(t\) (like a noise schedule)

Why masking?

• Natural corruption for text
• BERT-style MLMs already trained for this task
• Simpler than other discrete diffusion operations

MDLM: Training Objective

Training reduces to a weighted average of standard masked language modeling:

What this means:

• For each masked position \(i\), predict the original token
• Weight the loss by time-dependent factor \(w(t) = \frac{\alpha_t - \alpha_s}{1 - \alpha_t}\)
  • Where \(\alpha_t\) is the cumulative unmasking rate at time \(t\)
  • Comes from a theoretical lower bound on the likelihood
• That’s it - just like training BERT but with time-dependent importance weights!

MDLM: Inference 2 – Semi-Autoregressive Generation

• If we parameterize \(p_\theta(\cdot \mid x_t)\) without explicit time \(t\) (time-independent MDLM), we can cache logits for unmasked positions
• Why caching helps: Unmasked positions don’t change due to “carry-over unmasking” property
  • Only recompute activations for newly unmasked positions at each step
  • Cached positions remain unchanged throughout generation
• Extends naturally to variable-length sequences
• Efficient 2-3× speedup with caching

MDLM: Efficient Sampling with Caching

Problem: Recomputing all tokens at each step is wasteful

Solution: Cache activations for unchanged tokens

def sample_with_cache(model, x_t, cache):
    # Only compute changed positions
    changed_pos = get_masked_positions(x_t)
    new_activations = model(x_t, positions=changed_pos)
    
    # Reuse cached activations for unchanged positions
    cache.update(changed_pos, new_activations)
    
    # Predict and sample next state
    predictions = decode(cache.activations)
    return sample_next(predictions)

Speedup: 2-3× faster generation

SEDD: Score Entropy Discrete Diffusion

Reference: Lou et al., ICML 2024

Key Innovation: A principled score-based framework for discrete diffusion using continuous-time Markov chains

Core Idea:

• Models discrete diffusion as continuous-time Markov chain (CTMC)
• Learns probability ratios between tokens: \(s_\theta(x)_y \approx \frac{p(y)}{p(x)}\)
• Achieves better perplexity than GPT-2 - first discrete diffusion to do so!

SEDD: What Are Probability Ratios?

In continuous diffusion: Learn the score function \(\nabla_x \log p(x)\) (gradient pointing to higher probability)

In discrete diffusion: Can’t take gradients! Instead learn probability ratios:

What this means:

• For each position, learn relative likelihood of all possible tokens
• Not just “cat is correct” but “cat is 3× more likely than dog, 10× more likely than the”
• Richer signal than just predicting the correct token

Why ratios? They’re the discrete analog of the continuous score function

SEDD: Training with Score Entropy

Training objective: Learn to predict probability ratios via denoising

Score Entropy Loss:

What this does:

• Trains model to match true probability ratios from the forward process
• The ratios \(\frac{q(y|x_0)}{q(x_t|x_0)}\) are known from how we add noise!
• Squared error between predicted and true ratios

Key insight: Denoising makes this tractable - we know ground truth ratios

SEDD: Sampling via τ-leaping

Core idea: Take larger discrete steps, updating all positions in parallel

The update rule: At each position \(i\), sample new token from:

What this means:

• With probability \(\approx 1 - \Delta t\): stay at current token \(x_t^i\)
• With probability \(\propto \Delta t \cdot Q_t(x_t^i, y) \cdot s_\theta(\vec{x}_t)_{i,y}\): transition to token \(y\)
• The learned ratios \(s_\theta\) determine which transitions are more likely

Key property: All positions updated independently - enables parallel sampling

Tradeoff: Larger \(\Delta t\) = fewer steps but lower accuracy (like Euler discretization)

SEDD: Tweedie Denoising

Motivation: Can we do better than τ-leaping if we use the learned ratios more carefully?

Discrete Tweedie’s theorem: If we know all ratios \(\frac{p_t(y)}{p_t(x_t)}\), we can optimally denoise

Practical approximation: We only know ratios between nearby sequences, so:

1. Estimate clean data: \(\hat{p}_0 = \exp(-\sigma_t^{\Delta t} Q) \cdot s_\theta(\vec{x}_t)\)
2. Re-noise forward: Sample from \(\hat{p}_0\) re-noised to time \(t - \Delta t\)

The update rule: Sample token at position \(i\) from:

where \(\sigma_t^{\Delta t} = \overline{\sigma}(t) - \overline{\sigma}(t - \Delta t)\)

Benefit: Optimal among all parallel sampling strategies (proven in paper)

SEDD: Conditional Generation

Question: How to generate with constraints (e.g., fixed prefix, suffix, or infill)?

Key insight: Probability ratios make conditioning exact via Bayes’ rule:

What this means:

• \(\Omega\) = positions to generate (unfilled)
• \(\overline{\Omega}\) = positions to condition on (filled)
• Conditional ratios = unconditional ratios with fixed positions

In practice:

1. Use the same model \(s_\theta\) trained unconditionally
2. During sampling, only update positions in \(\Omega\)
3. Keep positions in \(\overline{\Omega}\) fixed at their conditioned values

Examples: prefix continuation (\(\overline{\Omega} = \{1,...,c\}\)), suffix (\(\overline{\Omega} = \{d+1,...,n\}\)), infilling (any subset)

SEDD: Results

Perplexity Results (vs GPT-2 Small):

Key Achievements:

• First discrete diffusion to beat autoregressive baselines
• 32× fewer function evaluations for comparable quality
• Principled likelihood-based training and evaluation
• Strong theoretical foundation (CTMC + score matching)

Diffuser: Edit-Based Reconstruction

Reference: Reid et al., arXiv 2022

Key Innovation: Uses all four Levenshtein edit operations for both corruption and reconstruction

Core Idea:

• Models text generation as natural editing processes
• Uses INSERT, DELETE, KEEP, REPLACE operations
• Two-component architecture: tagger + generator
• Enables flexible editing in various contexts

Diffuser: Architecture

Two-Stage Process:

1. Tagger: Predicts edit operations \(p^{\text{tag}}_\theta(e_t \mid x_t)\)
2. Generator: Produces tokens for edits \(p^{\text{gen}}_\theta(x_{t-1} \mid x_t, e_t)\)

Edit Operations: INSERT, DELETE, REPLACE, KEEP

Training: Corrupt with reverse operations (e.g., train INSERT by adding deletions)

Corruption Parameters:

• \(E_t\): Edit type distribution (e.g., 60% keep, 20% replace, 10% delete, 10% insert)
• \(E_l\): Edit length distribution (e.g., Poisson λ=3)

Diffuser: Training Objectives

Both modules trained with cross-entropy losses:

Tagger Loss: Predicts edit operations

where \(e^*\) are the ground truth edit operations that generated \(x_t\) from \(x_{t-1}\)

Generator Loss: Predicts tokens at edited positions

Key Difference from Other Diffusion Models:

• Not using ELBO (Evidence Lower Bound) like DDPM/MDLM/SEDD
• Instead: Two separate supervised learning objectives
• Simpler training, but less theoretical grounding
• Leverages edit distance structure for natural editing workflow

Diffuser: 2D Beam Search

Novel Decoding Strategy:

Standard beam search only searches over tokens. Diffuser introduces 2D beam search:

• Intra-revision dimension: Beam width \(b\) for token-level search
• Inter-revision dimension: Beam width \(r\) for diffusion steps

Algorithm:

1. Generate \(b\) candidates at current diffusion step
2. Keep top \(r\) based on log-likelihood
3. Feed all \(r\) to next diffusion step
4. Each produces \(b\) new candidates → \(r \times b\) total
5. Keep top \(r\) again

Result: Searches over both edit operations AND diffusion trajectories

Diffuser Results

Machine Translation (WMT’14 En-De):

Summarization (CNN/DM):

Key Findings:

• AR bootstrap best for translation (28.8 vs 27.3 baseline)
• Source bootstrap best for summarization (treats as editing)
• First non-AR model to beat AR on these tasks

Comparing Discrete Diffusion Approaches