Softmax backward

Background

The backward pass through softmax, done the right way. The softmax Jacobian has entries $J_{ij} = y_i(\delta_{ij} - y_j)$ — but materialising it as a $(V, V)$ matrix is a non-starter for real vocabularies (at $V=50257$ that's ~19 GB per example at fp32). Instead, the Jacobian-vector product collapses to a single $O(V)$ expression: one dot product and one elementwise multiply. This is exactly what PyTorch, JAX, and TensorFlow compute under the hood — the full matrix never exists.

Problem statement

Implement softmax_backward(grad_out, softmax_out). With $y = \text{softmax}(x)$ and upstream gradient $g =$ grad_out:

$$J_{ij} = y_i\,(\delta_{ij} - y_j), \qquad dx = y \odot \big(g - \langle y, g\rangle\big)$$

i.e. each $dx_i = y_i\big(g_i - \sum_j y_j g_j\big)$ — scale $g_i$ by how much it differs from the $y$-weighted average of $g$, then weight by $y_i$.

Input

• grad_out — 1-D np.ndarray of length V: the upstream gradient $\partial L/\partial y$.
• softmax_out — 1-D np.ndarray of length V: the saved output $y = \text{softmax}(x)$.

Output

Returns a 1-D np.ndarray of the same shape — $\partial L/\partial x$.

Examples

Example 1 — hand-checked

Input:  softmax_out = [0.2, 0.3, 0.5], grad_out = [1, 0, 0]
Output: [0.16, -0.06, -0.10]

Explanation: $\langle y, g\rangle = 0.2\cdot1 = 0.2$, so $dx = y\odot(g - 0.2) = [0.2\cdot0.8,\ 0.3\cdot(-0.2),\ 0.5\cdot(-0.2)] = [0.16, -0.06, -0.10]$. (This equals $(\operatorname{diag}(y) - yy^\top)\,g$, but computed in $O(V)$.)

Example 2 — a constant gradient gives zero

Input:  softmax_out = softmax([1,2,3]), grad_out = [0.5, 0.5, 0.5]
Output: [0, 0, 0]

Explanation: for a constant $g$, $\langle y, g\rangle = 0.5$ (the $y$'s sum to 1), so $dx = y\odot(0.5 - 0.5) = 0$. This reflects softmax's shift-invariance — adding a constant to every logit doesn't change the output, so a uniform upstream gradient produces no input gradient.

Constraints

• Use the closed form y * (grad_out - dot(y, grad_out)) — do not build the $(V, V)$ Jacobian; the cost is $O(V)$, not $O(V^2)$.
• softmax_backward consumes the softmax output $y$ (the cached forward result), not the input logits.
• A constant grad_out must give exactly 0; the analytic result must match both the explicit $J g$ and a finite-difference gradient (atol≈1e-3).

Notes

• Why it collapses. $\langle y, g\rangle$ is the $y$-weighted mean of the upstream gradient; subtracting it centres $g$ so the gradient is invariant to a constant shift — the algebraic shadow of softmax's input shift-invariance.
• Related. Pairs with softmax from scratch; when softmax feeds cross-entropy, the two backwards combine into the much simpler softmax − one-hot.