Representation-Space MMD for Diffusion Language Models

Ilya Drobyshevskiy; Ilia Sudakov; Maksim Semenov; Denis Kuznedelev; Maksim Ignatov; Pavel Temirchev; Nikita Balagansky; Viacheslav Meshchaninov

cs.CLOct 5, 2026

Representation-Space MMD for Diffusion Language Models

Authors: Ilya Drobyshevskiy, Ilia Sudakov, Maksim Semenov, Denis Kuznedelev, Maksim Ignatov, Pavel Temirchev, Nikita Balagansky, Viacheslav Meshchaninov, +2 more

Organizations: Yandex Research · HSE University · T-Tech · Constructor University · Applied AI Institute · AXXX

Abstract

We introduce a post-training method for diffusion language models (DLMs) that minimizes Maximum Mean Discrepancy (MMD) between generated and reference distributions in the feature space of a frozen pretrained DLM. To estimate MMD, we retain contextual features at individual token positions, obtaining multiple observations per sequence from a single extractor pass. We optimize this objective using policy gradients for discrete models and direct differentiation through generated latents for continuous models. In both cases, computing the loss directly from these features enables efficient post-training without full sampling trajectories or jointly trained auxiliary models. Experiments show lower generative perplexity at comparable entropy on OpenWebText and better accuracy-computation trade-offs on GSM8K. On 16B DMax-LLaDA2.0 models with hybrid masked-uniform diffusion, we increase decoding parallelism with similar or higher accuracy on math and code benchmarks.

Figures & tables

Figure 1: Representation-space MMD training. We illustrate the case of one reference sequence and $B$ generated sequences per MMD estimate. For discrete DLMs ( top left ), we independently draw $G$ such batches from the token distributions produced by a single denoiser pass and optimize the loss with grouped REINFORCE. For continuous DLMs ( bottom left ), $K$ bootstrap passes provide self-conditioning without gradients, followed by a differentiable generation pass. In both cases, token-level RBF MMD compares contextual features extracted by a frozen pretrained DLM ( right ).

Figure 2: Gen. PPL–Entropy trade-offs on OpenWebText. Top: MDLM-based models with temperature sweeps at

8

16

, and

32

sampling steps. Bottom: Continuous models evaluated at different sampling budgets. Low-entropy ELF and ELF ⋆ points are omitted. Note: entropy above the reference data entropy (

H{\approx}5.43

) does not imply better performance.

Figure 3: GSM8K accuracy vs sampling steps. Left: MDLM-MMD and MDLM-based baselines. Right: ELF-MMD and continuous baselines. For masked models, we vary the confidence threshold and report the average number of steps per sequence.

Method	GSM8K		MATH500		Minerva- Algebra		ASDIV		HumanEval- Instruct		MBPP- Instruct
Method	Acc. $\uparrow$	TPF $\uparrow$	Acc. $\uparrow$	TPF $\uparrow$	Acc. $\uparrow$	TPF $\uparrow$	Acc. $\uparrow$	TPF $\uparrow$	Acc. $\uparrow$	TPF $\uparrow$	Acc. $\uparrow$	TPF $\uparrow$
LLaDA-2.0-mini	92.6	2.04	75.8	2.58	91.4	3.01	92.8	2.03	84.2	4.38	80.6	2.71
dParallel SFT	92.3	2.79	75.8	3.42	91.4	3.91	93.0	2.72	76.8	5.12	74.7	3.66
DMax-Math	92.1	5.48	75.4	5.94	91.5	7.03	92.5	5.62	–	–	–	–
DMax-Math-MMD	92.1	6.15	76.0	6.84	92.1	8.19	92.9	6.20	–	–	–	–
DMax-Coder	–	–	–	–	–	–	–	–	83.5	7.36	79.2	5.86

Table 1: Scaling MMD post-training to 16B DMax. Accuracy (

\%

) and tokens generated per forward (TPF) on math and code benchmarks. Baseline results are taken from the original DMax paper.

Figure 4: MMD loss ablation on GSM8K. Accuracy vs. sampling steps for MDLM ( left ) and ELF ( right ), post-trained with different feature-based objectives. ELF results are obtained without IRD.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: MMD loss ablation on OpenWebText. Attraction-only and feature regression losses are shown only for MDLM, as both perform poorly on ELF in the OWT experiments.

