Basic matrix maths for neural networks

I thought it would be worth writing a post on how matrix multiplication is used to calculate the output of neural networks. We use matrices because they make the maths easier, and because GPUs can work with them efficiently, allowing us to do a whole bunch of calculations with a single step -- so it's really worth having a solid grounding in what the underlying operations are.

If you're an experienced ML practitioner, you should skip this post. But you might find it useful if you're a beginner -- or if, like me until I started working through this, you've coded neural networks and used matrix operations for them, but apart from working through an example or two by hand, you've never thought through the details.

In terms of maths, I'll assume that you know what a vector is, what a matrix is, and have some vague memories of matrix multiplication from your schooldays, but that's it -- everything else I will define.

In terms of neural networks, I'll assume that you are aware of their basic layout and how they work in a general sense -- but there will be diagrams for clarity and I'll define specific terms.

So, with expectations set, let's go!

Basic matrix maths

Let's start with some definitions. The dot product is an operation that works on two vectors of the same length. It simply means that you multiply the corresponding elements, then add up the results of those multiplications:

$$\left(\begin{array}{c}a\\ b\\ c\end{array}\right)·\left(\begin{array}{c}d\\ e\\ f\end{array}\right)=a·d+b·e+c·f$$

Or, more concretely:

$$\left(\begin{array}{c}2\\ 3\\ 4\end{array}\right)·\left(\begin{array}{c}1\\ 0\\ 2\end{array}\right)=2\times 1+3\times 0+4\times 2=2+0+8=10$$

Simple enough. Now let's define matrix multiplication. If you have a matrix $M$ which has $m$ rows and $n$ columns -- that is, it's an $m\times n$ matrix -- you can multiply it by any other matrix $N$, so long as $N$ has $n$ rows -- that is, it's $n\times x$, where $x$ can be anything. To put it another way, the number of columns in the first matrix in a matrix multiplication has to equal the number of rows in the second one.¹

The result of the matrix multiplication is a new matrix, with $m$ rows and $x$ columns -- that is, it has the same number of rows as the first one, and the same number of columns as the second one.

How do we fill in the values in that result matrix? Let's start with a formal definition, and then move on to an example. Formally, the value of $M_{i,j}$ -- that is, the element at row $i$, column $j$ in the output matrix -- is the dot product of row $i$ in the first matrix, taken as a vector, with column $j$ in the second matrix, also considered as a vector. ²

This example should help make it a bit clearer:

$$\left[\begin{array}{cc}a& b\\ c& d\\ e& f\end{array}\right]\left[\begin{array}{ccc}p& s& u\\ q& t& v\end{array}\right]=\left[\begin{array}{ccc}a·p+b·q& a·s+b·t& a·u+b·v\\ c·p+d·q& c·s+d·t& c·u+d·v\\ e·p+f·q& e·s+f·t& e·u+f·v\end{array}\right]$$

Looking at the element at $2,2$ in the result -- second row, second column, right in the middle -- you can see that it's the dot product of second row in the first matrix and the second column in the second:

$$\left(\begin{array}{c}c\\ d\end{array}\right)·\left(\begin{array}{c}s\\ t\end{array}\right)=c·s+d·t$$

There are lots of different ways of visualising this process -- I rather like this animation. But the formal definition using the dot product is the important one, at least for our purposes in this post.

Now, let's also note one extra thing that will become important later, a minimal case of matrix multiplication: if we take a $1\times n$ matrix and a $n\times 1$ matrix, and multiply them, we get a $1\times 1$ matrix that contains the dot product of the two original matrices regarded as vectors:

$$\left[\begin{array}{ccc}a& b& c\end{array}\right]\left[\begin{array}{c}d\\ e\\ f\end{array}\right]=\left[\begin{array}{c}a·d+b·e+c·f\end{array}\right]$$

Hopefully that's pretty obvious as a consequence of the definitions above. So on that note, let's move on to neural networks, and how to use these operations to calculate their results using matrices.

Neural networks

Here's a simple one-layer neural network:

On the left, we have three inputs, $x_1$, $x_2$, and $x_3$. They feed into two neurons, $n_1$ and $n_2$; each input goes to both neurons. On their way to the neurons, the inputs are multiplied by weights. The weight in question is written next to the connecting arrow -- for example, on its way to neuron $n_1$, input $x_2$ is multiplied by the weight $w_{1,2}$. You can see that each weight has a subscript which is the neuron number, then the input it relates to.

Each of the neurons adds together all of the weighted input values that it has received, adds on a bias ($b_1$ for $n_1$, $b_2$ for $n_2$), runs the resulting number through an activation function, and that provides the output for that neuron -- $y_1$ for $n_1$, $y_2$ for $n_2$.

