기술 요약: 추론 모델을 위한 가치 인식 확률적 KV 캐시 제거 (VASE)

1. 문제 정의

추론 모델(예: Qwen3, OpenAI의 o1)은 최종 답변을 내놓기 전 확장된 사고 체인(Chain of Thought)을 생성함으로써 높은 정확도를 달ะ합니다. 그러나 이러한 능력은 디코딩 단계에서 상당한 메모리 및 연산 병목 현상을 초래합니다. 시퀀스 길이가 길어짐에 따라, 모든 과거 토큰의 표현을 저장하기 위해 필요한 Key-Value(KV) 캐시는 막대한 오버헤드를 발생시킵니다.

기존 솔루션은 두 가지 범주로 나뉩니다:

• 선택 기반 방법 (Selection-based methods): 전체 KV 캐시는 유지하되 어텐션 계산 중에 일부 희소한 토큰 서브셋만을 활성화합니다. 정확도는 높지만, 메모리 사용량이 시퀀스 길이에 따라 선형적으로($O(T)$) 증가하여 메모리 병목 문제를 해결하지 못합니다.
• 제거 기반 방법 (Eviction-based methods): 정해진 예산(Budget)에 도달하면 중요도가 낮은 KV 쌍을 영구적으로 폐기하여, 정적인 메모리 풋프린트와 더 나은 처리량을 제공합니다. 그러나 현재의 제거 방법들은 선택 기반 대안들에 비해 추론 작업에서 상당한 정확도 저하를 겪으며, 이로 인해 모델이 반복적인 추론 루프에 빠지거나 터무니없는 출력을 생성하는 원인이 됩니다.

본 논문은 현재의 제거 전략이 두 가지 결정적인 요소를 간과하고 있음을 식별했습니다: **대규모 크기(Large-magnitude)의 가치 상태(Value states)**가 미치는 불균형한 영향력과, 유지되는 토큰들의 **확률적 다양성(Stochastic diversity)**에 대한 필요성입니다.

2. 방법론: VASE

저자들은 효율성과 정확도 사이의 간극을 메우기 위해 설계된 학습이 필요 없는(Training-free) 제거 프레임워크인 **가치 인식 확률적 KV 캐시 제거(Value-Aware Stochastic KV Cache Eviction, VASE)**를 제안합니다. VASE는 주기적인 제거 프레임워크(지속적인 예산 $K$와 최근 버퍼 $B$ 사용) 내에서 작동하며, 두 가지 핵심 메커니즘을 도입합니다.

A. 대규모 크기의 가치 상태 보호

저자들은 추론 모델의 가치 상태가 매우 편향된 분포를 보이며, 소수의 토큰이 비정상적으로 큰 벡터 크기(범위 $Range(v) = \max(v) - \min(v)$로 측정됨)를 가진다는 점을 관찰했습니다.

• 발견: 이러한 고크기(High-magnitude) 가치들을 제거하면 치명적인 정확도 붕괴(예: GSM8K에서 약 88%에서 14%로 하락)가 발생하며, 모델이 결론에 도달하지 못한 채 문맥을 끝없이 재검토하는 반복 루프를 유발합니다.
• 메커니즘: VASE는 특정 토큰 예산($N_v$)을 할당하여 가치 크기가 가장 큰 $N_v$개의 토큰을 무조건적으로 유지하도록 합니다. 이를 통해 가장 영향력 있는 가치 벡터들이 절대 버려지지 않도록 보장합니다.

B. 확률성 도입

현재의 제거 방법들은 종종 결정론적인 top-$k$ 선택을 사용하며, 이는 유지되는 캐시의 다양성 부족으로 이어질 수 있습니다.

• 발견: 확률성을 도입하면 전체 문맥에 대해 더 대표성 있는 커버리지를 확보함으로써 정확도를 향提升할 수 있습니다.
• 메커니즘: 결정론적으로 상위 점수 토큰을 선택하는 대신, VASE는 가중치 기반 확률적 샘플링을 채택합니다.
  • VASE-AttnV: 가치 인식 예약 방식과 (SnapKV에서 유도된) 어텐션 점수에 기반한 확률적 샘플링을 결합합니다.
  • VASE-DKV: 매 제거 단계마다 가우시안 투영 행렬 $G$를 재샘플링하여 CurDKV 방식을 적응시킵니다(CurDKV는 CUR 행렬 분해의 레버리지 점수를 사용함). 이는 특정 표현을 가진 토큰이 지속적으로 낮은 점수를 부여받고 영구적으로 제거되는 것을 방지합니다.

