기술 요약: KVpop — 예측 기반 온라인 프루닝을 통한 키-값 캐시 압축

1. 문제 정의

Transformer 기반 대규모 언어 모델(LLM)의 자기회귀 디코딩은 과거 토큰의 표현을 재계산하지 않기 위해 키-값(KV) 캐시에 의존합니다. 그러나 KV 캐시는 컨텍스트 길이에 따라 선형적으로 증가하며, 이는 긴 컨텍스트 추론 시 메모리와 대역폭 측면에서 심각한 병목 현상을 초래합니다.

기존의 비용 절감 솔루션은 일반적으로 두 가지 범주로 나뉩니다:

• 정적/휴리스틱 제거 (Static/Heuristic Eviction): Sliding Window 또는 Sink-Token 보존(예: StreamingLLM)과 같은 방법은 고정된 규칙에 의존하며, 점수 기반 방식(예: 누적 어텐션)은 온라인 프록시를 사용합니다. 이러한 방식은 미래 토큰의 유용성을 추적하는 데 실패하여, 관련성이 변할 때 취약한 제거 성능을 보입니다.
• 학습된 제거 (Learned Eviction): DMC(Dynamic Memory Compression)나 DMS(Dynamic Memory Sparsification)와 같은 접근법은 토큰을 병합하거나 버리는 정책을 학습합니다. 그러나 이러한 많은 방법은 토큰 삽입 시점에 제거 결정을 내리거나, 미래의 유용성에 기반한 구체적인 "유지 또는 삭제" 결정을 직접 감독하지 않는 미분 가능한 완화 기법에 의존합니다. 또한, 보호된 윈도우(protected windows)를 사용하는 방식은 토큰의 제거를 지연시키기는 하지만, 토큰이 보호 윈도우 내에 머무는 동안 축적되는 증거를 무시한 채 토큰이 들어오는 즉시 그 운명을 결정해 버리는 경우가 많습니다.

핵심 과제는 어떤 토큰이 미래의 쿼리에 유용할지 예측하는 것입니다. 로컬하게 두드러지는 토큰은 나중에 무의미해질 수 있는 반면, 즉각적인 어텐션은 낮더라도 나중에 중요해질 수 있는 토큰(예: 추론 과정에서의 토큰)이 존재하기 때문입니다.

2. 방법론: KVpop

KVpop(Predictive Online Pruning을 통한 KV 압축)은 기본 모델 아키텍처를 변경하지 않고 고정된 헤드당 KV 예산을 강제하도록 설계된 희소 어텐션 레트로핏(sparse-attention retrofit)입니다. 이 방식은 **미래 어텐션 감독(Future-Attention Supervision)**과 **지연된 상태 기반 스코어링(Delayed Stateful Scoring)**이라는 두 가지 주요 메커니즘을 통해 이전 연구의 한계를 해결합니다.

2.1 아키텍처 및 예산

KVpop은 $B = s + w + k$로 정의된 헤드당 제한된 캐시 예산을 유지합니다:

• $s$ Sink Tokens: 처음 몇 개의 토큰은 항상 유지됩니다.
• $w$ Protected Window: 가장 최근의 $w$개 토큰은 제거로부터 보호됩니다.
• $k$ Long-Range Top-K: 오래된 토큰들은 학습된 중요도 점수에 따라 고정된 $k$개의 슬롯을 두고 경쟁합니다.

2.2 미래 어텐션 타겟 감독

과거의 어텐션이나 로컬 통계에 기반하여 토큰을 스코어링하는 방식과 달리, KVpop은 토큰이 보호 윈도우를 벗어난 후 받는 **미래 어텐션 질량(future attention mass)**을 사용하여 스코러를 감독합니다.

