The softmax Jacobian

Worked solutions

Problem 1

1. $\operatorname{softmax}(z + c\mathbf{1})_i = \dfrac{e^{z_i + c}}{\sum_k e^{z_k + c}}$.Adding $c\mathbf{1}$ adds the same $c$ to every logit: the one in the numerator and every one in the sum.
2. $= \dfrac{e^c\, e^{z_i}}{e^c \sum_k e^{z_k}}$.$e^{a + b} = e^a e^b$, and $e^c$ is the same factor in every term of the sum, so it comes out of the sum.
3. $= \dfrac{e^{z_i}}{\sum_k e^{z_k}} = s_i$.$e^c > 0$, so it cancels. This holds for every $i$ and every real $c$.
4. Take $c = -\max_k z_k$. Then every exponent $z_k - \max_k z_k$ is $\le 0$, and it is $0$ for the largest logit.Step 3 allows any $c$, so choose the one that makes the exponentials safe. Unshifted, $e^{z_k}$ overflows to infinity once $z_k$ exceeds about $709$ in 64-bit floats (about $88$ in 32-bit), and the ratio becomes $\infty/\infty$, which is NaN.
5. $\operatorname{softmax}(z + c\mathbf{1}) = \operatorname{softmax}(z)$; subtracting $\max_k z_k$ keeps every exponent $\le 0$, so nothing overflows and the largest term is exactly $1$Every shifted exponential lies in $(0, 1]$, and because one of them is exactly $1$ the denominator is at least $1$: it cannot overflow, and it cannot underflow to $0$ however negative the other logits are.

Problem 2

1. Let $\Sigma = \sum_k e^{z_k}$, so $s_i = e^{z_i}/\Sigma$.Naming the denominator separates the two places $z$ appears: $z_i$ in the numerator, and every $z_k$ in $\Sigma$.
2. $\partial \Sigma/\partial z_j = e^{z_j}$.Only the $j$-th term of $\Sigma$ depends on $z_j$, and $e^t$ is its own derivative.
3. $i = j$: $\dfrac{\partial s_i}{\partial z_i} = \dfrac{e^{z_i}\,\Sigma - e^{z_i}\, e^{z_i}}{\Sigma^2} = s_i - s_i^2$.Both the numerator and the denominator contain $z_i$, so each contributes a term; then $e^{z_i}/\Sigma = s_i$ in both.
4. $i \neq j$: $\dfrac{\partial s_i}{\partial z_j} = -\dfrac{e^{z_i}\, e^{z_j}}{\Sigma^2} = -s_i s_j$.The numerator $e^{z_i}$ does not contain $z_j$, so only the denominator contributes, through step 2.
5. $s_i - s_i^2 = s_i(1 - s_i)$ and $-s_i s_j = s_i(0 - s_j)$.Factoring out $s_i$ leaves two expressions that differ only in the $1$ or $0$ in front of $s_j$, and that $1$ or $0$ is $\delta_{ij}$.
6. $\partial s_i/\partial z_j = s_i(\delta_{ij} - s_j)$$\delta_{ij}$ is $1$ exactly in the case $i = j$. Sanity check: summing over $i$ gives $s_j - s_j \sum_i s_i = 0$, as it must, because $\sum_i s_i = 1$ for every $z$, so no change in $z_j$ can change the total.

Problem 3

1. Entry $(i,j)$ of $\partial s/\partial z$ is $\partial s_i/\partial z_j = s_i\delta_{ij} - s_i s_j$.Numerator layout puts $s_i$ on the rows and $z_j$ on the columns; the value is Problem 2, multiplied out.
2. $s_i\delta_{ij}$ is entry $(i,j)$ of $\operatorname{diag}(s)$, and $s_i s_j$ is entry $(i,j)$ of the outer product $ss^\top$.$\operatorname{diag}(s)$ has $s_i$ at $(i,i)$ and zeros elsewhere; $(ab^\top)_{ij} = a_i b_j$, here with $a = b = s$.
3. Let $J = \partial s/\partial z = \operatorname{diag}(s) - ss^\top$, $n \times n$.Steps 1 and 2 agree entry by entry.
4. $J^\top = \operatorname{diag}(s)^\top - (ss^\top)^\top = \operatorname{diag}(s) - ss^\top = J$.A diagonal matrix is its own transpose, and $(ss^\top)^\top = (s^\top)^\top s^\top = ss^\top$.
5. $J\mathbf{1} = \operatorname{diag}(s)\,\mathbf{1} - s\,(s^\top\mathbf{1}) = s - s$.$\operatorname{diag}(s)\,\mathbf{1} = s \odot \mathbf{1} = s$; associativity lets $s^\top\mathbf{1}$ be computed first, and it is $\sum_i s_i = 1$.
6. $\partial s/\partial z = \operatorname{diag}(s) - ss^\top$; symmetric because both terms are; $(\operatorname{diag}(s) - ss^\top)\mathbf{1} = s - s(s^\top\mathbf{1}) = s - s = 0$$J\mathbf{1} = 0$ is Problem 1 in derivative form: moving $z$ along $\mathbf{1}$ does not change $s$, so the directional derivative in that direction is zero.

