Multi-Mask Diffusion Language Models for Few-Step Generation

2026-07-21

AI Summary

Existing masked diffusion language models suffer from insufficient terminal entropy in few-step generation due to the forward process collapsing into a single fully masked state, hindering high-quality rapid text synthesis. To address this limitation, this work proposes MultiMDM, a multi-mask diffusion model that maps each token to multiple masks and combines them into a structured mask set, thereby endowing the reverse process with draft-and-refine capabilities. Key innovations include preserving mask structure to enable efficient few-step generation, formulating a closed-form ELBO optimization objective, introducing a purely discrete-state consistency distillation method, and incorporating shared Gumbel coupling to reduce path entropy. Experiments demonstrate that MultiMDM serves as an effective foundation model for high-quality few-step text generation, supporting both continuous training from pretrained MDMs and efficient distillation.

Abstract

Masked diffusion models (MDMs) are a promising family of language generators, but achieving high-quality few-step generation remains challenging. In MDMs, all forward trajectories collapse to a single fully masked state, leaving no terminal entropy for consistency-style few-step generation. While recent few-step alternatives based on uniform-state diffusion avoid this degeneracy, it becomes harder to distinguish clean tokens from noise than MDMs, which usually harms modeling quality and training efficiency. In this work, we propose a multi-mask diffusion model (MultiMDM) that preserves the masking structure towards few-step generation. In the forward process, each clean token is first pushed towards a designated mask and then gradually mixes over the mask set. As a result, the backward process has a drafting capability by predicting a designated mask before refining to a clean token. We derive a closed-form ELBO training objective for MultiMDM that supports continual training from pretrained MDMs. In addition, we formulate a purely discrete-state consistency distillation scheme, with a shared-Gumbel coupling to reduce pathwise entropy. Experiments on pretraining and distillation show that MultiMDM provides an effective foundation for principled few-step generation.

Problem

• masked diffusion models
• few-step generation
• terminal entropy
• consistency distillation
• discrete-state diffusion

Innovation

• multi-mask diffusion
• few-step generation
• consistency distillation
• discrete-state diffusion
• ELBO training