Understanding DeepMind and Strassen algorithms | by Stefano Bosisio | Oct, 2022

An introduction to the matrix multiplication downside, with functions in Python and JAX

DeepMind lately revealed an fascinating paper that employed Deep Reinforcement Studying to seek out new Matrix Multiplication algorithms[1]. One of many targets of this paper is to decrease matrix multiplication computational complexity. The article has raised a number of feedback and questions on matrix multiplication — as you’ll be able to see from Demis Hassabis’ tweet.

Matrix multiplication is an intense analysis space in arithmetic [2–10]. Though matrix multiplication is a straightforward downside, the computational implementation has some hurdles to resolve. If we’re contemplating solely sq. matrices, the primary thought is to compute the product as a triple for-loop :

Such a easy calculation has a computational complexity of O(n³). Which means that the time for operating such a computation will increase because the third energy of the matrix measurement. This can be a hurdle to deal with, as in AI and ML we take care of big matrices for each mannequin’s step — neural networks are tons of matrix multiplications! Thus, given a continuing computational energy, we want an increasing number of time to run all our AI calculations!

DeepMind has introduced the matrix multiplication downside to a extra concrete step. Nevertheless, earlier than digging into this paper, let’s take a look on the matrix multiplication downside and what algorithms can assist us in decreasing the computational energy. Specifically, we’ll check out Strassen’s algorithm and we’ll then implement it in Python and JAX.

Keep in mind that for the remainder of the paper the dimensions of matrices might be N>>1000. All of the algorithms are to be utilized to block matrices.

The product matrix C, is given by the sum over the rows and the columns of matrices A and B, respectively — fig.2.

As we noticed within the introduction, the computational complexity for the usual matrix multiplication product is O(n³). In 1969 Volker Strassen, a German mathematician, smashed the O(n³) barrier, decreasing the matrix multiplication to 7 multiplications and 18 additions, attending to a complexity of O(n²·⁸⁰⁸)[8]. If we take into account a set of matrices A, B and C, as in fig.3, Strassen derived the next algorithm:

It’s price noticing a couple of issues about this algorithm:

• the algorithm works recursively on block matrices
• it’s simple to show the complexity O(n²·⁸⁰⁸). Given the matrix measurement n and the 7 multiplications, it follows that:

All of the steps in fig.4 are polynomials, subsequently the matrix multiplication might be handled as a polynomial downside.

• Computationally Strassen’s algorithm is unstable with floating precision numbers [14]. The numerical instability is attributable to rounding all of the sub-matrices outcomes. Because the calculation progresses the entire error sums up in a grave lack of accuracy.

From these factors we are able to translate the matrix multiplication to a polynomial downside. Every operation in fig.4 might be written as a linear mixture, for instance steps I is:

Right here α and β are the linear combos from matrices A and B’s parts, whereas H denotes a one-hot encoded matrix for addition/subtraction operations. It’s then potential to outline the product matrix C parts as a linear mixture and write:

As you’ll be able to see the complete algorithm has been lowered to a linear mixture. Specifically, the left-hand facet of the equation Strassenin fig.7 might be denoted by the matrices sizes, m, n, and p — which implies a multiplication between mxp and nxp matrices:

For Strassen <n, m, p> is <2,2,2>. Fig.8 describes the matrix multiplication as a linear mixture, or a tensor — that’s why typically the Strassen algorithm known as “tensoring” . The a, b, and c parts in fig.8 kind a triad. Following DeepMind’s paper conference the triad might be expressed as:

This triad establishes the target of discovering the most effective algorithm for minimizing the computational complexity of the matrix multiplication operation. Certainly, the minimal variety of triads defines the minimal variety of operations to compute the product matrix. This minimal quantity is the tensor’s rank R(t). Engaged on the tensor rank we’ll assist us to find new and extra environment friendly matrix multiplication algorithms — and that’s what DeepMind individuals have finished.

Ranging from Strasse’s work, between 1969 and right now there’s been a continuing creation of recent algorithms to resolve the matrix multiplication complexity downside (tab.1).

How did Pan, Bini, Schonage, and all of the researchers get to those good outcomes? One solution to remedy a pc science downside is to begin with the definition of an algebraic downside P. For the matrix multiplication downside, for instance, the algebraic downside P might be: “discover a mathematical mannequin to guage a set of polynomials”. From right here, scientists begin to scale back the issue and “convert” it to a matrix multiplication downside — right here is an efficient clarification of this method. In a nutshell, scientists have been in a position to show theorems in addition to steps that would decompose the polynomial analysis to a matrix multiplication algorithm. Finally, all of them acquired a theoretical algorithm that may be extra highly effective than the Strassen algorithm (tab.1)

