7 - La atención y el Transformer, desde cero

Antes de 2017, los modelos de lenguaje leían el texto palabra a palabra con redes recurrentes (RNN), arrastrando una memoria de una palabra a la siguiente. El libro cuenta cómo ese enfoque chocó con un muro en los textos largos, y cómo el artículo Attention Is All You Need lo sustituyó por el Transformer, cuya pieza central es la atención: cada token puede mirar directamente a todos los demás y decidir cuánto importa cada uno. En este cuaderno lo construimos pieza a pieza con NumPy, siguiendo las notas de Marta (secciones 7.7 y 7.8):

1. Importar las bibliotecas.
2. El muro de las redes recurrentes: recuerdos que se desvanecen.
3. La atención de producto escalar escalado, paso a paso.
4. Por qué dividir por $\sqrt{d_k}$.
5. La autoatención sobre una frase, y la máscara causal.
6. La atención multicabeza.
7. La codificación posicional con senos y cosenos.
8. Un bloque Transformer completo.

Paso 1: Importar las bibliotecas

---

import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=3, suppress=True)
rng = np.random.default_rng(0)

Paso 2: El muro de las redes recurrentes

---

Una RNN actualiza un estado oculto con cada token nuevo, $h_t = \tanh(W h_{t-1} + U x_t)$, así que la información sobre el primer token tiene que sobrevivir a todos los pasos posteriores. Durante el entrenamiento, el gradiente que indica a la red cómo afectó el primer token al resultado es un producto de una derivada por paso, y un producto de muchos números menores que uno se encoge exponencialmente: otra vez el desvanecimiento del gradiente del capítulo 6, ahora a lo largo de la secuencia. Medimos cuánto depende todavía el estado final del primer token para secuencias cada vez más largas:

d_hidden = 16
W = rng.normal(0, 0.9 / np.sqrt(d_hidden), (d_hidden, d_hidden))     # pesos recurrentes
U = rng.normal(0, 1 / np.sqrt(d_hidden), (d_hidden, d_hidden))       # pesos de entrada
inputs = rng.normal(0, 1, (60, d_hidden))                            # una secuencia de 60 vectores de tokens

h, jacobian, influence = np.zeros(d_hidden), np.eye(d_hidden), []
for t in range(60):
    h = np.tanh(W @ h + U @ inputs[t])
    jacobian = (np.diag(1 - h ** 2) @ W) @ jacobian                   # d h_t / d h_0, por la regla de la cadena
    influence.append(np.linalg.norm(jacobian))

for T in [1, 5, 10, 20, 40, 60]:
    print(f"Tras {T:>2} tokens, tamaño de d h_T / d h_0: {influence[T - 1]:.2e}")

Tras  1 tokens, tamaño de d h_T / d h_0: 2.63e+00
Tras  5 tokens, tamaño de d h_T / d h_0: 2.29e-01
Tras 10 tokens, tamaño de d h_T / d h_0: 1.44e-02
Tras 20 tokens, tamaño de d h_T / d h_0: 2.20e-04
Tras 40 tokens, tamaño de d h_T / d h_0: 7.85e-08
Tras 60 tokens, tamaño de d h_T / d h_0: 1.62e-13

Tras un token, la influencia del primer token es del orden de 1; tras 20 tokens ha caído a $2{,}2 \times 10^{-4}$, y tras 60 tokens a $1{,}6 \times 10^{-13}$, en la práctica cero. La red apenas puede aprender que una palabra del final de un texto largo depende de otra del principio. Como explica el libro, las LSTM y las GRU se desarrollaron para suavizar este problema, pero siguen leyendo un token cada vez, lo que además las hace lentas de entrenar.

Paso 3: La atención de producto escalar escalado

---

