Distribution Matching Distillation for Continuous Diffusion Language Models

1 INTRODUCTION

• Diffusion language models operate directly in discrete token space or through Gaussian diffusion in continuous embedding space, with recent continuous models reported as competitive with discrete counterparts.High-quality generation can still require hundreds of sequential network evaluations, motivating distillation to reduce evaluations while preserving quality.Source: §1 · p. 1
• Trajectory-based methods supervise predictions or transitions along teacher dynamics, whereas distributional methods guide the student toward the data distribution without prescribing generation trajectories.Distributional distillation has been explored for continuous and discrete data, including language.Source: §1 · p. 1
• Simplex-DMD uses continuous token relaxations and pathwise gradients, while Reinforce-DMD uses categorical sampling and REINFORCE with a learned density ratio.Both train the same student architecture for multi-step generation under a reverse-KL matching objective.Source: §1 · p. 2
• A unified framework compares noised student and data distributions through a generic discrepancy, recovers prior objectives as particular choices, and extends to multi-step generation.The framework is specialized to reverse-KL matching for the two proposed methods.Source: §1 · p. 2
• At 4 evaluations and unigram entropy 5.44, Simplex-DMD reduces Gen PPL from 90.2 for ReDi to 45.6; at 256 evaluations and entropy 5.00, Reinforce-DMD reduces it from 18.6 for D-MMD to 14.9.The evaluation uses 1,024-token sequences on OpenWebText.Source: §1 · p. 2

2 PRELIMINARIES: CONTINUOUS DIFFUSION FOR DISCRETE DATA

• Continuous diffusion models for discrete data embed sequences, reverse Gaussian corruption in embedding space, and apply a categorical decoder to produce token sequences.The embedding matrix may be fixed or learned jointly with the model.Source: §2 · p. 3
• The forward process corrupts embedded data with Gaussian noise on a decreasing schedule, with conditional marginals centered at α_t xE and covariance σ_t^2 I_d.The forward process is Markovian in the embedding space.Source: §2 · p. 3
• Reverse generation starts from terminal noise and applies learned transitions along a time grid, with a denoising network predicting distributions over clean tokens.The transition model extends the conditional law to simplex-valued denoiser inputs.Source: §2 · p. 3
• The diffusion loss trains embedded denoiser outputs toward the posterior mean, while direct cross-entropy training is an alternative and is not equivalent to the NELBO.With suitable weighting, cross-entropy can yield a likelihood bound.Source: §2 · p. 3

3 DISTRIBUTION MATCHING DISTILLATION FOR CONTINUOUS DIFFUSION ON DISCRETE DATA

• The shared distribution-matching objective yields Simplex-DMD for continuous relaxation and Reinforce-DMD for discrete sampling, with different gradient estimators.Both methods extend to multi-step generation.Source: §3 · p. 4

3.1 A UNIFIED VIEW OF ONE-STEP DISTRIBUTION MATCHING

• The one-step framework compares noised student and data marginals using a generic discrepancy; the methods specialize it to reverse KL.Different discrepancy choices recover objectives used in prior work.Source: §3.1 · p. 4

3.2 MULTI-STEP GENERATION

• Multi-step training matches renoised student samples to data marginals across starting noise levels, then composes transitions along decreasing noise levels.The transition family controls sampling stochasticity, from a deterministic update at ζt,T = 0 to other noise-injection choices.Source: §3.2 · p. 5

3.3 SIMPLEX-DMD: CONTINUOUS RELAXATION

• Simplex-DMD uses the student’s probability vectors as clean samples, enabling pathwise gradients through differentiable Gaussian noising.An auxiliary denoiser estimates the student score; the frozen teacher denoiser estimates the data score.Source: §3.3 · p. 6

3.4 REINFORCE-DMD: DISCRETE SAMPLING

• Reinforce-DMD samples discrete clean tokens and uses a score-function gradient, estimating its log-ratio cost with a discriminator trained to distinguish noised data from student samples.A leave-one-out baseline centers costs, and experiments include a KL anchor for stable training.Source: §3.4 · pp. 6–7

3.5 RELATED WORK

• Related work covers continuous diffusion for discrete data, trajectory distillation, and distributional distillation, whose methods differ mainly in local discrepancy and gradient estimation.The paper summarizes local discrepancies in Table 1.Source: §3.5 · pp. 7–8

4 EXPERIMENTS

