PaperlayerView exact revision on arXiv

Source-linked AI summary

Representation-Space MMD for Diffusion Language Models

Ilya Drobyshevskiy, Ilia Sudakov, Maksim Semenov, Denis Kuznedelev, Maksim Ignatov, Pavel Temirchev, Nikita Balagansky, Viacheslav Meshchaninov, Nikita Gushchin, Dmitry Baranchuk

arXiv:2610.06648v1cs.CLcs.LG

TL;DR

Diffusion language models’ conditional predictions do not necessarily yield high-quality samples in a few steps, and existing distillation methods require teacher-derived targets or jointly trained auxiliaries. This paper post-trains discrete and continuous DLMs by matching generated and reference contextual features from a frozen pretrained DLM. It reports improved generation and accuracy-computation trade-offs across text, math, and code evaluations.

  • Problem

    Conditional denoising estimates do not necessarily yield high-quality samples in a few steps, and existing distillation approaches require teacher-derived targets or jointly trained auxiliary models.

    Source: §1 · p. 1

  • Method

    The method matches generated and reference distributions of individual-position contextual features from a frozen pretrained DLM, using policy gradients for discrete models and direct differentiation through generated latents for continuous models.

    Source: §1 · pp. 1–2

  • Results

    The method improves generation in discrete and continuous DLMs, with lower generative perplexity at matched entropy on OpenWebText and better accuracy-computation trade-offs on GSM8K and 16B math and code benchmarks.

    Source: §1 · p. 2, §4.1 · p. 6, §4.4 · p. 9

  • Takeaways & Limitations

    Feature-space MMD provides a simple, efficient post-training procedure that requires neither full sampling trajectories nor jointly trained auxiliary models.

    Source: §5 · p. 9

    Matching token-feature distributions is not established to uniquely identify the underlying sequence distribution.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

Diffusion models’ conditional predictions do not necessarily produce high-quality samples in few steps, while existing distillation methods require teacher-derived targets or jointly trained auxiliaries. This work post-trains discrete and continuous DLMs by matching generated and reference contextual features, with reported gains across text generation and large-scale math and code benchmarks.

1 INTRODUCTION

  • Conditional denoising estimates guide iterative sampling but do not necessarily yield high-quality samples in a few steps, while existing distillation approaches require teacher-derived targets or jointly trained auxiliaries.The paper asks whether generated and reference samples can instead provide a direct distributional training signal.Source: §1 · p. 1
  • The method matches individual-position contextual features from generated and reference sequences using a frozen pretrained DLM, yielding multiple MMD observations per sequence in one extractor pass.The feature extractor is frozen; features are retained at token positions rather than pooled into one sequence representation.Source: §1 · pp. 1–2
  • For discrete DLMs, sampled sequences from a single denoiser pass are optimized with REINFORCE and a leave-one-out baseline.The method is instantiated as MDLM-MMD for masked-token prediction and DMax-MMD for hybrid masked–uniform diffusion.Source: §1 · p. 2
  • For continuous DLMs, MMD gradients pass through generated latents and a frozen ELF feature extractor; self-conditioning then extends the generator to multi-step sampling.The generator is trained to map Gaussian noise to clean latent sequences in one step.Source: §1 · p. 2
  • At 16B scale, DMax-MMD improves the accuracy-efficiency trade-off on established math and code benchmarks.On OpenWebText and TinyGSM, both realizations outperform evaluated distribution-matching and flow-map baselines in most setups.Source: §1 · p. 2

2 PRELIMINARIES

Discrete DLMs corrupt tokens and predict clean-token distributions, while continuous DLMs predict clean latents from noisy interpolations. MMD compares distributions through kernel similarities, with the selected feature space shaping what differences the comparison captures.

2 PRELIMINARIES

  • Discrete DLMs independently retain or replace tokens before predicting factorized clean-token distributions; masked diffusion reveals masked tokens, while uniform diffusion updates all positions and can revise earlier predictions.Masked diffusion uses confidence-threshold decoding to determine how many tokens to reveal per step.Source: §2 · p. 2
  • Continuous DLMs learn clean latents from interpolations of data latents and Gaussian noise, and ELF uses prior clean-latent predictions as self-conditioning during iterative sampling.The final latent sequence is decoded into tokens.Source: §2 · p. 2
  • MMD measures distributional difference through kernel similarities; a linear kernel matches feature means, while a characteristic RBF kernel can identify equal distributions in the space where it is evaluated.For generated features, the RBF within-distribution term penalizes average similarity among generated features, while the cross term rewards similarity to reference features.Source: §2 · p. 2

3 METHOD

The method matches contextual token-feature distributions with MMD, using estimators that account for dependence among positions in a sequence. Discrete models optimize sampled outputs with policy gradients, while continuous models backpropagate through generated latents and can refine predictions through self-conditioning.

