Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

Emiel Hoogeboom; David Ruhe; Jonathan Heek; Thomas Mensink; Tim Salimans

Source-linked AI summary

Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

Emiel Hoogeboom, David Ruhe, Jonathan Heek, Thomas Mensink, Tim Salimans

arXiv:2603.20155v1Mar 20, 2026cs.LGcs.CVstat.ML

TL;DR

Discrete diffusion is difficult to distill into few-step generators because sampling requires many iterations and factorized modeling can accumulate errors. D-MMD generalizes moment matching to discrete diffusion by using probability-based updates and stochastic soft samples. Across text and image experiments, distilled generators can outperform their teachers, though performance may later decline as sampling steps increase.

Problem
Discrete diffusion models require many sampling iterations, while factorized conditional modeling can accumulate errors during those iterations.
Method
D-MMD generalizes Moment Matching Distillation to discrete diffusion and uses soft probability vectors to obtain differentiable moment-matching updates.
Results
D-MMD generators outperform their teachers on text and image generation while using only a fraction of the denoising steps.
Takeaways & Limitations
D-MMD provides a principled approach for distilling discrete diffusion processes into few-step generators while retaining quality and diversity given sufficient sampling steps.
Takeaways & Limitations
When a student outperforms its teacher at some step count, its performance can degrade with more steps as it converges toward the teacher’s performance.

Abstract

from arXiv · show

It is currently difficult to distill discrete diffusion models. In contrast, continuous diffusion literature has many distillation approaches methods that can reduce sampling steps to a handful. Our method, Discrete Moment Matching Distillation (D-MMD), leverages ideas that have been highly successful in the continuous domain. Whereas previous discrete distillation methods collapse, D-MMD maintains high quality and diversity (given sufficient sampling steps). This is demonstrated on both text and image datasets. Moreover, the newly distilled generators can outperform their teachers.

1. Introduction

Discrete diffusion models can use blocks of tokens efficiently, but require many iterations because factorized conditional modeling accumulates errors. D-MMD adapts continuous moment-matching distillation to produce few-step discrete generators.

Motivation: Independent token modeling conditioned on previously generated tokens can accumulate errors across sampling iterations.
Motivation: Discrete diffusion sampling requires many iterations, increasing computing cost and FLOPs.Reducing iterations lowers cost.
Contribution: D-MMD generalizes Moment Matching Distillation to discrete diffusion processes.The method uses insights from continuous diffusion distillation.
Contribution: D-MMD distills few-step generators whose sample quality can surpass teachers on text and image generation.Figure 1 reports matching or better teacher performance with fewer function evaluations.
Method: The training procedure samples noisy data, predicts a probability vector, samples an intermediate state, and optimizes an auxiliary loss.The auxiliary model is trained using teacher and student-related targets.

2. Background

Discrete diffusion learns reverse transitions for categorical data using uniform or masked corruption, while diffusion training commonly targets conditional expectations. Moment Matching Distillation provides a continuous-domain framework based on matching these expectations.

Diffusion models: Diffusion models approximate reverse-process sub-steps, often assuming dimensional independence.
Diffusion models: Training objectives commonly reduce to estimating the conditional expectation of clean data given noisy data.Because this expectation is unavailable directly, models learn from diffused dataset samples.
Discrete diffusion: Discrete diffusion uses uniform or masked corruption processes for categorical data.Masked diffusion transforms tokens into a masking token; uniform diffusion transforms them toward a uniform distribution.
Discrete diffusion: The discrete process interpolates data with a factorized stationary distribution using a noise schedule.
Moment matching: Moment Matching Distillation seeks identical clean-data conditional expectations under the data and distilled sampling distributions.Practical optimization replaces unavailable generator expectations with teacher and auxiliary model outputs while keeping the teacher fixed.

3. Discrete MMD: A generalization of MMD

D-MMD rewrites moment-matching distillation for general diffusion processes and makes it practical for discrete variables by matching probability-based moments. Its factorized generator can produce correlated outputs through stochastic soft-sample generation.