Nevertheless, these theoretical algorithms can’t be coded up, until there are some heavy and powerful mathematical assumptions and restrictions that would have an effect on the algorithm’s effectivity.

Let’s now see how highly effective Strassen’s algorithm is and the way we are able to implement it in Python and JAX.

Right here is the repo with all the next codes. I run these assessments on a MacBook Professional, 2019, 2.6GHz 6-Core Intel Core i7, 16 GB 2667MHz DDR4 Reminiscence.

Within the important code we are able to comply with these steps:

• we’re going to create 2 sq. matrices A and B, initialised with random integers
• we’re going to take a look at the algorithms for various matrices’ sizes: 128, 256, 512, 768, 1024, 1280, 2048
• For every measurement will run numpy.matmul and Strassen’s algorithms thrice. At every run we’re recording the operating time in an inventory. Kind this listing we’re extracting the common time and the usual deviation to check each strategies (fig.10)

The core a part of the script is the recursivestrassen perform:

• Firstly, we examine enter matrix dimension. If the dimension is under a given threshold (or not divisible by 2) we are able to compute the remaining product with commonplace numpy, as this gained’t affect the ultimate computational price
• For every enter matrix the highest left, high proper, backside left and backside proper sub-matrices are extracted. In my code, I’m proposing a naive and easy resolution, so everyone can perceive what’s occurring. To additional take a look at and perceive the block matrix creation, attempt to manually compute the indices for a small matrix (e.g. 12×12).
• Within the ultimate step, the product matrix is reconstructed from all of the computed sub-elements ( C11, C12, C21 and C22 in fig.4)

Fig.10 compares the usual numpy.matmul and strassen algorithm. As you’ll be able to see for a dimension < 2000 ( Matrix Measurement < 2000 ) Strassen might be outperformed by the usual matrix multiplication. The true enchancment might be seen on greater matrices. Strassen completes the matrix multiplication for a 2048×2048 matrix in 8.16 +/- 1.56 s, whereas the usual strategies required 63.89 +/- 2.48 s. Doubling the matrix measurement, 4096 columns and rows, Strassen runs in 31.57 +/- 1.01 s, whereas the usual matrix multiplication takes 454.37 +/- 6.27 s.

In response to the equation in fig.9 we are able to additional decompose the Strassen algorithm in a tensor kind. The tensors u, v and w can then be utilized to the matrices’ blocks to acquire the ultimate product matrix. C.H. Huang, J. R. Johnson, and R. W. Johnson revealed slightly paper to indicate learn how to derive the tensor model of Strassen [18], adopted by one other formulation in 1994 [19] the place they explicitly wrote Strassen’s u, v and w tensors. For the detailed calculations you’ll be able to examine [18], whereas fig.12 studies the tensors values.

This can be a good place to begin for working with JAX and evaluating Strassen to the usual jax.numpy.matmul. For the JAX script I’ve adopted intently DeepMind’s implementation.

The script offers with 4×4 block matrices. The core perform, f, runs Strassen technique. On this case, all of the A and B block matrices are multiplied by the uand v tensors. The result’s multiplied by tensor w, acquiring the ultimate product (fig.13). Given JAX highly effective efficiency, the algorithm was examined on the next matrices’ dimensions: 8192, 10240, 12288, 14336, 16384, 18432, 20480

Lastly, within the final step the product matrix is reconstructed by concatenating and reshaping the product matrix from f perform (fig.14).

Fig. 15 compares JAX numpy matrix multiplication with the Strassen implementation. As you’ll be able to see JAX could be very highly effective, as 8192×8192 matrices multiplication might be run in 12 s (on common). For dimensinos underneath 12000×12000 there isn’t any actual enchancment and JAX commonplace technique takes a mean computational time of 60s on my laptop computer — whereas I’m operating another issues. Above that dimensions we are able to see a powerful 20% enchancment. For instance, for 18432×18432 and 20480×20480, the Strassen algorithm runs in 142.31+/-14.41 s and 186.07+/-12.80 s, respectively — and this was finished by operating on a CPU. A very good homework may very well be attempting this code including the device_put choice and operating on Colab’s GPU. I’m certain you’ll be flabbergasted!

In the present day we made slightly step ahead to get a whole understanding of DeepMind’s publication “Discovering sooner matrix multiplication algorithms with reinforcement studying” [1]. This paper proposes new methods to deal with the matrix multiplication downside, utilizing Deep Reinforcement Studying.

