기술 요약: Autonomy-of-Heads (AoH)

문제 정의

대규모 언어 모델(LLM)의 긴 문맥 추론(Long-context inference)은 현재 어텐션 메커니즘의 이차 함수적(quadratic) 계산 비용과 키-값(KV) 캐시의 선형적으로 증가하는 메모리 요구량에 의해 병목 현상을 겪고 있습니다. 희소 어텐션(sparse attention) 및 KV 압축 방법과 같은 기존 솔루션들은 일반적으로 어떤 토큰이나 헤드를 유지할지 결정하기 위해 입력 의존적(input-dependent) 신호에 의존합니다. 여기에는 런타임 어텐션 점수(예: Heavy-Hitter), 관찰 윈도우, 캘리브레이션 프롬프트 또는 학습된 게이팅 메커니즘 등이 포함됩니다. 이러한 방식들은 효과적이긴 하지만, 런타임 관찰, 추가 학습 또는 캘리브레이션 절차를 필요로 하므로 배포 비용이 높고 특정 입력 동작에 종속적이라는 단점이 있습니다.

저자들은 동결된 가중치(frozen weights) 그 자체만으로 런타임 점수, 캘리브레이션 프롬프트 또는 추가 학습 없이도 어텐션 헤드의 기능을 결정하는 유용한 사전 지식(prior)을 제공할 수 있는지 질문을 던집니다.

방법론: Autonomy-of-Heads (AoH)

AoH는 동결된 쿼리-키 투영(Query-Key projection)의 스펙트럼 기하학(spectral geometry)으로부터 직접 "리트리벌 헤드(retrieval heads, 전역적 문맥 필요)"와 "스트리밍 헤드(streaming heads, 로컬/싱크 토큰으로 충분)"를 식별하는 데이터가 필요 없고 학습이 필요 없는(data-free and training-free) 방법입니다.

1. 이론적 기초

이 방법은 각 헤드 $h$에 대한 **커널 어텐션 연산자(kernel attention operator)**를 다음과 같이 정의합니다:
$$M_h = W_K^{h\top} W_Q^h \in \mathbb{R}^{d_{model} \times d_{model}}$$
이 행렬은 헤드 $h$에 특화된 동결된 쿼리-키 매칭 방향을 나타냅니다. 저자들은 $M_h$의 스펙트럼 기하학이 헤드의 기능적 역할을 인코딩한다고 가설을 세웠습니다:

• 집중된 스펙트라(Concentrated Spectra): 소수의 지배적인 쿼리-키 매칭 방향을 나타내며, 이는 특정 문맥을 찾기 위해 전역적으로 탐색해야 하는 리트리벌 헤드와 관련이 있습니다.
• 확산된 스펙트라(Diffuse Spectra): 지배적인 전역 매칭 방향의 부재를 나타내며, 싱크 토큰(sink tokens)과 최근 윈도우에 집중하는 스트리밍 헤드와 관련이 있습니다.

2. 분류기로서의 유효 순위(Effective Rank)

AoH는 헤드 기능의 가중치 공간 측정치로서 $M_h$의 정규화된 특이값(singular values)에 대한 섀넌 엔트로피(Shannon entropy)의 지수값인 **유효 순위(Effective Rank, ER)**를 사용합니다:
$$\text{eff\_rank}(h) = \exp\left(-\sum_{k=1}^{r_h} \hat{\sigma}_k \log \hat{\sigma}_k\right)$$

• 낮은 ER: 집중된 스펙트라를 시사함 $\rightarrow$ 리트리벌 헤드 (전체 전역 어텐션 할당).
• 높은 ER: 확산된 스펙트라를 시사함 $\rightarrow$ 스트리밍 헤드 (제한된 싱크 및 최근 윈도우 어텐션 할당).

3. 효율적인 계산