3.1 REPRESENTATION-SPACE MMD TRAINING

  • The method uses a frozen pretrained DLM’s contextual features and retains individual scored positions, providing multiple feature observations per sequence in one pass.Features are extracted from clean samples because corruption at intermediate diffusion timesteps did not improve performance in the reported experiments.Source: §3.1 · p. 3
  • Within-distribution comparisons exclude pairs from the same sequence to preserve unbiasedness despite dependent token features; the real–generated term attracts toward reference features, while the generated–generated term discourages concentration.The estimator uses one reference and at least two independently sampled generated sequences per condition; with one generated sequence, attraction-only training performed worse for discrete models and rapidly collapsed for continuous models.Source: §3.1 · p. 4

3.2 DISCRETE DIFFUSION LANGUAGE MODELS

  • For discrete DLMs, the denoiser samples selected positions while holding other input tokens fixed; REINFORCE optimizes a negative MMD reward using groups of batches and a leave-one-out baseline.Sampled and reference sequences share context, and their contextual features at selected positions enter the MMD loss.Source: §3.2 · p. 5
  • MDLM-MMD applies the procedure to masked positions, while DMax-MMD uses MMD for masked prediction and refinement of the model’s own token predictions; original sampling procedures remain unchanged.Source: §3.2 · p. 5

3.3 CONTINUOUS DIFFUSION LANGUAGE MODELS

  • ELF-MMD initializes a one-step latent generator from pretrained ELF and differentiates MMD through generated latents and both inputs to the frozen ELF feature extractor.The generator starts from Gaussian noise with null self-conditioning; the extractor processes clean sequences at t = 1.Source: §3.3 · p. 5
  • Continuous predictions can be iteratively refined through self-conditioning, with preliminary no-gradient bootstrap steps followed by a differentiable MMD step; optional IRD distills a multi-step teacher prediction into one step.The method is also applicable to continuous DLM variants without self-conditioning, which could use alternative multi-step sampling schemes.Source: §3.3 · p. 6

4 EXPERIMENTS

Experiments evaluate representation-space MMD on OpenWebText and GSM8K, then test scaling to 16B DMax models and compare alternative MMD losses. Results report improved generation and accuracy–computation trade-offs across these settings.

4.1 UNCONDITIONAL GENERATION

  • On OpenWebText, MDLM-MMD yields ~17–21% lower generative perplexity than IDLM at reference entropy; ELF-MMD improves generation across most sampling budgets.At four steps, IRD reduces generative perplexity relative to ELF-MMD by ~30 points with T5 and ~10 points with GPT-2, while bringing entropy closer to the reference for both.Source: §4.1 · pp. 6–7, Figure 2 · p. 7

4.2 CONDITIONAL GENERATION

  • On GSM8K, MDLM-MMD reaches ~54% accuracy at ~49 steps, while ELF-MMD improves over ELF at every evaluated budget and ELF-MMD+IRD reaches 36.3% at 64 steps.At 4 and 8 steps, ELF-MMD+IRD reaches 20.8% and 32.5%, respectively, compared with ELF-PD’s 15.9% and 23.5%.Source: §4.2 · pp. 7–8, Figure 3 · p. 8

4.3 SCALING TO LARGE MODELS

  • On math benchmarks, DMax-Math-MMD increases tokens generated per forward by 10.3–16.5% with similar or higher accuracy; on code, DMax-Coder-MMD raises accuracy by 2.4 and 3.8 points while increasing TPF.The comparison uses the original DMax operating points; TPF measures decoding parallelism.Source: §4.3 · p. 8, Table 1 · p. 9

4.4 MMD LOSS ABLATION

  • On GSM8K, token-level RBF gives the strongest accuracy–computation trade-off, with larger gains for discrete than continuous models.The ablation compares linear-kernel mean matching, sequence-level RBF on mean-pooled embeddings, attraction-only loss, and feature regression; ELF results omit IRD.Source: §4.4 · pp. 8–9, Figure 4 · p. 9

5 DISCUSSION

The discussion reports that frozen pretrained DLM features support efficient MMD post-training, while representation design remains an open choice. Related work spans few-step distillation, distribution matching, and feature matching for diffusion and language models.

5 DISCUSSION

  • MMD on pretrained DLM features improves generation in discrete and continuous DLMs without full sampling trajectories or jointly trained auxiliary models.The authors identify representation design as an open direction and suggest combining MMD with other objectives for discrete DLMs.Source: §5 · p. 9

