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?


1 Answer 1


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)$$


$$\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


Your Answer

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

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