• 타겟 정의: 토겟은 토큰 $t$가 받는 미래 쿼리(위치 $d \ge t + w$)로부터의 집계된 로그 질량(log-mass)입니다.
• 효율적인 계산: 훈련 중 밀집 어텐션 맵을 생성하는 $O(S^2)$ 비용을 피하기 위해, KVpop은 **전치 어텐션 패스(transposed-attention pass)**를 사용합니다. 쿼리와 키의 역할을 바꾸고 희소 어텐션 커널에서 반환된 log-sum-exp (LSE) 정규화 인자를 재사용함으로써, 시스템은 밀집 행렬을 생성하지 않고도 미래 어텐션 타겟을 효율적으로 계산합니다.
• 손실 함수: 훈련 목적 함수는 경계 인식 쌍체 로지스틱 손실(boundary-aware pairwise logistic loss)입니다. 이는 새로 적격해진 토큰($t_{new}$)의 예측 점수를 교사 모델의 컷오프 토큰($t_{bnd}$)의 점수와 비교합니다. 이 손실은 학생 모델이 $t_{new}$를 $t_{bnd}$와 비교하여 올바르게 순위를 매기도록 유도하여, top-$k$ 선택이 교사의 미래 유용성 순위와 일치하도록 합니다.

2.3 상태 기반 스코어링 및 지연된 결정

KVpop은 지연된 메모리 기반 스코러(mLSTM으로 구현됨)를 도입하여 토큰이 제거 경계에 도달할 때까지 스코어링 결정을 미룹니다.

• 메커니즘: 스코러는 토큰 삽입 시 즉시 스코어를 매기는 대신, 토큰의 표현과 후속 컨텍스트를 통해 내부 상태를 업데이트합니다. 토큰이 보호 윈도우를 벗어날 때($q - w$), 스코러는 축적된 메모리를 읽어 중요도 점수를 생성합니다.
• 장점: 이를 통해 스코러는 (토큰이 보호 윈도우 내에 머무는 동안 축적되는 증거인) "근미래" 컨텍스트를 활용하여 유지 또는 삭제 결정을 내릴 수 있으며, 이는 삽입 시점에 스코어를 매기는 방식의 한계를 해결합니다.

2.4 추론 효율성

추론 시, KVpop은 고정된 예산 $B$를 강제합니다. 상위 $k$개의 장거리 토큰 선택은 정적 토큰 우선순위와 Fenwick 트리를 사용하는 러닝 스레숄드(running threshold)를 통해 최적화되며, 이를 통해 밀집 $S \times S$ 마스크를 저장하지 않고도 온더플라이(on-the-fly)로 희소 마스크를 생성할 수 있습니다 (예: FlexAttention 커널 내에서).

3. 주요 기여

1. 감독된 제거 정책: 저자들은 고정 예산 KV 제거를 미래 어텐션 타겟을 사용하여 제거 경계에서의 지도 학습 기반 예측 문제로 정식화했습니다.
2. 효율적인 타겟 계산: 밀집 어텐션 맵을 생성하지 않고 미래 어텐션 타겟을 계산하는 훈련 전용 전치 어텐션 패스를 도입하여 추론 시 오버헤드가 없음을 보장합니다.
3. 지연된 스코어링: 상태 기반 스코러가 토큰이 제거 가능해질 때까지 결정을 미룸으로써 근미래 컨텍스트를 활용할 수 있음을 입증하였으며, 이는 즉각적 스코어링 방식보다 우수한 성능을 보였습니다.
4. 실증적 검증: 수학적 추론(AIME, HMT) 및 도메인 외 작업(GPQA-D, LiveCodeBench)에서 평가를 수행하여, 학습된 제거 정책이 공격적인 압축 하에서도 높은 성능을 유지할 수 있음을 보여주었습니다.

4. 실험 결과

저자들은 Qwen3-4B 및 Qwen3-8B 모델을 대상으로 Nemotron-Math v2 데이터셋을 사용하여 시퀀스 길이 16,384에서 KVpop을 평가했습니다.

4.1 수학적 추론 (AIME & HMMT)

• 압축률 75% (예산 2048):
  • Qwen3-4B: KVpop은 풀 어텐션 교사 모델 성능의 **98%**를 유지했습니다.
  • Qwen3-8B: KVpop은 100%(교사 성능에 근접)를 유지했습니다.
  • 비교: KVpop은 훈련이 없는 베이스라인(StreamingLLM, TOVA) 및 학습된 베이스라인인 DMS를 크게 앞질렀습니다. 예를 들어, 75% 압축 시 StreamingLLM은 Qwen3-4B에서 상대적 성능이 ~76%로 떨어졌으나, KVpop은 98%를 유지했습니다.
