Beyond Fully Random Masking: Attention-Guided Denoising and Optimization for Diffusion Language Models
Jia Deng, Junyi Li, Xin Zhao, Jinpeng Wang, Hongyu Lu, Ji-Rong Wen
Abstract
Diffusion large language models (dLLMs) offer an efficient alternative to autoregressive models through parallel decoding, yet existing post-training methods largely rely on random masking strategies that overlook intrinsic token dependencies. In this work, we present an empirical analysis of attention in dLLMs and show that tokens attending more strongly to unmasked context exhibit greater generation stability and play a critical role in reasoning. Motivated by these findings, we propose AGDO, an attention-guided denoising and optimization framework that aligns both training and optimization with attention-derived dependencies. AGDO determines the denoising order based on attention structure and emphasizes attention-critical tokens during supervised fine-tuning and reinforcement learning. Experiments on mathematical and coding benchmarks demonstrate that AGDO consistently improves reasoning performance, outperforming state-of-the-art post-training methods for dLLMs.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on9
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- Efficient Streaming Language Models with Attention SinksGuangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han et al.ICLR 2024 · 1,714 citations
- Large Language Diffusion ModelsShen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang et al.NeurIPS 2025 · 949 citations
- Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel DecodingChengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu et al.ICLR 2026 · 428 citations
- DiffuCoder: Understanding and Improving Masked Diffusion Models for Code GenerationShansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu et al.ICLR 2026 · 198 citations
Related papers
- Advancing Reasoning in Diffusion Language Models with Denoising Process RewardsShaoan Xie, Lingjing Kong, Xiangchen Song, Xinshuai Dong et al.ACL 2026 · 5 citations
- d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement LearningSiyan Zhao, Devaansh Gupta, Qinqing Zheng, Aditya GroverNeurIPS 2025 · 191 citations
- Learnability-Informed Fine-Tuning of Diffusion Language ModelsShubham Parashar, Atharv Chagi, Jacob Helwig, Lakshmi Madhavarapu et al.ICML 2026
- MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward OptimizationChenglong Wang, Yang Gan, Hang Zhou, Chi Hu et al.NeurIPS 2025 · 4 citations
- Early Decisions Matter: Proximity Bias and Initial Trajectory Shaping in Non-Autoregressive Diffusion Language ModelsJiyeon Kim, Sungik Choi, Yongrae Jo, Moontae Lee et al.ICML 2026 · 4 citations
