Matmul, simplified: a deep learning essential

You’ve probably seen the torch.matmul function before, but do you really — and I mean, really — understand what it does?

In this article, we will dive into this amazing PyTorch functionality.

To start, matmul means matrix multiplication. This kind of operation is like a machine that can transform vectors, stretching, squashing, or moving them around.

But why vectors? In deep learning, what we’re doing is essentially projecting our data, organized as vectors, into higher dimensions so we can “untangle” it and learn relevant patterns from it. For that, matrix multiplication is essential.

The magic of matmul in PyTorch is that it still performs a linear projection, but it can handle applying this transformation to many instances at the same time.

Let’s get into it.

Single vector

Let’s start with a simple example of a single vector of 64 dimensions, x. If we want to project it to a higher dimension, e.g., 128, what should we do?

The standard practice is to multiply x by a transformation matrix A, with the dimensions 128 (our target dimension) by 64 (our current dimension).

Each row of matrix A represents a linear combination of the 64 features in x that maps it into the new 128-dimensional space.

So, the result of the matmul between A and x will be a new vector y, of 128 dimensions. Awesome, right?

Linear projection of a single vector with matmul.

Intuition: batch multiplication

However, in deep learning, things are usually not that straightforward. This is because there are often many vectors to be transformed at once.

So, instead of just one vector x, we might have, for example, 196 vectors.

Let’s translate this to a real case. What if we have 8 text passages, each with 196 tokens (e.g., words), and each represented as vectors of 768 dimensions?

Given this, how do we transform the 768-dimensional vectors into 128-dimensional representations?

Should we apply matrix A to each token individually and then aggregate the results across the 8 batches? That’s possible, but the computational cost would be too high.

So, what torch.matmul allows us to do seamlessly is: for every token in every batch, perform a multiplication of x and A!

In this case, A is usually called W (the weight matrix). When many instances are processed at once, this procedure is part of what we call a feedforward pass through a linear layer in a neural network.

Code

Here’s the batch example I mentioned in PyTorch:

X = torch.rand(8, 12, 196, 768)
W = torch.rand(768, 128)
Y = torch.matmul(X, W)  # torch.Size([8, 12, 196, 128])

See how the difference between X and Y is simply in the last dimension of the tensor? This is because the function keeps the batches and the tokens intact, and only applies the projection to the last dimension (the features dimension).

By the way, torch.matmul also automatically handles the earlier example of a single vector projection:

x = torch.rand(1, 64)
A = torch.rand(64, 128)
y = torch.matmul(x, A)  # torch.Size([1, 128])

See how it adjusts seamlessly? Magical.

See you next time.