Attention Sinks in Diffusion Language Models

TL;DR

This paper investigates attention sinks in diffusion language models (DLMs), a paradigm that uses forward masking and bidirectional denoising to generate text in parallel. Through empirical analysis of Dream-7B, LLaDA-8B, and MMaDA-8B, it shows that DLMs exhibit sinks whose positions move during denoising, and that these sinks have a markedly different role from those in autoregressive models. The authors introduce a metric to identify sinks and demonstrate that masking a subset of sinks results in only modest performance degradation, highlighting the robustness and distinct attention organization of DLMs. Overall, the work provides new insights into DLM internals, with implications for long-context modelling, interpretability, and potential efficiency gains in diffusion-based generation.

Abstract

Masked Diffusion Language Models (DLMs) have recently emerged as a promising alternative to traditional Autoregressive Models (ARMs). DLMs employ transformer encoders with bidirectional attention, enabling parallel token generation while maintaining competitive performance. Although their efficiency and effectiveness have been extensively studied, the internal mechanisms that govern DLMs remain largely unexplored. In this work, we conduct an empirical analysis of DLM attention patterns, focusing on the attention sinking phenomenon, an effect previously observed in various transformer-based architectures. Our findings reveal that DLMs also exhibit attention sinks, but with distinct characteristics. First, unlike in ARMs, the sink positions in DLMs tend to shift throughout the generation process, displaying a dynamic behaviour. Second, while ARMs are highly sensitive to the removal of attention sinks, DLMs remain robust: masking sinks leads to only a minor degradation in performance. These results provide new insights into the inner workings of diffusion-based language models and highlight fundamental differences in how they allocate and utilize attention compared to autoregressive models.

Figures (10)

• Figure 1: Incoming attention scores for each token in LLaDA-8B llada across denoising steps. Unlike autoregressive models, DLMs exhibit attention sinks that shift across the sequence as tokens are progressively unmasked.
• Figure 2: Snapshot of an inference step for different language models. ARMs and Dream-7B predict the next token, while MMaDA-8B and LLaDA-8B predict the current one. MMaDA-8B and LLaDA-8B perform semi-autoregressive block decoding, where only tokens in the current block are unmasked, while Dream-7B may unmask a token at any position.
• Figure 3: Distribution of attention scores in LLaDA-8B llada across denoising steps. Only a few tokens, the attention sinks, receive a very high attention score, while the majority of tokens in the sequence have scores close to zero.
• Figure 4: Moving sink in LLaDA-8B, and MMaDA-8B.
• Figure 5: Cumulative attention score for LLaDA-8B's sink across heads and layers. The variation of the model's main sink token is displayed across the different heads and layers, averaged through time. Brighter colours indicate higher attention score. In later layers there are usually fewer sinks and the attention score is therefore higher, as it is shared among fewer sink tokens.