Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

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

cs.LGMar 20, 2026

Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

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

Organizations: Google DeepMind Amsterdam

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: D-MMD text generators match or outperform their teacher using fewer function evaluations.

Panel of Figure 2 — Algorithm 1 Discrete MMD Training

Figure 3: The perplexity metric keeps improving with lower temperature sampling while the grad moment eventually degrades.

Model	4	8	16	32	64	128	256	512	1024
Uniform Teacher			36.3	17.1	10.7	8.6	7.9	7.6	7.5
Uniform D-MMD	7.1	5.0	4.1	3.7	3.8
Masked Teacher			122.9	47.1	20.0	11.1	7.8	6.7	6.4
Masked D-MMD	22.3	12.7	5.3	3.8	3.5

Table 1: Overview of D-MMD generators and base model sample quality for different NFEs on CIFAR10. Results are measured in FID with 50K samples compared to the train dataset. D-MMD models substantially outperform their teacher while using a fraction of the NFEs.

Model	8	16	32	64	128	256	512
Uniform Teacher ( $p=0.50$ )			0.375	0.326	0.330	0.324	0.313
Uniform D-MMD ( $p=0.70/0.70$ )	0.337	0.310	0.307	0.316
Masked Teacher ( $p=0.85$ )			0.402	0.307	0.297	0.275	0.275
Masked D-MMD ( $p=0.85$ )	0.456	0.236	0.225	0.231
AR Baseline	0.061

Table 2: Text D-MMD generators and base model sample quality for different NFEs measured in GPT-2 GM (see section 5 ). D-MMD models outperform their teacher at a fraction of the NFEs.

Model	16	256
256-Block Uniform Teacher ( $p=0.9$ )	-	0.225
256-Block Uniform D-MMD ( $p=0.7$ )	0.225	-

Table 3: Block auto-regressive diffusion model with block size 256. 16-step D-MMD matches the performance of the 256-step teacher.

Method	NFE	FID
Di4C Teacher	40	8.0
Di4C (hybrid)	20	9.5
Di4C	10	20.6
Uniform Teacher	512	7.6
	64	10.7
Uniform D-MMD	8	5.0

Table 4: Comparison with literature on CIFAR10. D-MMD considerably outperforms existing methods using fewer NFEs.

Method	NFE	GPT-2 GM $\downarrow$	GPT2 Perplexity $\downarrow$	Sample entropy
Duo + DCD	4		108.2	4.82
Duo + Di4C	4		150.7	4.81
MDLM + SDTT	4		339.7	5.38
MDLM + Di4C	4		239.3	5.40
FMLM	4		76.4	5.05
Masked Teacher	256	0.275	22.5	5.13

Table 5: Comparison with literature on OWT measured in GPT-2 GM (lower is better), generative perplexity (should not be too high) and sample entropy (should not be too low). D-MMD is able to achieve even better results in fewer steps.

D-MMD Masked		4	8	16	32	64
without noise	(FID)	151 .0	37.0	14.7	7.7	6.0
	(generator output entropy)	1.26	1.37	1.57	1.86	1.91
with noise	(FID)	22.3	12.7	5.3	3.8	3.5
	(generator output entropy)	1.01	1.29	1.53	1.76	1.83

Table 6: Noise input conditioning is important for masked distillation. Fewer steps require more generator output collapse, and generators with noise conditioning can collapse their factorized output distribution more.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: FID performance vs sampling temperature or top $p$ value in posterior sampling.

Panel of Figure 5 — Figure 5: FID performance vs teacher temperature while MMD’ing the student model.

Panel of Figure 6 — Figure 6: FID performance vs teacher top $p$ value while MMD’ing the student model.

Explore similar work

Distribution Matching Distillation for Continuous Diffusion Language Models
Sep 30, 2026Paul Le Van Kiem, Dario Shariatian, Umut Simsekli +1Diffusion Model Distillation Diffusion Language Models
Inria, PSL Research University · Cohere · CMAP, Ecole Polytechnique
One-Step Distillation of Discrete Diffusion Image Generators via Fixed-Point Iteration
May 20, 2026Chaoyang Wang, Yunhai TongDiffusion Model Distillation Few-Step Diffusion Sampling
Peking University
Continuous-Time Distribution Matching for Few-Step Diffusion Distillation
May 7, 2026Tao Liu, Hao Yan, Mengting Chen +8Diffusion Distillation Few-Step Diffusion Sampling
1VCIP, College of Computer Science, Nankai University · 2Alibaba Group · College of Artificial Intelligence, Jilin University

Sep 30, 2026J/K move · Enter open · S save

Distribution Matching Distillation for Continuous Diffusion Language Models

Paul Le Van Kiem, Dario Shariatian, Umut Simsekli, Alain Durmus

Inria, PSL Research University · Cohere · CMAP, Ecole Polytechnique

Diffusion Model Distillation Diffusion Language Models Continuous Diffusion Language Models Distribution Matching Distillation Diffusion Model Inference Acceleration Diffusion Language Model Distillation

Open paper AskPDF arXiv

Continuous diffusion language models generate all tokens in parallel, yet high-quality generation can still require hundreds of network evaluations (NFEs). We study how distributional distillation can reduce this cost by exploiting the student's probabilistic token outputs. Our unified formulation connects the student's output parameterization to the resulting gradient estimators and yields two methods with the same student architecture and reverse-KL matching objective: Simplex-DMD uses continuous token relaxations and pathwise gradients, while Reinforce-DMD uses categorical sampling and REINFORCE with a learned density ratio. We develop both methods for multi-step generation and investigate the training and sampling choices associated with each parameterization. On OpenWebText, for sequences of 1,024 tokens, Simplex-DMD achieves a generative perplexity of 45.6 at a unigram entropy of 5.44 nats in just 4 NFEs, a 49% reduction relative to the strongest evaluated diffusion baseline at matched entropy and sampling budget. Reinforce-DMD improves the frontier at larger budgets, reaching a generative perplexity of 14.9 at an entropy of 5.00 nats with 256 NFEs, a 20% reduction under the same comparison protocol.

Abstract

Explore similar work

Distribution Matching Distillation for Continuous Diffusion Language Models

One-Step Distillation of Discrete Diffusion Image Generators via Fixed-Point Iteration

Continuous-Time Distribution Matching for Few-Step Diffusion Distillation