ThinK : Cache de clés allégé par élagage piloté par les requêtes

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

Auteurs: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

cs.AI

Résumé

Les modèles de langage de grande taille (LLMs) ont révolutionné le domaine du traitement automatique du langage naturel, atteignant des performances sans précédent dans une variété d'applications grâce à l'augmentation de la taille des modèles et de la longueur des séquences. Cependant, l'augmentation associée des coûts de calcul et de mémoire pose des défis importants, notamment dans la gestion des longues séquences en raison de la complexité quadratique du mécanisme d'attention des transformateurs. Cet article se concentre sur le scénario de contexte long, en abordant les inefficacités dans la consommation de mémoire du cache KV lors de l'inférence. Contrairement aux approches existantes qui optimisent la mémoire en fonction de la longueur des séquences, nous révélons que la dimension des canaux du cache KV présente une redondance significative, caractérisée par une distribution de magnitude déséquilibrée et une structure de faible rang dans les poids d'attention. Sur la base de ces observations, nous proposons ThinK, une nouvelle méthode d'élagage du cache KV dépendante des requêtes, conçue pour minimiser la perte des poids d'attention tout en élaguant sélectivement les canaux les moins significatifs. Notre approche non seulement maintient ou améliore la précision du modèle, mais réduit également les coûts de mémoire de plus de 20 % par rapport aux méthodes classiques d'éviction du cache KV. Des évaluations approfondies sur les modèles LLaMA3 et Mistral à travers divers ensembles de données de longues séquences confirment l'efficacité de ThinK, établissant un nouveau précédent pour le déploiement efficace des LLMs sans compromettre les performances. Nous esquissons également le potentiel d'étendre notre méthode à l'élagage du cache de valeurs, démontrant la polyvalence et la large applicabilité de ThinK dans la réduction des surcharges de mémoire et de calcul.

English

Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.