For this post, I'm going to disregard the activation function, because that doesn't have any effect on what we use matrices for -- it's just a simple function taking one number and returning another, and is applied to the output value after all of the matrix-y stuff. I'm also going to put aside the bias for now, though we'll come back to that later.

Matrix maths for one neuron

So let's take a look at what is happening at neuron $n_1$. It's receiving the inputs, each one multiplied by its associated weight, and adding those numbers up -- this is sometimes called the weighted sum or the net input for that neuron. Now, multiplying things together and adding them up is a dot product! So let's write it out like that:

$$\left(\begin{array}{c}{w}_{1,1}\\ w_{1,2}\\ w_{1,3}\end{array}\right)·\left(\begin{array}{c}{x}_1\\ x_2\\ x_3\end{array}\right)=w_{1,1}·x_1+w_{1,2}·x_2+w_{1,3}·x_3$$

Let's convert that to one of those minimal matrix multiplications that we did above:

$$\left[\begin{array}{ccc}{w}_{1,1}& w_{1,2}& w_{1,3}\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\\ x_3\end{array}\right]=\left[\begin{array}{c}{w}_{1,1}·x_1+w_{1,2}·x_2+w_{1,3}·x_3\end{array}\right]$$

So we can see that this simple matrix calculation has worked out our weighted sum for $n_1$.

Multiple neurons

Now that we have it in matrix form, we can expand it. Remember that an $m\times n$ matrix multiplied by an $n\times x$ one produces an $m\times x$ one, where each element is the dot product of the matching row in the first input matrix taken as a vector and the matching column in the second, likewise as a vector. So if we add a second row to the first matrix in the above, we'll get a second row in our result:

$$\left[\begin{array}{ccc}{w}_{1,1}& w_{1,2}& w_{1,3}\\ w_{2,1}& w_{2,2}& w_{2,3}\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\\ x_3\end{array}\right]=\left[\begin{array}{c}{w}_{1,1}·x_1+w_{1,2}·x_2+w_{1,3}·x_3\\ w_{2,1}·x_1+w_{2,2}·x_2+w_{2,3}·x_3\end{array}\right]$$

That's pretty cool! By adding a second row to the first matrix containing the weights for the second neuron, we've been able to calculate the weighted sums for both $n_1$ and $n_2$ in a single matrix multiplication.

You can see that if we had more neurons, we could just tack their weights on as extra rows in the first matrix. If we had more inputs, we would just add on extra rows for the values in the second matrix -- and, of course, we'd need weights for them, which would go into extra columns in the first matrix.

Another nifty thing that might not be immediately obvious is that the result matrix has the outputs stacked up as a column, just like the inputs were. If we added the bias and applied the activation function to those, we'd have the outputs in exactly the same format as the inputs. That's useful, because if we had a multi-layer neural network, where the output of the first layer was the input for the second, and so on, then we could just do the calculations for the first and then feed the results directly into the calculations for the second.

But obviously we need to sort out the bias and activation functions first.

Batches

Before we do that, though: we've seen what we can do by adding on rows to the first matrix so that it changes from being the weights for just one neuron to being the weights for all of the neurons in the layer. What happens if we add on extra columns to the inputs matrix? We get the ability to do batches -- that is, to run a bunch of separate sets of inputs through the neural network in one go:

$$\left[\begin{array}{ccc}{w}_{1,1}& w_{1,2}& w_{1,3}\\ w_{2,1}& w_{2,2}& w_{2,3}\end{array}\right]\left[\begin{array}{cc}{x}_1^1& x_1^2\\ x_2^1& x_2^2\\ x_3^1& x_3^2\end{array}\right]=\left[\begin{array}{cc}{w}_{1,1}·x_1^1+w_{1,2}·x_2^1+w_{1,3}·x_3^1& w_{1,1}·x_1^2+w_{1,2}·x_2^2+w_{1,3}·x_3^2\\ w_{2,1}·x_1^1+w_{2,2}·x_2^1+w_{2,3}·x_3^1& w_{2,1}·x_1^2+w_{2,2}·x_2^2+w_{2,3}·x_3^2\end{array}\right]$$

That's getting a bit busy, but hopefully you can see that by providing two sets of inputs, one column for each, in our second matrix, we have -- with a single matrix multiplication -- evaluated the neural network for both inputs, and got our weighted sums.

Let's simplify the whole thing. We can call our matrix of all of the weights -- one row per neuron -- $W$, the matrix of our inputs -- one column per item in our batch -- $X$, and this intermediate weighted sum output we can call $Z$. The whole operation of this part of the evaluation of the neural network is just this:

$$Z=WX$$

That's it -- the bulk of the calculations for a neural network layer of arbitrary size, including batch processing, wrapped up in one simple equation!

But we still have that bias term. What can we do with that?

