So I've been tinkering around with the backpropagation algorithm and to try to get a better understanding of how it works and my calculus is quite rusty. I've derived the gradient for linear regression using a MSE loss function, but have nowhere to check it against.
So if $MSE=\frac{1}{2n}\sum_{i=1}^{n}\left ( \hat{y} - y \right )^2$ and $\hat{y}=\sum_{i=1}^{n}w_i x_i$ where $\hat{y}$ is the predicted value and $y$ is the true value. Then I can compute the gradient for the weights using the chain rule like this:
$$\frac{\partial MSE}{\partial w_i}=\sum_{j}\frac{\partial MSE}{\partial \hat{y}_j} \cdot \frac{\partial \hat{y}_j}{\partial w_i}$$
- Firstly, we compute $\frac{\partial MSE}{\partial \hat{y}_i}$. Since $MSE$ is a scalar, this will result in a $n$-dimensional gradient vector where $n$ is the number of instances: $$ \begin{equation} \begin{split} \frac{\partial MSE}{\partial \hat{y}_i} &= \frac{\partial}{\partial \hat{y}_i} \frac{1}{2n}\sum_{j=1}^{n}\left [\left ( \hat{y_j} - y_j \right )^2 \right ] \\ &=\frac{1}{2n}\sum_{j=1}^{n} \left [ \frac{\partial}{\partial \hat{y}_i}\left ( \hat{y}_j - y_j \right )^2 \right ] \\ &=\frac{1}{n}\sum_{j=1}^{n} \left [ \left ( \hat{y}_j - y_j \right )\frac{\partial}{\partial \hat{y}_i}\left ( \hat{y}_j - y_j \right ) \right ] \\ &=\frac{1}{n} \left [ \left ( \hat{y}_1 - y_1 \right )\frac{\partial}{\partial \hat{y}_i}\left ( \hat{y}_1 - y_1 \right ) + \dotsi + \left ( \hat{y}_n - y_n \right )\frac{\partial}{\partial \hat{y}_i}\left ( \hat{y}_n - y_n \right ) \right ] \\ &= \frac{1}{n}\left ( \hat{y}_i - y_i \right ) \end{split} \end{equation} $$
From this equation we get the gradient $$ \frac{\partial MSE}{\partial \hat{y}} = \frac{1}{n} \begin{bmatrix} \hat{y}_1 - y_1 \\ \hat{y}_2 - y_2 \\ \vdots \\ \hat{y}_n - y_n \end{bmatrix} $$
- Now we need to compute $\frac{\partial \hat{y}_j}{\partial w_i}$: $$ \begin{equation} \begin{split} \frac{\partial \hat{y}_j}{\partial w_i} &= \frac{\partial}{\partial w_i} \sum_{j=1}^m x_{ij}w_j \\ &= \sum_{j=1}^m \frac{\partial}{\partial w_i} x_{ij}w_j \\ &= \frac{\partial}{\partial w_i} x_{i1}w_1 + \frac{\partial}{\partial w_i} x_{i2}w_2 + \dotsi + \frac{\partial}{\partial w_i} x_{im}w_m \end{split} \end{equation} $$
This forms the Jacobian: $$ J = \begin{bmatrix} x_{11} & x_{21} & \dotsi & x_{n1} \\ x_{12} & x_{22} & \dotsi & x_{n2} \\ \vdots & \vdots & \ddots & \vdots \\ x_{1m} & x_{2m} & \dotsi & x_{nm} \\ \end{bmatrix} = X^T $$
Now that we've got both the elements, we just need to put them together: $$ \begin{equation} \begin{split} \frac{\partial MSE}{\partial w_i} &=\sum_{j}\frac{\partial MSE}{\partial \hat{y}_j} \cdot \frac{\partial \hat{y}_j}{\partial w_i} \\ &= \frac{1}{n} \sum_{j=1}^n \left ( \hat{y}_j -y_j \right ) x_{ij} \end{split} \end{equation} $$ Using matrices, this is equivalent to $$ \frac{\partial MSE}{\partial w_i} = \left [ \frac{\partial \hat{y}}{\partial w_i} \right ]^T \left [ \frac{\partial MSE}{\partial \hat{y}} \right ] $$
This gives the correct answer, however I have not found a detailed derivation of this anywhere, and cannot check my work against anything.
Moreover, the Jacobian $\frac{\partial \hat{y}_j}{\partial w_i}$ turns out to be $X^T$. I've never heard this mentioned. Is this right? Can I refer to $X^T$ as the Jacobian in this context?
Also, am I correct in assuming that if this was a linear layer in a neural net, then $W$ would not be a vector, but a matrix, therefore the Jacobian would not be a two dimensional matrix, but rather a three dimensional tensor? And if I were to try to derive that on paper, how would I write this down? What is the correct notation in such cases?
