31 · Atención multi-cabeza (multi-head)
Ejecuta h 'cabezas' de atención en paralelo sobre subespacios proyectados, luego concatena y proyecta — cada cabeza aprende un tipo distinto de relación.
La atención multi-cabeza ejecuta h atenciones en paralelo sobre subespacios de menor dimensión, así el modelo atiende a varias relaciones distintas con el MISMO costo total de cómputo que una sola cabeza.
Sin esto:
Una sola cabeza de atención solo captura un tipo de relación a la vez; la multi-cabeza es lo que da a los transformers su poder expresivo.
En la lección anterior construimos la auto-atención: cada posición calcula Attention(Q,K,V) = softmax(QKᵀ/√d_k)·V, mezclando información de toda la secuencia en un solo paso paralelo. Pero una sola cabeza de atención tiene una limitación. Cada cabeza produce UNA distribución softmax por posición de consulta — así que solo puede fijarse en UN tipo de relación. Quizás aprende a alinear verbos con sus sujetos. Pero, ¿qué hay de la relación entre un pronombre y su antecedente, o entre un token y el que tiene justo a su izquierda? Una sola cabeza no puede hacer todo esto a la vez: su único softmax debe llegar a un compromiso.
La atención multi-cabeza es la solución. En lugar de ejecutar una atención sobre los vectores completos de dimensión d_model, ejecutamos h cabezas en paralelo, cada una sobre un subespacio más pequeño de tamaño d_k = d_v = d_model / h. Cada cabeza tiene sus PROPIAS proyecciones aprendidas W_i^Q, W_i^K, W_i^V, de modo que cada cabeza ve una "vista" diferente de la entrada y puede especializarse en una relación distinta. Luego concatenamos las salidas de las h cabezas y las pasamos por una proyección de salida final W^O.
Formalmente, con d_model el ancho del modelo y h el número de cabezas:
head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)MultiHead(Q,K,V) = Concat(head_1, …, head_h) · W^O
donde cada W_i^Q, W_i^K ∈ ℝ^{d_model × d_k}, W_i^V ∈ ℝ^{d_model × d_v} y W^O ∈ ℝ^{h·d_v × d_model}. Recuerda que la operación central que ejecuta cada cabeza es exactamente el producto punto escalado de la lección 30: Attention(Q,K,V) = softmax(QKᵀ/√d_k)·V. La multi-cabeza simplemente hace h de estas sobre rebanadas más delgadas y vuelve a coser los resultados.
Lo ingenioso: como d_k = d_model / h, la cantidad total de cómputo a través de las h cabezas es la MISMA que una sola atención de ancho completo. Obtenemos la expresividad de muchas cabezas especializadas al precio de una.
Formas multi-cabeza: (seq, d_model) → dividir en (h, seq, d_k) → atención por cabeza → concatenar → W_O
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Pasa el cursor por una fila (token objetivo) para resaltar su distribución de atención. Cambia de cabeza arriba.
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Read-along de PyTorch: nn.MultiheadAttention con embed_dim, num_heads y need_weights para recuperar los mapas de atención
¿Por qué dividir la atención en múltiples cabezas en lugar de usar una sola atención de ancho completo?
- La atención multi-cabeza ejecuta h atenciones de producto punto escalado en paralelo sobre subespacios d_k = d_model/h, cada una con sus propias W_i^Q, W_i^K, W_i^V, luego concatena y aplica W^O.
- El número de parámetros es 4·d_model² (las cuatro proyecciones), independiente de h — añadir cabezas es un reshape, no nuevos pesos — y el cómputo total iguala a una sola cabeza de ancho completo.
- Las cabezas se especializan empíricamente (posicionales, sintácticas, de inducción), pero más cabezas no es estrictamente mejor: d_k se reduce al crecer h, así que hay un punto óptimo ajustable.
La atención multi-cabeza es el operador de mezcla central en todo transformer: BERT, GPT, T5, ViT, Whisper. Cuando lees 'GPT-3 tiene 96 cabezas', esa es la h de esta lección.
Si lo quitas: Leerías sobre 'cabezas de atención' y su especialización en papers de interpretabilidad sin entender mecánicamente cómo el ancho d_model se rebana, se proyecta y se recombina.