Balancing Understanding and Generation in Discrete Diffusion Models

In discrete generative modeling, two dominant paradigms demonstrate divergent capabilities: Masked Diffusion Language Models (MDLM) excel at semantic understanding and zero-shot generalization, whereas Uniform-noise Diffusion Language Models (UDLM) achieve strong few-step generation quality, yet neither attains balanced performance across both dimensions. To address this, we propose XDLM, which bridges the two paradigms via a stationary noise kernel. XDLM offers two key contributions: (1) it provides a principled theoretical unification of MDLM and UDLM, recovering each paradigm as a special case; and (2) an alleviated memory bottleneck enabled by an algebraic simplification of the posterior probabilities. Experiments demonstrate that XDLM advances the Pareto frontier between understanding capability and generation quality. Quantitatively, XDLM surpasses UDLM by 5.4 points on zero-shot text benchmarks and outperforms MDLM in few-step image generation (FID 54.1 vs. 80.8). When scaled to tune an 8B-parameter large language model, XDLM achieves 15.0 MBPP in just 32 steps, effectively doubling the baseline performance. Finally, analysis of training dynamics reveals XDLM’s superior potential for long-term scaling. Code is available at https://github.com/MzeroMiko/XDLM

💡 Research Summary

The paper tackles a fundamental tension in discrete diffusion models: masked‑diffusion language models (MDLMs) excel at semantic understanding and zero‑shot generalization, while uniform‑noise diffusion language models (UDLMs) achieve superior quality in few‑step generation. Neither paradigm delivers balanced performance across both dimensions, limiting practical applicability in tasks that demand both strong comprehension and efficient generation, such as multimodal synthesis.

To bridge this gap, the authors propose the Mixed Diffusion Language Model (XDLM). The core idea is a stationary noise kernel K that remains unchanged throughout the diffusion timeline. K is constructed as a convex combination of an all‑ones uniform matrix and absorbing matrices for special tokens (e.g., the mask token):
  K = k·N·J + ∑ μᵢ·Mᵢ,
where k controls the weight of uniform noise and μᵢ the strength of token‑specific absorption. The forward transition matrix is then defined as
  Qₜ|ₛ = αₜ|ₛ I + βₜ|ₛ K,
with α and β being scalar schedules satisfying α + β = 1. By fixing K, the diffusion process decouples noise structure from the schedule, avoiding the costly recomputation required by time‑varying approaches such as GIDD.

Theoretical analysis shows that XDLM subsumes both MDLM and UDLM as limiting cases: setting k = 0 yields a pure masking kernel (recovering MDLM), while k = 1 produces a pure uniform kernel (recovering UDLM). Lemma 3.1 formalizes this unification, and Lemma 3.3‑3.5 derive a scalar formulation for the posterior and KL‑divergence terms. By introducing token‑wise scalar primitives pₓ,ₑ (probability mass of token e) and a noise‑rate function r(e) = k/N + μ·δₑ,ₘ, the authors express the posterior as a product of simple scalar functions fₜ(x,e) = αₜ pₓ,ₑ + βₜ r(e). This eliminates the need for O(|V|²) matrix operations, reducing memory consumption to O(|V|) and enabling training on vocabularies of hundreds of thousands without GPU overflow. The continuous‑time limit further simplifies the auxiliary function hₜ, ensuring numerical stability.

Empirically, XDLM is evaluated on three fronts. First, zero‑shot language modeling: after pre‑training on the OpenWebText (OWT) corpus, the model is tested on seven external benchmarks (AG News, LAMBADA, LM1B, Penn Treebank, arXiv, PubMed, WikiText). XDLM attains an average perplexity of 54.11, outperforming UDLM by 5.4 points and matching MDLM’s 53.65 within a narrow margin, demonstrating that the added uniform noise does not degrade understanding. Second, image generation on ImageNet‑1K: in a 4‑step regime XDLM achieves an FID of 54.1, a substantial improvement over MDLM’s 80.8 and a modest gain over UDLM’s 68.3. At 16 steps it remains competitive, slightly surpassing UDLM. Third, large‑scale language modeling: the authors extend XDLM to an 8 B parameter LLaD‑A model (LLaD‑A‑XDLM) and perform continual pre‑training. In only 32 diffusion steps the model reaches an MBPP score of 15.0, more than double the baseline (6.8) and a 120 % relative improvement. Training dynamics reveal that while MDLM converges quickly then plateaus, XDLM continues to improve throughout training, indicating better scalability.

Overall, the paper delivers a principled, memory‑efficient framework that unifies the two dominant discrete diffusion paradigms. By leveraging a stationary mixing kernel and scalar reformulation, XDLM simultaneously advances zero‑shot understanding and few‑step generation, establishing a new Pareto frontier for discrete diffusion models across language and vision tasks.