AlphaTensor —AI discovering faster mathematical algorithms

Image Source: Author

What if an AI could make itself faster? 🤔

Well In 2022, DeepMind released Alpha Tensor, where an AI found more efficient ways to multiply matrices.

The way matrix multiplication algorithms have normally been discovered has been through human search, systematic exploration, and a few other techniques. This is not efficient.

This invention is crazy because it broke a longstanding 50-year streak from the previous algorithm that was thought to be the fastest way

Even crazier 2 days later, 2 random Austrian mathematicians made it even better. They also managed to train it to be faster for specific hardware.

Image Source: Deepminds Paper

It’s inspired by AlphaZero, an AI that can learn Chess, Go, and Shogi on its own.

The thing is, the search space is so large even the most efficient way of multiplying 3x3 matrices is still unknown. For comparison, remember AlphaGo playing against Lee Sedol? Well, the search space is 30 times bigger than the game of Go.

But nonetheless, they still found faster algorithms.

Now in 2023, they trained AlphaDev, an AI that found faster sorting algorithms 🔥.

So I got curious. And now I’m here to write an article to do a really deep dive article to dig deep into this explaining how it works and why.

What is a Tensor in AlphaTensor?

A tensor is a way to organize data. They come in different shapes and sizes.

Image Source: Author |

It’ll keep on going, 4D, 5D, 6D…

For simple definitions:

Think of 0D as a single number.

Think of 1D as numbers in a list/vector.

Think of 2D as numbers in a square box.

Think of 3D as numbers in a cube.

Tensors are used in ML because they’re useful for encoding multi-dimensional data. It’s a universal way of storing and manipulating data.

Onto the basics: Matrix Multiplication

Matrix Multiplication is a crucial operation for things like Neural Networks, Quantum Mechanics, and Computer Graphics.

Even a slight increase can yield significant results, giving the ability to tackle problems that were far too big before.

This is the algorithm we’re taught in high school:

Image Source: Geeksforgeeks

All you do is multiply the rows with the columns.

All of science🤯 | Image Source: Me

However, it’ll take an absurd amount of multiplications once you start increasing the matrix size. The time complexity of multiplying a square matrix is O(n³)

This means the # of multiplications needed increases to power 3 as the matrix size increases.

Multiplying two 2x2 matrices only takes 8 (2³) multiplications. But two 10x10 matrices would make it go from 8 -> 1000 (10³) multiplications.

It may seem like you HAVE to do all these multiplications. However…

In 1969, Strassen created an algorithm that needed fewer multiplications traditionally through the strategy “Divide and conquer”. His algorithm on multiplying two 2x2s, only uses 7 multiplications instead of 8. This might not seem much, but the speed really adds up as the matrix gets bigger.

1st step | Strassens Algorithm | Image Source: Wikipedia

2nd step | Strassens Algorithm | Image Source: Wikipedia

He’s traded the × for +/-, funny we’ve also learned this in high school. Take a look at this equation:

a² - b² = (a+b)(a-b)

1.
= 3² - 2²
= 9 - 4
= 5

2.
= (3+2)(3-2)
= (5)(1)
= 5

The 1st one uses 2 multiplication. But the 2nd one only uses 1? It’s smartly arranged in a way where only 1 is needed, to visualize it:

1. We use FOIL to distribute the 2 binomials
2. Then combine like terms
---------------------------------------
= (a+b)(a-b)                          | The faster formula
= (a)(a) + (a)(-b) + (a)(b) + (b)(-b) | FOIL to distribute
= a² - ab + ab - b²                   | Combining like terms
= a² + (-ab + ab) + (-b²)             | Notice (-ab + ab) cancels out 👀
= a² - b²                             | Original equation!

Props to my math teacher, thanks to her I learned it today.

This applies to matrix multiplication world

Look at the original matrix multiplication algorithm.

We’re going to calculate for c² | Image Source: Author

Then at Strassens algorithm.

The same process of what I did above | Author

It goes back to the original algorithm! This is why even though it seems impossible to speed up, all it’s doing is trading the × for +/-

Strassen's finding of this algorithm led to a continuous amount of people competing in improving the algorithm the latest being a time complexity of O(n².⁴³⁷¹⁵⁵²) in 2023.

Back to AlphaTensor, simply telling the AI to find an algorithm with the lowest amount of multiplication steps is a super vague request. We have to reframe the problem more specifically.

Tensor Decomposition to hone in on the process

The matrix multiplication algorithm can also be represented as a 3D Tensor (Cube). Meaning this:

The matrix multiplication algorithm | Image Source: Me

Can also be shown as.

3D Tensor of only 1 and 0’s | Image Source: AlphaTensors research paper

