Riduzione della dimensione della cache chiave-valore nei Transformer tramite attenzione cross-layer

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

May 21, 2024

Autori: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly

cs.AI

Abstract

La memorizzazione chiave-valore (KV) svolge un ruolo essenziale nell'accelerazione del processo di decodifica per i modelli linguistici di grandi dimensioni (LLM) autoregressivi basati su trasformatori. Tuttavia, la quantità di memoria necessaria per memorizzare la cache KV può diventare proibitiva per sequenze lunghe e dimensioni di batch elevate. Dalla creazione del trasformatore, due degli interventi più efficaci scoperti per ridurre le dimensioni della cache KV sono stati l'attenzione multi-query (MQA) e la sua generalizzazione, l'attenzione a query raggruppate (GQA). Sia MQA che GQA modificano il design del blocco di attenzione in modo che più teste di query possano condividere una singola testa chiave/valore, riducendo il numero di teste chiave/valore distinte di un fattore significativo, pur degradando solo minimamente l'accuratezza. In questo articolo, dimostriamo che è possibile spingere ulteriormente l'attenzione multi-query condividendo anche le teste chiave e valore tra strati adiacenti, ottenendo un nuovo design di attenzione che chiamiamo Attenzione Cross-Layer (CLA). Con CLA, abbiamo scoperto che è possibile ridurre ulteriormente le dimensioni della cache KV di un altro 2x mantenendo quasi la stessa accuratezza di MQA non modificato. Negli esperimenti di addestramento di modelli da 1B e 3B parametri da zero, dimostriamo che CLA offre un miglioramento di Pareto rispetto ai compromessi memoria/accuratezza possibili con il tradizionale MQA, consentendo inferenze con sequenze più lunghe e dimensioni di batch più grandi rispetto a quanto sarebbe altrimenti possibile.

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