Balancing Understanding and Generation in Discrete Diffusion Models

TL;DR

XDLM presents a principled unification of MDLM and UDLM through a stationary noise kernel, deriving a memory-efficient scalar posterior and a tractable training objective. The framework recovers MDLM and UDLM as limiting cases and demonstrates improved trade-offs between semantic understanding and generation quality across language and image modalities. Empirically, XDLM achieves a better Pareto frontier (e.g., zero-shot PPL improvements and lower FID in few-step image generation) and scales effectively to 8B-parameter LLMs, with strong long-horizon training dynamics. The work also provides detailed analyses of training dynamics, inference efficiency, and a sweet-spot mixing ratio k ≈ 0.1, along with public code for reproducibility.

Abstract

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

Figures (8)

• Figure 1: Left: XDLM combines the noise kernel of UDLM ($\mathbf{u}$) and MDLM ($\mathbf{m}$) to achieve a favorable trade-off between the two methods. [NORMAL] denotes normal tokens, while [MASK] represents the mask token. Right: The trade-off between understanding capability (zero-shot perplexity; lower is better) and generation capability (generation perplexity in 32 sampling steps; lower is better). The proposed XDLM with a mixing ratio of $k=0.1$ achieves the optimal balance, labeled as the 'Sweet Spot'.
• Figure 2: Language Generation Quality. Results on OWT (top) and LM1B (bottom) demonstrate that XDLM achieves a better balance across both few-step and multi-step regimes. For clarity, PPL and Generation Steps are reported in logarithmic scale.
• Figure 3: Evaluation of adapting LLaDA-8B to our XDLM formulation (LLaDA-XDLM): (a) LLaDA-XDLM consistently outperforms baselines across diverse benchmarks; (b) Improvements are particularly pronounced in code generation (MBPP), where the model substantially reduces generation failures.
• Figure 4: Training dynamics for (a) text and (b) image generation tasks. Colored regions indicate the dominant model during each phase, while transitions between colors mark points of performance crossover.
• Figure 5: Qualitative comparison of class-conditional generation on ImageNet-1K without CFG. In the absence of guidance, baseline models struggle significantly: MDLM and GIDD produce chaotic artifacts with poor structural coherence, while UDLM yields recognizable but over-smoothed images. XDLM, by effectively balancing the characteristics of MDLM and UDLM, generates the most coherent and semantically correct samples (e.g., the distinct rocket structure and pizza toppings) even without guidance.

Theorems & Definitions (18)

• Lemma 3.1: Construction of Mixing Kernel
• Definition 3.2: Scalar Primitives
• Lemma 3.3: Scalar Posterior
• Lemma 3.4: Scalar KL Divergence
• Lemma 3.5: Limiting Case
• Lemma 1: Construction of Mixing Kernel
• proof
• Definition 2.1
• Lemma 2: Scalar Posterior
• proof