¿Qué es la atención de consultas agrupadas (GQA)?

Ainslie et al concibieron la atención de consultas agrupadas como una optimización de la atención multicabezal (MHA), el innovador algoritmo de autoatención introducido en el artículo seminal de 2017 "Attention is All You Need" que estableció redes neuronales. Más concretamente, se propuso como una generalización y una aplicación más restringida de la atención a múltiples consultas (MQA), una optimización anterior de la MHA.

Aunque la atención multicabezal estándar catalizó un salto evolutivo en el machine learning, el procesamiento del lenguaje natural (PLN) y la IA generativa, es extremadamente exigente en recursos computacionales y ancho de banda de memoria. A medida que los LLM crecieron y se hicieron más sofisticados, estos requisitos de uso de memoria se convirtieron en un cuello de botella para el progreso, especialmente para los LLM de solo decodificador autorregresivo utilizados para la generación de texto, el resumen y otras tareas de IA generativa.

Las investigaciones posteriores se centraron en las técnicas para mejorar o agilizar la atención multicabezal. Algunas, como la atención flash y la atención en anillo, mejoran la forma en que las GPU utilizadas para entrenar y ejecutar los modelos gestionan los cálculos y el almacenamiento de la memoria. Otros, como GQA y MQA, exploraron cambios en la forma en que las arquitecturas del transformador procesan los tokens.

La atención de consultas agrupadas trata de equilibrar las ventajas y desventajas de la atención estándar a múltiples cabezas y la atención a múltiples consultas. La primera maximiza la precisión a costa de aumentar el ancho de banda de la memoria y reducir la velocidad. La segunda maximiza la velocidad y la eficiencia a expensas de la precisión.

Para comprender cómo la atención de consultas agrupadas optimiza los modelos de transformadores, es importante comprender primero cómo funciona la atención multicabezal en general. Tanto la GQA como la MQA simplemente refinan, en lugar de reemplazar, la metodología principal de la MHA.

La fuerza impulsora detrás de los LLM y otros modelos que utilizan la arquitectura del transformador es la autoatención, un marco para comprender las relaciones entre cada uno de los diferentes tokens en una secuencia. La autoatención permite a un LLM interpretar datos de texto no solo a través de definiciones de referencia estáticas, sino también del contexto proporcionado por otras palabras y frases.

En los LLM autorregresivos utilizados para la generación de texto, el mecanismo de atención ayuda al modelo a predecir el siguiente token de una secuencia determinando a qué tokens anteriores vale la pena "prestar atención" en ese momento. La información de los tokens que se considera más relevante recibe una mayor ponderación de atención, mientras que la información de los tokens que se considera irrelevante recibe una ponderación de atención cercana a 0.

El mecanismo de atención multicabezal que anima los modelos de transformador genera una información contextual rica a través del cálculo de la autoatención muchas veces en paralelo dividiendo las capas de atención en múltiples cabezales de atención.

Los autores de “Attention is All You Need” articularon su mecanismo de atención utilizando la terminología de una base de datos relacional: consultas, claves y valores. Las bases de datos relacionales están diseñadas para simplificar el almacenamiento y la recuperación de datos relevantes: asignan un identificador único ("clave") a cada dato, y cada clave se asocia con un valor correspondiente. El objetivo de una base de datos relacional es hacer coincidir cada consulta con la clave adecuada.

Para cada token de una secuencia, la atención multicabezal requiere la creación de tres vectores.

• Un vector de consulta, Q, que representa la información que el token está "buscando". Por ejemplo, el vector de consulta de un sustantivo podría representar una búsqueda de adjetivos que lo describan.
• Un vector clave, K, que representa la información que contiene el token. Las puntuaciones de alineación, que representan la relevancia del vector clave de cada token con respecto al vector de consulta de cada uno de los demás tokens, se utilizan para calcular los pesos de atención.
• Un vector de valor, V, que representa la información contextual que se actualizará mediante las contribuciones ponderadas por atención de los vectores clave de otros tokens.

Las interacciones matemáticas entre estos tres vectores, mediadas por el mecanismo de atención, son la forma en que un modelo ajusta su comprensión específica del contexto de cada token.
 

Para generar cada uno de estos tres vectores para un token determinado, el modelo comienza con la incrustación del vector original de ese token: una codificación numérica en la que cada dimensión del vector se corresponde con algún elemento abstracto del significado semántico del token. El número de dimensiones en estos vectores es un hiperparámetro predeterminado.

Los vectores Q, K y V de cada token se generan pasando la incrustación del token original a través de una capa lineal que precede a la primera capa de atención. Esta capa lineal se divide en tres matrices únicas de ponderaciones del modelo: W_Q, W_K y W_V. Los valores de peso específicos que contiene se aprenden mediante un preentrenamiento autosupervisado en un conjunto de datos masivo de ejemplos de texto.

Al multiplicar la incrustación del vector original del token por W_Q, W_K y W_V , se obtienen sus correspondientes vector de consulta, vector de clave y vector de valor, respectivamente. El número de dimensiones d que contiene cada vector está determinado por el tamaño de cada matriz de pesos. Q y K tendrán el mismo número de dimensiones, d_k.

