Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

ArXiv: 2603.20155

🎯 Pitch

Distilling discrete diffusion models was notoriously hard—previous attempts collapsed. D-MMD shows that adversarial moment matching not only makes distillation possible but lets the student vastly outperform the teacher: a 64-step image generator hits FID 3.5 vs. the teacher’s 6.4 at 1024 steps, and a 16-step text student beats a 256-step teacher.

---

1. Executive Summary

This paper introduces Discrete Moment Matching Distillation (D-MMD), a method for distilling discrete diffusion models into few-step generators by generalizing the continuous Moment Matching Distillation framework to discrete spaces. Evaluated on CIFAR-10 image generation and OpenWebText language modeling using PaLM 2-style discrete diffusion teachers, D-MMD operates through an adversarial min-max optimization between a student generator and an auxiliary model — the student minimizes a teacher loss while maximizing an auxiliary loss, and the auxiliary is trained to match both the student's outputs and the teacher's distribution — using soft probability vectors rather than hard categorical samples to maintain gradient flow. The distilled generators substantially outperform their teachers: masked diffusion on CIFAR-10 achieves an FID of 3.5 with 64 steps versus the teacher's 6.4 with 1024 steps, and text generation reaches a GPT-2 Gradient Moment of 0.236 with 16 steps versus the teacher's 0.275 with 256 steps, establishing that stochastic discrete distillation can simultaneously reduce sampling steps by over an order of magnitude while improving sample quality — a finding that holds for both masked and uniform diffusion processes, and generalizes to block-autoregressive settings where a 16-step student matches a 256-step teacher.

2. Context and Motivation

The Core Problem: Discrete Diffusion Models Are Slow to Sample

The fundamental problem this paper tackles is that discrete diffusion models require an impractically large number of sampling steps to produce high-quality generations. Unlike continuous diffusion models — which have been successfully distilled down to a handful of steps through methods like progressive distillation (Salimans and Ho, 2022), consistency models (Song et al., 2023), and moment matching distillation (Salimans et al., 2024) — discrete diffusion models have resisted similar distillation efforts. The paper states this starkly in its opening line:

"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."

This is not merely a matter of inconvenience. The number of sampling steps directly determines the computational cost of generation. As the paper explains, discrete diffusion models operate on blocks of tokens simultaneously, giving them high accelerator utilization — a genuine advantage over autoregressive models that process tokens sequentially. However, this advantage is undermined by the sheer number of iterations needed:

"these models tend to need many iterations to converge to a reasonable generation, leading to high computing costs and strictly higher FLOPs. The fewer iterations one takes, the lower the cost."

The tension is clear: discrete diffusion has architectural efficiency (parallel token processing) but requires superlinear computational effort due to extreme step counts. Closing this gap — reducing steps from hundreds to tens while preserving quality — would make discrete diffusion models practically competitive with autoregressive models for the first time.

Why This Problem Matters: The Deeper Technical Challenge

The difficulty of distilling discrete diffusion is not an implementation detail — it stems from a fundamental modeling constraint that the paper identifies with precision. Discrete diffusion models parameterize the denoising process in a factorized manner: at each step, each token is predicted independently conditioned on the previously generated tokens. The paper explains the consequence:

"As a result, errors from this assumed independence accumulate during the sampling iterations."

This independence assumption creates a compounding error problem. When you take many small steps, the model can gradually correct these errors — each individual step only needs to make a locally reasonable prediction, and the iterative refinement smooths out inconsistencies. But when you try to take fewer, larger steps, the factorized predictions must be substantially correct in a single shot, and the independence assumption means there is no mechanism for coordinating predictions across tokens to ensure global coherence.

This reveals why discrete distillation is fundamentally harder than continuous distillation. In continuous diffusion, the model outputs a vector of real numbers (e.g., pixel values or latent features), and errors in individual dimensions can be corrected by subsequent steps without any structural barrier. In discrete diffusion, the model outputs a distribution over discrete tokens, and the factorization means the model cannot directly represent correlations between tokens in a single step. Any distillation method must either break this factorization (as Di4C attempts with mixture distributions), work around it (as D-MMD does through adversarial optimization), or accept degraded quality.

The practical stakes are substantial. If discrete diffusion models could be sampled in 16-32 steps while maintaining quality, their hardware utilization advantages over autoregressive models would translate into genuine throughput gains. At current step counts (256-1024), the FLOP advantage of parallel token processing is consumed by iteration count. Effective distillation would tip this balance, making diffusion a viable alternative to autoregressive generation for the first time — not just in research benchmarks but in production settings where cost-per-token matters.

Prior Approaches and Where They Fall Short

The paper situates its contribution against four categories of prior work, each with identifiable limitations:

Deterministic Distillation Approaches (Progressive Distillation, Consistency Models)

The earliest distillation methods for diffusion models were deterministic, operating on the probability flow ODE. Progressive distillation (Salimans and Ho, 2022; Meng et al., 2022) iteratively halves the number of sampling steps by training a student to match two teacher steps in one. Consistency models (Song et al., 2023) take a different approach: the generator is trained to map any point along the diffusion trajectory directly to the clean data, enforcing self-consistency — the mapping from $z_t$ should equal the mapping from $z_{t-dt}$ after applying the denoising step.

Both approaches have been developed primarily for continuous diffusion. The paper notes their extension to discrete data: "flow-map or consistency-based distillation approaches have been applied to discrete data lifted to continuous space with standard diffusion models (Sahoo et al., 2025; Roos et al., 2026; Lee et al., 2026)." However, these methods operate on discrete data that has been embedded into continuous space, not on genuine discrete diffusion processes. The paper expresses skepticism about this approach:

"Currently, it remains to be seen whether these continuous models on discrete data can match the performance of discrete diffusion models. Furthermore, for both model classes it remains to be seen whether they can match the performance of standard autoregressive models."

This is a pointed observation: lifting discrete data to continuous space loses the natural structure of the discrete problem (the categorical nature of tokens, the factorization properties) and may impose performance ceilings that native discrete methods can exceed.

SDTT (Self-Distillation Through Time)

SDTT (Deschenaux and Gulcehre, 2025) represents the most direct prior attempt at discrete diffusion distillation. It takes an approach "reminiscent of progressive distillation but applied to discrete sampling." The paper identifies a fundamental limitation:

"Although the approach tends to produce improvements to limited degree, it is fundamentally limited. For example, perfectly correlated coin tosses of two coins cannot be approximated with a single step of this approach."

This example cuts to the heart of the problem. Imagine two coins that always land on the same side — a perfectly correlated pair. A factorized model that predicts each coin independently cannot represent this distribution in a single step: it can output (0.5, 0.5) for each coin, but the joint distribution over (heads, heads) and (tails, tails) requires coordination. SDTT's progressive distillation framework inherits this limitation because it operates within the factorized output space. The paper states the consequence bluntly:

"Due to the divergences chosen, SDTT will overcome the above mentioned limitation by directly dropping modes to achieve sampling speedups."

In other words, SDTT sacrifices diversity for speed. When it cannot represent the full distribution, it collapses to a subset of modes rather than producing a degraded approximation of the full distribution. This is a practical tradeoff — faster sampling at the cost of coverage — but it is a limitation of the method, not a fundamental property of discrete distillation.

The experimental results in Table 5 bear this out. SDTT (reimplementation) achieves a GPT-2 GM of 0.293 at 64 steps, which actually improves on the teacher's 0.307 at 64 steps. But this is misleading: as measured by GPT-2 GM, SDTT degrades over repeated distillation rounds. The paper notes that "even though SDTT improves upon the teacher model, it still degrades over repeated distillation rounds and is outperformed by D-MMD." The improvement at 64 steps likely reflects the mode-dropping behavior — the model produces samples that look better by some metrics (generative perplexity) but are less diverse and less representative of the full data distribution.

Di4C (Distillation of Discrete Diffusion through Dimensional Correlations)

Di4C (Hayakawa et al., 2024) takes a more direct approach to the factorization problem. Rather than accepting factorized outputs, it extends the model to output mixture distributions, allowing the model to represent correlations between tokens explicitly. The intuition is sound: a mixture of factorized distributions can represent correlations that any single factorized distribution cannot.

However, the paper identifies a scaling problem:

"Although effective to some degree, they tend to be limited in effect. One is often fighting an exponential of correlations between all tokens, and therefore the number of required mixtures also grows exponentially."

For a sequence of $D$ tokens, the number of possible joint configurations is exponential in $D$. A mixture model with $K$ components can represent at most $K$ distinct correlation patterns. As sequence length grows, the number of components needed to cover the meaningful correlations in the data grows combinatorially. Di4C provides partial relief — it can represent some correlations — but it cannot fundamentally escape the exponential scaling without an exponentially large mixture.

The CIFAR-10 comparison in Table 4 quantifies the gap. Di4C achieves an FID of 20.6 with 10 steps, 9.5 with a hybrid approach at 20 steps, and its teacher achieves 8.0 with 40 steps. D-MMD, in contrast, achieves 5.0 with 8 steps — dramatically better than Di4C at fewer steps. The comparison isn't perfectly apples-to-apples because Di4C uses a teacher trained with a different process (one that mimics Gaussian destruction), but the relative improvement is stark enough to suggest a qualitative difference in approach.

DiMO (Distilling Masked Diffusion Models into One-Step Generator)

DiMO (Zhu et al., 2025) is the closest precursor to D-MMD. The paper acknowledges this explicitly:

"Although derived differently via straight-through softmax sampling, the resulting algorithm is equivalent to the implementation of D-MMD for the one-step case."

This is a significant admission — D-MMD generalizes DiMO rather than replacing it. The contributions over DiMO are: (1) support for few-step generators, not just one-step; (2) generalization to other discrete processes beyond masked diffusion (specifically uniform diffusion); (3) application to text generation, where DiMO only demonstrated image generation.

The most important extension is the few-step capability. One-step generation is extremely aggressive — it asks the model to reverse the entire diffusion process in a single shot. While DiMO showed this is possible for image token generation, the paper's few-step results demonstrate that intermediate step counts (4-64 steps) provide a more practical operating point, balancing speed and quality. The text generation results in particular benefit from multiple steps: the 4-step D-MMD generator achieves a GPT-2 GM of 0.820, which improves to 0.236 at 16 steps and 0.225 at 32 steps (Table 2).

Concurrent Work: IDLM

The paper notes a concurrent submission, IDLM (Li et al., 2026), which proposes a similar framework with one key difference:

"The difference with IDLM is that the training algorithm generates the full $x$ and diffuses back to $z_t$, whereas our work samples from the posterior $q(z_s|z_t, x)$."

This is a technical distinction in how the intermediate states are sampled during training. Sampling from the posterior is computationally cheaper (it doesn't require forward diffusion from the clean sample) and more tightly couples the noisy state to the current generation, which may improve gradient signal. The paper views IDLM as complementary rather than competing.

The Broader Landscape: A Gap in the Distillation Toolkit

The paper identifies a clear pattern in the distillation literature: stochastic distillation outperforms deterministic distillation in few-step regimes. The paper states:

"When the generator is single-step, MMD (Salimans et al., 2024) is equivalent to the distribution matching approaches, but it tends to outperform them in few-step regimes."

This is an important empirical observation. Deterministic methods (progressive distillation, consistency models) commit to a specific trajectory through the diffusion process. Stochastic methods (distribution matching, MMD) optimize the generator's distribution to match the teacher's distribution without committing to a specific trajectory, which provides additional flexibility — the student can find a path through the diffusion process that is easier to learn than the teacher's path.

However, all of this insight comes from continuous diffusion. The discrete case has been largely untouched by stochastic distillation methods, for a simple reason: discrete sampling breaks gradient flow. In continuous diffusion, the generator outputs real-valued predictions that can be differentiated through directly. In discrete diffusion, the generator outputs a probability vector, and the natural next step — sampling a categorical token from that vector — is non-differentiable. This blocks the straightforward application of gradient-based stochastic distillation to discrete models.

D-MMD's core technical contribution is resolving this gradient flow problem while preserving the benefits of stochastic distillation. By using soft probability vectors ($\hat{x}_\eta(z_t)$) rather than hard categorical samples ($x \sim \text{Cat}(\hat{x}_\eta(z_t))$) in the generator loss, and carefully handling where hard samples are still needed (for the auxiliary model and posterior sampling), D-MMD maintains end-to-end differentiability for the generator while still operating in a discrete space.

Where the Paper Positions Itself

The paper explicitly frames D-MMD as a generalization of continuous MMD, not a completely new method:

"Our paper generalizes the formulation of Moment Matching Distillation (MMD) (Salimans et al., 2024) so that it can be used in more general settings. As our main focus is to distill discrete diffusion processes, we call this new algorithm Discrete-MMD (D-MMD)."

This positioning is both honest and strategic. By grounding D-MMD in the well-understood MMD framework, the paper inherits the theoretical guarantees and empirical track record of MMD while extending it to a new domain. The generalization is not trivial — it requires reformulating the loss in terms of cross-entropy rather than squared error, handling the non-differentiability of discrete sampling, and developing a mechanism for the generator to represent correlated outputs despite its factorized architecture — but the conceptual lineage is clear.

The paper's experimental positioning is also notable. Rather than claiming to beat autoregressive models (which it does not — the AR baseline in Table 2 achieves a GPT-2 GM of 0.061, far better than D-MMD's best of 0.225), the paper focuses on improving the Pareto frontier of discrete diffusion models. The goal is to show that discrete diffusion distillation is possible and effective, establishing a foundation that future work can build upon to close the remaining gap with autoregressive models. The block-autoregressive experiment in Section 6.3 hints at this trajectory: combining autoregressive context with diffusion-based block generation is where the practical value of discrete diffusion distillation may ultimately lie.

Finally, the paper positions the GPT-2 Gradient Moment metric as a contribution in its own right, motivated by the observation that standard evaluation metrics for discrete diffusion models are fundamentally flawed. Generative perplexity — the dominant metric in prior work — reports the perplexity of a reference AR model on generated samples. The paper argues this is deeply problematic:

"high density samples are often not typical... meaning that they are not actually similar to the data. An example failure case of the generative perplexity metric is assigning a good score to ungrammatical generated samples that feature many repeated words."

This is not just a theoretical concern. Figure 3 demonstrates concretely that perplexity and the proposed gradient moment metric diverge: as temperature or top-p decreases, generative perplexity continues improving (the samples become "easier" for GPT-2 to predict) while the gradient moment eventually degrades (the samples diverge from the data distribution). This demonstrates that generative perplexity can be "gamed" through mode-collapsing sampling strategies, exactly the behavior that distillation methods like SDTT exhibit. The gradient moment metric, by measuring whether an AR model would update its parameters on the generated data, provides a more principled measure of distributional similarity that is robust to mode collapse.

3. Technical Approach

3.1 Reader Orientation

This is a distillation paper — it teaches us how to take a slow, multi-step discrete diffusion model (the teacher) and train a much faster model (the student) that produces equally good or better samples using far fewer steps. The system is a training algorithm: given a trained teacher diffusion model, D-MMD produces a generator that can be sampled in 4–64 steps instead of 256–1024 steps, while actually improving sample quality as measured by FID (for images) or GPT-2 Gradient Moment (for text).

