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

대규모 언어 모델이 긴 문장을 추론할 때 발생하는 메모리 병목 현상을 획기적으로 해결하는 기술입니다. 기존 방식과 달리 데이터의 위치 정보를 삼각 함수로 예측하여 중요한 정보만 남김으로써, 일반 소비자용 그래픽 카드 한 장으로도 수만 개의 단어가 포함된 긴 대화를 끊김 없이 처리할 수 있게 합니다.

Transformer 모델은 문장이 길어질수록 과거의 정보를 저장하는 KV 캐시가 무한정 커져 메모리 부족 문제를 일으킨다. 기존에는 최근 데이터만 보고 중요도를 판단했지만, 이는 위치에 따라 벡터를 회전시키는 RoPE(Rotary Position Embedding) 특성상 관찰 창이 너무 좁아 중요한 정보를 실수로 삭제하는 불안정성을 초래한다.

연구진은 벡터가 회전하기 전인 'Pre-RoPE' 공간을 분석한 결과, 대부분의 어텐션 헤드에서 쿼리와 키 벡터가 특정 방향으로 쏠려 있는 'Q/K 집중 현상'을 발견했다. 이 집중된 중심점들은 RoPE의 회전 공식과 결합될 때 특정 거리(예: 아주 가까운 거리나 아주 먼 거리)에 있는 토큰을 선호하게 만드는 일종의 '거리별 선호도 곡선'을 형성한다.

TriAttention은 이 원리를 이용해 실제 추론이 일어나기 전에 어떤 거리의 토큰이 중요할지를 삼각 함수(Trigonometric Series)로 미리 계산한다. 마치 안개 속에서 전체 지형을 다 보는 대신, 과거의 통계적 중심점이라는 나침반을 이용해 어떤 방향에 중요한 정보가 있을지 예측하는 것과 같다. 이를 통해 아주 적은 양의 메모리만 남기고도 모델의 논리적 사고 흐름(Chain-of-Thought)을 깨뜨리지 않고 유지할 수 있다.

TriAttention의 핵심은 오프라인 캘리브레이션을 통해 각 어텐션 헤드의 Q/K 중심점을 구하고, 이를 기반으로 온라인 추론 시 키의 중요도를 채점하는 것이다. 전체 스코어 S(k, Δ)는 삼각 함수 기반 점수(Strig)와 노름 기반 점수(Snorm)의 합으로 구성된다.

삼각 함수 점수 Strig는 쿼리의 중심점 E[qf]와 현재 캐시된 키 kf의 내적을 계산한다. [E[qf]와 kf의 크기 및 위상차를 입력으로] → [RoPE 회전각 ωfΔ를 포함한 코사인 함수 Σ ∥E[qf]∥∥kf∥ cos(ωfΔ + ϕf) 연산을 수행해] → [거리 Δ에 따른 예상 어텐션 값을 얻고] → [이 값은 해당 키가 미래의 쿼리로부터 받을 평균적인 주목도를 의미한다].

노름 기반 점수 Snorm은 벡터의 집중도가 낮은 예외적인 상황을 보완한다. [쿼리 노름의 기댓값 E[∥qf∥]와 키 노름 ∥kf∥를 입력으로] → [집중도 지표인 Rf를 이용해 (1-Rf) 가중치를 곱하는 연산을 수행해] → [보정된 노름 점수를 얻고] → [이는 방향성이 불분명할 때 크기 정보라도 활용하여 중요한 토큰의 탈락을 방지하는 역할을 한다].

최종적으로 128개 토큰마다 한 번씩 윈도우 단위로 점수를 계산하여 하위 점수의 키를 삭제하는 Pruning 과정을 거친다. GQA(Grouped-Query Attention) 구조에서는 여러 쿼리 헤드의 점수를 Z-score 정규화한 후 최댓값을 취해 통합 관리한다.

Qwen3-8B 모델을 사용한 AIME25 수학 추론 벤치마크에서 TriAttention은 32K 토큰 생성 시 Full Attention과 동일한 40.8%의 정확도를 유지하면서 KV 캐시 메모리를 10.7배 절감했다. 이는 기존의 최신 기법인 R-KV가 동일 효율에서 약 절반의 정확도밖에 내지 못한 것과 대조적이다.

MATH 500 데이터셋에서는 단 1,024개의 KV 캐시 예산만으로도 Full Attention(69.6%)에 근접한 68.4%의 정확도를 기록했다. 또한 Recursive State Query 벤치마크를 통해 복잡한 재귀적 추론 상황에서도 중간 단계의 상태 정보를 잃지 않고 끝까지 추론을 완수하는 뛰어난 메모리 보존 능력을 입증했다.

처리량(Throughput) 측면에서도 큰 이득을 보였다. AIME25 기준 Full Attention 대비 2.5배 빠른 속도를 기록했으며, MATH 500에서는 최대 6.3배의 속도 향상을 달성하여 긴 문맥 추론의 실용성을 크게 높였다.

TriAttention은 RoPE의 수학적 구조가 Pre-RoPE 공간의 정적 벡터를 거리 기반의 동적 어텐션 패턴으로 변환한다는 점에 착안했다. 대다수 LLM에서 관찰되는 Q/K 집중 현상은 어텐션 로그잇이 쿼리와 키의 위치 차이 Δ에만 의존하는 삼각 급수로 근사될 수 있음을 의미하며, 이는 푸리에 합성(Fourier Synthesis)과 유사한 원리로 동작한다.

구현 측면에서 TriAttention은 추가적인 학습 없이 소량의 데이터(약 50k 토큰)로 오프라인 캘리브레이션만 수행하면 된다. 실험 결과 캘리브레이션 데이터의 도메인(코딩 vs 일반 대화)에 관계없이 성능이 일정하게 유지되어 모델 고유의 특성을 잘 포착함을 보여주었다. 또한 최신 아키텍처인 MLA(Multi-head Latent Attention)에서도 Q/K 집중 현상이 더 강하게 나타남을 확인하여 범용성을 입증했다.

현재 구현은 표준 추론 엔진 위에서 동작하며, TriAttention의 삼각 함수 계산 및 Pruning 로직을 하드웨어 수준에서 가속화할 전용 커널 개발이 추가로 필요하다. 또한 현재는 수학, 코딩, 일반 대화 도메인 위주로 검증되었으며 더 넓은 도메인과 헤드별 가변 예산 할당 전략에 대한 연구가 향후 과제로 남아 있다.