English · Español

Lab 03 — `matmul`, `softmax`, `cross_entropy`: las tres ops de alto riesgo¶

Objetivo: implementar las tres ops de las que depende el resto del currículo. matmul es el caballo de batalla de toda capa; softmax es el caballo de batalla de todo clasificador; cross_entropy (desde logits) es la pérdida que dirige el MLP de la fase 9 y el transformer de la fase 17. Hazlas bien aquí — en FP64, con los dos oráculos en verde — y la fase 9+ hereda un cimiento fiable.

Tiempo estimado: 4–5 horas. Este es el lab más largo de la fase.

Prerrequisitos: Labs 00 + 01 + 02. Teoría 03-matmul-and-softmax-grads.md releída de punta a punta. stable_softmax y stable_cross_entropy de la fase 2 abiertos en otra ventana — reutilizaremos la intuición, no el código.

Lo que produces¶

Tres ops añadidas a src/minitorch/tensor.py: Tensor.matmul, Tensor.softmax, función a nivel de módulo cross_entropy(logits, targets, reduction='mean').
El operador @ cableado: A @ B llama a A.matmul(B).
tests/test_matmul.py, tests/test_softmax.py, tests/test_cross_entropy.py — cada uno con cross PyTorch + gradcheck + casos límite.
Un test de gradiente de punta a punta con sabor gramatical: (one_hot_person @ tense_logits → softmax → cross_entropy) coincide con PyTorch.

Por qué estas tres son "de alto riesgo"¶

Los bugs en add aparecen en el siguiente borde de test. Los bugs en matmul/softmax/cross_entropy aparecen como fallos de entrenamiento silenciosos tres días después: el modelo entrena, la pérdida decrece, la precisión de validación se estanca 5% por debajo de donde debería. La estrategia de testing de la fase 8 existe principalmente para detectar bugs en estas tres ops antes de que contaminen las fases 9–22.

Las tres operaciones que no te puedes permitir tener mal: matmul (sesgo silencioso en cualquier capa lineal), softmax (overflow/underflow numérico), cross-entropy (la fórmula bonita softmax - one_hot es trivial de derivar pero fácil de implementar mal en el caso batched). Si las tres pasan PyTorch cross-check y gradcheck, el resto del currículo descansa sobre cimientos sólidos.

TODOs¶

Bloque A — `matmul`¶

__matmul__(self, other) -> Tensor: forward self.data @ other.data. Backward por la derivación de theory/03:
self.grad += out.grad @ other.data.swapaxes(-1, -2)
other.grad += self.data.swapaxes(-1, -2) @ out.grad
Usa _unbroadcast en cada uno: matmul broadcastea las dims de batch. A.shape = (1, 3, 4) @ B.shape = (B, 4, 5) produce (B, 3, 5). El padre (1, 3, 4) necesita el gradiente sumado sobre la dim de batch.
Cablea __matmul__ para que A @ B funcione.
Matriz de tests: 2D × 2D, 2D × 1D (vector), 1D × 2D, batched 3D × 3D, batched 3D × 2D (con broadcast).

Bloque B — `softmax`¶

softmax(self, axis=-1) -> Tensor: el forward usa el truco de estabilidad max-subtraction (fase 2):
```
shifted = self.data - self.data.max(axis=axis, keepdims=True)
exps = np.exp(shifted)
out_data = exps / exps.sum(axis=axis, keepdims=True)
```
Backward por la forma producto Jacobiano-vector (teoría 03):
```
# out.data is `s` (the softmax output, shape == self.shape)
# out.grad is `dy`
s = out.data
dy = out.grad
# dx = s * (dy - (dy * s).sum(axis=axis, keepdims=True))
```
Deriva en papel por qué este es el Jacobiano denso colapsado a una fórmula vectorial. Dedica 20 minutos a la derivación; el resultado es una de esas identidades que se vuelven "obvias" una vez vistas.
Testea contra PyTorch en tensores rank-1, rank-2, rank-3. Stress test con logits de magnitud 1e3 (no debe haber overflow).

Bloque C — `cross_entropy` desde logits¶

cross_entropy(logits: Tensor, targets: IntArray, reduction: str = 'mean') -> Tensor (función a nivel de módulo):
targets es un np.ndarray plano de enteros (no un Tensor) — los targets no tienen gradientes.

El forward usa el truco log-sum-exp (la stable_cross_entropy de la fase 2):

# logits.shape = (B, V); targets.shape = (B,)
lse = log_sum_exp(logits.data, axis=-1)  # shape (B,)
target_logit = np.take_along_axis(logits.data, targets[:, None], axis=-1).squeeze(-1)
losses = lse - target_logit              # shape (B,)
if reduction == 'mean':
    out_data = losses.mean()
elif reduction == 'sum':
    out_data = losses.sum()
elif reduction == 'none':
    out_data = losses

Backward — la identidad preciosa:

# Cache the softmax probabilities `p` (NOT the logits' exps — recompute stably).
p = softmax(logits.data, axis=-1)        # stable
grad = p.copy()
np.add.at(grad, (np.arange(B), targets), -1.0)  # subtract one-hot
if reduction == 'mean':
    grad /= B
# multiply by upstream out.grad (scalar for mean/sum, vector for none)
logits.grad += grad * upstream

No implementes cross_entropy como (-target_log_softmax).mean() encadenado desde ops existentes. La op combinada es numéricamente estable y evita materializar log(softmax). Impleméntala como una única op fundida. (Esto refleja F.cross_entropy de PyTorch.)

Bloque D — test de gramática de punta a punta¶

Después de que los Bloques A–C funcionen aisladamente, construye el mini-grafo gramatical y verifica contra PyTorch.

def test_grammar_pipeline_matches_pytorch():
    # The §A13 baseline: select a person, project against tense logits, classify.
    rng = np.random.default_rng(42)
    person_onehot = np.array([0.0, 1.0, 0.0])           # "you" (2nd person)
    W = rng.standard_normal((3, 5))                     # person → tense
    targets = np.array([2])                             # past-simple
    # ours
    P = Tensor(person_onehot[None, :], requires_grad=False)  # (1, 3)
    Wt = Tensor(W, requires_grad=True)                       # (3, 5)
    logits = P @ Wt                                          # (1, 5)
    loss = cross_entropy(logits, targets)
    loss.backward()
    # pytorch
    Pt = torch.tensor(person_onehot[None, :], dtype=torch.float64)
    Wpt = torch.tensor(W, dtype=torch.float64, requires_grad=True)
    logits_t = Pt @ Wpt
    loss_t = torch.nn.functional.cross_entropy(logits_t, torch.tensor(targets, dtype=torch.long))
    loss_t.backward()
    np.testing.assert_allclose(Wt.grad, Wpt.grad.numpy(), rtol=1e-7)

Stress tests de estabilidad numérica¶

Estos tests deben pasar. Son la diferencia entre un autograd que funciona en el ejemplo de libro y uno que funciona en el transformer de la fase 17.

def test_softmax_large_logits():
    # Logits with magnitude 1e3 must not overflow.
    x = Tensor(np.array([1000.0, 1001.0, 999.0]), requires_grad=True)
    y = x.softmax()
    y.sum().backward()
    assert np.all(np.isfinite(y.data))
    assert np.all(np.isfinite(x.grad))
    assert np.isclose(y.data.sum(), 1.0)

def test_cross_entropy_confident_correct():
    # Confident, correct prediction → ~0 loss; gradient is ~0.
    logits = Tensor(np.array([[1000.0, 0.0, 0.0]]), requires_grad=True)
    targets = np.array([0])
    loss = cross_entropy(logits, targets)
    loss.backward()
    assert loss.data < 1e-10
    assert np.allclose(logits.grad, np.array([[0.0, 0.0, 0.0]]), atol=1e-9)