The core idea is an adversarial min-max game between three players: a student generator, a fixed teacher model, and a learnable auxiliary model. The student tries to produce samples that the teacher would classify as high-quality, while simultaneously trying to fool the auxiliary model (maximizing its loss). The auxiliary model, in turn, is trained to distinguish the student's outputs from the teacher's distribution. At equilibrium — when the auxiliary model cannot tell the difference and equals the teacher — the student generates samples from the teacher's distribution, but does so in fewer steps.

3.2 Big-Picture Architecture (Diagram in Words)

The D-MMD training loop has four major components:

1. The Teacher Model ($\hat{x}_\theta$): A fully trained discrete diffusion model (masked or uniform) that takes a noised sample $z_t$ at timestep $t$ and outputs a denoising prediction $\hat{x}_\theta(z_t)$, which approximates $\mathbb{E}_q[x|z_t]$ — the expected clean data given the noisy observation. This model is frozen during distillation and provides the target distribution.

2. The Student Generator ($g_\eta$, with predictions $\hat{x}_\eta(z_t)$): A few-step generator that takes a noised sample $z_t$ and produces a soft probability vector $\hat{x}_\eta(z_t)$ over possible clean tokens. Crucially, this is not a single denoising step — it is a learned generator that may internally take multiple steps, but is trained end-to-end to map noisy states to clean predictions much faster than the teacher would.

3. The Auxiliary Model ($\hat{x}_\phi$): A learnable model trained to match the student's output distribution. It takes a partially-denoised sample $z_s$ (where $s < t$, meaning it's cleaner than the input to the student) and predicts $\hat{x}_\phi(z_s)$. Its role is adversarial — it provides a loss signal that prevents the student from collapsing to a degenerate distribution.

4. The Posterior Sampler: Given a hard categorical sample $x \sim \text{Cat}(\hat{x}_\eta(z_t))$ drawn from the student's soft predictions, this samples an intermediate noisy state $z_s$ from the diffusion posterior $q(z_s|z_t, x)$. This creates a bridge between the student's current output and a slightly cleaner state, which the teacher and auxiliary model then evaluate.

Information flow (one training step):

• A noisy state $z_t$ is produced by diffusing a real data point, or sampled from the stationary distribution.
• The student generator processes $z_t$ and outputs $\hat{x}_\eta(z_t)$ — a soft probability vector.
• A hard sample $x$ is drawn from $\text{Cat}(\hat{x}_\eta(z_t))$ (used for the auxiliary model only).
• The posterior sampler produces $z_s \sim q(z_s|z_t, x)$ — a state between $z_t$ and the clean data, consistent with the sampled $x$. A stop-gradient is applied to $z_s$ with respect to the student.
• The teacher evaluates $\hat{x}_\theta(z_s)$ on the intermediate state.
• The auxiliary model evaluates $\hat{x}_\phi(z_s)$ on the same state.
• On even training steps: the student loss $L_{\text{GEN}}(\eta)$ is computed by comparing $\hat{x}_\eta(z_t)$ against both the teacher's prediction and the auxiliary model's prediction, and the student parameters $\eta$ are updated.
• On odd training steps: the auxiliary loss $L_{\text{AUX}}(\phi)$ is computed by comparing the auxiliary model's prediction $\hat{x}_\phi(z_s)$ against both the hard sample $x$ and the teacher's prediction $\hat{x}_\theta(z_s)$, and the auxiliary parameters $\phi$ are updated.

This alternation creates the adversarial dynamic: the student learns to produce outputs the teacher assigns high probability to but the auxiliary assigns low probability to, while the auxiliary learns to better match whatever distribution the student is currently generating.

3.3 Roadmap for the Deep Dive

I will walk through the technical content in this order:

• First, the generalized D-MMD objective (Equation 9), because it is the mathematical foundation that everything else instantiates — understanding the min-max formulation reveals why the adversarial dynamic works and what fixed point it converges to.

• Second, the equivalence to continuous MMD (Equations 10–12), which shows that D-MMD is a genuine generalization, not a new method, and clarifies why the cross-entropy form is the natural discrete analog of the squared-error form.

• Third, the discrete instantiation — cross-entropy losses and soft probability vectors (Equations 11–12), including the critical detail of when hard samples are required versus when soft probabilities suffice, and why masked diffusion permits a simplification not available to uniform diffusion.

• Fourth, how a factorized generator can learn correlated outputs (Section 3.1), because this is the most counterintuitive claim in the paper — understanding the two-step sampling mechanism and the entropy collapse phenomenon explains how D-MMD circumvents the exponential correlation problem that limits Di4C.

• Fifth, the bias correction for the auxiliary model (Section 3.2), which addresses a subtle distribution shift between the student's soft predictions and the hard samples used for posterior sampling.

• Sixth, temperature and top-p distillation (Section 3.3), which extends the framework to incorporate mode-seeking sampling strategies that improve practical generation quality.

• Seventh, the noise conditioning mechanism (Section 6.5), which interacts with the factorized architecture to enable entropy collapse in masked diffusion generators.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodology paper whose core idea is that the Moment Matching Distillation framework can be generalized from continuous diffusion (where the loss is squared error) to discrete diffusion (where the loss is cross-entropy) by reformulating the alternating optimization as a min-max game over arbitrary loss functions, and then instantiating that game with soft probability vectors to maintain gradient flow through the otherwise non-differentiable discrete sampling step.

---

The Generalized D-MMD Objective: A Min-Max Formulation

The paper's key mathematical insight is that the alternating optimization used in continuous MMD (Equations 7 and 8 in the paper) can be rewritten as a single min-max objective that is agnostic to the specific loss function. This generalization is what allows D-MMD to work with cross-entropy (discrete) just as MMD works with squared error (continuous).

The generalized objective is:

$L_{\text{D-MMD}}(\eta) = \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) - L_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s) \right]$

where $L_s(\cdot, \cdot, z_s)$ is an arbitrary loss function that measures the discrepancy between a target (first argument) and a prediction (second argument), conditioned on the noisy state $z_s$. The terms are:

• $L_s(x, \hat{x}_\theta(z_s), z_s)$: The loss of the teacher's prediction on the intermediate state, evaluated against the hard sample $x$ drawn from the student. The student minimizes this term — it wants the teacher to score its samples well.

• $-L_s(x, \hat{x}_\phi(z_s), z_s)$: The negative loss of the auxiliary model on the same sample. The student minimizes this (which means maximizing the auxiliary's loss, since there's a minus sign) — it wants to fool the auxiliary model, making it assign poor likelihood to the student's outputs.

• $-L_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s)$: A regularization term that penalizes the auxiliary model for deviating from the teacher. This term does not depend on the student's outputs at all — it only affects the auxiliary's optimization. It ensures the auxiliary stays close to the teacher, preventing degenerate solutions where the auxiliary drifts arbitrarily far just to maximize the student's loss.

What this computes: The expectation is taken over $g_\eta(z_t, x, s, z_s)$, which is the joint distribution induced by: (1) sampling a noisy state $z_t$ (diffused from data or drawn from the stationary distribution), (2) having the student generate soft predictions $\hat{x}_\eta(z_t)$, (3) sampling a hard $x \sim \text{Cat}(\hat{x}_\eta(z_t))$, (4) sampling a time $s$ and noise offset $\delta_t$ (see Algorithm 1), and (5) sampling $z_s \sim q(z_s|z_t, x)$ from the posterior. The inner expression evaluates three loss terms on this trajectory. The outer min-max says: the student $\eta$ chooses parameters to minimize this expression (making the teacher loss small and the auxiliary loss large), while the auxiliary $\phi$ chooses parameters to maximize it (making its own loss small, i.e., better matching the student's distribution, while staying regularized toward the teacher).

Why this form: This formulation is loss-function-agnostic — it makes no assumptions about whether $L_s$ is squared error, cross-entropy, KL divergence, or any other discrepancy measure. This is precisely what enables generalization from continuous to discrete settings. In continuous MMD, $L_s$ is the squared error $\|\cdot\|^2$, and the min-max reduces to the alternating optimization in Equations 7–8. In discrete D-MMD, $L_s$ becomes cross-entropy, and the same min-max structure yields a valid distillation algorithm. The key property is that the fixed point of this min-max game — where neither player can improve — occurs when the student's distribution exactly matches the teacher's distribution, at which point $\hat{x}_\phi(z_s) = \hat{x}_\theta(z_s)$ and the loss equals zero. The regularization term $L_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s)$ is essential because without it, the auxiliary could trivially maximize the student's loss by outputting arbitrary predictions far from the teacher, providing no useful training signal. With the regularization, the auxiliary is anchored to the teacher and can only deviate in ways that genuinely reflect differences between the student's distribution and the teacher's.

---

Equivalence to Continuous MMD: The Cross-Entropy Generalization

To demonstrate that D-MMD is a genuine generalization rather than a new method, the paper shows that substituting squared-error loss into the generalized objective recovers the exact gradients of continuous MMD (up to a constant factor).

Let $L_s(x, \hat{x}, z_s) = w(s)\|x - \hat{x}\|^2$ (the continuous case with a time-dependent weight $w(s)$), and let $x_\eta = \hat{x}_\eta(z_t)$ be the student's prediction. The gradient of the D-MMD objective with respect to student parameters $\eta$ is:

$\nabla_\eta L_{\text{D-MMD}}(\eta) = \nabla_\eta \left[ L_s(x_\eta, \hat{x}_\theta(z_s), z_s) - L_s(x_\eta, \hat{x}_\phi(z_s), z_s) \right]$

Expanding the squared error and differentiating:

$\nabla_\eta L_{\text{D-MMD}}(\eta) = 2 w(s) \frac{d\hat{x}_\eta}{d\eta} \left( \hat{x}_\phi(z_s) - \hat{x}_\theta(z_s) \right)$

where $\frac{d\hat{x}_\eta}{d\eta}$ is the Jacobian of the student's output with respect to its parameters, $\hat{x}_\phi(z_s)$ is the auxiliary model's prediction on the intermediate state, and $\hat{x}_\theta(z_s)$ is the teacher's prediction.

What this computes: The gradient direction for the student update. The factor $(\hat{x}_\phi(z_s) - \hat{x}_\theta(z_s))$ is a vector difference: if the auxiliary model thinks the clean data should be higher in some dimension than the teacher does, the student is pushed to increase its output in that dimension (moving away from the auxiliary, toward the teacher). If the teacher thinks the clean data should be higher, the student is pushed to increase its output (moving toward the teacher, away from the auxiliary). The $2 w(s)$ factor scales the update by the time-dependent weight. The Jacobian $\frac{d\hat{x}_\eta}{d\eta}$ translates this output-space direction into parameter-space updates.

Why this matches MMD: The continuous MMD gradient (from Salimans et al., 2024, Equation 7) is:

$\nabla_\eta L_{\text{MMD}}(\eta) = w(s) \frac{d\hat{x}_\eta}{d\eta} \text{sg}\left( \hat{x}_\phi(z_s) - \hat{x}_\theta(z_s) \right)$

The two gradients are identical up to a factor of 2, which is absorbed by the learning rate. The stop-gradient $\text{sg}(\cdot)$ in MMD — which prevents gradients from flowing through $z_s$ back to $\eta$ — is implicitly present in D-MMD because the paper applies stop-gradient to $z_s$ in Algorithm 1. The equivalence holds because: (1) the squared error loss is symmetric and its gradient with respect to the first argument is proportional to the difference between the arguments, and (2) the stop-gradient on $z_s$ makes the dependence of the auxiliary and teacher predictions on $\eta$ zero, so the differentiation only passes through $x_\eta$.

This equivalence is not merely a mathematical curiosity — it validates the generalization. By showing that D-MMD reduces to MMD in the continuous case, the paper establishes that the min-max formulation is not a new algorithm but a more general description of the same underlying principle. The power is that this description reveals how to extend the principle to new loss functions: simply substitute the appropriate $L_s$ for the new domain.

---

The Discrete Instantiation: Cross-Entropy and Soft Probabilities

The critical challenge in discrete diffusion is that the natural loss function — cross-entropy between predicted and true token distributions — requires categorical samples $x$, but sampling from a categorical distribution is non-differentiable with respect to the distribution parameters. The paper resolves this through a careful distinction: the student generator uses soft probability vectors ($\hat{x}_\eta(z_t)$) in its own loss, while the auxiliary model is trained on hard samples ($x \sim \text{Cat}(\hat{x}_\eta(z_t))$) drawn from those probabilities.

Student Generator Loss (soft targets):

Substituting cross-entropy for $L_s$ in the D-MMD objective and using the soft probability vector $\hat{x}_\eta$ (a vector of probabilities over the vocabulary for each token position) directly in place of the hard sample $x$, the generator loss simplifies to:

$L_{\text{GEN}}(\eta) = \text{CE}\left( \hat{x}_\eta \| \hat{x}_\theta(z_s) \right) - \text{CE}\left( \hat{x}_\eta \| \hat{x}_\phi(z_s) \right)$

where $\text{CE}(p\|q) = -\sum_c p_c \log q_c$ is the cross-entropy from distribution $p$ to distribution $q$, summed over all categories $c$ in the vocabulary.

Expanding the cross-entropy:

$L_{\text{GEN}}(\eta) = -\sum_c (\hat{x}_\eta)_c \left( \log \hat{x}_\theta(z_s) - \log \hat{x}_\phi(z_s) \right)_c$

where $(\hat{x}_\eta)_c$ is the student's predicted probability for category $c$, $\log \hat{x}_\theta(z_s)_c$ is the log-probability the teacher assigns to category $c$ given the intermediate state, and $\log \hat{x}_\phi(z_s)_c$ is the log-probability the auxiliary model assigns.

What this computes: For each token position and each category $c$, if the teacher assigns higher log-probability to $c$ than the auxiliary does (i.e., $\log \hat{x}_\theta(z_s)_c > \log \hat{x}_\phi(z_s)_c$), then the term $(\hat{x}_\eta)_c (\log \hat{x}_\theta - \log \hat{x}_\phi)_c$ is positive, and the student is penalized for assigning probability to $c$ — the gradient will push $(\hat{x}_\eta)_c$ downward. Conversely, if the auxiliary assigns higher log-probability than the teacher, the term is negative, and the student is rewarded for assigning probability to $c$ — the gradient pushes $(\hat{x}_\eta)_c$ upward. The net effect is that the student shifts probability mass toward categories the teacher favors (relative to the auxiliary) and away from categories the auxiliary favors (relative to the teacher). The soft probability vector $\hat{x}_\eta$ provides the weight for each category's contribution to the gradient.

Why this works with soft targets: The cross-entropy $\text{CE}(\hat{x}_\eta \| \hat{x}_\theta(z_s))$ is differentiable with respect to $\hat{x}_\eta$ (the first argument) even though $\hat{x}_\eta$ itself is a probability distribution output by a neural network. The gradient flows through the softmax layer that produces $\hat{x}_\eta$ from logits, through the cross-entropy computation, and back to the generator parameters. This completely avoids the non-differentiable categorical sampling step. The key insight is that for the generator's objective, we don't need a hard sample — we need the generator's distribution to match the teacher's distribution in expectation, and comparing soft probabilities directly achieves this.

Auxiliary Model Loss (hard samples):

The auxiliary model, by contrast, must be trained on hard samples $x$ drawn from the student's distribution. Its loss is:

$L_{\text{AUX}}(\phi) = \text{CE}\left( x \| \hat{x}_\phi(z_s) \right) + \text{CE}\left( \hat{x}_\theta(z_s) \| \hat{x}_\phi(z_s) \right)$

