Reducción del tamaño de la caché clave-valor en Transformers mediante atención entre capas
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
Resumen
El almacenamiento en caché de clave-valor (KV) desempeña un papel esencial en la aceleración de la decodificación para modelos de lenguaje grandes (LLMs) autoregresivos basados en transformadores. Sin embargo, la cantidad de memoria requerida para almacenar la caché KV puede volverse prohibitiva en secuencias largas y tamaños de lote grandes. Desde la invención del transformador, dos de las intervenciones más efectivas descubiertas para reducir el tamaño de la caché KV han sido la Atención de Múltiples Consultas (MQA) y su generalización, la Atención de Consultas Agrupadas (GQA). Tanto MQA como GQA modifican el diseño del bloque de atención para que múltiples cabezas de consulta compartan una sola cabeza de clave/valor, reduciendo el número de cabezas de clave/valor distintas en un gran factor mientras solo degradan mínimamente la precisión. En este artículo, demostramos que es posible llevar la Atención de Múltiples Consultas un paso más allá al compartir también cabezas de clave y valor entre capas adyacentes, lo que da lugar a un nuevo diseño de atención que llamamos Atención entre Capas (CLA). Con CLA, encontramos que es posible reducir el tamaño de la caché KV en otro 2x mientras se mantiene casi la misma precisión que MQA sin modificar. En experimentos de entrenamiento de modelos de 1B y 3B parámetros desde cero, demostramos que CLA proporciona una mejora de Pareto en los compromisos memoria/precisión que son posibles con MQA tradicional, permitiendo inferencias con secuencias más largas y tamaños de lote más grandes de lo que sería posible de otra manera.
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