확산 언어 모델을 위한 싱크 인지 프루닝

Sink-Aware Pruning for Diffusion Language Models

확산 언어 모델(DLM)은 반복적인 노이즈 제거 과정으로 인해 높은 추론 비용을 발생시키며, 이는 효율적인 프루닝(pruning)의 필요성을 동기부여한다. 자기회귀(AR) 대형 언어 모델(LLM)에서 주로 차용한 기존의 프루닝 휴리스틱은 AR 싱크(sink)가 안정적인 전역 앵커(global anchor) 역할을 하기 때문에 일반적으로 어텐션 싱크(attention sink) 토큰을 보존한다. 우리는 이러한 가정이 DLM에는 적용되지 않음을 보여준다. 전체 생성 궤적에 걸쳐 어텐션 싱크 위치는 (주요 싱크 위치가 타임스텝에 따라 어떻게 이동하는지 측정했을 때) 상당히 더 높은 분산을 보이며, 이는 AR 모델에 비해 싱크가 종종 일시적이고 구조적으로 덜 필수적임을 나타낸다. 이러한 관찰을 바탕으로, 우리는 DLM에서 불안정한 싱크를 자동으로 식별하고 프루닝하는 ${\bf \texttt{Sink-Aware Pruning}}$을 제안한다 (기존 연구들은 보통 AR LLM을 위해 싱크를 유지한다). 재학습 없이도 우리의 방법은 더 나은 품질-효율성 트레이드오프를 달성하며, 동일한 연산량 조건 하에서 기존의 강력한 프루닝 베이스라인보다 뛰어난 성능을 보인다. 코드는 https://github.com/VILA-Lab/Sink-Aware-Pruning 에서 확인할 수 있다.

Diffusion Language Models (DLMs) incur high inference cost due to iterative denoising, motivating efficient pruning. Existing pruning heuristics largely inherited from autoregressive (AR) LLMs, typically preserve attention sink tokens because AR sinks serve as stable global anchors. We show that this assumption does not hold for DLMs: the attention-sink position exhibits substantially higher variance over the full generation trajectory (measured by how the dominant sink locations shift across timesteps), indicating that sinks are often transient and less structurally essential than in AR models. Based on this observation, we propose ${\bf \texttt{Sink-Aware Pruning}}$, which automatically identifies and prunes unstable sinks in DLMs (prior studies usually keep sinks for AR LLMs). Without retraining, our method achieves a better quality-efficiency trade-off and outperforms strong prior pruning baselines under matched compute. Our code is available at https://github.com/VILA-Lab/Sink-Aware-Pruning.