Atención multi-query

El modelo relaciona el fragmento actual del texto con la conversación anterior mediante varios cálculos paralelos. Estos cálculos se llaman cabezas de atención, o attention. Si cada cabeza almacena sus propias descripciones de toda la conversación anterior, aumentan el coste de memoria y su lectura. ¿Se puede compartir parte de los datos manteniendo distintas formas de leerlos?

Multi-Query Attention (MQA) mantiene múltiples query —descripciones de las necesidades de las posiciones actuales—, pero les da un conjunto común de key y value. Key sirve para establecer correspondencias y value contiene la información que se transmitirá. Todos son listas de números.

Dos cabezas pueden usar el mismo conjunto key/value y aun así repartir los pesos de forma distinta, porque tienen query diferentes. Compartir los datos no obliga a obtener un resultado idéntico.

El ahorro afecta, entre otras cosas, a KV Cache, las descripciones almacenadas del texto anterior. Sin embargo, el cambio puede afectar a la calidad y no significa una aceleración proporcional de todo el modelo. Grouped-Query Attention comparte los datos dentro de grupos más pequeños.

Mecanismo y detalles

En la atención de múltiples cabezas habitual, cada cabeza tiene sus propias proyecciones de consulta, clave y valor. MQA elimina la dimensión de cabezas de K/V y de sus proyecciones, manteniéndola en las consultas. Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, §3 define así esta variante.

La misma memoria, otra pregunta

Nuestro ejemplo tiene dos claves: $k_1=(1,0)$ y $k_2=(0,1)$. Dos cabezas query reciben respectivamente $q_1=(2,0)$ y $q_2=(0,2)$. Scaled Dot-Product Attention divide los productos por $\sqrt{2}$ y aplica Softmax. La primera cabeza da a las claves pesos de aproximadamente 0,804 y 0,196; la segunda invierte estas proporciones. Con unos $v_1=(1,0)$ y $v_2=(0,1)$ comunes, esas son también las coordenadas de dos resultados distintos.

De dónde procede el ahorro

La caché KV almacena las K/V anteriores. Con ocho cabezas de consulta, MHA guarda ocho conjuntos de K/V, y MQA, uno. Para la misma longitud de secuencia, dimensión de cabeza y precisión, los propios tensores de la caché son entonces ocho veces menores. Es un cálculo del tamaño, no una aceleración de ocho veces de todo el modelo.

El análisis de Shazeer en §2.4.1 y §3.1 relaciona este cambio con la transferencia de memoria durante los pasos sucesivos de generación. Sin embargo, cambia la arquitectura y su capacidad de representación; no es una eliminación sin pérdida de copias idénticas de cualquier MHA entrenado. La atención de consultas agrupadas conserva varios conjuntos de K/V como variante intermedia.

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