• 압축률 88% (예산 4096):
  • Qwen3-4B: KVpop은 교사 성능의 **97%**를 유지했습니다.
  • Qwen3-8B: KVpop은 교사 성능의 **100%**를 유지했습니다.
  • 비교: 훈련이 없는 방법들은 크게 무너졌으나(예: StreamingLLM은 약 47%의 상대적 성능으로 하락), KVpop은 견고함을 유지했습니다.

4.2 도메인 외 일반화

GPQA-Diamond (과학적 추론) 및 LiveCodeBench (코드 생성)에 대한 평가:

• 수학적 추론 데이터로만 훈련된 KVpop 모델은 이러한 도메인에서도 잘 일반화되었습니다.
• 75% 압축 시, Qwen3-4B 기반 KVpop은 GPQA-D에서 밀집 교사 모델과 유사한 정확도(0.59 vs 0.59)를 달성했으며, LCB에서도 경쟁력 있는 결과를 보였습니다.

4.3 추론 효율성

• 메모리: 밀집 어텐션의 경우 VRAM 사용량이 생성 길이에 따라 선형적으로 증가하지만(131k 토큰에서 약 36GB 도달), KVpop과 DMS는 시퀀스 길이에 관계없이 거의 일정한 메모리 풋프린트(~19GB)를 유지했습니다.
• 지연 시간: KVpop은 긴 생성 길이(최대 131k 토큰)에서 DMS 및 밀집 어텐션보다 낮은 엔드 투 엔드 디코딩 지연 시간을 보여주었습니다. 저자들은 이를 KVpop의 고정된 헤드당 예산 덕분으로 보았으며, 이는 컴파일 및 실행을 복잡하게 만드는 DMS의 "래기드(ragged)"한 캐시 구조를 피할 수 있게 해줍니다.

4.4 절제 연구 (Ablation Studies)

• 지연된 스코어링: 지연을 제거하면(삽입 시 즉시 스코어링) 토큰 정확도가 0.2포인트 하락하여, 근미래 컨텍스트를 통합하는 것의 이점을 확인했습니다.
• 상태 기반 vs 상태 비기반: 상태 기반 변형(KVpop)이 일반적으로 상태 비기반 변형(KVpopmlp)과 비슷하거나 약간 더 나은 성능을 보였으며, 특히 높은 압축률에서 그러했습니다.

5. 의의 및 주장

본 논문은 미래 어텐션 신호로 제거를 감독하는 것이 정적 휴리스틱이나 프록시 점수보다 KV 캐시 압축을 위한 더 우수한 전략이라고 주장합니다. 미래의 쿼리에 대해 어떤 토큰이 중요할지 예측하도록 직접 훈련함으로써, KVpop은 다음을 달성합니다:

1. 높은 충실도: 다른 방법들이 실패하는 높은 압축률(최대 88%)에서도 교사 모델의 성능을 거의 그대로 유지합니다.
2. 강건성: 학습된 정책이 도메인 외 작업으로 효과적으로 전이됨을 보여주며, 이는 모델이 특정 훈련 분포에 과적합되는 것이 아니라 토큰 유용성의 일반적인 원칙을 학습함을 시사합니다.
3. 효율성: 엄격한 메모리 경계를 강제하여, 모델 품질을 희생하지 않으면서도 일정한 메모리 사용량과 향상된 처리량을 통해 안정적인 롱 컨텍스트 디코딩을 가능하게 합니다.

저자들은 KVpop이 적절히 미래 유용성에 의해 감독되고 근미래 컨텍스트를 활용할 수 있다면, 모델의 품질을 저하시키지 않고도 자기회귀 생성의 메모리 병목 현상을 효과적으로 해결할 수 있음을 입증했다고 결론짓습니다.