English · Español

03 — Las dos derivaciones de alto riesgo: matmul y softmax-CE¶

Las dos derivaciones más importantes de la fase. Matmul backward: aparece en cada capa lineal, en cada attention, en cada FFN. Softmax + cross-entropy fundidos en una op: la única forma de tener un gradiente numéricamente estable para la pérdida. Cuando puedas reproducirlas de memoria, has cruzado un hito.

Matmul backward¶

Forward¶

C = A @ B con A.shape = (M, K), B.shape = (K, N), así que C.shape = (M, N).

Elemento a elemento:

\[ C_{ij} = \sum_{k=0}^{K-1} A_{ik} B_{kj} \]

Derivación por índices¶

Queremos ∂L/∂A y ∂L/∂B, dado ∂L/∂C (el gradiente upstream, forma (shape) (M, N)).

∂L/∂A_{pq}: ¿cómo cambia L cuando movemos A_{pq}? A través de cada C_{ij} que depende de A_{pq}. Por la fórmula del forward, A_{pq} aparece en C_{pj} para todo j (con coeficiente B_{qj}). Por tanto:

\[ \frac{\partial L}{\partial A_{pq}} = \sum_{j=0}^{N-1} \frac{\partial L}{\partial C_{pj}} \cdot B_{qj} \]

Observa la estructura de índices: esto es ∂L/∂C (fila p, columnas j) producto interno con B (columna q, filas j — es decir, fila q de B.T). Así pues:

\[ \frac{\partial L}{\partial A} = \frac{\partial L}{\partial C} \cdot B^T \]

∂L/∂B_{rs}: análogamente, B_{rs} aparece en C_{is} para todo i, con coeficiente A_{ir}. Por tanto:

\[ \frac{\partial L}{\partial B_{rs}} = \sum_{i=0}^{M-1} \frac{\partial L}{\partial C_{is}} \cdot A_{ir} = \sum_{i} A_{ir} \cdot \frac{\partial L}{\partial C_{is}} \]

Esto es A.T (fila r, columnas i — es decir, columna r de A) producto interno con ∂L/∂C (columna s, filas i). Así pues:

\[ \frac{\partial L}{\partial B} = A^T \cdot \frac{\partial L}{\partial C} \]

Las dos fórmulas para memorizar¶

\[\frac{\partial L}{\partial A} = \frac{\partial L}{\partial C} \cdot B^T, \quad \frac{\partial L}{\partial B} = A^T \cdot \frac{\partial L}{\partial C}\]

En código:

# C = A @ B
def _backward():
    if A.requires_grad:
        A.grad = (A.grad or 0) + C.grad @ B.data.T
    if B.requires_grad:
        B.grad = (B.grad or 0) + A.data.T @ C.grad

Comprobación de formas (shapes): - C.grad @ B.T: (M, N) @ (N, K) → (M, K). Coincide con A.shape. ✓ - A.T @ C.grad: (K, M) @ (M, N) → (K, N). Coincide con B.shape. ✓

Dimensiones de batch¶

Para matmul batcheado, A.shape = (..., M, K), B.shape = (..., K, N), C.shape = (..., M, N). Los "..." son dims de batch que broadcastean igual que elementwise.

Backward: mismas fórmulas, pero T se convierte en "intercambiar los dos últimos ejes" (np.swapaxes(B, -1, -2)), y las dims de batch pueden necesitar unbroadcast si hubo broadcasting en las dims de batch.

En código:

def _backward():
    A_contrib = C.grad @ np.swapaxes(B.data, -1, -2)
    B_contrib = np.swapaxes(A.data, -1, -2) @ C.grad
    A.grad = (A.grad or 0) + unbroadcast(A_contrib, A.data.shape)
    B.grad = (B.grad or 0) + unbroadcast(B_contrib, B.data.shape)

Usa np.swapaxes(x, -1, -2) en lugar de x.T — esto último invierte todos los ejes en NumPy para arrays 2-D, pero tú quieres específicamente los dos últimos.

Comprobación de cordura con el ejemplo trabajado¶

Toma A = [[1, 2], [3, 4]] (2×2), B = [[5], [6]] (2×1). Entonces C = A @ B = [[17], [39]] (2×1).

Supón ∂L/∂C = [[1], [1]] (gradiente upstream — por ejemplo L = C.sum()).

∂L/∂A = ∂L/∂C @ B.T = [[1], [1]] @ [[5, 6]] = [[5, 6], [5, 6]]. Comprueba a mano: L = 17 + 39 = (A[0,0]·5 + A[0,1]·6) + (A[1,0]·5 + A[1,1]·6). Así que ∂L/∂A = [[5, 6], [5, 6]]. ✓

