[논문] Sink-Aware Pruning for Diffusion Language Models

Source: arXiv - 2602.17664v1

(번역할 텍스트가 제공되지 않았습니다. 번역이 필요한 내용을 추가로 알려주시면 도와드리겠습니다.)

개요

Diffusion Language Models (DLMs)는 고품질 텍스트 생성에 가능성을 보여주었지만, 반복적인 디노이징 과정 때문에 추론 비용이 많이 듭니다. 새로운 논문 **“Sink‑Aware Pruning for Diffusion Language Models”**는 이러한 모델을 프루닝하는 방식에 숨겨진 비효율성을 밝혀내고, 출력 품질을 유지하면서 불필요한 부분을 제거하는 경량화된, 재학습이 필요 없는 솔루션을 제안합니다.

주요 기여

• Empirical discovery: DLM에서 “attention‑sink” 토큰(다른 토큰들이 가장 많이 주목하는 토큰)은 불안정하며, 이는 autoregressive (AR) LLM의 안정적인 전역 앵커와 다릅니다.
• Sink‑Aware Pruning algorithm: 추가 fine‑tuning 없이 이러한 변동성이 큰 sink 토큰을 자동으로 감지하고 가지치기하는 체계적인 방법.
• Better quality‑efficiency trade‑off: 동일한 compute budget에서 기존 pruning baseline(크기 기반 및 구조 기반)보다 우수한 성능을 보여줍니다.
• Open‑source implementation: 전체 코드를 공개하여 커뮤니티가 결과를 재현하고 자신들의 diffusion‑based language model에 적용할 수 있도록 합니다.

Methodology

1. Analyzing sink stability – 저자들은 DLM의 모든 디노이징 타임스텝에서 지배적인 attention‑sink 토큰을 추적합니다. 그들은 sink 위치가 얼마나 자주 변하는지를 정량화하는 variance score를 계산합니다. 높은 분산은 일시적인 sink로, 신뢰할 수 있는 전역 컨텍스트 역할을 하지 않음을 나타냅니다.
2. Identifying prune‑worthy heads – variance score를 사용하여 attention head(또는 전체 레이어)를 sink가 “불안정”한 정도에 따라 순위 매깁니다. 지속적으로 변하는 sink를 가리키는 head는 낮은 유용성으로 표시됩니다.
3. Pruning without retraining – 표시된 head는 모델의 forward pass에서 단순히 0으로 만들거나(또는 제거)합니다. diffusion 모델은 이미 일정 수준의 노이즈를 견디므로, 이러한 공격적인 pruning은 비용이 많이 드는 사후 fine‑tuning을 필요로 하지 않습니다.
4. Evaluation protocol – 저자들은 pruning된 모델을 표준 언어 생성 작업(예: 스토리 연속, 요약)에서 벤치마크하고, perplexity, BLEU/ROUGE 점수, 그리고 wall‑clock inference 시간을 pruning되지 않은 베이스라인 및 기존 pruning 방법과 비교합니다.

결과 및 발견

• 분산 분석 결과 DLM의 70 % 이상인 어텐션 헤드가 타임스텝 간에 sink 위치가 세 포지션 이상 이동한다는 것이 확인되어 불안정성 가설을 뒷받침합니다.
• Sink‑Aware Pruning은 BLEU/ROUGE 감소폭이 더 작아 일관되게 높은 생성 품질을 유지하면서 테스트된 방법 중 가장 큰 속도 향상을 제공합니다.
• 이 접근법은 즉시 사용 가능합니다: 추가 학습 에포크, 하이퍼파라미터 탐색, 혹은 데이터 의존적인 보정이 전혀 필요하지 않습니다.

실용적 시사점

• 프로덕션 서비스의 빠른 추론 – 확산 기반 챗봇이나 텍스트‑투‑코드 어시스턴트를 배포하는 기업은 품질 손실을 최소화하면서 지연 시간을 약 30‑40 % 줄일 수 있어 클라우드 비용 절감으로 직접 연결됩니다.
• 엣지 배포 – 파라미터 수와 연산량이 감소함으로써 이전에 반복적인 디노이징이 제한적이었던 리소스가 제한된 디바이스(예: 스마트폰, IoT 게이트웨이)에서 DLM을 실행하는 것이 가능해집니다.
• 모델 유지보수 간소화 – 프루닝이 정적이며 파인‑튜닝을 필요로 하지 않으므로 팀은 CI 파이프라인에 통합할 수 있습니다: 한 번 프루닝하고, 다듬어진 바이너리를 배포하며, 지속적인 재학습 오버헤드를 피합니다.
• 향후 모델 설계에 대한 가이드 – DLM이 안정적인 전역 앵커를 갖지 못한다는 발견은 아키텍처 연구가 본질적으로 프루닝에 더 친화적인 대안적인 어텐션 메커니즘(예: 동적 라우팅)을 탐색할 수 있음을 시사합니다.

제한 사항 및 향후 연구

• 평가 범위 – 실험은 영어 중심 벤치마크에 초점을 맞추고 있으며, 다언어 또는 도메인‑특화 DLM은 다른 sink 동역학을 보일 수 있습니다.
• 세분화 수준 – 현재 방법은 헤드 수준에서 프루닝을 수행합니다; 보다 미세한(예: 토큰‑단위) 프루닝은 추가적인 이득을 제공할 수 있지만 탐색되지 않았습니다.
• 다른 압축 기법과의 상호 작용 – Sink‑Aware Pruning이 양자화, 지식 증류, 혹은 저‑랭크 팩터화와 어떻게 결합되는지는 아직 미해결 질문입니다.
• 이론적 이해 – 경험적 분산이 견고한 프록시이지만, diffusion attention sink가 불안정한 이유에 대한 보다 깊은 이론 모델은 보다 원칙적인 프루닝 기준을 제시할 수 있습니다.

핵심 요점: diffusion 언어 모델이 자동 회귀 모델과 같은 “끈적이는” 어텐션 앵커를 필요로 하지 않는다는 점을 인식함으로써, 저자들은 실용적이고 플러그‑인‑플레이 방식의 프루닝 기법을 제공하여 품질 손실을 최소화하면서 추론 속도를 높입니다—diffusion 기반 텍스트 생성을 실제 제품에 적용하려는 모든 사람에게 매력적인 이점이 됩니다.

저자

• Aidar Myrzakhan
• Tianyi Li
• Bowei Guo
• Shengkun Tang
• Zhiqiang Shen

논문 정보

• arXiv ID: 2602.17664v1
• 카테고리: cs.CL, cs.AI, cs.LG
• 출판일: 2026년 2월 19일
• PDF: Download PDF