• The study distills the LangFlow teacher on OpenWebText and evaluates 1,024-token sequences using Gen PPL–entropy frontiers across sampling budgets.Both students and auxiliary networks are initialized from the teacher; the teacher embedding matrix is fixed.Source: §4 · p. 8

4.1 GENERATION QUALITY ACROSS SAMPLING BUDGETS

• At 4 NFEs, Simplex-DMD cuts Gen PPL from 90.2 (ReDi) to 45.6 at entropy 5.44 (49%); at 256, Reinforce-DMD cuts 18.6 (D-MMD) to 14.9 at entropy 5.00 (20%).Figure 2 shows the generative frontiers at these budgets.Source: Figure 2 · p. 8, §4.1 · p. 8
• Simplex-DMD has the lowest matched-entropy Gen PPL at 2, 4, 8, and 16 evaluations; Reinforce-DMD matches D-MMD at 128 and leads at 256.At 2 evaluations, FMLM's lower raw Gen PPL is measured at substantially lower entropy, so it is not a matched-diversity comparison.Source: Table 2 · p. 9, §4.1 · p. 9

4.2 TRAINING AND SAMPLING CHOICES

• At NFE = 128, removing Reinforce-DMD's teacher KL anchor collapses generation to essentially zero unigram entropy; positive weights yield similar frontiers, with β = 0.5 modestly best.The authors use β = 0.5 in the default configuration.Source: Figure 3 · p. 9, §4.2 · p. 9
• Across the swept budgets, forward renoising gives both Simplex-DMD and Reinforce-DMD the best observed Gen PPL–entropy frontier.Forward renoising uses ζt,T = σt, adding the maximum amount of fresh Gaussian noise at each sampling step.Source: Figure 4 · p. 10, §4.2 · p. 10
• Additional ablations show that scratch initialization collapses, self-conditioning helps only Simplex-DMD, and Simplex-DMD's bridge-based training underperforms forward renoising.The authors therefore initialize from the teacher and use self-conditioning only for Simplex-DMD.Source: §4.2 · p. 10

5 CONCLUSION

• The paper introduces a distribution-matching framework for distilling continuous diffusion models on discrete data, with Simplex-DMD strongest at low budgets and Reinforce-DMD at larger budgets among evaluated methods.The methods share a reverse-KL objective but use simplex-valued outputs and pathwise gradients, or categorical sampling and score-function estimation.Source: §5 · p. 10
• The methods require a pretrained teacher and an auxiliary model during training; experiments are limited to academic model scale and Gen PPL–entropy evaluation.Scaling to substantially larger language models will require studying training overhead and extending evaluation beyond this protocol.Source: §5 · p. 10

A NOTATION

• The notation defines X as the space of length-L token sequences and represents discrete sequences with one-hot rows.The vocabulary is V and each token is identified with a canonical basis vector.Source: A NOTATION · p. 15

B.1 CONTINUOUS DIFFUSION ON DISCRETE DATA: OBJECTIVE AND DESIGN CHOICES

• Likelihood-based training uses a negative evidence lower bound, whereas direct cross-entropy supervises clean-token probabilities at sampled noise levels.The NELBO includes prior, reconstruction, and embedding-space diffusion terms; the unweighted cross-entropy is not the NELBO. Cross-entropy encourages distinguishable embeddings, while spherical normalization controls embedding scale but cannot prevent token collisions.Source: §B.1 · pp. 16–17
• Self-conditioning reuses the preceding denoising prediction at sampling, while training uses preliminary predictions on a random subset and stops their gradients.Other examples use a zero conditioning input so the network can denoise at the first sampling step; only the second pass receives gradients for self-conditioned examples.Source: §B.1 · p. 17
• Noise schedules control information loss through γ(t), and the experiments reuse the teacher’s pretrained schedule without changing it during distillation or sampling.LangFlow fits a scaled Gumbel CDF to token cross-entropy to allocate more noise levels where it changes rapidly; RePlaid learns schedule endpoints and a monotone interpolation to reduce estimator variance.Source: §B.1 · p. 17

B.2 DISCRETE DIFFUSION: THE BOUND, THE PLUG-IN AND THE CONCRETE SCORE

