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

Ainslie et al concibieron la atención de consulta agrupada como una optimización de la atención multicabezal (MHA), el innovador algoritmo de autoatención presentado en el documento seminal de 2017 "Attention is All You Need" que estableció redes neuronales transformadoras . Más específicamente, se propuso como una generalización y una aplicación más restringida de la atención multiconsulta (MQA), una optimización anterior de MHA.

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

Las investigaciones posteriores se centraron en técnicas para potenciar o agilizar la atención multicabeza. Algunas, como la atención flash y la atención en anillo, mejoran la forma en que las GPU empleadas para capacitar y ejecutar modelos gestionan los cálculos y el almacenamiento en memoria. Otros, como GQA y MQA, exploraron cambios en la forma en que las arquitecturas transformadoras procesan tokens.

La atención de consultas agrupadas tiene como objetivo equilibrar las compensaciones entre la atención estándar de múltiples cabezas y la atención de múltiples consultas. El primero maximiza la precisión a costa de una mayor sobrecarga de ancho de banda de memoria y una disminución de la velocidad. Este último maximiza la velocidad y la eficiencia a expensas de la precisión.

Para entender cómo la atención agrupada a las consultas optimiza los modelos de transformador, es importante entender primero cómo funciona la atención multicabezal en general. Tanto el GQA como el MQA se limitan a perfeccionar, y no a sustituir, la metodología básica de la MHA.

La fuerza impulsora detrás de los LLM y otros modelos que utilizan la arquitectura transformadora es la autoatención, un marco matemático para comprender las relaciones entre cada uno de los diferentes tokens en una secuencia. La autoatención permite que un LLM interprete los 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 empleados 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 merece más la pena "prestar atención" en ese momento. La información de los tokens que considera más relevantes recibe una mayor pesos de atención, mientras que la información de los tokens considerados irrelevantes recibe pesos de atención cercanos a 0.

El mecanismo de atención de múltiples cabezas que anima los modelos transformadores genera información contextual rica a través del cálculo de la autoatención muchas veces en paralelo al dividir las capas de atención en múltiples cabezas 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 está asociada 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 de múltiples cabezas requiere la creación de 3 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 otros tokens, se utilizan para calcular los pesos de atención.
  
  
• Un vector de valor, V, que representa la información contextual que será actualizada por las contribuciones ponderadas por la atención de los vectores clave de otros token.

Las interacciones matemáticas entre estos 3 vectores, mediadas por el mecanismo de atención, son cómo un modelo ajusta su comprensión específica del contexto de cada token.
 

Para generar cada uno de estos 3 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 para 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 está dividida en tres matrices únicas de pesos del modelo: W_Q, W_K y W_V. Los valores de peso específicos que contiene se aprenden mediante un entrenamiento previo autosupervisado en un conjunto de datos masivo de ejemplos de texto.

Al multiplicar la incrustación del vector original del token por_WQ, WK y WV se obtiene su correspondiente vector de consulta, vector clave y vector de valor, respectivamente. El número de dimensiones d que contiene cada vector viene determinado por el tamaño de cada matriz de pesos. Q y K tendrán el mismo número de dimensiones, d_k.

Estos 3 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. Esas puntuaciones de alineación se normalizan en pesos de atención mediante una función softmax.

Para cada token x en una secuencia, las puntuaciones de alineación se calculan calculando el producto escalar del vector de consulta Q _x de ese token con el vector clave K de cada uno de los otros tokens: en otras palabras, multiplicándolos entre sí. Si una relación significativa entre 2 tokens se refleja en similitudes entre sus respectivos vectores, multiplicarlos arrojará un gran valor. Si los 2 vectores no están alineados, multiplicarlos arrojará un valor pequeño o negativo. La mayoría de los modelos de transformadores utilizan una variante llamada atención de producto escalado , 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 escriben en una función softmax. Softmax normaliza todas las entradas a un valor entre 0 y 1, de modo que todas sumen 1. Los resultados de la función softmax son los pesos de atención, cada uno de los cuales representa la Compartir (de 1) de la atención del token xque se debe prestar a cada uno de los otros tokens. Si el peso de atención de un token es cercano a 0, se ignorará. Un peso 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.

Finalmente, el vector de valor de cada token se multiplica por su peso de atención. Estas contribuciones ponderadas por atención de cada token anterior se promedian juntas y se agregan a la incrustación vectorial original para el token x. Con esto, la incrustación del token xahora se actualiza para reflejar el contexto proporcionado por los otros tokens en la secuencia que son relevantes para él.

La incrustación vectorial actualizada se envía a otra capa lineal, con su propia matriz de_{peso WZ}, donde el vector actualizado por contexto se normaliza de nuevo a un número constante de dimensiones y luego se envía a la siguiente capa de atención. Cada capa de atención progresiva captura mayores matices contextuales.
 

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

