Writing Back-Propogation out in terms of Matrix Multiplication

Question

I am trying to write my own backpropogation algorithm for neural networks for a class. For any specific weight of my network I could easily take the derivative, but for computational speed I want to write the derivatives out in terms of matrix expressions that could be simply coded and executed. Let us use the following notation: let $x\in\mathbb{R}^p$ be our set of $p$ inputs to our network. Let us have $m$ hidden layers of nodes each consisting of $n$ nodes. We will write the values at the $i^{th}$ hidden layer as $z_i$, which are retieved from multipling the $i^{th}$ weight matrix $W_i$ by the previous layer a chosen function $f$ (sigmoid, RELU, etc). Finally, we consider only $1$ output node (since we could loop through the output nodes and do this for all of them, although it'd be nice to generalize).

As such we can say that our network is given as $$ z_1 = f(W_1 x) $$ $$ z_2 = f(W_2 z_1) $$ $$ \vdots $$ $$ z_n = f(W_n z_{n-1}) $$ $$ \hat{y} = f(W_{n+1} z_n) $$

Using Mean Squared Error we get that our objective function for a given point is $$ MSE = (\hat{y} - y)^2 $$

The first update matrix (which is actually a vector) was fairly easy to find the update for. I found that $$ \frac{\partial(MSE)}{\partial W_{n+1}} = 2(\hat{y}-y)f'(W_{n+1}z_n) z_n = k_n z_n $$

where $k_n=2(\hat{y}-y)f'(W_{n+1}z_n)$. From here the next layer is also special I found that if we let $$ \frac{\partial(MSE)}{\partial W_n} = k_n (W_{n+1}\circ f'(W_n z_{n-1}))\otimes z_{n-1} = k_{n-1}\otimes z_{n-1} $$

where $\circ$ represents the element wise multiplication and $\otimes$ represents the outer product. The final formula I get is the one that should (in theory) work for every layer following, however I am getting a code error on the final step that the matrix dimensions I have do not properly line up, which means I must have a problem with my formula. I find that if $$ k_i = (k_{i+1}^T W_i)\circ f'(W_i z_{i-1}) $$ then we can say that $$ \frac{\partial (MSE)}{\partial W_i} = k_i \otimes z_{i-1} $$

If anyone can help me find an error that'd be much appreicated!

Answer 1

First, note that your $W_{n+1}$ should be a vector to get 1d output. So we need to treat the cases where an index $=n+1$ anywhere differently from cases where all indices $<n+1$.

Using differentials we have $$dz_i=f'(W_iz_{i-1})\circ dW_iz_{i-1}+f'(W_iz_{i-1})\circ W_idz_{i-1}$$ Using the chain rule for differentials we have $$dMSE[z_{i+1}(z_i),dz_i]=dMSE[z_{i+1}, dz_{i+1}(z_i,dz_i)]=\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\cdot dz_{i+1}(z_i,dz_i)=$$ $$=\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\cdot f'(W_{i+1}z_i)\circ W_{i+1}dz_i=W_{i+1}^T\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\circ f'(W_{i+1}z_i)\cdot dz_i$$ $\cdot$ in the formulas is the dot product.

Note that in the formula above it is assumed that $W_{i+1}dz$ is a matrix-vector product, so works only for $i+1<n+1$. So the derivative is $$\frac{\partial MSE}{\partial z_i}=W_{i+1}^T\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\circ f'(W_{i+1}z_i)$$ Also from the chain rule and the formulas above we have $$dMSE[z_i(W_i),dW_i]=dMSE[z_i,dz_i(W_i,dW_i)]=$$ $$=W_{i+1}^T\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\circ f'(W_{i+1}z_i)\cdot dz_i(W_i,dW_i)=$$ $$=W_{i+1}^T\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\circ f'(W_{i+1}z_i)\circ f'(W_iz_{i-1})z_{i-1}^T\cdot dW_i$$ So the derivative is $$\frac{\partial MSE}{\partial W_i}=W_{i+1}^T\frac{\partial MSE(z_{i+1})}{\partial z_{i+1}}\circ f'(W_{i+1}z_i)\circ f'(W_iz_{i-1})z_{i-1}^T$$ To get the closed form formula we have $$dMSE(z_n, dz_n)=2(\hat{y}-y)f'(W_{n+1}z_n)W_{n+1}\cdot dz_n$$ $$\frac{\partial MSE(z_n)}{\partial z_n}=2(\hat{y}-y)f'(W_{n+1}z_n)W_{n+1}=k$$ So for $i<n$ we have $$\frac{\partial MSE}{\partial z_i}=W_{i+1}^T\dots W_n^Tk\circ f'(W_nz_{n-1})\circ\dots\circ f'(W_{i+1}z_i)$$ You can plug it into the formula for $\frac{\partial MSE}{\partial W_i}$ to get the closed form of the derivative.

However if you really want to implement backpropogation, this would be an incorrect approach. The correct approach is to implement backward pass for each of your module with the input of the derivative from above of your computational tree. For example, taking linear layer $l(W,h)=Wh$ and derivative from above $\frac{\partial MSE}{\partial l}=A$ you have $$dMSE(l(W,h);dW,dh)=A\cdot dl(W,h;dW,dh)=A\cdot dWh+A\cdot Wdh=$$ $$=h^TA\cdot dW+AW^T\cdot dh$$ So that on backpropagation step this layer gets $A$ and values from forward pass as an input, passes $AW^T$ further (or rather down the chain) and also gives you $\frac{\partial MSE}{\partial W}=h^TA$ for the gradient descent on $W$. The value of $h$ you know from the forward pass.