Deriving linear regression gradient with MSE 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?
 A: Using matrix calculus:
$$\text{MSE}=\frac{1}{2n}\epsilon^T\epsilon=\frac{1}{2n}(y-\hat y)^T(y-\hat y)$$
$$\hat y = Xw$$
$$\frac{\partial \text{MSE}}{\partial w}=
\frac{\partial \hat y}{\partial w}^T\frac{\partial \text{MSE}}{\partial \hat y}
$$
$$\frac{\partial \text{MSE}}{\partial \hat y}=\frac{2}{2n}(y-\hat y)$$
$$\frac{\partial \hat y}{\partial w}=X$$
$$\frac{\partial \text{MSE}}{\partial w}=
\frac{1}{n}X^T(y-\hat y)=\frac{1}{n}X^T(y-Xw)
$$

For multivatiate response, $Y$ and $\hat Y = X W$, we have instead:
$$\text{MSE}=\frac{1}{2n}\|E\|_F=\frac{1}{2n}\operatorname{tr}((Y-\hat Y)^T(Y-\hat Y))$$
$$E=Y-\hat Y$$
As you discovered, the derivatives involving matrices will lead to tensors (for example $d\hat Y /dW = \mathbb I \otimes X$). Performing matrix calculus in derivative form when involving matrices is not efficient and complicates the equations. It is much easier to use the differential form:
$$d\text{MSE}=\frac{-2}{2n}\operatorname{tr}(dE^TE)=\frac{1}{n}\operatorname{tr}(d\hat Y^TE)=\frac{1}{n}\operatorname{tr}(dWX^TE)$$
So
$$\frac {\partial \text{MSE}}{\partial W}=\frac{1}{n}X^T(Y - XW),$$
an analogous result, generalizing the previous case.

See: https://math.stackexchange.com/a/1968517
