Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

1. Problem Statement

Discrete diffusion models (used for text and image generation) suffer from high computational costs because they require many sampling steps to converge to high-quality outputs. Unlike continuous diffusion models, which have established distillation techniques to reduce sampling steps to a handful, discrete diffusion distillation remains challenging.

• The Bottleneck: Current discrete distillation methods often fail because they rely on factorized token generation (modeling each token independently). This assumption leads to error accumulation during iterative sampling.
• The Limitation of Existing Methods: Previous attempts (e.g., SDTT, Di4C) either collapse in quality, fail to capture correlations between tokens, or are limited to specific diffusion types (like masked diffusion).
• The Goal: To develop a distillation method that reduces the number of function evaluations (NFEs) for discrete diffusion models while maintaining or improving sample quality and diversity, potentially allowing student models to outperform their teachers.

2. Methodology: Discrete Moment Matching Distillation (D-MMD)

The authors propose D-MMD, a generalization of the Moment Matching Distillation (MMD) framework from continuous to discrete domains.

Core Concept

Instead of minimizing the KL-divergence directly, D-MMD aims to match the conditional expectations (moments) of the clean data distribution between the teacher and the student.

• Adversarial Setup: The method employs an alternating optimization between a Student Generator ($\hat{x}_\eta$) and an Auxiliary Model ($\hat{x}_\phi$), guided by a fixed Teacher Model ($\hat{x}_\theta$).
• The Objective:
  • The Student tries to minimize the loss against the Teacher while maximizing the loss against the Auxiliary model.
  • The Auxiliary model tries to minimize the loss against the Student while staying close to the Teacher (regularization).
• Discrete Adaptation:
  • In continuous MMD, the gradient is derived from squared error. In discrete settings, direct sampling from a categorical distribution is non-differentiable.
  • Soft Targets: The authors use the soft probability vector $\hat{x}_\eta(z_t)$ rather than hard categorical samples for the generator's update. This allows for gradient flow.
  • Log-Probability Matching: The update rule effectively becomes a delta in log-probabilities between the Teacher and the Auxiliary model, rather than a direct difference in output space.
  • Handling Bias: For masked diffusion, soft targets work well. For uniform diffusion, the auxiliary model must be trained on hard samples to correct for bias.

Algorithm Flow (Algorithm 1)

1. Sample data $x$ and diffuse to $z_t$.
2. Generate soft probabilities $\hat{x}_\eta(z_t)$ from the student.
3. Sample a hard token $x$ (or use soft target for masked diffusion) and diffuse to a new state $z_s$.
4. Even Steps: Update Student parameters $\eta$ to minimize $L(\hat{x}_\eta, \hat{x}_\theta) - L(\hat{x}_\eta, \hat{x}_\phi)$.
5. Odd Steps: Update Auxiliary parameters $\phi$ to minimize $L(x, \hat{x}_\phi) + L(\hat{x}_\theta, \hat{x}_\phi)$.

Key Technical Insight: Factorized Generators & Correlation

A major theoretical hurdle is how a factorized model (which assumes token independence) can learn correlated outputs.

• The authors explain that the generator is a composition of two steps: generating "soft samples" (probabilities) and then sampling hard tokens.
• While the hard sampling step is factorized, the soft probability vector $\hat{x}_\eta(z_t)$ is not. To minimize the moment matching loss, the generator learns to reduce the entropy of its soft output distribution. This "collapse" of the factorized distribution allows the model to implicitly learn correlations between tokens, effectively acting as a mixture distribution.

3. Key Contributions

1. D-MMD Algorithm: A novel, generalizable distillation algorithm for discrete diffusion that works for both Masked Diffusion and Uniform Diffusion.
2. Teacher Outperformance: Demonstrates that distilled student models can achieve higher sample quality than their teachers using significantly fewer sampling steps (e.g., 16 steps vs. 1024 steps).
3. New Evaluation Metric (Gradient Moment): The paper introduces the Gradient Moment (GM) metric to evaluate discrete diffusion models.
   • Problem: Standard Perplexity is flawed for diffusion models because it can be "gamed" by low-temperature sampling (mode collapse) and does not reflect true distributional similarity.
   • Solution: GM measures the squared norm of the gradient of a reference LLM (e.g., GPT-2) on generated samples. If the samples are indistinguishable from the training data, the gradient is zero. This metric is robust against mode collapse.
4. Noise Conditioning: Identified that adding an input noise source (Gaussian noise) is crucial for masked diffusion distillation to allow the generator to collapse its output distribution effectively, significantly improving FID scores.

4. Experimental Results

The authors evaluated D-MMD on CIFAR-10 (images) and Open Web Text (OWT) (text).

• CIFAR-10 (Images):

  • Uniform Diffusion: D-MMD achieved an FID of 3.7 in 32 steps, outperforming the 1024-step teacher (FID 7.5).
  • Masked Diffusion: D-MMD achieved an FID of 3.5 in 64 steps, outperforming the teacher (FID 6.4).
  • D-MMD consistently outperformed existing methods like Di4C and SDTT with fewer NFEs.

• Text Generation (OWT):

  • Masked D-MMD: Outperformed the teacher with only 16 steps (GPT-2 GM: 0.236 vs. Teacher's 0.275 at 256 steps).
  • Uniform D-MMD: Also showed significant improvements over teachers.
  • Block Autoregressive: In a setup generating 256-token blocks, a 16-step D-MMD matched the performance of a 256-step teacher.

• Comparison with Literature:

  • D-MMD significantly outperformed SDTT and Di4C. SDTT showed degradation over repeated distillation rounds, whereas D-MMD maintained or improved quality.
  • The Gradient Moment metric revealed that methods with low perplexity often suffered from poor diversity (high collapse), whereas D-MMD balanced quality and diversity.

5. Significance and Conclusion

• Efficiency: D-MMD drastically reduces the inference cost of discrete diffusion models, making them competitive with autoregressive models in terms of latency while retaining the parallel sampling benefits of diffusion.
• Quality: The ability of student models to outperform teachers challenges the conventional wisdom that distillation only approximates the teacher. This is attributed to the adversarial nature of the training, which pushes the model toward modes (reverse-KL behavior) without fully collapsing.
• Generalization: By generalizing MMD to discrete spaces and handling both masked and uniform diffusion, this work provides a unified framework for accelerating discrete generative models.
• Metric Innovation: The introduction of the Gradient Moment metric offers a more reliable way to evaluate discrete generative models, addressing the limitations of perplexity in the context of diffusion-based generation.

In summary, D-MMD represents a breakthrough in making discrete diffusion models practical for real-world applications by solving the speed-quality trade-off and providing a robust theoretical and empirical framework for their distillation.