Skip to content

English · Español

02 — Derivadas de ops tensoriales (con broadcasting reverso)

Para cada op vamos a deducir: (a) qué hace el forward, (b) qué Jacobiano local aplicar al gradiente upstream, © cómo deshacer el broadcast si lo hubo. El "deshacer el broadcast" es la idea estrella: si forward expandió (3,) a (4, 3), backward suma a lo largo del eje expandido para devolver un (3,).

El broadcasting reverso — la idea clave

El broadcasting en forward expande una forma (shape) más pequeña a una mayor repitiendo a lo largo de ciertos ejes (con trucos de stride-0 por debajo). El gradiente de la salida respecto a la entrada más pequeña es el gradiente sumado a lo largo de esos ejes expandidos.

Por qué la suma: el elemento de la entrada más pequeña contribuyó a múltiples elementos de salida (uno por "fila" del broadcast). Por la regla de la cadena, la contribución de gradiente es la suma de las contribuciones a través de cada elemento de salida al que influyó.

Concretamente: c = a + b con a.shape = (3,), b.shape = (4, 3). NumPy broadcastea a a (4, 3) (replicando a lo largo del eje 0). Forward:

c[i, j] = a[j] + b[i, j]      for i ∈ [0, 4), j ∈ [0, 3)

a[j] aparece en c[0, j], c[1, j], c[2, j], c[3, j]. Su contribución de gradiente es:

∂L/∂a[j] = ∑_i (∂L/∂c[i, j]) · (∂c[i, j]/∂a[j])
         = ∑_i (∂L/∂c[i, j]) · 1
         = (∂L/∂c).sum(axis=0)[j]

Es decir: suma el upstream a lo largo del eje 0. El resultado tiene forma (shape) (3,), coincide con a.shape. ✓

El algoritmo general: unbroadcast(grad, target_shape)

def unbroadcast(grad: np.ndarray, target_shape: tuple[int, ...]) -> np.ndarray:
    # 1. Sum away the extra leading dims.
    n_extra = grad.ndim - len(target_shape)
    for _ in range(n_extra):
        grad = grad.sum(axis=0)
    # 2. Sum away dims of size 1 in target_shape that were broadcast in.
    for i, dim in enumerate(target_shape):
        if dim == 1 and grad.shape[i] != 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad

Esa función vive en minitorch/tensor.py como un helper. El _backward de toda op elementwise la llama antes de asignar al .grad de un padre.

Ops elementwise

Add: c = a + b

  • Forward: c.data = a.data + b.data (NumPy hace broadcasting según necesite).
  • Jacobianos locales: ∂c/∂a = 1, ∂c/∂b = 1.
  • Backward: 
    a_contrib = unbroadcast(c.grad, a.data.shape)
b_contrib = unbroadcast(c.grad, b.data.shape)

Subtract: c = a - b

  • Forward: c.data = a.data - b.data.
  • Jacobianos locales: ∂c/∂a = 1, ∂c/∂b = -1.
  • Backward: 
    a_contrib = unbroadcast(c.grad, a.data.shape)
b_contrib = unbroadcast(-c.grad, b.data.shape)

Multiply: c = a * b

  • Forward: c.data = a.data * b.data.
  • Jacobianos locales: ∂c/∂a = b, ∂c/∂b = a.
  • Backward: 
    a_contrib = unbroadcast(b.data * c.grad, a.data.shape)
b_contrib = unbroadcast(a.data * c.grad, b.data.shape)

Nota: b.data * c.grad es a su vez un broadcast de NumPy — si b.data y c.grad tienen formas (shapes) diferentes (lo que puede pasar cuando el broadcasting fue implícito en forward), el resultado tiene la forma (shape) broadcasteada. unbroadcast la reduce entonces hasta a.shape. La composición es correcta.

Divide: c = a / b

  • Forward: c.data = a.data / b.data.
  • Jacobianos locales: ∂c/∂a = 1/b, ∂c/∂b = -a/b².
  • Backward: 
    a_contrib = unbroadcast(c.grad / b.data, a.data.shape)
b_contrib = unbroadcast(-a.data * c.grad / (b.data ** 2), b.data.shape)

Negate, exp, log, relu, tanh

