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

• URL: http://arxiv.org/abs/2606.12273v1
• Date: Wed, 10 Jun 2026 16:14:23 GMT
• Title: Beyond Fully Random Masking: Attention-Guided Denoising and Optimization for Diffusion Language Models
• Abstract summary: We present an empirical analysis of attention in Diffusion large language models (dLLMs)<n>We show that tokens attending more strongly to unmasked context exhibit greater generation stability and play a critical role in reasoning.<n>Motivated by these findings, we propose AGDO, an attention-guided denoising and optimization framework.
• Score: 129.8880969196575
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Early Decisions Matter: Proximity Bias and Initial Trajectory Shaping in Non-Autoregressive Diffusion Language Models [54.16797570104461]
  Diffusion-based language models (dLLMs) have emerged as a promising alternative to autoregressive language models.<n>We investigate non-autoregressive decoding in dLLMs by systematically analyzing its inference dynamics along the temporal axis.
  arXiv  Detail & Related papers  (2026-04-12T10:26:41Z)
• DARE: Diffusion Large Language Models Alignment and Reinforcement Executor [36.88819250898739]
  We present textbfDARE (textbfdLLMs textbfAlignment and textbfReinforcement textbfExecutor), an open framework for post-training and evaluating dLLMs.
  arXiv  Detail & Related papers  (2026-04-05T18:13:19Z)
• Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models [60.25003813232697]
  We propose Adaptive Retrieval-Augmented Masked Diffusion (ARAM) as a training-free adaptive guidance framework for Masked Diffusion Models (MDMs)<n>ARAM calibrates the guidance scale during denoising according to the Signal-to-Noise Ratio (SNR) of the distributional shift induced by retrieved context.<n>Experiments on multiple knowledge-intensive QA benchmarks show that ARAM improves overall QA performance over competitive RAG baselines.
  arXiv  Detail & Related papers  (2026-03-18T12:54:50Z)
• Attention-Based Sampler for Diffusion Language Models [16.008464195041984]
  We propose a new training-free decoding algorithm, termed Attn-Sampler.<n>We show that optimal sequence likelihood can be approximately achieved by decoding tokens in descending order of their attention column sums.
  arXiv  Detail & Related papers  (2026-03-18T07:49:13Z)
• Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages [7.728607011076579]
  Reinforcement learning has been effective for post-training autoregressive (AR) language models, but extending these methods to diffusion language models (DLMs) is challenging due to intractable sequence-level likelihoods.<n>We formulate a diffusion-based sequence generation as a finite-horizon Markov decision process over the denoising trajectory and derive an exact, unbiased policy gradient.<n>Experiments on coding and logical reasoning benchmarks demonstrate state-of-the-art results, with strong competitive performance on mathematical reasoning.
  arXiv  Detail & Related papers  (2026-03-13T01:38:44Z)
• Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective [85.06838178922791]
  Reinforcement Learning (RL) has proven highly effective for autoregressive language models.<n>But adapting these methods to diffusion large language models (dLLMs) presents fundamental challenges.<n>We propose a principled RL framework that treats entire sequence generation as a single action and uses the ELBO as a tractable sequence-level likelihood proxy.
  arXiv  Detail & Related papers  (2025-12-03T13:05:32Z)
• MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization [66.82303841930752]
  diffusion language models (DLMs) have presented a promising alternative to traditional autoregressive large language models (LLMs)<n>DLMs still lag behind LLMs in reasoning performance, especially as the number of denoising steps decreases.<n>We propose a Multi-Reward Optimization (MRO) approach, which encourages DLMs to consider the token correlation during the denoising process.
  arXiv  Detail & Related papers  (2025-10-24T13:57:59Z)
• AURORA: Augmented Understanding via Structured Reasoning and Reinforcement Learning for Reference Audio-Visual Segmentation [113.75682363364004]
  AURORA is a framework designed to enhance genuine reasoning and language comprehension in reference audio-visual segmentation.<n>AURORA achieves state-of-the-art performance on Ref-AVS benchmarks and generalizes effectively to unreferenced segmentation.
  arXiv  Detail & Related papers  (2025-08-04T07:47:38Z)
• Don't Take Things Out of Context: Attention Intervention for Enhancing Chain-of-Thought Reasoning in Large Language Models [32.71672086718058]
  Few-shot Chain-of-Thought (CoT) significantly enhances the reasoning capabilities of large language models (LLMs)<n>We observe that isolated segments, words, or tokens within CoT demonstrations can unexpectedly disrupt the generation process of LLMs.<n>We propose a Few-shot Attention Intervention method (FAI) that dynamically analyzes the attention patterns of demonstrations to accurately identify these tokens.
  arXiv  Detail & Related papers  (2025-03-14T07:46:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.