Les modèles de langage à attention linéaire simple équilibrent le compromis entre rappel et débit.
Simple linear attention language models balance the recall-throughput tradeoff
February 28, 2024
Auteurs: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher Ré
cs.AI
Résumé
Des travaux récents ont montré que les modèles de langage basés sur l'attention excellent en rappel, c'est-à-dire la capacité à ancrer les générations dans des tokens précédemment vus dans le contexte. Cependant, l'efficacité des modèles basés sur l'attention est limitée lors de l'inférence par la consommation mémoire agressive du cache KV. Dans ce travail, nous explorons si nous pouvons améliorer l'efficacité des modèles de langage (par exemple, en réduisant la consommation mémoire) sans compromettre le rappel. En appliquant des expériences et des théories à un large ensemble d'architectures, nous identifions un compromis clé entre la taille de l'état d'un modèle et sa capacité de rappel. Nous montrons que les alternatives efficaces à l'attention (par exemple, H3, Mamba, RWKV) maintiennent un état récurrent de taille fixe, mais peinent à effectuer un rappel efficace. Nous proposons BASED, une architecture simple combinant une attention linéaire et une attention par fenêtre glissante. En faisant varier la taille de la fenêtre de BASED et la dimension des caractéristiques de l'attention linéaire, nous pouvons ajuster la taille de l'état et parcourir la frontière de Pareto de la courbe de compromis rappel-mémoire, retrouvant la qualité complète de l'attention à une extrémité et la petite taille d'état des alternatives à l'attention à l'autre. Nous entraînons des modèles de langage jusqu'à 1,3 milliard de paramètres et montrons que BASED correspond aux meilleurs modèles sous-quadratiques (par exemple, Mamba) en termes de perplexité et les surpasse sur des tâches intensives en rappel du monde réel avec une précision supérieure de 6,22 points. Les implémentations de l'attention linéaire sont souvent moins efficaces que les implémentations optimisées de l'attention standard. Pour rendre BASED compétitif, nous développons des algorithmes conscients des E/S qui permettent un débit 24 fois supérieur à celui de FlashAttention-2 lors de la génération de 1024 tokens avec des modèles de 1,3 milliard de paramètres. Le code de ce travail est disponible à l'adresse suivante : https://github.com/HazyResearch/based.
English
Recent work has shown that attention-based language models excel at recall,
the ability to ground generations in tokens previously seen in context.
However, the efficiency of attention-based models is bottle-necked during
inference by the KV-cache's aggressive memory consumption. In this work, we
explore whether we can improve language model efficiency (e.g. by reducing
memory consumption) without compromising on recall. By applying experiments and
theory to a broad set of architectures, we identify a key tradeoff between a
model's state size and recall ability. We show that efficient alternatives to
attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but
struggle at recall. We propose BASED a simple architecture combining linear and
sliding window attention. By varying BASED window size and linear attention
feature dimension, we can dial the state size and traverse the pareto frontier
of the recall-memory tradeoff curve, recovering the full quality of attention
on one end and the small state size of attention-alternatives on the other. We
train language models up to 1.3b parameters and show that BASED matches the
strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them
on real-world recall-intensive tasks by 6.22 accuracy points. Implementations
of linear attention are often less efficient than optimized standard attention
implementations. To make BASED competitive, we develop IO-aware algorithms that
enable 24x higher throughput on language generation than FlashAttention-2, when
generating 1024 tokens using 1.3b parameter models. Code for this work is
provided at: https://github.com/HazyResearch/based.