A RELATED WORK

  • Few-step continuous-generation methods distill denoising transitions or flow maps, whereas this work matches latent samples against references in a fixed feature space.The method also uses self-conditioning for refinement and may distill those steps afterward.Source: A RELATED WORK · p. 15
  • Discrete few-step methods use teacher-step matching, consistency regularization, or auxiliary models; this work instead computes frozen-feature MMD rewards for policy gradients.Examples include SDTT and CDLM, and auxiliary-model approaches such as IDLM, D-MMD, and DiDi-Instruct.Source: A RELATED WORK · p. 15
  • Prior representation-space approaches include adversarial MMD-GAN, frozen-feature GFMN [dos Santos et al., 2019], and patch-level diffusion-feature matching in SwD.Related objectives include feature mean-and-covariance comparison and attraction–repulsion field regression.Source: A RELATED WORK · p. 15
  • Language feature matching has used position-wise moments or rollout-level features; this work compares token-feature distributions with an RBF kernel for parallel discrete predictions or continuous latents.EBFT [Jelassi et al., 2026] represents each rollout with one feature vector, while SeqGFMN matches position-wise means and variances.Source: A RELATED WORK · p. 15

B CONTINUOUS DLMS TRAINING DETAILS

The continuous-model appendix describes optional iterative refinement distillation after MMD training and a separate decoder fine-tuning step. IRD trains a student to reproduce the endpoint of a fixed-noise teacher refinement trajectory.

B.1 ITERATIVE REFINEMENT DISTILLATION

  • IRD distills multiple teacher refinement iterations into one student prediction to reduce the number of refinement steps.It adapts a simple knowledge-distillation approach and is optional after MMD training; the MMD stage itself uses no paired teacher trajectories.Source: §B.1 · p. 16
  • The teacher is the frozen ELF-MMD generator, which produces a fixed K-step trajectory using the same Gaussian noise at every iteration.The student uses that same noise and no self-conditioning to predict the teacher’s final output.Source: §B.1 · p. 16
  • The student learns to reproduce the output after K teacher refinement steps, while tested consistency-distillation alternatives showed no noticeable improvement.After distillation, the generator can again use iterative self-conditioning at inference.Source: §B.1 · p. 16

B.2 DECODER FINE-TUNING

  • Decoder fine-tuning reconstructs original tokens from perturbed data latents with cross-entropy, increasing global batch size by about 20% for added samples.Source: §B.2 · p. 16

C ABLATION STUDY

The appendix evaluates design choices affecting performance and omits IRD from the ELF-MMD experiments.

C ABLATION STUDY

C.1 MMD LOSS ABLATION ON OWT

The ablations favor token-level RBF features in several settings, show diffusion-based features leading representation comparisons, and find that most policy-gradient group-size gains arrive at G=2. Scaling results show MMD benefits persist from ELF-B to ELF-M.

C.1 MMD LOSS ABLATION ON OWT

  • At 8, 16, and 32 steps, token-level RBF gives discrete models the lowest generative perplexity at matched entropy; continuous-model gains vary by latent space.With T5 latents it beats linear and sequence-level RBF trade-offs; with GPT-2 latents its improvement over linear is modest.Source: §C.1 · p. 16

C.2 REPRESENTATION SPACES

  • Figure 6 compares GSM8K accuracy against sampling steps for MDLM and ELF, with setups tuned separately for each representation type.Source: §C.2 · p. 16, Figure 6 · p. 17

C.2.1 INPUT REPRESENTATIONS

  • Input-representation comparisons use token embeddings for discrete MDLM and identity features on continuous latents, which already encode context.Discrete token embeddings depend only on token identity, unlike contextual hidden states.Source: §C.2.1 · p. 17

C.2.2 EXTERNAL ENCODER

  • On GSM8K, diffusion-based features achieve the highest accuracy for both discrete and continuous DLMs, outperforming the compared representation alternatives.Source: §C.2.2 · p. 18

C.3 POLICY-GRADIENT GROUP SIZE

  • Increasing policy-gradient group size from G=1 to G=2 yields the largest improvement; G=2, 4, and 8 have similar trade-offs, so G=4 is the shared default.The comparison spans OWT perplexity–entropy trade-offs and GSM8K accuracy versus average sampling steps.Source: §C.3 · p. 18, Figure 7 · p. 18

C.4 MODEL-SIZE SCALING

  • ELF-M-MMD outperforms ELF-M at every evaluated sampling budget, while ELF-B-MMD exceeds ELF-M at 4–16 steps and is similar at 32 steps.Increasing model size improves accuracy before and after MMD across evaluated budgets.Source: §C.4 · p. 18

C.5 BOOTSTRAPPING

The bootstrap ablation finds that one step can improve generation trade-offs, while additional steps bring limited or inconsistent gains at added training cost.

