PaperlayerView exact revision on arXiv

Source-linked AI summary

Balancing Understanding and Generation in Discrete Diffusion Models

Yue Liu, Yuzhong Zhao, Zheyong Xie, Qixiang Ye, Jianbin Jiao, Yao Hu, Shaosheng Cao, Yunfan Liu

arXiv:2602.01362v1cs.CL

TL;DR

MDLMs and UDLMs excel at different aspects of discrete diffusion, motivating a model that balances understanding and few-step generation. XDLM uses a stationary noise kernel to unify them, simplifies posterior computation, and improves the trade-off across language and image tasks. Its results include stronger zero-shot performance than UDLM, better few-step image generation than MDLM, and doubled MBPP performance over the LLaDA baseline.

  • Problem

    MDLMs favor semantic understanding and zero-shot generalization, whereas UDLMs favor few-step generation, and neither balances both capabilities.

  • Method

    XDLM uses a stationary noise kernel combining uniform and masked noise, algebraically simplifying posterior probabilities while recovering MDLM and UDLM as limiting cases.

  • Results

    XDLM advances the understanding–generation Pareto frontier across modalities, surpassing UDLM in zero-shot benchmarks and outperforming MDLM in few-step image generation while improving 8B-model code generation.

  • Takeaways & Limitations

    A mixing ratio of k = 0.1 provides a strong balance between zero-shot likelihood, few-step generation quality, and scalable large-model performance.

  • Takeaways & Limitations

    XDLM has not yet been trained from scratch at large scale, and its performance crossover, domain-specific sampling, and joint text–vision balance remain incompletely investigated.

Abstract

from arXiv · show

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

1. Introduction

MDLMs favor semantic understanding and zero-shot generalization, while UDLMs favor few-step generation, leaving a gap in balanced performance. XDLM addresses this gap by combining their noise kernels and targets a favorable understanding–generation trade-off.

  • MDLMs achieve stronger likelihood modeling and zero-shot generalization, whereas UDLMs perform better in low-step generation.
  • 17.6%: UDLM outperforms MDLM on ImageNet-1K generation in an 8-step regime.
  • XDLM uses a stationary noise kernel that bridges uniform and masked noise while recovering MDLM and UDLM as limiting cases.
  • XDLM combines UDLM and MDLM noise kernels to balance understanding capability and generation capability.Figure 1 evaluates understanding with zero-shot perplexity and generation with 32-step generation perplexity.
  • 5.4 points: XDLM surpasses UDLM on averaged zero-shot language metrics while trailing MDLM by 0.45.

2. Preliminary

Discrete diffusion models define forward corruption and reverse denoising Markov processes over categorical states. Their training objective uses a variational lower bound decomposed into prior, diffusion, and reconstruction terms.

  • Discrete diffusion models use a forward process to corrupt categorical data and a parameterized reverse process to denoise latent states.
  • Forward transitions are governed by row-stochastic matrices, with uniform and absorbing-state transitions corresponding to UDLM and MDLM.
  • The reverse posterior q(z_s | z_t, x_0) is derived using Bayes’ rule and the Markov property.
  • A neural network predicts the clean-data distribution, which parameterizes the reverse transition as p_θ(z_s | z_t) = q(z_s | z_t, x̃_0).
  • The variational lower bound decomposes into prior, diffusion, and reconstruction terms, with the continuous-time limit reducing the objective to the diffusion term.

3. miXed Diffusion Language Modeling via Stationary Noise Kernels