Expanding:

$L_{\text{AUX}}(\phi) = -\sum_c \left( x + \hat{x}_\theta(z_s) \right)_c \log \hat{x}_\phi(z_s)_c$

where $x$ is a one-hot vector (a hard categorical sample from the student's soft probabilities), $\hat{x}_\theta(z_s)$ is the teacher's soft prediction, and $\hat{x}_\phi(z_s)_c$ is the auxiliary model's predicted probability for category $c$.

What this computes: The auxiliary model is trained to maximize the log-probability it assigns to the hard sample $x$ (the first cross-entropy term) while also staying close to the teacher's predictions (the second cross-entropy term). At the optimum, the auxiliary's prediction $\hat{x}_\phi(z_s)$ should be the average of the student's expected output and the teacher's expected output: $\hat{x}_\phi(z_s) = \frac{1}{2}(\mathbb{E}_{g_\eta}[x|z_s] + \hat{x}_\theta(z_s))$. When the student's distribution matches the teacher's, $\mathbb{E}_{g_\eta}[x|z_s] = \hat{x}_\theta(z_s)$, and the optimum becomes $\hat{x}_\phi(z_s) = \hat{x}_\theta(z_s)$ — the auxiliary converges to the teacher.

Why hard samples are necessary for the auxiliary: The auxiliary model needs to learn $\mathbb{E}_{g_\eta}[x|z_s]$ — the expected clean data under the student's distribution given the intermediate state $z_s$. This expectation requires the actual distribution of $x$ values produced by the student, not just the soft probabilities $\hat{x}_\eta$. The soft probabilities represent the student's belief about the clean data, but the actual samples $x \sim \text{Cat}(\hat{x}_\eta)$ are what the auxiliary will encounter at inference time. Training the auxiliary on soft targets $\hat{x}_\eta$ would be training on a different distribution than what it needs to evaluate, creating a mismatch. The hard samples ensure the auxiliary learns to predict the outcomes of the student's sampling process, not just the student's internal probabilities.

The fixed point of the discrete algorithm: At equilibrium, three conditions hold simultaneously:

1. $\mathbb{E}_{g_\eta}[x|z_s] = \hat{x}_\phi(z_s)$ — the auxiliary model correctly predicts the expected clean data under the student's distribution for all intermediate states $z_s$.

2. $\hat{x}_\phi(z_s) = \hat{x}_\theta(z_s)$ — the auxiliary model equals the teacher for all $z_s$ (enforced by the regularization term when the student's distribution matches the teacher's).

3. Therefore $\mathbb{E}_{g_\eta}[x|z_s] = \hat{x}_\theta(z_s)$ — the student's expected output matches the teacher's denoising prediction for all $z_s$.

When condition 3 holds for all $z_s \in [0, 1]$, Appendix A proves that this is sufficient for the student's sampling distribution to exactly equal the data distribution: $p_\eta(x) = q(x)$. The proof relies on the fact that the teacher's prediction $\hat{x}_\theta(z_s) = \mathbb{E}_q[x|z_s]$ fully characterizes the reverse diffusion process in the continuous-time limit, so matching these conditional expectations at every timestep guarantees matching the marginal distribution over $x$.

---

How a Factorized Generator Can Learn Correlated Outputs

This section addresses the most counterintuitive claim in the paper: how can a generator with a factorized output distribution — where each token is predicted independently given the noisy input — generate coherent, correlated sequences? The answer lies in the two-step sampling process and the resulting entropy collapse phenomenon.

The two-step mechanism: The student generator $g_\eta$ does not directly produce a sample. Instead, it operates in two stages:

1. Soft sampling: The generator processes the noisy state $z_t$ and outputs a soft probability vector $\hat{x}_\eta(z_t)$. This step is stochastic because $z_t$ contains noise (either from the diffusion process or from injected Gaussian noise, as in the noise conditioning mechanism). Different draws of $z_t$ produce different probability vectors $\hat{x}_\eta(z_t)$.

2. Hard sampling: A categorical sample $x \sim \text{Cat}(\hat{x}_\eta(z_t))$ is drawn from the soft probabilities. This step is factorized — each token position is sampled independently given the probabilities — but the soft probabilities themselves already encode correlations across positions because they come from a neural network that processes all positions jointly.

Correlation through entropy collapse: The paper explains the mechanism:

"Because the second step is factorized, the only way for the generator to minimize the moment matching loss is to correlate the soft samples $\hat{x}_\eta(z_t)$ and reduce their output entropy."

Here is the causal chain:

• The D-MMD loss pushes the generator to match the teacher's conditional expectations $\mathbb{E}_q[x|z_s]$, which encode the true correlations in the data (e.g., that certain words tend to co-occur, that grammatical structures require agreement across positions).

• The generator's output is factorized in the second step, meaning the hard samples $x$ will have independent noise across positions for any fixed $\hat{x}_\eta(z_t)$.

• To reduce the variance of the hard samples (making them more consistently reflect the correlations in $\mathbb{E}_q[x|z_s]$), the generator must reduce the entropy of its soft predictions $\hat{x}_\eta(z_t)$. Lower entropy means the probabilities are more peaked — closer to one-hot — so the factorized sampling introduces less randomness.

• Crucially, the overall entropy of the generator comes from both the randomness in $z_t$ (which varies across different noise inputs) and the randomness in the factorized sampling. By reducing the factorized sampling entropy, the generator shifts entropy to the $z_t$ variation, which is learned and can represent correlations.

Empirical evidence (Table 6): The paper measures the generator's output entropy (the average entropy of $\hat{x}_\eta(z_t)$ across token positions) and observes that it decreases with fewer sampling steps:

• For masked D-MMD with 4 steps, output entropy is 1.26 (very low — near-deterministic per-position predictions).
• For 64 steps, output entropy is 1.91 (higher — more uncertainty at each position).

In the 4-step case, the generator has collapsed its factorized output to near-deterministic predictions, relying almost entirely on the variation in $z_t$ to produce diverse samples. In the 64-step case, the generator can afford more per-step randomness because the multi-step denoising process can correct inconsistencies introduced by independent sampling. This is a direct manifestation of the bias-variance tradeoff in sampling: fewer steps require lower per-step variance, which forces the generator to learn more deterministic (and thus more correlated) soft predictions.

Comparison to Di4C: Di4C attempts to model correlations explicitly through mixture distributions, requiring an exponentially growing number of components to capture all correlations. D-MMD avoids this by pushing the correlation problem into the neural network's processing of $z_t$: the generator learns to produce soft predictions that already encode the necessary correlations, and the factorized sampling step contributes only controlled noise. The paper puts it succinctly: "Another perspective is that our entire generator has become the mixture distribution." The "mixture" arises from the different $z_t$ values, not from an explicit mixture output head.

---

Correcting the Bias of $\hat{x}_\eta$ for the Auxiliary Model

A subtle but important detail arises in the training of the auxiliary model: the intermediate state $z_s$ is sampled from the posterior $q(z_s|z_t, x)$, which is conditioned on the hard sample $x$, not on the soft prediction $\hat{x}_\eta$. This creates a potential mismatch.

The problem: The posterior $q(z_s|z_t, x)$ assumes that $x$ is the true clean data. When we sample $x \sim \text{Cat}(\hat{x}_\eta(z_t))$, $x$ is the student's estimate of the clean data, not the ground truth. Therefore, $z_s \sim q(z_s|z_t, x)$ is a state that is consistent with the hard sample $x$, not necessarily with the soft prediction $\hat{x}_\eta$. If we were to use $\hat{x}_\eta$ as the target for training the auxiliary model on $z_s$, the auxiliary would be learning to predict the soft probabilities from a state that was generated assuming a different sample — introducing a bias.

The paper's solution depends on the diffusion type:

• For masked diffusion: The masked state $z_s$ contains [MASK] tokens in positions where the diffusion has erased information. Crucially, "per dimension a masked $z_s$ does not provide information about $x$." This means that conditioning on $x$ versus conditioning on $\hat{x}_\eta$ produces the same distribution over $z_s$ — the mask token carries no information about what the original token was. Therefore, for masked diffusion, it is equally valid to use the soft $\hat{x}_\eta$ or the hard $x$ as the target for the auxiliary model. The paper exploits this by optionally using $\hat{x}_\eta$ as a soft target for the auxiliary model in Algorithm 1 ("Optional use soft target $x \leftarrow \hat{x}_\eta(z_t)$ (only possible for masked diffusion)").

• For uniform diffusion: The noised state $z_s$ contains tokens sampled from a uniform distribution over the vocabulary in positions that have been diffused. Unlike the mask token, a uniform-diffused token does carry information — specifically, which tokens were not sampled. Therefore, conditioning on $x$ produces a different posterior over $z_s$ than conditioning on $\hat{x}_\eta$. The paper states: "for uniform diffusion the auxiliary model always needs to be trained on the hard samples." There is no shortcut — the auxiliary must be trained on the actual hard samples $x$ to avoid bias.

Why this matters: If the auxiliary model is trained on the wrong targets (soft $\hat{x}_\eta$ when the posterior requires hard $x$), it learns to predict a distribution that doesn't match either the student's true output or the teacher's. This breaks the fixed-point property: even if the student perfectly matches the teacher's distribution, a biased auxiliary model would assign incorrect scores, creating spurious gradients that push the student away from the optimum. The paper's handling of this detail ensures that the auxiliary model converges to the correct fixed point regardless of the diffusion process type.

---

Temperature and Top-p Distillation

In practice, language models are often sampled with modified logits — lower temperature (which sharpens the distribution toward high-probability tokens) or top-p sampling (which truncates the distribution to the most probable tokens whose cumulative probability exceeds $p$). These strategies push samples toward the modes of the distribution, improving coherence and quality at the cost of some diversity. The paper extends D-MMD to distill generators that incorporate this mode-seeking behavior.

Temperature distillation: The procedure is straightforward. The teacher's logits are modified during distillation:

$s_\theta(z_s) = \frac{1}{\tau} \log \hat{x}_\theta(z_s)$

where $\tau$ is the temperature parameter and $\log \hat{x}_\theta(z_s)$ is the log-probability output of the teacher. Lowering $\tau$ (e.g., $\tau = 0.7$) makes the distribution sharper, concentrating probability on high-likelihood tokens.

The student is then trained with D-MMD using these temperature-modified teacher logits as the target. Since the cross-entropy loss $\text{CE}(\hat{x}_\eta \| \hat{x}_\theta)$ uses log-probabilities directly, substituting the temperature-scaled logits is a drop-in modification that requires no change to the algorithm. The student learns to match the temperature-sharpened teacher distribution, producing mode-seeking behavior at inference time without explicitly applying temperature at sampling time.

Top-p distillation and the gradient spike problem: Top-p sampling presents a more challenging issue. The standard implementation masks out low-probability tokens by setting their logits to an extremely small value, such as $-10^{20}$, which effectively zeroes out their probability after softmax. The paper identifies a critical problem with this approach in the distillation context:

"This however could lead to gradient spikes, as the teacher log-probability now is in the order of $-10^{20}$. Note that the softmax Jacobian of $\hat{x}_\eta$ is not sufficiently small to cancel this term out."

The mechanism is: the student's softmax Jacobian $\frac{\partial \hat{x}_\eta}{\partial s_\eta}$ (where $s_\eta$ are the student's logits) is proportional to $\hat{x}_\eta(1 - \hat{x}_\eta)$ for the diagonal terms. When $\hat{x}_\eta$ is moderate (not near 0 or 1), this Jacobian is moderate. The teacher's log-probability of $-10^{20}$ for masked-out categories enters the gradient through the cross-entropy term, and the product of the softmax Jacobian with this enormous log-probability produces a gradient spike — an update orders of magnitude larger than typical — causing training instability or divergence.

The paper's solution — dynamic logit lowering: Instead of masking to an extreme value like $-10^{20}$, the paper proposes lowering the logits of masked-out categories by a constant offset $\Delta$:

$s_\theta(z_s) \leftarrow s_\theta(z_s) - (1 - \text{mask}_{\text{top-p}}) \cdot \Delta$

where $\text{mask}_{\text{top-p}}$ is a binary mask (1 for included categories, 0 for excluded), and $\Delta = 2$ in experiments. This reduces the probability of masked-out categories by roughly a factor of $1/e^\Delta$ (ignoring the correction from softmax renormalization, which is small for low-probability events). With $\Delta = 2$, excluded tokens have their probability reduced by a factor of approximately $e^{-2} \approx 0.135$ — they are substantially downweighted but not completely zero, resulting in log-probabilities that are negative but not astronomically so.

Why this works: The gradient from the student's softmax times the teacher's log-probability now produces updates of normal magnitude because $\log \hat{x}_\theta(z_s)_c$ for excluded categories is approximately $-2$ (plus a constant offset from the log-sum-exp normalization), not $-10^{20}$. The student still learns to assign low probability to these categories (because the teacher assigns them low probability), but the training dynamics remain stable. The paper notes that "the precise constant does not really matter, as small log-probability differences will be discounted through the softmax Jacobian of $\hat{x}_\eta$ for low-probability events" — meaning that once a category's teacher log-probability is sufficiently negative that the student's predicted probability for that category is near zero, further discrimination is irrelevant because the softmax Jacobian is near zero in that regime anyway.

---

Noise Conditioning for Masked Diffusion Generators

The paper discovers an important asymmetry between masked and uniform diffusion: masked distillation performs much better with an extra noise source injected into the generator, while uniform diffusion does not require it.

The mechanism: The generator is conditioned on input noise in addition to the noisy state $z_t$. For images, this noise is "a projection of a 2D Gaussian noise pyramid to be added to the residual" — meaning Gaussian noise at multiple scales is projected into the model's feature space and added to intermediate activations. For text, it is "a projection of plain Gaussian noise" — a learned linear projection of random Gaussian vectors, added to the model's representations.

Why masked diffusion needs it: In masked diffusion, the noisy state $z_t$ consists of a mixture of original tokens and [MASK] tokens. As $t$ approaches 0 (cleaner states), fewer positions are masked, and $z_t$ contains less randomness. For few-step generators that operate at coarser time discretizations, $z_t$ may not provide sufficient stochastic variation to represent the full data distribution — there are only so many masking patterns, and the generator's output becomes too deterministic. The additional noise injection provides an independent source of randomness that the generator can use to diversify its outputs. This explains the results in Table 6: without noise conditioning, the 4-step masked generator has output entropy of 1.26 (already collapsed) but achieves an FID of 151.0 (terrible quality). With noise conditioning, the output entropy drops slightly further to 1.01, but FID improves dramatically to 22.3 — the noise injection enables the generator to use its collapsed factorized output effectively, routing diversity through the noise-dependent processing rather than through the per-token sampling.

Why uniform diffusion doesn't need it: In uniform diffusion, the noisy state $z_t$ contains tokens randomly resampled from the uniform distribution. Even at later timesteps, uniform noise provides significant randomness — each diffused position can be any token in the vocabulary. This inherent randomness in $z_t$ serves the same function as the injected Gaussian noise, giving the generator sufficient stochastic variation without an additional noise source. The paper notes: "As is the case with Gaussian diffusion, for uniform diffusion there may already be sufficient noise in $z_t$ that the generator is able to use." This parallels continuous Gaussian diffusion, where the noisy state itself provides enough randomness for stochastic distillation without additional noise conditioning (as found in Salimans et al., 2024).

