장문 추론을 위한 정보 기반 키-값 캐시 압축

Information-Aware KV Cache Compression for Long Reasoning

대규모 언어 모델(LLM)의 추론 능력은 빠르게 발전하고 있으며, 이로 인해 사전 학습 및 디코딩 단계에서 키-값(KV) 캐시의 크기가 점점 커지고 있습니다. 기존의 KV 캐시 압축 방법은 주로 어텐션 가중치를 사용하여 토큰의 중요도를 추정합니다. 어텐션은 문맥적 관련성을 효과적으로 파악하지만, 예측 불확실성과 토큰 정보량과 관련된 상호 보완적인 정보 이론적 신호를 간과합니다. 본 논문에서는 토큰의 중요도를 미래 지향적인 관점에서 재검토하고, 압축된 토큰이 향후 문맥에 미치는 영향을 측정하는 지표인 "Forward Influence"를 소개합니다. 분석 결과, 어텐션 점수에 의해 선택되는 토큰은 주로 인접한 문맥에 영향을 미치는 반면, 높은 예측 불확실성을 나타내는 토큰은 먼 미래의 문맥에 훨씬 더 큰 영향을 미칩니다. 이러한 관찰을 바탕으로, 정보 이론적 신호를 통합하는 엔트로피 기반 KV 캐시 압축 프레임워크인 "InfoKV"를 제안합니다. InfoKV는 토큰 수준의 예측 불확실성과 레이어별 표현 변화를 결합하여 생성된 엔트로피 점수를 추론 과정에서 어텐션 점수와 함께 활용합니다. Llama-3.1, Llama-3.2 및 DeepSeek-R1을 사용한 장문 추론 벤치마크 실험 결과, InfoKV는 기존의 어텐션 기반 KV 압축 방법보다 장문 사전 학습 및 디코딩 시나리오 모두에서 일관되게 우수한 성능을 보였습니다.

Reasoning capability has advanced rapidly in large language models (LLMs), leading to an increasing size of key-value (KV) cache in both prefilling and decoding stages. Existing KV cache compression methods mainly rely on attention weights to estimate token importance. While attention effectively captures contextual relevance, it overlooks complementary information-theoretic signals related to predictive uncertainty and token informativeness. In this paper, we revisit token importance from a forward-looking perspective and introduce \textit{Forward Influence}, a metric that measures how compressed tokens affect future contexts. Our analysis reveals that tokens selected by attention scores mainly influence nearby contexts, whereas tokens associated with high predictive uncertainty exhibit substantially stronger influence on distant future contexts. Based on the observation, we propose \textbf{InfoKV}, an entropy-aware KV cache compression framework that incorporates information-theoretic signals. It combines token-level predictive uncertainty with layer-wise representation evolution and integrates the resulting entropy scores with attention scores during reasoning. Experiments on long-context reasoning benchmarks with Llama-3.1, Llama-3.2, and DeepSeek-R1 demonstrate that InfoKV consistently outperforms existing attention-based KV compression methods in both long prefilling and decoding scenarios.