Backprop y atención desde cero

Objetivo de maestría

entender backpropagation y el mecanismo de atención implementándolos a mano, sin nn.Module mágicos. Si puedes escribir un autograd mínimo y una self-attention desde cero, nada en el resto del curso es una caja negra. Esta es la lección que convierte "uso PyTorch" en "sé qué hace PyTorch".

1.1Por qué construir desde cero

Puedes usar transformers toda la vida sin saber cómo fluye un gradiente. Pero la diferencia entre el percentil 50 y el percentil top es saber por qué algo no entrena: un gradiente que explota, una inicialización mala, una normalización en el sitio equivocado. Eso solo se internaliza construyéndolo. La referencia canónica es nn-zero-to-hero de Karpathy (micrograd → makemore → GPT); aquí destilamos el núcleo.

1.2Autograd mínimo (el motor de backprop)

Backprop es la regla de la cadena aplicada sistemáticamente sobre un grafo de operaciones. Un Value que recuerda cómo se calculó puede propagar gradientes hacia atrás:

python

1# lab_n5l1_micrograd.py — autograd escalar mínimo (estilo micrograd de Karpathy)
2import math
3
4class Value:
5    def __init__(self, data, _children=(), _op=""):
6        self.data = data
7        self.grad = 0.0                 # dL/d(este value); empieza a 0
8        self._backward = lambda: None   # cómo propagar el grad a los hijos
9        self._prev = set(_children)     # nodos de los que depende
10        self._op = _op
11
12    def __add__(self, other):
13        other = other if isinstance(other, Value) else Value(other)
14        out = Value(self.data + other.data, (self, other), "+")
15        def _backward():
16            # d(a+b)/da = 1, d(a+b)/db = 1  -> el grad fluye igual a ambos
17            self.grad  += 1.0 * out.grad
18            other.grad += 1.0 * out.grad
19        out._backward = _backward
20        return out
21
22    def __mul__(self, other):
23        other = other if isinstance(other, Value) else Value(other)
24        out = Value(self.data * other.data, (self, other), "*")
25        def _backward():
26            # d(a*b)/da = b, d(a*b)/db = a  (regla del producto)
27            self.grad  += other.data * out.grad
28            other.grad += self.data  * out.grad
29        out._backward = _backward
30        return out
31
32    def tanh(self):
33        t = math.tanh(self.data)
34        out = Value(t, (self,), "tanh")
35        def _backward():
36            # d(tanh)/dx = 1 - tanh^2
37            self.grad += (1 - t**2) * out.grad
38        out._backward = _backward
39        return out
40
41    def backward(self):
42        # 1) ordena el grafo topológicamente (un nodo se procesa tras sus dependencias)
43        topo, visited = [], set()
44        def build(v):
45            if v not in visited:
46                visited.add(v)
47                for child in v._prev: build(child)
48                topo.append(v)
49        build(self)
50        # 2) propaga desde la salida: dL/dL = 1, luego hacia atrás
51        self.grad = 1.0
52        for v in reversed(topo):
53            v._backward()
54
55# Demo: un "neuron" y = tanh(w*x + b), y backprop manual
56x, w, b = Value(2.0), Value(-3.0), Value(1.0)
57y = (w*x + b).tanh()
58y.backward()
59print(f"y={y.data:.4f}  dy/dw={w.grad:.4f}  dy/dx={x.grad:.4f}  dy/db={b.grad:.4f}")

Líneas no triviales explicadas:

Cada operación crea un nodo out que guarda cómo calcular el gradiente hacia sus hijos (_backward). Eso es todo lo que un framework de autograd hace, magnificado.
+= en los gradientes (no =): si un valor se usa en varios sitios, sus gradientes se acumulan (regla de la suma de derivadas parciales). Olvidar el += es el bug clásico de autograd casero.
El orden topológico garantiza que cuando procesas un nodo, todos los que dependen de él (más cerca de la pérdida) ya aportaron su gradiente.
self.grad = 1.0 al inicio: dL/dL = 1, la semilla de toda la retropropagación.

