SubGen : Génération de jetons en temps et mémoire sous-linéaires

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

Auteurs: Amir Zandieh, Insu Han, Vahab Mirrokni, Amin Karbasi

cs.AI

Résumé

Malgré le succès significatif des grands modèles de langage (LLMs), leurs importants besoins en mémoire posent des défis pour leur déploiement dans la génération de tokens à contexte long. L'empreinte mémoire substantielle des décodeurs de LLMs découle de la nécessité de stocker tous les tokens précédents dans le module d'attention, une exigence imposée par la mise en cache des paires clé-valeur (KV). Dans ce travail, nous nous concentrons sur le développement d'une technique de compression efficace pour le cache KV. Des preuves empiriques indiquent une tendance significative au regroupement au sein des embeddings de clés dans le module d'attention. En nous appuyant sur cette observation clé, nous avons conçu une nouvelle méthode de mise en cache avec une complexité sous-linéaire, utilisant un clustering en ligne sur les tokens de clés et un échantillonnage ell_2 en ligne sur les valeurs. Le résultat est un algorithme de décodage d'attention à la fois précis et efficace, nommé SubGen. Non seulement cet algorithme garantit une empreinte mémoire sous-linéaire et une complexité temporelle sous-linéaire, mais nous établissons également une borne d'erreur serrée pour notre approche. Les évaluations empiriques sur des tâches de réponse à des questions à contexte long démontrent que SubGen surpasse significativement les méthodes existantes et de pointe de compression du cache KV en termes de performance et d'efficacité.

English

Despite the significant success of large language models (LLMs), their extensive memory requirements pose challenges for deploying them in long-context token generation. The substantial memory footprint of LLM decoders arises from the necessity to store all previous tokens in the attention module, a requirement imposed by key-value (KV) caching. In this work, our focus is on developing an efficient compression technique for the KV cache. Empirical evidence indicates a significant clustering tendency within key embeddings in the attention module. Building on this key insight, we have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online ell_2 sampling on values. The result is a provably accurate and efficient attention decoding algorithm, termed SubGen. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach. Empirical evaluations on long-context question-answering tasks demonstrate that SubGen significantly outperforms existing and state-of-the-art KV cache compression methods in terms of performance and efficiency.