Generalization: D-MMD generalizes alternating moment-matching optimization into a min-max formulation for discrete diffusion.
Generalization: The generator minimizes teacher loss while maximizing auxiliary-model loss, whereas the auxiliary model minimizes generator loss and stays close to the teacher.
Generalization: Under the stated stop-gradient and independence assumption, D-MMD produces the same gradients as continuous MMD.
Discrete matching: For discrete diffusion, D-MMD replaces hard categorical samples with soft probability vectors to obtain a differentiable moment-matching update.The resulting update acts in log-probability space.
Fixed point: Perfect matching between teacher and auxiliary models guarantees sampling from the teacher distribution.The algorithm’s fixed point occurs when the generator induces the teacher distribution, though adversarial optimization can depend on hyperparameters.
Correlated outputs: The generator can learn correlated outputs because stochastic soft samples are correlated while only the subsequent hard-token sampling step is factorized.Observed output correlation is associated with reduced soft-sample output entropy.
Training variants: Soft targets are valid for masked diffusion but uniform diffusion requires hard samples for auxiliary-model training.

4. Related work

Prior discrete distillation methods have limitations in representing correlations or supporting broader diffusion settings. D-MMD extends moment matching to uniform diffusion and few-step generators while relating to concurrent approaches.

Deterministic distillation: Deterministic diffusion distillation learns trajectories progressively or uses consistency-style self-targeting.
Stochastic distillation: Stochastic distillation approximately minimizes divergence between a distilled generator and its teacher, while MMD performs better in few-step regimes than related distribution-matching approaches.
Discrete distillation: SDTT improves discrete sampling only to a limited degree and cannot approximate perfectly correlated coin tosses in one step.
Discrete distillation: Di4C uses mixture distributions for correlations, but the required number of mixtures can grow exponentially with the number of token correlations.
Discrete distillation: D-MMD is equivalent to DiMO in the one-step masked-diffusion case while extending to uniform diffusion and few-step generators.
Concurrent work: IDLM differs by generating the full sample and diffusing backward, whereas D-MMD samples from the posterior; the authors view the methods as complementary.

5. Evaluating discrete diffusion models using Gradient Moments

The paper argues that generative perplexity can misjudge discrete diffusion samples and proposes Gradient Moment as an alternative quality metric. Gradient Moment compares sample and data loss gradients through a reference language model, including conditional evaluation.

Generative perplexity is unavailable as a tractable likelihood for distilled discrete diffusion models and can reward ungrammatical samples with repeated words.The paper identifies high-density but atypical generations as a failure case for this metric.
Gradient Moment evaluates sample quality using the squared norm of a reference model’s loss gradient on generated samples.The metric is motivated by the expected zero loss gradient on the data distribution when the reference model has converged.
The metric is centered by the data loss gradient to account for imperfect reference-model training or distribution mismatch.This correction prevents nonzero gradients on the data from being attributed entirely to generated-sample quality.
Gradient Moment reaches zero when the sampling distribution equals the data distribution and increases when the reference model distinguishes generated samples.The paper estimates it stochastically using inner products of gradients from independent sample and training-data minibatches.
Unlike sampling-based metrics such as FID, the reference-model gradient norm remains valid for samples conditioned on prompts or prefixes.Conditional likelihoods replace unconditional likelihoods in the metric’s equations.

6. Experiments

Across image, text, block-autoregressive, and related-work comparisons, D-MMD produces strong generators with substantially fewer denoising steps than teachers and prior methods. The experiments also examine evaluation metrics, noise conditioning, and the possibility that distilled students outperform teachers at limited step counts.

