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

2026-06-10

AI Summary

Existing post-training methods for diffusion-based language models employ fully random masking, which disregards semantic dependencies among tokens and thereby limits reasoning capabilities. This work addresses this limitation by analyzing attention mechanisms and revealing that tokens receiving high attention are critical for generation stability and reasoning performance. Building on this insight, the authors propose the Attention-Guided Denoising Optimization (AGDO) framework, which integrates attention-guided mechanisms into both the denoising schedule and training optimization of diffusion models. AGDO enables dynamic, semantics-aware masking and reinforces key tokens through an approach that combines attention analysis, attention-aware denoising strategies, and keyword-weighted supervised fine-tuning with reinforcement learning. Evaluated on mathematical and code reasoning benchmarks, AGDO significantly outperforms current state-of-the-art post-training methods.

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.

Problem

• diffusion language models
• random masking
• token dependencies
• reasoning performance
• post-training

Innovation

• attention-guided denoising
• diffusion language models
• token dependency
• reasoning optimization
• non-random masking