XDLM constructs a stationary mixture of signal preservation and a fixed noise kernel, then simplifies posterior and KL calculations into scalar forms. Its uniform-noise and masking endpoints recover UDLM and MDLM objectives.

  • 3.1. The Forward Process with Stationary Kernels: XDLM defines the forward transition as a convex combination of the identity matrix and a stationary noise kernel.The scalar schedule satisfies α_t|s + β_t|s = 1.
  • 3.1. The Forward Process with Stationary Kernels: Stationarity keeps noise characteristics invariant across timesteps and decouples them from scheduling dynamics.
  • 3.1. The Forward Process with Stationary Kernels: The kernel maps every input state toward a target distribution π and can mix uniform noise with absorbing special-token masses.
  • 3.2. Efficient Sampling and Training via Scalar Formulation: XDLM replaces expensive matrix-based posterior and KL computations with helper functions and equivalent scalar expressions.The reformulation reduces computational complexity for large vocabularies and supports efficient, stable training and sampling.
  • 3.2. Efficient Sampling and Training via Scalar Formulation: The limiting-case substitutions produce a tractable scalar training loss that bypasses explicit posterior-distribution and KL-divergence computation.
  • 3.3. Relationship to MDLM and UDLM: k = 0 yields pure masking and the standard MDLM cross-entropy objective, while k = 1 yields uniform noise and the UDLM loss.

4. Experiment

Experiments evaluate XDLM across zero-shot likelihood, language and image generation, large-model adaptation, efficiency, and training dynamics. Across these settings, XDLM balances masked and uniform diffusion strengths while improving scalability and computational efficiency.

  • Zero-shot likelihood: XDLM achieves an average zero-shot PPL of 54.110 across seven datasets, near MDLM’s 53.650 and GIDD’s 53.384 while outperforming UDLM’s 59.574.
  • Language generation: In 8–32-step OWT generation, XDLM reaches quality comparable to UDLM and surpasses purely masked models, while in 512–1024-step regimes it outperforms UDLM.
  • Image generation: 25.77 FID: XDLM achieves the lowest FID at 16 steps on ImageNet-1K without classifier-free guidance, remaining highly competitive with UDLM.
  • Image generation: 13.55 FID at 4 steps and 8.96 FID at 8 steps: XDLM outperforms UDLM with classifier-free guidance, while MDLM reaches 6.73 FID at 16 steps.
  • Large-language-model adaptation: 15.0 MBPP: LLaDA-XDLM doubles the original LLaDA’s 6.8 score at 32 steps and reduces non-compilable code errors.
  • Training dynamics: During LM1B training, XDLM eventually outperforms initially stronger GIDD and MDLM models, while on ImageNet-1K it achieves the lowest FID from the start.
  • Efficiency: 396,398 tokens/s forward throughput and 31.4 GB memory: XDLM nearly doubles GIDD’s forward throughput and uses substantially less memory than UDLM and GIDD.

5. Related work

Related work developed masked and uniform-noise discrete diffusion paradigms, alongside approaches that mix or generalize their corruption processes. GIDD is the closest predecessor, whereas XDLM uses a stationary kernel and scalar formulation.

  • Masked diffusion: MDLMs use absorbing-state corruptions and optimize continuous-time variational objectives for masked diffusion.
  • Uniform-noise diffusion: UDLM transition kernels mix the identity with a uniform vocabulary distribution, extending earlier multinomial and D3PM formulations.
  • Bridging approaches: Mixed-corruption and position-specific schedule methods connect masked, autoregressive, and diffusion modeling through different training or noise assignments.
  • Closest related work: Unlike GIDD’s dynamically blended transition matrices, XDLM enforces a stationary noise kernel that enables an efficient scalar formulation for sampling and training.

6. Conclusion

XDLM unifies masked and uniform-noise diffusion theoretically and computationally, then demonstrates improved performance across understanding, generation, and large-model code benchmarks. The authors identify several unresolved scaling, sampling, multimodal, and inference questions.

  • Contributions: XDLM’s weighted row-stochastic forward process recovers MDLM and UDLM as special cases while reducing computational complexity for large-vocabulary training.
  • Empirical findings: At k = 0.1, XDLM reaches a reported sweet spot combining masking-based zero-shot likelihoods with uniform-noise few-step generation and sample diversity.
  • Limitations: The study does not train XDLM from scratch at large scale and leaves performance crossover, domain-specific sampling, multimodal balance, post-training, and acceleration incompletely investigated.