3. 주요 기여

1. 결정적 요인 식별: 본 논문은 (1) 대규모 크기의 가치 상태가 추론 진행을 유지하고 반복 루프를 방지하는 데 결정적이며, (2) 제거 결정에서의 확률성이 캐시 다양성을 높여 정확도를 크게 향상시킨다는 점을 확립했습니다.
2. VASE 프레임워크: 가치 상태 크기 보호와 확률적 샘플링을 통합한 새로운 학습 불필요 제거 레시피를 제안합니다. 이는 키 기반 스코어링, 가치 기반 스코어링, 그리고 다양성 증진을 결합한 최초의 제거 방법입니다.
3. 양자화와의 연결: 저자들은 대범위 가치 상태가 토큰별 KV 캐시 양자화에서 재구성 오차의 주요 원인임을 입증하였으며, 이는 VASE의 통찰이 다른 압축 기술에도 일반화될 수 있음을 시사합니다.

4. 실험 결과

저자들은 Qwen3-4B 및 Qwen3-14B를 대상으로 6가지 추론 작업(AIME25/26, HMMT25, GPQA-Diamond, MATH, LiveCodeBench-v6)에서 4배의 KV 캐시 압축 비율을 적용하여 VASE를 평가했습니다.

• 선택 방법 대비 정확도: VASE-AttnV는 두 모델 크기 모두에서 가장 강력한 선택 기반 방법인 SeerAttention-R보다 높은 평균 정확도를 달성하면서도 정적인 메모리 풋프린트를 유지했습니다.
  • Qwen3-4B: VASE-AttnV(59.09%)는 SeerAttention-R(58.81%)과 가장 강력한 제거 베이스라인인 R-KV(54.69%)보다 4.4% 더 높은 성능을 보였습니다.
  • Qwen3-14B: VASE-AttnV(65.81%)는 SeerAttention-R(65.37%)과 대등한 수준이었으며, R-KV(60.90%)보다 4.9% 더 우수했습니다.
• 절제 연구 (Ablation Studies):
  • 가치 인식 (Value Awareness): 대규모 크기의 가치를 위한 슬롯을 예약하는 것은 베이스라인 대비 GSM8K 정확도를 최대 16.2% 향상시켰습니다.
  • 확률성 (Stochasticity): CurDKV에 확률적 샘플링을 추가했을 때 Qwen3-14B에서 정확도가 9.2% 향상되었습니다.
• 효율성: VASE-DKV는 테스트된 모든 방법 중 가장 높은 처리량(16K 토큰 기준 풀 모델 베이스라인보다 3.1배 빠름)과 가장 낮은 피크 메모리 사용량을 달성했습니다.
• 코드 생성: LiveCodeBench에서 VASE 방법들은 도메인 변화에 어려움을 겪은 선택 기반 방식인 SeerAttention-R을 크게 앞질렀습니다.

5. 의의 및 주장

본 논문은 VASE가 역사적으로 제거 방법들을 괴롭혀온 효율성과 정확도 사이의 간극을 성공적으로 메웠다고 주장합니다. 대규모 크기의 가치 상태를 우선시하고 확률성을 도입함으로써, VASE는 풀 캐시 또는 선택 기반 방식에서 나타나는 정확도 손실 없이 정적인 메모리 풋프린트로 추론 모델을 운영할 수 있게 합니다.

저자들은 대규모 크기의 가치 상태의 중요성에 관한 발견이 제거를 넘어 KV 캐시 양자화에도 광범위한 시사점을 준다는 점을 강조합니다. 대범위 가치가 오차의 주요 원인으로 식별되었기 때문입니다. 이는 향후 메모리 효율적인 추론 방법들이 이러한 중요한 고크기 상태를 보호하는 혼합 정밀도(mixed-precision) 접근 방식을 고려해야 함을 시사합니다.

궁극적으로 VASE는 FlashAttention2를 지원하고 긴 체인의 추론 모델을 위한 확장 가능한 추론을 가능하게 하는 단순하고 효과적인 학습 불필요 레시피를 제공합니다.