One Token Is Enough: Improving Diffusion Language Models with a Sink Token
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Commonsense Reasoning | HellaSwag | Accuracy61.31 | 1460 | |
| Question Answering | ARC Challenge | Accuracy22.35 | 749 | |
| Commonsense Reasoning | PIQA | Accuracy68.17 | 647 | |
| Question Answering | ARC Easy | Normalized Acc40.91 | 385 | |
| Mathematical Reasoning | GSM8K | Accuracy (GSM8K)58.45 | 358 | |
| Physical Commonsense Reasoning | PIQA | Accuracy62.24 | 329 | |
| Question Answering | ARC-E | Accuracy68.18 | 242 | |
| Language Modeling | LAMBADA | Accuracy66.41 | 183 | |
| Question Answering | ARC-C | Accuracy43.43 | 166 | |
| Reading Comprehension | RACE | Accuracy37.61 | 151 |