Figure 6: Representation-space ablation on GSM8K. Accuracy vs sampling steps for MDLM ( left ) and ELF ( right ), comparing input representations, external encoders, and pretrained DLM features.

Figure 7: Ablation on policy-gradient group size on OWT. Generative perplexity–entropy trade-offs at 8, 16, and 32 sampling steps. Increasing

G

from

1

2

yields the largest improvement.

Figure 8: Ablation on policy-gradient group size on GSM8K. Accuracy vs. average number of sampling steps.

Figure 9: Model-size scaling on GSM8K. ELF-B and ELF-M before (dashed) and after (solid) MMD post-training.

Figure 10: Ablation on bootstrap training. We vary the number of bootstrap steps on OWT with T5 (a) and GPT-2 (b) encoders, and on GSM8K (c). A single bootstrap step improves the trade-off with T5 and accuracy on GSM8K, while results with GPT-2 are mixed.

Figure 11: Pass@

k

on GSM8K. We report pass@

k

accuracy for different numbers of samples per prompt and sampling budgets of 4-32 steps for both discrete and continuous setups.

Steps	8		16		32
Method	gPPL	Ent.	gPPL	Ent.	gPPL	Ent.
Masked D-MMD	—	—	17.2	5.00	19.4	5.05
Masked D-MMD ( $p=1.0$ )	—	—	67.7	5.57	72.1	5.57
TokenDrift	55.04 $\pm$ 1.96	5.30 $\pm$ 0.02	27.35 $\pm$ 0.20	5.16 $\pm$ 0.02	—	—
MDLM	282.47 $\pm$ 2.52	5.45 $\pm$ 0.00	127.31 $\pm$ 1.87	5.44 $\pm$ 0.02	73.63 $\pm$ 1.50	5.42 $\pm$ 0.01
DiDi-Instruct	149.01 $\pm$ 5.00	5.44 $\pm$ 0.03	104.14 $\pm$ 5.99	5.52 $\pm$ 0.01	57.42 $\pm$ 2.54	5.41 $\pm$ 0.02

Appendix

Table 2: MDLM results on OpenWebText. We report generative perplexity (gPPL,

\downarrow

) and unigram entropy (Ent.) across sampling budgets. For each method and budget, we select the evaluated point whose entropy is closest to the data entropy (

H=5.43

). The results are averaged by 5 seeds. Gray marks results from the original papers with entropy substantially different from the data entropy.

Steps		2		4		8		16		32
Encoder	Method	gPPL	Ent.	gPPL	Ent.	gPPL	Ent.	gPPL	Ent.	gPPL	Ent.
T5	ELF	34.02 $\pm$ 6.95	4.38 $\pm$ 0.30	116.38 $\pm$ 6.61	5.48 $\pm$ 0.05	91.77 $\pm$ 2.70	5.38 $\pm$ 0.01	57.77 $\pm$ 1.42	5.36 $\pm$ 0.01	45.13 $\pm$ 1.06	5.30 $\pm$ 0.01
	ELF-PD	120.58 $\pm$ 1.91	5.49 $\pm$ 0.00	88.29 $\pm$ 1.32	5.45 $\pm$ 0.00	63.73 $\pm$ 0.45	5.39 $\pm$ 0.00	59.82 $\pm$ 0.60	5.36 $\pm$ 0.00	58.45 $\pm$ 0.50	5.35 $\pm$ 0.00
	ELF-GAN	86.06 $\pm$ 1.48	5.31 $\pm$ 0.00	52.88 $\pm$ 0.92	5.17 $\pm$ 0.00	41.08 $\pm$ 0.47	5.07 $\pm$ 0.00	36.06 $\pm$ 0.35	5.02 $\pm$ 0.00	32.86 $\pm$ 0.17	5.00 $\pm$ 0.00
	ELF-MMD	192.28 $\pm$ 1.08	5.53 $\pm$ 0.00	110.07 $\pm$ 0.69	5.49 $\pm$ 0.00	56.63 $\pm$ 0.13	5.45 $\pm$ 0.00	43.01 $\pm$ 0.20	5.41 $\pm$ 0.00	39.79 $\pm$ 0.29	5.40 $\pm$ 0.00
	+ IRD	141.03 $\pm$ 0.73	5.47 $\pm$ 0.00	78.11 $\pm$ 0.54	5.44 $\pm$ 0.00	47.33 $\pm$ 0.29	5.39 $\pm$ 0.00	38.83 $\pm$ 0.23	5.35 $\pm$ 0.00	35.75 $\pm$ 0.15	5.33 $\pm$ 0.00