Verifica tus gradientes contra PyTorch (torch.autograd) para ganar confianza. Cuando esto te resulte obvio, entiendes backprop.

1.3Self-attention desde cero

La atención es "para cada token, mira a los demás y agrega su información ponderada por relevancia". En matemáticas: proyecciones Q, K, V; pesos = softmax de QKᵀ escalado; salida = pesos·V.

python

1# lab_n5l1_attention.py — self-attention causal desde cero en PyTorch (sin nn.MultiheadAttention)
2import torch, torch.nn.functional as F
3
4def self_attention(x, Wq, Wk, Wv, causal=True):
5    # x: [B, T, d]   Wq,Wk,Wv: [d, d_head]
6    B, T, d = x.shape
7    Q = x @ Wq                       # [B, T, d_head]  "qué busco"
8    K = x @ Wk                       # [B, T, d_head]  "qué ofrezco"
9    V = x @ Wv                       # [B, T, d_head]  "qué transmito"
10    d_head = Q.size(-1)
11    scores = Q @ K.transpose(-2, -1) / d_head**0.5   # [B, T, T] afinidad token-token
12    if causal:
13        mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
14        scores = scores.masked_fill(mask, float("-inf"))   # un token no ve el futuro
15    weights = F.softmax(scores, dim=-1)                     # [B, T, T] suman 1 por fila
16    return weights @ V                                      # [B, T, d_head] agregación ponderada
17
18# Demo
19B, T, d, dh = 2, 5, 16, 16
20x = torch.randn(B, T, d)
21Wq, Wk, Wv = (torch.randn(d, dh) for _ in range(3))
22out = self_attention(x, Wq, Wk, Wv)
23print(out.shape)   # [2, 5, 16]

Líneas no triviales explicadas:

/ d_head**0.5 (escalado): sin él, los productos QKᵀ crecen con la dimensión → el softmax se satura (gradientes minúsculos). El escalado mantiene la varianza controlada. Es el "scaled" de scaled dot-product attention.
La máscara causal (triu con -inf): pone a -inf las posiciones futuras antes del softmax → su peso se vuelve 0 → cada token solo atiende a sí mismo y al pasado. Es lo que hace al modelo autoregresivo (predecir el siguiente sin ver la respuesta).
weights @ V: cada token recibe una mezcla de los V de los tokens a los que atiende, ponderada por relevancia. Eso es "atención".
Multi-head = correr esto en paralelo con varias Wq/Wk/Wv (cada "cabeza" aprende un tipo de relación) y concatenar. Lo añades en L2.

1.4Ejercicios

E1. Añade __pow__ y exp() a la clase Value (con sus _backward) y verifica los gradientes contra torch.autograd en una expresión compuesta.

E2. Quita el escalado /d_head**0.5 y observa qué le pasa a los pesos de atención cuando d_head es grande (256). ¿Por qué se saturan?

E3. Implementa multi-head attention (h cabezas) reutilizando self_attention. Verifica que la salida concatenada tiene dimensión h × d_head.

Pista E3

Crea h tripletes (Wq,Wk,Wv) de dim [d, d_head], corre `self_attention` por cada uno, concatena en la última dim, y proyecta con una Wo de [h*d_head, d].

1.5Trampas comunes

Usar = en vez de += en los gradientes del autograd → gradientes incorrectos cuando un valor se reutiliza.
Olvidar el escalado de atención → saturación del softmax, no aprende.
Aplicar la máscara causal después del softmax (mal) en vez de antes (bien).

1.6Referencias

Karpathy, Neural Networks: Zero to Hero (nn-zero-to-hero): micrograd y makemore. "Attention Is All You Need" (Vaswani et al. 2017). CS336 (arquitecturas).