---

Algorithm 1 in Full Detail

The paper presents D-MMD as Algorithm 1, which I will now walk through step by step, connecting each line to the mathematical framework established above.

Inputs: Student generator $\hat{x}_\eta$, teacher model $\hat{x}_\theta$, auxiliary model $\hat{x}_\phi$, training step $i$, number of sampling steps $k$, dataset $\mathcal{D}$, weighting function $w(s)$, loss function $L(\cdot, \cdot, \cdot)$.

Step 1 — Time sampling: $s, \delta_t \sim U(0, 1), U(0, \tfrac{1}{k})$ $t = \min(1, s + \delta_t)$

Two uniform random variables are drawn: $s$ (the intermediate time) and $\delta_t$ (a small time offset, bounded by $1/k$ where $k$ is the number of student sampling steps). The time $t$ is set to $s + \delta_t$, capped at 1. This means the student always sees a noisier state ($z_t$) than the state the teacher and auxiliary evaluate ($z_s$), with the gap bounded by the student's step size $1/k$. For a 16-step student, $\delta_t \leq 1/16 = 0.0625$, so the student denoises by at most one step's worth of noise per training iteration.

Step 2 — Data diffusion: Sample a clean data point from the dataset $\mathcal{D}$ and diffuse to time $t$ to produce the noisy state $z_t$. This uses the forward diffusion process: $z_t \sim q(z_t|x)$, which means applying the noise schedule to $x$ (e.g., for masked diffusion, randomly replacing a fraction of tokens with [MASK] according to the schedule at time $t$).

Step 3 — Student forward pass: The student generator processes $z_t$ and outputs a soft probability vector $\hat{x}_\eta(z_t)$. This is a distribution over clean tokens for each position, predicted by the student's neural network. For few-step generators, this internally involves a small number of denoising steps (possibly just one if the generator is a single network mapping $z_t$ to the clean prediction).

Step 4 — Hard sampling: $x \sim \text{Categorical}(p = \hat{x}_\eta(z_t))$

A hard token is drawn from the student's predicted distribution for each position. This is the step that would be non-differentiable if we needed gradients through it — but as we will see, it is only used for the auxiliary model and posterior sampling, not for the student's own loss.

Step 5 — Posterior sampling: $z_s \sim q(z_s|z_t, x), \text{ with stop-gradient on } z_s$

The posterior $q(z_s|z_t, x)$ is the reverse of the forward diffusion: given the noisier state $z_t$ and the clean sample $x$, what intermediate state $z_s$ (at time $s < t$) would be consistent with both? For masked diffusion, this means replacing some [MASK] tokens with their clean values from $x$. For uniform diffusion, this means resampling some uniformly-noised tokens to match $x$. The stop-gradient is critical: it prevents gradients from the auxiliary and teacher losses (which depend on $z_s$) from flowing back through the posterior sampling into the student's parameters. This decouples the student's optimization from the randomness in the posterior sampling step and matches the independence assumption used in the MMD gradient derivation.

Step 6 — Student update (even steps): On even training steps ($i$ is even), the student generator is updated. The loss is:

$L_{\text{GEN}}(\eta) = L_s(\hat{x}_\eta(z_t), \hat{x}_\theta(z_s), z_s) - L_s(\hat{x}_\eta(z_t), \hat{x}_\phi(z_s), z_s)$

This is exactly the generalized D-MMD objective for the generator: minimize teacher loss, maximize auxiliary loss. The soft probability vector $\hat{x}_\eta(z_t)$ is used directly (not the hard sample $x$), maintaining differentiability. For the discrete case with cross-entropy, this expands to the form in Equation 11. The student parameters $\eta$ are updated via gradient descent on this loss.

Step 7 — Auxiliary update (odd steps): On odd training steps, the auxiliary model is updated. For masked diffusion, there is an optional step: the hard sample $x$ can be replaced with the soft target $\hat{x}_\eta(z_t)$ (see Section 3.2 on bias correction — this is valid because masked $z_s$ carries no information about $x$). For uniform diffusion, the hard sample $x$ must be used.

The auxiliary loss is:

$L_{\text{AUX}}(\phi) = L_s(x, \hat{x}_\phi(z_s), z_s) + L_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s)$

The first term trains the auxiliary to predict the student's samples. The second term regularizes the auxiliary toward the teacher. The auxiliary parameters $\phi$ are updated via gradient descent on this loss.

Why alternation? The alternating optimization (even steps update student, odd steps update auxiliary) creates the adversarial dynamic without requiring a full inner-loop optimization. If the student and auxiliary were updated simultaneously, the auxiliary could "chase" the student's changing distribution, never converging to a stable fixed point. By giving each model a dedicated update step, the auxiliary has time to approximately converge to $\mathbb{E}_{g_\eta}[x|z_s]$ between student updates, providing a meaningful adversarial signal. This is the same alternation strategy used in continuous MMD (Salimans et al., 2024) and is a standard technique in adversarial training (e.g., GANs).

---

Design Choices and Their Justifications

Soft probabilities for the generator, hard samples for the auxiliary: This asymmetric design is the key to making D-MMD work with discrete data. The generator needs differentiability, which soft targets provide. The auxiliary needs to learn the actual distribution of the student's samples (not just the student's internal probabilities), which requires hard samples. The posterior sampling step with stop-gradient bridges these two worlds: hard samples are drawn from the student's soft predictions, used to generate $z_s$ (which the auxiliary and teacher evaluate), but gradients flow only through the soft predictions in the generator loss.

Cross-entropy instead of squared error: The natural discrete analog of the continuous squared error is cross-entropy, because both are proper scoring rules for their respective domains (Gaussian vs. categorical). Cross-entropy has the property that its gradient with respect to the prediction is proportional to $(p_{\text{pred}} - p_{\text{target}})$, analogous to how squared error's gradient is proportional to $(\hat{x} - x)$. This preserves the structure of the MMD update (difference of teacher and auxiliary predictions) in the discrete case.

Alternating optimization instead of simultaneous: As discussed above, alternation provides a stable adversarial dynamic that converges to the fixed point where student matches teacher. Simultaneous optimization would require careful tuning of relative learning rates to prevent one model from dominating, and even then may oscillate rather than converge.

Posterior sampling instead of full forward diffusion: The paper uses $z_s \sim q(z_s|z_t, x)$ from the posterior, which is cheaper than diffusing $x$ forward from scratch to time $s$, and more tightly couples $z_s$ to the student's current output. This is the same choice made in continuous MMD and contrasts with concurrent work (IDLM) that generates the full $x$ and diffuses back.

Noise conditioning for masked but not uniform: As explained in Section 6.5, this reflects the different information content of the noisy states in each diffusion process. The paper's empirical finding that noise conditioning is essential for masked distillation but unnecessary for uniform distillation is a practical insight that would not be obvious from the theory alone.

4. Key Insights and Innovations

Innovation 1: The Adversarial Min-Max Formulation as a Loss-Agnostic Generalization of Stochastic Distillation

The paper's most foundational conceptual contribution is not D-MMD itself — it is the observation that the alternating optimization used in continuous MMD is equivalent to a single min-max objective over an arbitrary loss function, and that this reformulation is what unlocks discrete distillation. This is a genuinely theoretical advance rather than an incremental engineering fix.

Prior to this work, the dominant assumption in stochastic distillation was that the squared-error formulation of MMD (Salimans et al., 2024) was tied to the Gaussian diffusion setting. The MMD gradient $\nabla_\eta L_{\text{MMD}} \propto \frac{d\hat{x}_\eta}{d\eta}(\hat{x}_\phi - \hat{x}_\theta)$ arises naturally from squared error because the gradient of $\|a - b\|^2$ with respect to $a$ is $2(a - b)$. For a different loss function — say, cross-entropy — the gradient would have an entirely different form, and it was not obvious that the same adversarial dynamic would yield a valid distillation procedure. The field's de facto assumption was that stochastic distillation for discrete models would require developing new theory from scratch.

The paper's min-max reformulation in Equation 9 overturns this assumption. By expressing the objective as $\min_\eta \max_\phi \mathbb{E}[L_s(x, \hat{x}_\theta) - L_s(x, \hat{x}_\phi) - L_s(\hat{x}_\theta, \hat{x}_\phi)]$, the algorithm becomes parametric in the loss function $L_s$. The mathematical structure — student minimizes teacher loss minus auxiliary loss, auxiliary minimizes its own loss regularized toward the teacher — does not depend on whether $L_s$ is squared error, cross-entropy, KL divergence, or any other proper scoring rule. The fixed-point property (student distribution = teacher distribution when auxiliary = teacher) holds for any loss function where the optimum of $\mathbb{E}_{g_\eta}[L_s(x, \hat{x}_\phi)]$ over $\hat{x}_\phi$ is $\mathbb{E}_{g_\eta}[x|z_s]$, which is true for all proper scoring rules. The squared-error gradient derivation in Equation 10 is not the definition — it is a verification that the min-max form reproduces the known MMD algorithm, confirming the generalization rather than constructing it.

This is a fundamental conceptual shift, not an incremental refinement. It reframes stochastic distillation from "a specific algorithm that happens to work for continuous diffusion" to "a general principle — match conditional expectations through adversarial optimization — that can be instantiated with any domain-appropriate loss function." The practical consequence is that adapting MMD to a new domain (discrete tokens, but also potentially structured outputs, graphs, or other data types) reduces to substituting the appropriate loss function, not re-deriving the entire training procedure. The paper explicitly demonstrates this by substituting cross-entropy for squared error and obtaining a working discrete distillation algorithm with minimal additional machinery.

The significance extends beyond the current paper. This loss-agnostic framing opens a systematic research program: what other loss functions yield useful distillation properties? Could one use adversarial losses from GAN literature? Energy-based model score matching? The min-max formulation provides a template for exploring these questions without needing to re-justify the core distillation principle each time. The anchoring evidence is the equivalence derivation in Equation 10 and the CIFAR-10 results in Table 1, where the cross-entropy instantiation (D-MMD) achieves an FID of 3.5 versus the teacher's 6.4 — not merely matching but dramatically exceeding the teacher, confirming that the generalization is not just formally correct but practically effective.

---

Innovation 2: Soft Probability Vectors Resolve the Discrete Gradient Flow Problem Without Architectural Changes

The second distinctive conceptual move is the paper's solution to the discrete gradient flow problem — how to backpropagate through categorical sampling — which takes a fundamentally different approach than prior work by decoupling the generator's training signal from the sampling operation entirely.

Prior work on discrete diffusion distillation faced a structural dilemma. The generator must produce discrete tokens, but sampling from a categorical distribution is non-differentiable. The field's existing solutions fell into two categories. The first, represented by SDTT (Deschenaux and Gulcehre, 2025), simply accepted the non-differentiability and trained through discrete operations using teacher forcing on clean data — which fundamentally limits the method to factorized output distributions and causes mode dropping. The second, represented by Di4C (Hayakawa et al., 2024), modified the architecture to output mixture distributions, adding explicit correlation modeling at the cost of exponentially scaling complexity. Both approaches treated the gradient flow problem as a constraint that must be accommodated — either by limiting ambitions (staying within factorized space) or by paying a complexity cost (adding mixture components).

D-MMD's conceptual innovation is to sidestep the constraint rather than accommodating it. The key insight is that the adversarial min-max objective does not actually require the student to produce hard samples for its own loss. The generator can operate entirely on soft probability vectors $\hat{x}_\eta(z_t)$ — comparing them to the teacher's soft predictions and the auxiliary's soft predictions via cross-entropy — and hard samples are only needed for the auxiliary model, where differentiability is not required (the auxiliary's parameters are updated on the hard samples, but the hard samples are treated as fixed targets with no gradient flowing back to the generator). This is visible in Algorithm 1: the hard sampling step ($x \sim \text{Cat}(\hat{x}_\eta(z_t))$) is used only in the auxiliary update (odd steps) and posterior sampling, never in the generator loss (even steps), which uses $\hat{x}_\eta(z_t)$ directly.

This is more than a computational trick — it represents a reconceptualization of what the generator is optimizing. In prior distillation approaches, the generator is trained to produce samples that a verifier judges as good. In D-MMD, the generator is trained to produce a distribution that the teacher assigns high probability to and the auxiliary assigns low probability to. The soft probability vector is not a degraded approximation of a hard sample — it is the primary object of optimization, and the hard samples drawn from it are strictly for the auxiliary's benefit. This shifts the optimization landscape: the generator never faces the high-variance, sparse-gradient problem of discrete sampling because its loss is computed in the continuous space of probability distributions over tokens.

The significance is that this approach requires zero architectural changes to the generator. Unlike Di4C, which adds mixture distribution output heads, D-MMD uses exactly the same factorized categorical output as the teacher. The generator's factorized architecture remains unchanged — what changes is that the adversarial optimization forces the soft predictions to implicitly encode correlations through entropy collapse (Innovation 3 below). This means D-MMD can be applied to any existing discrete diffusion model without modifying the architecture, the training infrastructure, or the inference pipeline (beyond reducing the number of steps). The practical consequence is that distillation becomes a post-hoc training procedure applied to a fully trained teacher, not a co-design of architecture and training.

The anchoring evidence is the comparison with SDTT in Table 5. SDTT, which must operate within the factorized output space without D-MMD's soft-target adversarial signal, achieves a GPT-2 GM of 0.293 at 64 steps — which actually improves on the teacher's 0.307 but degrades with further distillation. D-MMD achieves 0.236 at 16 steps and 0.225 at 32 steps, substantially better with far fewer steps. This gap is not attributable to architecture differences (both use factorized outputs) or teacher quality (the same teacher is used) — it stems from the fundamentally different optimization objective that soft probability vectors enable.

---

Innovation 3: Entropy Collapse as the Mechanism for Implicit Correlation Learning in Factorized Generators

The third intellectual contribution is the paper's identification and empirical demonstration of entropy collapse — the phenomenon where a factorized generator, under adversarial pressure to match correlated teacher expectations, reduces the entropy of its per-token predictions to shift correlation modeling from the sampling step to the noise-conditioned processing of $z_t$. This is a genuinely new explanatory concept for how discrete generators can produce coherent outputs despite factorized architectures.

Prior work recognized the correlation problem explicitly. Di4C (Hayakawa et al., 2024) frames it as the central challenge: "distillation of discrete diffusion through dimensional correlations" is literally in the title. The field's working assumption was that factorized models cannot represent correlations — if each token is predicted independently given the noise, the joint distribution over tokens is the product of marginals, which cannot capture dependencies. The solution, per this assumption, must involve modifying the output parameterization to explicitly model correlations: mixture distributions (Di4C), autoregressive factorization, or energy-based scoring. This assumption was so deeply held that prior distillation methods (SDTT) accepted mode dropping as inevitable, and Di4C accepted exponential complexity as the price of correlation.

D-MMD's conceptual reframing is that correlations need not be represented in the output distribution at all — they can be represented in the mapping from noise to output. The distinction is between two sources of randomness in the generator:

1. Noise variation: Different $z_t$ inputs (with different realizations of the diffusion noise and injected Gaussian noise) cause the generator to produce different soft predictions $\hat{x}_\eta(z_t)$. This mapping is learned and can encode complex correlations — the neural network processes all token positions jointly, so the predictions at position $i$ and position $j$ for a given $z_t$ can be coordinated.

