Reduzierung der Größe des Transformer Key-Value-Caches mit Cross-Layer Attention
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
May 21, 2024
Autoren: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly
cs.AI
Zusammenfassung
Die Key-Value (KV)-Zwischenspeicherung spielt eine wesentliche Rolle bei der Beschleunigung der Decodierung für auf Transformer basierende autoregressive große Sprachmodelle (LLMs). Allerdings kann der Speicherbedarf für die KV-Zwischenspeicherung bei langen Sequenzlängen und großen Batch-Größen prohibitiv werden. Seit der Erfindung des Transformers wurden zwei der effektivsten Maßnahmen zur Reduzierung der Größe des KV-Zwischenspeichers entdeckt: die Multi-Query-Aufmerksamkeit (MQA) und ihre Verallgemeinerung, die Gruppen-Query-Aufmerksamkeit (GQA). MQA und GQA ändern beide das Design des Aufmerksamkeitsblocks so, dass mehrere Abfrageköpfe einen einzelnen Schlüssel-/Wertekopf gemeinsam nutzen können, wodurch die Anzahl der verschiedenen Schlüssel-/Werteköpfe um einen großen Faktor reduziert wird, während die Genauigkeit nur minimal beeinträchtigt wird. In diesem Papier zeigen wir, dass es möglich ist, die Multi-Query-Aufmerksamkeit einen Schritt weiter zu gehen, indem auch Schlüssel- und Werteköpfe zwischen benachbarten Schichten gemeinsam genutzt werden, was zu einem neuen Aufmerksamkeitsdesign führt, das wir Cross-Layer-Aufmerksamkeit (CLA) nennen. Mit CLA stellen wir fest, dass es möglich ist, die Größe des KV-Zwischenspeichers um das 2-fache zu reduzieren, während die Genauigkeit fast gleich bleibt wie bei unveränderter MQA. In Experimenten zur Schulung von 1B- und 3B-Parametermodellen von Grund auf zeigen wir, dass CLA eine Pareto-Verbesserung gegenüber den Speicher-/Genauigkeits-Kompromissen bietet, die mit traditioneller MQA möglich sind, und die Inferenz mit längeren Sequenzlängen und größeren Batch-Größen ermöglicht, als es sonst möglich wäre.
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