SnapKV : Les LLM savent ce que vous cherchez avant la génération

SnapKV: LLM Knows What You are Looking for Before Generation

April 22, 2024

Auteurs: Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, Deming Chen

cs.AI

Résumé

Les modèles de langage de grande taille (LLMs) ont réalisé des progrès remarquables dans le traitement de contextes étendus, avec le cache clé-valeur (KV) jouant un rôle essentiel dans l'amélioration de leurs performances. Cependant, la croissance du cache KV en réponse à l'augmentation de la longueur des entrées pose des défis en termes d'efficacité mémoire et temporelle. Pour résoudre ce problème, cet article présente SnapKV, une approche innovante et sans ajustement fin qui minimise efficacement la taille du cache KV tout en offrant des performances comparables dans des applications réelles. Nous découvrons que chaque tête d'attention du modèle se concentre de manière constante sur des caractéristiques spécifiques de l'attention du prompt pendant la génération. Par ailleurs, ce motif robuste peut être obtenu à partir d'une fenêtre d'« observation » située à la fin des prompts. En s'appuyant sur cette observation, SnapKV compresse automatiquement les caches KV en sélectionnant des positions KV importantes et regroupées pour chaque tête d'attention. Notre approche réduit significativement la surcharge computationnelle croissante et l'empreinte mémoire lors du traitement de longues séquences d'entrée. Plus précisément, SnapKV atteint une vitesse de décodage constante avec une augmentation de 3,6x de la vitesse de génération et une amélioration de 8,2x de l'efficacité mémoire par rapport à la référence lors du traitement d'entrées de 16K tokens. Parallèlement, il maintient des performances comparables aux modèles de référence sur 16 ensembles de données de séquences longues. De plus, SnapKV peut traiter jusqu'à 380K tokens de contexte sur un seul GPU A100-80GB en utilisant l'implémentation HuggingFace avec des modifications mineures, ne montrant qu'une baisse de précision négligeable dans le test Needle-in-a-Haystack. Des études complémentaires approfondies suggèrent le potentiel de SnapKV pour des applications pratiques.

English

Large Language Models (LLMs) have made remarkable progress in processing extensive contexts, with the Key-Value (KV) cache playing a vital role in enhancing their performance. However, the growth of the KV cache in response to increasing input length poses challenges to memory and time efficiency. To address this problem, this paper introduces SnapKV, an innovative and fine-tuning-free approach that efficiently minimizes KV cache size while still delivering comparable performance in real-world applications. We discover that each attention head in the model consistently focuses on specific prompt attention features during generation. Meanwhile, this robust pattern can be obtained from an `observation' window located at the end of the prompts. Drawing on this insight, SnapKV automatically compresses KV caches by selecting clustered important KV positions for each attention head. Our approach significantly reduces the growing computational overhead and memory footprint when processing long input sequences. Specifically, SnapKV achieves a consistent decoding speed with a 3.6x increase in generation speed and an 8.2x enhancement in memory efficiency compared to baseline when processing inputs of 16K tokens. At the same time, it maintains comparable performance to baseline models across 16 long sequence datasets. Moreover, SnapKV can process up to 380K context tokens on a single A100-80GB GPU using HuggingFace implementation with minor changes, exhibiting only a negligible accuracy drop in the Needle-in-a-Haystack test. Further comprehensive studies suggest SnapKV's potential for practical applications.