Matrix Chain Rule

If observations are recorded along rows (row-major), the following linear regression model holds and conventional definition of the partials:

$$\begin{array}{rl}\hat{\mathbf{Y}}& =\mathbf{XW}\\ \frac{\partial \hat{\mathbf{Y}}}{\partial \mathbf{W}}& =\frac{\partial {\hat{\mathbf{Y}}}_{ik}}{\partial {\mathbf{W}}_{jk}}\end{array}$$

For column-major the model is:

$$\begin{array}{rl}\hat{\mathbf{Y}}& =\mathbf{WX}\\ \frac{\partial \hat{\mathbf{Y}}}{\partial \mathbf{W}}& =\frac{\partial {\hat{\mathbf{Y}}}_{ki}}{\partial {\mathbf{W}}_{kj}}\end{array}$$

In either case, it will be shown that the chain rule will be:

$$\begin{array}{r}\frac{\partial \epsilon }{\partial \mathbf{W}}=\frac{\partial \hat{\mathbf{Y}}}{\partial \mathbf{W}}{\left(\frac{\partial \epsilon }{\partial \hat{\mathbf{Y}}}\right)}^{\mathrm{T}}\end{array}$$

Notation

The notation, $\left[j==l\right]$ is adapted from Knuth ¹ and defined as:

$${\delta }_{jl}={\mathbf{I}}_{jl}=\left[j==l\right]=\{\begin{array}{l}1,j=l\\ 0,j\ne l\end{array}$$

Solving the Product Derivatives

$$\begin{array}{rl}\mathbf{Y}& =\mathbf{W}\mathbf{X}\\ ⟹\frac{\partial {\mathbf{Y}}_{kl}}{\partial {\mathbf{W}}_{ij}}& =\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left(\sum _{p=1}^{\mathrm{ncol}\left(\mathbf{X}\right)}\left[{\mathbf{W}}_{kp}{\mathbf{X}}_{pl}\right]\right)\\ & =\sum _{p=1}^{\mathrm{ncol}\left(\mathbf{X}\right)}\left[\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{W}}_{kp}{\mathbf{X}}_{pl}\right)\right]\\ & =0+\cdots \frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{W}}_{kj}{\mathbf{X}}_{jl}\right)+\cdots 0\\ & =\{\begin{array}{ll}\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{W}}_{ij}{\mathbf{X}}_{jl}\right)& k=i\\ 0& k\ne i\end{array}\\ & =\left[k==i\right]{\mathbf{X}}_{jl}\\ & ={\delta }_{ki}{\mathbf{X}}_{jl}\end{array}$$

$$\begin{array}{rl}\mathbf{Y}& =\mathbf{X}\mathbf{W}\\ ⟹\frac{\partial {\mathbf{Y}}_{kl}}{\partial {\mathbf{W}}_{ij}}& =\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left(\sum _{p=1}^{\mathrm{ncol}\left(\mathbf{X}\right)}\left[{\mathbf{X}}_{kp}{\mathbf{W}}_{pl}\right]\right)\\ & =\sum _{p=1}^{\mathrm{ncol}\left(\mathbf{X}\right)}\left[\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{X}}_{kp}{\mathbf{W}}_{pl}\right)\right]\\ & =0+\cdots +\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{X}}_{ki}{\mathbf{W}}_{il}\right)+0\cdots \\ & =\{\begin{array}{ll}\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\mathbf{X}}_{ki}{\mathbf{W}}_{ij}\right)& j=l\\ 0& j\ne l\end{array}\\ & =\left[j==l\right]{\mathbf{X}}_{ki}\\ & ={\delta }_{jl}{\mathbf{X}}_{ki}\end{array}$$

So we have:

$$\begin{array}{rl}\frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\left[\mathbf{W}\mathbf{X}\right]}_{\left[k,l\right]}\right)& ={\mathbf{I}}_{jl}{\mathbf{X}}_{ki}\\ \frac{\partial }{\partial {\mathbf{W}}_{ij}}\left({\left[\mathbf{X}\mathbf{W}\right]}_{\left[k,l\right]}\right)& ={\mathbf{X}}_{jl}{\mathbf{I}}_{ki}\end{array}$$

Solving the Chain Rule

For the column-major example $\mathbf{Y}=\mathbf{WX}$ :

$$\begin{array}{rl}\frac{\partial \epsilon }{\partial {\mathbf{W}}_{ij}}& =\sum _q\left[\sum _p\left[\frac{\partial \epsilon }{\partial {\mathbf{Y}}_{pq}}\frac{\partial {\mathbf{Y}}_{pq}}{\partial {\mathbf{W}}_{ij}}\right]\right]\\ & \sum _q\left[\sum _p\left[\frac{\partial \epsilon }{\partial {\mathbf{Y}}_{pq}}\frac{\partial {\mathbf{Y}}_{pq}}{\partial {\mathbf{W}}_{ij}}{\mathbf{I}}_{pi}\right]\right]\\ & \text{When }p\ne i\text{ We have zero so:}\\ & =\sum _q\left[\frac{\partial \epsilon }{\partial {\mathbf{Y}}_{iq}}\frac{\partial {\mathbf{Y}}_{iq}}{\partial {\mathbf{W}}_{ij}}\right]\\ & \text{From Matrix Identities we have:}\\ & =\sum _q\left[\frac{\partial \epsilon }{\partial {\mathbf{Y}}_{iq}}{\mathbf{X}}_{jq}\right]\\ & \text{Transpose}\mathbf{X}:\\ & =\sum _q\left[\frac{\partial \epsilon }{\partial {\hat{\mathbf{Y}}}_{iq}}{{\mathbf{X}}^{\mathrm{T}}}_{qj}\right]\\ & \text{By definition we have:}\\ & =\frac{\partial \epsilon }{\partial \hat{\mathbf{Y}}}{\mathbf{X}}^{\mathrm{T}}\\ ⟹\frac{\partial \epsilon }{\partial \mathbf{W}}& =\frac{\partial \epsilon }{\partial \hat{\mathbf{Y}}}{\mathbf{X}}^{\mathrm{T}}\end{array}$$

This of course assumes column-major where the columns of $\mathbf{Y}$ represent observations and the rows are features, in the event that these were transposed we would have

$$\begin{array}{rl}{\mathbf{Y}}_{\mathrm{r}}& ={\mathbf{Y}}^{\mathrm{T}}\\ {\mathbf{X}}_{\mathrm{r}}& ={\mathbf{X}}^{\mathrm{T}}\\ {\mathbf{W}}_{\mathrm{r}}& ={\mathbf{W}}^{\mathrm{T}}\end{array}$$

and hence:

$$\begin{array}{rl}{\left(\frac{\partial \epsilon }{\partial \mathbf{W}}\right)}^{\mathrm{T}}& ={\left(\frac{\partial \epsilon }{\partial \hat{\mathbf{Y}}}{\mathbf{X}}^{\mathrm{T}}\right)}^{\mathrm{T}}\\ & =\mathbf{X}{\frac{\partial \epsilon }{\partial \hat{\mathbf{Y}}}}^{\mathrm{T}}\\ & \text{Swap for Row-Major}\\ \frac{\partial \epsilon }{\partial {\mathbf{W}}_{\mathrm{r}}}& ={\mathbf{X}}_{\mathrm{r}}^{\mathrm{T}}\frac{\partial \epsilon }{\partial {\hat{\mathbf{Y}}}_{\mathrm{r}}}\end{array}$$