Faster Inference for Conditional Masked Diffusion Language Models by Knowledge Distillation of Guidance and Trajectory

Masked diffusion language models (MDLMs) offer a promising paradigm for natural language generation, enabling fully parallel, non-autoregressive decoding through iterative unmasking. However, existing MDLMs incur substantial inference costs due to the large number of neural network function evaluations required, thereby limiting their practical applicability in conditional sequence-to-sequence generation settings. In this paper, we propose a two-stage distillation framework for conditional MDLMs that aims to distill knowledge of both (i) classifier-free guidance and (ii) the denoising trajectory from a teacher MDLM to a student MDLM. As a result, the student model can, during inference, (i) replace the two forward passes required by classifier-free guided outputs with a single pass, and (ii) significantly reduce the number of denoising steps. Evaluations across diverse conditional generation tasks on MDLMs up to 8 billion parameters demonstrate improved generation quality with significantly fewer function evaluations.