# Neural Network matrix calculus

Just to preface, I have my layer activations as row vectors so I write $$x_{i}W_{i}$$.

Consider a very simple network, with an input layer with 1 node, 1 hidden layer with 2 nodes, and an output layer with 2 nodes. Let $$W_{1}$$ represent the weights connecting the hidden layer to the output layer, $$x_{1}$$ represent the activation of the hidden layer and $$x_{2}$$ is the activation of the output layer. Let the activation function be $$f$$.

$$$$\frac{\partial E}{\partial W_{1}}=\frac{\partial E}{\partial x_{2}}\frac{\partial x_{2}}{\partial W_{1}}=\frac{\partial E}{\partial x_{2}}\frac{\partial f(x_{1}W_{1})}{\partial W_{1}}$$$$

There are 3 terms to consider: $$\frac{\partial E}{\partial x_{2}}$$, $$\frac{\partial x_{1}W_{1}}{\partial W_{1}}$$ and $$f'(x_{1}W_{1})$$. $$\frac{\partial E}{\partial x_{2}}$$ and $$f'(x_{1}W_{1})$$ are both $$1\times 2$$ matrices, and the derivative of a vector such as $$x_{1}W_{1}$$ with respect to a matrix like $$W_{1}$$ is not well agreed upon (according to Wikipedia). We know that the derivative of the loss function wrt $$W_{1}$$ has to be in the shape of $$W_{1}$$ which is $$2\times 2$$, but this shape is not possible with 2 $$1\times 2$$ matrices regardless of the shape of the last term. What happened?

There's indeed some confusion in the literature on which shape to use for derivative with respect to a matrix. This book helps makes sense of this -- "Matrix Differential Calculus with Applications in Statistics and Econometrics" by Jan R. Magnus, in particular, section 9.3 "Bad notation" and section 9.4 "Good notation"

The key point is that there are multiple ways to define derivatives when you start dealing with higher-rank inputs/outputs, with no single one that's significantly better.

There's one special case with an obvious "best" notation is in the case of derivative of scalar output with respect to vector or matrix input A. In such case it's best keep derivative in the same shape as A. The reason is that this lets you reuse standard formula for gradient descent:

$$W\leftarrow W-\alpha f'(W)$$

Personally, I think derivatives of vectors/matrices w.r.t matrix are best kept as tensors, otherwise you have to deal with flattening 4d tensor into 2d matrix, there is more than one way of doing it, and you have to think about the correct way each time. Magnus/Nuedecker advocate for a particular way of doing this, relying on vec operation to convert matrix to a vector. This operation is efficient for column-major memory layout, while modern deep learning systems like PyTorch default to row-major memory layout. Using vec operation in such system would give an inefficient representation.

To avoid higher rank tensors, or flattening, it's best to use notational technique called differentials. You can think of a differential as a linear function that approximates your target function. So it's very closely related to derivative, except derivative is a tensor object, whereas differential is a function.

For instance for function $$Ax+b$$, $$f(x)=Ax$$ is the differential, whereas $$A$$ is the derivative. For practical purposes, it doesn't matter if your final result is $$A$$ or $$Ax$$ because you can convert between the two forms.

The reason to keep it as differential is because it lets you do algebra with objects of lower rank. For instance consider the following matrix function of matrix $$x$$

$$f(x)_{ij}=\sum_{ijkl} A_{ijkl} x_{kl}$$

Its differential of this function, $$f(x)$$, has matrix shape, whereas its derivative, $$A$$, is a rank-4 tensor.

Differentials have an analogous chain rule to derivatives, and you can extract derivative out of final differential easily.

Let's get a sense of how your derivation would go in the world of differentials. We are interested in computing derivative of $$f(g(W)$$ where $$f$$ is our loss and $$g$$ is our predictor. Define functions $$f,g$$ and their differentials $$df,dg$$

$$\begin{array}{lll} f(y)&=&y^T y\\ df(y)&=&2y^T dy\\ g(W)&=&Wx\\ dg(W)&=&dW x \end{array}$$

Now compute differential of $$f(g(W))$$ using chain rule for differentials

$$\begin{array}{lll} d[f(g)]&=&df \cdot dg=\\ & = & 2y^T dW x \end{array}$$

Now we can see that differential is $$2y^T dW x$$, we need to extract derivative out of this. Standard trick is to write it as trace, use cyclic property of trace to turn it into form $$\text{tr}(A'dX)$$ and then extract derivative from this form.

$$\begin{array}{lll} 2 y^T dW x&=&\text{tr}(y^T dW x) = \\ & = & \text{tr}(x y^T dW) = \\ & = & \text{tr}(A' dW) \end{array}$$

In the last step, we defined $$A=yx^T$$. Now, once you bring your differential into form $$\text{tr}(A' dW)$$, the derivative is just $$A$$ and it has the same shape as $$W$$. So now we get that

$$E'=y x^T$$

To see the power of differentials, consider the ease with which various matrix derivatives are derived in "Collected matrix derivative results for forward and reverse mode AD" paper. Their dot notation $$\dot{A}$$ is equivalent to differential $$\mathrm{d}A$$. Derivatives of matrix inverse and determinant can be obtained in matrix form this way.

Below is a cheat-sheet of some standard matrix differential calculus results and exercises to get comfortable with notation, taken from "Matrix Differential Calculus with Applications in Statistics and Econometrics" (Magnus, Neudecker)

Let $$|F|$$ denote the determinant of $$F$$

• $+\tt1\;$ But you have a typo in the line $$df(W)=dW\,x$$ It should read $$dg(W)=dW\,x$$
– greg
Dec 24, 2021 at 15:59
• ----------fixed Dec 24, 2021 at 16:19

The trace has the advantage of being a familiar operation, however it is a clumsy and confusing notation for doing Matrix Calculus. $$\,$$ Inner product notations have far greater pedagogical value.

For example, given a matrix variable, function and gradient \eqalign{ \phi &= f(X),\qquad G = \frac{\partial \phi}{\partial X} \\ } there are a variety of ways to write the differential
\eqalign{ d\phi &= G_{ij}\,dX_{ij} \\ &= G:dX \\ &= G\bullet dX \\ &= \big\langle G,\,dX\big\rangle \\ &= {\rm trace}\big(G^TdX\big) \\ &= {\rm trace}\big(G\,dX^T\big) \\ } But take a closer look at those trace expressions.

1. Why is there a function invocation?
2. Why is there a transpose?
3. Does the transpose operation go on the first or second term?

The cognitive benefits of using infix operators versus functional notation are obvious.
Indeed, it's the reason we use algebraic notation like this $$A + BC$$ instead of this $${\rm sum}\Big(A,\,{\rm prod}(B,C)\Big) \\$$

• BTW, going to Hessians and higher derivatives, makes matrix notation completely untenable, however Ricci calculus notation still works. After playing around with it, I actually prefer Ricci calculus for regular matrix derivatives too, it obviates the "Trace" trick, or remembering Kronecker products rules, and has a nice graphical representation as tensor networks -- arxiv.org/abs/1708.00006 Dec 24, 2021 at 16:24
• @YaroslavBulatov Looks like an interesting paper, I'll have to read it over the holiday break. I usually use standard index notation for higher-order gradient calculations, but I always use the differential approach.
– greg
Dec 24, 2021 at 16:53
• This one by same author goes deeper into this. Equations like vec(ABC) in terms of kronecker products follow naturally in tensor network/ricci calculus notation arxiv.org/abs/1912.10049 Dec 24, 2021 at 17:18