1
$\begingroup$

I am trying to calculate the derivatives of the loss w.r.t weights, bias, and input for a single-layered neural network where the loss function is mean squared error. The derivation is as follows:-

$\bullet~$ Let the weight column vector, input data, bias,and output column vector be
$\mathbf{W} \in \mathbb{R}^n$, $\mathbf{X} \in \mathbb{R}^{m\times n}$, $\mathbb{b} \in \mathbb{R}$, $\mathbf{Y} \in \mathbb{R}^m$

$\bullet~$Let $\mathbf{Z} = \mathbf{X}\times\mathbf{W}+\mathbb{b}$ be the linear transformation

$\bullet~$$\hat{\mathbf{Y}} = \max(0,\mathbf{Z})$ be the relu activation

$\bullet~$$L = (\hat{\mathbf{Y}} - \mathbf{Y})^T(\hat{\mathbf{Y}} - \mathbf{Y})/\left|\mathbf{Y}\right|$ be the mean squared error \begin{align*} &\frac{\partial L}{\partial \hat{\mathbf{Y}}} = \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T \in \mathbb{R}^{1\times m}\\ &\frac{\partial \hat{\mathbf{Y}}}{\partial \mathbf{Z}} = \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg) \in \mathbf{R}^{m\times m},~ \text{ where } \frac{\partial \hat{Y}_i}{\partial Z_i} = \begin{cases} 0 & \text{if } Z_i \leqslant 0 \\ 1 & \text{otherwise} \end{cases} \\ &\frac{\partial L}{\partial \mathbf{Z}} = \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg) \in \mathbf{R}^{1\times m} \end{align*}


$\blacksquare~$For weights: $$\frac{\partial \mathbf{Z}}{\partial \mathbf{W}} = \mathbf{X} \in \mathbf{R}^{m\times n}$$ Hence, $$\frac{\partial L}{\partial \mathbf{W}} = \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg) \times \mathbf{X} \in \mathbf{R}^{1\times n}$$ In order the make the dimensions of $\dfrac{\partial L}{\partial \mathbf{W}}$ same as $\mathbf{W}$, we need to take transpose of the above equation. This makes the RHS - $$\mathbf{X}^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg)^T\times \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y}) \in \mathbf{R}^n$$ My first question is -$\color{blue}{\text{ Is the above derivation correct? Or am I missing something?}}$


$\blacksquare~$For bias:

$$\frac{\partial \mathbf{Z}}{\partial\mathrm{b}} = \mathbf{1} \in \mathbf{R}^m$$ Hence, $$\frac{\partial L}{\partial \mathrm{b}} = \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg) \times \mathbf{1} \in \mathbf{R}^{1\times 1}$$ My second question- $\color{blue}{\text{Is the above expression correct?}}$


$\blacksquare~$For data

$$\frac{\partial \mathbf{Z}}{\partial \mathbf{X}} = \mathbf{W} \in \mathbf{R}^n$$ Hence,

$$\frac{\partial L}{\partial \mathbf{X}} = \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg) \times \mathbf{W}$$ $\color{red}{\text{Which is not correct because of the mismatch of the dimensions}}$

The only way this will work is as follows: $$\frac{\partial L}{\partial \mathbf{X}} = \mathbf{W}\times \frac{2}{\left|\mathbf{Y}\right|}(\hat{\mathbf{Y}} - \mathbf{Y})^T\times \text{diag}\bigg(\frac{\partial \hat{Y}_1}{\partial Z_1},\dots ,\frac{\partial \hat{Y}_m}{\partial Z_m}\bigg)$$ $\color{magenta}{\text{This seems to be just wrong to me.}}$ Can you please help me in understanding what is going wrong here?


Thanks!

$\endgroup$
2
$\begingroup$

You were doing pretty good until you got to $\frac{\partial{\cal L}}{\partial X}$
The problem is that $\frac{\partial z}{\partial X}$ is not a matrix but rather a 3rd order tensor!

The simplest way to avoid such tensors in matrix calculus is to use differentials.
First a bit of notation $$\eqalign{ z &= Xw + {\tt1}\beta \\ dz &= dX\,w \qquad&({\rm the\,differential\,of\,}z) \\ {\cal H}(z_k) &= \begin{cases}1\quad{\rm if}\quad z_k>0\\0\quad{\rm otherwise} \end{cases} \qquad&({\rm Heaviside\,step\,function}) \\ h &= {\cal H}(z) \qquad&({\rm apply\,the\,function\,elementwise}) \\ H &= {\rm Diag}(h) \qquad&({\rm diagonal\,\{{\tt0},\!{\tt1}\}\,matrix}) \\ A:B &= {\rm Tr}(A^TB) \qquad&({\rm Frobenius\,product}) \\ \\ }$$ The Heaviside function affords a more succinct way to write one of the earlier gradients
$$\eqalign{ \frac{\partial\hat y}{\partial z} &= H }$$ Next rewrite one of the previously calculated gradients in differential form and then perform the change of variables $z\to X$
$$\eqalign{ d{\cal L} &= \left(\frac{\partial{\cal L}}{\partial z}\right):dz \\ &= 2\|y\|^{-1}H(\hat y-y):dz \\ &= 2\|y\|^{-1}H(\hat y-y):dX\,w \\ &= 2\|y\|^{-1}H(\hat y-y)w^T:dX \\ \frac{\partial{\cal L}}{\partial X} &= 2\|y\|^{-1}H(\hat y-y)w^T \\ }$$ And now the dimensions work out perfectly (although it appears that your preferred layout convention is the transpose of this).

The key is that the differential of a matrix is just another matrix and obeys all of the rules of matrix algebra. This is simply not true for tensors.

Not only that, but it's impossible to write tensor expressions unless/until you learn index notation.


Update

This update is to clear up some questions in the comments.

Here is a list of the sizes of the various variables and products which occur in the solution $$\eqalign{ \beta &\in {\mathbb R}^{1\times 1} \\ w &\in {\mathbb R}^{n\times 1} \\ h,y,\hat y,z &\in {\mathbb R}^{m\times 1} \\ X &\in {\mathbb R}^{m\times n} \\ H &\in {\mathbb R}^{m\times m} \\ Xw,\,Hy &\in {\mathbb R}^{m\times 1} \\ Hyw^T &\in {\mathbb R}^{m\times n} \\ }$$ The properties of the trace function permit the terms in a Frobenius product $(:)$ to be rearranged in a number of equivalent ways, e.g. $$\eqalign{ &A:B = B:A = B^T:A^T \\ &A:BC = AC^T:B = C^T:BA^T = etc \\ }$$ Note that the matrix on each side of the product symbol (i.e. the colon) is exactly the same size. This is the same requirement as the Hadamard product. In fact, the Frobenius product can be defined as a Hadamard product $(\odot)$ followed by summation. $$\eqalign{ A:B &= \sum_i\sum_j (A\odot B)_{ij} \\ }$$ Finally, a gradient and a differential are two ways of conveying the same information $$\eqalign{ df = G:dX\qquad\iff\qquad G=\left(\frac{\partial f}{\partial X}\right)\\ \\ }$$

Update #2

Here are the differentials of $z$. $$\eqalign{ z &= Xw + {\tt1}\beta \\ dz &= dX\,w \quad&({\rm wrt\,}X) \\ dz &= X\,dw \quad&({\rm wrt\,}w) \\ dz &= {\tt1}\,d\beta\quad&({\rm wrt\,}\beta) \\ }$$ This post has already established that $$\eqalign{ \hat y &= \max(z,0) \\ d\hat y &= H\,dz \\ \frac{\partial\hat y}{\partial z} &= H \;=\; H^T \quad ({\rm it's\,symmetric}) \\ }$$ Let's calculate $\frac{\partial{\cal L}}{\partial\hat y}$ $$\eqalign{ {\cal L} &= \|y\|^{-1}(\hat y-y):(\hat y-y) \\ d{\cal L} &= 2\|y\|^{-1}(\hat y-y):d\hat y \\ \frac{\partial{\cal L}}{\partial\hat y} &= 2\|y\|^{-1}(\hat y-y) \\ }$$ Substituting $\,d\hat y=H dz\,$ yields $$\eqalign{ d{\cal L} &= 2\|y\|^{-1}(\hat y-y):H\,dz \\ &= 2\|y\|^{-1}H^T(\hat y-y):dz \\ &= 2\|y\|^{-1}H(\hat y-y):dz \\ \frac{\partial{\cal L}}{\partial z} &= 2\|y\|^{-1}H(\hat y-y) \\ }$$ The other gradients are obtained by substituting $dz$ with the appropriate differential, e.g. $$\eqalign{ d{\cal L} &= 2\|y\|^{-1}H(\hat y-y):dz \\ &= 2\|y\|^{-1}H(\hat y-y):{\tt1}\,d\beta \\ &= 2\|y\|^{-1}{\tt1}^TH(\hat y-y):d\beta \\ \frac{\partial{\cal L}}{\partial\beta} &= 2\|y\|^{-1}{\tt1}^TH(\hat y-y) \\ }$$
| cite | improve this answer | |
$\endgroup$
  • $\begingroup$ So, just to be clear, the other two parts are correct? $\endgroup$ – kusur Aug 6 at 6:47
  • $\begingroup$ Three questions - 1. How did $w$ become $w^T$? 2. $:$ symbol in the solution is the Frobenius product as per your notation,right? 3. The dimension of $H(.)$ in the last equation is $1\times m$ and the dimension of $w^T$ is $1\times n$. How are these two compatible? $\endgroup$ – kusur Aug 6 at 7:55
  • $\begingroup$ thanks for the clarification. There is one question still lingering though. In the equation just after $d{\cal L} = \left(\frac{\partial{\cal L}}{\partial z}\right):dz$, $\frac{\partial{\cal L}}{\partial z}$ looks like transposed term of the product given in the question. I understand it is needed to make the dimensions equal. Does this mean you took the transpose in this step deliberately or is it the consequence of something else? Rest everything is making sense except this step. Thanks again $\endgroup$ – kusur Aug 6 at 18:11
  • $\begingroup$ But the sizes that I am using are the same as yours. Every vector is a column vector. That is the reason that $\frac{\partial L}{\partial z}$ is coming the way it is in my solution. I have been using the numerator notation given in explained.ai/matrix-calculus/index.html $\endgroup$ – kusur Aug 6 at 19:24
  • $\begingroup$ Your question states explicitly that $\frac{\partial L}{\partial z}\in{\mathbb R}^{1\times m}$ which is a row vector. My convention is to use column vectors for such gradients. $\endgroup$ – greg Aug 6 at 19:32

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.