Discrete Moment Matching Distillation

• The paper introduces D-MMD, a framework that adversarially matches the first moments of teacher and student distributions, enabling efficient few-step generation while preserving quality.
• It overcomes issues like non-differentiability and diversity collapse by using cross-entropy-based losses and an alternating min–max training strategy.
• Empirical evaluations on CIFAR-10 and OpenWebText show that D-MMD achieves lower FID scores and better sample quality with significantly fewer sampling steps than previous methods.

Discrete Moment Matching Distillation (D-MMD) is a framework for distilling discrete diffusion models into efficient, high-quality generators by adversarially matching moments of the teacher’s and student’s output distributions. D-MMD generalizes continuous multistep moment matching approaches to the discrete domain, enabling few-step generation of text and images from discrete-valued data while avoiding the diversity collapse characteristic of naïve distillation strategies. It defines a min–max training regime between a student and auxiliary network, using cross-entropy-based losses on probability vectors (logits), and demonstrates empirical dominance over prior discrete model distillation approaches in both image and language settings (Hoogeboom et al., 20 Mar 2026).

1. Discrete Diffusion Models and Distillation Challenges

In discrete diffusion, a forward process $q$ iteratively corrupts a discrete signal $x$ (such as text tokens or quantized pixels), producing a noised state $z_t$. Two canonical instantiations are masked diffusion (gradual replacement with a MASK symbol [Austin et al. 2021]) and uniform diffusion (replacement with samples from a fixed prior $\pi$ [Hoogeboom et al. 2021]). The noising operation at time $t\in[0,1]$ is defined as

$$z_t \sim \text{Cat}( \alpha_t x + (1-\alpha_t)\pi )$$

where $\alpha_t$ is a noise schedule and $\text{Cat}$ denotes the categorical distribution.

The standard learning objective at each step $t$ is a weighted cross-entropy:

$$L_t^{\text{disc}}(x, \hat x_\theta(z_t), z_t) = w(t) \; \text{CE}( x \parallel \hat x_\theta(z_t) )$$

where $\hat x_\theta(z_t)$ learns to approximate $E_q[x|z_t]$.

Distillation of discrete diffusion models poses unique challenges: the sampling operation $x\sim\text{Cat}(\hat x_\eta(z_t))$ is non-differentiable, the output factorizes over dimensions, and naïve Kullback-Leibler or mean-squared matching collapses the student’s output diversity, often leading to mode dropping (Hoogeboom et al., 20 Mar 2026). In contrast, continuous diffusion models (with smooth, tractable reverse-time posteriors) permit direct gradient-based distillation via conditional moment matching (Salimans et al., 2024).

2. Theoretical Foundations of D-MMD

D-MMD frames the distillation objective as adversarial moment matching on discrete probability distributions (the logits emitted by the student and teacher networks). The discrete reverse-time posterior for $s<t$ is given by:

$$q(z_s \mid z_t, x) = \mathrm{Cat}\left( \frac{[\alpha_{t|s} z_t + (1-\alpha_{t|s})\pi^\top z_t] \odot [\alpha_s x + (1-\alpha_s)\pi]} {\alpha_t z_t^\top x + (1-\alpha_t) z_t^\top\pi} \right)$$

where $\alpha_{t|s} = \alpha_t/\alpha_s$.

The adversarial D-MMD objective optimizes a game between the student generator and an auxiliary network:

$$\mathcal{L}_{\mathrm{D\text{-}MMD}}(\eta,\phi) = \min_{\eta} \max_{\phi} \; \mathbb{E}_{g_{\eta}(z_t,x,s,z_s)} \left[ L_s(x, \hat x_\theta(z_s), z_s) - L_s(x, \hat x_\phi(z_s), z_s) \right]$$

where $L_s$ is the cross-entropy loss at time $s$. At Nash equilibrium, the student’s conditional output matches the teacher’s moments, driving the loss gap to zero (Hoogeboom et al., 20 Mar 2026).

Because discrete samples are not differentiable, D-MMD operates directly on the soft outputs (the student’s probability vectors), replacing sampling with expectation-based losses:

• Generator loss:

$$\mathcal{L}_{\mathrm{GEN}}(\eta) = \mathrm{CE}( \hat x_\eta(z_t) \parallel \hat x_\theta(z_s) ) - \mathrm{CE}( \hat x_\eta(z_t) \parallel \hat x_\phi(z_s) )$$

• Auxiliary loss:

$$\mathcal{L}_{\mathrm{AUX}}(\phi) = \mathrm{CE}(x \parallel \hat x_\phi(z_s)) + \mathrm{CE}( \hat x_\theta(z_s) \parallel \hat x_\phi(z_s) )$$

3. D-MMD Algorithm and Training Procedure

D-MMD employs an alternating min–max scheme between generator and auxiliary updates. Each training iteration comprises:

1. Sampling a time segment $(s, t)$ (with $t = \min(1, s + \delta_t)$, $\delta_t \sim \text{Uniform}(0, 1/k)$).
2. Obtaining a noised input $z_t$ from the data distribution via the forward process.
3. Computing student logits $p = \hat x_\eta(z_t)$ and sampling $x \sim \text{Cat}(p)$.
4. Sampling $z_s \sim q(z_s|x, z_t)$ using the exact discrete posterior.
5. Alternating between:
   • Generator update: $$\eta \leftarrow \eta - \nabla_\eta \left[ \text{CE}(p \parallel \hat x_\theta(z_s)) - \text{CE}(p \parallel \hat x_\phi(z_s)) \right]$$
   • Auxiliary update: $$\phi \leftarrow \phi - \nabla_\phi \left[ \text{CE}(x \parallel \hat x_\phi(z_s)) + \text{CE}(\hat x_\theta(z_s) \parallel \hat x_\phi(z_s)) \right]$$

This process matches the first moments (means) of the teacher’s and student’s trajectories for every segment, using the known reverse-time posteriors for discrete categorical variables. The same network backbones are used for student and auxiliary, with no architectural modifications (Hoogeboom et al., 20 Mar 2026).

4. Empirical Evaluation in Image and Text Generation

D-MMD was evaluated on both CIFAR-10 image and OpenWebText text datasets:

• CIFAR-10: 1024-step teacher diffusion models (uniform or masked) are distilled into students for $k \in \{4, 8, 16, 32, 64\}$ steps. Uniform D-MMD at $k=32$ achieves FID $=3.7$ (vs. teacher FID $=7.5$), with students dominating the FID-vs-step Pareto front. Masked D-MMD at $k=64$ achieves FID $=3.5$ (teacher FID $=6.4$) (Hoogeboom et al., 20 Mar 2026).
• OpenWebText: Students distilled for $k \in \{4, 8, 16, 32, 64, 128\}$ are evaluated by the GPT-2 Gradient Moment (GM) metric. Masked D-MMD with $k=32$ achieves GM $=0.225$ (teacher at $512$ steps GM $=0.275$), indicating improved sample quality and reduced computation.

Ablations reveal that input noise conditioning is critical for masked diffusion students (failure to add noise results in FID $\approx$ 151 at $k=4$, whereas proper noise gives FID $=22.3$). Selective teacher guidance via temperature or top-$p$ masking further enhances mode-seeking behavior (Hoogeboom et al., 20 Mar 2026).

5. Practical Recommendations and Implementation

Key considerations for applying D-MMD in practical scenarios include:

• Always add input noise when distilling masked diffusion students.
• Optionally employ temperature squashing or top-$p$ masking on teacher logits for improved mode coverage.
• Alternate generator and auxiliary updates on every step or minibatch.
• For text evaluation, use the GPT-2 Gradient Moment metric instead of perplexity.
• D-MMD is directly compatible with uniform, masked, and flow-matching discrete diffusion models and supports distillation into very few-step generators for both images and text (Hoogeboom et al., 20 Mar 2026).

Hyperparameter choices are dataset-dependent, e.g., learning rate $10^{-4}$ for images with batch size $256$, and $5 \times 10^{-5}$ for text with batch size $128$. The same loss weighting schedule $w(s)$ as teacher training is used for distillation.

6. Comparison to Related Distillation Approaches

D-MMD extends Multistep Moment Matching Distillation (MMD; (Salimans et al., 2024)) to the discrete domain by replacing score matching with probability vector matching and adversarial min–max optimization. Unlike naïve discrete distillation methods (e.g., SDTT, Di4C, DCD), which tend to collapse output diversity or drop modes, D-MMD retains both quality and diversity by explicitly matching the first moments of the teacher’s distribution.

Continuous MMD methods, such as those by Salimans et al., rely on the tractable Gaussian structure of continuous diffusion processes to define multistep conditional expectation losses; D-MMD generalizes this to discrete settings through cross-entropy and adversarial training (Salimans et al., 2024, Hoogeboom et al., 20 Mar 2026).

Empirical comparisons demonstrate that D-MMD consistently yields lower FID and GM scores with roughly an order of magnitude fewer sampling steps compared to prior discrete distillation methods and can outperform even the original teacher models when properly guided during training (Hoogeboom et al., 20 Mar 2026).

7. Theoretical Properties and Limitations

Matching first moments under a factorized generator is theoretically sufficient to recover the joint output distribution in the infinitesimal step-size limit ($dt\rightarrow 0$). D-MMD’s min–max formulation overcomes non-differentiability, prevents degeneracies in the student distribution, and exploits the implicit mode-seeking behavior of the adversarial loss to improve over the teacher. The method is, however, dependent on careful noise conditioning and hyperparameter tuning—especially for masked diffusion settings. Empirical evidence confirms that first-moment matching is typically adequate, although extension to higher-order moments (e.g., covariances) is possible in principle (Hoogeboom et al., 20 Mar 2026).

---

References:

• "Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD" (Hoogeboom et al., 20 Mar 2026)
• "Multistep Distillation of Diffusion Models via Moment Matching" (Salimans et al., 2024)