6.1. CIFAR-10: D-MMD models substantially outperform their CIFAR10 teachers while using a fraction of the NFEs.Table 1 measures image quality with FID using 50K samples compared with the training dataset.
6.1. CIFAR-10: 3.7 FID is achieved by uniform diffusion in 32 steps, compared with 7.5 for its 1024-step teacher.For masked diffusion, the distilled generator outperforms its teacher with 16 steps and reaches 3.5 FID with 64 uniform denoising steps.
6.2. Text: Masked D-MMD reaches 0.236 GPT-2 GM in 16 steps and both text generators improve the teacher Pareto front.Text models generate 1024 tokens unconditionally while varying denoising steps; GPT-2 GM is used instead of generative perplexity.
6.2. Text: 16-step D-MMD matches the performance of a 256-step teacher in block-autoregressive diffusion with block size 256.This setup generates limited token blocks conditioned on an autoregressive encoder.
6.4. Comparison related work: D-MMD outperforms Di4C with 5.0 FID using only 8 steps with the uniform generator.Di4C’s teacher achieves 8.0 FID using 40 steps, while D-MMD also outperforms its own teacher models.
6.4. Comparison related work: GPT-2 GM exposes degradation across repeated SDTT distillation rounds, while D-MMD outperforms SDTT despite using the same teacher.The comparison also reports generative perplexity and sample entropy, with top-p selected at p = 0.85 for the masked teacher.

7. Conclusions

D-MMD distills discrete diffusion processes into few-step generators that tend to considerably outperform their teachers while using only a fraction of denoising steps.

D-MMD provides a principled technique for distilling discrete diffusion processes into few-step generators.
Generators tend to considerably outperform their teachers while using only a fraction of the denoising steps.

A. Sufficiency of matching first moments

The sufficiency argument uses an analytically available ground-truth posterior and conditional expectations to show equality between the learned and diffusion joint processes.

The argument starts with a diffusion process q(x, z_0, . . . , z_1) and its analytically available ground-truth posterior q(z_t−dt|z_t, x).
In the dt→0 limit, linearity in x lets the posterior expectation be evaluated at the conditional mean E_q[x|z_t].
Using p_η(z_1)=q(z_1)=Categorical(π) and equality of conditional expectations, the learned and target joint processes are equal.
The result follows by marginalization.

B. Sufficiency of matching factorized probabilities

Matching factorized distributions across all diffusion times is sufficient, under matching priors, to guarantee equality between the generator and target data distributions.

The argument considers a factorized product of the true posterior marginals q(x_i|z_t).
With matching priors, the generator must match these factorized distributions for every t∈[0, 1] to guarantee p_η(x)=q(x).
Figure 4 reports FID performance against sampling temperature or top p value in posterior sampling.

C. Extended Results: CIFAR10

The extended CIFAR10 results provide more detailed main-result analyses, including posterior-sampling adjustments and teacher-temperature settings.

The section provides more detailed results for the paper’s main results.
Figure 4 examines FID against posterior-sampling temperature or top-p value.
Evaluation adjusts posterior sampling using temperature scaling or top-p selection of probability mass.
Figure 5 examines FID against teacher temperature while MMD’ing the student model.
Figure 6 examines FID against teacher top-p value while MMD’ing the student model.

D. D-MMD for other discrete diffusion models

D-MMD extends moment matching distillation to discrete diffusion settings where the conditional expectation need not be optimal, deriving a loss with a teacher–auxiliary fixed point.

General discrete parametrizations: For some discrete diffusion parametrizations, including uniform diffusion, E_q[x|z_t] is not the optimal solution for the model prediction.The paper identifies this issue as a motivation for adapting D-MMD beyond the standard setting.
General discrete parametrizations: The simplified posterior parametrization covers uniform diffusion and processes interpolating from data x to a factorized stationary distribution π.The process is written as z_t ∼ Cat(z_s|α_t x + (1 − α_t)π).
Loss construction: Choosing a positive discretization step d_t or taking the limit d_t→0 provides two routes for formulating the discrete diffusion loss.The limit case requires subsequent algebraic manipulation.
Loss construction: Subtracting the two KL terms cancels the negative-entropy term and yields the simplified loss used in the discrete formulation.This cancellation involves the entropy term Σ π_{s−d_s}(x̂_η) log π_{s−d_s}(x̂_η).
Fixed point: At a fixed point, the student generator matches the teacher-induced data distribution, the auxiliary model equals the teacher, and both updates are zero.The auxiliary optimum likewise matches its transformed distribution to the teacher’s transformed distribution.