cs.LGMar 20, 2026
SaveBeyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD
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
Figure 1: D-MMD text generators match or outperform their teacher using fewer function evaluations.
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 ( ) | 0.375 | 0.326 | 0.330 | 0.324 | 0.313 | |||
| Uniform D-MMD ( ) | 0.337 | 0.310 | 0.307 | 0.316 | ||||
| Masked Teacher ( ) | 0.402 | 0.307 | 0.297 | 0.275 | 0.275 | |||
| Masked D-MMD ( ) | 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 ( ) | - | 0.225 |
| 256-Block Uniform D-MMD ( ) | 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 | GPT2 Perplexity | 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
Figure 4: FID performance vs sampling temperature or top value in posterior sampling.
Figure 5: FID performance vs teacher temperature while MMD’ing the student model.
Figure 6: FID performance vs teacher top value while MMD’ing the student model.
Explore similar work
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.
One-Step Distillation of Discrete Diffusion Image Generators via Fixed-Point Iteration
Discrete diffusion models excel at visual synthesis but rely on slow, iterative decoding. Existing single-step distillation methods attempt to bypass this bottleneck, either by training auxiliary score networks that effectively double compute, or by introducing specialized parameterizations and multi-stage pipelines that fragment optimization. In this paper, we introduce Fixed-Point Distillation (FPD), an end-to-end framework that constructs local correction targets by partially corrupting the student's one-step draft and refining it with a single teacher step. To compute the training objective in a semantically meaningful space, we lift discrete tokens into continuous features and apply a multi-bandwidth drift loss that iteratively accumulates these corrections. To backpropagate through the discrete bottleneck, we employ a straight-through estimator that feeds exact hard-sampled tokens to the teacher and decoder during the forward pass, ensuring that training and inference operate on the same codebook manifold, while routing continuous gradients back to the student logits. This fully differentiable pathway additionally accommodates an optional unconditional adversarial objective to enhance perceptual realism. Evaluations on both class- and text-conditional generation validate the effectiveness of our framework. FPD achieves competitive visual fidelity and structural alignment within a single inference step, narrowing the gap to multi-step teachers while outperforming existing discrete distillation baselines.
Continuous-Time Distribution Matching for Few-Step Diffusion Distillation
Step distillation has become a leading technique for accelerating diffusion models, among which Distribution Matching Distillation (DMD) and Consistency Distillation are two representative paradigms. While consistency methods enforce self-consistency along the full PF-ODE trajectory to steer it toward the clean data manifold, vanilla DMD relies on sparse supervision at a few predefined discrete timesteps. This restricted discrete-time formulation and mode-seeking nature of the reverse KL divergence tends to exhibit visual artifacts and over-smoothed outputs, often necessitating complex auxiliary modules -- such as GANs or reward models -- to restore visual fidelity. In this work, we introduce Continuous-Time Distribution Matching (CDM), migrating the DMD framework from discrete anchoring to continuous optimization for the first time. CDM achieves this through two continuous-time designs. First, we replace the fixed discrete schedule with a dynamic continuous schedule of random length, so that distribution matching is enforced at arbitrary points along sampling trajectories rather than only at a few fixed anchors. Second, we propose a continuous-time alignment objective that performs active off-trajectory matching on latents extrapolated via the student's velocity field, improving generalization and preserving fine visual details. Extensive experiments on different architectures, including SD3-Medium and Longcat-Image, demonstrate that CDM provides highly competitive visual fidelity for few-step image generation without relying on complex auxiliary objectives. Code is available at https://github.com/byliutao/cdm.