Given a neural network with one hidden sigmoid layer and softmax output layer, I want to derive the gradient of the cross entropy loss with respect to the first weight matrix. This is equivalent to the network described in this question, but in my derivation the dimensions mismatch.

\begin{align*} &J = CE(y,\hat{y})=-\sum_i y_i log(\hat{y}_i)\\ &\hat{y} = softmax(z_2)\\ &z_2 = hW_2+b_2\\ &h = sigmoid(z_1)\\ &z_1 = xW_1+b_1 \end{align*} Where $W_1 \in \mathbb{R}^{D_x \times H}$, $b_1 \in \mathbb{R}^H$, $W_2 \in \mathbb{R}^{H \times D_y}$ and $b_2 \in \mathbb{R}^{D_y}$.

When deriving the gradients with respect to $W_1$, I use the chain rule similar to this question, in which the gradient with respect to the input is calculated. The result should be: $$ \frac{\partial J}{\partial \boldsymbol{W_1}} = \frac{\partial J}{\partial \boldsymbol{z_2}} \cdot \frac{\partial \boldsymbol{z_2}}{\partial \boldsymbol{h}} \cdot \frac{\partial \boldsymbol{h}}{\partial \boldsymbol{z_1}} \cdot \frac{\partial \boldsymbol{z_1}}{\partial \boldsymbol{W_1}} $$ The individual gradients for the first three derivations are:
$$ \frac{\partial J}{\partial \boldsymbol{z_2}} = \left( \hat{\boldsymbol{y}} - \boldsymbol{y} \right) $$ $$ \frac{\partial \boldsymbol{z_2}}{\partial \boldsymbol{h}} = \frac{\partial}{\partial \boldsymbol{h}} \left[ \boldsymbol{h}W_2 + \boldsymbol{b_2}\right] = W_2^T $$ $$ \frac{\partial \boldsymbol{h}}{\partial \boldsymbol{z_1}} = h \circ \left(1-h\right) $$ Plugging them in into the original equation yields: $$ \frac{\partial J}{\partial \boldsymbol{W_1}} = \left( \hat{\boldsymbol{y}} - \boldsymbol{y} \right) \cdot W_2^T \circ h \circ (1 -h) \cdot \frac{\partial \boldsymbol{z_1}}{\partial \boldsymbol{W_1}} $$ Since in my understanding $ \hat{\boldsymbol{y}} - \boldsymbol{y} \in \mathbb{R}^{1 \times D_y}$ and $W_2^T \in \mathbb{R}^{D_y \times H}$ and the multiplication with $h$ and $1-h$ is elementwise the result should be $\in \mathbb{R}^{1 \times H}$. No matter what the result of $\frac{\partial \boldsymbol{z_1}}{\partial \boldsymbol{W_1}}$ is, I can not see it to fit the dimensions, since the final result has to be in the dimensions of $W_1$. Where did I go wrong in my derivation.



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.