장기 추론 모델을 위한 KV 캐시 최적화와 미래 기여 예측

초록

ForesightKV는 장기 추론을 수행하는 대형 언어 모델의 KV 캐시를 효율적으로 관리하기 위해, 미래의 어텐션 영향을 예측하는 스코어링 모델을 학습한다. Golden Eviction 알고리즘으로 최적의 삭제 후보를 라벨링하고, 페어와이즈 랭킹 손실로 지도학습한 뒤, MDP 기반 강화학습(GRPO)으로 저엔트로피 토큰에서 발생하는 손실 급증을 완화한다. 실험 결과, 2K~4K 토큰 예산에서 기존 방법 대비 92%~99% 성능을 유지하면서 메모리 사용량을 절반 수준으로 줄였다.

상세 분석

본 논문은 KV 캐시가 선형적으로 증가하는 문제를 해결하고자, KV 쌍의 장기 기여도를 예측하는 두 단계 학습 프레임워크를 제안한다. 첫 단계인 Golden Eviction은 전체 생성 시퀀스의 미래 어텐션 스코어를 이용해, 각 eviction 시점에서 가장 낮은 미래 블록 스코어를 가진 KV 쌍을 삭제 후보로 선정한다. 이를 위해 어텐션 매트릭스를 고정 길이 블록으로 나누고, 블록‑단위 평균 풀링을 수행해 각 KV 쌍의 미래 스코어를 계산한다. 이렇게 얻은 “골든” 라벨은 페어와이즈 랭킹 손실을 통해 경량 MLP 스코어러(πθ)를 학습시키는 지도학습 데이터가 된다.

두 번째 단계는 KV 캐시 삭제를 마코프 결정 과정(MDP)으로 모델링하고, GRPO(Generalized Reward‑Weighted Policy Optimization) 알고리즘을 적용한다. 여기서 보상은 저엔트로피 토큰이 캐시 삭제 후 겪는 손실 증가량을 MSE 형태로 측정한 값이며, 높은 보상을 받는 정책은 이러한 급격한 손실 상승을 최소화한다. 중요한 점은 LLM 자체 파라미터는 전혀 업데이트되지 않고, 오직 스코어링 모델만 학습된다는 점이다.

실험에서는 Qwen‑3‑4B 기반의 세 가지 추론 모델을 대상으로 AIME2024·2025 수학 벤치마크에서 평가하였다. 캐시 예산을 2K와 4K 토큰으로 제한했을 때, ForesightKV는 각각 원본 성능의 92%와 99%를 유지하면서, 기존 SnapKV·R‑KV 등과 비교해 처리량을 1.8배 이상 향상시켰다. 또한, 저엔트로피 토큰(숫자·기호·엔터티)에서 발생하는 오류를 크게 감소시켜, 장기 추론 과정에서의 누적 오류 전파를 억제하였다.

핵심 인사이트는 (1) KV 쌍의 중요도가 시간에 따라 크게 변동하며, 특히 의미‑의존적 패턴은 블록‑단위 어텐션 구조를 통해 포착할 수 있다는 점, (2) 미래 어텐션 스코어를 라벨링으로 활용하면 지도학습이 효과적으로 수행될 수 있다는 점, (3) 저엔트로피 토큰에 대한 손실 급증을 보상 설계에 포함시키면 강화학습 단계에서 정책이 보다 안정적으로 수렴한다는 점이다. 이러한 설계는 기존 규칙 기반 혹은 단일 단계 학습 방식이 놓치기 쉬운 복합적인 어텐션 의존성을 효과적으로 보완한다.