Atención de consultas agrupadas

El modelo lee el texto anterior de varias formas en paralelo. Cada uno de esos cálculos es una cabeza de atención, o attention. Quieres reducir la memoria que ocupan esas cabezas, pero compartir por completo un único conjunto de datos podría limitar demasiado el modelo. Hay una variante intermedia: algunas cabezas usan un conjunto común y otras uno distinto.

Grouped-Query Attention (GQA) divide las cabezas query en grupos que comparten key y value. Query es una descripción numérica que sirve para leer el contexto, key para establecer correspondencias y value para transmitir información. Cada cabeza sigue calculando sus propios pesos.

Ocho cabezas query divididas en dos grupos de cuatro significan dos conjuntos key/value, no ocho. Un único grupo da MQA; un grupo por cabeza corresponde al MHA habitual.

Con dimensiones compatibles, menos conjuntos reducen el coste de almacenar KV. Es un compromiso entre diseño y calidad, no un interruptor que garantice una velocidad concreta. El tiempo también depende del hardware y de los demás cálculos.

Mecanismo y detalles

Para $H$ cabezas de consulta y $G$ grupos obtenemos $G$ conjuntos de K/V. Cuando $G=1$, es atención multi-query. Cuando $G=H$, volvemos a la atención de múltiples cabezas. Ainslie et al., GQA, §2.2 describen esta familia y la notación GQA-g.

Calcula la memoria antes de adivinar la velocidad

Nuestro cálculo para KV Cache supone batch size 1, 32 capas, 1024 tokens almacenados, una dimensión de 128 para cada cabeza key/value y 2 bytes por número. Solo los elementos K/V ocupan:

$$2\cdot32\cdot1024\cdot G\cdot128\cdot2\ \text{bytes}$$

El primer dos representa K y V. Con ocho cabezas de consulta, el resultado es 128 MiB en MHA, 32 MiB con dos cabezas KV y 16 MiB en MQA. Un MiB son $2^{20}$ bytes. No contamos los pesos del modelo, los búferes, la alineación de memoria ni la replicación entre dispositivos.

Reducir el número de cabezas KV no acorta aquí el contexto ni elimina cabezas de consulta. Limita el número de representaciones K/V separadas. El cálculo de memoria no predice la calidad ni el tiempo de ejecución; estos también dependen del entrenamiento, del hardware y de la implementación.

Al convertir el checkpoint, los autores promediaban las proyecciones de K/V dentro de los grupos y después continuaban el preentrenamiento (§2.1–2.2). Los resultados de §3 corresponden a variantes concretas de T5 y a tareas determinadas. No establecen un número universal de grupos ni garantizan conservar la calidad con solo cambiar la configuración.

Utilizo contenido generado por IA como parte de mi proceso de aprendizaje diario.