Explicación de la Atención de Múltiples Consultas

Guía para Manejar Varios Preguntas Simultáneamente

La Atención de Multi-Consulta (MQA) es un tipo de mecanismo de atención que puede acelerar la velocidad de generación de tokens en el decodificador mientras se asegura el rendimiento del modelo.

Se utiliza ampliamente en la era de los modelos de lenguaje grandes, muchos LLM adoptan MQA, como Falcon, PaLM, StarCoder y otros.

Atención Multi-Cabeza (MHA)

Antes de presentar MQA, primero veamos el mecanismo de atención por defecto del transformer.

La Atención Multi-Cabeza es el mecanismo de atención por defecto del modelo transformer, como se muestra en la Figura 1:

• Revisión del Creador de Colossyan ¿El Mejor Generador de Videos de IA?
• Métodos para modelar el valor de por vida del cliente Las cosas buenas y las trampas
• 10 Formas en las que la Inteligencia Artificial está dando forma al Desarrollo Seguro de Aplicaciones

Sin embargo, hay un problema con los modelos de lenguaje auto-regresivos basados en los decodificadores del transformer cuando se trata de la generación de texto.

Durante el entrenamiento, tenemos acceso a la secuencia objetivo verdadera y podemos implementar eficientemente la paralelización.

Sin embargo, durante la inferencia, la consulta de cada posición se aplica a todos los pares clave-valor generados en o antes de esa posición. En otras palabras, la salida de la capa de auto-atención en una posición específica afecta a la generación del siguiente token. Debido a la incapacidad de realizar cálculos en paralelo, la decodificación se vuelve más lenta.

A continuación se muestra el proceso de decodificación de una capa de auto-atención en un modelo de lenguaje auto-regresivo basado en decodificadores del transformer:

def MHAForDecoder(x, prev_K, prev_V, P_q, P_k, P_v, P_o):    q = tf.einsum("bd, hdk−>bhk", x, P_q)    new_K = tf.concat([prev_K, tf.expand_dims(tf.einsum ("bd, hdk−>bhk", x, P_k), axis = 2)], axis = 2)    new_V = tf.concat([prev_V, tf.expand_dims(tf.einsum("bd, hdv−>bhv", x, P_v), axis = 2)], axis = 2)    logits = tf.einsum("bhk, bhmk−>bhm", q, new_K)    weights = tf.softmax(logits)    O = tf.einsum("bhm, bhmv−>bhv", weights, new_V)    Y = tf.einsum("bhv, hdv−>bd", O, P_o)    return Y, new_K, new_V

Variables:

• x: el tensor de entrada en el paso actual, que es el paso m+1, con una forma de [b, d]
• P_q, P_k: los tensores de proyección de consulta y clave, con una forma de [h, d, k]
• P_v: el tensor de proyección de valor, con una forma de [h, d, v]
• P_o: las proyecciones lineales aprendidas, con una forma de [h, d, v]
• Prev_K: el tensor Clave del paso anterior, con una forma de [b, h, m, k]
• Prev_V: el tensor Valor del paso anterior, con una forma de [b, h, m, v]
• new_K: el tensor Clave con la adición del paso actual, con una forma de [b, h, m+1, k]
• new_V: el tensor Valor con la adición del paso actual, con una forma de [b, h, m+1, v]

Dimensiones:

• m: el número de pasos anteriores realizados
• b: tamaño del lote
• d: dimensión de la entrada y salida
• h: número de cabezas
• k: otra dimensión de los tensores Q, K
• v: otra dimensión del tensor V

Atención de Multi-Consulta (MQA)

La Atención Multi-Consulta es una variación de la atención de múltiples cabezas.

El enfoque de MQA consiste en mantener el número original de cabezas para Q, pero tener solo una cabeza para K y V. Esto significa que todas las cabezas de Q comparten el mismo conjunto de cabezas de K y V, de ahí el nombre de Multi-Consulta, como se muestra en la Figura 2:

El código del proceso de decodificación para MQA es esencialmente el mismo que el código para MHA, excepto que se elimina la letra “h” que representa la dimensión de las cabezas de la ecuación tf.einsum para K, V, P_k y P_v:

