Reduzindo o Tamanho do Cache de Chave-Valor em Transformers com Atenção Inter-Camadas
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
May 21, 2024
Autores: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly
cs.AI
Resumo
O cache de chave-valor (KV) desempenha um papel essencial na aceleração da decodificação de modelos de linguagem grandes (LLMs) autoregressivos baseados em transformadores. No entanto, a quantidade de memória necessária para armazenar o cache KV pode se tornar proibitiva em sequências longas e grandes tamanhos de lote. Desde a invenção do transformador, duas das intervenções mais eficazes descobertas para reduzir o tamanho do cache KV foram a Atenção Multi-Query (MQA) e sua generalização, a Atenção de Grupo-Query (GQA). Tanto MQA quanto GQA modificam o design do bloco de atenção para que múltiplas cabeças de consulta compartilhem uma única cabeça de chave/valor, reduzindo o número de cabeças de chave/valor distintas por um grande fator, enquanto degradam a precisão apenas minimamente. Neste artigo, mostramos que é possível levar a Atenção Multi-Query um passo adiante, compartilhando também cabeças de chave e valor entre camadas adjacentes, resultando em um novo design de atenção que chamamos de Atenção Trans-Camada (CLA). Com a CLA, descobrimos que é possível reduzir o tamanho do cache KV em mais 2x, mantendo quase a mesma precisão da MQA não modificada. Em experimentos treinando modelos de 1B e 3B parâmetros do zero, demonstramos que a CLA proporciona uma melhoria de Pareto nas compensações entre memória/precisão que são possíveis com a MQA tradicional, permitindo inferência com sequências mais longas e tamanhos de lote maiores do que seria possível de outra forma.
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