DOS: Dependency-Oriented Sampler for Masked Diffusion Language Models

1. Problem Statement

Masked Diffusion Language Models (MDLMs) have emerged as a promising alternative to autoregressive (AR) models, offering flexible generation orders and the ability to perform parallel decoding. However, existing decoding strategies for pre-trained MDLMs suffer from two critical limitations:

• Reliance on Token-Level Uncertainty: Current methods (e.g., Confidence, Entropy, Margin) select tokens for unmasking based solely on the uncertainty of individual token predictions. They fail to explicitly model inter-token dependencies and sequence-level information. This leads to a mismatch between the marginal token distributions predicted by the model and the true joint distribution of the sequence.
• Sensitivity to Block Structures: Many parallel decoding strategies rely on heuristic block partitions (e.g., block-wise diffusion). If the imposed block structure does not align with the underlying dependency structure of the data, it disrupts the global organization of the sequence. This makes generation quality highly sensitive to block size choices and hinders the extension to pure single-block diffusion settings.

The core challenge is to develop a training-free decoding strategy that explicitly accounts for inter-token dependencies to guide the decoding order, thereby recovering the target joint distribution more effectively without relying on rigid block structures.

2. Methodology: Dependency-Oriented Sampler (DOS)

The authors propose DOS, a training-free decoding strategy that leverages the internal attention mechanisms of the transformer to guide token updates.

Core Concept

Instead of using uncertainty metrics, DOS uses the attention matrix from the transformer blocks as a proxy for inter-token dependencies. The underlying hypothesis is that the attention weights (specifically the scaled dot-product attention $softmax(QK^\top/\sqrt{d})$) encode how much a token depends on other tokens in the sequence.

Algorithm Steps

1. Attention Extraction: During the forward pass of the MDLM, the attention matrices are extracted from a specific transformer block.
2. Dependency Scoring:
   • Let $M$ be the set of masked positions and $U$ be the set of unmasked (observed) positions.
   • For each masked token $m \in M$, the dependency score is calculated by summing the attention weights from $m$ to all unmasked tokens $u \in U$:
     $$\text{dep}(m) = \sum_{u \in U} \text{Attn}(m, u)$$
   • This score quantifies how strongly the prediction at the masked position relies on the currently available context.
3. Token Selection: At each denoising step, the masked tokens with the highest dependency scores are selected for unmasking. This ensures that tokens most supported by the current context are resolved first, creating a dependency-aligned decoding order.
4. Integration: DOS can be used as a standalone scoring mechanism (replacing Confidence/Entropy) or integrated with existing parallel sampling methods (like the EB-Sampler) to further improve efficiency.

Key Technical Advantage

Since the attention matrices are already computed during the standard forward pass of the transformer, DOS introduces no additional training cost and negligible computational overhead.

3. Key Contributions

• Dependency-Aware Decoding: The paper identifies that token-level uncertainty is insufficient for parallel decoding and proposes using attention weights to explicitly model and leverage inter-token dependencies.
• Training-Free Strategy: DOS requires no fine-tuning or additional parameters; it utilizes the existing architecture's attention mechanisms.
• Robustness to Block Size: Unlike existing methods that degrade significantly as block sizes increase, DOS maintains consistent performance across both single-block (pure diffusion) and multi-block settings.
• Seamless Integration: DOS acts as a general criterion that can be combined with acceleration techniques (e.g., EB-Sampler, Fast-dLLM) to improve both generation quality and inference speed.

4. Experimental Results

The authors evaluated DOS on two open-source MDLMs (LLaDA-Instruct-8B and Dream-v0-Instruct-7B) across four benchmarks: HumanEval and MBPP (code generation), and GSM8K and MATH500 (mathematical reasoning).

Performance Highlights

• Superior Accuracy: DOS significantly outperformed uncertainty-based baselines (Confidence, Entropy, Margin) and block-based methods (Fast-dLLM, KLASS).
  • On LLaDA-Instruct-8B (Single-block, 512 tokens), DOS achieved 84.31% accuracy on GSM8K, compared to ~33% for the best baseline (Margin).
  • On Dream-v0-Instruct-7B, DOS achieved 59.15% on HumanEval and 45.00% on MATH500, surpassing all baselines.
• Efficiency: When integrated with the EB-Sampler (DOS+EB), the method achieved higher accuracy while maintaining competitive or better inference efficiency (measured in Number of Forward Evaluations, NFE) compared to other accelerated samplers.
• Robustness:
  • Block Size Sensitivity: While baselines like Fast-dLLM and KLASS showed significant performance drops as block sizes increased (e.g., from 32 to 256), DOS remained stable and robust.
  • Layer Selection: The method is robust across different transformer layers, with shallow layers (1–10) often providing sufficient dependency signals.

5. Significance and Conclusion

This work fundamentally shifts the paradigm of decoding for Masked Diffusion Language Models from uncertainty-driven to dependency-driven.

• Theoretical Insight: It demonstrates that recovering the target joint distribution in parallel decoding requires respecting the underlying dependency structure between tokens, which attention mechanisms naturally encode.
• Practical Impact: DOS enables MDLMs to generate high-quality long sequences (especially in reasoning and coding tasks) without the need for complex block-wise heuristics. It bridges the gap between the flexibility of diffusion models and the coherence of autoregressive generation.
• Future Directions: The authors suggest that while attention is a good proxy, future work could explore explicit dependency structures (e.g., DAGs) or multi-layer aggregation to further refine dependency modeling.

In summary, DOS provides a simple yet powerful mechanism to unlock the full potential of masked diffusion models, making them more viable for complex, long-horizon generation tasks.