Distribution Matching Distillation for Continuous Diffusion Language Models

2026-09-30

AI Summary

This work addresses the prohibitive inference costs of continuous diffusion language models, which typically require hundreds of network evaluations for high-quality generation. We propose a distribution distillation-based acceleration framework that unifies student probability outputs via parameterization and gradient estimators, employing reverse KL divergence as the matching objective. Specifically, we introduce two methods: Simplex-DMD, which integrates continuous relaxation with pathwise gradients, and Reinforce-DMD, which combines categorical sampling, REINFORCE, and learned density ratios. Evaluated on OpenWebText, Simplex-DMD achieves a perplexity of 45.6 in only four steps, representing a 49% reduction, while Reinforce-DMD attains a perplexity of 14.9 within 256 steps, yielding a 20% reduction. These results demonstrate significant improvements in both training efficiency and sampling quality for multi-step generation.

Abstract

Continuous diffusion language models generate all tokens in parallel, yet high-quality generation can still require hundreds of network evaluations (NFEs). We study how distributional distillation can reduce this cost by exploiting the student's probabilistic token outputs. Our unified formulation connects the student's output parameterization to the resulting gradient estimators and yields two methods with the same student architecture and reverse-KL matching objective: Simplex-DMD uses continuous token relaxations and pathwise gradients, while Reinforce-DMD uses categorical sampling and REINFORCE with a learned density ratio. We develop both methods for multi-step generation and investigate the training and sampling choices associated with each parameterization. On OpenWebText, for sequences of 1,024 tokens, Simplex-DMD achieves a generative perplexity of 45.6 at a unigram entropy of 5.44 nats in just 4 NFEs, a 49% reduction relative to the strongest evaluated diffusion baseline at matched entropy and sampling budget. Reinforce-DMD improves the frontier at larger budgets, reaching a generative perplexity of 14.9 at an entropy of 5.00 nats with 256 NFEs, a 20% reduction under the same comparison protocol.

Problem

• Continuous diffusion language models
• Distribution matching distillation
• Network evaluations
• Text generation efficiency

Innovation

• Distribution Matching Distillation
• Continuous Diffusion Language Models
• Simplex-DMD
• Reinforce-DMD
• Reverse-KL Matching