Reading Notes

Personal reading notes. Original paper: arXiv:2606.12273v1.

arxiv_id: 2606.12273v1

Beyond Fully Random Masking: Attention-Guided Denoising and Optimization for Diffusion Language Models Reading Note

1. Paper Overview

This paper proposes a framework named AGDO (Attention-Guided Denoising and Optimization), designed to address the mismatch between masking strategies and intrinsic token dependencies in the post-training phase of Diffusion Large Language Models (dLLMs). The authors point out that existing dLLM post-training methods primarily rely on fully random masking strategies, ignoring the intrinsic dependencies formed between tokens via attention mechanisms. Through empirical analysis, the paper finds that tokens with stronger attention connections to unmasked contexts exhibit higher generation stability and play a critical role during inference. Based on this, the AGDO framework utilizes attention structures to guide the denoising order and emphasizes attention-critical tokens during Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) stages. Experiments demonstrate that this method significantly improves reasoning performance on math and code benchmarks, outperforming existing dLLM post-training methods.

2. Background and Motivation

Diffusion Large Language Models (dLLMs) provide an efficient alternative to Autoregressive (AR) models through parallel decoding. Unlike AR models that generate tokens sequentially from left to right, dLLMs iteratively denoise and decode tokens in parallel, offering significant efficiency advantages during the inference phase. Recent models such as LLaDA and Dream have demonstrated performance comparable to SOTA AR models of similar scale.

However, effective post-training for dLLMs still faces challenges. Existing full-attention dLLM post-training methods (e.g., diff-GRPO and wd1) typically rely on randomly masking tokens and optimizing masked positions. While this approach is simple and efficient, it fails to align with the actual inference dynamics of dLLMs, leading to a mismatch between training and inference. Although recent works have attempted to introduce different masking strategies (such as block-level SFT or following a natural left-to-right generation order), they still impose externally defined decoding orders. This ignores a key attribute of full-attention dLLMs: under bidirectional attention, token dependencies are not strictly determined by positional order but emerge dynamically through attention interactions.

3. Empirical Analysis of Attention Mechanism

To gain a deeper understanding of these intrinsic dependencies, the authors conducted an empirical analysis of attention patterns within dLLMs. The analysis focused primarily on the last layer, where semantic dependencies are usually most evident, using the Dream-v0-Instruct-7B model as the main subject.

The paper reveals two key observations:

1. Sparsity and Temporal Consistency: Attention distributions exhibit strong sparsity and temporal consistency across denoising steps, indicating that each token consistently relies on a smaller and stable set of context tokens.
2. Generation Stability: Tokens that attend more strongly to already denoised tokens demonstrate significantly higher probability stability during generation, suggesting that the denoising order is crucial for maintaining generation reliability.

To visually demonstrate these findings, the paper provides a visualization analysis of attention dynamics. As shown below, Figure 1 illustrates the attention dynamics during the denoising process on Dream-v0-Instruct-7B.

Figure 1 Interpretation:

• Subfigures (a) to (d): Show the attention distribution between tokens at randomly selected denoising steps. It can be observed that the attention matrix is not uniformly distributed but presents specific sparse patterns.
• Subfigures (e) to (h): Illustrate the changes in attention from randomly selected tokens to other tokens throughout the entire denoising process.
• Red Dashed Line: Indicates the step at which the corresponding token is unmasked.

These findings suggest that effective dLLM training strategies should explicitly align the denoising trajectory with attention-induced token dependencies.

4. Method: AGDO Framework

Based on the above insights, the authors propose the AGDO framework, a two-stage post-training framework designed to explicitly align the denoising trajectory of dLLMs with their intrinsic attention structure.

4.1 Attention-Guided Denoising Order

AGDO first derives an attention-guided denoising order from effective attention scores, which determines the sequence in which tokens are unmasked during training. By ensuring tokens are denoised only when sufficient attention-supported context is available, this denoising order forms the foundation of AGDO. This method avoids the training-inference mismatch caused by random masking, making the training process more consistent with the model’s actual attention dependency mechanisms.

4.2 Supervised Fine-Tuning and Reinforcement Learning

Based on the denoising order, AGDO performs Supervised Fine-Tuning (AGDO-SFT) and Reinforcement Learning (AGDO-RL).

• Attention-Critical Token Emphasis: Through an attention-based re-weighting strategy, attention-hub tokens are further emphasized.
• Policy Optimization: Addressing the characteristic that dLLMs cannot decompose tokens sequentially like AR models, the paper employs mean-field approximations to calculate the probability ratios and KL divergence terms required for policy optimization, thereby supporting efficient single-pass estimation.

5. Experiments and Conclusion

The paper conducted extensive experiments on challenging math and code benchmarks. Results show that AGDO significantly improves the reasoning capabilities of masked dLLMs and consistently outperforms existing post-training methods.

Summary of Key Contributions:

1. Empirical Analysis: Revealed the sparsity and temporal consistency of attention in dLLMs, as well as the relationship between attention strength and generation stability.
2. Method Innovation: Proposed the AGDO framework, which utilizes attention structures to guide the denoising order and optimizes critical tokens during SFT and RL stages.
3. Performance Improvement: Validated the importance of aligning training dynamics with intrinsic attention dependencies, proving that attention-guided denoising is a principled choice for dLLM design.

This work emphasizes the importance of considering intrinsic dependencies between tokens in dLLM post-training, providing a new direction for improving the reasoning capabilities of diffusion language models.