💬 NLP

ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution

ForesightKV는 골든 에빅션(Golden Eviction) 알고리즘을 지도 학습 및 강화 학습(GRPO)과 결합하여 가장 중요한 KV 쌍을 예측하고 유지함으로써, 긴 추론 작업에서 높은 성능을 유지하면서도 메모리 비용을 크게 절감하도록 추론 모델의 KV 캐시 제거를 최적화하는 훈련 기반 프레임워크입니다.

원저자: Zican Dong, Peiyu Liu, Junyi Li, Zhipeng Chen, Han Peng, Shuo Wang, Wayne Xin Zhao

게시일 2026-02-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zican Dong, Peiyu Liu, Junyi Li, Zhipeng Chen, Han Peng, Shuo Wang, Wayne Xin Zhao

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

당신이 아주 길고 복잡한 미스터리를 해결하려는 천재적인 탐정(AI)이라고 상상해 보세요. 이를 위해 당신은 책상 위에 방대한 양의 단서가 담긴 노트(KV 캐시)를 두어야 합니다. 이야기가 길어질수록 당신의 노트는 점점 더 두꺼워집니다. 결국, 책상은 종이들로 너무 어지러워져서 손을 움직일 수도 없고, 다음 추론을 하기 위해 특정 단서를 찾는 데 한참의 시간이 걸리게 됩니다. 이것이 바로 거대 언상 모델(LLM)에서 발생하는 "메모리 과부하" 문제입니다.

전통적으로, 책상을 정리하기 위해 사람들은 단순한 규칙을 사용합니다: "가장 오래된 종이를 버려라" 또는 "잉크가 가장 적게 묻은 종이를 버려라"와 같은 식이죠. 하지만 ForesightKV 논문은 이러한 규칙들이 너무 멍청하다고 주장합니다. 때로는 오래된 종이나 잉크가 흐릿한 종이가 이야기의 후반부에서 미스터리를 풀기 위한 가장 결정적인 단서가 될 수도 있기 때문입니다. 이를 버리는 것은 탐정이 실수를 저지르게 하여 수사 전체를 망쳐버리는 결과를 초래합니다.

ForesightKV가 이 문제를 해결하는 방법은 다음과 같이 세 단계로 설명할 수 있습니다.

1. 문제점: "원패스(One-Pass)"의 실수

당신이 여행 가방을 싸고 있다고 상상해 보세요. 표준적인 방법은 "마지막에 챙긴 10개의 아이템만 남기고 나머지는 버려라"라고 말할 것입니다. 하지만 만약 당신이 가장 처음에 챙긴 아이템(예: 지도)이 여행 내내 필수적인 것이라면 어떻게 될까요? 비록 그 지도가 가방 깊숙이 들어있더라도 말이죠.

기존 방식들은 단순한 패턴(예: "최신 것을 유지하라" 또는 "가장 인기 있는 것을 유지하라")을 기반으로 어떤 단서를 남길지 예측하려고 합니다. 이 논문은 복잡한 추론(예: 수학 문제)에서 단서들이 세 가지 유형의 '성격'을 가지고 있다는 것을 발견했습니다:

글로벌 스타 (The Global Star): 언제나, 무엇보다 중요한 존재.
로컬 이웃 (The Local Neighbor): 짧은 시간 동안만 중요한 존재.
시맨틱 카멜레온 (The Semantic Chameleon): 이들이 바로 까다로운 존재입니다. 지금은 쓸모없어 보일 수 있지만, 50단계 뒤에는 전체 퍼즐을 푸는 열쇠가 되는 단서입니다. 단순한 규칙들은 이러한 "카멜레온"들을 놓치기 쉽습니다.

2. 해결책: "시간 여행을 하는" 코치

저자들은 ForesightKV라는 새로운 시스템을 만들었습니다. 경직된 규칙 책을 사용하는 대신, 그들은 미래를 내다볼 줄 아는 작은 스마트한 조수(스코어링 모델)를 훈련시켰습니다.

