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

TL;DR

This paper introduces a simple method to stabilize diffusion language models by adding a single dedicated sink token, which improves performance and robustness by controlling information flow during text generation.

Contribution

The paper proposes a novel approach of adding an extra sink token with a modified attention mask to address sink token instability in diffusion language models.

Findings

Introducing a sink token stabilizes attention sinks.

The sink token improves model performance significantly.

The sink token's effectiveness is position-independent and semantically negligible.

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.