Bias: version 1

Well, the bias values are just a set of numbers that we need to add on to the weighted sums that we've already calculated, so we could just do that directly. Let's take the last example, but replace the worked-out dot products with $z$s for sanity's sake:

$$\left[\begin{array}{ccc}{w}_{1,1}& w_{1,2}& w_{1,3}\\ w_{2,1}& w_{2,2}& w_{2,3}\end{array}\right]\left[\begin{array}{cc}{x}_1^1& x_1^2\\ x_2^1& x_2^2\\ x_3^1& x_3^2\end{array}\right]=\left[\begin{array}{cc}{z}_1^1& z_1^2\\ z_2^1& z_2^2\end{array}\right]$$

That is, the weighted sums for our input $(x_1^1,x_2^1,x_3^1)$ are $(z_1^1,z_2^1)$, and the weighted sums for the input $(x_1^2,x_2^2,x_3^2)$ are $(z_1^2,z_2^2)$.

We can create a matrix of the same size as the outputs containing the bias terms as columns, repeated once for each example in the batch, and add that to our $Z$ matrix:

$$\left[\begin{array}{cc}{z}_1^1& z_1^2\\ z_2^1& z_2^2\end{array}\right]+\left[\begin{array}{cc}{b}_1& b_1\\ b_2& b_2\end{array}\right]=\left[\begin{array}{cc}{z}_1^1+b_1& z_1^2+b_1\\ z_2^1+b_2& z_2^2+b_2\end{array}\right]$$

...and that gives us our pre-activation values -- that is, the weighted sums plus the biases -- ready to be run through the activation function. We can write out the whole thing like this (using $\hat{Z}$ to represent the pre-activation values):

$$\hat{Z}=WX+B$$

This is actually quite close to how it's normally done in practice (with a couple of wrinkles that I'll get into later). But, at least to me, it feels slightly ugly. After decades of computer programming, having the size of $B$ have to dynamically change to adjust to the batch size feels kind of wrong -- it has the same kind of smell as an issue with the typing of a parameter, in some way I can't quite put my finger on. ³

Bias: version 2

There's another method, which feels more elegant to me, that is often used in more mathematical descriptions. Let's go back to our simple example from earlier, with a batch size of one:

$$\left[\begin{array}{ccc}{w}_{1,1}& w_{1,2}& w_{1,3}\\ w_{2,1}& w_{2,2}& w_{2,3}\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\\ x_3\end{array}\right]=\left[\begin{array}{c}{w}_{1,1}·x_1+w_{1,2}·x_2+w_{1,3}·x_3\\ w_{2,1}·x_1+w_{2,2}·x_2+w_{2,3}·x_3\end{array}\right]$$

The bias term is kind of like a weight. It goes into the neuron just like the product of the weights and the inputs do, and -- when we're training a neural network -- it's a trainable parameter just like the weights. The only difference is that it's not multiplied by an input. Well, multiplying a number by one is the same as not multiplying it by anything, so what happens if we extend the input matrix with some $1$s and just treat the bias terms as weights?

$$\left[\begin{array}{cccc}{w}_{1,1}& w_{1,2}& w_{1,3}& b_1\\ w_{2,1}& w_{2,2}& w_{2,3}& b_2\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\\ x_3\\ 1\end{array}\right]=\left[\begin{array}{c}{w}_{1,1}·x_1+w_{1,2}·x_2+w_{1,3}·x_3+b_1·1\\ w_{2,1}·x_1+w_{2,2}·x_2+w_{2,3}·x_3+b_2·1\end{array}\right]$$

We've done the whole pre-activation calculation in a single step! And we can do the same trick with batches of multiple inputs.

That, to me, is a very mathematically elegant way to do it. If we call the version of $X$ that has been adjusted to have a row of $1$s at the bottom $\bar{X}$. and the weights adjusted to have the bias terms in a new column on the right $\bar{W}$, our pre-activation function calculations come down to this:

$$\hat{Z}=\bar{W}\bar{X}$$

For some reason, "hacking" the input and weight matrices to add on an extra row/column feels less inelegant to me than changing the size of the bias matrix to match the batch size. So I personally prefer this formulation. It also puts all of the learnable parameters for the neural network in one variable, which feels nice and clean.

Wrapping up

So, that's how we can use matrix maths to simplify the calculation of a neural network -- even when we're calculating batches of inputs in one go -- down to a one-line formula.

As I said, though, in real-world ML code, we actually use a formula that is similar to these, but differs in important ways. It normally looks like this:

$$\hat{Z}=X{W}^T+B$$

You can see that we're using the calculation I found less elegant, but we've swapped around $X$ and $W$, and $W$ is also transposed (that's what the superscript $T$ is for). Why? I'll post about that next.