• A factorized reverse model’s transition error separates token-marginal mismatch from irreducible position dependence; the bridge plug-in recovers the exact marginals using a leave-one-out denoiser.For uniform replacement, this is the only simplex-valued field with that property; for mask replacement, the ordinary denoiser also has it at masked positions.Source: §B.2 · pp. 18–19
• In continuous time, the discrete diffusion bound compares concrete scores through a rate-weighted Bregman divergence, and SEDD parameterizes that score directly.The data score is unavailable, so training averages the score over the denoising posterior; replacing it by a posterior draw changes the divergence only by a constant in θ.Source: §B.2 · p. 19

C DISTRIBUTIONAL DISTILLATION: DISCREPANCIES AND GRADIENTS

• The appendix derives the objectives, gradients, and identities for the distillation methods in the one-step setting, with the student starting from the diffusion terminal law.Its clean output induces a sequence distribution whose noised marginals are used in the derivations.Source: C DISTRIBUTIONAL DISTILLATION:… · p. 19

C.1 CONTINUOUS DISTILLATION BASELINES: SID, FGM AND MOMENT MATCHING

• FGM, SiD, and one-step moment matching are members of one gradient family, differing in how they combine two losses and whether they freeze the sampled state.FGM gives the exact gradient and is SiD at α = 1/2; one-step moment matching is the α = 0 case. The reported best SiD α values lie beyond α = 1/2, and at α = 1 its pathwise term disappears (Huang et al., 2024).Source: §C.1 · p. 21

C.2 DISCRETE DISTILLATION AS A BREGMAN DIVERGENCE BETWEEN CONCRETE SCORES

• Discrete distillation’s reverse-KL objective becomes a Bregman divergence between teacher and student concrete scores, evaluated on states visited by the student.The teacher and a student-sample-fitted auxiliary model stand in for the unavailable scores; subtracting their per-sample training losses cancels a term independent of the score.Source: §C.2 · pp. 21–22
• IDLM writes the integrand as teacher loss on student samples minus its minimum over auxiliary fields, while D-MMD also pulls the auxiliary toward the teacher.Under the mask, the divergence takes a simpler form, derived in Section C.3.Source: §C.2 · p. 22

C.3 MASKED DIFFUSION DISTILLATION: COLLAPSE TO A POSTERIOR KL

• Under the mask law, only masked positions are scored, where the concrete score is the data denoiser up to a time-only factor.The corresponding student score uses the denoiser of the noised student distribution, not the student generator.Source: §C.3 · p. 22
• In practice, an auxiliary denoiser and the teacher approximate the denoisers in the masked discrepancy; under the uniform law, every position is scored and the denominator keeps the head in the loss.The passage identifies the resulting masked discrepancy with DiMO’s listed discrepancy and D-MMD’s masked loss.Source: §C.3 · p. 22

• The appendix derives the gradient used to train the objective.Source: D METHOD: DERIVATIONS AND PROOFS · p. 22
• It proves that the vertex constraint is recovered.Source: D METHOD: DERIVATIONS AND PROOFS · p. 22
• It derives a discriminator that supplies the log-ratio and an alternative baseline objective.Source: D METHOD: DERIVATIONS AND PROOFS · p. 22

D.1 DISTRIBUTION MATCHING DISTILLATION: GRADIENT DERIVATION

• Simplex-DMD’s reverse-KL gradient includes a term through the generated latent; the explicit density-derivative term has zero expectation.Source: §D.1 · pp. 22–23
• The pathwise derivation assumes fixed embeddings, noise schedule, and sampling distribution, with sufficient regularity to interchange differentiation and integration.Source: §D.1 · p. 22
• The one-step gradient is recovered by setting T = 1 and sampling t uniformly.Source: §D.1 · p. 23
• In practice, a frozen teacher and auxiliary denoiser estimate the scores; the detached score difference sends gradients only through the reparameterized latent.Source: §D.1 · p. 23

D.2 RECOVERING THE DISCRETE DATA DISTRIBUTION

• Zero distribution-matching loss recovers the clean data law, including when the student outputs continuous probability vectors.Source: §D.2 · pp. 23–24
• For multi-step generation, matching at a noise level identifies the clean law and yields transports from each starting noise level to the data distribution.Source: §D.2 · p. 24

D.3 THE OPTIMAL DISCRIMINATOR AND THE LOG-RATIO

• The population-optimal discriminator recovers the log-density ratio between the noised data and student distributions for almost every sampled noise-level pair.Source: §D.3 · p. 25

D.4 AUXILIARY SCORE ESTIMATION WITH BRIDGE RENOISING

