How can MMD in representation space improve diffusion language model post-training?

MMD compares features, not token probabilities

The training signal here is a kernel-based distance between two sets of features: those extracted from generated sequences and those extracted from reference sequences. Both sets come from the same frozen pretrained diffusion language model, so the comparison space is a learned representation rather than raw token space. The paper's stated motivation is that simple distances in high-dimensional spaces need not capture meaningful differences between outputs, while learned representations offer a promising comparison space for discrete sequences and continuous outputs alike.

The kernel used is an RBF kernel, exp(-||a-b||^2/(2*sigma^2)), which the paper notes is characteristic, so zero MMD identifies equal distributions in the space where the kernel is evaluated. When Q is the generated distribution, its within-distribution term penalizes average RBF similarity between generated features, while the cross term rewards similarity to reference features. Two choices therefore shape the objective: the RBF bandwidth and the representation space used for matching. The discussion names both as the main dependencies of performance.

One extractor pass yields many token-level observations

To estimate MMD, the method retains contextual features at individual token positions, which produces multiple observations per sequence from a single extractor pass. This is the efficiency argument of the paper: the loss is computed directly from these features, so post-training does not require full sampling trajectories or jointly trained auxiliary models. The paper contrasts this with diffusion distillation approaches that either train a student to reproduce several teacher denoising updates in fewer steps or obtain distribution-matching gradients from an auxiliary model trained on the generator's evolving outputs.

The two model families are optimized differently. For discrete models, sequences are sampled from the token distributions produced by a single denoiser pass, and the resulting MMD reward is optimized with REINFORCE using a leave-one-out baseline; the paper instantiates this as MDLM-MMD for masked-token prediction and as DMax-MMD for hybrid masked-uniform diffusion, where MMD trains both masked prediction and refinement of the model's own token predictions. For continuous DLMs, a generator is initialized from a pretrained Embedded Language Flows model and trained to map Gaussian noise to clean latent sequences in a single step, with MMD gradients propagating directly through the frozen feature extractor and the generated latents. At inference, self-conditioning iteratively refines the single-step predictions into multi-step sampling.

Lower generative perplexity at matched entropy

For unconditional generation, models are trained on OpenWebText with sequences packed to length 1024, and 1,000 generated samples are evaluated using generative perplexity under pretrained GPT-2 Large plus average unigram entropy. The paper uses the data entropy of about 5.43 as a reference rather than treating higher entropy as uniformly better, and it evaluates each method at 8, 16, and 32 sampling steps while varying sampling temperature to trace the perplexity-entropy trade-off.

MDLM-MMD achieves lower generative perplexity than all baselines at matched entropy. Interpolating the curves at the reference-data entropy gives roughly 17 to 21 percent lower generative perplexity than IDLM across the three sampling budgets. The comparison set includes DiDi-Instruct, IDLM, and IDLM-REINFORCE, which the paper says helps separate the effect of the training objective from the effect of the optimization procedure.

For continuous models, ELF-MMD improves generation across most sampling budgets for both the T5-small and GPT-2 Large latent spaces. At 8 steps it achieves lower generative perplexity and entropy closer to the reference value than ELF-PD; at 32 steps it outperforms ELF on both metrics. Adding iterative refinement distillation further improves few-step results: at 4 steps it reduces generative perplexity relative to ELF-MMD by about 30 and 10 points with the T5 and GPT-2 encoders respectively, while bringing entropy closer to the reference value for both.

Accuracy gains concentrate at small sampling budgets

Conditional generation is trained on TinyGSM with sequences packed to a maximum length of 512 and evaluated by final-answer accuracy on the GSM8K test set. For masked models, the confidence threshold is swept and the resulting average number of steps per sequence is reported, because confidence decoding determines the number of model evaluations adaptively. MDLM-MMD reaches higher accuracy at moderate and high decoding budgets, about 54 percent accuracy at about 49 steps.

For continuous models, ELF-MMD improves accuracy over ELF at every evaluated sampling budget and over ELF-PD at most of them, and IRD improves accuracy at every evaluated budget. At 4 and 8 steps, accuracy rises from 14.2 percent to 20.8 percent and from 27.5 percent to 32.5 percent respectively, compared with 15.9 percent and 23.5 percent for ELF-PD. At 64 steps, ELF-MMD+IRD reaches its highest accuracy of 36.3 percent, compared with 35.2 percent for ELF-MMD and 31.6 percent for ELF.

