멀티폴 어텐션으로 긴 문맥 추론 효율화

초록

본 논문은 대규모 추론 모델(LRM)이 수천 토큰에 달하는 체인‑오브‑생각을 생성할 때 발생하는 KV 캐시 메모리 부담을 완화하기 위해, 의미적으로 유사한 키들을 클러스터링하고 클러스터 중심을 이용해 중요 토큰은 정확히, 나머지는 근사적으로 처리하는 “Multipole Attention”을 제안한다. 빠른 온라인 클러스터 업데이트와 계층적 센트로이드 비교를 통해 추론 정확도를 유지하면서 최대 4.5배의 어텐션 속도 향상을 달성한다.

상세 분석

Multipole Attention은 기존 희소 어텐션 방식이 초래하는 정확도 저하 문제를 해결하기 위해 두 단계의 핵심 아이디어를 결합한다. 첫 번째는 키 벡터를 의미적 유사성 기반 k‑means 클러스터링으로 묶어 각 클러스터의 중심(센트로이드)을 계산하고, 현재 쿼리와 센트로이드를 비교해 어느 클러스터가 높은 어텐션 점수를 가질지 예측한다는 점이다. 이렇게 추정된 점수가 높은 클러스터에 속한 실제 키들은 “중요 토큰”으로 간주되어 기존과 동일하게 정확한 어텐션 연산을 수행한다. 반면 점수가 낮은 클러스터는 센트로이드 자체와 값(Value) 센트로이드를 이용해 어텐션 기여도를 근사한다. 이때 클러스터 내 토큰 수(N_i)를 가중치로 사용해 전체 어텐션 합을 보정함으로써, 중요한 컨텍스트 정보를 손실 없이 유지한다.

두 번째 핵심은 온라인 클러스터 업데이트 메커니즘이다. 체인‑오브‑생각을 자동 생성하는 과정에서 새로운 토큰이 지속적으로 KV 캐시에 추가되므로, 매 단계마다 전체 시퀀스를 재클러스터링하면 계산 비용이 폭증한다. 이를 방지하기 위해 저자는 블록 단위(k‑W 토큰) 클러스터링과 슬라이딩 윈도우 방식을 도입한다. 새 토큰이 추가되면 마지막 블록만 재클러스터링하고, 블록이 일정 크기를 초과하면 앞쪽 W 토큰을 고정 블록으로 이동시켜 재사용한다. 또한 새 토큰에 대해 무작위 초기 센트로이드를 샘플링하고, 배치형 k‑means(맥퀸 알고리즘) 방식으로 빠르게 할당·업데이트한 뒤, 전체 클러스터에 대해 소수의 정제 단계만 수행한다. 이러한 절차는 클러스터 품질을 유지하면서도 연산 오버헤드를 최소화한다.

계층적 멀티폴 근사는 거리(쿼리와 센트로이드 간) 기반으로 점점 더 거친 센트로이드를 사용해 멀리 떨어진 토큰을 빠르게 근사한다. 즉, 가까운 클러스터는 세밀한 센트로이드를, 먼 클러스터는 더 큰 범위의 센트로이드를 적용해 연산량을 단계적으로 감소시킨다.

실험에서는 Qwen‑8B, DeepSeek‑R1‑Distil‑Qwen2.5‑14B 등 최신 LRM을 대상으로 복잡한 수학·코딩 문제에서 기존 희소 어텐션 대비 정확도 손실이 거의 없으며, KV 캐시 예산을 10% 이하로 줄여도 성능을 유지한다. 커스텀 Triton 커널 구현을 통해 어텐션 단계에서 평균 4.5배, 전체 디코딩 파이프라인에서는 2.8배 정도의 속도 향상을 보고한다.

이 논문은 의미 기반 클러스터링과 센트로이드 근사를 결합한 새로운 희소 어텐션 프레임워크를 제시함으로써, 긴 컨텍스트 추론에서 메모리·연산 효율성을 크게 개선하고, 추론 정확도를 보존하는 실용적인 솔루션을 제공한다.