C.5 BOOTSTRAPPING

  • Figure 8 compares GSM8K accuracy with average sampling steps across policy-gradient group sizes.Source: Figure 8 · p. 18
  • Figure 9 compares ELF-B and ELF-M on GSM8K before and after MMD post-training.Source: Figure 9 · p. 18
  • A single bootstrap step improves the trade-off with T5 and GSM8K accuracy, while GPT-2 results are mixed.Additional steps yield limited or inconsistent improvements, and each adds training cost.Source: §C.5 · pp. 18–19, Figure 10 · p. 19

D ADDITIONAL DMAX RESULTS

These additional results compare MMD-tuned models with baselines across DMax, masked-DLM, and continuous-DLM settings. The results also report higher accuracy at every evaluated sampling budget for ELF⋆-MMD and FMLM+-MMD on GSM8K.

E.1 MASKED DLM BASELINES

  • Masked baselines include reproduced IDLM, IDLM-REINFORCE with its teacher and auxiliary-denoiser training retained, and DiDi-Instruct; D-MMD and TokenDrift use reported OWT results.DiDi-Instruct is evaluated without additional training on OWT and adapted for conditional TinyGSM generation.Source: §E.1 · p. 19

E.2 CONTINUOUS DLM BASELINES

  • Continuous baselines reuse or adapt ELF-family code and checkpoints; ELF-GAN adversarially post-trains pretrained ELF-B checkpoints using a discriminator on frozen-ELF token features.ELF-GAN uses five discriminator updates per generator update.Source: §E.2 · p. 20

F.1 ELF⋆-MMD

  • ELF⋆-MMD achieves higher GSM8K accuracy than ELF⋆ at every evaluated sampling budget.The model retains self-conditioning during inference and uses the same iterative-refinement procedure as ELF-MMD.Source: §F.1 · p. 20

F.2 FMLM+-MMD

  • FMLM+-MMD achieves higher GSM8K accuracy than FMLM+ at every evaluated sampling budget.It is evaluated using the FMLM+ protocol, with the reported step count denoting refinement steps.Source: §F.2 · p. 20

G HYPERPARAMETERS

D ADDITIONAL DMAX RESULTS

  • The DMax plots compare accuracy–TPF trade-offs across decoding thresholds; MMD results average five runs, with error bars showing ±1 sample standard deviation.Stars mark selected operating points, which differ slightly from Tab. 1 because the plots average five runs.Source: D ADDITIONAL DMAX RESULTS · p. 19, Figure 12 · p. 23, Figure 13 · p. 23

H.1 UNCONDITIONAL GENERATION ON OPENWEBTEXT

The appendix provides unconditional ELF-MMD samples on OpenWebText at each sampling step, alongside generative perplexity and entropy. The supplied examples include passages with repetition and malformed text.

H.1 UNCONDITIONAL GENERATION ON OPENWEBTEXT

  • Unconditional ELF-MMD samples on OpenWebText are provided for each sampling step, together with generative perplexity and entropy.Source: §H.1 · p. 26
  • One example repeats terms and contains malformed phrases within a passage about government surveillance and alleged misconduct.Source: §H.1 · p. 26
  • Other examples contain disjointed claims and quotations, including passages mentioning political fraud and government actions.Source: §H.1 · p. 26

H.2 CONDITIONAL GENERATION ON TINYGSM

The TinyGSM examples show a correct tree-height answer at four sampling steps and errors at two and three steps. The appendix also provides ELF-MMD OpenWebText samples and T5-encoder output examples.

H.2 CONDITIONAL GENERATION ON TINYGSM

  • In the tree-height example, the four-step sample is marked correct with answer 34, while the two- and three-step samples are marked incorrect with errors.Source: §H.2 · p. 27

I.1 UNCONDITIONAL GENERATION ON OPENWEBTEXT

  • ELF-MMD unconditional OpenWebText samples are provided at each sampling step, together with generative perplexity and entropy.Source: §I.1 · p. 28

I.1.1 T5 ENCODER

  • The T5-encoder examples contain passages with inconsistent claims and malformed or repetitive phrases.Source: §I.1.1 · p. 28

I.1.2 GPT-2 ENCODER

The GPT-2 encoder subsection presents generated-text examples, while the TinyGSM subsection shows outputs for a tutoring-earnings question at different sampling-step counts.

I.1.2 GPT-2 ENCODER

  • The GPT-2 examples mix product listings and silver or gold claims with shifting political and war narratives.Source: §I.1.2 · p. 29

I.2 CONDITIONAL GENERATION ON TINYGSM

  • At 1 and 2 sampling steps the TinyGSM outputs report errors and accuracy 0; at 4 steps the output is 130 with accuracy 1.Source: §I.2 · p. 30
  • The TinyGSM question asks for earnings at $10 per hour over 5 hours and 8 hours, with code combining the hours and multiplying by the rate.Source: §I.2 · p. 30
Loading 2610.06648v1…