# 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}$$

1. 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}$$

1. 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?

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.