Estos tres vectores se pasan a la capa de atención.

En la capa de atención, los vectores Q, K y V se utilizan para calcular una puntuación de alineación entre cada token en cada posición de una secuencia. A continuación, esas puntuaciones de alineación se normalizan en ponderaciones de atención mediante una función softmax.

Para cada token x en una secuencia, las puntuaciones de alineación se obtienen calculando el producto escalar del vector de consulta Q_x de ese token con el vector clave K de cada uno de los demás tokens: en otras palabras, multiplicándolos. Si una relación significativa entre dos tokens se refleja en las similitudes entre sus respectivos vectores, al multiplicarlos se obtiene un valor grande. Si los dos vectores no están alineados, al multiplicarlos se obtiene un valor pequeño o negativo. La mayoría de los modelos de transformadores utilizan una variante llamada atención al producto de puntos escalados, en la que QK se escala, es decir, se multiplica, por $\frac{1}{\sqrt{D_k}}$  para mejorar la estabilidad del entrenamiento.

Estas puntuaciones de alineación de claves de consulta se introducen en una función softmax. Softmax normaliza todas las entradas a un valor entre 0 y 1, de modo que todas sumen 1. Las salidas de la función softmax son las ponderaciones de atención, cada una representa la parte (de 1) de la atención del token x que se debe prestar a cada uno de los otros tokens. Si la ponderación de atención de un token es cercano a 0, se ignorará. Una pondración de atención de 1 significaría que un token recibe toda la atención de x y todos los demás serán ignorados.

Por último, el vector de valor de cada token se multiplica por su ponderación de atención. Estas contribuciones ponderadas por la atención de cada token anterior se promedian juntas y se suman a la incrustación vectorial original del token x. Con esto, la incrustación del token x se ha actualizado para reflejar el contexto que proporcionan los demás tokens de la secuencia que le son relevantes.

A continuación, la incrustación de vectores actualizada se envía a otra capa lineal, con su propia matriz de peso W_Z, donde el vector actualizado al contexto se normaliza a un número constante de dimensiones y, después, se envía a la siguiente capa de atención. Cada capa de atención progresiva capta mayores matices contextuales.
 

Utilizar los promedios de las contribuciones ponderadas por la atención de otros tokens en lugar de contabilizar cada pieza del contexto ponderado por la atención individualmente es matemáticamente eficiente, pero resulta en una pérdida de detalle.

Para compensar, las redes de transformadores dividen la incrustación del token de entrada original en h piezas de tamaño uniforme. También dividen W_Q, W_K y W_V en h subconjuntos llamados cabezales de consulta, cabezales de clave y cabezales de valor, respectivamente. Cada cabezal de consulta, cabezal de clave y cabezal de valor recibe una parte de la incrustación de token original. Los vectores producidos por cada uno de estos tripletes paralelos de cabezales de consulta, cabezales clave y cabezales de valor se introducen en un cabezal de atención correspondiente. Finalmente, las salidas de estos circuitos paralelos h se concatenan de nuevo para actualizar la incrustación de tokens completa.

En el entrenamiento, cada circuito aprende ponderaciones distintas que captan un aspecto separado de los significados semánticos. Esto, a su vez, ayuda al modelo a procesar las distintas formas en que las implicaciones de una palabra pueden verse influidas por el contexto de otras palabras que la rodean.

Lo malo de la atención multicabezal estándar no es tanto la presencia de algún fallo crucial como la falta de optimización. MHA fue el primer algoritmo de su clase y representa la ejecución más compleja de su mecanismo general para el cálculo de la atención.

La mayor parte de la ineficacia de MHA se debe a la abundancia de cálculos y parámetros del modelo. En el MHA estándar, cada cabeza de consulta, cabeza de clave y cabeza de valor en cada bloque de atención tiene su propia matriz de pesos. Así, por ejemplo, un modelo con ocho cabezales de atención en cada capa de atención, mucho menos que la mayoría de los LLM modernos, requeriría 24 matrices de ponderaciones únicas sólo para las cabezas Q, K y V de la capa. Esto implica un gran número de cálculos intermedios en cada capa.

Una consecuencia de esta configuración es que es costosa desde el punto de vista computacional. Los requisitos de cálculo de la MHA varían cuadráticamente con respecto a la longitud de la secuencia: duplicar el número de tokens de una secuencia de entrada requiere cuadruplicar la complejidad. Esto pone límites prácticos al tamaño de las ventanas de contexto.

MHA también ejerce una gran presión sobre la memoria del sistema. Las GPU no disponen de mucha memoria integrada para almacenar las salidas de la ingente cantidad de cálculos intermedios que deben recuperarse en cada paso de procesamiento posterior. En su lugar, estos resultados intermedios se almacenan en la memoria de gran ancho de banda (HBM), que no se encuentra en el propio chip de la GPU. Esto implica una pequeña cantidad de latencia cada vez que las claves y los valores deben leerse de la memoria. A medida que los modelos de transformadores comenzaron a escalar a muchos miles de millones de parámetros, el tiempo y el cálculo necesarios para entrenar y ejecutar la inferencia se convirtieron en un cuello de botella en el rendimiento del modelo.

