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

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Diffusion Language Models Latent Diffusion Model
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
Distribution Matching Distillation for Continuous Diffusion Language Models
Sep 30, 2026Paul Le Van Kiem, Dario Shariatian, Umut Simsekli +1Diffusion Language Models Diffusion Models
Inria, PSL Research University · Cohere · CMAP, Ecole Polytechnique
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Dataset Distillation
NVIDIA

May 8, 2026, cs.CLJ/K move · Enter open · S save

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

Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin+3

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

Diffusion Language Models Latent Diffusion Model Diffusion Models Latent Space

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Distribution Matching Distillation for Continuous Diffusion Language Models

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling