🤖 machine learning

MomentKV: Closing the Directional Gap in KV Cache Eviction for Long-Context Inference

MomentKV는 제거되는 토큰의 기하학적 규칙성을 보장하기 위해 컴팩트한 모멘트 통계량을 유지하고 이들의 어텐션 기여도에 대한 폐쇄형 보정(closed-form correction)을 제공함으로써 롱 컨텍스트 KV 캐시 제거 시 발생하는 방향성 불일치 병목 현상을 해결하며, 이를 통해 다양한 벤치마크와 압축 수준에서 기존 방법들을 크게 능가합니다.

원저자: Yu Li, Binxu Li, Tian Lan

게시일 2026-06-02

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yu Li, Binxu Li, Tian Lan

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 아주 긴 이야기를 기억하며 글을 계속 써 내려가려고 노력하고 있다고 상상해 보세요. 이를 위해 당신의 뇌는 지금까지 읽은 모든 내용을 담은 "연습장"(KV 캐시)을 계속 보관합니다. 문제는 이야기가 길어질수록 이 연습장이 거대해져서, 결국 책상 전체를 가득 채우게 되어 더 이상 작업하는 것이 불가능해진다는 점입니다.

이를 해결하기 위해, 현재의 방법들은 "쓰레기통" 전략을 사용합니다. 그들은 이야기를 살펴보고, 책상 위에 남겨둘 가장 중요한 문장들을 골라낸 뒤 나머지는 버립니다. 그들은 만약 "최고의" 문장들을 남겨둔다면, 나머지 쓰레기는 무시해도 된다고 가정합니다.

논문의 핵심 발견: "방향"의 문제
이 논문의 저자들인 MOMMENTKV는 이 "쓰레기통" 접근 방식에 숨겨진 결함이 있다는 것을 깨달았습니다.

당신이 책상 위에 남겨둔 몇 개의 잎사귀를 바탕으로 바람의 방향을 추측하려고 한다고 가정해 봅시다.

기존 방식: 가장 세게 흔들리는 잎사귀(가장 "중요한" 것들)를 남기고 나머지는 쓰레기통에 버립니다. 그런 다음 책상 위의 잎사귀들만을 사용하여 바람의 방향을 추측합니다.
문제점: 당신이 버린 잎사귀들이 당신이 남겨둔 것들과는 완전히 다른 방향(수직인 방향)으로 흔들리고 있을 수도 있습니다. 설령 90%를 버렸더라도, 남은 10%가 북쪽을 향하고 있고 쓰레기통의 잎사귀들이 동쪽을 향하고 있다면, 당신의 추측은 크게 틀릴 것입니다. 단순히 남은 북쪽 잎사귀들을 다시 계산(재정규화)한다고 해서 사라진 동쪽 바람을 바로잡을 수는 없습니다. 오류는 얼마나 많은 양을 버렸느냐의 문제가 아니라, 어떤 방향을 잃어버렸느냐의 문제입니다.

해결책: MOMENTKV
MOMENTKV는 단순히 쓰레기를 버리고 잘 되기를 바라는 대신, 문장이 쓰레기통으로 들어가기 전에 그 쓰레기에 대한 아주 작고 초소형의 "요약 노트"를 보관합니다.

이렇게 생각해 보세요:

요약 노트 (모멘트 통계량): 문장을 버리기 전에, 시스템은 그 문장에 대해 몇 가지 간단한 수치를 빠르게 계산합니다: "평균적인 의미는 무엇이었는가?" 그리고 "이 의미가 다음 단어와 함께 어떻게 변하는가?" 문장 전체를 저장하는 것이 아니라, 이 몇 가지 "모멘트"(통계적 지문과 같은 것)만을 저장합니다.
더 똑똑한 버리기: 무엇을 버릴지 결정할 때, 시스템은 다음과 같이 묻습니다: "이 문장은 나의 요약 노트에 의해 이미 잘 표현되어 있는가?" 만약 그렇다면, 버려도 안전합니다. 만약 그렇지 않다면(요약 노트가 포착하지 못하는 고유한 방향을 가지고 있다면), 그 문장을 남겨둡니다. 이는 "쓰레기"가 기하학적으로 규칙적이고 예측 가능하도록 유지해 줍니다.
마법 같은 교정: 모델이 다음 단어를 써야 할 때, 모델은 단순히 책상 위의 문장들만 보는 것이 아닙니다. 모델은 쓰레기통에 있는 쓰레기에 대한 그 작은 "요약 노트"를 사용하여, 사라진 문장들이 기여했을 내용을 수학적으로 재구성합니다. 이는 본질적으로 "내가 이것들을 버렸지만, 내 노트를 보니 이것들이 이야기를 특정 방향으로 밀어붙였다는 것을 알 수 있으니, 그 효과를 다시 더하겠다"라고 말하는 것과 같습니다.

