Attention Sinks in Diffusion Language Models

We explore attention sinks and their diverse manifestations in Diffusion Language Models, an emerging paradigm in text-generation.

Structure and reproducibility

🔨 Environment Setup

To reproduce our exact environment, use the provided environment.yml.
It includes the precise Python version and all pinned dependencies.

🔧 Installation

# Create the environment
conda env create -f environment.yml

# Activate it
conda activate dlms-sinks

🪐 Jupyter Files

There are files in the repository in a JupyText format, meaning that their extension is .py, but they are meant to be linked with a Jupyter Notebook using JupyText. We include these for an easy experience running and experimenting with the repository. In the main folder you can find attention_various_datasets.py as an example.

🧩 Inference and Visualization Scripts

We provide two helper scripts to reproduce some of the graphs we introduced in the paper.

inference_model.py

Runs inference for a selected model (LLaDA, MMaDA, or Dream) on GSM8K, saving attention weights and related outputs inside the runs/ directory.
Each run is timestamped and follows the format:

T{MMDD}_{HHMMSS}-S{steps}-L{gen_len}-B{block_len}

Arguments

Example

python inference_model.py LLaDA --num_samples 10 --gen_len 256

plots_creation.py

Generates heatmaps and videos for the saved inferences of a single model.

Arguments

Example

python plots_creation.py LLaDA ./runs

💡 Tip: Keep runs from different models in separate directories (e.g. ./runs/LLaDA, ./runs/MMaDA) to avoid conflicts.

Running evaluation

We provide evaluation scripts to test Diffusion Language Models (DLMs) (namely LLaDA-8B-Instruct, Dream-8B-Instruct, and MMaDA-8B) on reasoning and programming benchmarks. In order to investigate the role of attention sinks, we run evaluations under a sink-masked regime. This involves identifying high-attention sink tokens at each diffusion step and masking the top-k of them before the next step’s denoising phase. Example evaluation of LLaDA-8B-Instruct on the HumanEval-Instruct benchmark with 1 masked sink per step:

accelerate launch eval_llada.py --model llada_dist --tasks humaneval_instruct --batch_size 1 --model_args model_path=GSAI-ML/LLaDA-8B-Instruct,is_check_greedy=False,mc_num=128,cfg=0,steps=512,gen_length=512,block_length=32,remasking=low_confidence_sink_masked,epsilon=0.0,exclude_self=False,mask_sink_ratio=1.0,sink_top_k=1 --confirm_run_unsafe_code --apply_chat_template

Example evaluation of Dream-8B-Instruct on HumanEval-Instruct with 1 masked sink per step:

accelerate launch eval_dream.py --model dream --model_args pretrained=Dream-org/Dream-v0-Instruct-7B,trust_remote_code=True,max_new_tokens=768,diffusion_steps=768,dtype="bfloat16",temperature=0.1,top_p=0.9,alg="entropy",enable_sink_masking=True,sink_epsilon=0.0,sink_top_k=1 --tasks humaneval_instruct --device cuda --batch_size 1 --num_fewshot 0 --output_path output_reproduce/humaneval --log_samples --confirm_run_unsafe_code --apply_chat_template

FAQ

Here, we address some common questions about our work.

0. What is an "attention sink"?

Attention sinks are a widely studied phenomenon in the NLP literature. They were first noted in Xiao et al., and used to improve model efficiency. They can be easily recognised as vertical bright lines in the attention map. E.g.

Note that, in contrast with autoregressive models, diffusion-based ones exhibit one more generation dimension, i.e. time. That is why you see the vertical lines moving, as the sinks move across denoising steps.

Being more precise, we define an attention sink as:

j \text{ is a sink token if } \bar{A}^{(t,l,h)}_{j} > \frac{1}{S-1} \sum_{k \neq j} \bar{A}^{(t,l,h)}_{k} + \epsilon

1. What is a Diffusion Language Model?

Diffusion Language Models adapt the diffusion framework to discrete text generation. There have been both approaches to achieving this through a continuous formulation (e.g. Jo et al.) and a discrete one (e.g. D3PM). In this work, we explore how attention sinks manifest in the latter models. These models use a process reminiscent of BERT-style masking, where, during training, a subsection of the sequence tokens gets masked and the model has to predict the correct token at each position. The percentage of masking tokens is sampled uniformly between 0 and 1 (meaning that the whole sequence may be masked, in contrast with BERT).

During inference, these models start-off with a fixed generation length and number of denoising steps - which have to be proportional. The model then, at each denoising step, provides a prediction for each token in the sequence, and only the most probable k tokens are unmasked, where k is scheduled and may vary from model to model. This is how, after T denoising steps, one would get a full sequence of all-unmasked tokens.

2. How do Diffusion Language Models differ from Autoregressive ones?

The first key distinction among the two paradigms is that they use different kinds of Transformer architectures. Diffusion Language Models use an encoder-only architecture, while Autoregressive ones leverage a decoder-only one. This, in turn, has a great implication in how the attention mechanism works, as in diffusion-based ones, the model has access to the full sequence of tokens to trace complex relationships, while in the autoregressive ones, the causal mask is employed, yielding the well-known triangular matrix.

The other key difference lies in the generation capabilities. While autoregressive models need the token i to generate i+1, their diffusion counterparts can unmask a token at any position at any diffusion step, providing great generation flexibility and tuning.

3. How did you choose those values for $\epsilon$?

We chose the values for $\epsilon$ that provided us with a sufficiently strict metric to filter tokens that were not sinks, and at the same time be tolerantly precise to spot sinks in the generation process. We empirically showed that the selected ones are rigorous enough to just detect 5% of tokens as sinks:

Citation

@article{rulli2025attention,
  title={Attention Sinks in Diffusion Language Models},
  author={Rulli, Maximo Eduardo and Petruzzi, Simone and Michielon, Edoardo and Silvestri, Fabrizio and Scardapane, Simone and Devoto, Alessio},
  journal={arXiv preprint arXiv:2510.15731},
  year={2025}
}