Grouped Query Attention (GQA)

La atención de consultas agrupadas (GQA) es un tipo de atención multicabeza para Transformers en el que grupos de cabezas de consulta comparten una cabeza de clave y otra de valor. Por ejemplo, Llama 3.1 8B combina 32 cabezas de consulta con 8 cabezas de clave-valor en cada capa. Se sitúa entre la atención multicabeza (MHA), donde cada cabeza de consulta tiene sus propias claves y valores, y la atención multi-consulta (MQA), donde todas las cabezas de consulta comparten una única cabeza de clave y otra de valor. Joshua Ainslie y sus colegas de Google Research la presentaron en GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, en EMNLP 2023, y demostraron que los modelos GQA alcanzan una calidad cercana a MHA y una velocidad de decodificación cercana a MQA. La usan modelos de lenguaje grandes (LLM) de código abierto como Llama 2 70B, la familia Llama 3, Mistral 7B y Gemma 2. También la heredan en sus decodificadores de lenguaje los modelos de visión y lenguaje (VLM) basados en ellos, como PaliGemma 2 de Google, basado en Gemma 2.

Cómo funciona GQA#

En la autoatención, cada token se proyecta en una consulta, una clave y un valor para cada cabeza, y cada cabeza compara sus consultas con sus claves. GQA conserva el conjunto completo de cabezas de consulta, pero proyecta menos cabezas de clave y valor; cada grupo de cabezas de consulta usa las claves y valores que comparte con su grupo. Con 8 cabezas de consulta y 2 cabezas de clave-valor, cada cabeza de clave-valor sirve a 4 cabezas de consulta:

El cálculo de atención de cada cabeza sigue siendo la atención de producto escalar escalado sin cambios. Con un grupo, la capa se convierte en MQA; con tantos grupos como cabezas de consulta, se convierte en MHA. PyTorch admite directamente esta disposición mediante el argumento enable_gqa de scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Imprime Output (1, 8, 16, 64), same as copied heads: True y 256 with GQA, 1024 with MHA: las 8 cabezas de consulta generan salidas con una cuarta parte de las claves y valores que almacenaría MHA.

Por qué GQA reduce la caché KV#

Durante la generación de texto, un modelo conserva las claves y los valores de todos los tokens anteriores en su caché KV y lee toda la caché para cada token nuevo, así que la velocidad de decodificación depende más del ancho de banda de memoria que de los cálculos. El tamaño de la caché aumenta con el número de cabezas de clave-valor. Llama 3.1 8B tiene 32 capas, cada una con 8 cabezas de clave-valor de 128 dimensiones, así que almacena 128 KiB por token con precisión de 16 bits, o 16 GiB para su ventana de contexto completa de 131.072 tokens. Con MHA, el mismo modelo necesitaría 512 KiB por token y 64 GiB para la longitud máxima. La memoria liberada se puede dedicar a contextos más largos, lotes más grandes para el procesamiento continuo por lotes y menor latencia de inferencia por token. La interfaz LLM de Ultralytics acepta un base_url para cualquier servidor compatible con OpenAI, como una instancia autohospedada de Llama 3.1 8B. Sus 8 cabezas de clave-valor permiten al servidor mantener cuatro veces más tokens de conversación en una cantidad determinada de memoria de caché que una disposición MHA de 32 cabezas.

Conversión de modelos a GQA#

GQA no requiere entrenar desde cero. Los autores convirtieron un punto de control MHA existente calculando la media de las cabezas de clave y valor de cada grupo, y después continuaron el preentrenamiento con aproximadamente el 5 % del cómputo original, un proceso que denominan entrenamiento adicional. MQA, el caso extremo con una sola cabeza compartida, fue presentada por Noam Shazeer en 2019 y es la disposición más rápida, pero reducía la calidad y podía desestabilizar el entrenamiento. Un número intermedio de grupos recupera la mayor parte de esa calidad.

GQA y técnicas relacionadas#

• Atención latente multicabeza (MLA): Presentada con DeepSeek-V2, comprime las claves y los valores de cada token en un pequeño vector latente compartido.
• PagedAttention: Asigna la caché KV en bloques de tamaño fijo para evitar la fragmentación, independientemente del número de cabezas que la ocupen.
• Atención de ventana deslizante: Limita el número de tokens almacenados en caché por capa; Mistral 7B la combina con GQA.
• Atención Flash: Kernel de GPU que calcula la atención exacta con mayor rapidez y admite de forma nativa las disposiciones GQA.