La atención toma tres matrices construidas a partir de los tokens: las consultas $Q$ (queries, lo que busca cada token), las claves $K$ (keys, lo que ofrece cada token para ser encontrado) y los valores $V$ (values, la información que transmite cada token):

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Sigamos la intuición de Marta con vectores pequeños diseñados a mano. Cada token tiene 4 características: cardiaco, localización corporal, persona y palabra funcional. De momento, las consultas, las claves y los valores son los propios vectores, y miramos a qué atiende el token "dolor":

tokens = ['el', 'paciente', 'tiene', 'pecho', 'dolor', 'brazo']
#                  cardiaco localización persona funcional
X_toy = 2 * np.array([[0.0,    0.0,      0.0,    1.0],    # el
                      [0.1,    0.0,      1.0,    0.0],    # paciente
                      [0.0,    0.0,      0.2,    0.8],    # tiene
                      [0.9,    0.8,      0.0,    0.0],    # pecho
                      [1.0,    0.3,      0.0,    0.0],    # dolor
                      [0.5,    0.9,      0.0,    0.0]])   # brazo

def softmax(z, axis=-1):
    e = np.exp(z - z.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, mask=None):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)                  # similitud de cada consulta con cada clave
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)      # las posiciones prohibidas reciben peso cero
    weights = softmax(scores, axis=-1)               # cada fila suma 1
    return weights @ V, weights                       # media ponderada de los valores

q = X_toy[tokens.index('dolor')]
scores = X_toy @ q / np.sqrt(4)
weights = softmax(scores)
for tok, s, w in zip(tokens, scores, weights):
    print(f"dolor -> {tok:<8} puntuación {s:5.2f}   peso {w:.3f}")
output = weights @ X_toy
print(f"\nNuevo vector para 'dolor' (media ponderada de los valores): {output}")

dolor -> el       puntuación  0.00   peso 0.038
dolor -> paciente puntuación  0.20   peso 0.046
dolor -> tiene    puntuación  0.00   peso 0.038
dolor -> pecho    puntuación  2.28   peso 0.369
dolor -> dolor    puntuación  2.18   peso 0.334
dolor -> brazo    puntuación  1.54   peso 0.176

Nuevo vector para 'dolor' (media ponderada de los valores): [1.516 1.107 0.107 0.136]

"Dolor" atiende sobre todo a "pecho" (0,369), a sí mismo (0,334) y a "brazo" (0,176), y casi nada a "el", "paciente" o "tiene", los tokens sin características cardiacas ni de localización corporal. Su nuevo vector es una mezcla de esos tres, así que tras la atención "dolor" lleva información sobre dónde duele. Como dice Marta, la softmax decide el foco.

Paso 4: Por qué dividir por $\sqrt{d_k}$

---

Marta explica que, a medida que crece la dimensión $d_k$, los productos escalares se hacen muy grandes y empujan la softmax a zonas donde los gradientes son diminutos. Podemos medirlo: con consultas y claves aleatorias, calculamos los pesos de atención de una consulta sobre 10 claves, con y sin el escalado, para dimensiones crecientes.

for d_k in [4, 64, 512]:
    q, K = rng.normal(size=d_k), rng.normal(size=(10, d_k))
    raw = K @ q
    w_raw, w_scaled = softmax(raw), softmax(raw / np.sqrt(d_k))
    print(f"d_k = {d_k:>3}: dispersión (desv. típica) de las puntuaciones {raw.std():6.2f} | mayor peso sin escalar {w_raw.max():.3f}, "
          f"escalando {w_scaled.max():.3f}")

d_k =   4: dispersión (desv. típica) de las puntuaciones   0.64 | mayor peso sin escalar 0.228, escalando 0.160
d_k =  64: dispersión (desv. típica) de las puntuaciones   6.13 | mayor peso sin escalar 0.831, escalando 0.222
d_k = 512: dispersión (desv. típica) de las puntuaciones  32.88 | mayor peso sin escalar 1.000, escalando 0.326

