Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD
| Authors | Emiel Hoogeboom et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2603.20155 |
| Download | |
| Categories | cs.LG, cs.CV, stat.ML |
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.
Engineering Breakdown
Plain English
This paper solves the problem of compressing discrete diffusion models—which generate text and images by iteratively denoising—into faster student models that require far fewer sampling steps. Prior discrete distillation methods failed because they collapsed quality and diversity, whereas continuous diffusion models have well-established distillation techniques. The authors propose Discrete Moment Matching Distillation (D-MMD), which adapts moment matching ideas from continuous diffusion to the discrete setting and maintains both quality and diversity while dramatically reducing inference steps. Remarkably, the distilled student models sometimes outperform their teacher models, demonstrating that the approach not only speeds up inference but can actually improve the underlying model.
Core Technical Contribution
The core innovation is adapting Maximum Mean Discrepancy (MMD)—a powerful statistical technique for matching distributions—from continuous diffusion distillation into the discrete domain where previous methods failed. D-MMD works by directly matching the distributions of token sequences generated by teacher and student models across different timesteps, rather than trying to collapse individual tokens which caused prior methods to degrade. This is fundamentally different from prior discrete distillation attempts because it treats the full sequence distribution as the learning target rather than pointwise token predictions. The method elegantly sidesteps the discrete optimization challenges that plagued earlier approaches by operating at the distribution level using a well-understood kernel-based metric.
How It Works
The D-MMD algorithm operates by training a student discrete diffusion model to match the empirical distribution of sequences produced by a teacher model at various noise levels. At each training step, the student generates candidate sequences through its own denoising process while the teacher generates reference sequences; MMD then computes a distance metric between these two distributions without requiring explicit alignment of individual tokens. The key insight is computing this distance in a learned feature space rather than raw token space, allowing the student to capture the high-level distributional properties the teacher has learned. During inference, the student uses fewer denoising steps than the teacher (potentially just a handful), creating dramatic speedups. The algorithm uses importance weighting to account for the different number of steps between teacher and student, ensuring the distribution matching remains valid across different sampling schedules.
Production Impact
For production systems, D-MMD enables deploying discrete diffusion models (text generation, image generation) with 10-50x inference speedup compared to teachers, which directly reduces latency and compute costs in high-volume serving scenarios. Instead of requiring 1000 diffusion steps at inference, you might need only 20-50, cutting both latency (milliseconds per request) and GPU memory footprint dramatically. The fact that distilled models sometimes outperform teachers is particularly valuable: you get a speed-quality pareto improvement rather than the typical distillation trade-off. Integration requires retraining your model once (which is computationally expensive but one-time), then deploying the compact student; no changes needed to downstream serving infrastructure. The main trade-off is that you lose some of the iterative refinement properties of diffusion models (fewer steps means less opportunity for the model to correct mistakes), so this works best when you have sufficient steps (authors note this requirement) and when latency is a primary constraint.
Limitations and When Not to Use This
The paper requires sufficient sampling steps for D-MMD to work effectively—it's not a magical solution that works with 1-2 steps like some continuous distillation methods achieve, so you still need a meaningful student architecture with reasonable step counts. The method assumes you have access to a well-trained teacher model and the compute budget to run teacher inference during training, which creates a dependency and one-time cost. D-MMD's theoretical guarantees rely on the MMD metric being well-behaved in your feature space, which may not hold if your student architecture is drastically smaller than the teacher or if you're working with novel discrete domains beyond text and images. The paper doesn't deeply explore what happens when discrete and continuous tokens are mixed, or how to handle variable-length sequences in highly constrained settings, leaving practical edge cases for future work.
Research Context
This work directly builds on the success of moment matching and progressive distillation in continuous diffusion models (prior work by Song et al., Luhman & Luhman, and others), translating those techniques to the discrete setting where they were previously inapplicable. Discrete diffusion for text and images is an emerging area with papers like Hoogeboom et al.'s earlier discrete diffusion work and Antal et al.'s work on discrete image generation; D-MMD advances the practical utility of these models by making them deployable. The paper positions itself against prior discrete distillation attempts that tried to match individual token probabilities (which collapsed quality) and shows that distribution-level matching is the right inductive bias. This opens up a new research direction: optimizing the feature space and MMD kernel for discrete domains, combining D-MMD with other efficiency techniques like quantization, and extending it to multimodal and very large sequence models.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::