Representation-Space MMD for Fast Discrete Diffusion Language Models

We introduce MDLM-MMD, a kernel-based method for adapting pretrained masked diffusion language models to few-step generation. The student learns to complete masked sequences by minimizing Maximum Mean Discrepancy (MMD) between the token representations of sampled completions and their clean references. We use a frozen pretrained language model to extract features and REINFORCE with a leave-one-out baseline to optimize the resulting MMD-based reward directly on discrete token samples. Unlike many prior methods, this approach requires neither teacher-generated trajectories nor a jointly trained discriminator or auxiliary denoiser. On OpenWebText, MDLM-MMD improves the generative perplexity–entropy trade-off over IDLM-based baselines at 8, 16, and 32 sampling steps. Preliminary GSM8K experiments also show improvements in the accuracy–NFE trade-off. Applied to SDAR-1.7B and SDAR-4B, MDLM-MMD generates more tokens per forward (TPF) at comparable or higher scores on math and code benchmarks than the supervised fine-tuned baselines.