Une stratégie simple et efficace basée sur la norme L_2 pour la compression du cache KV

A Simple and Effective L_2 Norm-Based Strategy for KV Cache Compression

June 17, 2024

Auteurs: Alessio Devoto, Yu Zhao, Simone Scardapane, Pasquale Minervini

cs.AI

Résumé

Le déploiement des grands modèles de langage (LLM) est souvent entravé par les exigences mémoire importantes du cache clé-valeur (KV), en particulier à mesure que les longueurs de contexte augmentent. Les approches existantes pour réduire la taille du cache KV impliquent soit un ajustement fin du modèle pour apprendre une stratégie de compression, soit l'utilisation des scores d'attention pour réduire la longueur de la séquence. Nous analysons les distributions d'attention dans les modèles basés sur des Transformers à décodeur uniquement et observons que les schémas d'allocation de l'attention restent cohérents à travers la plupart des couches. Étonnamment, nous trouvons une corrélation claire entre la norme L_2 et les scores d'attention sur les paires KV mises en cache, où une faible norme L_2 d'une clé d'embedding conduit généralement à un score d'attention élevé lors du décodage. Cette découverte indique que l'influence d'une paire KV est potentiellement déterminée par l'embedding de la clé elle-même avant d'être interrogée. Sur la base de cette observation, nous compressons le cache KV en fonction de la norme L_2 des embeddings de clés. Nos résultats expérimentaux montrent que cette stratégie simple peut réduire la taille du cache KV de 50 % sur des tâches de modélisation du langage et de recherche d'une aiguille dans une botte de foin, et de 90 % sur des tâches de récupération de passkey, sans perte de précision.

English

The deployment of large language models (LLMs) is often hindered by the extensive memory requirements of the Key-Value (KV) cache, especially as context lengths increase. Existing approaches to reduce the KV cache size involve either fine-tuning the model to learn a compression strategy or leveraging attention scores to reduce the sequence length. We analyse the attention distributions in decoder-only Transformers-based models and observe that attention allocation patterns stay consistent across most layers. Surprisingly, we find a clear correlation between the L_2 and the attention scores over cached KV pairs, where a low L_2 of a key embedding usually leads to a high attention score during decoding. This finding indicates that the influence of a KV pair is potentially determined by the key embedding itself before being queried. Based on this observation, we compress the KV cache based on the L_2 of key embeddings. Our experimental results show that this simple strategy can reduce the KV cache size by 50% on language modelling and needle-in-a-haystack tasks and 90% on passkey retrieval tasks without losing accuracy.