새로운 “어텐션 싱크 디텍터(Attention Sink Detector)”는 거대 언어 모델(LLM) 추론의 프리필(prefill) 단계에서 어텐션 질량(attention mass)의 대부분을 차지하는 토큰을 찾아냅니다. 이 도구는 KV 캐시에서 이러한 싱크 토큰을 제거하면 모델 성능이 급격히 저하된다는 것을 보여주며, 이 토큰들이 안정적인 생성을 위한 필수적인 앵커(anchor)임을 입증합니다.

프리필 단계가 중요한 이유

대부분의 LLM 연구는 토큰 단위의 생성 루프에 집중하지만, 첫 번째 토큰이 나오기 전 단계인 프리필(prefill) 작업이 이후의 모든 과정을 결정합니다. 프리필 단계에서 모델은 전체 프롬프트를 처리하고, KV 캐시를 구축하며, 모든 위치에 어텐션을 분산시킵니다. 어텐션 가중치를 생성하는 소프트맥스(softmax)의 합은 반드시 1이어야 하므로, 모델은 "관련 있는 것이 아무것도 없다"라고 말할 수 없습니다. 따라서 모델은 남은 확률 질량을 처리하기 쉬운 토큰, 대개 시퀀스의 첫 번째 토큰에 몰아넣습니다. 이 토큰이 바로 **어텐션 싱크(attention sink)**가 됩니다.

어텐션 싱크란 무엇인가 (쉬운 설명)

트랜스포머(Transformer)에서 각 헤드는 모든 토큰 쌍에 대해 가중치를 계산합니다. 분포가 심하게 왜곡될 경우, 단일 토큰이 어텐션 가중치의 불균형적으로 큰 비중을 차지할 수 있습니다. 이 현상은 버그가 아니라 소프트맥스 제약 조건에서 비롯된 결과입니다. 첫 번째 토큰(종종 문장 시작 마커)은 다른 곳에 할당할 수 없는 잔여 확률을 받아내는 "방출 밸브" 역할을 합니다.

KV 캐시 관리의 중요성

KV 캐시는 처리된 모든 토큰의 키(key)와 값(value) 벡터를 저장하여 생성 과정에서 빠른 조회를 가능하게 합니다. 긴 컨텍스트 시나리오에서 실무자들은 GPU 메모리 제한을 준수하기 위해 캐시를 프루닝(pruning)합니다. 새로운 디텍터는 "중요하지 않아" 보이는 토큰을 무분별하게 삭제하면 모델이 의존하는 싱크 토큰까지 함께 삭제되어 성능이 급락한다는 것을 보여줍니다. 캐시에 싱크 토큰을 고정(pinning)하면 모델의 내부 평형을 유지하고 생성을 안정적으로 유지할 수 있습니다.

디텍터의 작동 원리

  • Eager attention: 전체 어텐션 행렬을 압축하는 FlashAttention과 같은 빠른 커널을 우회하여, 각 헤드에 대한 가공되지 않은(raw) 어텐션 점수를 기록합니다.
  • Layer-wide aggregation: 모든 레이어와 헤드에 걸쳐 점수를 평균하여 토큰당 단일 어텐션 질량 프로필을 생성합니다.
  • Log-median absolute deviation (MAD): 어텐션 가중치는 오른쪽으로 심하게 치우쳐 있기 때문에, 단순한 평균-표준편차 테스트를 사용하면 일반적인 변동을 이상치(outlier)로 오분류할 수 있습니다. 가중치에 로그 변환을 적용하여 분포를 정규화한 후, MAD를 적용하여 어텐션 질량이 일반적인 범위를 초과하는 토큰을 찾아냅니다.

밝혀진 두 가지 싱크 클래스

  1. True sinks (진정한 싱크) – 프롬프트 내용과 관계없이 문장 시작(BOS) 토큰이 지속적으로 막대한 양의 어텐션을 흡수합니다.
  2. Structural sinks (구조적 싱크) – ChatML에서 사용되는 마커(<im_start>, <im_end>)와 같은 시스템 레벨 프롬프트에 속하는 토큰들은 어텐션을 끌어당기는 작은 클러스터를 형성합니다. 이들은 논리적 경계 역할을 하여 모델이 사용자 메시지와 시스템 지침을 구분하도록 돕습니다.

반론: 정말로 가공되지 않은 행렬이 필요한가?

저자는 가공되지 않은 수치(raw numbers) 없이는 싱크 현상이 숨겨지게 되며, 불완전한 데이터에 기반한 어떠한 캐시 프루닝 정책도 모델을 불안정하게 만들 위험이 있다고 지적합니다.

향후 주목할 점

이 디텍터는 생성 품질을 위한 엔트로피 트래커로 시작된 4부작 시리즈 중 두 번째 작업입니다. 다음 파트에서는 싱크 탐지를 **엔트로피 가이드 기반 투기적 디코딩(entropy-guided speculative decoding)**과 통합할 예정입니다. 마지막 연구는 실증적 연구가 될 것입니다.

요약: 어텐션 싱크는 모호한 특이 현상이 아니라, 프리필 단계 동안 트랜스포머의 어텐션 분포를 안정적으로 유지하는 구조적 지주입니다. 이를 무시하는 어떠한 KV 캐시 프루닝 전략도 모델의 안정성을 해칠 것입니다. 이러한 토큰을 탐지하고 보존하는 것은 긴 컨텍스트 추론을 신뢰할 수 있고 효율적으로 만들 수 있는 저비용의 안전장치입니다.