Con $d_k = 4$ las puntuaciones son pequeñas y los pesos siguen repartidos con o sin escalado. Pero la dispersión de las puntuaciones crece con la dimensión (6,13 con 64 dimensiones, 32,88 con 512) y, sin escalar, la softmax se satura: con $d_k = 512$ una sola clave se lleva todo el peso (1,000) y las demás reciben prácticamente cero, así que sus gradientes se desvanecen. Dividir por $\sqrt{d_k}$ devuelve las puntuaciones a una dispersión comparable en cualquier dimensión.

Paso 5: La autoatención sobre una frase, y la máscara causal

---

En un Transformer real, $Q$, $K$ y $V$ no son los propios vectores de los tokens: cada una se obtiene multiplicando la matriz de vectores de los tokens $X$ por una matriz aprendida, $Q = XW^Q$, $K = XW^K$, $V = XW^V$. Con la frase entera a la vez, la atención produce una matriz de pesos: la fila $i$ dice cuánto atiende el token $i$ a cada token $j$. Aquí las proyecciones son aleatorias, como antes del entrenamiento.

Los modelos de lenguaje como GPT generan texto de izquierda a derecha, así que un token no debe mirar los tokens que vienen después. Eso se consigue con una máscara causal, que da peso cero a todas las posiciones futuras:

sentence = ['el', 'paciente', 'tiene', 'dolor', 'torácico', 'que', 'irradia', 'al', 'brazo']
n, d_model, d_k = len(sentence), 8, 8
X = rng.normal(size=(n, d_model))                                   # embeddings de los tokens (sin entrenar)
W_Q, W_K, W_V = (rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_k)) for _ in range(3))

out_full, w_full = attention(X @ W_Q, X @ W_K, X @ W_V)
causal = np.tril(np.ones((n, n), dtype=bool))                       # el token i solo puede ver los tokens 0..i
out_causal, w_causal = attention(X @ W_Q, X @ W_K, X @ W_V, mask=causal)
print(f"Forma de la salida: {out_full.shape} (un vector nuevo por token)")
print(f"Cada fila de pesos suma 1: {np.allclose(w_full.sum(axis=1), 1)}")
print(f"Con la máscara, peso dado a los tokens futuros: {w_causal[~causal].sum():.1f}")

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, w, title in [(axes[0], w_full, 'Autoatención (cada token ve todos los tokens)'),
                     (axes[1], w_causal, 'Autoatención causal (sin mirar hacia delante)')]:
    ax.imshow(w, cmap='Blues', vmin=0, vmax=1)
    ax.set_xticks(range(n), sentence, rotation=45, ha='right')
    ax.set_yticks(range(n), sentence)
    ax.set_xlabel('atiende a (claves)')
    ax.set_ylabel('token (consultas)')
    ax.set_title(title)
plt.tight_layout()
plt.show()

Forma de la salida: (9, 8) (un vector nuevo por token)
Cada fila de pesos suma 1: True
Con la máscara, peso dado a los tokens futuros: 0.0

Con pesos aleatorios el patrón todavía no significa nada: el entrenamiento es lo que ajusta $W^Q$, $W^K$ y $W^V$ para que los tokens relacionados produzcan puntuaciones altas. Lo que importa aquí es la estructura. Cada token está conectado con todos los demás en un solo paso, sea cual sea la distancia entre ellos, que es justo lo que le faltaba a la RNN. Y la máscara causal deja solo el triángulo inferior.

Paso 6: La atención multicabeza

---

Una sola atención calcula una única noción de "relevancia". La atención multicabeza calcula $h$ a la vez, cada una con sus propias proyecciones, concatena los resultados y los mezcla con una matriz final $W^O$:

$$\text{head}_i = \text{Attention}(XW_i^Q, XW_i^K, XW_i^V) \qquad \text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\,W^O$$

Con la elección habitual $d_k = d_v = d_{\text{model}} / h$, la concatenación vuelve a tener $d_{\text{model}}$ columnas.

