Self-attention causal, multi-cabeza

Hasta ahora cada token es un vector que no sabe nada de sus vecinos. La self-attention (auto-atención) es el mecanismo que permite que cada posición lea información de las demás. En un modelo de lenguaje autoregresivo, con una restricción clave: la posición $i$ solo puede mirar las posiciones $\le i$. Eso es la atención causal.

Queries, keys y values

Cada token proyecta su vector $x_i$ en tres roles mediante matrices aprendidas:

• Query (consulta) $q_i = x_i W_Q$: "¿qué estoy buscando?"
• Key (clave) $k_j = x_j W_K$: "¿qué ofrezco para ser encontrado?"
• Value (valor) $v_j = x_j W_V$: "si me eligen, esto es lo que aporto".

El puntaje de afinidad entre $i$ y $j$ es el producto punto $q_i \cdot k_j$. Normalizamos con softmax y promediamos los values. En forma matricial, para toda la secuencia:

donde $M$ es la máscara causal: $M_{ij} = 0$ si $j \le i$ y $M_{ij} = -\infty$ si $j > i$. Como $e^{-\infty} = 0$, el softmax asigna peso exactamente cero al futuro.

Escribe una oración y observa la matriz de pesos con máscara causal. Cada fila es un token "mirando" hacia atrás; el triángulo superior queda vacío:

¿Por qué dividir por $\sqrt{d_k}$?

Si las componentes de $q$ y $k$ son aproximadamente independientes con media 0 y varianza 1, su producto punto tiene varianza $d_k$. Con $d_k = 64$, los puntajes tienen desvío estándar ~8; el softmax de valores tan dispersos se satura en un one-hot y sus gradientes casi desaparecen. Dividir por $\sqrt{d_k}$ devuelve la varianza a ~1.

Múltiples cabezas

Una sola atención produce un único promedio ponderado por posición. Pero un token puede necesitar varias cosas a la vez: el sujeto de la oración, el token anterior, la apertura de un paréntesis. La multi-head attention (atención multi-cabeza) parte la dimensión $d$ en $h$ cabezas de tamaño $d/h$, corre atención en paralelo en cada una y concatena los resultados:

varias atenciones en paralelo, cada una atenta a algo distinto; pegamos sus resultados y $W_O$ los mezcla en un solo vector

El costo en parámetros es el mismo que una cabeza de dimensión $d$: solo reorganizamos los tensores.

Implementación

Calculamos Q, K y V con una sola capa lineal de salida $3d$ (más eficiente que tres capas), reordenamos para tener las cabezas como dimensión de batch, y aplicamos la máscara con masked_fill.

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


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0
        self.n_head = cfg.n_head
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd)
        self.attn_drop = nn.Dropout(cfg.dropout)
        self.resid_drop = nn.Dropout(cfg.dropout)
        mask = torch.tril(torch.ones(cfg.block_size, cfg.block_size, dtype=torch.bool))
        self.register_buffer("mask", mask.view(1, 1, cfg.block_size, cfg.block_size),
                             persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        hs = C // self.n_head
        q, k, v = self.qkv(x).split(C, dim=2)                 # 3 x (B, T, C)
        q = q.view(B, T, self.n_head, hs).transpose(1, 2)     # (B, nh, T, hs)
        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
        v = v.view(B, T, self.n_head, hs).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(hs)       # (B, nh, T, T)
        att = att.masked_fill(~self.mask[:, :, :T, :T], float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)

        y = att @ v                                           # (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C)      # concatena cabezas
        return self.resid_drop(self.proj(y))

Detalles que importan:

• register_buffer(..., persistent=False): la máscara viaja con el módulo a la GPU (.to(device)) pero no es un parámetro entrenable ni se guarda en el state_dict.
• self.mask[:, :, :T, :T]: recortamos la máscara al largo real de la secuencia, que puede ser menor que block_size.
• .contiguous() antes del view: transpose devuelve una vista no contigua en memoria, y view lo exige.

Un test que vale oro: ¿se filtra el futuro?

Los bugs de máscara son silenciosos: el modelo entrena "demasiado bien" porque puede leer la respuesta. Un test simple: si cambias el último token, las salidas de las posiciones anteriores no deben moverse.

from dataclasses import dataclass

@dataclass
class Cfg:
    n_embd: int = 32
    n_head: int = 4
    block_size: int = 16
    dropout: float = 0.0

torch.manual_seed(0)
attn = CausalSelfAttention(Cfg()).eval()

x = torch.randn(1, 10, 32)
x2 = x.clone()
x2[:, -1] = torch.randn(32)          # perturbamos solo la última posición

y, y2 = attn(x), attn(x2)
assert torch.allclose(y[:, :-1], y2[:, :-1], atol=1e-6), "¡fuga de información!"
assert not torch.allclose(y[:, -1], y2[:, -1])
print("La máscara causal funciona")

Agrega este test a tu suite: atrapa errores de indexado, transposiciones invertidas y máscaras mal recortadas.

En producción: la versión fusionada

PyTorch 2 incluye F.scaled_dot_product_attention, que implementa lo mismo en un kernel fusionado (FlashAttention cuando el hardware lo permite) sin materializar la matriz $T \times T$ completa en memoria:

y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                   dropout_p=self.attn_drop.p if self.training else 0.0)

Escribir la versión explícita primero vale la pena: es la que puedes inspeccionar, visualizar y depurar. Luego reemplázala y usa el test de arriba para verificar que da los mismos resultados.

Un último número para tener en mente: la matriz de atención tiene $T^2$ entradas por cabeza, así que el costo de esta capa crece cuadráticamente con el largo del contexto. Es la razón de fondo de buena parte de la investigación en contextos largos.

Lo que te llevas

• Atención = promedio de values ponderado por afinidad query–key, normalizado con softmax.
• La máscara causal pone $-\infty$ en el futuro; el escalado $1/\sqrt{d_k}$ evita saturar el softmax.
• Multi-cabeza = varios canales de lectura en paralelo, al mismo costo de parámetros.
• Siempre testea que no haya fuga de información del futuro.