OBJECTIF : Optimisation de la compression du cache clé-valeur dans la génération de contexte long.

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

December 18, 2024

Auteurs: Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan He, Deyu Zhou

cs.AI

Résumé

Le cache clé-valeur (KV) est devenu un goulot d'étranglement pour les LLMs dans la génération de contexte long. Malgré les nombreux efforts dans ce domaine, l'optimisation de la phase de décodage est généralement ignorée. Cependant, nous croyons que cette optimisation est cruciale, en particulier pour les tâches de génération de sortie longue, sur la base des deux observations suivantes : (i) une compression excessive pendant la phase de préremplissage, qui nécessite un contexte complet spécifique, nuit à la compréhension de la tâche de raisonnement ; (ii) une déviation des éléments les plus importants se produit dans les tâches de raisonnement avec des sorties longues. Par conséquent, SCOPE, un cadre simple mais efficace qui effectue séparément l'optimisation du cache KV pendant les phases de préremplissage et de décodage, est introduit. Plus précisément, le cache KV pendant la phase de préremplissage est préservé pour maintenir les informations essentielles, tandis qu'une nouvelle stratégie basée sur le glissement est proposée pour sélectionner les éléments les plus importants essentiels pour la phase de décodage. L'utilisation de la mémoire et le transfert de mémoire sont en outre optimisés en utilisant des stratégies adaptatives et discontinues. Des expériences approfondies sur LongGenBench montrent l'efficacité et la généralisation de SCOPE ainsi que sa compatibilité en tant que module complémentaire à d'autres méthodes de compression KV uniquement pour le préremplissage.

English

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.