프리스코어링으로 효율적인 어텐션: 핵심 키 우선순위 지정

초록

본 논문은 쿼리와 무관한 전역 중요도 사전점수를 키에 부여한 뒤, 계층적 근사 어텐션에 적용하는 프리스코어링 프레임워크를 제안한다. K‑means 기반 클러스터링 또는 레버리지 스코어를 이용해 구조적으로 중요한 키를 선별하고, 이를 HyperAttention과 결합해 동일한 연산 예산 하에서 퍼플렉시티를 크게 낮춘다. 언어 모델(ChatGLM, 131k 토큰)과 비전 트랜스포머(ViT‑Large) 실험에서 클러스터링 기반 선택이 레버리지 기반보다 일관되게 우수함을 보이며, 플랜티드 서브스페이스 모델 하에서 이론적 회복 보장을 제공한다.

상세 분석

프리스코어링은 기존 효율 어텐션이 쿼리‑의존적인 로컬리티(LSH, 블록 분할 등)에만 의존해 전역적으로 중요한 토큰을 놓치는 문제를 보완한다. 핵심 아이디어는 키 행렬 K에 대해 쿼리와 무관하게 전역 중요도 점수를 사전에 계산하고, 이 점수가 높은 키만을 “우선키 집합” S에 포함시킨 뒤, 이후 단계에서 HyperAttention의 해시 기반 블록 매칭을 수행한다. 두 가지 점수 산출 방식이 제안되는데, 첫 번째는 K‑means(또는 K‑median) 클러스터링을 이용해 d+1개의 클러스터(임베딩 차원 d와 잡음 전용 하나)를 만든 뒤, 각 클러스터 중심에 가장 가까운 s개의 키를 선택한다. 이는 키 공간의 기하학적 구조를 직접 활용해 고중량 키를 효과적으로 포착한다. 두 번째는 LevScore와 유사한 레버리지 스코어를 빠르게 근사해 상위 s개의 키를 선택하는 방식이다. 실험적으로는 동일한 키 예산(s) 하에서 클러스터링이 레버리지보다 높은 재현율을 보이며, 특히 ViT‑Large에서 128개의 선택된 키가 원본 정확도의 84.46%를 유지해 이미지 분류 성능 저하가 최소임을 확인했다.

이 프레임워크를 HyperAttention에 통합한 결과, “Pre‑Score + HyperAttention” 조합이 기존 HyperAttention(퍼플렉시티 17.54) 대비 9.5까지 퍼플렉시티를 감소시켰으며, 프리스코어링만 적용해도 10.38으로 큰 개선을 이끌어냈다. 이는 프리스코어링이 연산 예산을 고정하면서도 고중량 키를 효과적으로 회수한다는 증거다. 또한, 플랜티드 서브스페이스 모델을 도입해 이론적 보장을 제공한다. 모델에서는 d개의 신호 클러스터와 잡음 클러스터를 가정하고, 각 클러스터 내 분산은 O(σ²)이며 클러스터 간 거리 Δ는 Ω(1)이다. 이 가정 하에 K‑means가 모든 ϵ‑heavy 키를 확률적으로 1−exp(−c·n) 수준으로 회수함을 정리 4.4, 4.5에서 증명한다. 레버리지 스코어 역시 동일한 회수율을 보이지만, 클러스터링은 실제 데이터에서 더 강건한 구조적 정보를 활용한다는 점에서 실용적 우위를 가진다.

연산 복잡도 측면에서 프리스코어링은 O(ndk·k·I) (I는 클러스터링 반복 횟수, 보통 ≤10) 혹은 O(ndk·log dk) (레버리지) 정도로, 전체 어텐션 레이어당 한 번만 수행된다. 역전파는 클러스터링 파라미터를 학습하지 않으므로 추가 비용이 없으며, 자동 회귀 디코딩 시에는 미리 계산된 S를 재사용하거나 일정 주기마다만 업데이트해 O(n) 비용을 회피한다.

제한점으로는 키 벡터의 ℓ₂ 정규화가 필요하다는 점이다. 논문 부록에서는 정규화되지 않은 키가 클러스터링을 방해하는 사례를 제시하고, 실제 구현에서는 LayerNorm/RMSNorm 후 ℓ₂ 정규화를 적용해 이 문제를 해결한다. 또한, 현재는 Softmax 어텐션에 대한 이론적 보장은 제공되지 않으며, 실험적 결과에 의존한다는 점이 향후 연구 과제로 남는다.

전반적으로 프리스코어링은 쿼리‑독립적인 전역 중요도 사전선택과 기존 쿼리‑의존적 로컬 어텐션을 결합함으로써, 연산 효율성을 유지하면서도 전역적인 정보 손실을 최소화하는 실용적인 접근법이다.