Representation-Space MMD for Diffusion Language Models

Representation-Space MMD for Diffusion Language Models. By Ilya Drobyshevskiy and nine others, from Yandex Research, HSE University, and four other institutions. Published on the fifth of October. Its code repository is public and real: fourteen source files.

Post-training a diffusion language model needs some measure of how far its generations sit from reference text. The method here is built to get that measure cheaply: without full sampling trajectories, and without auxiliary models trained jointly alongside the one being improved. Those are the two costs it sets out to avoid.

The objective is Maximum Mean Discrepancy, or MMD, between the generated and reference distributions. It is a comparison of two distributions as wholes, rather than of individual samples. And it is measured not on raw text but in the feature space of a frozen, pretrained diffusion language model, which stays fixed throughout.

To estimate MMD, contextual features are kept at individual token positions. A single pass of the extractor therefore yields many observations per sequence rather than one, which is what makes the estimate affordable. The objective is optimized with policy gradients for discrete models, and by differentiating directly through generated latents for continuous ones. Either way, the loss comes straight from those features.

On OpenWebText, generative perplexity is lower at comparable entropy. On GSM8K, the trade-off between accuracy and computation improves. And on sixteen billion parameter DMax-LLaDA2.0 models with hybrid masked-uniform diffusion, decoding parallelism goes up with similar or higher accuracy on math and code benchmarks.

Compactly: a frozen model's features become the measuring stick, every token position counts as a sample, and that is enough to post-train both discrete and continuous diffusion models efficiently, at scales up to sixteen billion parameters.