def init_multihead(d_model, n_heads, rng):
    d_k = d_model // n_heads
    heads = [tuple(rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_k)) for _ in range(3)) for _ in range(n_heads)]
    W_O = rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_model))
    return heads, W_O

def multihead(X, heads, W_O, mask=None):
    outputs = [attention(X @ Wq, X @ Wk, X @ Wv, mask)[0] for Wq, Wk, Wv in heads]
    return np.concatenate(outputs, axis=1) @ W_O

heads, W_O = init_multihead(d_model=8, n_heads=2, rng=rng)
mh = multihead(X, heads, W_O, mask=causal)
n_params = sum(w.size for h in heads for w in h) + W_O.size
print(f"2 cabezas de dimensión {heads[0][0].shape[1]}; forma de la salida {mh.shape}; parámetros {n_params}")

2 cabezas de dimensión 4; forma de la salida (9, 8); parámetros 256

Paso 7: La codificación posicional

---

Las notas de Marta señalan una propiedad sorprendente: la atención, por sí sola, es ciega al orden de las palabras. Si barajamos los tokens, las salidas se barajan exactamente igual, pero no cambian. Comprobémoslo:

perm = rng.permutation(n)
out = multihead(X, heads, W_O)                       # sin máscara, para que el orden no lo imponga la máscara
out_shuffled = multihead(X[perm], heads, W_O)
print(f"Barajar los tokens solo baraja las salidas: {np.allclose(out[perm], out_shuffled)}")

Barajar los tokens solo baraja las salidas: True

"El paciente atendió al doctor" y "el doctor atendió al paciente" parecerían iguales. El Transformer lo resuelve sumando a cada vector de token una codificación posicional hecha de senos y cosenos de frecuencias distintas:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \qquad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$$

Marta demuestra que, para cualquier desplazamiento $k$, la codificación de la posición $pos + k$ es una rotación de la codificación de $pos$, con una matriz que solo depende de $k$. Construimos la codificación, la dibujamos y comprobamos esa propiedad numéricamente:

def positional_encoding(n_positions, d_model):
    pos = np.arange(n_positions)[:, None]
    omega = 1 / 10000 ** (np.arange(0, d_model, 2) / d_model)       # una frecuencia por cada pareja de dimensiones
    pe = np.zeros((n_positions, d_model))
    pe[:, 0::2] = np.sin(pos * omega)
    pe[:, 1::2] = np.cos(pos * omega)
    return pe, omega

PE, omega = positional_encoding(100, 16)
plt.figure(figsize=(10, 4))
plt.imshow(PE.T, aspect='auto', cmap='RdBu')
plt.xlabel('posición en la secuencia')
plt.ylabel('dimensión')
plt.title('Codificación posicional sinusoidal: frecuencias rápidas arriba, lentas abajo')
plt.colorbar()
plt.show()

k = 3
R = np.zeros((16, 16))
for i, w in enumerate(omega):                        # matriz diagonal por bloques de rotaciones 2x2
    R[2 * i:2 * i + 2, 2 * i:2 * i + 2] = [[np.cos(k * w), np.sin(k * w)], [-np.sin(k * w), np.cos(k * w)]]
error = np.max(np.abs(PE[k:] - PE[:-k] @ R.T))
print(f"Mayor diferencia entre PE(pos + {k}) y R_{k} PE(pos), en todas las posiciones: {error:.1e}")

PE8 = positional_encoding(n, 8)[0]                  # la codificación para nuestros tokens de 8 dimensiones
out_pe = multihead(X + PE8, heads, W_O)
out_pe_shuffled = multihead(X[perm] + PE8, heads, W_O)
print(f"Con codificación posicional, barajar solo baraja las salidas: {np.allclose(out_pe[perm], out_pe_shuffled)}")

Mayor diferencia entre PE(pos + 3) y R_3 PE(pos), en todas las posiciones: 2.4e-15
Con codificación posicional, barajar solo baraja las salidas: False

