TriAttention: 삼각 함수 기반 KV 압축을 통한 효율적인 긴 문맥 추론

LLM이 복잡한 추론을 수행할 때 발생하는 방대한 KV 캐시 메모리 병목 문제를 해결하는 새로운 압축 기법이다. 기존 방식과 달리 RoPE 회전 이전의 안정적인 벡터 공간을 활용하여, 메모리 사용량을 획기적으로 줄이면서도 추론의 정확도를 유지할 수 있음을 증명했다.

Transformer 모델은 문장이 길어질수록 과거의 정보를 저장하는 KV 캐시가 기하급수적으로 커져 메모리 부족 문제를 일으킨다. 기존의 압축 방식은 최근의 Attention 점수를 보고 중요도를 판단하지만, RoPE라는 위치 인코딩 기술 때문에 Query 벡터가 위치마다 계속 회전하므로 관찰 창이 좁아지면 중요한 정보를 놓치기 쉽다는 한계가 있다.

이 논문은 RoPE 회전이 적용되기 전인 'Pre-RoPE' 공간에 주목했다. 분석 결과, 이 공간에서 Query와 Key 벡터들은 무작위로 흩어져 있는 것이 아니라 특정 방향(중심점)을 향해 뭉쳐 있는 'Q/K Concentration' 현상을 보였다. 이 중심점들은 모델이 학습 과정에서 획득한 고유한 특성으로, 어떤 거리의 토큰을 선호하는지를 결정하는 유전 정보와 같다.

연구진은 이 고정된 중심점들을 삼각 함수 공식에 대입하면, 미래에 어떤 위치에서 Query가 들어오더라도 특정 거리(예: 아주 가까운 곳 혹은 아주 먼 곳)에 있는 토큰이 얼마나 많은 관심을 받을지 수학적으로 예측할 수 있음을 발견했다. 즉, 현재의 Attention 점수만 보는 것이 아니라 모델의 내재된 선호도를 파악해 미래에도 계속 중요할 토큰을 정확히 골라낼 수 있게 된 것이다.

TriAttention의 핵심은 오프라인 캘리브레이션과 온라인 중요도 산출의 결합이다. 먼저 소량의 데이터셋을 사용하여 각 Attention Head별로 Pre-RoPE 공간에서의 Query 중심점 E[q]와 Key 중심점 E[k]를 미리 계산한다. 이 중심점들은 모델의 내재적인 거리 선호도를 나타내는 지표로 활용된다.

실제 추론 시에는 두 가지 신호를 결합하여 각 토큰(Key)의 점수를 매긴다. 첫 번째는 삼각 급수 점수 Strig로, [Query 중심점과 현재 Key 벡터를 입력으로] → [삼각 함수 기반의 내적 근사 연산을 수행해] → [미래의 특정 거리 Δ에서 받을 예상 Attention 값을 얻고] → [해당 토큰이 구조적으로 얼마나 중요한지]를 판단한다. 두 번째는 Norm 기반 점수 Snorm으로, 벡터의 크기가 클수록 정보량이 많다는 가정을 반영한다.

최종 점수 S는 Q/K 응집도 지표인 Mean Resultant Length(R)를 가중치로 사용하여 두 점수를 적응적으로 혼합한다. [응집도 R이 1에 가까우면] → [삼각 급수 점수에 높은 가중치를 부여하고] → [R이 낮으면] → [Norm 기반 점수의 비중을 높여] → [예측의 신뢰도를 보정한다]. 이후 128토큰마다 윈도우 단위로 점수가 낮은 토큰을 캐시에서 제거하여 메모리 사용량을 일정하게 유지한다.

Qwen3-8B 모델을 사용한 AIME25 수학 추론 벤치마크에서, TriAttention은 Full Attention의 정확도(40.8%)를 그대로 유지하면서 KV 캐시 메모리를 10.7배 줄이는 데 성공했다. 이는 기존의 선도적 기법인 R-KV가 동일한 효율성에서 절반 수준의 정확도만 보인 것과 대조적이다.

처리량 측면에서도 괄목할만한 성장을 보였다. 동일한 정확도 기준에서 TriAttention은 Full Attention 대비 2.5배 높은 토큰 생성 속도를 기록했다. 특히 MATH 500 데이터셋에서는 KV 캐시를 단 1,024개로 제한했음에도 불구하고 Full Attention(69.6%)에 근접한 68.4%의 정확도를 달성했다.

추가로 진행된 '재귀적 상태 쿼리(Recursive State Query)' 실험을 통해, TriAttention이 복잡한 추론 과정에서 중간 단계의 정보를 잃지 않고 끝까지 유지하는 능력이 기존 방식보다 월등히 뛰어남을 입증했다. 이는 긴 체인 오브 쏘트(CoT)가 필요한 추론 작업에서 TriAttention의 안정성을 뒷받침한다.

TriAttention은 RoPE의 수학적 구조를 분석하여 Attention Logit이 Pre-RoPE 벡터들의 위상 차이에 의한 삼각 급수로 표현될 수 있음을 이론적으로 증명했다. 구체적으로 logit(Δ) ≈ Σ ||q_f|| ||k_f|| cos(ω_f Δ + φ_f) 형태의 공식을 도출했으며, 여기서 ω_f는 RoPE의 주파수, Δ는 토큰 간 거리를 의미한다.

이 기법의 차별점은 '정적 중심점'의 활용이다. 기존의 H2O나 SnapKV가 동적인 Attention Score에 의존하여 '관찰 창(Observation Window)'의 한계에 갇혔던 것과 달리, TriAttention은 모델의 가중치에 내재된 통계적 중심을 사용함으로써 훨씬 긴 범위의 미래 Attention 패턴을 안정적으로 예측한다.

구현 측면에서는 GQA(Grouped-Query Attention) 구조를 지원하기 위해 각 Query Head의 점수를 Z-score로 정규화한 뒤 최대값(Max-pooling)을 취하는 방식을 채택하여 여러 Query가 공유하는 Key의 중요도를 공정하게 평가한다. 또한 128토큰 단위의 배치형 프루닝(Window-based Pruning)을 통해 매 토큰마다 계산이 발생하는 오버헤드를 최소화했다.

현재 구현은 범용적인 최적화 수준이며, TriAttention의 고유한 연산 구조에 최적화된 하드웨어 가속 커널(Dedicated Inference Kernel)이 개발된다면 더 높은 속도 향상을 기대할 수 있다. 또한 현재는 수학 및 코딩 도메인 위주로 검증되었으므로 더 넓은 일반 도메인에 대한 추가 검증이 필요할 수 있다.