Какие заголовки важны для рассуждений? Сжатие кэша ключей и значений с использованием обучения с подкреплением

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

October 9, 2025

Авторы: Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang

cs.AI

Аннотация

Крупные языковые модели, демонстрирующие сложные рассуждения, реализуют их через расширенную генерацию цепочки мыслей, что создает беспрецедентную нагрузку на кэш ключей-значений (KV) на этапе декодирования. Существующие методы сжатия KV-кэша показывают низкую эффективность для моделей рассуждений: методы отбрасывания токенов нарушают целостность рассуждений, удаляя критически важную информацию, а методы перераспределения голов ошибочно сжимают головы, важные для рассуждений, поскольку они разработаны для задач поиска, что приводит к значительному ухудшению производительности с увеличением степени сжатия. Мы предполагаем, что головы KV в моделях рассуждений демонстрируют функциональную гетерогенность: некоторые головы критически важны для согласованности цепочки мыслей, в то время как другие поддаются сжатию. Чтобы проверить и использовать это наблюдение, мы предлагаем RLKV — новый фреймворк для идентификации голов, критически важных для рассуждений, который использует обучение с подкреплением для прямой оптимизации связи между использованием кэша каждой головы и качеством рассуждений. Поскольку RLKV генерирует награды на основе реальных сгенерированных образцов во время обучения, он естественным образом идентифицирует головы, связанные с поведением рассуждений. Затем мы выделяем полный KV-кэш для этих голов, применяя сжатый постоянный KV-кэш к остальным для эффективного вывода. Наши эксперименты показывают, что лишь небольшая доля голов внимания критически важна для рассуждений, что позволяет нашему подходу к сжатию 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.