Pass@k results follow the same pattern. For MDLM-based models, MDLM-MMD outperforms baselines at 4 to 8 sampling steps for k up to 8, and at 16 to 32 steps it achieves higher pass@k for most evaluated values of k with smaller differences at larger sample budgets. For continuous DLMs, ELF-MMD also improves pass@k at smaller sample budgets while the gaps narrow as k increases.

16B DMax gains parallelism without losing accuracy

The scaling test starts from released 16B DMax-Math and DMax-Coder checkpoints. DMax is a hybrid masked-uniform diffusion model tuned from LLaDa2.0-Mini that learns to recover target tokens from both masked inputs and its own token predictions; the authors follow the original training implementation and use MMD to train on both, yielding DMax-Math-MMD and DMax-Coder-MMD. Training runs 400 steps, taking about 13 minutes for DMax-Math and about 19 minutes for DMax-Coder on eight NVIDIA H100 GPUs, roughly 1.7 and 2.5 GPU-hours per run, covering training only and excluding initialization and evaluation.

Evaluation uses the original dInfer implementation with Soft Parallel Decoding, a block size of 32, and a maximum generation length of 2048 tokens, reporting accuracy and tokens generated per forward as a measure of decoding parallelism. The selected thresholds are 0.85 for DMax-Math-MMD and 0.9 for DMax-Coder-MMD. Across the four math benchmarks, DMax-Math-MMD increases tokens per forward by 10.3 to 16.5 percent over the reported DMax-Math operating points while achieving similar or higher accuracy. On code benchmarks, DMax-Coder-MMD improves accuracy by 2.4 percentage points on HumanEval-Instruct and 3.8 points on MBPP-Instruct while also increasing tokens per forward on both.

These numbers are operating points selected from accuracy-TPF curves evaluated at eleven decoding thresholds, so the threshold is part of the result rather than a fixed setting. The MMD points average five training runs with different seeds, and error bars show plus or minus one sample standard deviation in both coordinates.

The kernel and representation choices carry the result

The ablation compares alternative loss formulations on GSM8K, with OpenWebText results in the appendix. It starts from linear-kernel MMD, computed as mean squared error between generated and reference feature means, which isolates how much the RBF kernel and the full discrepancy contribute beyond simple mean matching. The discussion states that performance depends on choices such as the RBF bandwidth and the representation space used for distribution matching, which makes representation design a natural direction for further exploration.

The paper also notes that it uses features from a single layer on clean inputs, and suggests that combining features across layers or noise levels, or using an ensemble of extractors, could provide complementary information. Alongside representation design, future work could explore combining MMD with other training objectives for discrete DLMs, as suggested by the gains from applying IRD after ELF-MMD post-training in the continuous setting.

What the reported gains do and do not cover

The evidence supports a specific claim: MMD on pretrained DLM features improves generation in both discrete and continuous DLMs, and does so with a training procedure that requires neither full sampling trajectories nor jointly trained auxiliary models. The reported gains are measured on OpenWebText generative perplexity at matched entropy, GSM8K final-answer accuracy and pass@k at given sampling budgets, and accuracy and tokens per forward on math and code benchmarks for 16B DMax models.

Several boundaries qualify those gains. The 16B results come from 400 training steps on released DMax checkpoints and are reported at selected decoding thresholds, so they describe operating points on an accuracy-TPF curve rather than a threshold-free improvement. The unconditional and conditional results come from specific training sets, sequence lengths, and evaluation setups, and the paper's own discussion attributes performance to the RBF bandwidth and the representation space, meaning the reported numbers are tied to those choices. The paper also uses features from a single layer on clean inputs, so the effect of combining layers, noise levels, or extractors is untested here.

The practical priority this paper establishes is a trade-off rather than a single score: at small sampling budgets, representation-space MMD buys accuracy and lower generative perplexity at matched entropy, and on 16B DMax it buys decoding parallelism at similar or higher accuracy. The decision-relevant variable is the decoding threshold, since the reported math and code gains are read off accuracy-TPF curves at thresholds of 0.85 and 0.9. Readers comparing post-training methods should therefore ask which sampling budget and which threshold a claimed gain is measured at, because the same method can look different at 4 steps, at 64 steps, and at a different point on the parallelism curve.

Sources used in this answer