하나의 토큰이면 충분하다: 싱크 토큰을 활용한 디퓨전 언어 모델 개선

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

디퓨전 언어 모델(DLM)은 자동 회귀 방식의 대안으로 부상하며, 경쟁력 있는 성능으로 병렬 텍스트 생성을 가능하게 합니다. 그러나 DLM에는 '이동하는 싱크 현상'이라는 중요한 불안정성이 존재합니다. 저희의 분석에 따르면, 싱크 토큰은 트랜스포머의 값 공간에서 낮은 정규화 값을 가지며, 이러한 이동하는 싱크 현상은 DLM에서 과도한 정보 혼합을 방지하는 보호 메커니즘 역할을 합니다. 하지만 이들의 위치가 디퓨전 과정에서 예측 불가능하게 변화하기 때문에 추론의 안정성을 저해합니다. 이러한 문제를 해결하기 위해, 저희는 수정된 어텐션 마스크를 통해 구현된 간단하지만 효과적인 추가 싱크 토큰을 제안합니다. 구체적으로, 이 특별한 토큰은 오직 자신에게만 어텐션을 수행하도록 제한되지만, 동시에 다른 모든 토큰에게는 전역적으로 노출됩니다. 실험 결과, 하나의 추가적인 토큰을 도입함으로써 어텐션 싱크를 안정화시키고 모델 성능을 크게 향상시킬 수 있음을 확인했습니다. 더욱 중요한 점은, 추가 분석을 통해 이 토큰의 효과가 그 위치에 독립적이며, 의미론적 내용이 미미하다는 것을 확인하여, 이 토큰이 견고하고 헌신적인 구조적 싱크 역할을 한다는 것을 검증했습니다.

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.