If you look at the position of each variable in the algorithm, they’ll actually match up with the image.

Next, we “decompose” from 3D into a 1D Tensors. Just like taking apart a cube into smaller pieces. Basically, the algorithm turns into multiple vectors lined up side to side. For example, on the right is the decomposed version of Strassen’s algorithm.

Note: All AlphaTensor does is interact with the algorithm. Not with anything else. | Image source: AlphaTensors Paper

I’ll break down what the right side means. Let's take m¹ = (a¹ + a⁴)(b¹+b⁴) from the left.

Image Source: Me

Notice the numbers in each column of “U and V” show where to add what. The 1st Column of U and V is:

U¹ = [1, 0, 0, 1] — Corresponds to (a¹ + a⁴)

V¹ = [1, 0, 0, 1] — Corresponds to (b¹ + b⁴)

Once these 2 are multiplied, it’ll represent the variable M¹.

So if this is 1 step, notice every vector of U,V dictates a multiplication step. Meaning 7 columns = 7 multiplications.

W shows when the M variables get used, each row shows where to add or subtract M to get the final values in the matrix.

W¹ = [1, 0, 0, 1, -1, 0, 1] — Corresponds to m¹ + m⁴ — m⁵ + m⁷

Do this for all of them, and you’ll get exactly the final algorithm. See how decomposing U,V,W can represent it?

To know if it’s a valid decomposition or not. We can reconstruct it back together to see if it matches with the actual one.

This can be done by

1. Taking the outer product of each vector in each column and adding them all together.
2. Subtracting the sum from the original 3D Tensor and see if you get a 0.

Why 0? If we subtract a valid algorithm 3D Tensor from the original 3D Tensor, it should be the same and therefore be 0.

Think of the outer product like doing matrix multiplication, except the second term is transposed. For example, the steps for calculating the outer product of M¹ would be like this:

A is the (a1 + a4). B is the (b1 + b4). W is the desired location it’ll add up to

The outer product of A and B results as:

Can also be thought of as A ⋅ B but B is Transposed | Image Source: Author

And then taking the outer product of that matrix with W, it becomes a 3D Tensor.

I used a cube instead to represent it as it’s easier to see. The colored blocks are 1’s and the non-colored are 0’s. | Image Source: Author

We’ll do this for all the M columns. Add them up to get one final 3D Tensor, and if subtracted with the original 3D Tensor to get a 0, it’s valid.

Mathematical representation for this | Image Source: AlphaTensor’s Paper

TLDR; So long as the decomposed Tensor adds back up to the original 3D Tensor. It’s a valid algorithm.

Notice that the fewer 1D Tensors there are, the fewer multiplications are used. The goal is to use the least amount of 1D Tensors to decompose the original 3D Tensor.

Finding the fastest algorithm through TensorGame

Now it can be reframed as a Reinforcement Learning problem. Which is just a way to learn by itself through trial and error using feedback from its actions. DeepMind designed this single-player game for AlphaTensor to play, calling it TensorGame.

Decomposing a 3D Tensor itself is hard, but constructing one using vectors is really simple. So the agent is trained to guess tensors to subtract from the original one until it reaches 0.

It works like this:

1. The current state is the target tensor to decompose, so the current state is the original 3D Tensor.
2. A set of 3 vectors (U, V, W) is selected then take the outer product of these, and then subtracted from the current state. Every time it does this, it gets a penalty (negative reward) to encourage the AI to use the least amount of steps.
3. Repeat until the result is a zero-tensor, if it goes into the negatives it’ll get even more penalty.

To encourage MCTS to reach an end state and not just waste time, rewards are also given based on the closeness to the closeness goal.

The way they increased the speed on a given hardware, they added an extra reward based on the runtime. That’s all it needs to know to train a speedup.

Training the AI to play the game — Act and Learn

The main question now is, how does the AI know which set of Vectors to subtract to get a valid decomposition as fast as possible?

Well, there are 2 steps it repeats to do so:

1. Acting
2. Learning

This is the entire process:

Acting:

1. TensorGame starts, input is initialized.
2. Data augmentation is done on the starting Tensor for more diversity.
3. Monte Carlo Tree Search ( A tree search algorithm ) with reinforcement learning will decide what to do next.
4. Once a decomposition is found or the maximum number of steps is reached. All the states of this game + the chosen actions will be added to the list of played games for the Learning phase to start.

Learning:

1. The AI goes into training mode by randomly sampling from the list of played games / synthetic data and learning how to predict which moves are better in the future.
2. The updated model now has a more accurate way of predicting which moves are better to search/choose from in the tree search. That way it’ll be able to find better paths.
3. Go back into the Acting phase to try it out and keep improving.