∂L/∂B = A.T @ ∂L/∂C = [[1, 3], [2, 4]] @ [[1], [1]] = [[4], [6]]. Comprueba a mano: L = (1+3)·B[0,0] + (2+4)·B[1,0] = 4·B[0,0] + 6·B[1,0]. Así que ∂L/∂B = [[4], [6]]. ✓

Esta es la derivación más importante de la fase 8. Memoriza las dos fórmulas. Vuelve a derivarlas desde cero si alguna vez las olvidas — la manipulación de índices es corta.

Softmax backward (en solitario) — y por qué la evitamos¶

Forward¶

Para un vector x de longitud N, softmax es:

\[ s_i = \frac{e^{x_i}}{\sum_{j} e^{x_j}} \]

La salida es una distribución de probabilidad: todos los s_i ∈ (0, 1), suman 1.

Jacobiano¶

El Jacobiano de softmax tiene entradas fuera de la diagonal:

\[ \frac{\partial s_i}{\partial x_j} = s_i (\delta_{ij} - s_j) \]

donde δ_ij = 1 si i==j si no 0.

Deriva: por la regla del cociente sobre s_i = exp(x_i) / Z con Z = Σⱼ exp(xⱼ):

\[ \frac{\partial s_i}{\partial x_j} = \frac{\partial}{\partial x_j} \left[ \frac{e^{x_i}}{Z} \right] \]

Caso i = j: ∂(e^{x_i})/∂x_i = e^{x_i}; ∂Z/∂x_i = e^{x_i}. Por la regla del cociente: $$ \frac{\partial s_i}{\partial x_i} = \frac{e^{x_i} \cdot Z - e^{x_i} \cdot e^{x_i}}{Z2} = \frac{e^{x_i}}{Z} - \left(\frac{e^{{x_i}}{Z}\right)}2 = s_i - s_i^2 = s_i(1 - s_i) $$

Caso i ≠ j: ∂(e^{x_i})/∂x_j = 0; ∂Z/∂x_j = e^{x_j}. Así pues: $$ \frac{\partial s_i}{\partial x_j} = \frac{0 \cdot Z - e^{x_i} \cdot e^{x_j}}{Z2} = -s_i s_j $$

Combinado: ∂s_i/∂x_j = s_i(δ_ij - s_j). ✓

Backward en forma matricial¶

Dado el upstream ∂L/∂s (un vector de longitud N), la contribución a ∂L/∂x es:

\[ \frac{\partial L}{\partial x_j} = \sum_i \frac{\partial L}{\partial s_i} \cdot \frac{\partial s_i}{\partial x_j} = \sum_i \frac{\partial L}{\partial s_i} \cdot s_i (\delta_{ij} - s_j) \]

\[ = s_j \left( \frac{\partial L}{\partial s_j} - \sum_i s_i \cdot \frac{\partial L}{\partial s_i} \right) \]

En código:

# s = softmax(x)
def _backward():
    # s.shape = x.shape = (..., N)
    weighted_sum = (s.data * s.grad).sum(axis=-1, keepdims=True)
    x.grad = (x.grad or 0) + s.data * (s.grad - weighted_sum)

Por qué no solemos usar softmax en solitario en entrenamiento¶

Porque el Jacobiano es denso, y el uso típico es alimentar softmax a cross-entropy, y el gradiente combinado simplifica enormemente.

Softmax en solitario sigue siendo útil (para inferencia, para pesos de attention). La implementamos. Pero proporcionamos una op cross_entropy(logits, targets) fundida separada para entrenamiento.

Cross-entropy desde logits (la op fundida)¶

Forward¶

Dados los logits x ∈ R^{B × C} (batch B, C clases) y los objetivos enteros y ∈ Z^B (índice de etiqueta por ejemplo), la pérdida cross-entropy es:

\[ L = -\frac{1}{B} \sum_{b=0}^{B-1} \log\left( \text{softmax}(x_b)_{y_b} \right) \]

La media sobre el batch es estándar; algunos frameworks devuelven sum (con reduction='sum'). Por defecto devolvemos mean.

Cómputo numéricamente estable: nunca materialices log(softmax(x)) directamente — para logits muy negativos, softmax hace underflow a 0 y log(0) = -inf. En su lugar:

\[ \log \text{softmax}(x_b)_c = x_{b,c} - \log\sum_{j} e^{x_{b,j}} = x_{b,c} - \text{logsumexp}(x_b) \]

Y logsumexp(x_b) = m + log(sum(exp(x_b - m))) donde m = max(x_b) — el truco estándar de la fase 2.

En código, forward:

# x.shape = (B, C); y.shape = (B,) int
m = x.data.max(axis=-1, keepdims=True)
log_sum_exp = m + np.log(np.exp(x.data - m).sum(axis=-1, keepdims=True))
log_softmax = x.data - log_sum_exp  # shape (B, C)
nll = -log_softmax[np.arange(B), y.data]  # shape (B,) — pick the target class's log-softmax
L_data = nll.mean()