이 코치는 단순히 현재의 단서만을 보는 것이 아니라, 미래에 어떤 단서가 가장 중요해질지를 예측하도록 학습됩니다. 코치는 다음의 2단계 훈련 과정을 통해 이를 수행합니다:

1단계: "골든 스탠다드" (지도 학습)

먼저, 연구진은 아무것도 버리지 않은 상태에서 AI에게 수학 문제를 풀게 했습니다. 그리고 AI의 주의력(Attention)을 관찰하며 다음과 같이 물었습니다: "만약 우리가 지금 당장 몇몇 단서를 버려야 한다면, 최종 정답에 가장 적은 피해를 줄 수 있는 것은 무엇인가?"
그들은 **골든 에비션(Golden Eviction)**이라는 방법을 사용하여 남겨두어야 할 완벽한 단서 세트를 찾아냈습니다. 그런 다음, 이 작은 조수가 이 "완벽한" 의사결정 과정을 모방할 수 있도록 가르쳤습니다. 이는 마치 학생에게 정답지를 보여주며 "어떻게 하면 옳은 답을 고를 수 있는지 배워라"라고 말하는 것과 같습니다.

2단계: "실전 시뮬레이션" (강화 학습)

하지만 AI가 실제로 문제를 풀고 있을 때는, AI의 사고 과정이 변화하기 때문에 "완벽한" 정답지가 항상 완벽하지는 않습니다.
따라서 두 번째 단계는 비디오 게임 시뮬레이션과 같습니다. 조수에게 다음과 같은 명령이 내려집니다: "자, 단서들을 마음껏 버려봐. 만약 AI가 특정 유형의 단어(예: 틀리기 쉬운 숫자나 기호)에서 실수를 하면 너는 벌점을 받을 거야. 만약 AI가 계속해서 정답을 맞히면 너는 보상을 받을 거야."
조수는 이를 통해 더욱 똑똑해지는 법을 배웁니다. 즉, 특정 "지루한" 단어들(저엔트로피 토큰)을 유지하는 것이 나중에 AI가 숫자를 잘못 말하는(환각 현상) 것을 방지하는 데 매우 중요하다는 사실을 깨닫게 됩니다.

3. 결과: 더 똑똑하고 가벼워진 책상

연구진은 이 모델을 세 가지 서로 다른 AI 모델에 적용하여 어려운 수학 벤치마크(AIME 2024 및 2025)로 테스트했습니다.

주장: ForesightKV는 전체 용량을 가진 무거운 노트만큼 잘 문제를 풀면서도, 메모리 공간은 절반만 사용할 수 있습니다.
비유: 이는 마치 배낭이 50% 더 가벼워졌음에도 불구하고, 게임에서 이기기 위해 필요한 모든 단서를 여전히 기억하고 있는 것과 같습니다.
속도: 배낭이 가벼워졌기 때문에, AI는 더 빠르게 생각할 수 있고 더 많은 사람을 동시에 처리할 수 있습니다(높은 처리량).

요약

ForesightKV는 AI의 메모리를 관리하는 새로운 방법입니다. 단순히 오래된 노트를 무작정 버리는 대신, 장기적인 해결책에 어떤 노트가 결정적일지를 예측하도록 훈련된 스마트한 조수를 활용합니다. "완벽한 예시"를 통한 훈련 단계와 "직접 해보며 배우는" 게임 단계를 결합함으로써, 복잡한 추론 퍼즐을 풀기 위해 필요한 중요한 세부 사항을 잊지 않으면서도 AI를 빠르고 효율적으로 유지합니다.

기술 요약: ForesightKV

문제 정의

