Distribution Matching Distillation for Continuous Diffusion Language Models

Paul Le Van Kiem; Dario Shariatian; Umut Simsekli; Alain Durmus

cs.LGSep 30, 2026

Distribution Matching Distillation for Continuous Diffusion Language Models

Authors: Paul Le Van Kiem, Dario Shariatian, Umut Simsekli, Alain Durmus

Organizations: Inria, PSL Research University · Cohere · CMAP, Ecole Polytechnique

Abstract

Continuous diffusion language models generate all tokens in parallel, yet high-quality generation can still require hundreds of network evaluations (NFEs). We study how distributional distillation can reduce this cost by exploiting the student's probabilistic token outputs. Our unified formulation connects the student's output parameterization to the resulting gradient estimators and yields two methods with the same student architecture and reverse-KL matching objective: Simplex-DMD uses continuous token relaxations and pathwise gradients, while Reinforce-DMD uses categorical sampling and REINFORCE with a learned density ratio. We develop both methods for multi-step generation and investigate the training and sampling choices associated with each parameterization. On OpenWebText, for sequences of 1,024 tokens, Simplex-DMD achieves a generative perplexity of 45.6 at a unigram entropy of 5.44 nats in just 4 NFEs, a 49% reduction relative to the strongest evaluated diffusion baseline at matched entropy and sampling budget. Reinforce-DMD improves the frontier at larger budgets, reaching a generative perplexity of 14.9 at an entropy of 5.00 nats with 256 NFEs, a 20% reduction under the same comparison protocol.

Figures & tables

Figure 1: Gen PPL on OpenWebText against number of function evaluations (NFE), at the matched unigram entropy

H_{\mathrm{NFE}}

(top axis), defined as in Table 2 . Each method’s log Gen PPL is linearly interpolated at

H_{\mathrm{NFE}}

between the two sampling temperatures whose entropies bracket it. LangFlow and MDLM are undistilled models; FMLM, ReDi and D-MMD are distillation methods. A method that does not reach

H_{\mathrm{NFE}}

is shown at the closest measured entropy, with its value in parentheses. The dashed purple line marks the Gen PPL of the data, with its entropy in parentheses.

Criterion	Local discrepancy $\mathsf{D}$	Representative methods
Continuous diffusion
Reverse KL	$\log p_{t}^{\eta}(z_{t})-\log p_{t}(z_{t})$	Diff-Instruct ( Luo et al., 2023 ) , DMD ( Yin et al., 2024 )
Score / Velocity / Moment matching	$\big\lVert\nabla_{z_{t}}\log p_{t}^{\eta}(z_{t})-\nabla_{z_{t}}\log p_{t}(z_{t})\big\rVert^{2}$	SiD ( Zhou et al., 2024 ) ; FGM ( Huang et al., 2024 ) ; Moment matching ( Salimans et al., 2024 )
Discrete diffusion
Bregman divergence	$\displaystyle\sum_{d_{\mathrm{H}}(z^{\prime},z_{t})=1}d_{F}\!\left(\frac{p_{t}^{\eta}(z^{\prime})}{p_{t}^{\eta}(z_{t})},\,\frac{p_{t}(z^{\prime})}{p_{t}(z_{t})}\right)$	IDLM ( Li et al., 2026 ) ; D-MMD ( Hoogeboom et al., 2026 )
Posterior $f$ -divergence	$\displaystyle\sum_{\ell}D_{f}\!\left(p_{\mathbf{x}\mid t}^{\eta,\ell}(\cdot\mid z_{t}),\,p_{\mathbf{x}\mid t}^{\ell}(\cdot\mid z_{t})\right)$	DiMO ( Zhu et al., 2025 )

Table 1: Local discrepancies for distributional distillation. Score, velocity and moment discrepancies agree up to time-dependent factors. Definitions and derivations are given in Appendix C .

Figure 2: Generative frontiers on OpenWebText at

\mathrm{NFE}=4

and

256

