English · Español

04 — Forward + backward desarrollado a través de `Linear → GELU → Linear`¶

🇪🇸 Aquí desarmamos un MLP de 2 capas con activación GELU en medio y mostramos cada matmul, cada bias broadcast y cada gradiente con shapes reales. Es el ejercicio que convierte la "fórmula de cadena" de la Fase 7 en un objeto numérico que cabe en la memoria de trabajo. Lo hacemos sobre un batch de §A13 (4 ejemplos, 23-d one-hot de verbo⊕persona; salida 5 logits).

Anclas: LYNX_CORTEX.md §4 / PHASE 9; LYNX_CORTEX_ADDENDUM.md §A13. Fase 4 §02 regla de la cadena; Fase 8 §03 gradiente de matmul.

Setup¶

Un MLP de dos capas con GELU en medio, escrito en la API minimodel de la Fase 9:

mlp = Sequential(
    Linear(23, 16),   # W1: (16, 23), b1: (16,)
    GELU(),
    Linear(16, 5),    # W2: (5, 16),  b2: (5,)
)

El batch de entrada X tiene forma (B, 23) con B = 4. Cada fila es one_hot(verb_id) ⊕ one_hot(person_id) — ocho verbos en este micro-grid (4 regulares: work, play, walk, talk; 4 irregulares: be, have, do, go) rellenados a 20 + 3 personas = 23 features. Los targets y son índices de tiempo en {0, 1, 2, 3, 4}.

Rastrearemos shapes a través del forward y del gradiente en modo reverso computado por el autograd de la Fase 8, pero anotaremos lo que cada op backward hace a mano.

Forward pass con shapes explícitos¶

Capa 1 — `Linear(23, 16)`¶

X        : (B, 23)        = (4, 23)
W1       : (16, 23)
b1       : (16,)
Z1 = X @ W1.T + b1
   shape : (4, 16)

El coste del matmul es B · 16 · 23 = 4 · 368 = 1472 multiply-adds. El bias b1 hace broadcast de (16,) a (4, 16) a lo largo del eje 0.

Activación — `GELU`¶

H1 = GELU(Z1)
   shape : (4, 16)

Punto a punto, sin cambio de forma. Usamos la aproximación tanh:

\[ \text{GELU}(z) \approx 0.5\,z\,\bigl(1 + \tanh\bigl[\sqrt{2/\pi}\,(z + 0.044715\,z^{3})\bigr]\bigr) \]

Para una entrada Z1[i, j] = 1.2, GELU(1.2) ≈ 1.0617; para Z1[i, j] = -0.7, GELU(-0.7) ≈ -0.1671. (Valores de sanity; la Fase 17 deriva la aproximación.)

Capa 2 — `Linear(16, 5)`¶

H1       : (4, 16)
W2       : (5, 16)
b2       : (5,)
logits = H1 @ W2.T + b2
   shape : (4, 5)

Coste: B · 5 · 16 = 320 multiply-adds.

Pérdida¶

Cross-entropy sobre los 5 tiempos:

p = softmax(logits, axis=-1)      # (4, 5)
loss = -mean(log(p[arange(B), y])) # scalar

El camino forward completo es dos matmuls, un GELU punto a punto, un softmax, un gather, una media — 6 ops en total. Memoriza esta cuenta; es idéntica (módulo factores) para cada MLP de dos capas que escribirás jamás.

Backward pass — regla de la cadena, op por op¶

Propagamos dL/d• desde la pérdida hacia los parámetros. Símbolo: \(\nabla_x L \equiv \partial L / \partial x\).

Paso 1 — gradiente fundido de softmax + cross-entropy¶

Esta es la identidad más elegante de todo el MLP. La Fase 4 §02 la deriva; la Fase 8 §03 la implementa.

\[ \nabla_{\text{logits}} L = \frac{1}{B}\,(p - \mathrm{onehot}(y)) \]

Shape: (4, 5). Concretamente, si el tiempo verdadero de la fila 0 es 2 y p[0] = [0.1, 0.2, 0.5, 0.1, 0.1], entonces (p - onehot(y))[0] = [0.1, 0.2, -0.5, 0.1, 0.1], dividido por B = 4. La entrada negativa es el gradiente empujando hacia arriba el logit de la clase verdadera.

Paso 2 — backward a través de `Linear(16, 5)`¶

El forward era logits = H1 @ W2.T + b2. Así que:

\[ \nabla_{W2} L = (\nabla_{\text{logits}} L)^{\!\top} \cdot H1 \qquad \nabla_{b2} L = \sum_{i=1}^{B} (\nabla_{\text{logits}} L)_{i,\cdot} \qquad \nabla_{H1} L = (\nabla_{\text{logits}} L) \cdot W2 \]

Shapes:

∇_W2     = (5, 4) @ (4, 16) → (5, 16)   matches W2
∇_b2     = sum_{axis=0} (4, 5) → (5,)   matches b2
∇_H1     = (4, 5) @ (5, 16) → (4, 16)   matches H1