Unarios; no hay broadcasting entre args. Pero aun así:

  • Forward: aplica la primitiva de NumPy.
  • Backward: Jacobiano local × upstream. No hace falta unbroadcast (las formas (shapes) ya coinciden).
# c = exp(a)
def _backward():
    a.grad = (a.grad or 0) + c.data * c.grad   # since exp(a) = c

# c = log(a)
def _backward():
    a.grad = (a.grad or 0) + (1 / a.data) * c.grad

# c = relu(a)
def _backward():
    mask = (a.data > 0).astype(a.data.dtype)
    a.grad = (a.grad or 0) + mask * c.grad

# c = tanh(a)
def _backward():
    a.grad = (a.grad or 0) + (1 - c.data ** 2) * c.grad

GELU: c = gelu(a)

GELU (Gaussian Error Linear Unit) es la activación usada por GPT-2, BERT, MiniGPT. Definición:

\[ \text{gelu}(x) = x \cdot \Phi(x) \]

donde Φ es la CDF de la normal estándar. Dos implementaciones habituales:

Exacta (usa erf): 

import scipy.special  # we can avoid the dep; numpy has erf via math.erf elementwise
gelu = 0.5 * x * (1 + special.erf(x / np.sqrt(2)))

Aproximación tanh (más rápida, usada en el GPT-2 original): 

gelu_approx = 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))

Derivada (forma exacta):

\[ \frac{d}{dx}\text{gelu}(x) = \Phi(x) + x \cdot \phi(x) = \Phi(x) + \frac{x}{\sqrt{2\pi}} e^{-x^2/2} \]

donde φ es la PDF de la normal estándar. Implementación:

# c = gelu(a)
def _backward():
    pdf = (1 / np.sqrt(2 * np.pi)) * np.exp(-a.data ** 2 / 2)
    cdf = 0.5 * (1 + special.erf(a.data / np.sqrt(2)))
    a.grad = (a.grad or 0) + (cdf + a.data * pdf) * c.grad

Derivada con la aproximación tanh: deriva por separado, más tediosa pero no requiere erf. Escoge una para el BLUEPRINT.

Ops de reducción

Estas reducen un tensor a lo largo de uno o más ejes. Su backward debe broadcastear el gradiente upstream de vuelta a la forma (shape) de entrada (el inverso de sumar).

Sum: c = a.sum(axis=axes, keepdims=keepdims)

  • Forward: c.data = a.data.sum(axis=axes, keepdims=keepdims).
  • Jacobiano local: cada elemento de entrada contribuye 1 a su salida sumada.
  • Backward: 
    if not keepdims:
    # Reshape upstream to add singleton dims along reduced axes.
    shape = list(a.data.shape)
    for ax in (axes if isinstance(axes, tuple) else (axes,)):
        shape[ax] = 1
    upstream_reshaped = c.grad.reshape(shape)
else:
    upstream_reshaped = c.grad
a.grad = (a.grad or 0) + np.broadcast_to(upstream_reshaped, a.data.shape)

Parte peliaguda: si axes es None (sum sobre todo), el resultado es un escalar y broadcasteas el escalar a la forma (shape) de entrada. Si axes es un int, normaliza a una tupla. Trata keepdims=True (el upstream ya tiene dims singleton; no hace falta reshape).

Mean: c = a.mean(axis=axes, keepdims=keepdims)

  • Forward: c.data = a.data.mean(axis=axes, keepdims=keepdims).
  • Jacobiano local: cada elemento de entrada contribuye 1/N a la media, donde N es la cuenta de elementos reducidos.
  • Backward: igual que sum, pero divide por N.
N = np.prod([a.data.shape[ax] for ax in axes_normalized])
# ... compute upstream_reshaped as in sum ...
a.grad = (a.grad or 0) + np.broadcast_to(upstream_reshaped / N, a.data.shape)

Ops de forma (shape)

Estas no cambian valores — solo re-visualizan el tensor. Su backward invierte la operación de forma (shape).

Reshape: c = a.reshape(new_shape)

  • Forward: c.data = a.data.reshape(new_shape). Jacobiano local: identidad (tras permutación).
  • Backward: a.grad += c.grad.reshape(a.data.shape).