. Temperature varies from

0.8

1.1

; stars mark temperature

1.0

and diamonds the data. Gen PPL uses an inverted log scale. Dashed GPT-2 and OPT use

1{,}024

evaluations. Simplex-DMD is omitted at

\mathrm{NFE}=256

due to low diversity. Full sweeps are in Section E.2 .

	Network evaluations
Method	2	4	8	16	128	256
$H_{\mathrm{NFE}}$	5.65	5.44	5.38	5.20	5.27	5.00
LangFlow ( Chen et al., 2026b )	1077	317	147	58.6	39.1	21.3
MDLM ( Sahoo et al., 2024 )	1909	678	267	72.8	44.3	23.8
FMLM ( Lee et al., 2026 )	142 (5.32)	128 (5.40)	94.7	55.2	22.6 (4.87)	13.5 (4.50)
D-MMD ( Hoogeboom et al., 2026 )	1752	431	135	41.7	27.3	18.6

Table 2: Gen PPL on OpenWebText at matched unigram entropy

H_{\mathrm{NFE}}

; lower is better.

H_{\mathrm{NFE}}

is the Simplex-DMD entropy closest to the data for

\mathrm{NFE}=2

–

16

, and the corresponding Reinforce-DMD entropy for

\mathrm{NFE}=128,256

. Log Gen PPL is linearly interpolated at

H_{\mathrm{NFE}}

. If a curve does not reach the target entropy, its closest point is reported with entropy in parentheses. Bold indicates the lowest matched value.

Figure 4: Sampler comparison for

c\in\{0,0.25,0.5,0.75,1,1.5\}

and forward renoising.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Teacher	Teacher updates	Distillation updates	Student params.
SDTT ( Deschenaux & Gulcehre, 2025 )	MDLM	1M	$r\times 10$ k $r=1,\ldots,7$	$\approx$ 170M
IDLM ( Li et al., 2026 )	MDLM	1M	1M	169.63M
FMLM ( Lee et al., 2026 )	FLM	1.5M	1M	$\approx$ 170M
ReDi ( Yoo et al., 2025 )	DUO	$\sim$ 500k	$\leq$ 1M	$\approx$ 165M
Simplex-DMD (ours)	LangFlow	1M	10k	$\approx$ 170M
Reinforce-DMD (ours)	LangFlow	1M	10k	$\approx$ 170M

Appendix

Table 3: Training budgets and model sizes for the distilled checkpoints. Teacher updates are pretraining steps; distillation updates are additional optimizer updates. For our methods, the 10k total includes 8k auxiliary and 2k student updates. SDTT uses successive round checkpoints, indexed by

r

. The ReDi student budget is an upper bound. Parameter counts refer to one sampling model, including embeddings, and exclude auxiliary training networks.

Component	Simplex-DMD	Reinforce-DMD
Student initialization	teacher	teacher
Auxiliary initialization	teacher	teacher
Renoising kernel	$q_{t\mid\mathbf{x}}$	$q_{t\mid\mathbf{x}}$
Self-conditioning	on	off
KL anchor $\beta$	–	$0.5$
Sampler	forward	forward

Appendix

Table 4: Training and sampling configuration of each student. A dash marks a component the objective does not have. The last row is applied at inference.

Figure 6: Design space on OpenWebText, one component changed per variant, read as in Figure 2 . Each student is shown at the two budgets where it is competitive.

Figure 7: Auxiliary loss for Simplex-DMD, read as in Figure 2 : the soft target of ( 17 ) (CE soft), the squared error of ( 83 ) (L2) and the sampled-token target of ( 84 ) (CE hard). The three variants trace one frontier and differ in where their temperature-

1.0

point sits on it, the sampled-token target lowest in entropy and squared error highest.

Figure 8: Forward noising

q_{t\mid\mathbf{x}}

against the bridge

q_{t\mid T,\mathbf{x}}

as the training renoising kernel for Simplex-DMD, read as in Figure 2 . Each variant is a separate training run.

