DeepMind Unveils First AI to Discover Faster Matrix Multiplication Algorithms

Can artificial intelligence (AI) create its own algorithms to accelerate matrix multiplication, one of the most fundamental tasks of machine learning? Today, in a paper published in Nature, revealed DeepMind AlphaTensor, the “first artificial intelligence system for discovering new, efficient and demonstrably correct algorithms.” Google’s lab said the research “sheds light” on a 50-year-old math open question about finding the fastest way to multiply two matrices.

Since the Strassen algorithm was published in 1969, computer science has been looking to exceed the speed of multiplying two matrices. While matrix multiplication is one of the simplest operations of algebra taught in high school math, it is also one of the most fundamental arithmetic tasks and, as it turns out, one of the most important math operations in neural networks today.

Matrix multiplication is used to process smartphone images, understand voice commands, generate computer graphics for computer games, data compression, and more. Today, companies use expensive GPU hardware to increase the efficiency of matrix multiplication, so any extra speed would be a huge change in terms of cost reduction and energy savings.

AlphaTensor, according to a DeepMind blog post, builds on AlphaZero, an agent who has delivered superhuman feats in board games such as chess and Go. This new work takes the AlphaZero journey further, moving from playing games to tackling unsolved math problems.

DeepMind uses AI to improve computer science

This research looks at how AI can be used to improve computer science itself, Pushmeet Kohli, chief of AI for science at DeepMind, said at a news conference.

“If we can use AI to find new algorithms for fundamental computational tasks, it has huge potential because we may be able to go beyond the algorithms currently in use, which could lead to improved efficiency,” he said.

This is a particularly challenging task, he explained, because the process of discovering new algorithms is so difficult, and automating algorithmic discovery using AI requires a long and difficult reasoning process — from forming intuition about the algorithmic problem. to actually writing a new algorithm and proving that the algorithm is correct in specific cases.

“This is a tough set of steps and AI hasn’t been very good at it so far,” he said.

An ‘Intriguing, Baffling Problem’

DeepMind took on the challenge of matrix multiplication because it’s a known problem in computation, he said.

“It’s also a very intriguing, perplexing problem, because matrix multiplication is something we learn in high school,” he said. “It’s an extremely simple operation, but we don’t currently know how best to multiply these two sets of numbers. So that is also enormously stimulating for us as researchers to better understand this.”

According to DeepMind, AlphaTensor discovered algorithms that are more efficient than the state of the art for many matrix sizes and outperform humans designed.

AlphaTensor starts out without any knowledge of the problem, Kohli explained, and then gradually learns what’s happening and improves over time. “It first finds this class algorithm that we were taught, and then it finds historical algorithms like Strassen’s and at some point it surpasses them and discovers completely new algorithms that are faster than before.”

Kohli said he hopes this article inspires others to use AI to guide algorithmic discoveries for other fundamental competitive tasks. “We think this is an important step towards really using AI for algorithmic discovery,” he said.

DeepMind’s AlphaTensor uses AlphaZero

According to Thomas Hubert, staff research engineer at DeepMind, it’s really AlphaZero that runs behind the scenes of AlphaTensor as a single-player game. “It’s the same algorithm that taught chess that was applied here for matrix multiplication, but it had to be extended to handle this infinitely large space — but many of the components are the same,” he said.

In fact, according to DeepMind, this game is so challenging that “the number of possible algorithms to consider far exceeds the number of atoms in the universe, even for small cases of matrix multiplication.” Compared to Go, which has been an AI challenge for decades, the number of possible moves is 30 orders of magnitude greater.

“The game is about zeroing the tensor, with some allowed moves that actually represent some algorithmic operations,” he explained. “This gives us two very important results: one is that if you can parse zero from the tensor perfectly, you are guaranteed to have a demonstrably correct algorithm. Second, the number of steps required to factor this tensor indicates the complexity of the algorithm. So it’s very, very clean.”

DeepMind’s paper also pointed out that AlphaTensor is discovering a richer space of matrix multiplication algorithms than previously thought – up to thousands for each size.

According to the blog post, the authors said they modified AlphaTensor to specifically find algorithms that are fast on a particular hardware, such as Nvidia V100 GPU and Google TPU v2. These algorithms multiply large matrices 10-20% faster than commonly used algorithms on the same hardware, demonstrating AlphaTensor’s flexibility in optimizing arbitrary targets,” the blog post reads.

Greater impact of AI on science and mathematics

Back in July, researchers showed that DeepMind’s AlphaFold tool could predict the structures of more than 200 million proteins from about one million species, encompassing nearly every known protein on Earth. Kohli said AlphaTensor shows the potential that AI has not only in science, but also in mathematics.

“To see AI deliver on that promise of going beyond what human scientists have been able to do over the past 50 years is incredibly exciting personally,” Kohli said. “It just shows how much impact AI and machine learning can have.”