Transpose: c = a.transpose(axes)

  • Forward: c.data = a.data.transpose(axes).
  • Backward: transpón por la permutación de ejes inversa. 
    inverse_axes = np.argsort(axes)
a.grad += c.grad.transpose(inverse_axes)

broadcast_to: c = broadcast_to(a, target_shape)

  • Forward: NumPy crea una vista (view) con ejes de stride-0.
  • Backward: unbroadcast(c.grad, a.data.shape). Misma maquinaria que las ops elementwise.

getitem: c = a[indices]

  • Forward: c.data = a.data[indices].
  • Backward: dispersa (scatter) c.grad de vuelta a un tensor lleno de ceros en las posiciones indexadas. 
    contribution = np.zeros_like(a.data)
np.add.at(contribution, indices, c.grad)   # handles repeated indices
a.grad += contribution

Usa np.add.at en lugar de contribution[indices] = c.grad para acumular correctamente cuando los índices se repiten.

cat / stack

cat: concatenar a lo largo de un eje existente. stack: apilar a lo largo de un eje nuevo.

Forward: primitivas de NumPy np.concatenate y np.stack.

Backward: divide el gradiente upstream a lo largo del eje cat/stack y contribuye cada trozo al padre correspondiente.

# cat([a, b, c], axis=0), shapes (Na, K), (Nb, K), (Nc, K) → (Na+Nb+Nc, K)
def _backward():
    splits = np.split(out.grad, indices_or_sections=[Na, Na+Nb], axis=0)
    a.grad = (a.grad or 0) + splits[0]
    b.grad = (b.grad or 0) + splits[1]
    c.grad = (c.grad or 0) + splits[2]

Lo que falta (diferido a 03-matmul-and-softmax-grads.md)

  • matmul — el grande. La manipulación de índices por elemento es la derivación correcta pero larga.
  • softmax — el Jacobiano es denso (N × N para una entrada de longitud N), por eso necesitamos la op combinada cross_entropy(logits, targets) para evitar materializarlo.
  • cross_entropy — combinada con softmax por estabilidad numérica y simplicidad del gradiente.

Esas tres viven en la página siguiente.

Escollos comunes

  1. Olvidar llamar a unbroadcast. El contrato de forma (shape) se viola; la siguiente op se atraganta o el paso del optimizador queda silenciosamente mal.
  2. Usar += sobre a.grad cuando es None. Reserva lazy: a.grad = a.grad + ... if a.grad is not None else .... O reserva siempre np.zeros_like en el primer acceso. La BLUEPRINT de la fase 8 escoge lazy.
  3. Modificar in-place el c.grad upstream. El _backward de un hijo puede ser llamado por múltiples caminos aguas abajo en el grafo (el caso diamante de la fase 7). Si haces c.grad *= something has cambiado el valor que ven otros caminos. Lee siempre c.grad y escribe a los padres; nunca escribas de vuelta a c.grad.
  4. Bug de forma (shape) con reduce(keepdims=False). Olvidar reañadir dims singleton antes del broadcast → desajuste de forma (shape).
  5. np.add.at vs contribution[indices] = c.grad. El último sobreescribe en índices repetidos; el primero acumula. Para el backward de getitem, usa siempre np.add.at.
  6. Desajuste de dtype entre data y grad. Si data es fp32 y calculas grad desde un 1 / x donde x es fp64, grad acaba en fp64. Los tests fallarán de forma intermitente. Castea explícitamente.

Recapitulación en un párrafo

El backward de toda op elementwise calcula un producto Jacobiano-local-por-gradiente-upstream y luego llama a unbroadcast(contribution, parent.shape) para tratar el broadcasting reverso. El backward de las ops de reducción broadcastea el gradiente upstream de vuelta vía np.broadcast_to (tras reshape para añadir dims singleton si keepdims=False). El backward de las ops de forma (shape) aplica la transformación de forma (shape) inversa. Las ops basadas en índices usan np.add.at para acumular correctamente bajo índices repetidos. Las ops de alto riesgo — matmul, softmax, cross_entropy — viven en la página siguiente.

Siguiente: 03-matmul-and-softmax-grads.md