• Under bridge renoising, the usual auxiliary score estimator is approximate; estimating the exact bridge score requires an additional prediction of E[ZT | Zt].Source: §D.4 · p. 26

D.5 MULTI-STEP GENERATION

• For forward renoising, exact distribution matching at each starting level implies the required clean-marginal condition under Theorem 1’s geometric assumption.Source: §D.5 · p. 26

E.1 IMPLEMENTATION DETAILS

• Students receive 10,000 optimizer updates—8,000 auxiliary and 2,000 student updates—and use 170.8M-parameter LangFlow architectures with frozen, shared embeddings.The auxiliary model retains the LangFlow architecture; the Reinforce-DMD discriminator replaces the vocabulary head with a scalar output per position.Source: §E.1 · p. 27
• At inference, sampling traverses a uniform time grid from Tn = 1 to T1, followed by a final discrete decoder step, for n total NFEs.Temperature is applied to the output logits before softmax.Source: §E.1 · p. 27
• The Reinforce-DMD discriminator uses the teacher’s transformer backbone on noised latents and averages its per-position scalar logits into one sequence logit.Its scalar head is initialized separately from the teacher-initialized backbone.Source: §E.1 · p. 27
• Evaluation uses 1,024-token OpenWebText sequences, sampling budgets from 1 to 1,024 NFEs, and one random seed.Generative perplexity is computed with GPT-2-large; diversity is mean per-sequence Shannon entropy of empirical GPT-2 token frequencies.Source: §E.1 · p. 28

E.2 FULL GENERATIVE FRONTIERS

• The evaluation compares generative frontiers across all evaluated sampling budgets.Generative perplexity is paired with unigram entropy, and Figure 5 extends the comparison to every budget.Source: §E.2 · p. 28, §E.1 · p. 28

E.3 DESIGN SPACE

• Training Simplex-DMD from scratch collapses to near-uniform tokens, with unigram entropy about 6.9 and Gen PPL above 6 × 10^4.Without self-conditioning, the curve reaches slightly higher Gen PPL at the same entropy; changing the noise schedule shifts the temperature-1.0 point along the recipe’s curve.Source: §E.3 · p. 28
• With β = 0, Reinforce-DMD collapses to zero unigram entropy; positive β weights differ little, with β = 0.5 best among them.Adding self-conditioning moves the frontier to higher Gen PPL at every entropy.Source: §E.3 · p. 28
• Figure 6 shows one-component-at-a-time variants, each at the two budgets where that student is competitive.The bridge-renoising ablation for Simplex-DMD is reported separately.Source: §E.3 · p. 28, Figure 6 · p. 30

E.4 AUXILIARY LOSS

• Squared-error regression and cross-entropy on a sampled one-hot token share the soft-target objective’s population minimizer, leaving the gradient identity unchanged.The sampled-token estimator adds variance from the categorical draw.Source: §E.4 · p. 31
• The soft-target, squared-error, and sampled-token auxiliary losses trace one frontier; the sampled-token target has the lowest Gen PPL at the lowest entropy, while squared error is more diverse and less fluent.From NFE 32 to 128, unigram entropy falls by 1.03 nats for the soft target, 1.06 for squared error, and 1.17 for the sampled-token target; the soft target is retained because faster collapse can lower perplexity while reducing diversity.Source: §E.4 · p. 31, Figure 7 · p. 31

E.5 BRIDGE RENOISING FOR SIMPLEX-DMD

• At every displayed budget, forward noising yields lower Gen PPL than bridge renoising at the same entropy, with the gap narrowing as budget grows.The comparison is for Simplex-DMD; bridge renoising is not evaluated during Reinforce-DMD training.Source: §E.5 · p. 31, Figure 8 · p. 32

E.6 GENERATED SAMPLES

• For each model, the displayed temperature is chosen for unigram entropy closest to the data’s 5.44, and the first three stored samples are shown without selection, truncated to 180 words.The examples are from 1,024-token evaluations; <eos> marks an end-of-sequence token within a sequence.Source: §E.6 · p. 32
• Simplex-DMD at 4 NFEs has Gen PPL 45.6 and entropy 5.44, while Reinforce-DMD at 128 NFEs has Gen PPL 26.9 and entropy 5.27.The display also includes LangFlow at 1,024 NFEs with Gen PPL 85.6 and entropy 5.49, and MDLM at 1,024 NFEs with Gen PPL 65.5 and entropy 5.49.Source: §E.6 · pp. 32–35