SlowGuess/NeurIPS_2023_OCR huggingface dataset:
MONARCH MIXER: A Simple Sub-Quadratic GEMM-Based Architecture
Daniel Y. Fu $^{1}$ , Simran Arora $^{*,1}$ , Jessica Grogan $^{*,2}$ , Isys Johnson $^{*,2}$ , Sabri Eyuboglu $^{*,1}$ , Armin W. Thomas $^{*,3}$ , Benjamin Spector $^{1}$ , Michael Poli $^{1}$ , Atri Rudra $^{2}$ , Christopher Ré $^{1}$
*Equal Contribution.  ${}^{1}$  Department of Computer Science,Stanford University.
$^{2}$ Department of Computer Science and Engineering, University at Buffalo, SUNY.
$^{3}$ Department of Psychology, Stanford University.
danfu@cs.stanford.edu, simarora@stanford.edu,
{jrgrogan,isysjohn}@buffalo.edu，{eyuboglu,athms,bfs,poli}  $@$  stanford.edu,
atri@buffalo.edu, chrismre@cs.stanford.edu
Machine learning models are increasingly being scaled in both sequence length and model dimension to reach longer contexts and better performance. However, existing architectures such as Transformers scale quadratically along both these axes. We ask: are there performant architectures that can scale sub-quadratically along sequence length and model dimension? We introduce MONARCH MIXER (M2), a new architecture that uses the same sub-quadratic primitive along both sequence length and model dimension: Monarch matrices, a simple class of expressive structured matrices that captures many linear transforms, achieves high hardware efficiency on GPUs, and scales sub-quadratically. As a proof of concept, we explore the performance of M2 in three domains: non-causal BERT-style language modeling, ViT-style image classification, and causal GPT-style language modeling. For non-causal BERT-style modeling, M2 matches BERT-base and BERT-large in downstream GLUE quality with up to  $27\%$  fewer parameters, and achieves up to  $9.1 \times$  higher throughput at sequence length 4K. On ImageNet, M2 outperforms ViT-b by  $1\%$  in accuracy, with only half the parameters. Causal GPT-style models introduce a technical challenge: enforcing causality via masking introduces a quadratic bottleneck. To alleviate this bottleneck, we develop a novel theoretical view of Monarch matrices based on multivariate polynomial evaluation and interpolation, which lets us parameterize M2 to be causal while remaining sub-quadratic. Using this parameterization, M2 matches GPT-style Transformers at 360M parameters in pretraining perplexity on The PILE—showing for the first time that it may be possible to match Transformer quality without attention or MLPs.<sup>1</sup>
Machine learning models in natural language processing and computer vision are being stretched to longer sequences and higher-dimensional representations to enable longer context and higher quality, respectively [6, 10, 62, 84]. However, existing architectures exhibit time and space complexities that grow quadratically in sequence length and/or model dimension—which limits context length and makes scaling expensive. For example, attention and MLP in Transformers scale quadratically in sequence length and model dimension [15]. In this paper, we explore a natural question: can we find a performant architecture that is sub-quadratic in both sequence length and model dimension?
$^{1}$ Code is available at https://github.com/HazyResearch/m2.
Figure 1: Monarch matrices are a simple, expressive, and hardware-efficient class of sub-quadratic structured matrices. MONARCH MIXER (M2) uses Monarch matrices to mix inputs first along the sequence dimension and then along the model dimension. See the Appendix for PyTorch implementation of an M2 layer.
In our exploration, we seek a sub-quadratic primitive for both the sequence length and model dimension. Our framing takes inspiration from work such as MLP-mixer [73] and ConvMixer [74], which observed that many machine learning models operate by repeatedly mixing information along the sequence and model dimension axes, and used a single operator for both axes. Finding mixing operators that are expressive, sub-quadratic, and hardware-efficient is challenging. For example, the MLPs in MLP-mixer and convolutions in ConvMixer are expressive, but they both scale quadratically in their input dimension [73, 74]. Several recent studies have proposed sub-quadratic sequence mixing with long convolutions or state space models [27, 64, 77]—both computed using the FFT—but these models have poor FLOP utilization (3-5% [28]) and maintain quadratic scaling in model dimension. Meanwhile, there has been promising work in sparsifying dense MLP layers without losing quality, but some of the models can actually be slower than their dense counterparts, due to low hardware utilization [7, 8, 14, 26, 35].
We turn to an expressive class of sub-quadratic structured matrices called Monarch matrices [14] (Figure 1 left) to propose MONARCH MIXER (M2). Monarch matrices are a family of structured matrices that generalize the fast Fourier transform (FFT) and have been shown to capture a wide class of linear transforms including Hadamard transforms, Toeplitz matrices [32], AFDF matrices [57], and convolutions. They are parameterized as the products of block-diagonal matrices, called monarch factors, interleaved with permutation. Their compute scales sub-quadratically: setting the number of factors to  $p$  results in computational complexity of  $O(pN^{(p + 1) / p})$  in input length  $N$ , allowing the complexity to interpolate between  $O(N\log N)$  at  $p = \log N$  and  $O(N^{3 / 2})$  at  $p = 2$ .<sup>2</sup>
M2 uses Monarch matrices to mix information along the sequence and model dimension axes. It is both simple to implement and hardware-efficient: the block-diagonal Monarch factors can be computed efficiently on modern hardware using GEMMs (generalized matrix multiply algorithms). Our proof-of-concept implementation of an M2 layer, written in less than 40 lines of pure PyTorch (including imports), relies only on matrix multiplication, transpose, reshape, and elementwise products (see pseudocode in Figure 1 middle) and achieves  $25.6\%$  FLOP utilization<sup>3</sup> for inputs of size 64K on an A100 GPU. On newer architectures such as the RTX 4090, a simple CUDA implementation achieves  $41.4\%$  FLOP utilization at the same size.
Non-Causal Settings As a first proof of concept of M2, we evaluate how it compares to Transformers in terms of speed and quality in non-causal settings such as BERT-style masked language modeling [21] and ImageNet classification. We introduce M2-BERT, which replaces the attention blocks in BERT with bidirectional gated convolutions implemented using Monarch matrices and replaces the dense matrices in the MLP with Monarch matrices. M2-BERT reduces parameter count but maintains quality—matching BERT-base and BERT-large in downstream GLUE quality with  $27\%$  and  $24\%$  fewer parameters, respectively. Sub-quadratic scaling in sequence length enables high throughput at longer sequences—up to  $9.1\times$  higher throughput at sequence length 4K than
2Monarch matrices were originally [14] parameterized with  $p = 2$ , but the general  $p$  case is a natural extension.
For context, the most optimized attention implementations achieve  $25\%$  FLOP utilization, while unoptimized implementations of attention can have as low as  $10\%$  FLOP utilization [15].
HuggingFace BERT, and  $3.1 \times$  higher throughput at sequence length 8K than BERT optimized with FlashAttention [15].
For image classification, we adapt HyenaViT-b [64], an attention-free vision transformer based on gated convolutions. We replace the convolution operation with M2 primitives and replace the MLP layers with an M2 block as well. These changes reduce the parameter count compared to a ViT-b [22] model with the same model width and depth by a factor of 2. Surprisingly, despite this parameter reduction, we find that M2 slightly outperforms ViT-b and HyenaViT-b baselines, achieving  $1\%$  higher accuracy on ImageNet [18].
Causal Settings Causal settings such as GPT-style [65] auto-regressive language modeling present a technical challenge: masking out the upper triangular elements in an attention matrix (or equivalent structure) introduces a quadratic bottleneck. To alleviate this quadratic bottleneck with Monarch matrices, we develop new theory to characterize which parameterizations of Monarch matrices maintain causality. To do so, we take a view of  $p$ -order Monarch matrix multiplication as  $p$ -variate polynomial evaluation and interpolation (e.g.,  $p = 2$  factors corresponds to bivariate polynomials, Figure 2 left). Using this view, we show that the M2 convolution shown in Figure 1 (middle) can be viewed as manipulation of modular polynomial multiplication. This result allows us to develop conditions (Theorem 3) under which M2 is causal. We can use this causal parameterization to outperform GPT-style language models on causal language modeling by 0.2 PPL points on the PILE at model size 360M—without using either attention or MLP blocks.
Summary Overall, our results present a potential path to building machine learning models with sub-quadratic primitives. We hope our work can serve as a starting point to explore models that are more efficient in both sequence length and model dimension.
In this section, we provide some background on the key components behind the cost of operations on GPUs, and then discuss the scaling characteristics of some common primitives used to mix information across the sequence dimension and model dimension in modern machine learning models.
GPU Accelerator Cost Model We provide a brief discussion of relevant factors affecting runtime performance of deep learning operations on GPUs. Depending on the balance of computation and memory accesses, operations can be classified as either compute-bound or memory-bound [44]. In compute-bound operations, the time accessing GPU memory is relatively small compared to the time spent doing arithmetic operations. Typical examples are matrix multiply with large inner dimension, and short convolution kernels with a large number of channels.
The speed of these operations is determined by the FLOP/s available on compute units, and the number of FLOPs necessary to complete the operation. In our paper, we exploit fast matrix multiply units such as tensor cores. On the A100, tensor cores can achieve 312 TFLOP/s in half-precision matrix multiply operations, while non-matrix multiply operations are limited to 19 TFLOP/s [59]. This trend began with tensor cores in the V100 [58], and is continuing into the next-generation H100 chips [60].
In memory-bound operations, the time taken by the operation is determined by the number of memory accesses, while time spent in computation is much smaller. Examples include most elementwise operations (e.g., activation, dropout) and reductions (e.g., sum, softmax, batch norm, layer norm).
The runtime of memory-bound operations is determined by the memory bandwidth of different layers of the memory hierarchy. GPU memory is large but relatively slow—up to 80 GB on A100, but with bandwidth of 2 TB/s [59]. Higher levels of the memory hierarchy such as caches are much smaller (20 MB) but an order of magnitude faster (1
Table 1: FLOP cost and utilization of various mixer layers, input dimension 64K on an RTX 4090.
Layer	FLOP Cost	Util
MLP	N2	95.5%
FlashAttn	N2	24.0%
FFT	N log N	3.0%
M2 Conv	N3/2	41.4%
Common Mixer Primitives To help contextualize our work, we provide scaling and hardware utilization characteristics for a few common operations that are used to mix information in machine learning models, summarized in Table 1.
Transformers [75] use attention to mix information across the sequence dimension, and MLP blocks to mix information across the model dimension. Both of these blocks scale quadratically in input length. MLP layers are compute-bound, so they have high FLOP utilization out of the box. Attention blocks are memory-bound, so even the most optimized implementations such as FLASHATTENTION [15] have relatively lower FLOP utilization.
Recent work has made progress towards attention-free models by replacing attention layers with long convolution layers, interleaved with elementwise gating [27, 28, 36, 54, 64, 67-69]. These layers are computed using FFT operations using the FFT convolution theorem:  $y = \mathbf{K} * \mathbf{X} = FFT^{-1}(FFT(\mathbf{X}) * FFT(\mathbf{K}))$ . While the FFT scales asymptotically well in  $O(N \log N)$ , it is often memory-bound and thus has low FLOP utilization. In our work, we aim to construct a mixer that has both sub-quadratic scaling and high FLOP utilization.
3 MONARCH MIXER
In this section, we recall Monarch matrices, introduce how M2 uses Monarch matrices to mix along the sequence and model dimensions, and benchmark a M2 convolution in terms of hardware utilization.
3.1 Monarch Matrices
Monarch matrices [14] are a sub-quadratic class of structured matrices that are hardware-efficient and expressive. They can represent many linear transforms, including convolutions, Toeplitz-like transforms, low-displacement rank transforms, and orthogonal polynomials. Directly implementing these different structured transforms on GPUs as dense matrices can be inefficient. In contrast, their Monarch decompositions can be computed by interleaving matrix multiplications with tensor permutations.
A Monarch matrix  $\mathbf{M} \in \mathbb{R}^{N \times N}$  of order  $p$  is defined by the following:
$$
\mathbf {M} = \left(\prod_ {i = 1} ^ {p} \mathbf {P} _ {i} \mathbf {B} _ {i}\right) \mathbf {P} _ {0}, \tag {1}
$$
where each  $\mathbf{P}_i$  is related to the 'base  $\sqrt[p]{N}$ ' variant of the bit-reversal permutation, and  $\mathbf{B}_i$  is a block-diagonal matrix with block size  $b$ . Setting  $b = \sqrt[p]{N}$  achieves sub-quadratic compute cost. For example, for  $p = 2, b = \sqrt{N}$ , Monarch matrices require  $O(N^{3/2})$  compute in sequence length  $N$ .
In this paper, we use Monarch matrices to construct architectures that are sub-quadratic in both sequence length  $N$  and model dimension  $d$ . We will often parameterize order-2 Monarch matrices, written as  $\mathbf{M} = \mathbf{P}\mathbf{L}\mathbf{P}\mathbf{R}\mathbf{P}$ , where  $\mathbf{L}$  and  $\mathbf{R}$  are block-diagonal matrices (for "left" and "right"), and  $\mathbf{P} = \mathbf{P}_2 = \mathbf{P}_1 = \mathbf{P}_0$  is a permutation that reshapes the input to 2D, transposes it, and flattens it to 1D. A common case is to set  $\mathbf{L} = \mathbf{R} = (\mathbf{I}_{\sqrt{N}} \otimes \mathbf{F}_{\sqrt{N}})$ , where  $\mathbf{F}_{\sqrt{N}}$  is a  $\sqrt{N}$  DFT matrix, and  $\otimes$  is the Kronecker product.
3.2 MONARCH MIXER Architecture
We describe how MONARCH MIXER uses Monarch matrices and elementwise operations to construct sub-quadratic architectures (Figure 1 middle). We take a mixer view of model architectures, where each layer is a sequence of mixing operations across the sequence and the model dimension axes. Each layer takes as input a sequence of embeddings  $\mathbf{X} \in \mathbb{R}^{N \times d}$ , and outputs a sequence  $\mathbf{Y} \in \mathbb{R}^{N \times d}$  where  $N$  is the sequence length, and  $d$  is the model dimension. For simplicity, we show the order-2 case here, though we can use higher-order blocks to scale to longer sequences and larger model dimensions.
Let  $\mathbf{M}_1, \mathbf{M}_2 \in \mathbb{R}^{N \times N}$  and  $\mathbf{M}_3, \mathbf{M}_4 \in \mathbb{R}^{d \times d}$  be order-2 Monarch matrices, let  $\mathbf{K}_1 \in \mathbb{R}^{N \times d}$ , let  $\sigma$  be an optional point-wise non-linearity (e.g. ReLU), and let  $\odot$  be elementwise multiplication. M2
Table 2: FLOP cost and utilization of M2 compared to dense MLP at different input sizes  $N$ , with block size  $\sqrt{N}$ , on an A100 and RTX 4090.
N	4K	16K	64K	256K
Dense Matmul TFLOP Cost	0.025	0.412	6.60	106.0
M2 TFLOP Cost	0.002	0.013	0.103	0.824
Dense FLOP Utilization (A100)	63.0%	78.0%	80.0%	OOM
M2 FLOP Utilization (A100)	4.78%	12.7%	25.6%	42.8%
Wall-Clock Speedup (A100)	1.2×	5.1×	20.6×	>55.0×
Dense FLOP Utilization (4090)	74.6%	96.7%	98.0%	OOM
M2 FLOP Utilization (4090)	11.1%	32.1%	41.4%	53.7%
Wall-Clock Speedup (4090)	2.2×	10.5×	27.0×	>69.1×
uses Monarch matrices to construct expressive architectures. For example, a convolutional block with a sparse MLP can be expressed as follows:
1. Mix along sequence axis:
$$
\tilde {\mathbf {X}} = \mathbf {M} _ {2} \left(\mathbf {K} _ {1} \odot \mathbf {M} _ {1} \mathbf {X}\right) \tag {2}
$$
2. Mix along embedding axis:
$$
\mathbf {Y} ^ {\top} = \mathbf {M} _ {4} \sigma \left(\mathbf {M} _ {3} \tilde {\mathbf {X}} ^ {\top}\right) \tag {3}
$$
When  $\mathbf{M}_1$  is set to the DFT and  $\mathbf{M}_2$  is set to the inverse DFT, Equation 2 exactly corresponds to a convolution with kernel  $\mathbf{K}_1$  parameterized in frequency space. Equation 3 corresponds to an MLP with the dense matrices replaced by Monarch matrices. More expressive layers are also easily expressible; for example, replacing Equation 2 with  $\mathbf{V} \odot \mathbf{M}_2(\mathbf{K}_1 \odot \mathbf{M}_1(\mathbf{Q} \odot \mathbf{K}))$ , where  $\mathbf{Q}, \mathbf{K}, \mathbf{V}$  are linear projections of  $\mathbf{X}$ , reproduces a gated convolution block, as in [27, 28, 64].
The basic M2 layer is simple to implement; pseudocode is shown in Figure 1 (middle), and the Appendix gives an efficient implementation of M2 in under 40 lines of pure PyTorch (including imports). The convolution case with Monarch matrices fixed to DFT and inverse DFT matrices also admits implementations based on FFT algorithms [11].
3.3 Architecture Benchmarks
We benchmark the efficiency of the  $\mathbf{M}(\mathbf{K} \odot \mathbf{MX})$  convolution operator (Equation 2) implemented in a simple CUDA kernel (calling standard cuBLAS sub-routines [61]), as the dimension  $N$  increases. Equation 3 scales similarly, as dimension  $d$  increases. We keep the block size  $b$  fixed to  $\sqrt{N}$ .
Table 2 shows the FLOP cost and utilization of a M2 operator as a function of the input size on an A100 as well as on an RTX 4090. On the A100, the operator is more dominated by the data movement costs of the permutation operations (see the Appendix for a roofline analysis). For longer inputs, the sub-quadratic scaling allows MONARCH MIXER to outperform dense matrix multiplication. On the RTX 4090, which has a larger and faster L2 cache than the A100, we can manually optimize an implementation to amortize data movement costs.
4 Theoretical Analysis: M2 as Polynomial Multiplication
In this section, we develop theory to make the M2 layer causal in the input  $\mathbf{X}$  -e.g., ensure that an output  $Y_{i}$  of the M2 should only depend on  $X_{1},\ldots ,X_{i}$ . Our approach involves interpreting Monarch matrix multiplication as multivariate polynomial evaluation and interpolation. We then show that an M2 convolution is equivalent to modular polynomial manipulation in a univariate basis.
The challenge is controlling the degrees of the resulting univariate polynomials, to prevent "underflow" under modular multiplication (see Figure 2 for an overview). Our key result is deriving sufficient conditions on the degrees of the bivariate polynomials defining the Monarch factors to prevent such underflow. We focus on the bivariate case (order  $p = 2$ ) in the body, and give the general multivariate case in the Appendix. We present proof sketches in the main body, and leave proofs and additional results for the Appendix.
Figure 2: Monarch multiplication can be interpreted as polynomial evaluation and interpolation. We derive sufficient conditions on the polynomial formulation of Monarch matrices for M2 to be causal.
Monarch Multiplication as Polynomial Evaluation First, we show that order-2 Monarch matrix-vector multiplication  $\mathbf{M} \cdot \mathbf{u}$  is equivalent to bivariate polynomial evaluation.
Fix a Monarch matrix  $\mathbf{M} \in \mathbb{R}^{N \times N} = \mathbf{PLPRP}$ , for two block-diagonal matrices  $\mathbf{L}$  and  $\mathbf{R}$  with blocks of size  $b = \sqrt{N}$ . We can interpret Monarch matrices as bivariate polynomial evaluation by setting  $A = \{\omega_0, \dots, \omega_{b-1}\}$  as a set of evaluation points (e.g., the  $b$ th roots of unity), and letting  $\{\ell_0(X,Y), \dots, \ell_{b-1}(X,Y)\}, \{r_0(Y), \dots, r_{N-1}(Y)\}$  be sets of basis polynomials with individual degrees of  $X, Y$  being  $< \sqrt{N}$ . The values of  $\{\ell_0(X,Y), \dots, \ell_{b-1}(X,Y)\}$  evaluated on  $A^2$  determine the entries of  $\mathbf{L}$ , and the values of  $\{r_0(Y), \dots, r_{N-1}(Y)\}$  evaluated on  $A$  determine the entries of  $\mathbf{R}$ . We give the mapping from  $\ell, r$ , and  $A$  to  $\mathbf{L}$  and  $\mathbf{R}$  in the Appendix.
Then, matrix-vector multiplication between  $\mathbf{M}$  and a vector  $\mathbf{u}$  is equivalent to polynomial evaluation of the basis functions  $\ell, r$  on the evaluation points  $A^2$ :
Theorem 1. Let  $m(j) = j \mod \sqrt{N}$ . For any vector  $\mathbf{u} \in \mathbb{R}^N$ ,  $\mathbf{M}\mathbf{u}$  is a bivariate polynomial  $u(X,Y)$  evaluated at  $A^2$ , with  $u(X,Y) = \sum_{j=0}^{N-1} u_j f_j(X,Y)$ , where  $f_j(X,Y) = \ell_{m(j)}(X,Y)r_j(Y)$ .
Monarch Inverse as Polynomial Interpolation Next, we exploit the fact that Monarch inverse multiplication  $\mathbf{M}^{-1} \cdot \mathbf{u}$  is equivalent to polynomial interpolation in the basis polynomials of  $\mathbf{M}$ .
Theorem 2. Let  $\mathbf{M}_0, \mathbf{M}_1, \mathbf{M}_2$  be Monarch matrices, and let  $A$  be the set of  $\sqrt{N}$  roots of unity. Then, the operation
$$
\mathbf {f} = \mathbf {M} _ {0} ^ {- 1} \left(\left(\mathbf {M} _ {1} \mathbf {k}\right) \odot \left(\mathbf {M} _ {2} \mathbf {u}\right)\right). \tag {4}
$$
is equivalent to representing the polynomial
$$
h (X, Y) = k (X, Y) u (X, Y) \mod (X ^ {\sqrt {N}} - 1, Y ^ {\sqrt {N}} - 1)
$$
in terms of the basis polynomials  $\ell, r$  corresponding to  $\mathbf{M}_0$ , and where  $k(X,Y)$  and  $u(X,Y)$  are the polynomials corresponding to  $\mathbf{M}_1\mathbf{k}$  and  $\mathbf{M}_2\mathbf{u}$ , respectively.
The above follows from Theorem 1 and the fact that Monarch matrix-vector multiplication with an inverse Monarch matrix is equivalent to polynomial interpolation in a given basis. The mod part comes from the fact that  $A$  is the set of roots of the polynomial  $Z^{\sqrt{N}} - 1$ .
Causal Monarch Maps Now, we give a class of Monarch matrices from which we can build a causal map. First, we define a polynomial with minimum degree  $j$ :
Definition 1. A polynomial of minimum degree  $j$  (and maximum degree  $N - 1$ ) is defined as  $\bar{q}_j(Z) = \sum_{a = j}^{N - 1}\bar{q}_j[a]Z^a$ .
To ensure causality, we first convert the bivariate polynomial basis into a univariate basis, and then we expand the degree of the univariate polynomial. The resulting univariate polynomial multiplication is naturally causal (exploiting similar properties as the causal FFT convolution).
We use the Kronecker substitution  $(X\gets Z,Y\gets Z^{\sqrt{N}})$  to convert the bivariate polynomial basis into a univariate basis:
$$
q _ {j} (Z) = \ell_ {m (j)} (Z) r _ {j} \left(Z ^ {\sqrt {N}}\right), \tag {5}
$$
where  $m(j)$  is defined as in Theorem 1.
Then, the following class of Monarch matrices (with the conversion to univariate polynomial basis as above) forms a causal map:
Theorem 3. Let  $\mathbf{u},\mathbf{k}\in \mathbb{R}^n$ , where  $n < N / 2$ . Let  $m(j)$  be as in Theorem 1, and  $k(j) = \left\lfloor j / \sqrt{N}\right\rfloor$ . Then define the basis polynomials  $\ell_{m(j)}$  to have minimum degree  $m(j)$ , basis polynomials  $r_j$  to have minimum degree  $k(j)$ , and all polynomials  $q_j(Z)$  to have maximum degree  $< N / 2$  for all  $j < N / 2$  and for  $N / 2\leq j < N$  have maximum degree  $N - 1$ . Let  $\mathbf{M}_N$  be defined by such basis polynomials via (5) where the evaluation points are now the  $N$ th roots of unity. Then, we have that
$$
\mathbf {u} \mapsto \left(\mathbf {M} _ {N} ^ {- 1} \left(\mathbf {M} _ {N} (\mathbf {k}, \mathbf {0} _ {N - n}) \odot \mathbf {M} _ {N} (\mathbf {u}, \mathbf {0} _ {N - n})\right)\right) [ 0: n - 1 ] \tag {6}
$$
gives a causal map in  $\mathbf{u}$ .