어떤 헤드가 추론에 중요한가? RL 기반 KV 캐시 압축

초록

추론 능력을 갖춘 대형 언어 모델은 확장된 사고 사슬(chain-of-thought) 생성을 통해 복잡한 추론 행동을 보이며, 디코딩 단계에서 전례 없는 키-값(Key-Value, KV) 캐시 오버헤드를 발생시킵니다. 기존의 KV 캐시 압축 방법들은 추론 모델에서 제대로 작동하지 않습니다: 토큰 드롭핑 방법은 중요한 정보를 제거함으로써 추론의 무결성을 깨뜨리고, 헤드 재할당 방법은 검색 작업을 위해 설계되었기 때문에 추론에 중요한 헤드를 잘못 압축하여 압축률이 증가함에 따라 성능이 크게 저하됩니다. 우리는 KV 헤드가 추론 모델에서 기능적 이질성을 보인다는 가설을 세웁니다—일부 헤드는 사고 사슬의 일관성을 위해 중요하지만, 다른 헤드는 압축 가능합니다. 이 통찰을 검증하고 활용하기 위해, 우리는 RLKV라는 새로운 추론-중요 헤드 식별 프레임워크를 제안합니다. 이 프레임워크는 강화 학습을 사용하여 각 헤드의 캐시 사용과 추론 품질 간의 관계를 직접 최적화합니다. RLKV는 훈련 중 실제 생성된 샘플로부터 보상을 생성하므로, 추론 행동과 관련된 헤드를 자연스럽게 식별합니다. 그런 다음, 이러한 헤드에는 전체 KV 캐시를 할당하고, 다른 헤드에는 압축된 상수 KV 캐시를 적용하여 효율적인 추론을 가능하게 합니다. 우리의 실험 결과, 추론에 필수적인 주의(attention) 헤드는 극히 일부에 불과하며, 이를 통해 우리의 KV 압축 접근법은 기준 방법들을 능가하면서도 20-50%의 캐시 감소를 달성하고 압축되지 않은 결과와 거의 손실 없는 성능을 보여줍니다.

English

Reasoning large language models exhibit complex reasoning behaviors through the extended chain-of-thought generation, creating unprecedented Key-Value (KV) cache overhead during the decoding phase. Existing KV cache compression methods underperform on reasoning models: token-dropping methods break reasoning integrity by discarding critical information, while head-reallocating methods mistakenly compress reasoning-critical heads since they are designed for retrieval tasks, resulting in significant performance degradation as compression rates increase. We hypothesize that KV heads exhibit functional heterogeneity in reasoning models-some heads are critical for chain-of-thought consistency while others are compressible. To validate and exploit this insight, we propose RLKV, a novel reasoning-critical head identification framework, which uses reinforcement learning to directly optimize the relationship between each head's cache usage and reasoning quality. As RLKV produces rewards from actual generated samples during training, it naturally identifies heads relevant to reasoning behaviors. We then allocate full KV cache to these heads while applying compressed constant KV cache to others for efficient inference. Our experiments reveal that only a small fraction of attention heads is essential for reasoning, enabling our KV compression approach to outperform baseline methods while achieving 20-50% cache reduction with near lossless performance compared to uncompressed results.

어떤 헤드가 추론에 중요한가? RL 기반 KV 캐시 압축

Which Heads Matter for Reasoning? RL-Guided KV Cache Compression

초록

Support