2. Sampling variation: Given a fixed $\hat{x}_\eta(z_t)$, independent categorical sampling at each position introduces per-token randomness that is, by construction, uncorrelated across positions.

The innovation is recognizing that the generator can shift variance from source 2 to source 1 by reducing the entropy of $\hat{x}_\eta(z_t)$. When $\hat{x}_\eta(z_t)$ is near-deterministic (entropy close to 0), the factorized sampling introduces almost no randomness — the generator's output is essentially a deterministic function of $z_t$. In this regime, all correlations in the output come from the learned $z_t \to \hat{x}_\eta(z_t)$ mapping, which can be arbitrarily complex. The factorized architecture stops being a limitation because the factorization is never exercised — the sampling step has negligible entropy and contributes negligible variance.

Table 6 provides the empirical grounding. The generator's output entropy decreases systematically with fewer sampling steps: from 1.91 at 64 steps (where the model can rely on multi-step denoising to clean up sampling noise) to 1.26 at 4 steps (where the model must get it right in one shot). Without noise conditioning, the 4-step masked generator achieves entropy of 1.26 but FID of 151.0 — it has collapsed but can't use the collapse effectively because $z_t$ alone doesn't provide enough variation. With noise conditioning, entropy drops further to 1.01 and FID improves to 22.3 — the injected noise provides the stochastic diversity, and the collapsed factorized output ensures coherence. The paper's description is precise: "Another perspective is that our entire generator has become the mixture distribution." The "mixture" arises from the different $z_t$ values (with their injected noise), not from an explicit mixture output head.

This is a fundamental conceptual advance rather than an incremental improvement because it reframes the correlation problem from an output representation problem to an optimization dynamics problem. The question is not "how do we parameterize correlations in the output?" but rather "how do we train the generator so that it learns to encode correlations in its noise-to-output mapping and suppress the per-token sampling noise?" The answer — adversarial pressure from the teacher and auxiliary models — is elegant because it requires no architectural changes and emerges naturally from the min-max objective. The broader significance is that this insight is not specific to discrete diffusion: any factorized generator with a stochastic input can, in principle, learn to produce correlated outputs by collapsing its per-output entropy, provided the training objective rewards correlation. This reframes factorized architectures from "inherently limited" to "can be coerced into correlation through appropriate training pressure."

---

Innovation 4: The Gradient Moment Metric as a Distribution-Level Evaluation for Discrete Generative Models

The fourth contribution is the introduction of the Gradient Moment metric (specifically, GPT-2 Gradient Moment) as a principled evaluation tool for discrete generative models — a conceptual advance in measurement rather than in model architecture or training.

The field's standard metric for evaluating discrete diffusion models — generative perplexity, where a reference autoregressive model's perplexity is computed on generated samples — is demonstrably flawed in ways that the paper exposes with unusual clarity. Generative perplexity measures how "predictable" the samples are under the reference model, not how similar they are to the data distribution. As the paper notes, "high density samples are often not typical" (citing Meister et al., 2022), and an "example failure case of the generative perplexity metric is assigning a good score to ungrammatical generated samples that feature many repeated words." A model that outputs the same high-probability phrase repeatedly achieves excellent generative perplexity (the reference model assigns high probability to common words) but produces catastrophically bad samples. This is not a theoretical edge case — it is precisely the failure mode that mode-collapsing distillation methods (SDTT) exhibit, and generative perplexity rewards rather than penalizes it.

The paper's insight is that while the value of the reference model's log-likelihood is a poor quality measure, its gradient with respect to the reference model's parameters is informative. The reasoning is elegant: if a reference model has been trained to convergence on the true data distribution $q(x)$, its expected gradient on that distribution is zero — $\mathbb{E}_{q(x)} \nabla_\theta \log p_{\text{LLM}}(x) = 0$ — because the model is at a stationary point of the training loss. If we evaluate generated samples $x_g$ from some model $g$, and compute $\nabla_\theta \log p_{\text{LLM}}(x_g)$, a large gradient norm indicates that the generated samples look like data points the reference model would update on — i.e., they are distinguishable from the training distribution. The proposed metric formalizes this as the squared norm of the difference between the expected gradient on generated samples and the expected gradient on real data:

$\|\mathbb{E}_g[\nabla_\theta \log p_{\text{LLM}}(x)] - \mathbb{E}_q[\nabla_\theta \log p_{\text{LLM}}(x)]\|^2$

When $g = q$ (the generator exactly matches the data distribution), this metric attains its minimum value of zero. When $g$ diverges from $q$, the metric increases, with the magnitude reflecting the degree of distributional mismatch. The paper's stochastic approximation using inner products of gradient differences on independent minibatches (Equation 14) provides a computationally tractable estimator.

This is a fundamental measurement innovation rather than an incremental refinement because it provides what the field has been missing: a single scalar metric that reliably tracks distribution-level sample quality for discrete generative models, without being gamed by mode collapse, repetition, or temperature reduction. Figure 3 demonstrates the problem it solves: as top-p decreases (more aggressive mode-seeking), generative perplexity steadily improves (goes down), masking the degradation in diversity. The gradient moment initially stays flat, then degrades sharply at very low top-p values — it correctly identifies the point where mode-seeking becomes mode-collapse. This divergence between the two metrics is not a minor calibration issue; it means that an entire line of prior work reporting improved generative perplexity may have been measuring mode collapse, not quality improvement.

The practical significance extends beyond this paper. The gradient moment metric is valid for conditional generation as well as unconditional — simply substitute conditional likelihoods $\log p_{\text{LLM}}(x | x_c)$ — making it applicable to prompted text generation, a setting where metrics like FID (which requires embedding spaces) are inapplicable. This is explicitly noted as a "meaningful advantage." For a field that has struggled to evaluate discrete generative models (unlike images, where FID and Inception Score provide reasonably reliable quality signals despite their flaws), a principled, mode-collapse-resistant metric is a genuine infrastructure contribution that could reshape how future work is evaluated and compared.

The anchoring evidence is Table 5, where methods are compared on both GPT-2 GM and generative perplexity. SDTT achieves perplexity of 26.9 at 64 steps (better than the teacher's 30.4 at 32 steps, suggesting improvement), but its GPT-2 GM of 0.293 is worse than D-MMD's 0.236 at 16 steps. The perplexity metric suggests SDTT is competitive; the gradient moment reveals the gap. This is precisely the diagnostic value the metric provides — it prevents the field from optimizing the wrong thing.

5. Experimental Analysis

Evaluation Methodology

• Dataset. For image generation, the paper uses CIFAR-10 — 50,000 32×32×3 images across 10 classes — trained unconditionally. For text generation, it uses OpenWebText (OWT) , reserving the final 2% of documents as a validation set. The CIFAR-10 training treats each of the 3,072 pixel values (32 × 32 × 3) as a discrete token from a vocabulary of 256 values, with no inductive bias that nearby pixel values are semantically related.

• Base model(s). The teachers are discrete diffusion models — masked diffusion (Austin et al., 2021) which transforms tokens into a special [MASK] token, and uniform diffusion (Hoogeboom et al., 2021) which transforms tokens toward a uniform distribution over the vocabulary. For CIFAR-10, the masked teacher achieves an FID of 6.4 with 1,024 denoising steps and the uniform teacher achieves 7.5. For OWT text generation, the masked teacher achieves a GPT-2 GM of 0.275 with 256 steps (top-p = 0.85) and 0.672 with 256 steps at top-p = 1.0. The paper notes these teachers are "representative" but acknowledges they underperform standard continuous diffusion models on CIFAR-10 (which achieve FID ~3). An autoregressive baseline achieves a GPT-2 GM of 0.061 on OWT — substantially better than any of the diffusion teachers or distilled students, a gap the paper is transparent about.

• Metrics. For images, the paper uses FID (Fréchet Inception Distance) computed on 50,000 generated samples compared to the CIFAR-10 training set. For text, the primary metric is GPT-2 Gradient Moment (GPT-2 GM) — the squared norm of the difference between the expected gradient of a pretrained GPT-2 model's log-likelihood on generated samples versus real data: $\|\mathbb{E}_g[\nabla_\theta \log p_{\text{LLM}}(x)] - \mathbb{E}_q[\nabla_\theta \log p_{\text{LLM}}(x)]\|^2$, estimated via the inner product of gradient differences on independent minibatches (Equation 14). A value of zero indicates the generator's distribution is indistinguishable from the real data distribution under GPT-2's parameters. The paper also reports generative perplexity (GPT-2 perplexity on generated samples) and sample entropy (average per-token entropy of generated sequences) for comparison with prior work, while explicitly arguing these metrics can be gamed by mode collapse — Figure 3 demonstrates that perplexity continues improving as top-p decreases while gradient moment eventually degrades, exposing the metric's flaw.

• Baselines.

  • Teacher models at various numbers of function evaluations (NFEs): the original discrete diffusion models sampled with the standard iterative denoising procedure. These establish the upper bound of what standard discrete diffusion achieves.
  • SDTT (Deschenaux and Gulcehre, 2025) : A reimplementation of Self-Distillation Through Time, the dominant prior discrete diffusion distillation method. Evaluated on OWT at 32 and 64 NFEs in Table 5. The paper also reports SDTT results from prior work for context (Table 5, rows labeled "MDLM + SDTT" achieving GPT-2 GM of 339.7 at 4 NFEs).
  • Di4C (Hayakawa et al., 2024) : Distillation of Discrete Diffusion through Dimensional Correlations. Evaluated on CIFAR-10 at 10, 20, and 40 NFEs in Table 4, using both a pure Di4C generator and a hybrid approach.
  • DiMO (Zhu et al., 2025) : The one-step masked distillation method that D-MMD generalizes. Not directly compared in tables, but the paper notes the one-step D-MMD case is equivalent to DiMO's algorithm.
  • Autoregressive baseline: A standard causal language model achieving GPT-2 GM of 0.061 on OWT — included in Table 2 to contextualize the performance gap that remains between diffusion-based and autoregressive generation.
  • Various methods from prior literature reported in Table 5: Duo + DCD, Duo + Di4C, FMLM, and results taken from Lee et al. (2026). These use different teachers and training procedures, so the comparison is approximate.

• Generation budget / compute accounting. The primary axis of comparison is Number of Function Evaluations (NFEs) — the number of forward passes through the neural network required to generate one sample. For the teacher, this equals the number of denoising steps (e.g., 256, 512, 1024). For the distilled student, this equals the number of sampling steps the student takes (e.g., 4, 8, 16, 32, 64). The paper sweeps NFE across powers of 2, typically from 4 to 1024. One NFE is a single forward pass for one complete set of token positions — the student and teacher use the same per-step computational cost, so NFE provides a fair comparison of total compute. The paper does not account for training cost in its comparisons (distillation is treated as a one-time training expense amortized over many inference calls) and does not report wall-clock time or latency, only NFE count. For the block-autoregressive experiment, 16 student NFEs are compared against 256 teacher NFEs within a block of 256 tokens.

• Cross-validation / statistical protocol. The paper does not report formal cross-validation for model selection or statistical significance testing. For top-p and temperature parameters, the paper states it "tuned the top-p value for the best GPT-2 GM" on the masked teacher, selecting p = 0.85 by sweeping different values. This tuning is done using the GPT-2 GM metric itself (or FID for images), meaning the metric used for selection is the same as the metric used for evaluation — the paper does not hold out a separate validation set for hyperparameter selection. The images in the appendix (Figures 4, 5, 6) show FID sweeps over temperature and top-p values for posterior sampling, teacher distillation temperature, and teacher top-p distillation, indicating these hyperparameters were selected based on full-dataset sweeps. For the text experiments, the last 2% of OWT is held out as a validation set, but it is unclear whether this was used for early stopping, hyperparameter selection, or evaluation of the distilled models — the main results in Tables 2, 3, and 5 report metrics without specifying which data split they were computed on. The lack of explicit cross-validation or held-out evaluation for hyperparameter selection is a limitation — the reported numbers may be optimistically biased if the same 2% validation set was used both for tuning and final reporting.

Main Quantitative Results

CIFAR-10 Image Generation: D-MMD Comprehensively Outperforms Teacher and Prior Distillation Methods

The headline result for image generation is in Table 1: D-MMD-distilled generators achieve substantially better FID than their teachers while using a small fraction of the denoising steps.

For masked diffusion, the teacher achieves FID of 6.4 at 1,024 NFEs. The D-MMD student:

• At 4 NFEs: FID 22.3
• At 8 NFEs: FID 12.7
• At 16 NFEs: FID 5.3 — already better (less than) the teacher's 6.4 at 1,024 NFEs, a 64x reduction in steps while improving quality
• At 32 NFEs: FID 3.8
• At 64 NFEs: FID 3.5 — the best result, close to what standard continuous diffusion achieves on CIFAR-10

The student's performance at 16 NFEs (FID 5.3) is notably better than the teacher at any budget level, including 1,024 NFEs. This is the fundamental efficiency gain: the student gets better quality using 16× fewer forward passes (16 vs. 256, the lowest NFE the teacher was evaluated at).

For uniform diffusion, the teacher achieves FID of 7.5 at 1,024 NFEs. The D-MMD student:

• At 4 NFEs: FID 7.1
• At 8 NFEs: FID 5.0
• At 16 NFEs: FID 4.1
• At 32 NFEs: FID 3.7
• At 64 NFEs: FID 3.8 (slightly worse than 32 NFEs)

The uniform student's best FID (3.7 at 32 NFEs) is less than half the teacher's best FID (7.5 at 1,024 NFEs), again representing a massive step reduction and quality improvement simultaneously.

An important pattern visible in Table 1: the teacher's FID improves monotonically with more NFEs (from 122.9 at 4 steps to 6.4 at 1,024 steps for masked). The student's FID improves sharply from 4 to 64 NFEs, then plateaus — for uniform, it actually slightly degrades from 3.7 (32 NFEs) to 3.8 (64 NFEs). This is the "paradoxical side-effect" discussed in Section 6.6: the student can outperform the teacher at intermediate step counts, but must eventually converge to the teacher's performance at high step counts. The plateau near 3.5–3.8 likely approaches the fundamental limit of discrete diffusion without continuous inductive biases.

Table 4 provides comparison with prior distillation work on CIFAR-10. Di4C's teacher achieves FID 8.0 at 40 NFEs (better than D-MMD's masked teacher with 40 NFEs — approximately 20.0 — because Di4C uses a different teacher trained with a Gaussian-mimicking discrete process). Nevertheless:

• D-MMD uniform at 8 NFEs: FID 5.0 vs. Di4C at 10 NFEs: FID 20.6
• D-MMD masked at 16 NFEs: FID 5.3 vs. Di4C (hybrid) at 20 NFEs: FID 9.5
• D-MMD masked at 32 NFEs: FID 3.8 vs. Di4C teacher at 40 NFEs: FID 8.0

The gap is dramatic — D-MMD achieves 4–5× better FID with fewer NFEs than Di4C. This comparison is imperfect because Di4C uses a different teacher, but the magnitude of improvement (FID 20.6 → 5.0 at comparable NFEs) suggests a qualitative difference in distillation effectiveness rather than merely a stronger teacher.

OWT Text Generation: D-MMD Reduces Steps While Improving Sample Quality

The headline text generation results appear in Table 2. The masked teacher achieves GPT-2 GM of 0.275 at 256 NFEs (top-p = 0.85). The D-MMD student:

• At 4 NFEs: GPT-2 GM 0.456
• At 16 NFEs: GPT-2 GM 0.236 — already better (less than) the teacher at 256 NFEs, a 16× step reduction
• At 32 NFEs: GPT-2 GM 0.225
• The trend continues improving from 16 to 32 NFEs, unlike images where it plateaus

The uniform teacher achieves GPT-2 GM of 0.313 at 256 NFEs (top-p = 0.50). The D-MMD uniform student achieves 0.337 at 8 NFEs and plateaus at 0.307–0.316 across 16–64 NFEs — basically matching the teacher with 8× fewer steps, but not exceeding it as dramatically as the masked student.

Why does masked D-MMD improve more than uniform? The paper does not provide a direct comparison with matched compute, but the numbers suggest masked distillation is more effective: the masked student achieves 0.236 at 16 NFEs vs. teacher 0.275 at 256 NFEs, while the uniform student achieves 0.310 at 16 NFEs vs. teacher 0.326 at 128 NFEs — the relative improvement is larger for masked. This may relate to the noise conditioning asymmetry (masked benefits from injected noise; uniform already has sufficient noise in $z_t$) or to the bias correction flexibility (masked allows soft targets for the auxiliary; uniform requires hard targets).

The block-autoregressive result in Table 3 demonstrates practical scaling. In a setting where a 256-token block is generated by diffusion conditioned on an autoregressive context, the 16-step D-MMD generator matches the 256-step teacher exactly (GPT-2 GM 0.225 for both). This is a 16× reduction in diffusion steps without quality loss, and it represents the most realistic deployment scenario — combining autoregressive efficiency with diffusion parallelism within blocks.

Table 5 provides the most comprehensive comparison with prior work on OWT. Several methods are drawn from the literature (using different teachers), making exact comparison difficult, but the patterns are informative:

• SDTT (reimplementation, using the same masked teacher as D-MMD) achieves GPT-2 GM 0.293 at 64 NFEs and 0.340 at 32 NFEs. D-MMD achieves 0.236 at 16 NFEs and 0.225 at 32 NFEs — better quality with 2–4× fewer steps.
• At top-p = 1.0 (no mode-seeking), the pattern holds: SDTT at 64 NFEs achieves GPT-2 GM 0.293; D-MMD at 32 NFEs achieves 0.578 — but wait, this is actually worse. The paper's top-p = 1.0 results (bottom section of Table 5) show D-MMD achieving 0.558–0.719 across NFEs, while the teacher achieves 0.672–0.781. D-MMD does outperform the teacher slightly at 16 NFEs (0.558 vs. 0.711), but the improvement is much less dramatic than at top-p = 0.85. This suggests mode-seeking distillation (temperature/top-p) is crucial to the student-over-teacher effect — without it, the student approximately matches the teacher's distribution rather than exceeding it.
• Prior methods using different teachers: Duo + DCD achieves GPT-2 GM 108.2 at 4 NFEs; FMLM achieves 76.4. D-MMD at 4 NFEs achieves 0.820 — incomparably better. However, these prior methods are evaluated at 4 NFEs (presumably one-step generation), where D-MMD achieves its worst result (0.820). The paper's strength is in the few-step regime (16–32 NFEs), not one-step.

The sample entropy column in Table 5 is revealing. Real data has entropy 5.44. The masked teacher (top-p = 0.85) has entropy 5.13–5.19 across NFEs — slightly lower than data, consistent with mild mode-seeking from top-p. D-MMD at 16 NFEs has entropy 5.00 and at 32 NFEs has 5.05 — lower than both data and teacher, indicating additional mode-seeking beyond what top-p alone provides, consistent with the paper's claim that adversarial distillation "may move more density towards modes without fully collapsing" (Section 6.6). SDTT at 64 NFEs has entropy 5.17 — closer to the teacher's entropy, suggesting SDTT preserves the teacher's diversity better than D-MMD. Priormethods at 4 NFEs have substantially lower entropy (Duo + DCD at 4.82, FMLM at 5.05), indicating more aggressive mode collapse at very low step counts.

The GPT-2 Perplexity column demonstrates the metric's unreliability. Real data has perplexity 15.4. D-MMD at 16 NFEs achieves 17.2 — reasonably close. But the prior methods (Duo + DCD, FMLM) achieve perplexities of 4.82–5.05, which are better (lower) than real data's 15.4. This is impossible for a model that matches the true data distribution (where perplexity should equal the data entropy plus a constant), and it reflects mode collapse — the models produce repetitive, high-probability text that is trivially predictable but not representative. The GPT-2 GM correctly penalizes these models: Duo + DCD achieves GM 108.2, FMLM achieves 76.4, both far worse than D-MMD's 0.236 at 16 NFEs.

Figure 2 provides a qualitative sanity check: a random 1024-token sample from the 16-step masked D-MMD generator shows coherent English text with proper quotation syntax, named entities ("Klineen"), and sensible (if somewhat repetitive) discourse structure. This is not cherry-picked per the caption, suggesting the generator produces generally legible text.

Noise Conditioning Ablation: Essential for Masked, Irrelevant for Uniform

Table 6 reports the effect of input noise conditioning for masked diffusion on CIFAR-10. With noise conditioning disabled, the FID values are catastrophic at low NFEs:

• 4 NFEs: FID 151.0 without noise vs. 22.3 with noise
• 8 NFEs: FID 37.0 vs. 12.7
• 16 NFEs: FID 14.7 vs. 5.3
• 32 NFEs: FID 7.7 vs. 3.8
• 64 NFEs: FID 6.0 vs. 3.5

The gap narrows as NFEs increase — at 64 NFEs, 6.0 vs. 3.5 is a meaningful but not catastrophic difference. This is consistent with the entropy collapse mechanism: at low NFEs, the generator must collapse its factorized output distribution to low entropy to produce coherent samples, but without noise conditioning, $z_t$ alone doesn't provide sufficient stochastic variation to represent the data distribution's diversity — the generator collapses to a near-deterministic mapping that produces poor samples. With noise conditioning, the injected Gaussian noise provides the missing diversity, and the collapsed factorized output ensures coherence. At higher NFEs, the generator can afford higher per-step entropy because multi-step denoising can clean up sampling inconsistencies, reducing the need for both entropy collapse and noise conditioning.

The generator output entropy rows in Table 6 quantitatively confirm this mechanism. Without noise conditioning, output entropy increases from 1.26 (4 NFEs) to 1.91 (64 NFEs). With noise conditioning, entropy is consistently lower: 1.01 (4 NFEs) to 1.83 (64 NFEs). At 4 NFEs, the noise-conditioned generator achieves entropy of 1.01 — nearly deterministic per-position predictions — while achieving FID 22.3. Without noise conditioning, the generator at 4 NFEs has entropy 1.26 (slightly higher, meaning more per-token randomness) but achieves terrible FID 151.0 — confirming that the quality improvement comes from routing diversity through the noise conditioning rather than through the factorized sampling step. The paper states that for uniform diffusion, "we did not observe any meaningful improvements" from noise conditioning, consistent with uniform noise in $z_t$ already providing sufficient stochastic variation.

Ablation Studies and Robustness Checks

Posterior sampling temperature and top-p (Appendix C, Figure 4): The paper sweeps sampling hyperparameters during evaluation to find optimal FID. For masked diffusion, FID is relatively flat across temperature 0.7–1.0 (FID ~6–7) but degrades sharply at low temperature (<0.5) and high temperature (>2.0). Top-p shows a U-shaped curve with optimum around p=0.9 (FID ~6) and degradation below 0.7 and above 0.95. For uniform diffusion, temperature has a narrower sweet spot near 1.0, and top-p shows a shallow optimum near p=0.5. These sweeps justify the specific sampling hyperparameters used in the main experiments.

Teacher temperature during distillation (Appendix C, Figure 5): When the teacher uses temperature-scaled logits during D-MMD training, the student's FID varies with the teacher's temperature. For masked diffusion, FID is best (lowest) at teacher temperature around 0.7–0.8 and degrades at both lower and higher temperatures. For uniform diffusion, the optimum is broader but centered near 0.9. This suggests that distilling a moderately mode-seeking teacher (temperature < 1.0) improves the student's sample quality — consistent with the Section 6.6 discussion of mode-covering vs. mode-seeking behavior.

Teacher top-p during distillation (Appendix C, Figure 6): Similar pattern: distilling a teacher with top-p < 1.0 improves student FID. For masked diffusion, FID is best at teacher top-p ~0.8–0.85 and degrades above 0.95. For uniform diffusion, the optimum is broader but best near p=0.5–0.7. The key insight: the teacher is trained with maximum likelihood (mode-covering), but teaching the student from a slightly mode-seeking teacher distribution (via temperature or top-p) produces better generative models — the adversarial D-MMD dynamic itself induces some mode-seeking, and this can be amplified or guided by the teacher's sampling parameters.

SDTT comparison with matched teacher (Table 5): The reimplementation uses the same masked teacher as D-MMD. SDTT at 64 NFEs achieves GPT-2 GM 0.293 — this is actually an improvement over the teacher at 64 NFEs (0.307), suggesting SDTT does provide some distillation benefit. However, D-MMD at 16 NFEs achieves 0.236 and at 32 NFEs achieves 0.225 — substantially better with fewer steps. The critical distinction is that SDTT's improvement may come at the cost of diversity (mode dropping), whereas D-MMD's lower GPT-2 GM and sample entropy suggest better distribution matching. The paper notes that SDTT "degrades over repeated distillation rounds" — Table 5 shows SDTT at 32 NFEs achieving GM 0.340, worse than its 64-step result (0.293), implying that pushing SDTT to lower step counts causes quality degradation that D-MMD avoids.

Top-p = 1.0 vs. top-p = 0.85 (Table 5, bottom vs. top): At top-p = 1.0 (no mode-seeking in the teacher or student), the teacher achieves GPT-2 GM 0.672 at 256 NFEs, degrading to 0.781 at 64 NFEs. D-MMD achieves 0.719 at 4 NFEs (worse than teacher at 256 but better than teacher at 64) and improves to 0.558 at 16 NFEs. The relative improvement over the teacher is modest — D-MMD at 16 NFEs achieves GM 0.558 vs. teacher at 64 NFEs achieving 0.781 (a clear improvement) and vs. teacher at 256 NFEs achieving 0.672 (better but not dramatically). With top-p = 0.85, the improvement is much starker: D-MMD at 16 NFEs achieves 0.236 vs. teacher at 256 NFEs achieving 0.275. This demonstrates that the student-over-teacher effect is substantially amplified by mode-seeking distillation — the adversarial D-MMD dynamic alone provides some improvement (top-p = 1.0 results), but the combination with temperature/top-p distillation produces the dramatic gains reported in the abstract.

Mixture distribution vs. factorized output (Table 5 vs. Table 1): No direct ablation compares factorized D-MMD against a D-MMD variant with mixture distributions — this experiment was not run. The comparison with Di4C (which uses mixtures) in Table 4 is confounded by different teachers. A controlled experiment keeping the teacher fixed and varying whether the student uses factorized or mixture outputs would directly test the paper's claim that entropy collapse can substitute for explicit mixture modeling, but this ablation is absent.

Posterior sampling vs. full forward diffusion (no direct ablation): The paper mentions that IDLM (Li et al., 2026) differs in using full forward diffusion from generated $x$ back to $z_t$ rather than posterior sampling from $q(z_s|z_t, x)$. No ablation compares these two approaches for D-MMD. The posterior sampling choice is justified by computational efficiency and tighter coupling to the student's output, but whether it is necessary or merely convenient is untested.

Soft vs. hard targets for auxiliary model in masked diffusion (Algorithm 1, step 7): The algorithm notes that for masked diffusion, the auxiliary can optionally be trained on the soft target $\hat{x}_\eta(z_t)$ instead of the hard sample $x$. The paper states this is valid because masked $z_s$ provides no information about $x$, but no ablation compares the two variants. This is a missing piece — does using soft targets improve auxiliary convergence? Reduce gradient variance? The paper reports results without specifying which variant was used for the main experiments.

Effect of alternation frequency (no ablation): The alternating optimization updates the student on even steps and the auxiliary on odd steps — effectively a 1:1 ratio. No sweep over different ratios (e.g., 1:2, 2:1, 5:1) is reported. In GAN literature, the balance between generator and discriminator updates significantly affects training stability and final quality. This dimension of D-MMD's sensitivity is unexplored.

Critical Assessment

Claim 1: D-MMD distills discrete diffusion models into few-step generators that maintain high quality and diversity, whereas previous discrete distillation methods collapse.

The experiments convincingly demonstrate that D-MMD achieves better quality (lower FID, lower GPT-2 GM) than both SDTT and Di4C at comparable or lower NFEs. For CIFAR-10, D-MMD achieves FID 5.0 at 8 NFEs (uniform) and 5.3 at 16 NFEs (masked), while Di4C achieves FID 9.5 at 20 NFEs (hybrid) and 20.6 at 10 NFEs (pure). For text, D-MMD achieves GPT-2 GM 0.236 at 16 NFEs while SDTT (reimplementation with same teacher) achieves 0.293 at 64 NFEs and 0.340 at 32 NFEs.

The "diversity" part of the claim is harder to assess. The sample entropy values in Table 5 show D-MMD's entropy (5.00–5.05 at 16–32 NFEs) is lower than the teacher's (5.13–5.19) and lower than real data (5.44). This suggests some diversity loss relative to the teacher. SDTT's entropy at 64 NFEs (5.17) is closer to the teacher's (5.19) — SDTT may actually preserve more diversity than D-MMD while achieving worse GPT-2 GM. The paper frames this as a feature (mode-seeking toward higher-quality modes) rather than a bug, but the claim of "maintaining high quality and diversity" needs qualification: D-MMD trades some diversity (lower entropy) for quality (lower GM), and whether this tradeoff is acceptable depends on the application. The paper does not report metrics that directly quantify diversity independent of quality (e.g., distinct n-gram counts, self-BLEU, or coverage metrics), making it difficult to verify the "diversity" half of the claim.

A genuine weakness: the comparison with prior methods is not fully controlled. Di4C uses a different teacher (trained with a Gaussian-mimicking process achieving 40-step FID 8.0 vs. D-MMD's teacher at 40 steps achieving ~20.0). SDTT's reimplementation uses the same teacher as D-MMD, which is the cleanest comparison, but prior SDTT results in the literature (MDLM + SDTT achieving GM 339.7 at 4 NFEs) use a different teacher (MDLM). The "collapse" behavior the paper attributes to prior methods is thus partly confounded by teacher quality — a stronger teacher might yield better SDTT results. The paper's SDTT reimplementation partially addresses this (showing D-MMD still wins with matched teacher), but only at 32–64 NFEs for text and not at all for images.

Claim 2: The distilled generators can outperform their teachers.

This claim is strongly supported for the specific metric-mode-seeking combination used. For CIFAR-10, the masked D-MMD student at 16 NFEs achieves FID 5.3 vs. teacher FID 6.4 at 1,024 NFEs — clearly better. For text, the masked D-MMD student at 16 NFEs achieves GPT-2 GM 0.236 vs. teacher 0.275 at 256 NFEs (top-p = 0.85) — also clearly better.

However, the claim requires careful scoping. At top-p = 1.0 (no mode-seeking), the student's improvement over the teacher is much smaller: GPT-2 GM 0.558 at 16 NFEs vs. teacher 0.672 at 256 NFEs — the student is better, but the gap is ~17% rather than ~14% in the mode-seeking case. The "outperforming" effect is thus partly attributable to mode-seeking distillation (Section 3.3) rather than the core D-MMD algorithm alone. The paper acknowledges this in Section 6.6: "many distillation approaches such as D-MMD have an adversarial component... which both are reminiscent of reverse-KL optimization. D-MMD may move more density towards modes without fully collapsing." This is an honest characterization — the student outperforms the teacher by being different (more mode-seeking), not by being a strictly better approximation of the data distribution.

A deeper question: is "outperforming" on FID and GPT-2 GM the right criterion? Both metrics measure distributional similarity (distance between generated and real distributions). If a student achieves better FID than the teacher, it means the student's distribution is closer to the real data distribution in the FID sense. This could be because the teacher's maximum-likelihood training produces a mode-covering distribution that includes low-quality modes, and the student's adversarial training produces a mode-seeking distribution that better matches the Inception embedding statistics. FID is not a perfect metric (it's sensitive to feature extraction and has known biases), but it is a standard benchmark. The paper's claim is defensible within this framework, with the caveat that "outperforming" means "achieves a better score on distribution-comparison metrics" — the student's samples are not necessarily more diverse or more representative of rare modes.

