Post-training diffusion language models with MMD in frozen feature space: MDLM-MMD and ELF-MMD lower generative perplexity at matched entropy, and 16B DMax raises decoding parallelism on math and code benchmarks

It proposes using contextual token features of a frozen pretrained diffusion language model as the MMD comparison space and retaining per-position features, so a single extractor pass yields multiple observations per sequence for distribution matching. Earlier MMD-style objectives in visual generation relied on pretrained feature spaces, while language-side feature matching such as EBFT represents a rollout by a single feature vector; here distributions of token-level features are compared with an RBF kernel, with an estimator that excludes same-sequence token pairs to preserve unbiasedness. The paper gives the sequence-level kernel average in Eq. (4) and the batch estimator in Eq. (5), noting that same-sequence token pairs can be dependent and that retaining them generally introduces bias; ablations show token-level RBF outperforming the linear kernel, sequence-level RBF, attraction-only, and feature-regression losses.

For discrete diffusion language models it optimizes an MMD reward with REINFORCE and a leave-one-out baseline formed from independent batches, instantiated as MDLM-MMD and DMax-MMD. Unlike distillation methods that rely on teacher denoising targets or jointly trained auxiliary models such as IDLM and DiDi-Instruct, the reward is computed directly from MMD in frozen features, without full sampling trajectories or auxiliary models. On OpenWebText, initialized from MDLM with group size g=4, it reports lower generative perplexity than DiDi-Instruct, IDLM, and IDLM-REINFORCE at matched entropy; on TinyGSM/GSM8K, sweeping confidence thresholds, it reports higher accuracy at moderate and high decoding budgets and fewer steps to reach a given accuracy.

For continuous diffusion language models it differentiates directly through generated latents, instantiated as ELF-MMD, and extends to multi-step sampling through self-conditioning refinement, with optional iterative refinement distillation (IRD). Compared with progressive distillation and flow-map approaches that need teacher trajectories, the MMD stage uses no paired teacher trajectories; IRD exists separately as an optional trajectory-supervised stage. It reports improved perplexity-entropy trade-offs across sampling budgets for both T5 and GPT-2 latent spaces; on GSM8K it reports accuracy gains at 4 and 8 steps, and the highest accuracy of that group at 64 steps for ELF-MMD+IRD; appendices report gains for ELF⋆-MMD and FMLM-MMD over their respective baselines at every evaluated budget.

It tests scaling on 16B DMax-Math and DMax-Coder: MMD trains both masked inputs and the model's own token predictions, and after 400 training steps decoding parallelism (TPF) rises while accuracy is maintained or improved. It moves token-level MMD post-training from small and mid-scale experiments to released 16B hybrid masked-uniform diffusion checkpoints, reporting both accuracy and tokens per forward as a parallelism metric. It reports 400 training steps on eight NVIDIA H100 GPUs, taking minutes for DMax-Math and minutes for DMax-Coder; across four math benchmarks TPF increases relative to the reported DMax operating points with similar or higher accuracy; on HumanEval-Instruct and MBPP-Instruct accuracy improves by percentage points while TPF also rises.

The result targets the post-training stage of diffusion language models: both discrete models (MDLM, DMax) and continuous models (ELF) can use it, and training needs only a frozen pretrained diffusion language model as feature extractor plus reference samples, without full sampling trajectories or jointly trained auxiliary models. Settings include unconditional text generation (OpenWebText) and conditional mathematical reasoning (trained on TinyGSM, evaluated on GSM8K), as well as 16B-scale math and code benchmarks. At inference the method keeps each model's original sampling procedure, and continuous models can add self-conditioning iterative refinement and optional IRD.

Performance depends on the RBF bandwidth and the chosen representation space, and the paper lists representation design as a direction for further exploration; it uses features from a single layer on clean inputs, while combining layers or noise levels is not yet validated. The discrete MMD reward relies on policy gradients and group size, and the paper reports that increasing group size from 1 to 4 gives the main gain with limited further benefit. Iterative refinement and IRD are optional stages for continuous models, and their benefit versus training cost still needs per-task balancing. In addition, several table values are incomplete in the supplied text, so some specific numbers cannot be checked item by item here and readers should consult the original tables.