Desmitificando GQA: atención de consultas agrupadas para una formación previa eficiente en LLM
La variante de atención multicabezal que impulsa LLM como LLaMA-2, Mistral7B, etc.
En el artículo anterior sobre entrenamiento de modelos a gran escala, analizamos LoRA. En este artículo, examinaremos otra estrategia adoptada por diferentes modelos de lenguajes grandes para una capacitación eficiente: Atención de consultas agrupadas (GQA). En resumen, la atención de consultas agrupadas (GQA) es una generalización de la atención de múltiples cabezas (MHA) y la atención de múltiples consultas (MQA), siendo cada una de ellas un caso especial de GQA. Por lo tanto, antes de sumergirnos en la atención de consultas agrupadas, revisemos la atención tradicional de múltiples cabezas propuesta por Vaswani et al. en el artículo fundamental “La atención es todo lo que necesitas”. A continuación, exploraremos la atención de consultas múltiples y cómo aborda los desafíos con MHA. Finalmente, responderemos las preguntas “¿Qué es GQA?” y “¿Cómo nos brinda lo mejor de ambos mundos?”
La atención de múltiples cabezas es un componente crítico de los modelos Transformer, ya que les permite procesar y comprender de manera eficiente secuencias complejas en tareas como traducción de idiomas, resúmenes y más. Para comprender sus complejidades, debemos profundizar en los fundamentos matemáticos y comprender cómo funcionan las múltiples cabezas del mecanismo de atención.
El mecanismo de atención básico calcula una suma ponderada de valores, cuyos pesos dependen de una consulta y un conjunto de claves. Matemáticamente esto se expresa como:
Esto se conoce como atención del producto escalado. En esta ecuación, Q (Consulta) y K (Clave) son matrices que representan las consultas y las claves. V (Valor) es la matriz de valores. “d_k” es la dimensionalidad de las claves, que se utiliza para escalar.
Expandiéndose con atención de múltiples cabezas (MHA)
La atención de múltiples cabezas emplea múltiples ‘cabezas’ de capas de atención, lo que permite que el modelo preste atención a información de diferentes subespacios de representación. En cada encabezado, hay un conjunto independiente de capas lineales (matrices de proyección) para la consulta, la clave y los valores (este es un punto importante que revisaremos en GQA). Para cada cabeza (numerada h):
cabezaʰ = Atención(Q.Wqʰ,K.Wkʰ,V.Wvʰ)
Concatenación de salidas de cabezales
Las salidas de los cabezales individuales se concatenan y luego se transforman linealmente.
Cabezal múltiple(q,k,V) = Concat(cabeza¹,cabeza²,…,cabezaʰ) .W.ᵒ
Wᵒ es otra matriz de peso que transforma linealmente el vector concatenado a la dimensión de salida final.
La intuición detrás de la atención de múltiples cabezas es que al aplicar el mecanismo de atención varias veces en paralelo, el modelo puede capturar diferentes tipos de relaciones en los datos.
Sin embargo, MHA permite una comprensión matizada de las relaciones entre las diferentes partes de la entrada. Sin embargo, esta complejidad tiene un costo: una demanda significativa de ancho de banda de memoria, especialmente durante la inferencia del decodificador.
El desafío del ancho de banda de la memoria en la atención de múltiples cabezas
El quid de la cuestión reside en la sobrecarga de memoria. Cada paso de decodificación en modelos autorregresivos como Transformers requiere cargar los pesos del decodificador junto con todas las claves y valores de atención. Este proceso no sólo requiere un uso intensivo de computación sino también un uso intensivo del ancho de banda de la memoria. A medida que crecen los tamaños de los modelos, estos gastos generales también aumentan, lo que hace que la ampliación sea una tarea cada vez más ardua.
Aparición de la atención multiconsulta (MQA)
La atención multiconsulta (MQA) surgió como una solución para mitigar este cuello de botella. La idea es simple pero efectiva: use múltiples encabezados de consulta pero solo un encabezado de clave y valor. Este enfoque reduce significativamente la carga de memoria, mejorando la velocidad de inferencia. Se ha empleado en múltiples modelos a gran escala, como PaLM, StarCoder y Falcon.
En la atención de consultas múltiples, promediamos los encabezados de claves y valores para que todos los encabezados de consulta compartan la misma clave y valor. Esto se logra replicando el “cabezal” agrupado medio H veces, donde H es el número de cabezas de consulta.
Una pregunta interesante que cabe plantearse aquí es: ¿cómo se convierte un modelo de atención de cabezales múltiples previamente entrenado en un modelo de atención de consultas múltiples (MQA)? La creación de un modelo de atención de consultas múltiples a partir de un modelo de múltiples cabezales existente implica un proceso de dos pasos: conversión de la estructura del modelo y posterior entrenamiento previo. [1]
Conversión de punto de control: este paso transforma la estructura de un modelo de múltiples cabezales en un modelo de múltiples consultas. Se logra fusionando (agrupación media) las matrices de proyección (capas lineales) para claves y valores de los múltiples cabezales del modelo original en matrices de proyección únicas para claves y valores. Se ha descubierto que este enfoque de agrupación de medias es más eficaz que seleccionar uno de los encabezados de clave y valor existentes o inicializar nuevos encabezados de clave y valor desde cero. La estructura resultante tiene una proyección de clave y valor consolidada, característica del modelo multiconsulta.
Entrenamiento previo del modelo convertido: después de la transformación estructural, el modelo se somete a un entrenamiento adicional. Esta capacitación no es tan extensa como la capacitación modelo original; es una fracción (denotada como α) de los pasos de entrenamiento del modelo original. El propósito de esta fase de preentrenamiento es permitir que el modelo ajuste y optimice su rendimiento de acuerdo con su nuevo mecanismo de atención simplificado. La formación sigue la misma receta que la original, garantizando coherencia en la dinámica de aprendizaje.
Sin embargo, MQA no está exento de inconvenientes. La complejidad reducida puede conducir a la degradación de la calidad y la inestabilidad del entrenamiento.
Atención de consultas agrupadas
La atención de consultas agrupadas (GQA) es un enfoque simple que combina elementos de atención de múltiples cabezas (MHA) y atención de múltiples consultas (MQA) para crear un mecanismo de atención más eficiente. El marco matemático de GQA se puede entender de la siguiente manera:
División en grupos: en GQA, los encabezados de consulta (Q) de un modelo tradicional de múltiples encabezados se dividen en grupos G. A cada grupo se le asigna un único cabezal de clave (K) y valor (V). Esta configuración se denomina GQA-G, donde G representa el número de grupos.
Casos especiales de GQA:
- GQA-1 = MQA: con un solo grupo (G = 1), GQA se vuelve equivalente a MQA, ya que solo hay una única clave y encabezado de valor para todos los encabezados de consulta.
- GQA-H = MHA: cuando el número de grupos es igual al número de cabezas (G = H), GQA se comporta como MHA tradicional, y cada cabeza de consulta tiene su clave y valor únicos.
Combinamos las matrices de proyección de valores y claves de los cabezales originales dentro de cada grupo para convertir un modelo de múltiples cabezales en un modelo GQA. Esta técnica promedia las matrices de proyección de cada cabeza de un grupo, lo que da como resultado una única proyección de clave y valor para ese grupo.
Al utilizar GQA, el modelo mantiene un equilibrio entre la calidad MHA y la velocidad MQA. Debido a que hay menos pares clave-valor, se minimizan las necesidades de ancho de banda de memoria y carga de datos. La elección de G presenta una compensación: más grupos (más cercanos a MHA) dan como resultado una mayor calidad pero un rendimiento más lento, mientras que menos grupos (cercanos a MQA) aumentan la velocidad a riesgo de sacrificar la calidad. Además, a medida que crece el tamaño del modelo, GQA permite una disminución proporcional en el ancho de banda de la memoria y la capacidad del modelo, correspondiente a la escala del modelo. Por el contrario, para modelos más grandes, la reducción a una sola clave y valor puede ser excesivamente severa en MQA.
Conclusión
En esta publicación, analizamos por primera vez la atención multicabezal tradicional (MHA) y su variante, la atención multiconsulta. Luego analizamos una formulación GQA más genérica, que muchos modelos LLM utilizan para una capacitación previa eficaz. GQA combina atención multicabezal (MHA) con atención multiconsulta (MQA), proporcionando un equilibrio justo entre calidad y velocidad. GQA minimiza las demandas de ancho de banda de memoria al agrupar los encabezados de consulta, lo que lo hace apropiado para ampliar modelos. GQA se ha utilizado en lugar de la típica atención de cabezales múltiples en modelos recientes como LLaMA-2 y Mistral7B.
Referencias:
[1] GQA: Entrenamiento de modelos de transformadores de consultas múltiples generalizados a partir de puntos de control de múltiples cabezales:https://arxiv.org/pdf/2305.13245.pdf
[2] MQA: Decodificación rápida de transformadores: todo lo que necesita es un cabezal de escriturahttps://arxiv.org/abs/1911.02150
[3] MHA: Atención es todo lo que necesita: https://arxiv.org/abs/1706.03762
Desmitificando GQA: atención de consultas agrupadas fue publicado originalmente en Hacia la ciencia de datos en Medium, donde las personas continúan la conversación resaltando y respondiendo a esta historia.