def test_cross_entropy_confident_wrong():
    # Confident, wrong prediction → loss ≈ 1000; gradient is (1, 0, -1).
    logits = Tensor(np.array([[1000.0, 0.0, 0.0]]), requires_grad=True)
    targets = np.array([2])
    loss = cross_entropy(logits, targets)
    loss.backward()
    assert np.isclose(loss.data, 1000.0, atol=1e-3)
    expected = np.array([[1.0, 0.0, -1.0]])
    np.testing.assert_allclose(logits.grad, expected, atol=1e-9)

Restricciones¶

softmax usa siempre el truco del max. Nada de exp / sum(exp) ingenuo en ninguna parte del forward.
cross_entropy nunca materializa log(softmax). Usa LSE.
El backward de matmul usa swapaxes(-1, -2), no .T. .T solo transpone limpiamente arrays 2D; para matmul batcheado necesitamos intercambiar las dos últimas dims preservando las dims de batch.
Los targets a cross_entropy son enteros, no Tensors. Rechaza targets one-hot en la frontera del API — registra un error ruidoso.

Patrones de test¶

Para cada op, escribe: 1. Cross-check contra PyTorch FP64 en los pares de forma (shape) típicos. 2. Gradcheck en una forma (shape) pequeña (≤ (3, 4)). 3. Caso límite (dim de tamaño 1, reducción escalar, etc.). 4. Test de estrés / estabilidad (logits de gran magnitud para softmax / CE).

Para cross_entropy, testea además reduction='sum' y reduction='none'.

Condiciones de parada¶

Hecho cuando:

matmul, softmax, cross_entropy implementadas y testeadas.
Cross-check PyTorch FP64 verde para toda combinación de forma (shape) de op.
Gradcheck verde para cada op con eps=1e-6, atol=1e-4.
Stress tests de estabilidad todos verdes.
Test de gramática de punta a punta verde.
mypy --strict limpio en src/minitorch/.
Puedes re-derivar ∂L/∂x = softmax(x) - one_hot(y) en una página en blanco, por índices, en menos de 5 minutos.

Escollos¶

Confusión en la fórmula del backward de softmax. Dos formas equivalentes: (a) Jacobiano completo ∂s_i/∂x_j = s_i(δ_{ij} - s_j) y luego multiplicar matricialmente por dy; (b) la forma vectorial s * (dy - sum(dy * s)). La forma (b) es O(N); la (a) es O(N²). Usa (b). Deriva (b) a partir de (a) una vez en papel.
El backward de cross_entropy olvida el / B para reduction='mean'. Te pilla al comparar con PyTorch — gradientes desviados exactamente por el tamaño del batch.
Transposición en matmul batcheado. np.matmul((B, M, N), (B, N, P)) → (B, M, P). La "transposición" en el backward es .swapaxes(-1, -2). Si escribiste .T, los casos batcheados producirán formas (shapes) erróneas silenciosamente. Añade un test 3D batcheado.
Softmax con axis != -1. Todo el broadcasting en la fórmula del backward asume que axis es el último; para axis general necesitas keepdims=True de forma consistente. Testéalo con axis=0 en un tensor (3, 4).
cross_entropy con un único ejemplo. logits.shape = (1, V), targets.shape = (1,). Caso límite habitual; testéalo.
Dtype de los targets. PyTorch quiere torch.long; nuestro cross_entropy debería aceptar np.int64 (o int32). Documéntalo.

Cuándo consultar `solutions/`¶

Después de que los cuatro bloques pasen todos los tests, incluyendo el de gramática de punta a punta. solutions/03-matmul-softmax-ce-ref.md (al abrir la fase) compara tus tres closures y la op fundida cross_entropy contra las implementaciones canónicas. Incluye además una página de rescate de un folio "si has pasado más de 90 minutos depurando el backward de softmax, lee esto primero".

Este es el último lab de la fase 8. Cuando los tres labs estén verdes, el MLP de experiments/08-tense-classifier/ se vuelve factible — y ese experimento es el que cierra la fase.

Lab 03 — matmul, softmax, cross_entropy: las tres ops de alto riesgo¶