Appendix

Table 3: Continuous DLMs on OpenWebText. For each number of sampling steps we report generative perplexity (gPPL,

\downarrow

) and unigram entropy (Ent.). Bold highlights the lowest perplexity at comparable entropy. The results are averaged by 5 seeds.

Method	1	2	4	8	16	32	64
ELF	$0.38_{\pm 0.05}$	$2.00_{\pm 0.33}$	$5.99_{\pm 0.20}$	$13.84_{\pm 0.84}$	$23.26_{\pm 0.47}$	$28.84_{\pm 0.91}$	$31.63_{\pm 0.94}$
ELF-PD	$\textbf{2.24}_{\pm 0.33}$	$6.31_{\pm 0.25}$	$15.95_{\pm 0.36}$	$23.52_{\pm 0.56}$	$27.16_{\pm 1.24}$	$28.32_{\pm 0.98}$	$28.48_{\pm 0.55}$
ELF-GAN	$0.73_{\pm 0.08}$	$3.38_{\pm 0.17}$	$9.73_{\pm 0.38}$	$20.27_{\pm 0.85}$	$29.11_{\pm 0.35}$	$32.77_{\pm 0.18}$	$34.36_{\pm 0.80}$
ELF-MMD	$0.71_{\pm 0.20}$	$4.00_{\pm 0.15}$	$14.21_{\pm 0.45}$	$27.49_{\pm 0.53}$	$34.16_{\pm 0.44}$	$34.28_{\pm 0.40}$	$35.24_{\pm 0.87}$
+ IRD	$1.33_{\pm 0.12}$	$\textbf{7.46}_{\pm 0.67}$	$\textbf{20.76}_{\pm 0.77}$	$\textbf{32.46}_{\pm 0.68}$	$\textbf{35.62}_{\pm 0.89}$	$\textbf{35.97}_{\pm 0.62}$	$\textbf{36.30}_{\pm 0.60}$
ELF ⋆	$0.60_{\pm 0.10}$	$4.14_{\pm 0.43}$	$12.69_{\pm 0.29}$	$25.31_{\pm 0.43}$	$30.25_{\pm 0.59}$	$32.40_{\pm 0.55}$	$33.59_{\pm 0.59}$

Appendix

Table 4: Continuous DLMs on GSM8K. For each number of sampling steps we report final-answer accuracy (

\%

). The results are averaged by 5 seeds.

Figure 12: Accuracy–TPF trade-offs on math benchmarks. DMax-Math and DMax-Math-MMD are evaluated at eleven decoding thresholds. MMD points average five training runs with different seeds. Error bars show

\pm 1

sample standard deviation in both coordinates. Lines connect points in threshold order, and stars mark the operating points in Tab. 1 .

Figure 13: Accuracy–TPF trade-offs on code benchmarks. DMax-Coder and DMax-Coder-MMD are evaluated at eleven decoding thresholds. Accuracy denotes pass@1. MMD points average five training runs with different seeds. Error bars show

\pm 1

sample standard deviation in both coordinates. Lines connect points in threshold order, and stars mark the operating points in Tab. 1 .

Table 18

Explore similar work

May 8, 2026cs.CL

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being

2{\text -}13\times

faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.

Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4

Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2

Diffusion Language Models Latent Diffusion Model Diffusion Models Latent Space

Ask PDF

Sep 30, 2026cs.LG

Distribution Matching Distillation for Continuous Diffusion Language Models

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.

Paul Le Van Kiem, Dario Shariatian, Umut Simsekli +1

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

Diffusion Language Models Diffusion Models Dataset Distillation Distribution Matching Distillation Single-Token Output Distributions Simplex

Ask PDF

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