Going into more detail.

Synthetic Data + Diversifying Target Tensor

Powerful AI models need lots of data, though it’s a lot harder to find in math problems.

So they did 3 things to increase data:

1. Linear Algebra property: Change of basis. Think of it as a different way to represent the same 3D Tensor. They would generate random bases up to 100,000 and would play them all in Parallel 🤯. The purpose is if the AI finds an algorithm in ANY of these bases. It’ll be able to convert all the way back to a valid algorithm.

Solving the game is really hard as most outcomes would end up an invalid decomposition. To help speed up the process:

2. They used Synthetic Data to help with training. Remember how decomposing a 3D Tensor is really hard, but constructing one is really elementary? They constructed a synthetic game because now they’d know what the outcome was and the steps needed to get there. This is to give the AI more success cases to learn from.

3. Swapping move positions within the game to get more data. Imagine there’s the number five. 5–2–3=0. If you swap the order of subtractions: 5–3–2=0, you’d still get the same answer. They could swap it for each & every move, but that’s too much data. So they just swapped the last move with an intermediate move.

Monte Carlo Tree Search with Reinforcement Learning

MCTS is a tree search algorithm that, unlike the other ones that perform an exhaustive search, will only search the most promising routes. It does this by being guided by an AI. Specifically with 2 kinds of neural networks:

1. Policy network: outputs a list of probabilities on which path to take (learns based on what actions were chosen in the game and optimizes for that)
2. Value network: estimates the reward of the current state (learns based on the outcome of the game and optimizes for that)

The possible actions are massive though. 1 action from MCTS = a choice of 3 vectors (u,v,w). So they created a separate policy network for u, v, and w.

Each path it chooses is influenced by the reward that’s given, the AI will think “How can I choose a path that’ll maximize the reward?”

It can learn because it already knows the outcome and what happened during the game, so it’ll start training by randomly sampling a state from the list of played games/synthetic data to learn from.

When it first starts off, it’ll practically be guessing random options, when there’s more training down. The Policy & Value network will be better at predicting which paths/nodes are better to choose from in the future based on what happened on the data.

The AI can actually transfer what it learned onto other matrices too. They’d start training with a 25x25. But if you wanted to matrix multiply a 3x3, put it in the 25x25 and the rest will be 0's.

Image Source: Author

So it’s able to transfer what It learned unless it goes beyond 25x25.

Transformer-based architecture guiding the tree search

(A transformer is what ChatGPT uses, some of its components have been used here.)

This is the actual AI part of it while training, composed of a Torso (encoder) and a Policy and Value network.

The input is initialized which consists of the current state, previous states, and the chosen actions.

1. Input goes through the torso which is based on a special kind of attention called Axial Attention, it’s interesting because the original AlphaZero used CNNs. The purpose is to create representational embedding of the input that’s useful for both networks. It’s a way to turn something like words/pictures/data into data a computer can easily understand.
2. Embeddings go through the policy head model, which will pay attention to what happened before (previous time steps) to make its predictions.
3. Embeddings go through the value head, which is a simple multilayer perceptron. (A fully connected feed-forward neural network.)

Resources Used

For Learning: TPU v3, 64 TPU Cores, trained for 600k iterations

For Acting: TPU v4, 1.6k actors.

TPUs are like GPUs but tailored specifically for training AI. It took a week before it started getting to convergence. (A point where it reached its lowest loss and can’t seem to find more improvement)

The Results

Within a few minutes of training, it rediscovered Strassen’s algorithm, and then… it went further beyond.

Image Source: AlphaTensors paper

This is the 4x4 algorithm they found. It uses 47 multiplications rather than Strassen’s 49.

Image Source: AlphaTensors Paper

They found a ton more faster algorithms listed in the Paper. And the improvement difference seems to increase as the matrix size gets bigger.

And yep that's about it!

I will see you in my next one.

[1] Wikipedia contributors. “Strassen algorithm.” *Wikipedia, The Free Encyclopedia*. Wikipedia, The Free Encyclopedia, 12 Oct. 2023. Web. 13 Oct. 2023.

[2]: Wikipedia contributors. “Computational complexity of matrix multiplication.” *Wikipedia, The Free Encyclopedia*. Wikipedia, The Free Encyclopedia, 3 Sep. 2023. Web. 13 Oct. 2023.

[3] Fawzi, A., Balog, M., Huang, A. et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature 610, 47–53 (2022). https://doi.org/10.1038/s41586-022-05172-4

[4] This is a game-changer! (AlphaTensor by DeepMind explained) YouTube . Yannic Kilcher .Oct 7, 2022 https://www.youtube.com/watch?v=3N3Bl5AA5QU