def MQAForDecoder(x, prev_K, prev_V, P_q, P_k, P_v, P_o):    q = tf.einsum("bd, hdk−>bhk", x, P_q)    new_K = tf.concat([prev_K, tf.expand_dims(tf.einsum ("bd, dk−>bk", x, P_k), axis = 2)], axis = 2)    new_V = tf.concat([prev_V, tf.expand_dims(tf.einsum("bd, dv−>bv", x, P_v), axis = 2)], axis = 2)    logits = tf.einsum("bhk, bmk−>bhm", q, new_K)    weights = tf.softmax(logits)    O = tf.einsum("bhm, bmv−>bhv", weights, new_V)    Y = tf.einsum("bhv, hdv−>bd", O, P_o)    return Y, new_K, new_V

Rendimiento

¿Cuánto puede mejorar realmente MQA la velocidad? Veamos el gráfico de resultados proporcionado en el artículo original:

Según la tabla anterior, se puede observar que la mejora de velocidad de MQA en el codificador no es muy significativa, pero es bastante significativa en el decodificador.

También hay experimentos sobre la calidad en el artículo, que muestran que MQA solo tiene un rendimiento ligeramente inferior en comparación con la línea de base. Consulta el artículo para obtener más detalles, el enlace se encuentra al final de este artículo.

¿Por qué MQA puede lograr aceleración de inferencia?

Más eficiente en memoria

En MQA, el tamaño de los tensores de clave y valor es b * k y b * v, mientras que en MHA, el tamaño de la clave y el valor es b * h * k y b * h * v, donde h representa el número de cabezas.

Menor complejidad computacional

Al usar la memoria caché KV, el costo computacional de calcular los tensores de Clave y Valor en cada paso de MQA es de 1/h del MHA, donde h representa el número de cabezas.

Resumen

En general, MQA logra aceleración de inferencia a través de los siguientes métodos:

• El tamaño de la memoria caché KV se reduce en un factor de h (número de cabezas), lo que significa que los tensores que deben almacenarse en la memoria GPU también se reducen. El espacio ahorrado se puede usar para aumentar el tamaño del lote, mejorando así la eficiencia.
• Se reduce la cantidad de datos leídos de la memoria, lo que reduce el tiempo de espera de las unidades de cálculo y mejora la utilización computacional.
• MQA tiene una memoria caché KV relativamente pequeña que puede caber en la caché (SRAM). MHA, por otro lado, tiene una memoria caché KV más grande que no se puede almacenar por completo en la caché y debe leerse desde la memoria GPU (DRAM), lo que consume tiempo.

Conclusión

Vale la pena mencionar que MQA fue propuesto en 2019 y su aplicación no era tan extensa en ese momento. Esto se debe a que los modelos anteriores no necesitaban preocuparse por estos aspectos, por ejemplo, LSTM solo necesitaba mantener un estado sin necesidad de mantener ninguna memoria caché.

Cuando se propuso inicialmente el transformador, se utilizaba principalmente en tareas Seq2Seq, específicamente en modelos Codificador-Decodificador. Sin embargo, los modelos no eran muy grandes en escala y no había mucha demanda práctica para ellos, por lo que MQA no atrajo mucha atención.

Más tarde, el modelo representativo BERT, que también se basa en la estructura del codificador del transformador, realizó un pase directo hacia adelante.

Solo cuando los modelos de lenguaje grandes recientes basados en el decodificador del transformador, como GPT, ganaron una aplicación generalizada se descubrió el cuello de botella de la inferencia. Como resultado, las personas volvieron a visitar trucos de hace unos años y los encontraron muy útiles. En otras palabras, se debe principalmente a la demanda práctica de modelos generativos de estilo GPT a gran escala.

Finalmente, si hay algún error u omisión en este texto, no dude en señalarlo.

Referencias

MQA paper: Fast Transformer Decoding: One Write-Head is All You Need

Attention Is All You Need

https://paperswithcode.com/method/multi-query-attention

We will continue to update Zepes; if you have any questions or suggestions, please contact us!