6 Atención multi-cabeza
Dónde estamos. Ya sabemos cómo una atención mezcla información entre palabras. Pero una sola atención es miope: solo puede fijarse en un tipo de relación a la vez. Este capítulo explica por qué los transformers ejecutan varias atenciones en paralelo —las cabezas— y qué gana el modelo con ello. Es un cambio pequeño en la fórmula con una consecuencia enorme en lo que el modelo puede entender.
6.1 La idea en una frase
En lugar de una sola atención, el modelo corre varias a la vez (las cabezas), cada una mirando la frase con un criterio distinto, y luego junta lo que cada una encontró.
6.2 Conceptos clave y su papel en el transformer
Antes de entrar en detalle, definimos los términos de este capítulo y para qué sirve cada uno dentro de un transformer:
- Cabeza (head). Definición: una atención completa e independiente (la del Cap. 4) que trabaja sobre una porción del vector, con sus propias tablas \(W_i^Q, W_i^K, W_i^V\). En el transformer: es uno de los canales de atención en paralelo; cada cabeza puede seguir un tipo distinto de relación a la vez.
d_k(dimensión por cabeza). Definición: el tamaño del trozo de vector que recibe cada cabeza,d_k = d_model / h. En el transformer: reparte el vector entre lashcabezas; al ser más pequeño, hace que muchas cabezas cuesten casi lo mismo que una sola atención grande.- Concatenación (Concat). Definición: pegar las
hsalidas una tras otra para reconstruir un vector del tamaño original. En el transformer: reúne en un solo sitio las notas que tomó cada cabeza por separado. - \(W^O\) (mezcla final). Definición: una tabla aprendida que reproyecta el vector concatenado. En el transformer: combina lo que aportó cada cabeza en un único resultado coherente que devolver al flujo residual.
- Especialización emergente. Definición: los roles de cada cabeza (posicional, sintáctica, de palabras raras…) no se programan: surgen del entrenamiento. En el transformer: explica por qué las cabezas no se parecen entre sí y por qué se reparten el trabajo de leer la frase.
- Induction head. Definición: una cabeza que empareja prefijos y copia lo que vino después ([A][B] … [A] → [B]). En el transformer: se cree que es el motor de buena parte del in-context learning —aprender de los ejemplos del propio prompt—.
- Retrieval head. Definición: una cabeza especializada en encontrar un dato enterrado en un contexto larguísimo. En el transformer: un grupo diminuto (<5%) sostiene la capacidad de “aguja en el pajar”; lo retomamos en la Parte II.
- Poda de cabezas. Definición: eliminar cabezas en inferencia con poca pérdida. En el transformer: revela que las cabezas son parcialmente redundantes y que unas pocas especializadas hacen el trabajo pesado.
- MQA / GQA. Definición: esquemas que comparten claves/valores entre cabezas (una sola, o por grupos). En el transformer: recortan la memoria del KV-cache en inferencia conservando casi la calidad.
Con esto en mente, veamos por qué una sola atención no basta.
6.3 Para qué sirve (su papel en el transformer)
¿Por qué no basta con una atención? Porque una sola produce un único conjunto de pesos: obliga a elegir qué tipo de relación resaltar. Si el modelo usa esa única mirada para seguir la concordancia gramatical, no le queda con qué seguir, a la vez, a qué se refiere un pronombre tres frases atrás. Los propios autores lo dicen: con una sola cabeza “el promediado lo inhibe” —las relaciones distintas se mezclan y se pierden—.
Su función: dar al modelo varios canales de atención en paralelo, para que pueda atender a distintas cosas en distintos sitios al mismo tiempo —una cabeza a la sintaxis, otra a un referente lejano, otra a una palabra clave rara—. Sin multi-cabeza, el modelo solo podría rastrear una clase de relación por capa.
🧩 Analogía. Imagina un panel de especialistas leyendo la misma frase a la vez: el gramático marca quién concuerda con quién, la experta en referencias rastrea a qué apunta “ella”, el vigía de tecnicismos no pierde de vista el término raro. Cada uno toma sus notas en paralelo (las cabezas) y al final se juntan todas y se resumen en un único informe (eso es la concatenación + la mezcla final).
6.4 La mecánica
El truco es elegante: en vez de una atención sobre el vector completo, se parte el vector en h trozos (las cabezas), cada trozo de tamaño d_k = d_model / h. Cada cabeza hace su propia atención sobre su trozo; luego se concatenan (se pegan) las h salidas y se pasan por una última tabla aprendida, W^O:
\[ \text{MultiCabeza}(Q,K,V) = \text{Concat}(\text{cabeza}_1,\dots,\text{cabeza}_h)\,W^O \]
\[ \text{cabeza}_i = \text{Atención}(Q W_i^Q,\; K W_i^K,\; V W_i^V) \]
Qué hace cada pieza:
- \(\text{cabeza}_i\) = una atención completa e independiente (la del Cap. 4), pero sobre la porción
idel vector. Cada una aprende sus propias tablas \(W_i^Q, W_i^K, W_i^V\) → cada una “busca y anuncia” cosas distintas. (Aquí \(Q, K, V\) son los vectores de entrada; cada cabeza los proyecta con sus propias tablas, igual que las \(W_Q, W_K, W_V\) del Cap. 4.) - Concat(…) = pegar las
hsalidas una tras otra para reconstruir un vector del tamaño original. - \(W^O\) = una tabla aprendida que mezcla y reproyecta esas notas concatenadas en un único vector de salida. Para qué sirve: combinar lo que aportó cada cabeza en algo coherente que devolver al flujo residual.
Uno pensaría que h cabezas cuestan h veces más. No. Como cada cabeza trabaja sobre un trozo más pequeño (d_k = d_model/h), el cómputo total es prácticamente el mismo que el de una sola atención grande (cita literal de Vaswani et al. (2017)). Por eso el multi-cabeza es casi gratis: más miradas, mismo coste.
6.5 Números típicos
d_k = d_model / h no es una elección libre: sale de repartir el vector entre las cabezas. Algunos casos reales:
| Modelo | d_model | nº de cabezas (h) | d_k por cabeza |
|---|---|---|---|
| Transformer original (2017) | 512 | 8 | 64 |
| GPT-2 small | 768 | 12 | 64 |
| LLaMA-2-7B | 4096 | 32 | 128 |
d_k = 64 es un tamaño por cabeza muy habitual.
6.6 Lo fascinante: las cabezas se especializan solas
Aquí viene lo interesante. Nadie programa a una cabeza para un rol: los roles emergen del entrenamiento. Al analizar modelos ya entrenados, se encuentran cabezas con trabajos sorprendentemente nítidos (Voita et al. 2019): cabezas posicionales (siempre miran al token anterior o siguiente), sintácticas (siguen relaciones gramaticales) y de palabras raras (vigilan los tokens menos frecuentes).
El rol más célebre es la induction head (Olsson et al. 2022). Hace dos cosas en equipo: empareja prefijos (busca atrás dónde apareció antes el token actual) y copia (sube la probabilidad de lo que vino después la vez anterior). Es decir, implementa el patrón [A][B] … [A] → [B]: “la última vez que vi A, le seguía B, así que predigo B”. Lo asombroso: estas cabezas aparecen de golpe en una ventana estrecha del entrenamiento —un pequeño bache en la curva de pérdida marca ese instante— y, justo ahí, la capacidad del modelo da un salto. Se cree que son el mecanismo de buena parte del in-context learning —la capacidad de los modelos grandes de aprender de los ejemplos que les das en el propio prompt, sin reentrenar—.
6.7 ¿Hacen falta todas las cabezas?
Curiosamente, no todas. Se puede podar (eliminar) muchas cabezas en inferencia con apenas pérdida: Voita et al. (2019) quitaron 38 de 48 cabezas de un traductor perdiendo solo 0.15 BLEU; las últimas en caer son las especializadas. La lectura honesta: las cabezas son parcialmente redundantes, y unas pocas especializadas “hacen el trabajo pesado”. (Ojo: que sean podables después no significa que sobren durante el entrenamiento.)
La capacidad de un modelo para encontrar un dato enterrado en un contexto larguísimo (la “aguja en el pajar”) vive en un grupo diminuto de cabezas: menos del 5%. Si las anulas, el modelo alucina y no encuentra la aguja; si anulas cabezas al azar, no pasa casi nada. Son universales (todo modelo de contexto largo las tiene) y ya están en el modelo base. Lo retomamos en la Parte II (contexto largo). (Wu et al. 2024)
6.8 Un apunte hacia adelante: GQA y MQA
En inferencia, guardar las claves y valores de todas las cabezas llena la memoria (el KV-cache). Por eso los modelos modernos comparten claves/valores entre cabezas: MQA (una sola compartida) y GQA (grupos de cabezas comparten) reducen esa memoria conservando casi la calidad. Lo veremos al hablar de eficiencia y contexto largo (Parte II y VI).
6.9 En código
import torch
n, d_model, h = 5, 768, 12 # 5 tokens, GPT-2 small
d_k = d_model // h # 64
x = torch.randn(n, d_model)
# partir en h cabezas: cada token, en h trozos de tamaño d_k
heads = x.view(n, h, d_k) # (5, 12, 64)
print(heads.shape) # cada cabeza hará su atención sobre su trozoCada una de esas 12 porciones pasa por la atención del Cap. 4; las 12 salidas se concatenan de vuelta a tamaño 768 y se mezclan con W^O.
En un visualizador como BertViz puedes ver el mapa de atención de cada cabeza por separado. Comprobarás que no se parecen entre sí: unas miran al token anterior, otras saltan lejos, otras se clavan en el primer token. Esa diversidad es el porqué del multi-cabeza.
6.10 Resumen
- Una sola atención es miope (un tipo de relación); multi-cabeza corre varias en paralelo, cada una con su criterio aprendido.
- Función: varios canales de atención a la vez → el modelo atiende a cosas distintas en sitios distintos simultáneamente.
- Mecánica: partir el vector en
hcabezas (d_k = d_model/h), una atención por cabeza, concatenar y mezclar conW^O. Cuesta casi lo mismo que una sola. - Los roles emergen (no se diseñan): posicionales, sintácticos, de palabras raras, y las induction heads (motor del in-context learning).
- Muchas cabezas son podables → redundancia parcial; unas pocas especializadas hacen el trabajo pesado.
Siguiente (Capítulo 6): la atención mueve información entre palabras; pero cada palabra también necesita procesar lo que ha recogido. De eso se encarga la red feed-forward.
6.11 Ejercicios
- El reparto. Si
d_model = 1024y quieresh = 16cabezas, ¿cuánto valed_k? ¿Y conh = 8? - Coste. Explica con tus palabras por qué 12 cabezas no cuestan 12 veces más que una sola. (Pista: ¿de qué tamaño es el trozo de cada cabeza?)
- Induction head. Dada la secuencia
... gato perro ... gato ?, ¿qué predeciría una induction head, y con qué dos pasos (emparejar + copiar)? - Poda. Si puedes quitar 38 de 48 cabezas perdiendo casi nada, ¿significa que el multi-cabeza es inútil? Argumenta (piensa en entrenamiento vs inferencia).