Source-linked AI summary
Attention Sinks in Diffusion Language Models
Maximo Eduardo Rulli, Simone Petruzzi, Edoardo Michielon, Fabrizio Silvestri, Simone Scardapane, Alessio Devoto
TL;DR
The internal attention mechanisms of diffusion language models remain relatively unexplored. This paper empirically studies attention sinks in three DLMs and finds that their sinks shift during denoising, while the models are more robust to sink masking than autoregressive models.
Problem
How attention works internally in DLMs remains largely unexplored, despite their bidirectional, parallel generation process.
Source: §1 · p. 2Method
The authors empirically analyze attention patterns, including sink behavior, in Dream-7B, LLaDA-8B, and MMaDA-8B.
Source: §1 · p. 2Results
DLM sinks shift during denoising, and DLMs are more robust than autoregressive models to sink removal, continuing to work with a performance drop when sinks are masked.
Source: §1 · p. 2, §5.2 · p. 8Takeaways & Limitations
DLMs exhibit dynamic sinks and remain functional under sink masking, distinguishing their attention behavior from autoregressive models.
Source: §6 · p. 9The analysis covers instruct DLMs, not their corresponding base versions, and does not examine how training-procedure changes affect sink behavior.
Source: §8 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
The paper empirically examines attention sinks in masked diffusion language models, whose internal attention mechanisms remain relatively unexplored. It finds that their sinks shift during denoising and that the models are more robust to sink masking than autoregressive models.
1 Introduction
- DLMs iteratively denoise a fully masked sequence with bidirectional attention, but how this architecture affects their internal mechanisms remains largely unexplored.Source: §1 · p. 2
- The study examines attention patterns in three open-source masked DLMs: Dream-7B, LLaDA-8B, and MMaDA-8B.The models differ in initialization and training: Dream-7B starts from an autoregressive model, LLaDA-8B is trained from scratch, and MMaDA-8B is multimodal and trained from LLaDA-8B.Source: §1 · p. 2
- DLM attention sinks are unstable and shift during denoising, while DLMs are more robust to sink removal than autoregressive models.The authors attribute this robustness to DLMs’ decoding strategy.Source: §1 · p. 2
- The paper reports empirical evidence that attention sinks consistently emerge in DLMs.Source: §1 · p. 2
- The authors characterize sink disappearance and movement during inference and introduce a metric to track sink intensity and location across denoising steps.Source: §1 · p. 2
2 Related Work
Related work describes the development of discrete diffusion language models and prior attention-sink research. The paper positions its study as the first investigation of attention sinks in DLMs.
2.1 Diffusion Language Models
- Discrete DLMs operate directly on token vocabularies, iteratively refining masked sequences; earlier approaches faced scaling challenges.Continuous-space diffusion instead embeds tokens into continuous vectors before denoising and mapping them back to discrete tokens.Source: §2.1 · p. 2
- Discrete DLMs now include open-source models at 7 billion parameters and beyond, narrowing the performance gap with autoregressive models.Source: §2.1 · p. 2
- The study investigates attention patterns in large discrete DLMs that operate directly in vocabulary space.Source: §2.1 · p. 2
2.2 Attention Sink in Transformers
- Figure 2 contrasts next-token prediction in ARMs and Dream-7B with current-token prediction in LLaDA-8B and MMaDA-8B.LLaDA-8B and MMaDA-8B unmask tokens in the current block, whereas Dream-7B may unmask a token at any position.Source: Figure 2 · p. 3
- Prior work describes attention sinks and studies their properties and explanations across transformer architectures; this paper presents the first investigation of the phenomenon in DLMs.An attention sink is a small subset of tokens receiving disproportionate attention from other sequence tokens.Source: §2.2 · pp. 2–3
3 Background on Masked Discrete Diffusion
Autoregressive models factor sequence probability into left-to-right conditional predictions. Masked discrete DLMs instead corrupt sequences with masks and learn to reverse that process through denoising.
3 Background on Masked Discrete Diffusion
- ARMs factor a sequence’s probability into left-to-right token predictions conditioned on preceding tokens.Source: §3 · p. 3
- Masked discrete DLMs use a forward Markov process that progressively replaces clean-sequence tokens with [MASK], followed by a learned reverse denoising process.The model is often trained to predict clean data from a noisy input, then uses that prediction to approximate the reverse transition.Source: §3 · p. 3
- LLaDA-8B and MMaDA-8B predict masked tokens and use semi-autoregressive block diffusion, while Dream-7B uses a shift-operation objective initialized from an autoregressive model.LLaDA-8B and MMaDA-8B unmask blocks from left to right; Dream-7B is trained to predict the next token when a token is masked.Source: §3 · pp. 3–4
4 Analysis of Attention Sinks in Masked Diffusion Language Models
The analysis finds attention sinks in masked diffusion language models, but their positions can shift or vanish during denoising. Unlike the tested autoregressive model, the tested DLMs show only modest degradation when sinks are masked.
4.1 Definition of Attention Sink
- The attention-score distribution in LLaDA-8B shows a few tokens receiving very high scores while most sequence tokens have scores close to zero.The authors report similar patterns for Dream-7B and MMaDA-8B.Source: Figure 3 · p. 4, §4 · p. 4
- The analysis defines a sink by its cumulative attention score: the average attention a token receives from all tokens at a denoising step, exceeding the average for other tokens by a threshold.The experiments use threshold ϵ = 3, selected to filter out at least 96% of sequence tokens.Source: §4.1 · p. 4
4.2 Sink Patterns
- DLM sinks may shift position or vanish across denoising steps, unlike the static sinks commonly observed in autoregressive models.Some LLaDA-8B sinks persist for consecutive steps before disappearing, while others last only one timestep.Source: §4.2 · p. 5
- In deeper LLaDA-8B layers, sinks decrease to one or two per layer, and masked and unmasked tokens can maintain separate sinks that switch gradually.Some heads show a one-step sink that disappears at the next denoising step.Source: §4.2 · p. 5, Figure 6 · p. 6
- LLaDA-8B sinks often form on punctuation, whitespace, or end-of-sequence tokens, whereas Dream-7B sinks tend to follow position rather than token content.Dream-7B sinks often begin at the rightmost masked token and shift left as tokens are unmasked.Source: §4.2 · pp. 5–6
- MMaDA-8B sinks are mostly static, while sinks tend to shift right in LLaDA-8B and move right-to-left in Dream-7B.MMaDA-8B has fewer evident sink patterns; its sinks often remain at their initial positions.Source: §4.2 · pp. 5–6
4.3 Robustness of DLMs to Masking Sinks
- Masking one sink degrades each tested DLM by less than 1%, while masking even one sink causes severe performance drops in LLama-3.1-8B.Degradation becomes substantial only when masking 10 sinks, mostly in MMaDA-8B.Source: Table 1 · p. 7, §4.3 · p. 7
5 Discussion
DLM sinks shift during generation and may align with content or structure rather than fixed positions. DLMs remain functional after sink masking, and their moving sinks may support long-context generation and past-context removal.
5.1 Dynamic Sinks and Positional Encoding
- Sink positions vary across denoising steps and can split between masked and unmasked tokens, often landing on semantic or structural markers.The top-five sink-token frequencies include whitespace and [MASK] across models; MMaDA also frequently sinks on the prompt-template assistant token.Source: §5.1 · p. 8, Table 2 · p. 8
5.2 Robustness to Masking Sinks
- DLMs continue functioning after sink masking, unlike ARMs, which would fail completely when sinks are removed.The authors believe bidirectional attention and iterative denoising together provide stability; they suggest affected tokens may remain masked when confidence falls.Source: §5.2 · p. 8
5.3 Long Context Modelling
- Moving sinks might help DLMs maintain long-range connections, and their future-position sinks allow past context to be discarded without significant performance degradation.The authors suggest future-position anchors may benefit long reasoning and planning, where a reference anchor in the future could be useful.Source: §5.3 · p. 9
6 Conclusion
The analysis finds that DLMs consistently exhibit attention sinks, but their behavior differs from that in autoregressive models. Their sinks shift and align with semantic or structural tokens, while sink masking causes only minor degradation.
6 Conclusion
- DLMs consistently exhibit attention sinks, with behavior that differs from autoregressive models.The analysis presents these findings as its first empirical study of attention sinks in DLMs.Source: §6 · p. 9
- DLM sinks often shift across denoising steps and align with semantic or structural tokens rather than fixed positions.The conclusion describes these behaviors as evidence of flexible attention organization in diffusion models.Source: §6 · p. 9
- DLMs remain robust to sink masking, suggesting their bidirectional and iterative generation avoids reliance on single anchor tokens.Source: §6 · p. 9
7 Future Work
The authors identify open questions about what future-position sinks store, whether sinks could support acceleration or compression, and how observed sink behaviors arise.
7 Future Work
- The information stored in sinks corresponding to future positions remains unclear; mechanistic analysis such as the Logit Lens is suggested.Source: §7 · p. 9
- Whether DLM sinks could be exploited for acceleration or compression remains an open question.The authors relate this possibility to the original use case of sinks in Xiao et al. (2023).Source: §7 · p. 9
- The study did not explain the observed sink behaviors in detail because interpretability-focused investigation lies beyond its primarily empirical scope.Source: §7 · p. 9
8 Limitations
The analysis covers three instruct DLMs but not their corresponding base versions, and it does not test how training-procedure changes affect sink behavior.
8 Limitations
- The analysis covers instruct DLMs only; corresponding base-model versions were not tested.Source: §8 · p. 9
- The study focuses on pre-trained models and does not examine how changes to training procedures affect sink behavior.The authors note that this aspect has been investigated for ARMs.Source: §8 · p. 9
A Additional plots
Additional plots show that only a few tokens receive disproportionately high attention scores.
- Only a few tokens receive disproportionately high attention scores in the additional plots.Source: A Additional plots · p. 12
B Selection of Sink Threshold
The investigation uses ϵ = 3 as a balanced sink threshold; at that setting, all analyzed DLMs filter out at least 96% of tokens.
- At ϵ = 3, all analyzed DLMs filter out at least 96% of tokens when selecting sinks.Source: B Selection of Sink Threshold · p. 12, Figure 10 · p. 12
- Figure 10 plots the percentage of tokens not considered sinks across ϵ values for a sequence of 64 tokens.Source: Figure 10 · p. 12