Para seguir avanzando se necesitaban métodos que redujeran el número de pasos computacionales sin mermar la capacidad de los transformadores para aprender y reproducir patrones lingüísticos intrincadamente complejos. En este contexto se introdujeron la MQA y, posteriormente, la GQA.

La atención de consultas múltiples (MQA) es un mecanismo de atención más eficiente desde el punto de vista computacional que simplifica la atención de varios cabezales para reducir el uso de memoria y los cálculos intermedios. En lugar de entrenar una cabeza de clave y una cabeza de valor únicas para cada cabeza de atención, MQA utiliza un cabezal de clave y un cabezal de valor única en cada capa. Por lo tanto, los vectores clave y los vectores de valor se calculan solo una vez; este conjunto único de vectores clave y de valor se comparte entre todos los cabezales de atención h.

Esta simplificación reduce en gran medida el número de proyecciones lineales que el modelo debe calcular y almacenar en la memoria de gran ancho de banda. Según el documento de 2019 que introdujo la MQA, la MQA permite un almacenamiento de pares clave-valor entre 10 y 100 veces más pequeño (o caché KV) y una inferencia de decodificador 12 veces más rápida. El uso reducido de memoria de la MQA también acelera significativamente el entrenamiento al permitir un tamaño de lote más grande.

A pesar de sus beneficios, la MQA presenta varios inconvenientes inevitables.

• Degradación del rendimiento: como era de esperar, la reducción del número de parámetros únicos y entrenables del modelo reduce la capacidad del modelo para el conocimiento y los matices. La MQA implica una caída significativa en la precisión en comparación con la MHA estándar, lo que lo hace inadecuado para ciertas situaciones y casos de uso.
  
  
• Debe entrenarse desde cero: un modelo entrenado con la MHA estándar no puede adaptarse simplemente a MQA, sino que debe entrenarse con la MQA desde cero. Esto significa que la MQA no puede utilizarse para optimizar los modelos existentes y conlleva un considerable coste de oportunidad cuando se experimenta con la MQA para nuevos modelos.
  
  
• Redundancias en el paralelismo de tensores: uno de los principales beneficios de entrenar modelos de transformadores en GPU es la capacidad de realizar múltiples operaciones de tensores complejas en paralelo. Los valores K y V deben estar presentes en cada nodo del clúster de GPU que realiza estas operaciones, lo que significa que en la práctica deben replicarse para cada nodo. Este no es un uso óptimo de los recursos informáticos a pesar de ser más eficiente que la MHA estándar.

La atención de consultas agrupadas es una formulación más general y flexible de la atención de consultas múltiples que divide los encabezados de consulta en varios grupos que cada uno comparte un conjunto de claves y valores, en lugar de compartir un conjunto de claves y valores en todos los encabezados de consulta.

Tras la publicación de "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" en mayo de 2023, muchos LLM adoptaron rápidamente GQA. Por ejemplo, Meta adoptó por primera vez GQA para sus modelos Llama 2 en julio de 2023 y mantuvo GQA en los modelos Llama 3 lanzados en 2024. Mistral AI utilizó GQA en el modelo Mistral 7B que lanzó en septiembre de 2023. Del mismo modo, los modelosGranite 3.0 de IBM emplean GQA para una inferencia rápida.

En teoría, la GQA puede considerarse una generalización del espectro entre la MHA estándar y la MQA completa. La GQA con el mismo número de grupos de cabezas clave-valor que cabezas de atención es el equivalente de MHA estándar; GQA con un grupo principal es el equivalente de MQA.

En la práctica, la GQA casi siempre implica algún enfoque intermedio, en el que el número de grupos es en sí mismo un hiperparámetro importante.

La atención de consultas agrupadas ofrece varias ventajas que han llevado a su adopción relativamente generalizada por parte de los principales LLM.

• Uso eficiente de la GPU: la distribución de pares clave-valor de GQA se beneficia del paralelismo de tensores, lo que reduce la cantidad de cálculo que se "desperdicia" replicando valores redundantes.
  
  
• Compromiso eficaz: la GQA ofrece un equilibrio ideal entre la velocidad de inferencia del decodificador y la precisión del rendimiento, ya que es casi tan preciso como la MHA y casi tan rápido como la MQA.
  
  
• Sobrecarga de ancho de banda de memoria reducida: como la MQA, la GQA reduce significativamente el número de cálculos intermedios que deben ser computados, almacenados y recuperados en el momento de la inferencia.
  
  
• Formación flexible: a diferencia de MQA, la atención de consultas grupales no requiere que los modelos se entrenen desde cero mediante el enfoque. Los modelos previamente entrenados con la MHA estándar se pueden adaptar para usar la GQA mediante un proceso de ajuste fino llamado "mejora".