왜 작동하는가
이 논문은 두 가지 유명한 AI 모델(LLaMA 및 Qwen)을 사용하여 두 가지 벤치마크(LongBench 및 RULER)에서 테스트를 진행했습니다.

결과: MOMENTKV는 모든 다른 방법들을 지속적으로 능가했으며, 특히 "책상 공간"이 매우 좁을 때(공격적인 압축 시) 더욱 뛰어난 성능을 보였습니다.
비유: 이는 마치 단순히 가장 인기 있는 책들을 선반에 꽂아두는 것이 아니라, 지하실에 있는 모든 책에 대한 작은 인덱스 카드를 보관하는 사서와 같습니다. 당신이 질문을 던지면, 사서는 그 인덱스 카드를 사용하여 지하실에 있는 책들의 본질을 즉각적으로 떠올려 냄으로써, 지하실을 완전히 무시했을 때보다 훨씬 더 나은 답변을 제공합니다.

요약하자면
현재의 AI 방식들은 오래된 맥락을 버리고 남은 조각들이 충분하기를 바랍니다. MOMENTKV는 버려진 조각들의 방향 또한 그 조각들 자체만큼이나 중요하다는 것을 깨달았습니다. "쓰레기"에 대한 작고 똑똑한 요약을 유지하고 이를 통해 AI의 기억을 수학적으로 수정함으로써, 이 방식은 모델이 위치를 놓치거나 실수를 하지 않고 훨씬 더 긴 이야기를 처리할 수 있게 해줍니다.

기술 요약: MOMENTKV

문제 정의

Transformer 기반 언어 모델의 자기회귀 디코딩(Autoregressive decoding)은 이전에 계산된 투영(projection)을 디코딩 단계마다 재사용하기 위해 키-값(Key-Value, KV) 캐시에 의존합니다. 시퀀스 길이가 수만 또는 수십만 토큰으로 확장됨에 따라, KV 캐시의 메모리 점유율은 선형적으로 증가하며 긴 문맥(long-context) 추론의 주요 병목 현상이 됩니다.

KV 캐시 제거(Eviction)는 고정된 크기의 키-값 쌍 부분 집합을 유지하고 나머지는 버림으로써 이 문제를 해결합니다. 기존 방법들은 누적 어텐션 가중치나 가치 인식 스코어링(value-aware scoring) 등을 기반으로 한 토큰 선택 전략을 개선해 왔으나, 다음과 같은 근본적인 한계를 공유합니다: 제거 후의 추론은 오직 유지된 토큰들에 대해서만 어텐션을 재정규화(renormalize)하며, 제거된 집합은 완전히 무시합니다.

저자들은 출력 성능 저하의 주된 원인이 단순히 제거된 토큰에 남겨진 잔여 어텐션 질량(residual attention mass) 때문이 아니라, **유지된 토큰 집합과 제거된 토кен 집합 사이의 방향성 불일치(directional mismatch)**에 있다는 점을 식별했습니다. 실증적 분석 결과, 실제로 제거된 토큰들은 유지된 토큰들과 **거의 직교(near-orthogonal)**하는 경향이 있음이 밝혀졌습니다. 결과적으로, 표준적인 재정규화로는 복구할 수 없는 상당한 방향성 오류를 아주 적은 양의 제거된 어텐션 질량조차 만들어낼 수 있습니다. 이는 "방향성 갭(directional gap)"을 생성하며, 토큰 선택 능력만을 높이는 것은 수확 체감의 법칙에 직면하게 만듭니다. 왜냐하면 더 선택적인 유지는 유지된 집합을 특정 쿼리에 정렬된 방향으로 더욱 집중시켜, 제거된 집합과의 격차를 더욱 벌리기 때문입니다.