On this first article, we began to scratch the floor of matrix multiplication. We realized what’s the computational price for this operation and we noticed the Strassen algorithm.

From right here we outlined how the Strassen algorithm is made and what are its mathematical implications. Since its publication, researchers have discovered higher and higher options to the matrix multiplication downside. Nevertheless, not all of those strategies might be carried out in code.

Lastly, we performed a bit with Python and JAX to learn the way highly effective the algorithm is. We realized that Strassen is a superb instrument to make use of when we now have to take care of very massive matrices. We noticed the facility of JAX in dealing with massive matrix multiplications and the way simple is to implement such an answer, with out utilizing GPUs or additional reminiscence choices.

Within the subsequent paper, we’ll see extra particulars from DeepMind’s paper. Specifically, we’ll deal with the deep reinforcement algorithm, in addition to the paper findings. Then, we’ll implement the brand new DeepMind algorithms and run them in JAX on a GPU occasion.

I hope you loved this text 🙂 and thanks for studying it.

Please, be at liberty to ship me an e-mail for questions or feedback at: stefanobosisio1@gmail.com or straight right here in Medium.

1. Fawzi, Alhussein, et al. “Discovering sooner matrix multiplication algorithms with reinforcement studying.” Nature610.7930 (2022): 47–53.
2. Bläser, Markus. “Quick matrix multiplication.” Concept of Computing (2013): 1–60.
3. Bini, Dario. “O (n2. 7799) complexity for nxn approximate matrix multiplication.” (1979).
4. Coppersmith, Don, and Shmuel Winograd. “On the asymptotic complexity of matrix multiplication.” SIAM Journal on Computing 11.3 (1982): 472–492.
5. Coppersmith, Don, and Shmuel Winograd. “Matrix multiplication through arithmetic progressions.” Proceedings of the nineteenth annual ACM symposium on Concept of computing. 1987.
6. de Groote, Hans F. “On kinds of optimum algorithms for the computation of bilinear mappings II. Optimum algorithms for two× 2-matrix multiplication.” Theoretical Pc Science 7.2 (1978): 127–148.
7. Schönhage, Arnold. “A decrease sure for the size of addition chains.” Theoretical Pc Science 1.1 (1975): 1–12.
8. Strassen, Volker. “Gaussian elimination just isn’t optimum.” Numerische mathematik 13.4 (1969): 354–356.
9. Winograd, Shmuel. “On multiplication of two× 2 matrices.” Linear algebra and its functions 4.4 (1971): 381–388.
10. Gentleman, W. Morven. “Matrix multiplication and quick Fourier transforms.” The Bell System Technical Journal 47.6 (1968): 1099–1103.
11. Alman, Josh, and Virginia Vassilevska Williams. “A refined laser technique and sooner matrix multiplication.” Proceedings of the 2021 ACM-SIAM Symposium on Discrete Algorithms (SODA). Society for Industrial and Utilized Arithmetic, 2021.
12. Le Gall, François. “Powers of tensors and quick matrix multiplication.” Proceedings of the thirty ninth worldwide symposium on symbolic and algebraic computation. 2014.
13. Williams, Virginia Vassilevska. “Multiplying matrices sooner than Coppersmith-Winograd.” Proceedings of the forty-fourth annual ACM symposium on Concept of computing. 2012.
14. Bailey, David H., King Lee, and Horst D. Simon. “Utilizing Strassen’s algorithm to speed up the answer of linear methods.” The Journal of Supercomputing 4.4 (1991): 357–371.
15. Pan, V. Ya. “Strassen’s algorithm just isn’t optimum trilinear strategy of aggregating, uniting and canceling for developing quick algorithms for matrix operations.” nineteenth Annual Symposium on Foundations of Pc Science (sfcs 1978). IEEE, 1978.
16. Schönhage, Arnold. “Partial and complete matrix multiplication.” SIAM Journal on Computing 10.3 (1981): 434–455.
17. Davie, Alexander Munro, and Andrew James Stothers. “Improved sure for complexity of matrix multiplication.” Proceedings of the Royal Society of Edinburgh Part A: Arithmetic 143.2 (2013): 351–369.
18. Huang, C-H., Jeremy R. Johnson, and Rodney W. Johnson. “A tensor product formulation of Strassen’s matrix multiplication algorithm.” Utilized Arithmetic Letters 3.3 (1990): 67–71.
19. Kumar, Bharat, et al. “A tensor product formulation of Strassen’s matrix multiplication algorithm with reminiscence discount.” Scientific Programming4.4 (1995): 275–289.