English · Español

Break — Saltarse el `unbroadcast` (suma por ejes) en el backward de add¶

Bug clásico de la fase 8: olvidar que el broadcasting en forward implica suma a lo largo de los ejes replicados en backward. Sin esa suma, el gradiente sale con forma errónea — y o crashea ruidosamente, o (peor) la forma encaja por accidente y los números son basura.

Objetivo: cualquier op add en src/minigrad/tensor.py (o tu equivalente).

Hipótesis¶

El aprendiz predice: "Sin el paso explícito unbroadcast(grad, parent.shape) en el backward de una op que broadcasteó en forward, aparecen dos regímenes: (1) cuando el eje de broadcasting fue creado (un operando tenía un menor número de dimensiones que el otro), la forma (shape) del gradiente upstream difiere de la del padre y o bien crashea o bien se corrompe silenciosamente vía el broadcasting silencioso de NumPy en la asignación; (2) cuando solo se expandió un eje de longitud 1, la forma (shape) del gradiente da la casualidad de que broadcastea de vuelta, pero los valores son incorrectos por un factor de axis_length."

La rotura¶

En tu Tensor.__add__:

 def _backward():
-    a.grad += unbroadcast(out.grad, a.shape)
-    b.grad += unbroadcast(out.grad, b.shape)
+    a.grad += out.grad
+    b.grad += out.grad

(unbroadcast es el helper que suma a lo largo de los ejes en los que el operando fue broadcasteado.)

Procedimiento de ejecución¶

Dos casos de test que ejercitan ambos regímenes:

uv run python -c "
import numpy as np
from src.minigrad.tensor import Tensor

# --- Regime 1: dimension-creation broadcast (3,) + (4, 3) ---
a = Tensor(np.array([1., 2., 3.]), requires_grad=True)       # shape (3,)
b = Tensor(np.array([[10.,20.,30.],[40.,50.,60.],[70.,80.,90.],[100.,110.,120.]]), requires_grad=True)  # (4, 3)
c = a + b                                                     # (4, 3)
loss = c.sum()
loss.backward()
print('--- Regime 1 ---')
print(f'a.grad.shape = {a.grad.shape}  expected (3,)')
print(f'a.grad       = {a.grad}        expected [4., 4., 4.]')
print(f'b.grad.shape = {b.grad.shape}  expected (4, 3)')

# --- Regime 2: length-1 axis broadcast (3, 1) + (3, 5) ---
a2 = Tensor(np.array([[1.],[2.],[3.]]), requires_grad=True)   # (3, 1)
b2 = Tensor(np.ones((3, 5)), requires_grad=True)              # (3, 5)
c2 = a2 + b2                                                  # (3, 5)
loss2 = c2.sum()
loss2.backward()
print('--- Regime 2 ---')
print(f'a2.grad.shape = {a2.grad.shape}  expected (3, 1)')
print(f'a2.grad       = {a2.grad.ravel()}  expected [5., 5., 5.]')
"

Modo de fallo esperado¶

Régimen 1 — broadcast con creación de dimensión:

Sin unbroadcast, out.grad tiene forma (shape) (4, 3) pero a.grad (la hoja) está inicializada a forma (shape) (3,). El += o bien: - Crashea con ValueError: non-broadcastable output operand with shape (3,) doesn't match the broadcast shape (4, 3) (NumPy 1.20+), o - Broadcastea silenciosamente la suma, produciendo una a.grad con forma (shape) (4, 3) que ya no coincide con a.data.shape. El a.data -= lr * a.grad del optimizador entonces crashea o se corrompe silenciosamente en el siguiente paso.

Régimen 2 — broadcast de eje de longitud 1:

La forma (shape) da la casualidad de que se alinea: out.grad es (3, 5), el almacenamiento de a.grad es (3, 1). El += broadcastea y acumula el total equivocado: a.grad[i, 0] es la suma de out.grad[i, :] por columnas — [5., 5., 5.] — pero solo se supone que cae una de las cinco contribuciones por elemento. Sin unbroadcast, también obtienes [5., 5., 5.] porque la asignación broadcastea el mismo valor. Disfrazado como correcto en este caso.

Ahora cambia la pérdida a una no uniforme:

loss2 = (c2 * Tensor(np.arange(15).reshape(3, 5).astype(np.float64))).sum()

Con la rotura en su sitio, a2.grad difiere del valor analítico. Compara contra PyTorch.

Diagnóstico¶

Solo desde los logs:

Crash en el Régimen 1. Fácil de pillar pero solo si tienes un test con un broadcast con creación de dimensión.
Print de forma (shape) después del backward. assert tensor.grad.shape == tensor.data.shape después de cada backward() es una invariante de una línea que detecta la corrupción silenciosa.
Gradcheck. Compara contra diferencias finitas sobre una pérdida no uniforme; el grad analítico discrepa del numérico para cualquier tensor que se broadcasteó en forward.
Cross-check contra PyTorch. torch.autograd.grad sobre la misma expresión da la respuesta correcta; haz diff de tu .grad frente al de PyTorch.

Lección¶

El broadcasting en forward es replicación virtual; la regla de la cadena dice que el gradiente respecto al operando replicado es la suma sobre las copias replicadas. NumPy no hace esta suma automáticamente — debes llamar a grad.sum(axis=broadcast_axes, keepdims=...) en _backward.

Esta es la nueva familia de bugs que la fase 8 introduce (la fase 7 no tenía broadcasting). El helper unbroadcast(grad, target_shape) es tan importante que se gana su propia función utilitaria testeada aisladamente. Los módulos de la fase 9 y las capas de normalización de la fase 10 lo reutilizan para cada op que admita broadcasting.

Referencias¶

El recorrido por los internals de PyTorch de Karpathy ("zero to hero" — construyendo micrograd → tensor autograd) cubre esta transición.
Goodfellow, Bengio, Courville, Deep Learning, §6.5.3 (backprop a través de grafos computacionales generales, incluyendo broadcasting) — la regla formal suma-sobre-replicados.

Break — Saltarse el unbroadcast (suma por ejes) en el backward de add¶