Impact Statement

XDLM unifies masked and uniform diffusion, with each recovered as a special case, while its generation quality improves as the sampling budget increases. Its trajectories show early rejection and re-masking of poor tokens, followed by increasingly coherent refinement and diminishing returns beyond 64 steps.

  • XDLM reduces to MDLM at k = 0 and to UDLM at k = 1, formally connecting the two diffusion paradigms.
  • Sampling-budget effects: At T = 4, XDLM fails to converge to the data manifold, producing valid tokens without syntactic binding.At T = 8, superficial syntax emerges, but semantic content remains illogical.
  • Sampling-budget effects: At T = 16 and T = 32, outputs progress from grammatical structures with thematic drift to cohesive clauses with clear subject-verb-object relationships.
  • Sampling-budget effects: At T = 64, XDLM produces semantically robust text handling complex entities and domain terminology, while T = 128 adds only marginal improvements.The reported trajectory identifies approximately 64 steps as a threshold for fine-grained semantic detail, with diminishing returns thereafter.
  • Internal generation dynamics: XDLM begins from masks and random uniform-noise tokens, then rejects many poor initializations through re-masking before combining generation, substitution, and error correction.By T = 24 → T = 32, re-masking disappears and the process focuses on filling gaps and fine-grained lexical substitutions.

G. Details of Zeroshot Capability of XDLM Trained on OWT

XDLM’s zero-shot capability remains close to MDLM and improves substantially over UDLM, with performance depending strongly on the mixing ratio k. Its advantage is visible throughout training rather than only at convergence.

  • XDLM’s zero-shot evaluation measures perplexity on seven external datasets after training on OpenWebText.The evaluation uses validation splits from AG News, LAMBADA, LM1B, Penn Treebank, scientific papers, and WikiText.
  • Mixing-ratio sensitivity: Lower mixing ratios preserve validation performance near MDLM, whereas increasing k toward 0.9 degrades performance toward UDLM.At k = 1e-3 and k = 0.1, XDLM achieves PPL values of 23.495 and 24.097; at k = 0.9, performance degrades to 25.731.
  • Zero-shot results: 54.110 average zero-shot PPL for XDLM (k = 0.1) is close to MDLM’s 53.650 and below UDLM’s 59.574.
  • Training dynamics: XDLM (k = 0.1) follows MDLM’s zero-shot PPL trajectory from 64.85 early in training to 54.11 at convergence, while UDLM remains higher.The alignment persists across 100k-step intervals through 1M training steps.

H. Detailed Language Generation Results

XDLM’s mixing ratio k controls the trade-off between few-step efficiency and high-fidelity language generation, interpolating between UDLM-like and MDLM-like behavior. Higher k is favored with limited sampling steps, while lower k supports stronger performance with larger budgets.

  • OWT: Higher k improves XDLM’s few-step language generation, aligning it with UDLM’s efficiency.On OWT, XDLM with k = 0.9 reaches perplexity 189.750 at 8 steps, near UDLM’s 183.991 and far below MDLM’s 711.382.
  • OWT: Lower k preserves XDLM’s ability to achieve high-fidelity generation when many sampling steps are available.With k = 0.1, XDLM reaches perplexity 52.609 at 1024 steps without sacrificing entropy-based diversity.
  • Cross-dataset pattern: Across OWT and LM1B, tuning k lets XDLM combine uniform-noise efficiency at small budgets with masking-based quality at larger budgets.
  • LM1B: At 4 LM1B sampling steps, XDLM with k = 0.5 achieves perplexity 232.702, effectively matching UDLM’s 232.403 and outperforming MDLM’s 377.177.

I. Detailed Image Generation Results