Este es el patrón estándar. La Fase 8 §03 deletrea por qué el gradiente de W recoge la activación del lado de la entrada transpuesta: sale de d(x @ W.T)/dW = x y la regla de la cadena.

Paso 3 — backward a través de GELU¶

Multiplicación punto a punto por la derivada de GELU:

\[ \nabla_{Z1} L = \nabla_{H1} L \;\odot\; \text{GELU}'(Z1) \]

Con la aproximación tanh, GELU'(z) se puede expresar en forma cerrada (deberes — el Lab 01 lo implementa). Numéricamente, GELU'(1.2) ≈ 1.046, GELU'(0.0) = 0.5, GELU'(-0.7) ≈ 0.151. Nota GELU'(z) > 0 en todas partes, así que a diferencia de ReLU no hay máscara de gradiente cero duro.

Shape: (4, 16). Igual que H1.

Paso 4 — backward a través de `Linear(23, 16)`¶

\[ \nabla_{W1} L = (\nabla_{Z1} L)^{\!\top} \cdot X \qquad \nabla_{b1} L = \sum_{i=1}^{B} (\nabla_{Z1} L)_{i,\cdot} \qquad \nabla_{X} L = (\nabla_{Z1} L) \cdot W1 \]

Shapes:

∇_W1     = (16, 4) @ (4, 23) → (16, 23)   matches W1
∇_b1     = sum_{axis=0} (4, 16) → (16,)   matches b1
∇_X      = (4, 16) @ (16, 23) → (4, 23)   matches X

∇_X no es un gradiente de parámetro, pero la capa de embedding de la Fase 13 lo consumirá (el backward del embedding necesita el gradiente con respecto a su salida, que es X aquí).

Memoria en el momento del backward¶

La Fase 18 (bucle de entrenamiento) se preocupará de que el backward mantenga vivas las activaciones del forward. Para nuestro MLP de 2 capas a B = 4:

X — (4, 23) — 92 floats
Z1 — (4, 16) — 64 floats (guardado para el backward de GELU)
H1 — (4, 16) — 64 floats (guardado para el backward de Linear(16, 5))
logits — (4, 5) — 20 floats (guardado para el backward de softmax-CE)
Parámetros — W1, b1, W2, b2 — 368 + 16 + 80 + 5 = 469 floats

Total de activaciones a mantener: ~240 floats. A FP32 son 960 bytes. Más pequeño que un único set de línea de caché L1 en el i5-8250U. Este es el dividendo de microscopic-scope de §A13.

Cuenta de parámetros vs FLOPs forward vs FLOPs backward¶

Fase	FLOPs (aprox)
Forward `X @ W1.T`	`2 · B · 16 · 23`
Forward `H1 @ W2.T`	`2 · B · 5 · 16`
Backward `∇_W1`	`2 · 16 · B · 23`
Backward `∇_X`	`2 · B · 16 · 23`
Backward `∇_W2`	`2 · 5 · B · 16`
Backward `∇_H1`	`2 · B · 5 · 16`

Súmalo: el backward es ~2× el coste del forward para una cadena de linears. Este es el ratio titular que conduce cada presupuesto de memoria/cómputo que harás en la Fase 18+ (reportes de ejecución de mlflow), Fase 22 (KV cache), Fase 27 (Flash attention).

Un bug común que este ejercicio previene¶

Transponer la matriz incorrecta en ∇_W. El patrón es ∇_W = (∇_out)^T @ x_in, no ∇_W = ∇_out @ x_in^T. Ambos producen una matriz de la forma correcta para problemas cuadrados, así que un test unitario en Linear(4, 4) pasa silenciosamente. Ejecuta el lab en Linear(23, 16) y el desajuste de shape revienta inmediatamente. Esta es la razón por la que los tests de la Fase 9 usan shapes no cuadradas.

Referencias¶

Goodfellow, Bengio, Courville, Deep Learning, MIT Press, 2016. Cap. 6 §6.5 "Back-Propagation and Other Differentiation Algorithms" desarrolla la misma cadena a mano.
La identidad del gradiente fundido softmax-CE también aparece en la nota a pie §A.1 de Vaswani et al. 2017 (y la precede en décadas — Bridle 1990).

Recapitulación de un párrafo¶

Un MLP de 2 capas sobre una entrada (B, 23) hace dos matmuls de forma (B,23)×(23,16) y (B,16)×(16,5) con un GELU punto a punto en medio, costando aproximadamente 1800 multiply-adds por ejemplo para forward, el doble para backward. El backward pass propaga (p - onehot(y))/B a través de la cadena, con cada Linear contribuyendo el patrón ∇_W = ∇_out^T @ x_in, ∇_b = sum(∇_out, axis=0), ∇_in = ∇_out @ W. GELU contribuye un multiplicador positivo punto a punto — sin gradiente cero duro. La memoria total de activaciones en B=4 es menor de 1 KB. Memoriza la aritmética de shapes de esta página: cada fase posterior (atención, MoE, LoRA) reutiliza el mismo patrón de regla de la cadena.

Prev: 03-optimizers.md

04 — Forward + backward desarrollado a través de Linear → GELU → Linear¶