SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation (2412.13649v3)

Abstract: Key-Value (KV) cache has become a bottleneck of LLMs for long-context generation. Despite the numerous efforts in this area, the optimization for the decoding phase is generally ignored. However, we believe such optimization is crucial, especially for long-output generation tasks based on the following two observations: (i) Excessive compression during the prefill phase, which requires specific full context impairs the comprehension of the reasoning task; (ii) Deviation of heavy hitters occurs in the reasoning tasks with long outputs. Therefore, SCOPE, a simple yet efficient framework that separately performs KV cache optimization during the prefill and decoding phases, is introduced. Specifically, the KV cache during the prefill phase is preserved to maintain the essential information, while a novel strategy based on sliding is proposed to select essential heavy hitters for the decoding phase. Memory usage and memory transfer are further optimized using adaptive and discontinuous strategies. Extensive experiments on LongGenBench show the effectiveness and generalization of SCOPE and its compatibility as a plug-in to other prefill-only KV compression methods.

• The paper introduces SCOPE, a framework optimizing key-value cache compression for large language models by decoupling management during prefill and decoding phases.
• SCOPE employs a novel sliding strategy during decoding to dynamically select essential heavy hitters and mitigate memory pressure through adaptive and discontinuous methods.
• Experiments show SCOPE achieves near-full performance equivalence with a 35% compression rate, demonstrating its memory efficiency and adaptability as a plug-in solution.

SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation

The paper introduces SCOPE, a framework designed to enhance the efficiency of key-value (KV) cache compression in LLMs with a focus on long-context tasks. The bottleneck of KV cache in such scenarios is a well-recognized issue, leading to significant memory consumption, particularly affecting long-output generation tasks.

Key Observations and Motivation

The authors identify two critical inefficiencies: the impact of excessive compression during the prefill phase on comprehension, and the deviation of heavy hitters during the decoding phase. Heavy hitters, which represent critical information captured by pivotal tokens, tend to deviate in tasks requiring long outputs. These deviations suggest the necessity of a differentiated approach to KV cache management.

SCOPE Framework

SCOPE proposes a decoupled compression strategy addressing the prefill and decoding phases separately:

1. Prefill Phase: Essential information is preserved by maintaining the KV cache. This phase ensures that crucial context needed for task comprehension remains intact.
2. Decoding Phase: A novel sliding strategy is introduced to dynamically select essential heavy hitters. Adaptive and discontinuous strategies are employed to optimize memory usage and transfer further. These strategies help mitigate memory pressure by dynamically and efficiently reallocating cache resources during decoding.

Numerical Results

Experiments on LongGenBench demonstrate SCOPE's effectiveness across long-context tasks, confirming its adaptability as a plug-in solution to existing prefill-only methods. SCOPE achieves near-full performance equivalence with a 35% compression rate, providing strong evidence of its memory efficiency. The experiments substantiate that both the adaptive and discontinuous strategies outperform previous methods by maintaining a balanced allocation between the KV cache generated at different phases.

Implications and Future Directions

The introduction of SCOPE has both practical and theoretical implications:

• Practical Implications: SCOPE alleviates the bottleneck of GPU memory usage in LLM inference, making it more feasible to deploy such models in real-time applications dealing with long-context scenarios, such as detailed question-answering systems or complex summarization tasks.
• Theoretical Implications: The findings suggest potential avenues for further improving how attention mechanisms can be aligned with the memory architecture to enhance efficiency without compromising task performance. Future research can explore more dynamic strategies for cache management, extending beyond LLMs, to other transformer-based models or even multi-modal settings.

Conclusion

In summary, SCOPE represents a significant advancement in the management of KV caches for long-context LLMs by effectively managing them through phase separation and strategic optimization. This work offers a defensible step toward resolving KV cache inefficiencies and opens a corridor for further innovation in memory optimization strategies in AI models.