대규모 언어 모델(LLM), 특히 추론 모델은 긴 추론 흔적(reasoning traces)을 생성하는 데 탁월한 능력을 보여주었습니다. 그러나 시퀀스 길이가 길어짐에 따라 Key-Value(KV) 캐시가 선형적으로 확장되어, 메모리 제한적인 디코딩 과정에서 상당한 메모리 오버헤드와 지연 시간 증가를 초래합니다. 기존의 KV 캐시 제거(eviction) 방식들은 주의력 점수(attention scores), 위치 또는 특징과 같은 휴리스틱 규칙을 기반으로 덜 중요한 KV 쌍을 폐기하여 이를 완화하려고 시도하지만, 서로 다른 어텐션 헤드 간의 복잡하고 동적인 의존성을 포착하는 데는 종종 실패합니다. 구체적으로, 이러한 학습이 필요 없는(training-free) 방식들은 컨텍스트의 진화에 따라 KV 쌍의 중요성이 동적으로 변하는 "의미 의존적(semantic-dependent)" 패턴을 처리하는 데 어려움을 겪습니다. 또한, 중요한 KV 쌍을 제거하면 저엔트로피 토큰(결정론적인 추론 부분)에서 손실(loss)이 급격히 증가하여, 이후의 추론 단계를 왜곡하는 사실적 오류를 유발할 수 있습니다.

방법론

저자들은 KV 쌍의 장기적 기여도를 학습하고 생성 중에 덜 중요한 것들을 적응적으로 제거하도록 설계된 학습 기반 프레임워크인 ForesightKV를 제안합니다. 이 프레임워크는 고정된 메모리 예산 내에서 작동하며, 각 KV 쌍의 중요도를 예측하는 경량 스코어링 모델( $\pi_\theta$ )을 최적화하기 위해 2단계 학습 패러다임을 채택합니다. 결정적으로, 전체 과정 동안 LLM 파라미터는 동결된 상태로 유지되며, 오직 스코어링 모델만 업데이트됩니다.

1. 동적 제거 프로세스 (Dynamic Eviction Process)

제거는 주기적으로(매 $L$ 토큰마다) 트리거됩니다. 캐시 크기가 $B + L$ ( $B$ 는 목표 예산)에 도달하면, 시스템은 가장 최근의 $L$ 개 토큰을 유지하고 스코어링 모델을 사용하여 나머지 캐시를 $B - L$ 로 압축합니다. 스코어링 모델은 키( $k_n$ ), 값( $v_n$ ), 그리고 어텐션 특징( $a_n$ )의 결합을 입력으로 받아 중요도 점수( $\phi_n$ )를 예측합니다. 제거 결정은 무작위성을 도입하기 위해 Top-K 이후 다항 샘플링(multinomial sampling) 전략을 통해 이루어집니다.

2. 1단계: "골든 이비션(Golden Eviction)"을 통한 지도 학습

스코어링 모델을 학습시키기 위해, 저자들은 감독 레이블을 구축하는 Golden Eviction 알고리즘을 도입했습니다.

메커니즘: 전체 추론 흔적에 대해 어텐션 행렬을 블록 단위로 분할합니다. 알고리즘은 블록 내부 및 블상(across) 어텐션 헤드 간의 어텐션 점수를 집계하여 "블록 점수"를 계산합니다.
미래 인지(Future Awareness): 각 제거 단계에 대해, 알고리즘은 각 KV 쌍에 대해 모든 미래 단계에서의 최대 블록 점수를 식별합니다. 이 "미래 점수"는 미래의 어텐션 계산에 미치는 영향을 최소화하는 최적의 유지 전략을 나타냅니다.
학습: 스코어링 모델은 골든 이비션 알고리즘으로부터 도출된 미래 점수와 예측된 두 KV 쌍의 점수가 일치하도록 하는 Pairwise Ranking Loss를 사용하여 학습됩니다.

3. 2단계: 강화 학습 (RL)

지도 학습만으로는 제거 중에 발생하는 분포 변화(distribution shift)를 설명하지 못할 수 있습니다. 이 간극을 메우기 위해, 저자들은 KV 캐시 제거를 **마르코프 결정 과정(MDP)**으로 정식화하고 GRPO(Group Relative Policy Optimization) 알고알고리즘을 적용합니다.