Claim 3: D-MMD generalizes to both masked and uniform diffusion, and to block-autoregressive settings.

Supported. Tables 1 and 2 show results for both diffusion types. Uniform D-MMD achieves FID 3.7 (vs. teacher 7.5) and GPT-2 GM 0.310 (vs. teacher 0.326) — the uniform results are less dramatic than masked but still show improvement. The block-autoregressive result in Table 3 demonstrates a third configuration: diffusion within blocks conditioned on autoregressive context, where the 16-step student matches the 256-step teacher.

However, the generalization is demonstrated only within the scope of CIFAR-10 and OWT, using the specific model architectures described for each diffusion type. The paper does not test D-MMD on other discrete diffusion variants (e.g., absorbing-state diffusion with different noise schedules, discrete flow matching with alternative interpolation paths, or multinomial diffusion with different transition matrices). The claim of generality is supported for the two most common discrete diffusion processes, but the space of possible discrete diffusion models is larger than these two examples.

Claim 4: The GPT-2 Gradient Moment metric reliably measures sample quality without being gamed by mode collapse, unlike generative perplexity.

The theoretical argument is strong and the empirical demonstration in Figure 3 is clear: as top-p decreases, perplexity improves monotonically while gradient moment eventually degrades. Table 5 shows that methods with suspiciously good perplexity (Duo + DCD at 4.82, FMLM at 5.05, below data's 15.4) are correctly flagged as poor by the gradient moment (108.2 and 76.4 respectively). The metric appears to capture what it claims.

Weaknesses: The metric depends on the choice of reference model (GPT-2 in this case). If GPT-2 has its own biases (e.g., it was trained on web text with certain stylistic patterns), the gradient moment will penalize deviations from those patterns even if the generated text is objectively high-quality. A generator of formal academic prose might score poorly because GPT-2 was trained on informal web text — the metric cannot distinguish between "different from data" and "worse than data." The paper acknowledges this implicitly by stating the metric measures whether the reference model would update on the generated data, but doesn't discuss reference model bias as a limitation. Additionally, the metric requires computing gradients through a full autoregressive model, which is computationally expensive (though the paper's stochastic approximation makes it tractable). The metric has not been validated against human judgments or downstream task performance — we know it agrees with FID in calling D-MMD better than SDTT, but we don't know whether it tracks what users actually care about.

Missing experiments that would have strengthened the paper:

1. Human evaluation of text quality. The paper relies entirely on automatic metrics (GPT-2 GM, perplexity, entropy). For text generation, human judgments of fluency, coherence, and relevance are standard. The absence of human evaluation makes it impossible to know whether the GM improvement from 0.275 (teacher) to 0.225 (D-MMD 32-step) translates to perceptible quality differences, or whether the reduced entropy (5.05 vs. 5.19) corresponds to noticeably less diverse outputs.

2. D-MMD with a stronger teacher. The teachers achieve FID 6.4–7.5 on CIFAR-10 — far from state-of-the-art continuous diffusion (~FID 3). If D-MMD were applied to a stronger discrete diffusion teacher (e.g., MD4 from Shi et al., 2024), would the student still outperform the teacher? The plateau at FID ~3.5 in Table 1 suggests D-MMD may be hitting a floor — perhaps the factorized architecture imposes a fundamental quality limit that even optimal distillation cannot exceed. This would be important to know.

3. Scaling to larger models and longer sequences. All text experiments use 1024-token unconditional generation. How does D-MMD perform with longer sequences (2048, 4096 tokens) where correlations span greater distances? The entropy collapse mechanism should still work in principle, but the required degree of collapse may become harder to achieve as sequence length grows. Larger model scales would test whether the student-over-teacher effect persists.

4. Direct combination of D-MMD with Di4C's mixture outputs. If D-MMD's adversarial training and Di4C's mixture distribution architecture are complementary, combining them might yield further improvements. The paper positions entropy collapse as an alternative to mixture modeling, but doesn't test whether the two approaches are substitutes or complements.

5. Measuring diversity directly. Reporting distinct n-gram counts, self-BLEU, or vocabulary usage statistics would quantify the diversity loss that the lower sample entropy suggests, allowing an explicit quality-diversity tradeoff analysis.

6. Ablation on soft vs. hard auxiliary targets for masked diffusion. The paper notes this is possible but doesn't report results. Given that the auxiliary model's role is central to the adversarial dynamic, this choice could significantly impact training dynamics and final quality.

Conditions where the claims hold:

• The student-over-teacher effect is strongest with mode-seeking distillation (top-p ≈ 0.85, temperature ≈ 0.7–0.8) and relatively weaker at top-p = 1.0.
• Masked distillation benefits substantially from noise conditioning; uniform distillation does not require it.
• The improvement over prior methods is clearest at 8–32 NFEs. At 4 NFEs, D-MMD's advantage narrows (GPT-2 GM 0.820 vs. FMLM 76.4 — D-MMD is better but not by the same margin as at 16 NFEs). At very high NFEs (64+), the student must converge to the teacher, eliminating the advantage.
• The block-autoregressive result is demonstrated only for one block size (256 tokens) and one NFE pair (16 vs. 256). Generalization to other block sizes or NFE ratios is untested.
• The GPT-2 GM metric is validated as superior to generative perplexity for detecting mode collapse, but its correlation with human judgment is unknown.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Not Amortized in the Reported Efficiency Gains

The assumption or constraint. The compute-optimal framework rests on estimating each prompt's difficulty before selecting a test-time strategy, and the paper uses a procedure — generating 2,048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — that is extraordinarily expensive. The paper acknowledges this in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. The headline 4× efficiency gains (Figure 4: 16 generations matching best-of-N at 64; Figure 8: 64 generations matching best-of-N at 256) are computed after difficulty is already known, without amortizing the cost of learning it. In a deployment setting, the total compute would be difficulty estimation + strategy execution. The estimation step (2,048 generations) alone exceeds the largest test-time budgets studied (256–512 generations). A practitioner deploying this method would find that the true per-query cost is dominated by difficulty estimation, potentially eliminating or even reversing the reported gains. For interactive applications where each query is unique, the amortization argument (estimate difficulty once per question, reuse on subsequent calls) does not apply — every query pays the full estimation cost.

What evidence exists in the paper. The paper does not report any experiment that includes difficulty estimation cost in the compute budget. The compute-optimal scaling curves in Figures 4 and 8 are plotted against the strategy execution budget only. The paper does acknowledge the problem (Section 3.2, Section 8) but provides no empirical characterization of how the tradeoff plays out — no experiment showing at what query volume the amortized cost becomes acceptable, and no comparison of alternative cheaper difficulty estimation methods (e.g., using fewer than 2,048 samples, or a learned difficulty predictor).

Mitigation status. The paper flags this as "a key avenue for future work" (Section 3.2) and suggests two directions: "pretraining or finetuning models to directly predict difficulty of a question" and adaptive estimation that amortizes difficulty assessment into the problem-solving process. Neither is developed or evaluated. The reported 4× efficiency figure should therefore be understood as an upper bound on achievable deployment efficiency, not a realized gain.

---

Hard Problems Gain Essentially Nothing from Test-Time Compute

The assumption or constraint. The paper's entire framework assumes the base model has some non-trivial probability of producing correct solutions — test-time compute amplifies existing capability but does not create it. Section 7 makes this explicit:

"Test-time compute is powerful when problems are within the base model's reach (it already produces correct solutions at some non-trivial rate), but it cannot compensate for fundamental capability gaps that larger pretraining would address."

The consequence. Across every method studied — PRM search (Figure 3, right), iterative revisions (Figure 7, right), compute-optimal allocation of both (Figures 4 and 8), and the FLOPs-matched comparison (Figure 9) — performance on the hardest questions (difficulty bin 5) remains near zero regardless of compute budget. In Figure 3, bin 5 accuracy sits at 1–3% for all search methods at all budgets up to 256 generations. In Figure 7, bin 5 shows roughly 2–3% accuracy across all sequential-to-parallel ratios at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% while the 14× larger model achieves substantially better performance on these problems at all three values of R.

This is a hard capability boundary: if the base model cannot produce a correct solution at any meaningful rate (pass@1 near zero), no inference-time strategy — beam search, lookahead, sequential revisions, or any combination — will help. For practitioners, this means test-time compute is only useful on problems where the model is already roughly capable, which requires knowing (or estimating) the difficulty distribution of the deployment workload.

What evidence exists in the paper. The evidence is consistent and unambiguous across all difficulty-bin analyses. Bin 5 in Figures 3, 7, and 9 shows near-zero improvement across all methods and budgets. The FLOPs-matched comparison in Figure 9 quantifies the gap: on bin 5, the 14× larger model's greedy performance (stars) is substantially above the compute-optimal scaling curve at all values of R, confirming that pretraining — not inference-time compute — is the only path to improvement on these problems.

Mitigation status. The paper is transparent about this limitation. The Section 7 takeaway explicitly states the boundary condition. However, the paper does not provide guidance on how to estimate whether a given problem or task distribution falls on the "within reach" side of this boundary without already having access to a larger model for comparison.

---

Single Benchmark and Single Model Family Limit Generality

The assumption or constraint. All experiments use the MATH benchmark (12,000 training, 500 test questions) with PaLM 2-S* as the base model. The paper states in Section 4 that it "believes this model is representative of the capabilities of many contemporary LLMs," but no evidence is provided for this belief. The paper does not test on other reasoning benchmarks (code generation, logical reasoning, scientific QA), other model families (LLaMA, GPT, Claude), or other model scales within the same family.

The consequence. Several aspects of the findings could be model-specific or benchmark-specific:

• PRM quality and over-optimization behavior depend on PaLM 2-S*'s output distribution and calibration. A model with different error patterns — e.g., one that produces more diverse or more peaked solution distributions — could exhibit different difficulty-dependent scaling curves and different over-optimization thresholds.
• Revision model effectiveness depends on the base model's ability to learn from in-context incorrect examples during fine-tuning, which varies substantially across model families and scales. A model with stronger or weaker in-context learning might show different sequential-to-parallel optimal ratios.
• The MATH benchmark consists of competition-level math problems requiring symbolic reasoning and step-by-step deduction. The difficulty-dependent patterns (beam search hurting easy problems, sequential revisions helping easy problems, no method helping hard problems) may not generalize to other reasoning domains — code generation has different error patterns (syntax errors vs. logic errors), logical reasoning may benefit more from search, and factual QA may depend more on knowledge retrieval than inference strategy.

What evidence exists in the paper. No evidence. The paper contains zero replication on other benchmarks or other model families. The test set of 500 questions, split into five difficulty quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin. The paper does not report confidence intervals or standard errors on the compute-optimal scaling curves, making it impossible to assess statistical reliability at this sample size.

Mitigation status. The paper acknowledges scope limitation in Section 8: "extending these results to other benchmarks and model families is an important direction for future work." No mitigation is attempted within the paper.

---

The FLOPs-Matched Pretraining Baseline Is Not Compute-Optimal

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time strategies against a model with approximately 14× more parameters, trained on the same number of tokens — i.e., scaling only model size, not training data. The paper explicitly acknowledges this departs from compute-optimal pretraining:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding — it receives no test-time compute budget of its own. This is acknowledged in Section 7.

The consequence. The reported crossover points — where test-time compute with the smaller model outperforms pretraining with the larger model — are potentially overstated relative to a stronger baseline. A compute-optimal pretrained model (scaling both parameters and data according to Hoffmann et al., 2022) would achieve better perplexity per total FLOP than a parameter-only-scaled model, making the pretraining baseline stronger. Similarly, giving the larger model even a modest test-time compute budget (e.g., best-of-8 or majority voting) would substantially improve its performance, since the paper's own results show that test-time compute improves all models, not just small ones.

The specific numbers that would likely change under a stronger baseline:

• The +27.8% relative improvement for revisions on easy questions at R ≪ 1 (Figure 1 top-right) could shrink or reverse.
• The -52.9% disadvantage for PRM search on hard questions at R ≫ 1 (Figure 1 bottom-right) would likely become even more negative.
• The difficulty crossover points (where test-time compute stops being preferable) might shift, making pretraining more favorable on a larger fraction of the problem distribution.

What evidence exists in the paper. The paper does not provide any comparison against a Chinchilla-optimal larger model or against a larger model receiving any test-time compute augmentation. The 14× scaling factor and the per-R crossover points are therefore specific to a baseline that is weaker than what a practitioner deploying at scale would actually use. The numbers are honest — the paper states its assumptions clearly — but the practical interpretation (when to choose test-time compute over pretraining) may be overly optimistic.

Mitigation status. The paper explicitly flags this as future work in Section 8: "analyzing the FLOPs-matched comparison with compute-optimal pretraining where both data and parameters are scaled." No sensitivity analysis is provided to bound how much the results would change under a stronger baseline.

---

Sequential Revision Strategies Impose Latency Costs Not Captured by Generation Count

The assumption or constraint. The paper measures test-time compute in "generations" — the total number of complete solutions sampled — as a proxy for total computational cost. This is reasonable for FLOP accounting but ignores wall-clock latency: a strategy that allocates 128 generations as 64 sequential revisions × 2 parallel chains takes approximately 64× longer wall-clock time than one that runs 128 fully parallel samples, assuming sufficient hardware for parallelism.

The consequence. The compute-optimal policies selected in Figures 4 and 8 favor sequential-heavy strategies on easy-to-medium problems. For instance, Figure 7 shows that on bin 3 (medium difficulty) at 128 generations, the optimal sequential-to-parallel ratio is around 2:1 to 8:1 — meaning the majority of the budget is spent on sequential revisions that cannot be parallelized. This creates a tension between throughput (total FLOPs per correct answer) and latency (time to first token or time to final answer) that the paper does not characterize.

For latency-sensitive applications — interactive assistants, real-time decision-making, or any setting where the user is waiting for a response — a strategy that takes 64× longer to produce an answer may be unacceptable regardless of its FLOP efficiency. The paper's exclusive focus on generation count as the cost metric obscures this tradeoff and could lead practitioners to deploy strategies that are optimizer-unfriendly in production.