La propiedad de la rotación se cumple con una diferencia de $2{,}4 \times 10^{-15}$, es decir, exactamente salvo errores de redondeo: avanzar 3 posiciones es la misma rotación en cualquier punto de la secuencia. Y una vez sumada la codificación posicional, barajar los tokens sí cambia las salidas, así que el modelo ya puede distinguir "el paciente atendió al doctor" de "el doctor atendió al paciente".

Paso 8: Un bloque Transformer completo

---

Un bloque Transformer lo combina todo, como en las notas de Marta: la atención multicabeza, que mezcla información entre tokens; una red feed-forward aplicada a cada token por separado, $\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2$; y, alrededor de cada una, una conexión residual más una normalización de capa, $\text{LayerNorm}(x + \text{Subcapa}(x))$. La conexión residual es una "autopista" para los gradientes, y la normalización da a cada vector de token media 0 y varianza 1. Los grandes modelos de lenguaje apilan decenas de estos bloques.

def layer_norm(x, eps=1e-5):
    return (x - x.mean(axis=-1, keepdims=True)) / np.sqrt(x.var(axis=-1, keepdims=True) + eps)

def transformer_block(x, heads, W_O, W1, b1, W2, b2, mask=None):
    x = layer_norm(x + multihead(x, heads, W_O, mask))            # atención + Add & Norm
    ffn = np.maximum(0, x @ W1 + b1) @ W2 + b2                    # feed-forward, token a token
    return layer_norm(x + ffn)                                     # FFN + Add & Norm

d_model, d_ff = 8, 32
W1, b1 = rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_ff)), np.zeros(d_ff)
W2, b2 = rng.normal(0, 1 / np.sqrt(d_ff), (d_ff, d_model)), np.zeros(d_model)

x_in = X + PE8
y = transformer_block(x_in, heads, W_O, W1, b1, W2, b2, mask=causal)
n_block = n_params + W1.size + b1.size + W2.size + b2.size
print(f"Entrada {x_in.shape} -> salida {y.shape}")
print(f"Media de cada token de salida: {y.mean(axis=1).round(3)}")
print(f"Varianza de cada token de salida: {y.var(axis=1).round(3)}")
print(f"Parámetros de este bloque pequeño: {n_block}")

Entrada (9, 8) -> salida (9, 8)
Media de cada token de salida: [-0.  0. -0. -0.  0. -0.  0.  0. -0.]
Varianza de cada token de salida: [1. 1. 1. 1. 1. 1. 1. 1. 1.]
Parámetros de este bloque pequeño: 808

El bloque conserva la forma de la secuencia: entran 9 tokens de 8 dimensiones y salen 9 tokens de 8 dimensiones, cada uno mezclado ahora con su contexto. Tras la última normalización de capa, cada token tiene media 0 y varianza 1. Este bloque pequeño tiene 808 parámetros; en un gran modelo de lenguaje la misma estructura se repite en decenas de capas con miles de dimensiones, y así es como los modelos llegan a miles de millones de parámetros.

Conclusión

---

• Las redes recurrentes pasan la información de token en token, y su influencia se desvanece exponencialmente con la distancia. La atención conecta cada token con todos los demás en un solo paso.
• La atención de producto escalar escalado compara las consultas con las claves, convierte las puntuaciones en pesos con softmax y mezcla los valores con esos pesos. Dividir por $\sqrt{d_k}$ evita que la softmax se sature.
• La máscara causal impide que un token mire al futuro, que es lo que permite a un modelo generar texto de izquierda a derecha. La atención multicabeza calcula varias nociones de relevancia en paralelo.
• La atención es ciega al orden, así que el Transformer añade una codificación posicional, cuyos senos y cosenos convierten un desplazamiento de $k$ posiciones en una rotación.
• Un bloque Transformer añade una red feed-forward, conexiones residuales y normalización de capa. Aquí todo tiene pesos aleatorios: el entrenamiento con cantidades enormes de texto, como explica el libro, es lo que convierte esta maquinaria en un modelo de lenguaje.