Image-generation results show that XDLM’s best mixing ratio depends on the dataset, sampling setting, and guidance regime. Intermediate k values can outperform pure masked and uniform-noise baselines, while increasing k improves performance on CIFAR-10.

  • Cross-setting conclusion: The optimal k is not fixed across image-generation scenarios; tuning within 0 < k < 1 is important for maximizing performance.
  • ImageNet-1K: With standard ImageNet-1K conditioning, XDLM at k = 0.5 achieves FID 23.417 at 16 steps, outperforming the low-mixing variant and UDLM’s 26.242.
  • ImageNet-1K with CFG: With CFG, XDLM at k = 0.1 achieves FID 8.625 at 16 steps, surpassing UDLM’s 8.980 and XDLM at k = 0.5 with 8.790.The guidance setting shifts the optimal operating point toward lower k.
  • CIFAR-10: On CIFAR-10, increasing k from 10^-3 to 0.5 improves 32-step FID from 164.040 to 56.299, narrowing the gap with UDLM’s 41.027.

J. Detailed LLaDA Continual Pretraining Results

Continual pretraining and refinement sampling extend XDLM to LLaDA’s 8B setting, where it improves low-compute reasoning and code generation. The gains include stronger MBPP performance and fewer syntactically invalid outputs than the baselines and ablations.

  • Generation strategy: The LLaDA-XDLM generation strategy adds a refinement branch for k > 0.0, allowing previously generated tokens to be revised during denoising.This changes generation from purely additive mask filling into a dynamic error-correction process.
  • Benchmark results: LLaDA-XDLM outperforms the baselines across reasoning and code-generation benchmarks, particularly at low sampling budgets.The evaluation covers GSM8K, MATH, BBH, HumanEval, and MBPP across varying sampling steps.
  • Code generation: At 32 sampling steps, LLaDA-XDLM achieves MBPP 15.00, versus 6.80 for LLaDA and 4.40 for the continued MDLM control.The result effectively doubles the base LLaDA performance.
  • MBPP failure modes: LLaDA-XDLM reduces failed MBPP cases from 429 to 304 and increases passed outputs from 34 to 75.Failed cases are defined as non-compilable or syntactically invalid code.
  • Controls and ablations: LLaDA-XDLM outperforms the continued-pretraining MDLM control and the inference-only ablation, supporting gains from the learned XDLM formulation rather than extra training or sampling tricks alone.

K. Detailed Computational Efficiency Analysis

XDLM’s scalar reformulation improves computational efficiency among methods using uniform-noise components. It achieves the highest throughput in the evaluated forward, forward-backward, and sampling regimes while avoiding the higher memory costs of GIDD and UDLM.

  • Throughput: XDLM achieves the highest throughput among uniform-noise methods in forward, forward-backward, and sampling evaluations.The evaluations represent perplexity estimation, training, and generation, respectively.
  • Mechanism: XDLM’s efficiency is attributed to a scalar reformulated sampling and training strategy that avoids expensive matrix operations.
  • Baseline comparison: MDLM remains the overall throughput and memory leader because of its simpler absorbing noise kernel.
  • Memory: Among methods with uniform-noise components, XDLM has the lowest memory cost relative to GIDD and UDLM.

L. Detailed Training Dynamics

The section provides numerical training-dynamics data for text and image generation, supplementing the visual analysis with exact metrics across training trajectories. XDLM shows a decisive final-checkpoint advantage on ImageNet-1K, while UDLM achieves the lowest baseline perplexity late in LM1B training.

  • Exact Perplexity, Entropy, FID, and IS values supplement the training-dynamics visualization across the reported trajectories.The numerical data covers language modeling and image generation experiments.
  • By 1M LM1B training steps, UDLM reaches the lowest baseline perplexity, 96.385, under a fixed budget of 128 sampling steps.MDLM performs better initially, whereas UDLM scales strongly during later training stages.
  • XDLM maintains a decisive performance advantage during ImageNet-1K training evaluated over 500k steps with 16 sampling steps.The table reports FID and IS throughout the training trajectory.
  • At the final ImageNet-1K checkpoint, XDLM with k = 0.5 records the lowest reported FID of 23.417 and highest IS of 46.620.These are the reported best final-checkpoint values in the quantitative training-dynamics table.
Loading 2602.01362v1…