Réduction de la taille du cache clé-valeur des Transformers grâce à l'attention inter-couches
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
May 21, 2024
Auteurs: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly
cs.AI
Résumé
La mise en cache clé-valeur (KV) joue un rôle essentiel dans l'accélération du décodage pour les grands modèles de langage (LLM) autoregressifs basés sur les transformateurs. Cependant, la quantité de mémoire nécessaire pour stocker le cache KV peut devenir prohibitive pour des longueurs de séquence importantes et des tailles de lot élevées. Depuis l'invention du transformateur, deux des interventions les plus efficaces découvertes pour réduire la taille du cache KV ont été l'attention multi-requête (MQA) et sa généralisation, l'attention par requêtes groupées (GQA). MQA et GQA modifient toutes deux la conception du bloc d'attention afin que plusieurs têtes de requête puissent partager une seule tête clé/valeur, réduisant ainsi le nombre de têtes clé/valeur distinctes d'un facteur important tout en ne dégradant que minimalement la précision. Dans cet article, nous montrons qu'il est possible d'aller encore plus loin avec l'attention multi-requête en partageant également les têtes clé et valeur entre les couches adjacentes, ce qui donne lieu à une nouvelle conception d'attention que nous appelons l'attention inter-couches (CLA). Avec CLA, nous constatons qu'il est possible de réduire la taille du cache KV d'un facteur supplémentaire de 2 tout en maintenant une précision presque identique à celle de MQA non modifiée. Dans des expériences de formation de modèles de 1 milliard et 3 milliards de paramètres à partir de zéro, nous démontrons que CLA offre une amélioration de Pareto par rapport aux compromis mémoire/précision possibles avec la MQA traditionnelle, permettant ainsi une inférence avec des longueurs de séquence plus longues et des tailles de lot plus importantes que ce qui serait autrement possible.
English
Key-value (KV) caching plays an essential role in accelerating decoding for
transformer-based autoregressive large language models (LLMs). However, the
amount of memory required to store the KV cache can become prohibitive at long
sequence lengths and large batch sizes. Since the invention of the transformer,
two of the most effective interventions discovered for reducing the size of the
KV cache have been Multi-Query Attention (MQA) and its generalization,
Grouped-Query Attention (GQA). MQA and GQA both modify the design of the
attention block so that multiple query heads can share a single key/value head,
reducing the number of distinct key/value heads by a large factor while only
minimally degrading accuracy. In this paper, we show that it is possible to
take Multi-Query Attention a step further by also sharing key and value heads
between adjacent layers, yielding a new attention design we call Cross-Layer
Attention (CLA). With CLA, we find that it is possible to reduce the size of
the KV cache by another 2x while maintaining nearly the same accuracy as
unmodified MQA. In experiments training 1B- and 3B-parameter models from
scratch, we demonstrate that CLA provides a Pareto improvement over the
memory/accuracy tradeoffs which are possible with traditional MQA, enabling
inference with longer sequence lengths and larger batch sizes than would
otherwise be possible