What evidence exists in the paper. The paper provides no latency measurements, no wall-clock timing, and no discussion of the throughput-latency tradeoff. The generation budget axis in all figures (NFE on x-axis) is a parallelizable compute metric, not a latency metric.

Mitigation status. Not addressed. The paper does not mention latency as a consideration, does not propose latency-aware allocation strategies (e.g., capping the sequential chain length), and does not discuss how a practitioner should trade off the accuracy gains from sequential revisions against the latency cost.

---

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate with Only Patchwork Mitigation

The assumption or constraint. The revision model was trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target (Section 6.1). At test time, when the model encounters a correct answer in its own revision chain (produced in an earlier step), it has no training signal for what to do — and frequently revises correct answers into incorrect ones. The paper reports:

"approximately 38% of correct answers get converted back to incorrect ones"

The consequence. The sequential revision strategy can actively damage performance if deployed naively — the model may produce a correct answer at step 3 of a 16-step chain, then incorrectly revise it at step 4, losing the correct solution. The paper's mitigation is to use within-chain selection (majority voting or verifier-based scoring) that picks the best answer from any point in the chain rather than taking the final revision output. However, this selection mechanism still wastes the computational effort spent on revisions that overwrite correct answers — those steps consume compute without contributing to quality and can introduce incorrect answers that the selection mechanism must then distinguish from the correct one.

Additionally, the revision model's fragility is demonstrated by the ReST^EM experiment (Appendix K, Figure 16): when the authors attempted to further optimize the revision model using on-policy reinforcement learning, performance substantially degraded. The paper hypothesizes this is due to spurious correlations in on-policy revision data. This suggests the positive revision results depend on specific offline data construction choices (edit-distance-based pairing, fixed incorrect-to-correct ratio) that may not transfer to other training paradigms or model families.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1 (no figure/table citation). The ReST^EM degradation is shown in Appendix K, Figure 16. The paper does not report:

• How the reversion rate varies with problem difficulty (it may be higher or lower on the easy problems where sequential revisions are deployed).
• Whether the reversion rate decreases with longer training or different data construction (e.g., including correct-to-correct examples in training).
• How much the within-chain selection mechanism recovers — i.e., the accuracy after selection vs. the accuracy if reversions didn't occur.

Mitigation status. Partial. The within-chain selection mechanism (majority voting or verifier-based selection) mitigates the consequence of reversions but not the occurrence — the model still wastes compute producing incorrect revisions from correct answers. A more principled solution — such as training the model with a "no-change-needed" token or including correct-to-correct examples in the training data — is not explored. The paper does not flag this as future work.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the discrete diffusion landscape from pessimism about distillation feasibility to a demonstrated capability with clear remaining challenges. Before D-MMD, the default assumption in the field was that discrete diffusion models could not be effectively distilled — the factorized output architecture, the non-differentiable sampling step, and the exponential correlation problem each seemed like fundamental barriers. SDTT showed limited progress but at the cost of mode dropping. Di4C showed promise but required architectural changes with exponential scaling costs. The paper's Section 2 framing captured this stagnation: "It is currently difficult to distill discrete diffusion models."

D-MMD changes this conversation in three specific ways:

First, it demonstrates that stochastic distillation works for discrete diffusion, and does so without architectural modifications. The CIFAR-10 results in Table 1 — where a 16-step masked D-MMD generator achieves FID 5.3 versus the teacher's 6.4 at 1,024 steps — are not incremental. They represent a ~64× step reduction with simultaneous quality improvement. For a field where prior methods struggled to match teacher quality at any step count, this is a qualitative threshold crossing. The finding that students can outperform their teachers (discussed in Section 6.6) recontextualizes distillation from "lossy compression" to "distributional refinement" — the student is not merely approximating the teacher but finding a better operating point on the quality-diversity Pareto frontier.

Second, it provides a loss-agnostic formulation of stochastic distillation (Equation 9) that decouples the distillation principle from any specific domain. The min-max objective $\min_\eta \max_\phi \mathbb{E}[L_s(x, \hat{x}_\theta) - L_s(x, \hat{x}_\phi) - L_s(\hat{x}_\theta, \hat{x}_\phi)]$ is parametric in the loss function $L_s$. The verification that substituting squared error recovers continuous MMD (Equation 10) and substituting cross-entropy yields working discrete distillation (Equation 11) establishes that this is not a one-off trick but a genuine generalization. This reframes the research question from "can we make distillation work for discrete models?" to "what loss functions and generator architectures yield the best distillation outcomes under this framework?" — a much more productive question that invites systematic exploration rather than bespoke solutions for each domain.

Third, the GPT-2 Gradient Moment metric (Section 5) provides a tool the field has been missing. Prior work relied on generative perplexity, which Figure 3 demonstrates is actively misleading — it improves as models mode-collapse. The gradient moment rests on a principled foundation (a converged model has zero expected gradient on its training distribution; deviations from zero indicate distributional mismatch) and correctly penalizes the mode-collapsing methods that perplexity rewards (Table 5: Duo + DCD perplexity 4.82 but GM 108.2; D-MMD perplexity 17.2 but GM 0.236). The metric's validity for conditional generation (noted in Section 5) extends its utility beyond the unconditional setting studied in this paper. This is infrastructure: future discrete diffusion work can and should report gradient moment alongside perplexity, preventing the field from optimizing for the wrong signal.

The research directions this work makes less attractive:

• Pure deterministic distillation for discrete diffusion (SDTT-style progressive approaches). The comparison in Table 5 shows SDTT achieving GM 0.293 at 64 NFEs while D-MMD reaches 0.236 at 16 NFEs — and SDTT degrades with further distillation rounds. The paper's theoretical critique (Section 4: factorized outputs cannot represent correlated distributions like perfectly correlated coin tosses) identifies a fundamental ceiling. Deterministic distillation in the factorized space appears to hit an information-theoretic barrier that D-MMD's adversarial formulation circumvents.

• Explicit mixture distribution modeling as the primary solution to the correlation problem (Di4C-style). D-MMD achieves FID 5.0 at 8 NFEs versus Di4C's FID 20.6 at 10 NFEs (Table 4) without any output distribution modifications. The entropy collapse mechanism (Section 3.1, Table 6) shows that adversarial training pressure, combined with noise conditioning, can induce implicit correlation learning that outperforms explicit mixture modeling at lower step counts. This doesn't mean mixture distributions are obsolete — they may still be useful in combination with D-MMD — but the paper demonstrates that they are not necessary.

Follow-Up Research This Work Enables

Systematic exploration of loss functions in the D-MMD framework. The paper shows that cross-entropy works, and verifies that squared error recovers continuous MMD, but the space of possible loss functions is vast. A natural experiment: replace the cross-entropy in Equation 11 with a contrastive loss (e.g., InfoNCE-style, where the student maximizes similarity to teacher predictions while minimizing similarity to auxiliary predictions and negative examples drawn from the data). Would contrastive losses provide stronger gradient signal for rare tokens, where cross-entropy gradients can be vanishing? A systematic study training D-MMD on CIFAR-10 with squared error, cross-entropy, KL divergence, Jensen-Shannon divergence, and contrastive losses — all within the same generator architecture — would map the relationship between loss function choice and distillation quality, providing engineering guidance and potentially revealing whether certain losses interact better with the entropy collapse mechanism.

Scaling D-MMD to model sizes and sequence lengths where the factorized limitation should bite hardest. The paper's experiments use CIFAR-10 (3072 tokens) and OWT (1024 tokens). The entropy collapse mechanism argues that as sequence length grows, the generator must collapse its factorized output further to maintain coherence, because the number of spurious independence violations grows with sequence length. At what sequence length does this collapse become so extreme that the generator can no longer represent the data distribution's diversity? A stress-test: apply D-MMD to train 512-token, 2048-token, and 4096-token generators on OWT or similar text data, measuring both GPT-2 GM and explicit diversity metrics (distinct n-gram counts, self-BLEU, vocabulary coverage). If diversity degrades sharply beyond some sequence length, that identifies a fundamental scaling limit for factorized discrete diffusion that architectural changes (e.g., hybrid autoregressive-diffusion models) would need to address.

Combining D-MMD adversarial training with Di4C-style mixture distribution outputs. The paper positions entropy collapse as an alternative to mixture distributions, but the two approaches address different aspects of the correlation problem. Mixture distributions provide explicit capacity to represent multi-modal joint distributions. Adversarial training provides the optimization pressure to use that capacity effectively. An experiment: train D-MMD generators where the student outputs $K$ mixture components (as in Di4C), varying $K$ from 1 (factorized, equivalent to standard D-MMD) to 16 or 32, on CIFAR-10 with matched compute budgets. Does the mixture capacity allow the generator to maintain higher output entropy (less collapse) while achieving the same or better FID? Does the adversarial training make the mixture components more useful than they are under standard maximum-likelihood training? This would clarify whether D-MMD and Di4C are substitutes or complements.

Characterizing and closing the gap between D-MMD-distilled models and autoregressive models. The AR baseline achieves GPT-2 GM 0.061 on OWT (Table 2), while the best D-MMD generator achieves 0.225 at 32 NFEs — a ~3.7× gap. Is this gap attributable to the factorized output architecture (fundamental), the teacher quality (potentially closable with stronger teachers and further distillation), or the step count (solvable with more distillation research)? A controlled experiment: train a masked diffusion teacher of comparable parameter count and training FLOPs to the AR baseline, distill it with D-MMD at various step counts, and measure how the gap changes with teacher quality and student steps. If the GM gap shrinks proportionally with teacher quality improvements, the AR gap is an artifact of current teacher strength and will close naturally as discrete diffusion training improves. If the gap persists at arbitrarily strong teachers, the factorized architecture imposes a ceiling — and the practical value proposition of D-MMD-distilled diffusion shifts from "replacing AR models" to "complementing AR models in block-autoregressive settings."

D-MMD for discrete flow matching processes beyond masked and uniform diffusion. The paper demonstrates D-MMD for two diffusion processes but derives the generalized objective agnostically. The appendix (Section D) sketches the extension to processes where $\mathbb{E}_q[x|z_t]$ is not the optimal denoiser, requiring the posterior-aware loss $\text{KL}(\pi_{z_s}(x, z_t) \| \pi_{z_s}(\hat{x}_\theta, z_t))$. A systematic validation: implement D-MMD for discrete flow matching with alternative interpolation paths (e.g., linear interpolation with different base distributions, or score-based discrete processes from Lou et al., 2023). Do the distillation gains (student over teacher, 16×+ step reduction) transfer across process types? Are there diffusion processes where the gap between student and teacher is larger or smaller? This would establish whether D-MMD is a universal discrete diffusion distillation tool or requires process-specific tuning.

Cheap on-the-fly difficulty estimation to make the compute-optimal framework deployment-ready. (Note: this direction is flagged in the paper's Section 3.2 and Section 8 as critical future work, and the prior sections have highlighted its importance.) A concrete experiment: train a lightweight classifier — e.g., a small transformer or even a linear probe on top of the base LLM's final hidden state — to predict the difficulty quintile of a MATH question from the question text alone (or from the question text plus the base model's initial logprobs on a single greedy decode). Train on the 12,000-question MATH training set using oracle difficulty labels (or PRM-based difficulty labels), evaluate difficulty prediction accuracy on the 500-question test set, and then measure the compute-optimal scaling curve when using this cheap difficulty estimator instead of the 2,048-sample PRM-based estimator. If the cheap estimator achieves comparable strategy selection to the expensive estimator, the 4× efficiency gains become practically realizable. If not, the gap between predicted and oracle difficulty in Figures 4 and 8 gives an upper bound on what a perfect difficulty estimator could achieve.

Practical Applications and Downstream Use Cases

Cost-efficient text generation for batch processing. For organizations generating large volumes of text (content generation, synthetic data creation, document processing), D-MMD's 16× step reduction translates directly to cost savings. A masked diffusion model requiring 256 forward passes per generation would cost ~16× more in inference compute than a 16-step D-MMD generator producing better-quality output (GPT-2 GM 0.236 vs. 0.275, Table 2). In a batch setting processing millions of documents, this difference determines economic viability. The block-autoregressive configuration (Section 6.3) is particularly relevant here: combine a standard autoregressive model for context encoding with D-MMD-distilled diffusion blocks for efficient parallel token generation, amortizing the diffusion cost over 256-token blocks rather than per-token.

On-device image generation with reduced latency. The CIFAR-10 results (Table 1) demonstrate that D-MMD generators achieve high quality at very low step counts: FID 5.3 at 16 NFEs, competitive with continuous diffusion models that typically require specialized hardware. For on-device image generation (mobile photo editing, AR filters, sketch-to-image), where model size and inference latency are tightly constrained, a 16-step discrete diffusion generator could run on-device at interactive speeds, whereas a 1,024-step teacher would be completely impractical. The uniform diffusion results (FID 7.1 at 4 NFEs, 5.0 at 8 NFEs) are especially relevant for latency-critical applications — 4-step generation can complete in under 100ms on modern mobile hardware, making real-time generation feasible.

Verifier development for generative models using the gradient moment. The GPT-2 Gradient Moment metric has immediate practical utility beyond academic evaluation. For teams developing discrete generative models, the metric can serve as a continuous quality monitor during training: periodically compute the gradient moment on a held-out validation set, and use degradation in the metric (increase from the minimum) as an early stopping signal or a detection mechanism for training instability. Because the metric is principled (zero at distribution match) and robust to mode collapse (unlike perplexity), it provides a reliable signal that the model is diverging from the data distribution — something perplexity-based monitoring would miss. The stochastic approximation (Equation 14) makes this computationally practical even for large reference models.

When to Prefer This Method

The paper articulates a clear tradeoff between D-MMD and prior discrete distillation approaches, grounded in the factorized output limitation and the mode-dropping behavior of deterministic methods:

• Prefer D-MMD when you need few-step generation (4–64 NFEs) from a discrete diffusion model and sample quality (measured by FID or GPT-2 GM) is the primary objective. The 16×–64× step reductions with simultaneous quality improvements (Tables 1, 2) make it the strongest available method in this regime. D-MMD is especially indicated when you can afford input noise conditioning (for masked diffusion), when you want mode-seeking behavior in the generator (via temperature or top-p distillation, Section 3.3), and when your application can tolerate the adversarial training complexity (alternating optimization, tuning of auxiliary model capacity and update ratio).

• Prefer SDTT or progressive distillation when you need deterministic distillation (predictable quality-step tradeoff without adversarial training instability), when you are operating at very high step counts (64+ NFEs) where D-MMD's advantage over the teacher narrows and SDTT's mode-dropping behavior is less severe, or when preserving the teacher's exact diversity profile (sample entropy) is more important than improving generation quality metrics. The sample entropy results in Table 5 show SDTT at 64 NFEs maintains entropy 5.17 (close to teacher's 5.19), while D-MMD at 32 NFEs drops to 5.05 — SDTT preserves more of the teacher's diversity.

• Prefer Di4C-style mixture distributions when you need explicit, interpretable control over output correlations (e.g., for controllable generation where different mixture components correspond to different output modes), or when the factorized architecture imposes an unacceptable quality ceiling at your target sequence length. The paper does not demonstrate D-MMD at sequence lengths beyond 3,072 tokens (CIFAR-10 pixels), and the entropy collapse mechanism may become harder to achieve as sequence length grows — explicit mixture modeling provides a more direct (if more expensive) path to correlation.