La identidad preciosa¶

El gradiente combinado respecto a los logits:

\[ \frac{\partial L}{\partial x_{b,c}} = \frac{1}{B} \left( \text{softmax}(x_b)_c - \mathbb{1}[c = y_b] \right) \]

Es decir: (probabilidad predicha) menos (objetivo one-hot), dividido por el tamaño de batch.

Derivación¶

El Jacobiano de softmax en solitario: ∂s_i/∂x_j = s_i (δ_ij - s_j).

La derivada del logaritmo de softmax: $$ \frac{\partial \log s_i}{\partial x_j} = \frac{1}{s_i} \cdot s_i(\delta_{ij} - s_j) = \delta_{ij} - s_j $$

CE con objetivo y: L_b = -log(s_y). Por tanto: $$ \frac{\partial L_b}{\partial x_{b,j}} = -(\delta_{y,j} - s_j) = s_j - \delta_{y,j} $$

Ese es el gradiente por ejemplo. Promediando sobre el batch da (s - one_hot(y))/B.

La implementación¶

# L = cross_entropy(x, y)
def forward():
    B, C = x.data.shape
    m = x.data.max(axis=-1, keepdims=True)
    log_sum_exp = m + np.log(np.exp(x.data - m).sum(axis=-1, keepdims=True))
    log_softmax = x.data - log_sum_exp
    nll = -log_softmax[np.arange(B), y.data]
    L_data = nll.mean()
    return L_data, log_softmax  # keep log_softmax for backward

def _backward(log_softmax):
    B, C = x.data.shape
    softmax = np.exp(log_softmax)
    one_hot = np.zeros_like(softmax)
    one_hot[np.arange(B), y.data] = 1.0
    grad = (softmax - one_hot) / B
    x.grad = (x.grad or 0) + grad * L.grad   # L.grad is scalar (we called backward on the loss)

Por qué fundir¶

Estabilidad numérica. log(softmax(x)) es x - logsumexp(x). Nunca calculamos log(0). Lo mismo para el gradiente: softmax - one_hot está bien condicionado.
Economía de cómputo. Una pasada en lugar de dos.
Memoria. No hay tensor intermedio (B, C) de softmax en el grafo (se calcula en _backward a partir del log_softmax cacheado).
Claridad pedagógica. La identidad grad = softmax - one_hot es uno de los hechos más elegantes del ML. Vale la pena verlo en código.

F.cross_entropy(logits, targets) de PyTorch es exactamente esta op fundida, con la misma matemática. Nuestro minitorch.tensor.cross_entropy lo refleja.

Patrón de cross-check¶

def test_cross_entropy():
    rng = np.random.default_rng(42)
    x_data = rng.standard_normal((4, 3)).astype(np.float64)
    y_data = np.array([0, 2, 1, 0], dtype=np.int64)

    x = Tensor(x_data, requires_grad=True)
    y = Tensor(y_data)
    L = cross_entropy(x, y)
    L.backward()

    tx = torch.tensor(x_data, dtype=torch.float64, requires_grad=True)
    ty = torch.tensor(y_data, dtype=torch.long)
    tL = torch.nn.functional.cross_entropy(tx, ty)
    tL.backward()

    assert np.allclose(L.data, tL.item(), atol=1e-9)
    assert np.allclose(x.grad, tx.grad.numpy(), atol=1e-9)

Si los dos gradientes coinciden a 1e-9, la op es correcta.

Tabla resumen¶

Op	Forma (shape) del forward	Jacobiano local	Fórmula del backward
`matmul(A, B)`	`(..., M, K) @ (..., K, N) → (..., M, N)`	con forma tensorial	`dA = dC @ B.T`, `dB = A.T @ dC`
`softmax(x)`	preserva la forma (shape)	Jacobiano denso `s_i(δ_ij - s_j)`	`s · (ds - sum(s · ds))`
`cross_entropy(logits, y)`	`(B, C), (B,) → ()` (escalar)	combinado	`(softmax(logits) - one_hot(y)) / B`

Estas tres derivaciones son el contenido más importante de la fase 8. Interiorízalas.

Recapitulación en un párrafo¶

El backward de matmul es dA = dC @ B.T y dB = A.T @ dC — deriva por índices una vez, memoriza. La softmax en solitario tiene un Jacobiano denso (s_i(δ_ij - s_j)) pero rara vez se usa en entrenamiento porque la op combinada cross-entropy-desde-logits se simplifica a la preciosa identidad grad = softmax - one_hot, dividido por el tamaño de batch. La op fundida es además numéricamente estable (no hay log(0)) y PyTorch coincide con nuestra implementación a 1e-9. Estas tres derivaciones son el entregable de alto riesgo de la fase 8; todo lo demás es directo.

Siguiente: 04-gradcheck-and-property-tests.md