방법론: MOMENTKV

이 방향성 갭을 해결하기 위해, 저자들은 제거된 토큰 집합에 대해 압축된 소규모 모멘트 통계량(moment statistics)을 유지하는 프레임워크인 MOMENTKV를 제안합니다. 이 통계량은 두 가지 목적을 수행합니다: 제거 결정을 안내하고, 추론 중 어텐션 출력을 보정하는 것입니다.

1. 압축된 모멘트 통계량

MOMENTKV는 제거된 토큰을 직접 저장하는 대신, 매 제거 시마다 점진적으로 업데이트되는 네 가지 실행 합계(running sums)를 헤드당 유지합니다:

개수 ( $n_e$ ): 제거된 토큰의 수.
키 합 ( $s_k$ ): $\sum_{i \in E} k_i$ .
값 합 ( $s_v$ ): $\sum_{i \in E} v_i$ .
외적 합 ( $S$ ): $\sum_{i \in E} v_i k_i^\top$ .

이러한 통계량을 통해 제거된 토큰들을 다시 방문하지 않고도 키 평균( $\bar{k}$ ), 값 평균( $\bar{v}$ ), 그리고 **값-키 공분산( $\tilde{S}$ )**을 계산할 수 있습니다. 총 저장 오버헤드는 컨텍스트 길이에 독립적으로 헤드당 $O(d^2)$ 입니다.

2. 제거된 출력의 1차 근사 (First-Order Approximation)

추론 중에 MOMENTKV는 제거된 어텐션 출력 $\hat{f}_E(q)$ 에 대한 폐쇄형(closed-form) 1차 근사치를 도출합니다. 평균 키를 중심으로 로짓(logits)을 재중심화하고 소프트맥스 함수에 1차 테일러 전개(Taylor expansion)를 적용함으로써, 저자들은 다음을 도출합니다:
$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$

$\bar{v}$ 항은 쿼리와 무관한 베이스라인 역할을 하며, 제거된 하위 출력의 평균 방향을 복구합니다.
$\tilde{S} q$ 항은 쿼리 적응형 보정 역할을 하며, 베이스라인만으로는 놓칠 수 있는 수직적 변이(방향성 콘텐츠)를 포착합니다.

이 근사식은 모델이 손실된 제거된 기여도를 재구성할 수 있게 하여, 표준적인 제거 방식이 말살해버리는 방향성 정보를 효과적으로 복원합니다.

3. 모멘트 기반 제거 (Moment-Informed Eviction)

모멘트 통계량은 제거 과정 자체를 가이드합니다. 각 유지된 토큰 $j$ 에 대해, 알고리즘은 **모멘트 잔차(moment residual)**를 계산합니다:
$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$
이 잔차는 토큰 $j$ 가 축적된 요약 정보 외에 얼마나 새로운 방향성 정보를 담고 있는지를 측정합니다. 제거 점수는 어텐션 가중치와 잔차 노름(norm)의 곱으로 정의됩니다: $\text{score}(j) = \alpha_j \cdot \|r_j\|$ .

잔차가 낮은 토큰은 모멘트 모델에 의해 잘 예측되므로 제거 우선순위가 높아집니다.
잔차가 높은 토큰(고유한 방향성 정보를 가진 토큰)은 어텐션 가중치가 낮더라도 유지됩니다.

이는 **상호 강화 루프(mutually reinforcing loop)**를 형성합니다: 모멘트 기반 제거는 제거된 집합을 "기하학적으로 규칙적"(낮은 로짓 분산 $\sigma$ )으로 유지하여, 추론 보정을 위한 근사 경계(approximation bound)를 좁힙니다. 결과적으로, 정확한 보정은 향후 제거 결정을 위한 더 나은 잔차를 제공합니다.

4. 정규화 보정 추론 (Normalization-Corrected Inference)