상태(State): 현재 남아 있는 KV 캐시.
행동(Action): 유지할 KV 쌍의 하위 집합을 선택함.
보상(Reward): 보상 함수는 저엔트로피 토큰에서의 성능 저하를 완화하도록 설계되었습니다. 구체적으로, 원래는 낮은 엔트로피를 가졌으나 제거 후 손실이 크게 증가한 토큰들에 집중합니다. 보상은 이 특정 토큰 집합에 대한 손실 증가의 평균 제곱 오차(MSE)의 음수 값으로 정의됩니다.
최적화: 정책 모델은 시퀀스 전반에 걸쳐 누적 보상을 최대화하도록 최적화되어, 결정론적인 추론 단계에 필요한 핵심 컨텍스트를 효과적으로 보존하는 법을 학습합니다.

주요 기여

ForesightKV 프레임워크: KV 쌍의 장기적 기여도를 예측하는 법을 학습하여 정적인 휴리스틱 규칙을 넘어선 새로운 2단계 학습 방법(지도 학습 + RL)을 제시합니다.
Golden Eviction 알고리즘: 스코어링 모델을 위한 고품질의 감독 신호를 제공하기 위해 미래 어텐션 정보를 사용하는 최적의 제거 궤적 생성 방법을 제공합니다.
제거를 위한 MDP 정식화: 제거 과정을 MDP로 모델링하고, 결정론적 추론의 정확성에 중요한 저엔트로피 토큰의 손실 급증을 타겟팅하기 위해 GRPO를 활용합니다.
효율성: 베이스 LLM에 대한 역전파 연산이나 파라미터 업데이트 없이 경량 스코어링 모델만을 학습시키는 프레임워크입니다.

실험 결과

저자들은 AIME2024 및 AIME2025 수학 벤치마크를 사용하여 세 가지 추론 모델(Qwen3-1.7B, Qwen3-4B, DeepSeek-R1-Distill-Qwen-7B)에 대해 ForesightKV를 평가했습니다.

성능: ForesightKV는 기존 베이스라인(SnapKV, H2O, R-KV)보다 일관되게 우수한 성능을 보였습니다. 특히, 단 절반의 KV 캐시 예산만을 사용하면서도 전체 캐시 베이스라인과 대등하거나 더 뛰어난 성능을 달 기록했습니다. 예를 들어, Qwen3-4B를 사용한 AIME2024에서 ForesimentKV는 1K 예산으로 2K 예산을 사용한 R-KV보다 높은 성능(54.5 vs 44.8 Pass@1)을 달성했습니다.
처리량(Throughput): 이 방법은 처리량을 크게 개선했습니다. 32K 시퀀스에서 1K 예산을 사용할 때, ForesightKV는 전체 KV 캐시 대비 최대 9.79배의 처리량 이득을 달성했습니다.
일반화: 주로 수학 작업에 대해 학습되었음에도 불구하고, ForesightKV는 GPQA(과학) 및 LiveCodeBench(코딩) 벤치마크에서 절반의 캐시 예산을 가진 베이스라인들을 능가하며 강력한 일반화 능력을 입증했습니다.
절제 연구(Ablation Studies): 실험 결과, 지도 학습 단계(초기 예측 능력 제공)와 강화 학습 단계(저엔트로피 토큰 보존을 위한 정책 정교화)가 모두 성공에 필수적임을 확인했습니다. 특히 큰 손실 증가를 보이는 저엔트로피 토큰에 집중하는 특정 보상 설계가 가장 효과적인 것으로 나타났습니다.

의의 및 주장

본 논문은 ForesightKV가 긴 컨텍스트 추론을 위한 적응형 KV 캐시 관리에 있어 유망한 새로운 방향을 제시한다고 주장합니다. 규칙 기반 휴리스틱에서 벗어나 KV 쌍의 장기적 기여도를 명시적으로 모델링하고 저엔로피 토큰의 보존을 목표로 하는 학습 기반 접근 방식으로 전환함으로써, 이 방법은 메모리 효율성과 생성 품질 사이의 우수한 트레이드오프를 달성합니다. 저자들은 자신들의 방식이 특정 예산 제약에 과적합되지 않고 KV 캐시 항목의 본질적인 중요성을 포착함으로써, 다양한 예산 설정과 작업 도메인 전반에 걸쳐 잘 일반화될 수 있음을 강조합니다.