Computing the gradient of Cross Entropy Loss The categorical cross entropy loss is expressed as:
$$L(y,t) = -\sum_{k=1}^{K}t_k\ln{y_k}$$
where $t$ is a one-hot encoded vector. $y_k$ is the softmax function defined as:
$$y_k = \frac{e^{z_k}}{\sum_{j=1}^{K}e^{z_j}}$$
I want to compute the gradient, $\nabla_z$, of the loss function with respect to the input of the output node. What I know: I understand how to compute the partial derivative of L with respect to a selected node (say, $z_k$). This yields the following expression:
$$\frac{\partial L}{\partial z_k} = y_k - t_k$$
But I am not sure how to generalize this to the entire vector, $z$. In essence, I know how to compute $\frac{\partial L}{\partial z_k}$ when $k = j$ and $k \neq j$, but I don't know how to calculate the gradient, $\nabla_z$.
 A: $\def\p#1#2{\frac{\partial #1}{\partial #2}}\def\D{{\rm Diag}}\def\o{{\tt1}}$Given
an independent vector $z$, define the variables
$$\eqalign{
p &= \exp(z)
 \quad&\implies\quad P=\D(p) 
 \quad&\implies &\quad dp&=p\odot dz &= P\,dz
\\
 &&&&\o^Tdp&=p^Tdz
\\
y &= \frac{p}{\o^Tp}
 \quad&\implies\quad Y=\D(y) 
 \quad&\implies\; &\;Ydz&=\frac{P\,dz}{\o^Tp}&=\frac{dp}{\o^Tp}
\\
}$$
Now calculate the differential of the elementwise softmax function
$$\eqalign{
dy &= \frac{(\o^Tp)dp-p(\o^Tdp)}{(\o^Tp)^2} = \Big(Y-yy^T\Big)\,dz \\
}$$
and substitute it into the differential of the loss function
$$\eqalign{
{\cal L}  &= -t^T \log(y) \\
d{\cal L} &= -t^T\Big(d\log(y)\Big) \\
 &= -t^T\Big(Y^{-1}dy\Big) \\
 &= -t^TY^{-1}\Big(Y-yy^T\Big)\,dz \\
 &= -t^T\Big(I-\o y^T\Big)\,dz \\
 &= \Big(y(t^T\o) - t\Big)^T\,dz \\
 &= \Big(y - t\Big)^T\,dz \\
\p{\cal L}{z} &= y-t \\
}$$
Taking the components of this vector-valued gradient recovers your solution.
A: Assume the basis at $z_k$ direction are denoted as $\hat {z_k}$.
In order to generalize it to the entire vector, let's first compute $\frac{\partial {z_k}}{\partial z}$.
We hope to get $$\frac{\partial {z_k}}{\partial {z}}\cdot\hat{z_j}=\left\{\begin{matrix}
1 \ \ \  if \ \ \  k=j\\
0 \ \ \  if \ \ \  k\neq j
\end{matrix}\right.$$
since $z_k$ only change if we are changing in the $\hat{z_k}$ direction.
$$\frac{\partial {z_k}}{\partial z}=\hat{z_k}$$
can satisfy our want quite well.
After getting this, we can apply the equation $$\frac{\partial f(x_1(t),x_2(t),...,x_n(t))}{\partial t}=\sum_{i=1}^n{\frac{\partial {f}}{\partial x_i}\frac{\partial {x_i}}{\partial t}}$$
of total differential to your problem.
