Découvrir les joyaux dans les premières couches : Accélérer les LLM à long context avec une réduction de 1000 fois du nombre de jetons en entrée

Discovering the Gems in Early Layers: Accelerating Long-Context LLMs with 1000x Input Token Reduction

September 25, 2024

Auteurs: Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, Shafiq Joty

cs.AI

Résumé

Les grands modèles de langage (LLM) ont démontré des capacités remarquables pour traiter de longues entrées contextuelles, mais cela se fait au détriment de ressources computationnelles et de latence accrues. Notre recherche présente une approche novatrice pour surmonter le goulot d'étranglement du contexte long afin d'accélérer l'inférence des LLM et réduire la consommation de mémoire GPU. Notre recherche montre que les LLM peuvent identifier les jetons pertinents dans les premières couches avant de générer des réponses à une requête. En exploitant cette observation, nous proposons un algorithme qui utilise les premières couches d'un LLM comme filtres pour sélectionner et compresser les jetons d'entrée, réduisant ainsi significativement la longueur du contexte pour le traitement ultérieur. Notre méthode, GemFilter, présente des améliorations substantielles à la fois en termes de vitesse et d'efficacité mémoire par rapport aux techniques existantes, telles que l'attention standard et SnapKV/H2O. Notamment, elle permet d'obtenir une accélération de 2,4 fois et une réduction de 30\% de l'utilisation de la mémoire GPU par rapport aux méthodes de pointe. L'évaluation sur la tâche Needle in a Haystack montre que GemFilter surpasse significativement l'attention standard, SnapKV et démontre des performances comparables sur le défi LongBench. GemFilter est simple, ne nécessite pas d'entraînement, et est largement applicable à différents LLM. De manière cruciale, elle offre une interprétabilité en permettant aux humains d'inspecter la séquence d'entrée sélectionnée. Ces résultats offrent non seulement des avantages pratiques pour le déploiement des LLM, mais améliorent également notre compréhension des mécanismes internes des LLM, ouvrant la voie à de nouvelles optimisations dans la conception et l'inférence des LLM. Notre code est disponible sur https://github.com/SalesforceAIResearch/GemFilter.

English

Large Language Models (LLMs) have demonstrated remarkable capabilities in handling long context inputs, but this comes at the cost of increased computational resources and latency. Our research introduces a novel approach for the long context bottleneck to accelerate LLM inference and reduce GPU memory consumption. Our research demonstrates that LLMs can identify relevant tokens in the early layers before generating answers to a query. Leveraging this insight, we propose an algorithm that uses early layers of an LLM as filters to select and compress input tokens, significantly reducing the context length for subsequent processing. Our method, GemFilter, demonstrates substantial improvements in both speed and memory efficiency compared to existing techniques, such as standard attention and SnapKV/H2O. Notably, it achieves a 2.4times speedup and 30\% reduction in GPU memory usage compared to SOTA methods. Evaluation on the Needle in a Haystack task shows that GemFilter significantly outperforms standard attention, SnapKV and demonstrates comparable performance on the LongBench challenge. GemFilter is simple, training-free, and broadly applicable across different LLMs. Crucially, it provides interpretability by allowing humans to inspect the selected input sequence. These findings not only offer practical benefits for LLM deployment, but also enhance our understanding of LLM internal mechanisms, paving the way for further optimizations in LLM design and inference. Our code is available at https://github.com/SalesforceAIResearch/GemFilter.