Para compensar, las redes transformadoras dividen la incrustación del token de entrada original en h trozos de tamaño uniforme. Asimismo, dividen _WQ, _WK y_WV en h subconjuntos denominados cabezas de consulta , cabezas de clavey cabezas de valor, respectivamente. Cada cabeza de consulta, cabeza de clave y cabeza de valor recibe un fragmento de la incrustación de token original. Los vectores producidos por cada una de estas tripletas paralelas de cabezas de consulta, cabezas de clave y cabezas de valor se introducen en una cabeza de atención correspondiente .Finalmente, los resultados de estos h circuitos paralelos se concatenan para actualizar el token completo.

En el entrenamiento, cada circuito aprende distintos pesos que capturan un aspecto separado de los significados semánticos. Esto, a su vez, ayuda al modelo a procesar las diferentes formas en que las implicaciones de una palabra pueden verse influenciadas por el contexto de otras palabras a su alrededor.

La desventaja de la atención multicabezal estándar no es tanto la presencia de algún defecto crucial, sino más bien la falta de optimización. MHA fue el primer algoritmo de su tipo 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 ineficiencia de MHA se debe a la abundancia de cálculos y parámetros del modelo. En el MHA estándar, cada encabezado de consulta, encabezado de clave y encabezado de valor en cada bloque de atención tiene su propia matriz de ponderaciones. Así, por ejemplo, un modelo con 8 cabezales de atención en cada capa de atención, muchos menos que la mayoría de los LLM modernos, requeriría 24 matrices de peso únicas solo para los cabezales 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 computacionalmente costosa. Los requisitos de cálculo para MHA escalan cuadráticamente con respecto a la longitud de la secuencia: duplicar el número de tokens en una secuencia de entrada requiere cuadruplicar la complejidad. Esto pone límites prácticos estrictos al tamaño de las ventanas de contexto.

MHA también ejerce una gran presión sobre la memoria del sistema. Las GPU no tienen mucha memoria integrada para almacenar las salidas de la cantidad masiva de cálculos intermedios que deben ser recordadas en cada paso de procesamiento posterior. En cambio, estos resultados intermedios se almacenan en la memoria de alto ancho de banda (HBM), que no se encuentra en el 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ómputo necesarios para entrenar y ejecutar la inferencia se convirtieron en un cuello de botella en el rendimiento a escala del modelo.

El progreso adicional requería métodos para reducir el número de pasos computacionales sin reducir la capacidad de los transformadores para aprender y reproducir patrones lingüísticos intrincadamente complejos. Fue en este contexto que se introdujeron MQA y, posteriormente, GQA.

La atención multiconsulta (MQA) es un mecanismo de atención computacionalmente más eficiente que simplifica la atención de múltiples cabezales para reducir el uso de memoria y los cálculos intermedios. En lugar de entrenar un cabezal clave único y un cabezal de valor para cada cabeza de atención, MQA utiliza un solo cabezal clave y un cabezal de valor único 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 entonces en 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 presentó MQA, 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 MQA también acelera significativamente el entrenamiento al permitir un tamaño de lote más grande.

A pesar de sus beneficios, MQA tiene varias desventajas inevitables.

• Degradación del rendimiento: como era de esperar, la reducción del número de parámetros de modelo únicos y entrenables reduce la capacidad del modelo para el conocimiento y los matices. MQA implica una caída significativa en la precisión en comparación con MHA estándar, lo que lo hace inadecuado para ciertas situaciones y casos de uso.
  
  
• Debe entrenarse desde cero: un modelo entrenado con MHA estándar no puede adaptarse simplemente a MQA, sino que debe entrenarse con MQA desde cero. Esto significa que MQA no se puede utilizar para optimizar los modelos existentes y conlleva un costo de oportunidad considerable al experimentar con MQA para nuevos modelos.
  
  
• Redundancias en el paralelismo tensor: uno de los principales beneficios de entrenar modelos transformadores en GPU es la capacidad de realizar múltiples operaciones complejas de tensor 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 seguir siendo 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 comparten 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ó GQA por primera vez 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 modelos Granite 3.0 de IBM emplean GQA para una inferencia rápida.

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

En la práctica, 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 para los principales LLM.

• Uso eficiente de la GPU: la distribución de pares clave-valor de GQA usa el paralelismo tensorial, lo que reduce la cantidad de cómputo que se "desperdicia" al replicar valores redundantes.
  
  
• Compromiso eficaz: 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 MHA y casi tan rápido como MQA.
  
  
• Reducida la sobrecarga del ancho de banda de la memoria: Al igual que MQA, GQA reduce significativamente la cantidad de cálculos intermedios que deben calcularse, almacenarse y recuperarse en el momento de la inferencia.
  
  
• Entrenamiento flexible: a diferencia de MQA, la atención a las consultas en grupo no requiere que los modelos se entrenen desde cero empleando el enfoque. Los modelos previamente entrenados con MHA estándar se pueden adaptar para usar GQA a través de un proceso de ajuste llamado "entrenamiento".