Explore similar work

May 22, 2026cs.LG

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Diffusion language models intrinsically fail to capture correlations between decoded tokens, which leads to a harsh trade-off between sampling quality and throughput. To solve this issue, we propose DiLaDiff, a variant of masked diffusion language models with three components: (1) a continuous latent space with semantic capabilities, learned by an auto-encoder fine-tuned from an existing masked diffusion language model; (2) a latent diffusion model learning the prior over the encoder distribution; (3) a consistency model distilling the learned prior into a few-step latent generative model. We show that, even without distillation, our latent-guided diffusion model outperforms the masked diffusion baseline while significantly accelerating inference. Consistency distillation further lowers the computational overhead of continuous diffusion, such that the latent is generated in negligible time compared to discrete decoding.

Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3

NVIDIA

Diffusion Language Models Dataset Distillation Language Modeling

Ask PDF

Oct 1, 2026cs.CL

Scaling and Distilling Text Embeddings for Better Diffusibility

Diffusion language models (DLMs) offer a promising alternative to autoregressive (AR) language generation. Recent advances in continuous DLMs, which apply latent diffusion to continuous text embeddings, raise a practical question: which embedding makes the best latent space, i.e., the most diffusible? To answer this, we search through different embeddings and find that scaling the embedding model to stronger ones within the same family (T5 to T5Gemma-1 to T5Gemma-2) greatly improves generative performance. But the raw T5Gemma-2 embeddings are still not optimal. They are so discriminative that even the embeddings of plausible alternative words are separated, which makes the generation vulnerable to imperfect sampling. Consequently, continuous diffusion often fails to reach any of them and ends up at an invalid embedding instead. To address this, we distill T5Gemma-2 into a student encoder that learns the teacher's decoded probabilities as soft labels. Learning from such soft labels makes the student pull the alternative embeddings closer while maintaining the encoding-decoding mechanism. The distilled embeddings form a more connected and diffusible latent space, improving over the vanilla T5Gemma-2 embeddings. As a result, our medium-sized DLM achieves Gen. PPL 17.8 (against real-text PPL 15.4) at real-text entropy on OpenWebText, outperforming GPT-2-M on Gen. PPL.

Zekai Zhang, Yunjie Tian, Yanjin He +4

University of Michigan

Diffusion Language Models Text Embeddings Diffusion Dynamics Latent Space Scaling Entropy

Ask PDF

Apr 30, 2026cs.LG

Consistent Diffusion Language Models

Diffusion language models (DLMs) are an attractive alternative to autoregressive models because they promise sublinear-time, parallel generation, yet practical gains remain elusive as high-quality samples still demand hundreds of refinement steps. In continuous domains, consistency training along the probability-flow ODE is a popular recipe to accelerate diffusion. For discrete diffusion, no analogous sample-space ODE exists, making direct adaptation ill-defined. We argue that the right discrete substitute is the exact posterior bridge, the closed-form conditional law linking any two noise levels, which is available for broad corruptions including masked and uniform diffusion. Building on this observation, we introduce Multi-Path Discrete Consistency (MPDC), a new principle that trains a denoiser to be path-invariant in expectation across these stochastic bridges, and instantiate it as the Consistent Diffusion Language Model (CDLM), a single-stage training framework that does not require an already trained teacher model. Our CDLM objective recovers masked diffusion, continuous consistency models, and progressive or discrete distillation as analytic limits or empirical approximations of one common view. Empirically, CDLM establishes a new state of the art on both conditional and unconditional text-generation, consistently outperforming strong base discrete diffusion models and often even multi-stage distilled baselines across sampling budgets, with the largest gains in the few-step regime. Together, these results position CDLM as a principled and scalable foundation for the next generation of fast, high-fidelity discrete generative modeling.

Hasan Amin, Yuan Gao, Yaser Souri +4

Department of Computer Science, Purdue University · Microsoft

Diffusion Language Models Diffusion Models

Ask PDF