Vectorization of Cross Entropy Loss I am dealing with a problem related to finding the gradient of the Cross entropy loss function w.r.t. the parameter $\theta$ where:
$CE(\theta) = -\sum\nolimits_{i}{y_i*log({\hat{y}_{i}})}$
Where, $\hat{y}_{i} = softmax(\theta_i)$ and $\theta_i$ is a vector input.
Also, $y$ is a one hot vector of the correct class and $\hat{y}$ is the prediction for each class using softmax function. 
Hence, for example lets have $y_i = \begin{pmatrix}0\\0\\0\\1\\0\end{pmatrix}$
and $\hat{y}_{i} = \begin{pmatrix}0.10\\0.20\\0.10\\0.40\\0.20\end{pmatrix}$
To find the partial derivative $\frac{\partial{CE(\theta)}}{\partial{\theta{ik}}} = -{y_{ik} - \hat{y}_{ik}}$
Taking from there for each $i$ the individual partial gradients will be 
$\frac{\partial{CE(\theta)}}{\partial{\theta{i}}} = \begin{pmatrix}y_{i1} - \hat{y}_{i1}\\y_{i2} - \hat{y}_{i2}\\y_{i3} - \hat{y}_{i3}\\y_{i4} - \hat{y}_{i4}\\y_{i5} - \hat{y}_{i5}\end{pmatrix}$
But this is not true because the gradients should actually be 0 for all other rows except for the 4th row because we have used the property of the one hot vector. So actual gradient should be
$\frac{\partial{CE(\theta)}}{\partial{\theta{i}}} = \begin{pmatrix}0\\0\\0\\y_{i4} - \hat{y}_{i4}\\0\end{pmatrix}$
And hence the gradients for all $i$ should be
$\frac{\partial{CE(\theta)}}{\partial{\theta}} = \left( \begin{array}{ccc}
0 & 0 & 0 & y_{i4} - \hat{y}_{i4} & 0 \\
0 & 0 & y_{i3} - \hat{y}_{i3} & 0 & 0 \\
... \\
0 & y_{i2} - \hat{y}_{i2} & 0 & 0 & 0 \end{array} \right)$
But this is not equal to $\hat{y} - y$. So we should not call the gradient of the cross entropy function a vector difference between predicted and original. 
Can some one clarify on this ?
UPDATE: Fixed my derivation
$\theta = \left( \begin{array}{c}
\theta_{1} \\
\theta_{2} \\
\theta_{3} \\
\theta_{4} \\
\theta_{5} \\
\end{array} \right)$
$CE(\theta) = -\sum\nolimits_{i}{y_i*log({\hat{y}_{i}})}$
Where, $\hat{y}_{i} = softmax(\theta_i)$ and $\theta_i$ is a vector input.
Also, $y$ is a one hot vector of the correct class and $\hat{y}$ is the prediction for each class using softmax function. 
$\frac{\partial{CE(\theta)}}{\partial{\theta{i}}} = - (log(\hat{y}_{k}))$
UPDATE: Removed the index from $y$ and $\hat{y}$
Hence, for example lets have $y = \begin{pmatrix}0\\0\\0\\1\\0\end{pmatrix}$
and $\hat{y} = \begin{pmatrix}0.10\\0.20\\0.10\\0.40\\0.20\end{pmatrix}$
UPDATE: Fixed I was taking derivative w.r.t. $\theta_{ik}$ it should be only w.r.t. $\theta_{i}$.
To find the partial derivative $\frac{\partial{CE(\theta)}}{\partial{\theta{i}}} = -{y_{k} - \hat{y}_{k}}$
Taking from there for each $i$ the individual partial gradients will be 
$\frac{\partial{CE(\theta)}}{\partial{\theta}} = \begin{pmatrix}y_{1} - \hat{y}_{1}\\y_{2} - \hat{y}_{2}\\y_{3} - \hat{y}_{3}\\y_{4} - \hat{y}_{4}\\y_{5} - \hat{y}_{5}\end{pmatrix}$
The above happens because $CE(\theta) = -(y_k*log({\hat{y}_{k}}))$
And,  $\hat{y}_{k} = log(softmax(\theta_k)) = \theta_k - log(\sum\nolimits_{j}{exp(\theta_j)})$
Taking the partial derivative of $CE(\theta)$ w.r.t. $\theta_i$ we get:
$\frac{\partial{CE(\theta)}}{\partial{\theta{i}}} = - (\frac{\partial{\theta_k}}{\partial{\theta{i}}} - softmax(\theta_i))$
MAIN STEP:
The fact that $\frac{\partial{\theta_k}}{\partial{\theta{i}}} = 0, i \neq k$ and $\frac{\partial{\theta_k}}{\partial{\theta{i}}} = 1, i = k$ makes the vector $\frac{\partial{CE(\theta)}}{\partial{\theta}} = \hat{y} - y$ which completes the proof.
 A: No, the gradients should not be zero for the other components.  If your prediction is $\hat y_{ij}$ for some $i,j$ and your observation $y_{ij}=0$, then you predicted too much by $\hat y_{ij}$.
A: The following is the same content as the edit, but in (for me) slightly clearer step-by-step format: 
We are trying to proof that:
$\frac{\partial{CE}}{\partial{\theta}} = \hat{y} - y$
given 
$CE(\theta) = -\sum\nolimits_{i}{y_i*log({\hat{y}_{i}})}$
and
$\hat{y}_{i} = \frac{exp(\theta_i)}{\sum\nolimits_{j}{exp(\theta_j)}}$
We know that $y_{j} = 0$ for $j \neq k$ and $y_k = 1$, so:
$CE(\theta) = -\ log({\hat{y}_{k}})$
$= - \ log(\frac{exp(\theta_k)}{\sum\nolimits_{j}{exp(\theta_j)}})$
$ = - \ \theta_k + log(\sum\nolimits_{j}{exp(\theta_j)}) $
$\frac{\partial{CE}}{\partial{\theta}} = - \frac{\partial{\theta_k}}{\partial{\theta}}  + \frac{\partial}{\partial{\theta}} log(\sum\nolimits_{j}{exp(\theta_j))}$
Use the fact that $ \frac{\partial{\theta_k}}{\partial{\theta_k}} = 1 $
 and $ \frac{\partial{\theta_k}}{\partial{\theta_q}} = 0 $ for $q \neq k$, to show that. 
$ \frac{\partial{\theta_k}}{\partial{\theta}} = y $
For the second part we write out the derivative for each individual element of $\theta$ and use the chain rule to get:
$\frac{\partial}{\partial{\theta_i}} log(\sum\nolimits_{j}{exp(\theta_j))} = \frac{exp(\theta_i)}{\sum\nolimits_{j}{exp(\theta_j)}} = \hat{y}_{i}$
Hence,
$\frac{\partial{CE}}{\partial{\theta}} = \frac{\partial}{\partial{\theta}} log(\sum\nolimits_{j}{exp(\theta_j))} - \frac{\partial{\theta_k}}{\partial{\theta}}  = \hat{y}$ - y