Problem 4

1. Let $J = \operatorname{diag}(s) - ss^\top$ (Problem 3), and for any $v \in \mathbb{R}^n$ let $\bar v = \sum_i s_i v_i$. Then $v^\top J v = \sum_i s_i v_i^2 - \bar v^2$.$v^\top \operatorname{diag}(s)\, v = \sum_i s_i v_i^2$, and $v^\top ss^\top v = (s^\top v)^2 = \bar v^2$ by associativity.
2. $v^\top J v = \sum_i s_i (v_i - \bar v)^2 = \operatorname{Var}_s(v)$, the variance of the entries of $v$ when entry $i$ is drawn with probability $s_i$.Expanding the square gives $\sum_i s_i v_i^2 - 2\bar v^2 + \bar v^2$, because $\sum_i s_i v_i = \bar v$ and $\sum_i s_i = 1$.
3. $\operatorname{Var}_s(v) \ge 0$, with equality if and only if $v_i = \bar v$ for every $i$, that is, $v$ is a multiple of $\mathbf{1}$.Each term is $\ge 0$, and because every $s_i > 0$ the sum is $0$ only if every $v_i - \bar v$ is $0$. Conversely, $v = a\mathbf{1}$ gives $\bar v = a$.
4. If $Jv = 0$ then $v^\top J v = 0$, so $v \in \operatorname{span}\{\mathbf{1}\}$; and $J\mathbf{1} = 0$. So the null space is exactly $\operatorname{span}\{\mathbf{1}\}$, and $\operatorname{rank} J = n - 1$.Multiply by $v^\top$ and apply step 3; Problem 3 gives the reverse inclusion; then rank–nullity for an $n \times n$ matrix.
5. If $\operatorname{softmax}(z) = \operatorname{softmax}(z')$, then $z' = z + (\log\Sigma' - \log\Sigma)\,\mathbf{1}$, where $\Sigma = \sum_k e^{z_k}$ and $\Sigma' = \sum_k e^{z'_k}$.Taking logs, $\log s_i = z_i - \log\Sigma = z'_i - \log\Sigma'$, and the correction is the same for every $i$. Step 4 is about first-order changes; this covers finite ones.
6. $\operatorname{rank} = n - 1$; null space $\operatorname{span}\{\mathbf{1}\}$: logits are determined by the probabilities only up to a common shift, which is Problem 1 againOnly the differences $z_i - z_j = \log(s_i/s_j)$ are fixed by $s$. Because $J$ is symmetric, its column space is the orthogonal complement of its null space: every first-order change $J\,dz$ in $s$ has entries summing to zero.

Problem 5

1. Let $c$ be the index with $y_c = 1$. Then $L = -\log s_c$.Every other term of the sum has $y_i = 0$.
2. With $\Sigma = \sum_k e^{z_k}$, $\log s_c = z_c - \log\Sigma$, so $L = -z_c + \log\Sigma$.The log of a quotient is a difference, and $\log e^{z_c} = z_c$. This form never divides by $s_c$.
3. $\partial L/\partial z_j = -\delta_{jc} + e^{z_j}/\Sigma = -\delta_{jc} + s_j$.$z_c$ contains $z_j$ only when $j = c$; the derivative of $\log\Sigma$ is $(1/\Sigma)\,\partial\Sigma/\partial z_j$, and only the $j$-th term of $\Sigma$ depends on $z_j$.
4. $\delta_{jc} = y_j$.$y$ is $1$ at $c$ and $0$ elsewhere, which is the definition of $\delta_{jc}$ as a function of $j$.
5. $\nabla_z L = s - y$Entry $j$ of the gradient is $\partial L/\partial z_j = s_j - y_j$, and the result is $n \times 1$, the shape of $z$. Sanity check: its entries sum to $1 - 1 = 0$, as they must, because $L$ depends on $z$ only through $s$, which ignores a shift along $\mathbf{1}$.

Problem 6

1. With $\Sigma = \sum_k e^{z_k}$, $\partial \ell_i/\partial z_j = \delta_{ij} - e^{z_j}/\Sigma = \delta_{ij} - s_j$.$\partial z_i/\partial z_j = \delta_{ij}$, and $\partial \log\Sigma/\partial z_j = s_j$ as in Problem 5, step 3; the second term is the same for every $i$.
2. $\delta_{ij}$ is entry $(i,j)$ of $I$, and $s_j$ is entry $(i,j)$ of $\mathbf{1}s^\top$.$(\mathbf{1}s^\top)_{ij} = 1 \cdot s_j$: every row of $\mathbf{1}s^\top$ is $s^\top$, matching a term that does not depend on $i$.
3. $\partial \ell/\partial z = I - \mathbf{1}s^\top$, $n \times n$.Step 1 entry by entry. It is not symmetric: entry $(i,j)$ subtracts $s_j$ and entry $(j,i)$ subtracts $s_i$. Multiplying by $\operatorname{diag}(s)$ on the left gives $\operatorname{diag}(s) - ss^\top$, Problem 3, because $s = e^{\ell}$ elementwise has Jacobian $\operatorname{diag}(s)$ with respect to $\ell$.
4. $L = -y^\top \ell$, so $\nabla_z L = (\partial\ell/\partial z)^\top \nabla_\ell L = -(\partial\ell/\partial z)^\top y$.$L = -\sum_i y_i \ell_i$ is linear in $\ell$ with gradient $-y$; then the chain rule for gradients.
5. $(I - \mathbf{1}s^\top)^\top y = y - s\,\mathbf{1}^\top y$.$(\mathbf{1}s^\top)^\top = s\,\mathbf{1}^\top$, and associativity lets $\mathbf{1}^\top y$ be computed first.
6. $\partial \ell/\partial z = I - \mathbf{1}s^\top$; then $\nabla_z L = -(\partial\ell/\partial z)^\top y = -(y - s\,\mathbf{1}^\top y) = s - y$$\mathbf{1}^\top y = \sum_i y_i = 1$ for a one-hot $y$. This is why cross-entropy is usually computed through log-softmax: it is evaluated from $z$ directly, as $z - \log\Sigma\,\mathbf{1}$, so $\log s_i$ never has to be taken of an $s_i$ that has underflowed to $0$.

Problem 7

1. $L = \tfrac1N \sum_k L_k$ with $L_k = -\sum_i Y_{ki}\log S_{ki}$.Group the double sum by rows; $L_k$ is the loss of example $k$.
2. $L_k$ depends on $Z$ only through row $k$.The softmax is taken row by row, so row $k$ of $S$ is computed from row $k$ of $Z$ alone, and $L_k$ uses only row $k$ of $S$ and $Y$.
3. $\partial L/\partial Z_{kj} = \tfrac1N\, \partial L_k/\partial Z_{kj} = \tfrac1N (S_{kj} - Y_{kj})$.By step 2, only the $k$-th term of the mean contains $Z_{kj}$; the value is Problem 5 applied to row $k$, whose softmax is row $k$ of $S$.
4. $\nabla_Z L = \tfrac1N (S - Y)$, $N \times n$Entry $(k,j)$ of a matrix gradient is $\partial L/\partial Z_{kj}$, and step 3 fills every entry. The shape is that of $Z$.

Problem 8

1. With $\Sigma = \sum_k e^{z_k}$, $L = -\sum_i y_i (z_i - \log\Sigma) = -\sum_i y_i z_i + \big(\sum_i y_i\big)\log\Sigma$.$\log s_i = z_i - \log\Sigma$, and $\log\Sigma$ does not depend on $i$, so it leaves the sum with coefficient $\sum_i y_i$. Nothing about $y$ has been assumed yet.
2. $\partial L/\partial z_j = -y_j + \big(\sum_i y_i\big)s_j$.Only the $j$-th term of $\sum_i y_i z_i$ contains $z_j$, and $\partial\log\Sigma/\partial z_j = s_j$ (Problem 5, step 3).
3. $\nabla_z L = (\sum_i y_i)\,s - y$, which is $s - y$ when the labels sum to $1$Step 2 for every $j$. One-hot labels are the case $\sum_i y_i = 1$, so Problem 5 is a special case. If the labels sum to $a \neq 1$, $a > 0$, the gradient is $a s - y$, which is zero only at $s = y/a$: the loss pulls $s$ towards the normalised labels, and $s - y$ is not its gradient.

Problem 9

1. Let $u = z/T$, so $s = \operatorname{softmax}(u)$ and $\partial u/\partial z = \tfrac1T I$.Each $u_i = z_i/T$ depends only on $z_i$; what remains is the softmax of Problem 3.
2. $\partial s/\partial z = \big(\operatorname{diag}(s) - ss^\top\big)\,\tfrac1T I = \tfrac1T\big(\operatorname{diag}(s) - ss^\top\big)$, with $s = \operatorname{softmax}(u)$.Chain rule, outer Jacobian (Problem 3 at $u$) on the left.
3. $T \to \infty$: $s \to \mathbf{1}/n$, and every entry of $\partial s/\partial z$ is at most $1/(4T)$ in size, so the Jacobian $\to 0$.$u \to 0$ and the softmax is continuous. The bracket's entries are $s_i(1 - s_i) \le \tfrac14$ and $-s_i s_j$ with $s_i + s_j \le 1$, so $s_i s_j \le \tfrac14$.
4. $T \to 0^+$, unique largest logit $z_m$, gap $g = z_m - \max_{k \neq m} z_k > 0$: for $i \neq m$, $s_i \le e^{-g/T}$, and $1 - s_m \le (n - 1)e^{-g/T}$.The denominator of $s_i$ contains $e^{z_m/T}$, so $s_i \le e^{(z_i - z_m)/T}$; then $1 - s_m = \sum_{k \neq m} s_k$.
5. So $s \to$ the one-hot vector at $m$, a hard argmax, and every entry of $\partial s/\partial z$ is at most $(n - 1)e^{-g/T}/T$ in size, which $\to 0$.Entry $s_i(\delta_{ij} - s_j)$ is bounded by $s_i$, by $s_j$, or (if $i = j = m$) by $1 - s_m$; step 4 bounds the one whose index is not $m$.
6. $T \to 0^+$, $r \ge 2$ logits tied for the largest: each tied $s_i \to 1/r$, so the diagonal entry $\tfrac1T s_i(1 - s_i) \to \infty$.For tied $i$, $s_i = 1/\big(r + \sum_k e^{(z_k - z_i)/T}\big)$, summed over untied $k$, whose terms $\to 0$; so $s_i(1 - s_i) \to \tfrac1r(1 - \tfrac1r) > 0$.
7. $\partial s/\partial z = \tfrac1T\big(\operatorname{diag}(s) - ss^\top\big)$ with $s = \operatorname{softmax}(z/T)$; as $T \to \infty$, $s \to \mathbf{1}/n$ and the Jacobian $\to 0$; as $T \to 0^+$, when the largest logit is unique the softmax approaches a hard argmax and the Jacobian $\to 0$, while at a tie for the largest logit entries of the Jacobian grow without boundSteps 2, 3, 5 and 6.

Problem 10

1. $\partial L/\partial s_i = -y_i/s_i$, so $\nabla_s L = -y/s$, $n \times 1$.Treat $L$ as a function of $s$: only the $i$-th term contains $s_i$, and the derivative of $\log t$ is $1/t$. Every $s_i > 0$, so the division is defined.
2. $\nabla_z L = J^\top \nabla_s L$ with $J = \partial s/\partial z = \operatorname{diag}(s) - ss^\top$.The chain rule for gradients: the gradient with respect to the input is the transposed Jacobian times the gradient with respect to the output.
3. $J^\top = J$.Problem 3.
4. $\operatorname{diag}(s)\,(-y/s) = -y$.A diagonal matrix acts as an elementwise product: $s_i \cdot (-y_i/s_i) = -y_i$.
5. $-ss^\top(-y/s) = s\,\big(s^\top (y/s)\big) = s\,\big(\textstyle\sum_i y_i\big) = s\,(\mathbf{1}^\top y)$.Associativity lets the scalar $s^\top(y/s) = \sum_i s_i\, y_i/s_i$ be computed first, and each $s_i$ cancels.
6. $\nabla_z L = (\operatorname{diag}(s) - ss^\top)(-y/s) = -y + s\,(\mathbf{1}^\top y)$.Steps 4 and 5 are the two terms of $J(-y/s)$.
7. $\nabla_z L = J^\top \nabla_s L = s - y$, the same as Problem 5$\mathbf{1}^\top y = 1$ for a one-hot $y$. The $1/s_i$ in $\nabla_s L$ cancels against the $s_i$ in $J$, so the large entry $-1/s_c$ of $\nabla_s L$, at the target index $c$ when the target's probability $s_c$ is small, becomes a bounded $\nabla_z L$: every entry of $s - y$ lies in $[-1, 1]$.