전체 $d_{model} \times d_{model}$ 행렬 $M_h$를 구축하는 것은 계산 비용이 많이 듭니다. 저자들은 효율적인 $d_{head}$ 차원의 계산법을 도출했습니다. 실베스터의 행렬식 정리(Sylvester's determinant theorem)를 활용하여, $M_h$의 0이 아닌 특이값들을 더 작은 프록시 행렬 $C_h$로부터 계산할 수 있습니다:
$$C_h = (W_Q^h W_Q^{h\top})(W_K^h W_K^{h\top}) \in \mathbb{R}^{d_{head} \times d_{head}}$$
이를 통해 복잡도를 $O(d_{model}^2)$에서 $O(d_{head}^2 \cdot d_{model})$로 줄여, 헤드 레이블의 사전 계산을 가능하게 합니다.

4. 배포 전략

• 헤드 분류: 추론 전에 ER에 따라 헤드 순위를 매깁니다. 상위 $k$개의 헤드(가장 낮은 ER)를 리트리벌 헤드로 지정하고, 나머지는 스트리밍 헤드로 지정합니다.
• KV-캐시 관리:
  • 리트리벌 헤드: 전체 전역 KV 상태를 유지합니다.
  • 스트리밍 헤드: 고정된 크기의 "싱크(sink)" 캐시(예: 처음 128개 토큰)와 "최근(recent)" 윈도우(예: 마지막 256개 토큰)만을 유지합니다.
• 구현: 이 방법은 그룹 쿼리 어텐션(GQA)과 호환되며, 헤드 수준의 점수를 그룹 수준으로 집계(평균 집계 사용)합니다. 또한 특수한 커널을 요구하지 않고 표준 청크 프리필(chunked prefill) 및 FlashAttention-2와 통합됩니다.

주요 기여

1. 데이터 프리 헤드 선택: 동결된 쿼리-키 기하학으로부터 캘리브레이션 데이터나 런타임 관찰 없이 직접 리트리벌 및 스트리밍 헤드를 식별하는 방법인 AoH를 도입했습니다.
2. 유효 순위 분류기: $M_h$의 유효 순위에 기반한 분류기와 프록시 행를 통한 스펙트럼 계산 알고리즘을 개발했습니다.
3. 희소 어텐션 정책: 정확도를 유지하면서 레이턴시와 메모리를 크게 줄이는 단순하고 학습이 필요 없는 희소 어텐션 정책으로서 AoH를 구현했습니다.

실험 결과

저자들은 여러 모델(Qwen2.5-7B, Qwen3-8B, Llama3.1-8B-Instruct)에 대해 LongBench에서 AoH를 평가했습니다.

• 정확도: 50% 희소성(50%의 헤드를 리트리벌로 설정)에서 AoH는 평균적으로 Full Attention 성능의 **96.5%**를 유지합니다.
  • Qwen3-8B의 경우, AoH는 Full Attention보다 단 0.32포인트 뒤처집니다.
  • Llama3.1-8B-Instruct의 경우, 격차는 1.10포인트입니다.
  • AoH는 H2O, SnapKV, Quest, DuoAttention과 같은 베이스라인은 물론, 무작위(random) 및 역순(reversed) 헤드 선택 절제 실험보다 일관되게 우수한 성능을 보였습니다.
• 효율성:
  • 레이턴시: 256K 문맥 길이에서 75% 희소성일 때, AoH는 Llama3.1-8B 기준으로 Full Attention 대비 최대 3.24배의 프리필(prefill) 속도 향상과 9.14배의 디코드(decode) 속도 향상을 달 achieves 합니다.
  • 메모리: 75% 희소성에서 KV-캐시 메모리 사용량을 최대 3.98배 줄이며, 256K 토큰에서 50% 희소성일 때 **50.0%**까지 절감합니다.

의의 및 주장

본 논문은 **헤드가 무엇을 아는지 이미 알고 있다(heads "know what they know")**고 주장합니다. 즉, 헤드의 기능적 역할(리트리벌 vs 스트리밍)이 동결된 가중치의 스펙트럼 기하학에 인코딩되어 있다는 것입니다. AoH는 희소 어텐션을 위한 단순하고 확장 가능하며 학습이 필요 없는 사전 지식을 제공합니다.

• 독립성: 기존 연구와 달리, AoH는 입력 의서적 신호, 학습된 게이트 또는 캘리브레이션 프롬프트에 의존하지 않습니다. 이는 런타임 관찰 이후가 아니라 첫 번째 토큰(프리필)부터 희소 어텐션 전략을 적용할 수 있게 합니다.
• 호환성: GQA를 포함한 기존 아키텍처 및 표준 추론 엔진인 FlashAttention과 호환되므로 실제 배포에 실용적입니다.
• 강건성: 저자들은 유효 순위(effective-rank) 순위가 다양한 문맥 길이에 대해 안정적이며, RoPE(Rotary Positional Embeddings) 회전을 고려하더라도 매우 일관되게 유지됨을 입증했습니다.

결론적으로, AoH는 추가적인 학습이나 복잡한 런타임 휴리스틱의 오버헤드 없이 긴 문맥 추론의 효율성을 개선할 수 있는 실행 가능한 경로를 제공합니다.