Multi-head attention: muchas miradas en paralelo

El problema de una sola atención

Una capa de atención produce, para cada posición, un promedio ponderado de los values. Pero una palabra suele relacionarse con otras por motivos muy distintos a la vez. En "El abogado que contrató la empresa ganó el juicio", el verbo "ganó" necesita encontrar su sujeto ("abogado"), su objeto ("juicio") y quizás información de tiempo verbal del contexto. Si todo eso tiene que salir de una sola distribución de pesos, el promedio mezcla señales que deberían mantenerse separadas.

El paper lo dice con precisión: con una única cabeza, el promediado inhibe la capacidad de atender a distintos subespacios de representación en distintas posiciones.

La idea: proyectar, atender, concatenar

En vez de una atención sobre vectores de dimensión $d_{\text{model}}$, el Transformer ejecuta $h$ atenciones en paralelo, cada una sobre una proyección distinta y más chica:

cada cabeza $i$ hace la misma atención de siempre, pero sobre su propia versión achicada de $Q$, $K$ y $V$

con $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ y $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$.

En el modelo base del paper: $d_{\text{model}} = 512$, $h = 8$ y $d_k = d_v = d_{\text{model}}/h = 64$.

Cada cabeza tiene sus propias proyecciones, así que puede aprender su propia noción de "relevancia": una puede especializarse en relaciones sintácticas, otra en la palabra anterior, otra en correferencias. La proyección de salida $W^O$ recombina lo que encontraron.

El costo: casi gratis

Lo notable es que esto no cuesta más que una sola atención de dimensión completa. Cada cabeza trabaja con $d_k = d_{\text{model}}/h$, así que el total de parámetros en las proyecciones es:

$h$ cabezas chicas suman la misma cantidad de pesos que una sola grande: $d_{\text{model}}^2$

por cada una de $Q$, $K$ y $V$: exactamente lo mismo que una proyección cuadrada. El cómputo de los puntajes $QK^\top$ es $O(n^2 d_k)$ por cabeza, y $h$ cabezas suman $O(n^2 d_{\text{model}})$, igual que una cabeza grande.

Hay un costo oculto, eso sí: las matrices de atención. Cada cabeza produce su propia matriz $n \times n$, así que la memoria para guardar los pesos escala como $h \cdot n^2$. Esto será relevante cuando hablemos de contextos largos.

Implementación: el truco del reshape

En la práctica nadie implementa $h$ proyecciones separadas. Se usa una proyección de $d_{\text{model}} \to d_{\text{model}}$ y luego se "corta" el resultado en $h$ trozos con un view. Matemáticamente es equivalente: una matriz $d_{\text{model}} \times d_{\text{model}}$ es la concatenación de $h$ matrices $d_{\text{model}} \times d_k$.

Sigue las formas con atención, porque es donde más bugs aparecen:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int = 512, n_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.d_k = d_model // n_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def _split(self, x):
        # (B, n, d_model) -> (B, h, n, d_k)
        B, n, _ = x.shape
        return x.view(B, n, self.h, self.d_k).transpose(1, 2)

    def forward(self, query, key, value, mask=None):
        B, n_q, _ = query.shape
        q = self._split(self.w_q(query))     # (B, h, n_q, d_k)
        k = self._split(self.w_k(key))       # (B, h, n_k, d_k)
        v = self._split(self.w_v(value))     # (B, h, n_k, d_k)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)  # (B, h, n_q, n_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = self.dropout(F.softmax(scores, dim=-1))

        out = attn @ v                                           # (B, h, n_q, d_k)
        out = out.transpose(1, 2).contiguous().view(B, n_q, -1)  # (B, n_q, d_model)
        return self.w_o(out)

Dos detalles que vale la pena notar:

• El transpose(1, 2) pone la dimensión de cabezas antes de la secuencia, para que la multiplicación de matrices se haga en batch sobre (B, h). Las $h$ atenciones corren en paralelo como si fueran un batch más grande.
• El .contiguous() antes del view final es necesario porque transpose cambia los strides sin mover memoria.

Tres usos de la misma pieza

El Transformer usa multi-head attention en tres lugares, cambiando solo de dónde vienen $Q$, $K$ y $V$:

1. Self-attention en el encoder: $Q$, $K$, $V$ vienen de la salida de la capa anterior del encoder. Cada posición de la entrada atiende a todas las posiciones de la entrada.
2. Self-attention enmascarada en el decoder: igual, pero cada posición solo puede atender a posiciones anteriores (lo veremos en la lección de arquitectura).
3. Atención encoder-decoder (cross-attention): las queries vienen del decoder; las keys y values, de la salida del encoder. Es el heredero directo de la atención de Bahdanau: el decoder "consulta" la oración de entrada.

¿Qué aprenden las cabezas?

El paper incluye en su apéndice visualizaciones de cabezas que parecen capturar estructura sintáctica, como dependencias de larga distancia y resolución de anáforas (por ejemplo, conectar "its" con el sustantivo al que se refiere). Trabajos posteriores profundizaron esto:

• Algunas cabezas atienden casi siempre a la palabra anterior o la siguiente: posicionales puras.
• Otras se especializan en relaciones sintácticas concretas (sujeto-verbo, objeto directo).
• Muchas atienden fuertemente a tokens especiales o de puntuación, una especie de "no hacer nada" cuando la cabeza no encuentra lo que busca.
• Estudios de pruning (Michel et al., 2019, "Are Sixteen Heads Really Better than One?") mostraron que en muchas capas se puede eliminar gran parte de las cabezas en inferencia con poca pérdida: hay redundancia.

Cuidado con sobreinterpretar: que un mapa de atención "se vea" sintáctico no prueba que el modelo use esa información de la forma que suponemos. Los pesos de atención son una pista, no una explicación causal. Esta tensión es una de las semillas de la interpretabilidad mecanicista, donde las cabezas se estudian como circuitos (por ejemplo, las induction heads).

Resumen

• Multi-head ejecuta $h$ atenciones en paralelo sobre subespacios de dimensión $d_{\text{model}}/h$.
• El costo en parámetros y FLOPs es el mismo que una atención de dimensión completa; la memoria de las matrices $n \times n$ se multiplica por $h$.
• Se implementa con una proyección grande y un view/transpose.
• Las cabezas se especializan, pero hay redundancia y los mapas de atención no son explicaciones causales.

Falta una pieza crítica: hasta ahora, nada en la atención sabe en qué orden están las palabras. Eso es lo que resolvemos en la próxima lección.