One Token Is Enough: Improving Diffusion Language Models with a Sink Token

TL;DR

This work tackles the moving-sink instability in diffusion language models by analyzing sink tokens in value space and showing they occupy low-$L_2$ norm regions to absorb excess attention. It proposes a simple, position-stable extra sink token with a constrained attention mask that acts as a dedicated structural sink, improving stability and generation quality across initialization regimes and model scales. The method delivers consistent gains on a broad suite of benchmarks, remains effective even when the sink token carries negligible semantic content, and is robust to sink placement and quantity, making it a practical augmentation for DLMs. Overall, the paper advances understanding of attention dynamics in diffusion models and offers a lightweight architectural fix with tangible performance benefits.

Abstract

Diffusion Language Models (DLMs) have emerged as a compelling alternative to autoregressive approaches, enabling parallel text generation with competitive performance. Despite these advantages, there is a critical instability in DLMs: the moving sink phenomenon. Our analysis indicates that sink tokens exhibit low-norm representations in the Transformer's value space, and that the moving sink phenomenon serves as a protective mechanism in DLMs to prevent excessive information mixing. However, their unpredictable positions across diffusion steps undermine inference robustness. To resolve this, we propose a simple but effective extra sink token implemented via a modified attention mask. Specifically, we introduce a special token constrained to attend solely to itself, while remaining globally visible to all other tokens. Experimental results demonstrate that introducing a single extra token stabilizes attention sinks, substantially improving model performance. Crucially, further analysis confirms that the effectiveness of this token is independent of its position and characterized by negligible semantic content, validating its role as a robust and dedicated structural sink.

Figures (3)

• Figure 1: Moving attention sinks and their low-norm property in diffusion language models. Attention maps show representative patterns of LLaDA-Instruct (top) and Dream-Instruct (bottom) at two denoising steps ($step{=}0$ and $step{=}64$). Bright vertical stripes mark sink tokens that attract concentrated attention; their positions shift across steps, demonstrating moving sinks. Norm scatter plots compare the $L_2$ norms of value vectors for 10K sink tokens and 10K normal tokens sampled during inference on GSM8K gsm8k and HumanEval humaneval. Sink tokens exhibit significantly lower mean $L_2$ norms compared to normal tokens (2.36 vs. 7.60 for LLaDA; 3.15 vs. 7.79 for Dream). This consistent disparity suggests that DLMs preferentially select low-norm tokens as sinks during generation.
• Figure 2: Overview of the moving-sink phenomenon in diffusion language models and our stable sink token.Top: Diffusion inference iteratively denoises a partially masked sequence; at each step, the model predicts masked positions with bi-directional attention. Middle (left): In vanilla DLMs, attention can concentrate on low-norm tokens (often a [MASK]) that act as an implicit sink; since the masked set changes across steps, the sink position shifts over time (moving sink), which may increase inference instability. Middle (right): We prepend an extra [SINK] token, turning the moving sink into a stable sink. Bottom: Legend of the symbols used in the figure.
• Figure 3: Sink token analysis across transformer layers for DLMs initialized from Qwen2.5-0.5B and Qwen2.5-1.5B. (a) and (c): mean value-space $L_2$ norm per transformer layer, comparing the extra sink token against all other tokens. (b) and (d): the proportion of attention mass received by the sink token per layer. After the initial few layers and before the final layers, across the intermediate transformer layers that are most closely related to information integration and transformation, the sink token exhibits substantially smaller value-space $L_2$ norms than other tokens while receiving a large share of global attention.