최종 어텐션 출력은 유지된 출력( $f_R$ )과 근사된 제거된 출력( $\hat{f}_E$ )의 볼록 조합(convex combination)으로 재구성됩니다:
$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$
혼합 가중치 $\hat{w}_R$ 은 젠슨의 부등식(Jensen's inequality)을 통해 도출된 제거된 파티션 함수의 하한(lower bound)을 사용하여 추정됩니다. 이를 통해 모델이 제거된 질량의 크기를 고려하면서 동시에 그 방향을 보정할 수 있도록 보장합니다.

핵심 기여

방향성 갭 식별: 본 논문은 KV 캐시 제거의 지배적인 에러 원인이 단순히 제거된 어텐션 질량의 크기가 아니라, 유지된 부분 출력과 제거된 부분 출력 사이의 거의 직교하는 관계임을 입증했습니다.
압축된 모멘트 통계량: 헤드당 $O(d^2)$ 의 저장 공간만을 사용하여 제거된 토큰의 기하학적 특성(평균 및 공분산)을 추적하는 메모리 효율적인 방법을 도입했습니다.
폐쇄형 보정: 평균 방향과 쿼리 의존적 변이를 모두 복구하는 제거된 어텐션 출력에 대한 1차 근사식을 도출했습니다.
이중 역할 메커니즘: 모멘트 통계량이 제거 결정(중복 토큰 식별)과 추론 출력 보정(손실된 방향성 정보 복구)을 동시에 최적화하는 통합된 프레임워크를 제시했습니다.

실험 결과

저자들은 LLaMA-3.1-8B-Instruct와 Qwen3-4B-Instruct를 사용하여 LongBench 및 RULER 벤치마크에서 MOMENTKV를 평가했습니다.

성능: MOMENTKV는 모든 캐시 예산( $L \in \{128, 256, 512, 1024\}$ )에 대해 모든 베이스라인(H2O, SnapKV, PyramidKV, Ada-KV)을 일관되게 능가했습니다.
공격적인 압축: 가장 큰 이득은 공격적인 압축 상황에서 관찰되었습니다. LLaMA-3.1-8B의 LongBench ( $L=128$ )에서 MOMENTKV는 가장 강력한 베이스라인인 Ada-KV 대비 평균 +1.35의 점수 향상을 달성하며, 전체 캐시 성능의 **94.1%**를 유지했습니다.
검색 작업: 정밀한 토큰 검색을 테스트하는 RULER에서, MOMENTKV는 LLaMA와 Qwen 각각에 대해 Ada-KV 대비 각각 +3.4와 +3.5의 비례적인 이득을 보여주며 검색 집약적 환경에서의 효과를 입증했습니다.
절제 연구(Ablation Studies):
- **정규화 보정 추론(NC)**은 모멘트 기반 제거(MI) 단독(+0.57)보다 훨씬 더 큰 기여(+2.41)를 했으며, 이는 제거된 방향을 복구하는 것이 성능 향상의 주된 동력임을 확인시켜 줍니다.
- 두 메커니즘의 결합은 시너지 효과를 내어(+총 3.17 이득), 개별 이득의 합을 초과했습니다.
효율성: 이 방법은 토큰당 단 1.6 ms의 지연 시간만을 추가하며, 4.1 MB의 보조 메모리(유지된 캐시의 6.4%)를 요구하는 반면, 전체 캐시 대비 토큰당 지연 시간을 65% 이상 줄였습니다.

의의

본 논문은 긴 문맥 추론을 더욱 발전시키는 길은 단순히 더 나은 토큰을 선택하는 것이 아니라, 현재의 제거 전략이 말살해버리는 방향성 정보를 복구하는 것에 있다고 주장합니다. MOMENTKV는 제거된 집합에 대한 압축된 통계량을 유지함으로써 방향성 갭을 좁힐 수 있으며, 이를 통해 제거의 "줄일 수 없는 비용"을 교정 가능한 근사 오류로 전환할 수 있음을 보여줍니다. 이 접근 방식은 고정된 예산의 KV 캐시를 사용하면서도 긴 문맥 언어 모델을 효율적으로 배포할 수 있는 실용적이고 확장 가능한 솔루션을 제공합니다.