Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture
Authors: Daniel Y Fu
Simran Arora
Jessica Grogan
Isys Johnson
Sabri Eyuboglu
Armin W Thomas
Benjamin Frederick Spector
Michael Poli
Atri Rudra
Christopher Re
Venue: NeurIPS 2023 oral (2023)
Keywords: structured matrices
transformers
efficiency
TL;DR: We present Monarch Mixer, a new simple architecture that scales sub-quadratically in sequence length and model dimension.

Abstract: 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.

MONARCH MIXER: A Simple Sub-Quadratic GEMM-Based Architecture
Jessscaa Daniel Y. Fu', Simran Arora Jessica Grogan" x:2 Isys Johnson 2 Sabri Eyuboglu":
Armin W. Thomas" Benjamin Spector', Michael Poli', Atri Rudra", Christopher Ré
"Equal Contribution. Department of Computer Science, Stanford University.
Department of Computer Science and Engineering. University at Buffalo, SUNY. *Department of Psychology, Stanford University.
danfuêcs fjrgrogan, isysj ohn Obuffall edu, feyuboglu, athms, bfs @stanford edu, atri@buffalo. edu, chrismreCcs. stanf www.elecssssaafoooooooooooeooooe
simaroradstanfor edu,
edu
Abstract
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 ex- pressive structured matrices that captures many linear transforms, achieves high hardware efficiency GPUs, and scales sub-quadratically. As proof of concept, we explore the performance of M2 in three domains: non-causal BERT-style lan- guage modeling, ViT-style image classification, and causal GPT-style language modeling. For non- non-causal BERT-style modeling, M2 matches BERT-base and BERT-large in downstream GLUE quality with up 27% fewer parameters, and achieves up to 9.1 higher throughput at sequence length 4K. On ImageNet, M2 outerforrms 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 novel theoretical view of Monarch matrices based on multivariate polynomial evaluation and interpolation, which lets us parameterize M2 to 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 match Transformer quality without attention MLPs.
Introduction
Machine learning models in natural language processing and computer vision are being stretched to longer sequences and higher-dimenssonall 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-whhe 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 performant architecture that is sub sub-quadratic in both sequence length and model dimension?
limits context length and
'Code available at https://giithmn
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

P
I
def M2_ Layer(X):
M
P. P
# mix sequence Z = M @ (k (M X))
1=1
# mix channels
N
Subquadratic: O(pNe
Y M @ a(M Z.T))
Hardware-Efficcen (GEMMs)
Expressive (generalizes FFT)
return
Efficient Mixing on Sequence, Dimension
Order-p Monarch Matrices
Simple Layers
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 Py PyTorch implementation of M2 layer.
Figure I: Monarch matrices are a simple, expressive, and
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 single operator for both axes. Finding mixing operators that expressive, qua-qquaratic, and hardware-efficien is challenging. For example, the MLPs in MLP-mixer and convolutions in Conv Mixer 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- these models have poor FLOP utilization (3-5% [28]) and maintain quadratic scaling in model dimension. Meanwhile, there has been promising work sparsifying dense MLP layers without losing quality, but some of the models can actually be slower than their dense counterparts, due low hardware utilization [7, 8, 14, 26, 35].
We turn to expressive class of sub-quadratic structured matrices called Monarch matrices [14] (Figure 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 wide class of linear transforms including Hadamard transforms, Toeplitz matrices [32], AFDF matrices [57], and convolutions. They are parameterized 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 (ppN(p+)) in input length N, allowing the complexity to interpolate between ( log N) at p = log N and O(N3/2) at = 2.2
M2 uses Monarch matrices to mix information along the sequence and model dimension axes. It both simple implement and harfware-efficientt the block-diagonal Monarch factors can be computed efficiently on modern hardware using GEMMs (generalized matrix multiply algorithms). Our proof-of-concep implementation of an M2 layer, written less than 40 lines of pure Py Torcch (including imports), relies only on matrix multiplication, transpose, reshape, and elementwise products (see pseudocode in Figure middle) and achieves 25.6% FLOP utilization' for inputs of size 64K on an A 100 GPU. On newer architectures such as the RTX 4090, simple CUDA implementation achieves 41 .4% FLOP utilization at the same size.
Non-Causal Settings As first proof of concept of M2, we evaluate how it compares to Trans- formers 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 BERT-base and BERT-large in downstream GLUB quality with
but maintains guality-matching 27% and 24% fewer parameters, respectively. Sub-quadratic scaling in sequence length enables high throughput at longer sequences-uu up to 9.1x higher throughput at sequence length 4K than
Monarch matrices were originally [14] parameterized with p 2. but the general p case is natural extension.
For context, the most optimized attention inplementations achieve 25% FLOP utilization, while unoptimized inplementations of attention have low 10% FLOP utilization [15].

---

HuggingFace BERT, and 3.1 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 M2 block as well. These changes reduce the parameter count compared to ViT-b [22] model with the same model width and depth by a factor of Surprisingly. despite this parameter reduction, find that M2 slightly outperforms ViT-b and Hyena ViT-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 quadratic bottleneck. To alleviate this quadratic bottleneck with Monarch matrices, we develop new theory to characterize which parameterizations of Monarch matrices maintain causality. 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 left). Using this view, we show that the M2 convolution shown Figure (middle) can be viewed as manipulation of modular polynomial multiplication. This result allows 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 the PILE at model size withhut using either attention or MLP blocks.
Summary Overall, our results present a potential path building machine learning models with sub-quadratic primitives. We hope our work can serve as starting point to explore models that are more efficient both sequence length and model dimension.
Preliminaries
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 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 the balance of computation and memory accesses, operations can be classified as either compute-boure bound or memory bound [44]. In compute-bound operations, the time accessing GPU memory relatively small compared to the time spent doing arithmetic operations. Typical examples are matrix multiply with large inner dimension, and short convolution kemelss with 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 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- next-generation H100 chips [60].
In memory-boune operations, the time taken by the oper- ation is determined by the number of memory accesses, while time spent in computation is much smaller. Exam- ples include most clementwise operations (c.g., activation, dropout) and reductions (e.g., sum, softmax, batch norm, layer norm).
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%
The runtime of memory bouud operations determined by the memory bandwidth of different layers of the memory hierarchy. GPU memory 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 (19 TB/s).
FlashAut
N
24.0%
FFT
N Vog N
3.0%
M2 Conv
3/2
41.4%
3

---

Common Mixer Primitives To help contextualize work, 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 memory-bound,
even the most optimized implementations such as FLASHATTENTION [15] have relatively lower FLOP utilization.
Recent work has made progress towards attention-fre: 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 K * X = (FF * FFT( While the FFT scales asymptotically well in O log N), it is often memory memory-bound and thus has low FLOP utilization. In our work, we aim to construct 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 M2 convolution in terms of hardware utilization.
3.1 Monarch Matrices
Monarch matrices [14] are sub-quadratic class of structured matrices that serdwarre-fficieen and expressive. They can represent many linear transforms, including convolutions, Toeplitz-like transforms, low-displacemmen 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 M IR NxNN of order p is defined by the following:
P
M = [[P.B.
Po.
where each P, is related to the "base V variant of the bit-reversal permutation, and B; is block-diagonal matrix with block size b. Setting V/N achieves sub-quadratic compute cost. For example, for p 2, VN, Monarch matrices require O(N3/2) compute in sequence length N.
In this paper, we use Monarch matrices construct architectures that are sub-quadratic in both sequence length N and model dimension d. We will often parameterize order-2 Monarch matrices, written M PLPRP, where and R block-diagone matrices (for "left" and "right"), and P= P2 P1 = Po is permutation that reshapes the input to 2D, transposes it, and flattens it to ID. A common case is to set L R (I the Kronecker product.
X F ) where F VN is VN DFT matrix, and @ is
3.2 MONARCH MIXER Architecture
We describe how MONARCH MIXER uses Monarch matrices and clementwise operations to construct sub-quadratic architectures (Figure middle). We take a mixer view of model architectures, where each layer is sequence of mixing operations across the sequence and the model dimension axes. Each layer takes as input a sequence embeddings X E N and outputs a sequence E R.Nxx 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 seale to longer sequences and larger model dimensions.
Let M, M2 E N and M3, M. E Rddd be order-2 Monarch matrices, K₁ Nxd let be an optional point- wise non-linearity (e.g. ReLU), and let O be elementwise multiplication. M2

---

Table 2: FLOP cost and utilization of M2 compared to dense MLP at different input sizes N, with block size VN, 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.2x
5.1x
20.6x
>55.0x
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.2x
10.5x
27.0x
>69.1
uses Monarch matrices construct expressive architectures. For example, convolutional block with a sparse MLP can be expressed as follows:
1. Mix along sequence axis:
X M2(K, M,X)
2. Mix along embedding axis:
Y M
When M1 is set the DFT and M2 is set the inverse DFT, Equation 2 exactly corresponds to convolution with kernel Ki parameterized in frequency space. Equation corresponds to an MLP with the dense matrices replaced by Monarch matrices. More expressive layers also easily expressible; for example, replacing Equation 2 with V @ M2(Ki @ M, @ K)), where Q, K, V linear projections of X, reproduces a gated convolution block, as [27, 28, 64].
The basic M2 layer is simple to implement; pseudocode is shown in Figure (middle), and the Appendix gives 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 M(K @ MX) convolution operator (Equation 2) implemented in simple CUDA kernel (calling standard cuBLAS sub-routines [61]], the dimension N increases. Equation 3 scales similarly, as dimension d increases. We keep the block size fixed to vN.
Table 2 shows the FLOP cost and utilization of M2 operator function of the input size on an A100 as well as an RTX 4090. On the A100, the operator is more dominated by the data movement costs of the permutation operations (see the Appendix for 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 1.2 cache than the A 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 X-c.g., ensure that an output Y; of the M2 should only depend on X1, ... X₁. 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
bivariate polynomials defining the Monarch factors to prevent such underflow. We focus on the bivariate case (order = 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.
5

---

Causal Parameterization
P
P
P
M'(MuoMK)
e(x, Y)
r(Y)
f(Z) g(Z) mod Z
Conditions on Bivariate
Univariate Multiplication deg(f), deg(g) < <
Polynomial Degrees
Causal Map
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 be causal.
Monarch Multiplication as Polynomial Evaluation First, we show that order-2 Monarch matrix- vector multiplication M is equivalent to bivariate polynomial evaluation.
Fix Monarch matrix M E RNXN = PLPRP, for two block-diagonal matrices L and with blocks of size b VN. We can interpret Monarch matrices as bivariate polynomial evaluation by setting A {wo..... W6-1) } as a set of evaluation points (e.g.. the bth roots of unity), and letting })}. (ro(Y),...............,, ~N-1(Y) (Y)} be sets of basis polynomials with individual degrees of X, Y being < VN. The values (eu(X,Y), .)..................... (X,, Y)} evaluated on A² deter- mine the entries of L, and the values of (ro(Y)...- PN- (Y)) evaluated A determine the entries of R. We give the mapping from r, and A to and R. in the Appendix.
{Lo(X, (o.,X,)
Then, matrix-vector multiplication between M and vector u is equivalent to polynomial evaluation of the basis functions l, r on the evaluation points A
Theorem 1. Let m(j) = j mod VN. For any vector u E RN,, Mu is a bivariate poly- nomial u( u(X,Y) evaluated at A with w(X,Y) N-1 (X);(X,Y), where Jj(XY) (X,) (Y)
Monarch Inverse as Polynomial Interpolation Next, we exploit the fact that Monarch inverse multiplication M-1.. is equivalent to polynomial interpolation in the basis polynomials of M.
Theorem 2. Let Mo, M1, M2 be Monarch matrices, and let be the of / roots of unity. Then, the operation
f M= ((Mik) (M2u)).
is equivalent to representing the polynomial
h(X,Y) = k(X,
mod - -
in terms of the basis polynomials É, r corresponding to Mo. and where k(X, Y) and u(X, Y) are the polynomials corresponding M,k and Mzu, respectively.
The above follows from Theorem and the fact that Monarch matrix-vector multiplication with an inverse Monarch matrix is equivalent polynomial interpolation in given basis. The mod part comes from the fact that A is the set of roots of the polynomial ZvN
Causal Monarch Maps Now, we give class of Monarch matrices from which we can build 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 (Z) =N EN' 7"
To ensure causality, we first convert the bivariate polynomial basis into univariate basis, and then expand the degree 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 Z,Y + ZVN to convert the bivariate polynomial basis
into a univariate basis:
(2) ( (2)r (ZVN),
(5)

---

where m(j) is defined as in Theorem 1.
Then, the following class Monarch matrices (with the conversion to univariate polynomial basis as above) forms a causal map:
Theorem 3. Let u R", where < N22. Let m(j) be as in Theorem 1, and k(j) = j/VN
Then define the basis polynomials m)()) to have minimum degree m(j), basis polynomials r; to have minimum degree k(j), and all polynomials qj (Z) to have maximum degree < 2 all j 2/2 Let My be defined by such basis polynomials
and N S N have maximum degree N via (5) where the evaluation points are now the Nih roots of unity. Then, we have that
u u (M,' (M Ox @ Mx Ox Ox-n))) [0
(6)
gives a causal map in u.
Theorem 3 gives causal map that can be computed entirely using Monarch matrices enforeing causality with sub-quadratic scaling. The main technical ingredient in proving the above result is that the product qj (Z)qz (Z) can be written
linear combination of (Z) for + j' S < N (this uses the above specified properties on the minimum and maximum degrees of (2)2). This in turn implies that the term kyrugu (Z)gx(() only contributes the coefficients of "higher order" basis polynomials ga (Z) for a 2 j + + in the product k(Z )u(2). 2). which is needed for causality. Figure 2 gives an example of restricted polynomials generating causal map.
Experiments
We compare MONARCH MIXER to Transformers three tasks where Transformers have been dominant: BERT- style non-causal masked language modeling. ViT-style image classification, and GPT-style causal language mod- match Transformers
Monarch
eling. In each, we show that in quality using neither attention nor MLPs. We addition- ally evaluate wall-clock speedups against strong Trans- former baselines in the BERT setting. Additional exper- iments on speech and alternative architectures are given in Appendix B, and experimental details are given in Ap- pendix C.
Monarch long conv
GeLU
short conv
short conv
short
GLU
5.1 Non-Causal Language Modeling
Sequence Mixer
Dimension Mixer
We introduce M2-BERT, an M2-based architecture for non-causal language modeling. M2-BERT acts as a drop- in replacement for BERT-style language models [21], which are workhorse application of the Transformer architecture [1. 39, 40, 45, 48, 49, 52, 56, 85, 89]. We train M2-BERT using masked language modeling over C4 [66] with the sertt-bas-uuccaaae tokenizer.
Figure 3: M2-BERT uses Monarch ma-
trices to create a bidirectional gated long convolution in the sequence mixer, and uses Monarch matrices to replace the lin- ear layers in the dimension mixer.
M2-BERT starts with Transformer backbone and replaces the attention and MLPs with M2 layers, shown in Figure 3. In the sequence mixer, we replace attention with bidirectional gated convolutions with residual convolution (Figure 3 left). To recover convolutions, we the Monarch matrices to DFT and inverse DFT matrices. Following [27, 64], we also add short depthwise convolutions after the projections. In the dimension mixer, we replace the two dense matrices in MLPs with learned block-diagone matrices (Monarch matrix of order 1, b = 4). We pretrain two M2- 12-BEET--base models, 80M and 11M, and two M2-BERT-large models, at 260M and 341M. These are equivalent to BERT-base and BERT-large, respectively.
Downstream GLUE Scores First, we evaluate M2-BERT models on downstream fine-tuning compared to BERT-base and BERT-lagee from [20]. We take the pretrained models and fine-tune them on BERT, following the procedure in [38]. Table shows performance for BERT-base equivalent models, and Table 4 shows performance for BERT-large equivalent models. M2-BERT-bbas can
7

---

Table 3: Average GLUE Score for M2-BERT-base compared to BERT-base [20], along with change
in parameters and GLUE score.
Model
GLUE Score
A Params
A GLUE Score
BERT-base (110M)
79.6
-0%
+0.0
11-BBERT-bae (80M)
79.9
+0.3 -27%
12-BBRT-bare base (1 10M)
80.9
-0%
+1.3
Table 4: Average GLUE Score for M2-BERT-large compared BERT-large [20], along with change
in parameters and GLUE score.
Model
GLUE Score
A Params
GLUE Score
BERT-large (340M)
82.1
-0%
+0.0
M2-BERT- (260M)
82.2
-24%
+0.1
M2-BERRT-arge (341M)
82.8
+0.2%
+0.7
match BERT-base in GLUE quality with 27% fewer parameters- -r outperform BERT-base in quality 1.3 points when parameter matched. M2- M2-BERT-large matches BERT-large with 24% fewer parameters, and outperforms by 0.7 points when parameter matched.
GPU Throughput by Sequence Length Next, we evaluate throughput of M2-BERT models by sequence length, compared to HuggingFace implementations of BERT, as well as optimized implementations of BERT running FlashAttention [15]. Table 5 shows forward throughput for BERT- base equivalent models, and the appendix shows throughput for BERT-large (where the performance trends are similar). Inference times are reported in tokens/ms on an A100-40GB GPU. M2-BERT- base achieves higher throughput than even highly-optimized BERT models, and up to 9.1x faster throughput than standard HuggingFace implementation at sequence length 4K.
CPU Inference Latency Finally, we report CPU inference latency for M2-BERT-base (80M) compared to BERT-base, running direct PyTorch implementations for both. In short sequences, the impacts of data locality still dominate the FLOP reduction, and operations such filter generation (which are not present in BERT) pay a higher cost. Starting at sequences 1K and longer, M2-BERT- base starts to have speedup over BERT-base, up 6.5x at sequence length 8K. We believe further optimization and applying 10-awan principles can further improve CPU performance.
5.2 Image Classification
To validate that our methods generalize to images as well as language for non-causal modeling, we next evaluate M2 on image classification. We compare M2 to ViT-style models and recent work, HyenaViT-b [64], which uses gated long convolutions to replace the attention layers in ViT-b. In our work, M2-VIT builds off HyenaViT-b and replaces the long convolutions with the M2 operator in Equation 2 (again setting the Monarch matrices to the DFT and inverse DFT). We replace the MLP blocks in Hyena ViT-b with fock-ddagonnal matrices, similarly M2-BERT. Appendix additionally compares M2 to the Swin-family of architectures [50, 51].
Table 7 shows the performance of MONARCH MIXER against ViT-b, HyenaViT-b,, and ViT-b- Monarch (which replaces the MLP blocks of standard VIT-b with Monarch matrices) on ImageNet- 1k. MONARCH MIXER outperforms the other models with only half the parameters of the original ViT-s model. Surprisingly, MONARCH MIXER also outperforms ResNet- 152, with fewer parameters-eve though the latter was explicitly designed for ImageNet performance.
5.3 Causal Language Modeling
GPT-style causal language modeling is a critical application for Transformers [6, 31, 43). We introduce M2-GPT, M2-based architecture for causal language modeling. For the sequence mixer, M2-GPT combines the convolutional filter from Hyena [64]. the state-of-the-art aitention-free

---

Table 5: Throughput in tokens/ms by context length for M2-BERT- (80M) compared to BERT-
base
Model
512
1024
2048
4096
8192
HF BERT-base (110M)
206.
130.8
71.3
39.0
OOM
FlashAttention BERT-base (110M)
367.4
350.1
257.2
179.1
102.4
M2-BERT-base (80M)
386.3
380.7
378.9
353,9
320.1
M2 Speedup over HF BERT-base (110M)
1.9x
2.9x
5.2x
9.1x
Table 6: CPU inference latency in milliseconds with batch size at varied input sequence lengths. Measurements averaged over 10 examples a 48 vCPU, 96 GB RAM instance from the GCP 2t-standard-48 series, which runs Intel Cascade Lake processors. This is based on the protocol in [29].
Model
512
1024
2048
4096
8192
BERT-base (110M)
182
389
918
2660
11820
M2-BERT-base (80M)
289
361
651
948
1820
Speedup
0.6x
1.1x
1.4x
2.8
6.5x
language model, with parameter sharing across multiple heads from H3 [27]. We use the causal parameterization of Equation 2 to replace the FFT in these architectures, and remove the MLP layers entirely. The resulting architecture is entirely attention- and MLP-free.
We pretrain M2-GPT on the PILE, standard dataset for causal language modeling. Following prior work [28, 64J, we train models at two model sizes, with varying amounts training data-decaying the learning rate appropriately for each experiment. Table 8 shows the results. Even though our model attention- and MLP-free, it outperforms both Transformers and Hyena perplexity on pretraining. These results suggest that radically different architectures than Transformers may be performant causal language modeling.
6 Related Work
Long Convolutions Recent work proposes to use long convolution layers as a replacement for the Transformer attention layers in sequence modeling [28, 64, 67-69]. Many of these models rely the FFT convolution theorem to compute the long convolutions. We build on the insights in many of these architectures in constructing our M2 architectures, and additionally replaces the FFT operations with Monarch matrices.
Our work is also related to a rich literature in convolutions in other bases, such as Chebyshev bases [82] or orthogonal polynomial bases [34]. These approaches have analogues our multivariate analysis; replacing the basis polynomials of the Monarch matrices in MONARCH MIXER may be able to approximate some of these operations. An interesting question for future work would to study how well our techniques and concerns about causality and hardware utilization translate to these alternative convolution bases.
Optimization of deep learning primitives There is rich history of the optimization of deep learning primitives, as accelerating their performance can yield substantial savings in compute and cost for large models. There many approaches speed up these operations, but they usually either reduce data movement or compute.
Reducing data movement: In many applications, the major bottleneck is the storage and movement of large amounts of memory. One popular approach to reducing data movement is checkpointing, wherein one stores fewer intermediate results and recomputes the others on-the-fly where they are needed, trading additional compute for memory [46, 78]. Another approach is kernel fusion, wherein algorithms initially described as sequential steps can often fused in ways that improve their properties. For example, it is generally faster to implement dot-product through multiply-

---

Table 7: Accuracy on ImageNet- Ik. ResNet-152 provided for reference.
Model
Top-1%
Top-5%
Description
ResNet- 152 (60M)
78.6
94.3
ConvNet, MLP
ViT-b (87M)
78.5
93.6
Attention, MLP
ViT-b + Monarch (33M)
78.9
94.2
Attention, MLP-Free
HyenaViT-- (88M)
78.5
93.6
Attention-Free Free, MLP
M2-ViT-b (45M)
79.5
94.5
Attention-Firee, MLP-Free
Table Perplexity on the PILE when trained for different numbers of tokens.
Model
5B
10B
15B
Description
Transformer (125M)
13.3
11.9
11.2
Attention, MLP
Hyena (155M)
13.1
11.8
11.1
Attention-Free MLP
M2-GPT (145M)
12.9
11.6
10.9
Attention-Free, MLP-Free
Transformer (355M)
11.4
9.8
9.1
Attention, ML.
Hyena (360M)
11.3
9.8
9.2
Attention-Free, MLP
M2-GPT (360M)
11.0
9.6
9.0
Attention-Frec, MLP-Frec
accumulate rather than first multiplying and then accumulating. Recently, libraries such as PyTorch 2.0 [63] have added kernel fusion capabilities, although the very best performance usually still arises from handwritten kernels. Third, in order to better exploit memory locality, load small blocks of memory, do intensive computation them, and then write the results tile at a time [83]. Finally, many algorithms also have hund-opiimizaiinn that can remove unnecessary computation memory accesses [55].
often fastest to
Efficient algorithms usually make use of combination of these techniques. For example, FlashAt- tention [15] uses all four to dramatically decrease both the latency and memory consumption of multi-head attention. Though we have made a modest effort to implement MONARCH MIXER efficiently, we think it likely that MONARCH MIXER could be further optimized by these techniques.
Reducing flops: A first target for optimization is the multi-layer perceptron (MLP), owing to its ubiquity. A variety of structured sparse factorizations exist, many of which draw in this work [7, 11, 14, 16, 17, 19, 26, 91]. Attention is also popular target for optimization. Recently, a plethora of sub-quadratic approximations of attention have emerged, that aim approximate attention to reduce its quadratic complexity. Some methods rely on sparsification, relying on the fact that the attention matrix is extremely sparse long sequence lengths [2, 23, 24, 42, 53]. Others use low-rank approximations of the attention matrix [13, 79, 91] or kernel methods instead [9, 41J. A subset use combination of these techniques, such as [8, 72]. Finally, third category of methods [27, 64] aim to replace attention entirely, relying on spaee models [33].
7 Discussion and Conclusion
We explore MONARCH MIXER (M2), a new architecture that is sub-quadratic in both sequence length and model dimension and is hardware-efficient modern accelerators. We motivate M2 from both theoretical and systems performance perspectives and conduct a preliminary proof-of- concept investigation into performance on masked language modeling, image classification, and causal language modeling.
While our initial results are promising, our work is only first step in this direction. The M2 layer can likely be further optimized with systems optimization techniques such as kernel fusion. Our work has also not been optimized for inference like more well-established models such Transformers, even more recent models such as state space models. It also remains to be seen whether M2 layers have as widespread applicability as Transformers. We hope that these be fruitful directions for future work.
A discussion of broader impacts can be found in the Appendix.
10

---

Acknowledgments
We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (SDH) and PA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Supervision): the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Microsoft Azure, Swiss Rc, Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments Stanford Graduate Fellowship in Science and Engineering. and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys. The U.S. Government authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or rec- ommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government. JG and AR's work is supported by NSF grantit CCF-2247014. IJ's work is supported by an NSF Graduate Fellowship.
References
[1] Iz Beltagy, Kyle Lo, and Arman Cohan. Scibert: A pretrained language model for scientific text. arXiv preprint woXiv:/903. 10676, 2019.
[2] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.00 005000 2020.
[3] Emily M. Bender, Timnit Gebru, Angelina McMillan- Majorr, and Shmargaret Shmitchell, On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, FAccT' '21, page 610-623, New York, NY, USA, 2021. Association for Computing Machinery.
[4] Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. Better plain vit baselines for imagenet- Ik. arXiv preprint wrXh:2200.001500, 2022.
[5] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunitics and risks of foundation models. arXiv preprint arXiv:2108. 072558, 2021.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, er Language models are few-shot learners. Advances in neural information processing systems, 33:18 33:1877-1901, 2020.
[7] Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang. Zhao Song. Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. 2021.
[8] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[9] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
[10] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02331 2022.
[11] James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19900), 1965.

---

[12] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops. pages 702-703, 2020.
[13] Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. tunnee-transformrr: Filtering sequential redundancy for efficient language processing. Advances in neural information processing systems, 33:4271-4225 2020.
[14] Tri Dao, Beidi Chen, Nimit S Sohoni, Arjum Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference Machine Learning. PMLR, 2022.
[15] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-effcient exact attention with 10-awareness. In Advances in Neural Information Processing Systems, 2022.
[16] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algo- rithms for linear transforms using butterfly factorizations, 2020.
[17] Tri Dao, Nimit S. Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps, 2021.
[18) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pages 248 255. Ieee, 2009.
[19] Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing performance. arXiv preprint arXiv:1907.04480 2019.
[20] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
[21] Jacob Devlin, Ming-Wer Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In 1rxiv::810004805. 2019.
[22] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale, arXiv preprint arXiv:2010.11991 2020.
[23] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts In International Conference Machine Learning, pages 5547-5569. PMLR, 2022.
[24] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-5270, 2022.
[25] Wikimedia Foundation. Wikimedia downloads.
[26] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.03635. 2018.
[27] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. International Conference Learning Representations, 2023.
12

---

[28] Daniel Y. Fu, Elliot I.. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling. International Conference on Machine Learning, 2023.
[29] Morgan Funtowicz. Scaling up bert-like model inference modern epu part 1, 2021.
[30] Jonas Geiping and Tom Goldstein. Cramming: Training language model single gpu in one day. MwXi::2212.14334v1 2022.
[31] Google. Bard, httpp::/ //bard://bbrd..../ http://aard...goleeeee e. com/ 2023.
[32] Robert Gray al. Toeplitz and circulant matrices: A review. Foundations and Trends@ in Communications and Information Theory, 2(3):155-239) 2006.
[33] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:211 2021.
[34] Nicholas Hale and Alex Townsend. An algorithm for the convolution of legendre series. SIAM Journal on Scientific Computing,
2014.
[35] Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural net- works with pruning. trained quantization and huffman coding. arXiv preprint drXiv:1510.00149 2015.
[36] Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint wwXiv:2209.12951 2022.
[37] Dan Hendrycks, Norman Mu, Ekin D Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshmi- narayanan. Augmix: A simple data processing method to improve robustness and uncertainty. arXiv preprint arXiv:1912.02788 0207811, 2019.
[38] Peter Izsak, Moshe Berchansky, and Omer Levy. How train bert with an academic budget. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10644-10652, 2021.
[39] Di Jin, Zhijing Jin, Joey Tianyi Zhou, and Peter Szolovits. Is bert really robust? strong baseline for natural language attack on text classification and entailment. In Proceedings of the AAAI conference artificial intelligence, volume 34, pages 8018-8025, 2020.
[40] Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S Weld, Luke Zeutlemoyer, and Omer Levy. Spanbert: Improving pre-training by representing and predicting spans. Transactions of the Association for Computational Linguistics, 8:64-77, 77, 2020.
[41] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
[42] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint wxXiv:2001.04451, 2020.
[43] Jan Kocoñ, Igor Cichecki, Oliwier Kaszyca, Mateusz Kochanek, Dominika Szydło, Joanna Baran, Julita Bielaniewicz, Marcin Gruza, Arkadiusz Janz, Kamil Kanclerz, et al. Chatgpt: Jack of all trades, master of none, arXiv preprint 2302:22302.17224 2023.
[44] Elias Konstantinidis and Yiannis Cotronis. A practical performance model for compute and memory bound gpu kernels. In 2015 23rd Euromicro International Conference on Parallel, Distributed, and Kenwork-Based Processing, pages 651-658. IEEE, 2015.
[45] MV Koroteev. Bert: a review of applications natural language processing and understanding. arXiv preprint
2021.
13

---

[46] Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-efficient backpropagation. Advances in Neural Information Processing Systems, 32, 2019.
[47] Lagrange polynomial. Lagrange polynomial - Wikipedia, the free encyclopedia, 2005. https: //en. wikipedia.
[48] Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang. Biobert: a pre-trained biomedical language representation model for biomedical text mining. Bioinformatics, 36(4): 12e1112442222 1240,, 2020.
[49] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1 1692. 2019.
[50] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo. Swin transformer v2: Scaling up capacity and resolution. In International Conference Computer Vision and Pattern Recognition (CVPR), 2022.
[51] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[52] Xiaofei Ma, Zhiguo Wang, Patrick Ng. Ramesh Nallapati, and Bing Xiang. Universal text representation from bert: An empirical study. arXiv preprint urXiv:/910.07990 2019.
[53] Xuezhe Ma. Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34:2441-2453, 2021.
[54] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209. 100099/0655 10655, 2022.
[55] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint 1805 022667, 2018
[56] Derek Miller. Leveraging bert for extractive text summarization on lectures. arXiv preprint 1rXiv:1906.04115, 2019.
[57] Marcin Moczulski. Misha Denil, Jeremy Appleyard, and Nando de Freitas. Acde: A structured efficient linear layer. arXiv preprint 1rXiv:1511.059M 2015.
[58] NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
[59] NVIDIA. Nvidia A 100 tensor core GPU architecture, 2020.
[60] NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
[61] NVIDIA. cuBLAS, 2023.
[62]
OpenAI. Gpt-4 technical report, 2023.
[63]
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[64] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. International Conference Machine Learning, 2023.
14

---

[65] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. 2018.
[66] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. arXiv preprint arXiv: 10683, 2019.
[67] David W Romero, R Bruintjes, Erik J Bekkers, Jakub M Tomezak, Mark Hoogendoorn, and JC van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes, In 10th International Conference on Learning Representations, 2022.
[68] David W Romero, David M Knigge, Albert Gu, Erik J Bekkers, Efstratios Gavves, Jakub M Tomezak, and Mark Hoogendoom. Towards a general purpose enn for long range dependencies in {N} d. arXiv preprint arXiv:2206.0333 03398, 2022.
[69] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub Mikolaj Tomczak, and Mark Hoogen- doorn. Ckconv: Continuous kernel convolution for sequential data. In International Conference on Learning Representations, 2021.
[70] Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2 10270, 2021.
[71] G. Szegö. Orthogonal Polynomials. Number v.23 in American Mathematical Society colloquium publications. American Mathematical Society, 1967.
[72] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, ACM Computing Surveys, 55(6):1-28, 2022.
[73] Iyya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp- mixer: An all-mlp architecture for vision. Advances neural information processing systems, 34:24261-24222 222722, 2021.
[74] Asher Trockman and J Zico Kolter. Patches are all you need? Transactions on Machine Learning Research, 2023.
[75] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. volume 30, 2017.
[76] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding.
arXiv:1804.0741,, 2018.
[77] Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint 12..:222222100444, 2022.
[78] Qipeng Wang, Mengwei Xu, Chao Jin, Xinran Dong, Jinliang Yuan, Xin Jin, Gang Huang, Yunxin Liu, and Xuanzhe Liu. Melon: Breaking the memory wall for resource-efficient on- device machine learning. In Proceedings of the 20th Annual International Conference on Mobile Systems, Applications and Services, pages 450-463, 2022.
[79] Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint ar arXiv:2006.04766,, 2020.
[80] Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, al. Ethical and social risks of harm from language models. arXiv preprint 2t..::22122.433 2021.
[81] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Perric Cistac, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: Of-te- art natural language processing. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, 2020.
15

---

[82] Kuan Xu and Ana F. Loureiro. Spectral approximation of convolution operators. SIAM Journal on Scientific Computing, 40(4):A2336-A2355 2018.
[83] Yufan Xu, Saurabh Raje. Atanas Rountev, Gerald Sabin, Aravind Sukumaran Makumaran-RRaam, and Sadayappan. Training of deep learning pipelines on memory-constraine gpus via segmented fused-tiled execution. In Proceedings of the st ACM SIGPLAN International Conference on Compiler Construction, pages 104-116, 2022.
[84] Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers, 2023.
[85] Shanshan Yu, Jindian Su, and Da Luo. Improving bert-based text classification with auxiliary sentence and domain knowledge. IEEE Access, 7:176600- 176612, 2019.
[86] Li Yuan, Yumpeng Chen, Tao Wang, Weihao Yu. Yujun Shi, Zi- Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to- to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF international conference computer vision, pages 558-567, 2021.
[87] Sangdoo Yun, Dongyoon Han, Scong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference computer vision, pages 6023-6032, 2019.
[88] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 1rXiv:1710.09412. 2017.
[89] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi Bertscore: Evaluating text generation with bert. arXiv preprint arXiv:I904 190..09900999 096.5,967577 2019.
[90] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 13001-13008, 2020.
[91] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems, 34:17723-17736, 2021.
[92] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
16

---

Author Contributions
D.Y.F. Conceptualized the research; coordinated collaborations; developed M2 architectures; led experimental and implementation efforts; assisted in development of theoretical results; coordinated writing.
S.A.
Worked on developing M2 architectures; worked on implementing and conducting BERT experiments; worked on implementing and performing CPU experiments; assisted in writing and framing the work.
J.G.
Led development of theory and causal algorithms; wrote Appendix D.
I.J.
Led development of theory and causal algorithms; wrote Appendix D. Worked on writing and framing the work, assisted in development of M2 architectures; assisted in the optimized M2 implementation; conducted mixer benchmarks; assisted with BERT experiments; conducted Swin ImageNet experiments. Conducted ViT experiments; assisted in writing.
S.E.
A.W.T.
B.S.
Assisted in optimized M2 implementation; conducted mixer benchmarks; assisted in writing.
M.P.
Assisted in development of M2-GPT architecture. Supervised theory development; developed proofs; reviewed manuscript. Supervised research; reviewed manuscript.
A.R.
C.R.
Simran Arora, Jessicu Grogan, Isys Johnson, Sabri Eyuboglu, and Armin Thomas contributed equally to this work.
Appendix
Appendix A discusses broader impacts of our work. Appendix B presents additional experiments. Appendix C gives details for the experiments, including model architectures and hyperparameters. Appendix D gives missing details and proofs for the theoretical analysis, well as generalizations to broader results. Appendix E gives a Py Torch code listing of an M2 layer.
A Broader Impacts
Our work seeks to understand the fundamental capabilities and limitations of newly-emerging model architectures. As the amount of data and model size grows, we also seek understand how make training these models more efficient, both in terms of the amount of training context and the model size. This potentially connects to energy savings during model development and deployment, as well as making machine learning models accessible to larger population of people.
However, with any machine learning models, developing new techniques may impact wide range of applications, each with potential benefits and harms. Making language model training cheaper and longer context may make it cheaper spread disinformation. Similarly, improving the efficiency of model training may not reduce the overall environmental footprint training, since the same resources may be used train more models, or train the same models for longer. While our work makes partial progress on the fronts of efficiency and understanding. it does not explicitly address the issues of fairness and bias in language models. In addition, our work demonstrates a proof-of-concept; it has eиror modes, and we recognize the inherent risks of training and using machine learning models, including language models. Detailed discussions of these risks are in [3, 5, 80].
B Additional Experiments
B.1 Per-Task GLUE Numbers
We report full GLUE numbers for M2-BERT-base and M2-BERT-large in Table 9,
B.2 Additional Throughput Results
We report the throughput of M2-BERT-base (80M) compared to BERT models of the same size (BERT-base with fewer parameters), as well as the throughput of M2-BERT-large (260M) compared to BERT-large.
17

---

Table 9; Fine-tuning performance on GLUE [76]. We report the standard metrics FI scores for QQP and MRPC, Matthew's 's correlation for CoLA, Spearman's correlation for STS-B, BB, and accuracy for the remaining tasks, following the procedure from [38].
mm)
QNII
QQP
SST2
Col.A
MRPC
Average
M2 Muse (80M)
78.4 78.6
68.5
84.6
86.7
92.0
86.3
53.0
89.8
79.9
M2. BEBB--RRSSSSS (110M)
79.6/ 80.5
69.3
86.0
87.0
92.3
86.9
56.0
89.2
80.9
M2- large (260M)
72.8
93.3
M2. BERT large (341M)
82.2/82.3
75.0
87.0
87.7
92.4
88.3
59.6
90.1
82.8
Table 10: Throughput in tokens/ms by context length for M2-BERT-base (80M) compared to 80M
BERT models.
Model
512
1024
2048
4096
8192
HF BERT (79M)
248.4
157.3
86.0
46.8
0OM
FlashAttention BERT (79M)
433.3
425.1
335.2
217.4
122.6
M2-BERT-bass (80M)
386.3
380.7
378.9
353.9
320.1
M2 Speedup over HF BERT (80M)
1.6x
2,4x
4.4x
7.5x
Table 10 compares the performance of M2-BERTI base (80M) to BERT models murameter-matcher to 80M parameters. M2 is slower than Attenton sequence lengths 512 and I K, but outperforms FlashAttention starting at sequence length 2K. We believe further optimization of the M2 kernel can close the gap to FtashAttention for short sequences.
Table 11 compares M2-BERT-large (260M) to BERT-large. Trends are mostly similar to comparisons against BERT-base; M2 nearly matches FlashAttention sequence length 512, and outperforms for sequence length 1K and longer. We also see up to 4.3x speedup over HuggingFace BERT-large sequence length 2K.
B.3 ImageNet Comparison against Swin
Table 12 reports the results of replacing attention and MLP in Swin V2 using M2 as drop-in replace- ment. Surprisingly, Swin-M2 outperforms Swin-MLP-B, competitive with Swin-V1-B, V1-B, and comes within point of -2-nn even without any hyperparameter tuning or architecture adjustment from the ViT formula. We expect that performance may improve further with hyperparameter tuning specific to M2.
B.4 Speech Applications
Table 13 presents the performance of M2 on Speech Commands- 10, speech classification task over raw I-second clips sampled at 16 kHz. M2 is competitive with state-of-the-art architectures on this task.
B.5 CIFAR10
Table 14 shows the performance of MONARCH MIXER on CIFAR10. The trends are largely the same as on ImageNet.
B.6 Learnable Monarch Matrices in Sequence Mixer
In most of our models, we have used fixed Monarch matrices for the sequence mixer, and learnable Monarch matrices for the dimension mixer. Table presents an experiment evaluating using learnable Monarch matrices for the sequence mixer on the sequential CIFAR task. We use non- gated convolutional architecture based off long convolutions, as presented [28]. Learning the Monarch matrices in the sequence mixer yields 1.5 points of lift.
18

---

Table 11: Throughput in tokens/ms by context length for M2- M2-BERT-large (260M) compared to BaRge
Model
512
1024
2048
4096
8192
HF BERT-large (340M)
75.4
47.1
25.2
OOM
OOM
FlashAttention BERT-large (340M)
125.0
111.9
91.6
54.5
OOM
M2-BERT-large (260M)
122.5
118.6
109.4
94.5
75.0
M2 Speedup over HF BERT-large (340M)
1.6x
2.5x
4.3x
Table 12:
ImageNet accuracy of Swin models.
Model
ImageNet (acc@ 1) ImageNet (acc@5)
Swin-MLP-B
81.3
95.3
Swin- V1-B
83.5
96.5
Swin-V2-B
84,2
96.9
M2-Swin-B
83.5
96.7
B.7 Roofline Analysis
Figure 4 shows a Roofline analysis of a simple PyTorch implementation of single M2 operator M- MM- @ Mk on an A100 GPU, with 4K input length. The operation is more dominated by the data movement operations, which helps explain why performance is higher on newer architectures like RTX 4090 (which have faster and larger 1.2 cache).
B.8 Associative Recall
In Table 16, we present simple experiment demonstrating the causal parameterization of M2 on associative recall, synthetic language designed to test in-context learning. The model demonstrates in-context learning abilities in sequences up to 128K tokens, but Transformers do not scale past 8K.
B.9 BERT Experiments with Alternative Architecture
Here, we report results using an older version of the M2-BERT architecture, that uses non-gated convolutions and is trained on English Wikipedia [25] and English Bookcorpus [92]. For clarity, refer to this model as M1-BERT.
We found that M1-BERT could match Transformers MLM quality, but underperformed down- stream fine-tuning. We attribute this gap in performance to sub-optimal training hyperparameters (optimized for throughput using NVIDIA MIPerf hyperparameters) as well as sub-optimal archi- tecture. We report results here for completeness, but refer to the gated convolution architecture in the main body as the proper M2-BERT model.
These models followed the reference implementations and hyperparameters from Hugging Face Transformers examples [81] and Nvidia Deep Learning examples (https: //tp:///tthhuls cwm/ COm/NVIDIA/ DeepLearn ingnnninnEexsss...
particular, we use the LAMB optimizer with a learning rate of 5e For each sequence length. we use as large minibatch size as possible that fits on the GPU (A1 (A100- 80GB in Table 17 and V100 in Table 18). We set the gradient accumulation to reach a global batch size of 65, 536 sequences. To investigate the effect of sequence length, each model is trained for fixed sequence length in single phase of training (in contrast to some training protocols, which train the model in multiple phases, each at different sequence lengths).
Time Fixed Pretraining Quality on 8xA100 We compare time to fixed pretraining quality, training MI-BERT-base on English Wikipedia [25] and English Bookcorpus [92]. We compare against BERT-base trained with FLASIIATTNNION [15], as well as the Monarch-BERT--oe mentation from the original Monarch paper [14]. We measure wall-clock time for M1-BERT and the base Transformer reach 50% in masked language modeling accuracy on 8xA 100 Nvidia GPUs with 80GB memory each. Table 17 summarizes results. In short sequence lengths, MI-BERT is
imple-
19

---

Table 13: Accuracy Speech-Commands 10. An "x" means that the model did not fit in memory.
M2
S4
WaveGan-D Transformer Performer
CKConv
97.9
97.5
96.3
30.8
71.7
Table 14: Accuracy CIFAR-10.
Model
Top-1 1%
Description
VIT (1.2M)
78.6
Attention MLP
VIT + Monarch (607K)
79.0
Attention, MLP-Free
Hyena ViT (1.3M)
80.6
Attention-Free MLP
HyenaViT-M2 ViT-M2 (741K)
80.8
Attention-Free + MLP Free
comparable to FLASHATTENTION, even without using a heavily-optimized fused kernel. In longer sequence lengths, the FLOP savings make MI-BERT more efficient- up 2.4x faster than BERT with FLASHATTENTION at sequence length 4096.
BERT in a Day Inspired by recent work focusing on training under limited resource constraints [30], we measure how far we can get when training on a single V100 GPU in 12 hours. In Table 18, we report the masked language modeling accuracy achieved by the same of models and sequence lengths (except for the PIAASHATTENTION baseline, which is not supported on V100). We observe MI-BERT both achieves higher accuracy within the time limit and can be trained at longer sequence lengths than the baseline architectures.
Downstream Fine- Tane-Tunnng We evaluate the quality of M1-BERT-base models on the GLUE benchmark [76]. Table 19 shows fine-tuning performance the GLUE tasks, using the same hyperparameters and 5 epochs for all tasks and both models. MI-BERT-bass is competitive with Transformers trained using MLPerf hyperparameters on Bookcorpus and Wikitext, but underperforms fully-trained transformers and M2-BERT-base.
C Experiment Details
C.1 Model Architectures
In this section, we describe the exact model architectures we used for each task, including the design of the block (residuals and gating). We additionally release our code for reproducibility,
BERT Language Modeling The M2-BERT architectures use a standard BERT backbone, but replace the attention with bidirectional gated convolutions and replace the linear layers in the MLPs with block-diagonal matrices. All the M2-BERT architectures use an expansion factor of four. M2-BERT-Be (80M) has model width of 768 and 12 layers; M2-BERT-bare (1 10M) has model width of 960 and 12 layers; M2-BER- Mare (260M) has a model width of 1536 and 12 layers; and M2-BERT- (341M) has model width of 1792 and 12 layers. We train all these models for 70,000 steps, with sequence length 128, and global batch size 4096 sequences. For all the models, use decoupled Adam with learning rate 8e-4 and decoupled weight decay le-5. We use linear learning rate decay with a warmup of 6% of the steps, and we use MLM masking percentage of 30%.
C4
For GLUE fine-tuning, we do small search of learning rate, weight decay, and number of epochs. Following [38], we fine-tune RTE, MRPC, and STS-B from the MNLI checkpoint. We fine-tune all tasks with sequence length 128. For some tasks, we also pool the embeddings of all the non-padding tokens instead of using the CLS token.
The final hyperparameters for M12-BET--bal (80M) are decoupled Adam W with learning rate 5e-5
and weight decay 5e-6 for epochs for MNLI; AdamW with learning rate 5e-5 and weight decay 0.01 for 6 epochs for RTE; AdamW with learning rate 3e-5 and weight decay 0.01 for 10 epochs on QQP:
Adam W with learning rate 5e-5 and weight decay le-5 for 10 epochs with average pooling for QNLI;
decoupled Adam W with learning rate 3e-5 and weight decay 3ed-6 for 3 epochs for SST-2; AdamW
20

---

Table 15: Accuracy on sequential CIFAR for fixed vs. learnable Monarch in the sequence mixer.
Model
SCIFAR Accuracy
M2, Fixed Monarch
91.0
M2, Learnable Monarch
92.5
M2 PyTorch performance Roolline
FP36 Tensor Cores]
A100 hcoline
HD Pyltrch implementaiion
1011
Peroorment
1040
1371
104
100
10'
107
10'
Antomebi intensity SLOPPDlyte
Figure 4: Roofline plot of PyTTrrc implementation of a single M2 operator M- @ Mk).
with learning rate 7e-5 and weight decay 0.01 for 10 epochs for STS-B; Adam W with learning rate 5e-5 and weight decay 0.01 for 10 epochs for MRPC: and decoupled AdamW with learning rate 5c-5 and weight decay 5e-6 for 10 epochs for COLA.
For M2-BERT-ber (110M), the hyperparameters are decoupled AdamW with learning rate 5e-5 and weight decay 5e-6 for 3 epochs for MNLI; decoupled Adam with learning rate le-5 and weight decay le-6 for epochs for RTE; decoupled AdamW with learning rate 3e-5 and weight decay 3e-6 for 5 epochs on QQP; decoupled AdamW with learning rate 5e-5 and weight decay le-5 for 10 epochs with average pooling for QNLI; decoupled A with learning rate 3e-5 and weight decay 3ed-6 for epochs for SST-2; decoupled Adam with learning rate 8e-5 and weight decay 3e-6 for 10 epochs for STS-B; decoupled AdamW with learning rate 8e-5 and weight decay 8e-5 for 10 epochs for MRPC; and Adam W with learning rate 8e-5 and weight decay 5e-6 for 10 epochs for COLA.
For M2-BERT-large (260M), the hyperparameters are decoupled AdamW with learning rate 5e-5 and weight decay 5e-6 for 3 epochs for MNLI; decoupled AdamW with learning rate le-5 and weight decay le-6 for epochs for RTE; decoupled Adam with learning rate 3e-5 and weight decay 3 for 5 epochs QQP; decoupled AdamW with learning rate 5e-5 and weight decay le-5 for 10 epochs for QNLI; decoupled AdamW with learning rate 3e-5 and weight decay 3ed-6 for 3 epochs for SST-2; decoupled A with learning rate 7e-5 and weight decay 3e-6 for 10 epochs for STS-B; decoupled AdamW W with learning rate 8e-5 and weight decay 8e-6 for 10 epochs for MRPC; and A with learning rate 5e-5 and weight decay 5e-6 for 10 epochs for COLA.
For M2-BERT-large (341M), the hyperparameters are decoupled Adam W with learning rate 5e-5 and weight decay 5e-6 for 3 epochs for MNLI; AdamW with learning rate 5e-5 and weight decay 1e-6 for 2 epochs for RTE; decoupled Adam W with learning rate 3e-5 and weight decay 3e-6 for epochs QQP: decoupled AdamW with learning rate 5e-5 and weight decay 1e-6 for 10 epochs for QNLI; decoupled Adam with learning rate 3e-5 and weight decay 3ed-6 for 3 epochs for SST-2; decoupled AdamW with learning rate 8e-5 and weight decay 3e-5 for 8 epochs for STS-B decoupled Adam V with learning rate 8e-5 and weight decay 8e-6 for 10 epochs for MRPC; and decoupled AdamW W with learning rate 5e-5 and weight decay 1e-6 for 10 epochs for COLA.
VIT We use a standard ViT model architecture as base [22]. In line with recent improvements to the ViT architecture [4, 70, 86), we use sinusoidal position embeddings and global average-pooling (GAP) instead of a class token.
21

---

Table 16: In-context learning performance on associative recall at various sequence lengths, vocab size 20, X indicates the Transformer did not finish in week.
Model
0.5K
2K
8K
32K
128K
Transformer 100.0
100.0
100.0
X
X
MONARCH MIXER
98.7
99.4
99.4
99.4
99.4
Table 17: Time in hours to reach 50% masked language modeling validation accuracy 8xA100 with different sequence lengths.
Model
512
1024
2048
4096
Architecture Details
BERT-base- ALSIHAATNTTATTTNIIO (110M)
2.7
3.8
5.7
13.2
Attention, MLP
BERT-base-Hugge Hugginggaene (110M)
3.3
5.6
13.1
26.7
Attention, MLP
HERT-Monarch-e (80M)
3.1
4.7
10.3
22.1
Attention, MLP-free
MI-BERT-base (55M)
2.5
3.5
4.0
5.5
Attention-Free MLP-free
Specdup
1.1x
1.1x
1,3x
2.4x
We adapt the ViT architecture by replacing its MLP and/or attention components with Monarch Matrices (similar to our adaptation of BERT):
We replace the MLP with randomly initialized Monarch Matrices of the same dimension the dense matrices of the MLP and learn those matrices during training, setting the number of blocks in the block-diagonal matrices to 4.
We replace attention with the recently introduced Hyena operator [64]. The Hyena operator represents a recurrence of two efficient sub-quadre -ub-uuadaaicc primitives, implicit long convolution and multiplicative element-win -wie gating of the projected input. Hyena operators apply the FFT algorithm to achieve fast long convolutions in sub-quadratic time. We further adapt the Hyena operator by replacing its long convolutions with the M2 operator and setting the Monarch Matrices to the DFT and inverse DFT.
VIT for ImageNet-lk In line with other work [4, 14, 64, 70], we use ViT-base architecture with 12 layers, hidden size of 768, 12 attention heads per layer, an intermediate size of the MLP projection of 3, 072, and patch size of 16 x 16 pixels. For optimization, we follow the training procedure of T2T- ViT [86], including augmentations such as RandAugment [12] (magnitude 9, magnitude-std = 0.5, layers 2), Mixup [88] (a 0.8), CutMix [87] (a = 1.0), Random erasing [90] (rate 0.25), and AugMix [37]. See Table for all other training settings.
VIT for CIFAR-10 We use ViT architecture with 6 layers, hidden size of 128, 8 attention heads per layer, an intermediate size of the MLPP projection of 512, and a patch size of x pixels. We further tune weight decay (0 0.1), stochastic depth rate (0 or 0.1), and base learning rate (1e 4 or 3e - 4 or 1e - 3) and report the test performance for the model variant that achieved the highest accuracy in separate held-out validation dataset (randomly selected 10% of training data). We also apply an early stopping rule such that training is stopped if the model's validation loss does not improve for 10 training epochs. See Table 20 for all other training settings.
GPT Causal Language Modeling Similarly to our ViT approach, we also replace attention with the Hyena operator, using the same architecture as in [64] starting point. The Hyena architecture has two convolutions, which can be computed using the FFT convolution theorem. In our architecture, we additionally replace these FFT operations with causal Monarch matrices.
In addition, we re-use the heads extension from the H3 architecture [27]. The heads extension groups the model dimension into heads, ties together the long convolution parameters in each head, and then computes the outer product between different input projections. An algorithmic listing adapted from the H3 paper [27] is provided in Listing 1, with updates to replace the SSM layers with Hyena convolutions. We use head dimension of 16. Setting the head dimension to and replacing the Monarch matrices with FFT is equivalent to the Hyena layer.
22

---

Table 18: Masked language modeling validation accuracy achieved on single V100 in 12 hours with different sequence lengths. X indicates the model does not on device with a batch size of 1.
Model
512
1024
2048
4096
8192
Architecture Details
BERT-base (110M)
11.5
7.8
6.8
X
X
Attention, MLP Attention, MLP-Free Attention-Free, MLP-Frree
SERT-Monarch--ass
6.9
8.5
6.8
X
X
MI-BERT-base
20.2
20.2
20.1
17.1
12.9
Table 19: Fine-tuning performance on the GLUE benchmark [76], after pretraining on Wikipedia and Bookcorpus. We report the standard metrics F1 scores for QQP and MRPC, Matthew's correlation for CoLA, Spearman's 's correlation for STS- B, and accuracy for the remaining tasks [21].
Model
MNLI (m/ mmn)
RTE
QNLI
QQP
SST2
STS-B
CoLA
MRPC
Architecture Details
BERT no pretrain
34.1
47.3
79.9
0.0
77.9
Attention, MLP
BERT base
74.5 74.7
55.6
69.3
81.8
83.9
19.8
12.1
74.2
Attention, MLP
53.1
85.2
75.4
Attention-
Algorithm M2 Hyena Layer with Heads
Input: Input sequence R Nxx from the previous layer, weight matrices W x1: W x2. Wv. Wo Bd causal Monarch matrix M, short convolution kernels K,, K2, K3, a Hyena convolution kernel Klong head dimension dy.
Output: Output sequence E R
Compute X1 u. x1, X2 & x2, = uWv Nxd
Pass X1, X2, V each through the short convolution using the causal Monarch matrices: X
X2,
M MMX,, @ MKK)). (MX2 @ MK2), M MMM---MMMM--
(MV @ M
MK3).
Split X2, V into H "heads" (X, (h) X2
(h)
A (h)
for h L...., H). each sequence of N vectors of
size da. d/H.
for S h S do
Take the batched outer product X² (h)
) E Nxd, xdn (batched in the N N-dimension) and pass
through the long convolution using the causal Monarch: (V( = - 'MX"X.""" (h) (v(h)) MK RR- Xdd xd,,
Batch-multiply by X,:
(h))
X (h) XV
X X 12 (batched in the N-
dimension).
Concatenate the output (h) of each head, and multiply by the output projection matrix Wo gdxd
Finally, we remove the MLP layers entirely (equivalent replacing the layer with an identity). and make the model wider to compensate (the depths match the equivalent Hyena models). The small model has model width of 1160 with 18 layers and uses learning rate of 0.0006, and the medium model has model width of 1344 with 40 layers and uses learning rate of 0.0008. All other hyperparameters match the Hyena models [64].
D Missing details from Section
This section contains all the missing details (including proofs) from Section 4.
In Appendix D.I, we review some definitions and results on murti-variate polynomials and set some notation needed for this section. In Appendix D.2, we explicitly connect Monarch matrices for p and bivariate polynomial evaluation. Specifically, we prove Theorem and Theorem 2. Then in Appendix D.3 we show how to instantiate the bivariate basis polynomials so that we get causal map. This includes converting the bivariate polynomials to univariate polynomials (with evaluations over the Nith roots of unity) and this proves Theorem We then show how this causal map can be implemented only using GEMMs (and O (N3/2) FLOPs) in Appendix D.4.
Next, we note that while our evaluations points are over complex numbers, our input and output to the Monarch convolution layers are over reals. Hence, it is natural to wonder if we can implement the entire layer just with operations over real numbers. One potential advantage of this is that theoretically only have to keep N real numbers for intermediate results (instead of 2N reals
23

---

Table 20: ViT training settings,
ImageNet-Ik 1k
CIFAR-10
Optimizer
AdamW
Optimizer momentum Learning rate schedule
31. B2 0.9, 0.999
Cosine decay w/ linear warmup 0 0.1
Dropout rate
Label smoothing
Image size
224 224
32 32
Base learning rate
le-3
(1e-4, 3e-4, le-3)
Batch size
1024
512
Training epochs
300
up to 500 5
Warmup epochs
10
Stochastic depth rate
0.1
{0, 0.1} {0, 0.13
Weight decay
0.05
numbers when we keep track of vectors in CC This can reduce the data movement costs. Further, multiplication of two complex numbers requires six operations over real numbers (four multiplication and two addition). Thus, moving to an implementation that only uses real numbers could potentially lead wall clock time speedup. We propose one such scheme Appendix D.5 that proves version of Theorem just over reals by moving to the Chebyshev basis (instead of the standard monomial basis). This creates new technical challenges, which we also address.
Finally, we generalize our results to arbitrary 2 2 Appendix D.6. We would like to point out that to get causal map (in Theorem 17) we need to 'embed' input vectors of size n into vectors of size N 2P n + O (n--1e (n)) For p 2, we avoided the blowup of 2 = 4 with blowup of 2 instead (via Theorem 3). Whether this is possible to do (i.e. have a blowup of 2 instead of 2P) for p > 2 an interesting direction for future work. Further, the matrices that lead to causal map can be represented with O (p.2/r) parameters while the matrices in Theorem 3 use more parameters. Extending the causal map for p > 2 that uses O (N!*)) parameters is an exciting direction for future work.
D.1 Background and Notation
We collect known facts and definitions about multi-variate polynomials in Appendix D.I.1 and recall some notation from [14] in Appendix D.1.2. These will be needed throughout this appendix section.
D.1.1 Multi-variate Polynomials
Basic Definitions Letp 2 be an integer. We recollect some definitions on p-variate polynomials
(over IR) in variables , When p [1,2}, we will use variables in {X, 2} for notational simplicity.
We will use X to denote the vector of variables (No...., Xp-1). Further for j E 20- we use the notation
x
X
a=0
XJ is (standard basis) monomial, where j
(jo- ............. 1p-1).
generic p-variate polynomial defined as (with standard monomial representation)
(()) = Σ Sj xJ
jE2P
where the coefficient lj E R.
We will need the following notion of degrees:
24

---

Definition 2 (Degree). Let 0 S <. p. The degree of Xa in X (with) Ljo....-Jp-1 (p)1)) Ja- The degree of Xa of (X)), denoted by degx, (q) is the maximum degree of Xa over all monomials XJ with Oj + 0.
Note that for p the above coincides with the usual notion of degree of univariate polynomial (()), in which case we just use deg(g(Z) (2)) to denote degz ())()(g()))
We will need the notion of taking mod of p-variate polynomial with p-tuple of polynomials, The notion of mod is well defined for a univariate polynomial (which we will assume as a given below) but in general for arbitrary p-variate polynomials q(X) and q' pXXX), the operation q (X) mod q' q((X) is not well defined. However, we will only need the following restricted operation:
Definition Let p 2 Fix a p-tuple of polynomials Ro (o(X)).... ( Then for any j 2h 20 we define
p-1
x mod (Ro (Xo) ....., Rp- (Xp-1))
II (x mod (Ra (Xa))). a=u
For a general polynomial p(X).
p(X) mod (Ro (X.)...., Ry-1 (Xp-1))
defined by extending the definition for X3 by linearity.
Polynomial Evaluation Given a p-variate polynomial ((X) and an point a E RP, the evaluation of q at denoted by g(a) is evaluation of as function
Given subsets S. C, we define (X) evaluated at x the vector of values g(a) overall E S
In this paper, we will in many cases evaluate polynomials at the appropriate roots of unity. Specifically for integer N, we will define
WN 201/NN
and note that the Nth roots of unity is the set {aly 0 si<
Polynomial Interpolation We now recall univariate and bivariate polynomial interpolation results (proved via the Lagrange basis), which we will use in later subsections.
Theorem 4. Let D 2 be an integer. Given ys for 0 S i < D and o; for 0 < i < D there exists a unique univariate polynomial P(X) with deg( P)) D, such that for all 0 < i D,
P(a:) =
Proof. This proof is based on the Wikipedia entry for Lagrange polynomials [47].
Given a sequence of values a for 0 S < stt.t a; * ay + j. the Lagrange basis for polynomials of degree < D for these values is the set of each polynomials {Po(X), P1 (X),... PD-1 (X)) each of degree D - 1. Each basis polynomial are defined as:
(X)
X
X X-0i-1
0i-1
X - 0.+1
X X-OD-1
II
X - 0j
(8)
ag au
a
0:
04 04+1
0; QD-1
0Sj<D 1
0; aj
By definition,
[1 for jorj=i
p:(a,) =
0
otherwise
The Lagrange interpolating polynomial for those nodes through the corresponding values yi for 0 < i < D is the linear combination:
D-1
P(X) =
Σ V/ (x().
(10)
1=0
25

---

By (9), for all 0 sii< D:
P(a.) =
(1I)
Finally, the interpolating polynomial unique. Assume there is another polynomial Mf (X) of degree
< D such that ((a)) = Vi for all 0 S i < D. Then the difference ( P(X) is 0 at D distinct points a, for S < D. And the only polynomials of degree < D with more than D roots the 0 polynomial. So, (() = P(X).
Theorem 5. Let Dx, Dy 2 1 be integers. Given values Vig for 0 Si jx < Dy and Dx distinct points (ar..... DDx-11) Dy distinct points (3n...., By there exists a unique bivariate
polynomial P(X,Y) with degx(P) < Dx degy Dy such that for all S i Dx 0 OSj Dy:
P(ay, Bj) Vij
(12)
Proof. Define
TR n
P(X,Y) ΣΣ Vi P:() p:()))p;)y))
(13)
where pt and py are Lagrange basis polynomials defined in the proof of Theorem 4 such that for 0 S i < Dx.
p:(ax) { for i = k otherwise
(14)
and for 0 j Dy
((BB) {1 for h = l
(15)
otherwise
From above, we have for all k,
[1
for i and j =
p((ax) (,,())
0
(16)
otherwise
Then, for all 0 0 <xiiDxx.o Dxi<Dxx Dy:
P(as, Bj) Vij.
(17)
By definition of Lagrange basis polynomials, degx (P) < Dx and degy < Dy.
Finally, the interpolating polynomial unique. Assume there is another polynomial M M Y) with degx x (M) 0 < j < Dy. Then the difference
Dx and degy (M) Dy such that M (as, By) yij for all 0 < i Dx and -XX,,, is 0 at Dx Dy distinct points, (a;, B,) for 0 S i < Dx,0 0 S j Dy. And the only polynomial with degx Dx and degy Dy that has Dx Dy roots is the 0 polynomial.
D.1.2 Notation
Here we recall notation we will use from [14].
1. The class of Monarch matrices is defined in appendix C of of [14] as ((b.N) which N x N matrices with block size for any integer 0 < 6 < N that divides N. When b /N drop from the notation giving (i1, io) and (f1,jo). For example, this is used Proof of Corollary 2.
2. Row index can be represented as (i1, io), Which gives = inb + io.
3. Similarly, column index can be represented as (jı, (j1,jo),- Which gives j = jib + jo- Note that when b VN, j1 k(j) and ju m(j). We choose to the (o) notation here since that notation is easier to generalize for > 2.
4. L E (B.,,N) is N x N matrix with b b blocks that are all diagonal matrices.
26

---

5. R. BB((,.) meaning it's block diagonal N N matrix with block size b b.
6. We have a class of permutation matrices defined as (b. (7) = in N + 11. This can be denoted by an N x N matrix, P(t.N). where the 1 row is (o(b.N)(1)-
7. We'll use pair notation 4 to denote the rows, and j or pair notation (f1,jo), to denote columns. It should be clear from context which one we're using.
For any 0 S ji VN. let (X,, be arbitrary bivariate polynomial with degx (li.). degy (e)) < VN.
For any S $o < VN, let in (Y) be an arbitrary univariate polynomial degree < vN. Let A = (ag..... -V--11 B =
each be sequence of distinct eval points. Note From the proof of Theorem in the Appendix C of the Monarch paper [14] we get,
that A and B need not disjoint.
P(.N) L. L.PToN))
= PONN) I. ..P. N4
Therefore,
L = Pl..) (b.N) L. P(IN)
(( L P(I,N)-
Define DB and BD as set of all such and matrices over IR N where in # ko
L [i. kol def (011, (k1, ko) =0
(18)
and if k1 + ji
Re [o] del R((kI, ko) (11.50) O.
(19)
Pictorially, L and R look as follows:
ky 30
Ag
4
PLP
R
In [14], Monarch matrices with block size b VN, M' L R. and thus for all 0 < i, i1.00 VN:
M (1 vo) (11, = E, [o, io] R jol-
(20)
We note that our definition Monarch matrix M in Section is slightly different in that M M'P with M' defined in [14].
27

---

D.2 Monarch Matrices and Bivariate Polynomial Evaluation
Given polynomials 2, (X, for - < VN, polynomials Jo (Y) for 0 Σ J1.j0 VN, evaluation points A (ag: a VVN-1) B (3p. .. .. / (as in Appendix D.1.2), define the matrices L E DBVN.N and R E BD NN
as:
For every 0 S j1, 11,00 < VN:
Li 1o. iol Lji (ai, Bio).
(21)
For every 5 $1.Jo, VN:
Righ [ol + To do (Big).
(22)
Note that all entries of L and R not specified above are set to 0.
Let f be the above function that maps coefficients of l (X, Y) (which are coefficients of monomials X i for all io VN and hence represented by matrix in RVNxVN and coefficients of di (Y) (which are coefficients of monomials Y'o for all 0 S in VN and hence represented by vector in I VN) for all S j1, < VN to pairs of matrices in DBVN,N BDVN,N
Theorem 6. Let be as defined above. Then is a bijection.
Proof. To prove is bijection we must show s' one-to-one and f-1 is one-to-one (and exists). To show f is one to one means each set of polynomials' coefficients given to f. will output a unique set of matrices (L, R) DBVN,N BDVN,N This follows from (21), (22) and the known fact that polynomial evaluation is a function.
Now, to show f-1 exists and is one-to-one, we must show that there is map from any pair (L, R.) E DBVN.N x BDVN.N to unique sets polynomials, E, F, with parameters as defined in Appendix D.1.2. Further, we need
Lighi iol = Egg (a:. 3r0)
(23)
and
R R Fo. = ji (3io)-
(24)
We will use Theorems and 4 to show the existence of E,, and Tii polynomials, giving us the mapping from the matrices to unique polynomials.
We first show the existence of the unique polynomials in (24). Fix J1. < VN. Then consider the values 0 < io < VN:
Via
Riihi Po Jol.
(25)
Then by Theorem 4, there exists a unique polynomial of degree < VN (call it f1.30 (Y) such that for all 0 < in < VN,
ii (Bog) = Vio
which by (25) shows (24).
Next we show the existence of the unique polynomials in (23). Fix 0 S < VN. Consider the values 0 Σ 11, in < VN:
Vi Ln [8o do]-
(26)
Then by Theorem 5, there exists a unique bi-variate polynomial of deg (call E (X,Y)) such that for all 0 < in in < VN.
N and degy N
Eg (0ri Bio) Wii
(27)
which by (26) shows (23).
Therefore f is bijection.
28

---

We conclude:
Corollary 1. For every matrix M' as defined in (20), there exists unique polynomials Ej (X, Y) and Tinio (Y), such that for all S i1, 10, 31.30 < VN.
[ 40) fo)
Es (to) dr (.to).
(28)
Proof. Follows from (20), (21), (22) and Theorem 6.
D.2.1 Proof of Theorem 1
We begin with an immediate consequence of Corollary 1:
Corollary 2. Let A, B C C such that |A| = =B = VN. Then the jth column of M' is the evaluation of the polynomial Es (X,Y) Jj (Y) over A x
Proof. Observe that for fixed jo.Ji the right hand side of (28) is b, (X, Jo (Y) evaluated at all (a, B) E x B. Thus, (j1,j0) column is evaluation of E₁ (X, Y) -Tisdo (Y) over points in x B, as desired.
Next, we state a generalization of Theorem that follows from Corollary 2: Corollary 3. Let A and B be as in Corollary 2. For any vector M' is ū(X, Y) evaluated at A Further,
x(X,))=
Σ
Ui (X, To
(29)
0ShJoKVN
where E and defined by M' in Corollary I.
Proof. Follows from Corollary 2 and definition of matrix vector multiplication.
In Theorem 1 and the following sections, consider the polynomial evaluated the basis polynomials defined by
M M
Corollary 4. Let A and B be as in Corollary 2. For any vector u, and M u is u(X, Y) evaluated at A x B. Further,
u(X,))
Σ
Ujo.h
Ejo (X, (X.Y) Y) -Jih (Y),
(30)
0ShJo<vN
where and are defined by M' as in Corollary 1.
Proof. Follows from Corollary 3 and definition of M.
Specifically, Theorem is a special case of Corollary 4 where l (X,Y) = (x))(( and (Y) -TooJi (Y).
D.2.2 Proof Theorem
By Corollary 4, for all 0 < Siii < VN,
6 (X, Y), Tjo (Y), (Y), (X,,) Y), (Y)
29

---

the basis polynomials corresponding M1 and M2. For the coefficient vector and similarly for u (1) <VN. we can construct two polynomials
x) = E
kinndo (X, Tjach (Y)
OSiIdosvA u(X,)) = Σ
Ji 60 (X.Y) 330.n (Y)
OShJosvN
2TLa
whose evaluation over (as, .,,)) = (wit, wh) where recall as in Appendix D.1.Iw = e VN by Theorem is equivalent to the products M1 k and M2 respectively. Taking the component product, y- k) @ (M2 u), the entry at = (0,, is given by
y((1,, = k(w", w00) u(w" w
Noting that the element of A, i.c. the - roots of unity, satisfy ZVN means that the above are evaluations of
MX, = k(X Y)= w(X,Y) mod (XvN - 1,YvN -
at A x A. Finally, Theorem and the fact that M7¹ exists implies Mo y is polynomial interpo- lation into basis polynomials corresponding to Mo. (Here we use the well known fact polynomial interpolation is the inverse of polynomial evaluation).
D.3 Proof of Theorem
We review some concepts in Appendix D.3.1. In Appendix D.3.2, we discuss square matrices and causality in terms of operations on univariate polynomials. This allows us define a general class of operators for causal ID convolution. In Appendix D.3.3, we give class of matrices suitable for perform causal Monarch convolution. Specifically, we prove Theorem 3.
D.3.1 Review
Consider the linear operation on an input vector uI:
y A u.
We say that the map is causal to mean the entry ii only depends on u(0], u[1),.. ufi]. This will be the case when A is lower triangular matrix (we index the top left entry of A (0, When A is lower triangular Toeplitz matrix with entries corresponding to some coefficient vector k, this operation exactly the ID convolution
...
y y= (F ((F22 k') (F2n u())) [0: n 1),
where k' = (k, 0,), n' = (u, 0,,), and F,, is the x n DFT matrix.
Definition 4. For matrix M E _axx, let us define the map
y y
M.
(31)
matrix convolution. When M is Monarch matrix, (31) is called Monarch convolution.
In this section, we are interested determining large subclasses of matrices M such that for any coefficient vector k, (31) is causal in u. We provide a class of matrices for which Monarch convolution
causal.
We note that for general Monarch matrix M, (31) is not causal in By Theorem 2, we have
y (X,Y) k(X,Y) un (X,Y) mod (XvN-1 (XVX-I,yvN-I)
This is not causal because the mod (X / yVN-1 term condenses higher order terms into lower order terms, hence the yii] wouldn't just depend on input information up to value i.
30

---

D.3.2 Univariate Matrix Convolutions
We start with a couple of notation assumptions.
Assumption 1. N a perfect square.
Assumption 2. We will
pair notation for this subsection since throughout have i
VV do and j + 1 + jo-
In order to discuss square matrices in terms of univariate polynomials, we give univariate analogs of Theohee and Theorem 2 for general univariate basis. With an eye toward towards performing causal convolution, we restrict our analysis to certain classes of univariate polynomials.
We first define matrices whose 7th columns the evaluation of a minimum degree (and maximum degree N 1) polynomial (recall Definition 1). We generalize Theorem 3 to such matrices.
N Lemma 1. For sequence of points A = (1,wN,--w'" N-1 where is the N root of unity, let M be defined as
M[i.j] = (wx)
(32)
where qj (Z) is defined in Definition 1. Then for any vector E IR,, M v is equivalent to evaluating the polynomial
N-1
x22) = Σ
0/ (Z)
(33)
j=
at (I, W ...NN..-WN- W
Proof. By our definition of M, the column M[; j] is exactly the evaluation of the polynomial 7, (2) each point in The claimed result comes from the definition of matrix vector multiplication and (33).
Note that M any M's in this sub-section are not necessarily Monarch matrices.
Next, we state the following intermediate result:
Proposition Let be the set of the N-th roots of unity. Then for M1, M2 defined as (32)
== (M1 k) (M2-u0 u)
the same as evaluating the polynomial
v(Z) = k(Z) u(z) mod (2 - -1)
over A where k(2), u(() are of the form (33), corresponding to M1 and M2, respectively. In other words, for any < N,
yld (wx)
Proof. This result follows from Lemma and the definition of the Hadamard product.
Next, state re- re-interpretation of M-'y:
Proposition 2. Let M be full rank matrix whose columns are the evaluations of the basis polyno-
mials ūj (Z) from Definition for < j < N, and let y E N be an arbitrary vector. If = M M-1y. then for all 0 <
y|i] = ul(w)
where ul (() is the same in Lemma 1 for M. In other words, M is the polynomial interpolaton problem for the polynomial basis q; for 0 < N.
Proof. This follows from Lemma and the fact that M is invertible.
From Propositions and 2, we get the following generalization of Theorem 2:
31

---

Theorem 7. For matrices Mo, M1, M2 as defined above, the operation
M, ((M, k) o (M2- u))
equivalent to representing the polynomial
f(2) k(2) u(Z) mod (Z 1)
in terms of the basis polynomials
for j j= 0......,
where k(Z), u(Z) defined as in Lemma I in terms of the basis polynomials corresponding M and M2, respectively, and (ū, ((,(2)))sss) corresponds to Mg.
Proof. Follows from Propositions 1 and 2.
Now give the class of matrices from which build causal map. Specifically we prove generalization of Theorem 3:
Theorem Let n 2 1, les N = [v2n| 2 Then define the basis polynomial T3 (Z) to have minimum
degree j and maximun degree 1 for 0 S j < n, and n S N. 93 (Z) has minimum degree j and maximum degree N 1.
For all MN with basis columns defined by (ūs(Z))osse) as above, the operation
11
(M (MN 'MM
Ox,0 MN (u, Ow-))) 00: 1]
(34)
gives a causal map.
Proof. To prove this is causal means each entry, f[i] dependent only u(0), u(1], u(i], where f M (MN (k, Ov- MN (u, 0N on)). By Theorem 7, have
((2) = k(Z) u(Z) mod (ZN-1). -1). 1),
where
k(Z) Σ
kygj (Z)
u( Σ Ug (Z).
Since deg(k(Z) u(Z)) < 2 N 2, this is equivalent to
f(2) -((Z) u((z)
n-1
Σ --'==
by - Uj (Z) (Z).
By our choice of qG we ensure that T - has minimum degree j j and deg(ū)-qz)) S 2n 2n-2 < N for any 0 < j.j n. Then by Lemma 2 (see below), there exists coefficients aj+J','' , such that,
N Σ
1(2) = Σ
ky
Uj
Osti (Z))
=
V=
N-1
Σ
Σ
O'+t'ss ky (Z))).
0
J-'=0 =0 J+
If we define
32

---

N-1
1(2) = Σ
Ji (2),
=0
then for 0 Si < n, we get:
n-1
f: = Σ
astr's uj.
3.'==
j4''
Note that fi only depends on u (0), u(1], ................ uli]. as desired.
We used the following lemma in the proof of Theorem 8.
Lemma 2. Let (Z) be defined as in Theorem 8. Then for any 0 0 j.j' <
N-1
(2())
Σ
ag4za. (Z))
(35)
i=j+
for some set of coefficients o'+jj',5-
Proof. We first note that by our choice of Øj. the minimum degree of (Z) (2) is j and deg(ū; () Tz (Z)) I 2n - 2 < N - 1. Our claim follows from that fact that any polyno- mial pa(2) Z) of minimum degree d and deg(pa) N can be expressed as linear combination of (Z(Z),
(Z).
D.3.3 Causal Monarch Convolutions
In this section we will prove Theorem 3. We will so by showing that the basis polynomials (;( as defined in Theorem 3 are special case of the basis polynomials, q (Z) defined in Theorem 8.
We start with couple of notation assumptions.
Assumption 3. In this sub-section we will using block size b = VN therefore. we are dropping the block size from index notation. For example, (01,30)<77 vN becomes (in, io) for this section.
Assumption 4. Permutation matrices in this subsection all the same P (VN,N) drop the subscript and just use P.
Definition 5. Define
del (Z) TjiJo (ZvR).
(36)
l (Z) has minimum degree j1, and Tjn.jo has minimum degree jo- All polynomials Ki have maximum degree S N 1.
(2)
Next we argue that the above basis polynomials have a specific minimum degree.
Lemma 3. Polynomial (63 (Z) as defined in equation (36) has minimun degree N ji.
Proof. We need to show that ( (Z)
(Z) is a minimum degree jo / + polynomial. Note that
(Z) = - rizdo ( VVV
= (N-1-h
Z(VN-IWN FVN-1-J0
1
7.
ZVN
'This claim be shown using downward induction on d N 1, 2.....
33

---

where & VV- (.) has degree VN- -1 1- and -N--130 ()) has degree VN W1 -1-jo- Simplifying we get
11-10 (
1
((ijo) (2) (2)= 2N-1, L ( 2 -
LVV
This claim follows since e 1111- (Y) (X) ijrwo (rvR)) has degree (vN -1- (VN -
jo) VN = -N (iov NNV++0 ti).
Note that the polynomial ((1,5o)) (Z) has minimum degree jovN + +j. This is not jivN V + jo as in defined in equation (1), we will talk more about this soon.
Next, we observe that the polynomials in (36) define matrix that satisfies (20).
Lemma 4. Let (os for S JivN N + jo < n be as in (36). Define
M'((i. (h..o)) (o).do)(W) (No0).
Then M' satisfies (20).
Proof. If we evaluate the polynomials d(jn.0o) (Z) at wiy for 0 si < N, get
((siddo (wx) = Eu (wx) rs 1e (W
Since
WVN N
= WiNNNioNN = w = w we get
M to) (f1, = E₁, (wij) ( Vi+t, Ts (w))))
The above corresponds to how we define (20) since we have
[(1 1o), (o): = Linii (io, dol R Jo]
with
Li [No, iol = Euj (uii
and
R [oo = Tiijo (
Recall from Lemma 3 that (5130) has minimum degree jo / +11- For causality we need di)) (Z)
to have degree minimum ji /N + jo (then the polynomials will satisfy the minimum degree require- ments from (I)). Therefore, we permute the columns of M',
M
(37)
Note that the above M PLPRP and is indeed a Monarch matrix defined in Section 3. Note that the basis polynomials of M are defined as,
U(n (Z) = (Z) Fjo.s (2VZ).
(38)
We note that the above is same as (Z) defined in (5) where j = 11 V + jo with the correction in (5) that (Z) Lm() r, (ZV where rY) = iin (Y). We are finally ready to prove Theorem 3. Corollary (Theorem restated). Let N [22n]² 2- Define Mx by ((r as in (38). Then, 1 u (M (MN ON ) MN (N)))) [0 [0:n-
gives a causal map.
Proof. Due to using q polynomials as in (38), by Lemma 3 the degree of the (J1.ju) th column is a polynomial with minimum degree jı N + jo This implies that these basis polynomials are subset of the more general causal maps in Theorem 8, which proves the claim.
It can also be verified that (0) has maximum degree N 1.
34

---

D.4 Block Algorithms for Complex Numbers and Block Size VN
In the following subsections we restate the results in Section D.3.3 in terms of block operations. As mentioned earlier, this so that operations (and simple data movement operations like permutations).
do computations Monarch matrices using only GEMM
In Appendix D.4.1 consider arbitrary Monarch matrices and in Appendix D.4.2 we consider the sub-class Monarch matrices corresponding to Theorem 3.
D.4.1 General Block Monarch Convolution
In this subsection we re-state general Monarch convolutions in terms of block operations.
Recall from equation (22) we defined the block diagonal matrix R as follows (0 S io f1.j1 < VN):
Fr (wt (w))))
R (oo jol
(39)
where deg(f) < VN.
To do so, will first work with M'y such that MM MMy P. We want express Monarch matrices, Mr, as univariate polynomial evaluation over (I, -113. Towards that end, define
Tu (Z) CsiJo (ZvM).
(51 (ivN+ia) VN -Ww//N who NN
By simple observation that w N
we have
Jo (w VN vi) =0 to (wi)x)
In other words we have,
R [o] = Tiido (wx).
Fix j1 < V so we're looking at the 2 column of block R going down the column we evaluate the polynomial ij.jo at points (I.w....... multiplication of Fourier matrix of size vN and a matrix of the coefficients of the polynomials.
WVN-1) Which equivalent to matrix
WN
VN
Ri.h*
F
Fniila
NN
ju
So we can think of the blocks of R. as a Fourier matrix times a coefficient matrix. In other words, let us define a matrix R NxvN which will hold VN coefficient blocks Ro, Ri.... R R/N-1
R RVNNVN such that,
R2 jol = F (a).
where
NN -1
(Y)=
Σ
in (a) ya.
a=0
Then we define
Righ
F FV Rr
35

---

Next, restate the above product of two block diagonal matrices:
N
R.0
Reo
R
F VN
R
VN-1
R M...
In other words, we have
R = diag (F ( F....... FvR) diag R ....... R. RV--))
(40)
Equation (21) defined L as evaluation of bivariate polynomials. now wish to do the same with univariate polynomials.
Recall that L PLP. We define the ig diagonal block (0 S i1, i0, < VN) for Las
L by. (WNrNNtH)
(41)
where deg(en) N.
Let us define matrix L NxvN that will hold the coefficients of polynomials l, (Z) ie.
[a, = Ej](a] [a]
where
N Σ a=0
(2)
&g [a] z
We can multiply this matrix with the Fourier matrix of size N x N to get the blocks of L (which we will need diagonalize). Specifically define
F
The rows of L" and the rows M' are both indexed by i. Meaning they're ordered in lexicographic ordering io), which is problem for the following reason. The block diagonal matrix L made from L" needs to be in lexicographic ordering (ig. i1) (see (41)) since it gets permuted on the left M PLP and right allowing M be ordered by (61,4)) Therefore, when composing from L" we must permute the rows by P. So we get,
P'=P.L".
Let
L'=
L'
Then
L diag(L/
L' )
(42)
Pictorially:
In
L
I'o
P
FN N
L
T LyR-n
L' JR-
VW-1
36

---

Let be the function that maps coefficient matrices L, R E to (L, R. where L and R are defined by (42) and (40).
Theorem 9. Let be defined above. Then is a bijection.
Proof. To prove f is a bijection we must show f is one-to-one and f one-to-one (and exists).
To show f one-to-one means both L and R given to f will output unique pair (L, This follows from (42) and (40), and the fact that polynomial evaluation is a function.
Now show f-1 exists and one-to-one, we must show that there's map for any (L, R.) to unique L, R E R N Then for any pair (L, R) where both L and R are block diagonal matrices, there's map unique sets of polynomials, E, where each coefficient is entry of L, R (thus giving unique mapping from (L, R)'s (I, Ry's).
We need to show the existence of e, (Z) and Jo (Y) such that: Lio [i, = Ej (w Vii
and
R ji (io. jol ii (w 'NN)
Fix 0 S j1,j0 < VN. Then consider the values 0 < i0 < VN:
Uha f Righ [ol-
Then by Theorem there exists a unique polynomial of degree < VN (call it (5..)) such that for all S i0 < VN:
J1
VN Vio
There will be N polynomials with VN coefficients each, meaning there's unique sets of coefficients make up the indices of R.
Now to get the entries of L from L, fix 0 S f1 < VN. Then consider the values 0 S i1, VN: 1/1.00 +- Lig [i,
Then by Theorem 4 there exists unique polynomial of degree < N (call it (j.) such that for all 0 S in, i0 < VN:
en (w fu N
Vir.io
There will be / polynomials with N coefficients each, meaning there's unique sets of coefficients to make up the indices of L.
Algorithm 2 is pseudo code for the map from L, R NxVN to block diagonal matrices L, R.
Algorithm 2 BLOCKY MONARCH(L, R.
Input: E, R NNVV
Output:
Block diagonal matrices L, R e CNN
b First, compute L from L
Let Fy be the Fourier transform Monarch matrix PL, P.
See Corollary
2: + P- FN L
3: for <- 0 to VN- do
4: I' < L'(av av + VN 1,-
5: L diag(L L' \
Now compute R. from R
6: for - 1 N
do
7: R - VN
R Rav/ a. + -1.
8: R + diag (Ro.... R
9: return L, R
37

---

3/N Lemma 5. Algorithm uses O(N312) FLOPs and 3v N GEMMs of two VN x VN malrices.
Proof. Lines 1 and 2 are multiplying a monarch matrix representing the Fourier transform times N x N matrix in block fashion again, giving O(N3/2) FLOPs.
Similarly, lines and 7 have two N matrices multiplied vN times. Giving 0(N3/2) FLOPs.
Lines 3, 4, 5, 8, and 9 don't count towards FLOPs
Therefore have O(N3/2) FLOPS and 3v N GEMMs of two VN VN matrices.
Now that we have the operations terms of blocks, we will make them causal in the following sub-section.
D.4.2 Causal Block Monarch Convolution Recall that a Monarch matrix is defined as
PLPRP,
then per equation (37) we have
M' PLPR.
(43)
Then if q jo) Jo) (Z) is the basis polynomial corresponding the (i1, column of M, by (38) the
basis polynomial corresponding to M is
(ii (2) ijo (ZVN)
(44)
with the minimum degree of ji VN + jo (recall that we pick fo and iooJs to have minimum degree jo and ji respectively) as desired.
Theorem 10. Let L, R + BLOCKY MONARCH(I, R). Let M be as in (37), Then for every 0 S i, j N. have:
M[ ((i..oo) o0) (WNeNNtio).
where (inijo)) is as in (44).
Proof. To prove this we need to show the (f1,50)** column is the basis polynomial (6r.jo)
evaluated at the tth roots of unity, where ((1)))) is as defined in (44). Note we have
96155) (w (Nto) Lj (w (wirNNtio) Foodis (wofi +i
ejo (w (wt) Pjoss (w (
(45)
By definition we have,
M (1 (61, (f1,)00) = Lin. [io. dol Rind (i. [o].
By (39) we have
R jish jo] = Pinvoo (wia (
and by (41) and the fact that
I -
we have
Li 1io, do} ejj (wii (Nttt)
Thus, have
M' 4o) (f1.Ju0l = Esj (wiy Finjo (w)xx)
Since
M. M-P=MM
38

---

have
io) (f1, M io), (o0.5i]] = -io (w*) (Noo)) Jjoss (w>x) 96i15o)(W ( N
where the last equality follows from (45).
Corollary 6. The DFT M as Theorem 10 when all blocks of and the top block L are the identity matrix of size /N x VN (the rest of is all O's).
Proof. Since only the diagonal of the top block of L will contain any non-zero values only index the first O...., VN- 1 rows. And since all blocks of R are the identity matrix we get
Lij,,11 = Lnr 1=111111 =1
and
R2 Jol = Psn Jo Lool =
All other entries of L and R are 0. Thus, we have
tx (Z) ~
and
(2) 2
As per Theorem 10,
(Z) = (Z) ijowi (2
-u.. znNN zottiiNN =z¹
Then by Theorem 10 note
(o) (w(x = = M(((1, (Gi)
which implies M is F v as desired.
Algorithm 3 is pseudo code for the Monarch convolution algorithm. It maps input space of L, R matrices, a kernel vector k, and input vector u to a vector
Algorithm 3 BLOCKMoNARCHCONVVi R. k. u)
Input: L, R e NNVN
k, u e N
Output: f N
I: L. R +- BLOCKY (L, R)
2: M - PLPRP
Get M' from BLOCKY MONARCH b Compute k,, uy from M
3: ky + M- k
4: u/ M
5: f M-1 (ky 3 uy)
Compute f b
6: return
We next outline how to make the map in Algorithm 3 causal. Towards that end, we observe:
39

---

Lemma 6. For any fixed 2 1, les N [V2m]]
Define L, R E NRNN such that for every
0 Sh < VNN,
Lo
Ro
L
and R
V
R VN--
Define L', R' Nx with corresponding coefficient blocks
(Ih)
RRy
OSKEVN
oSkeeN
R.
0 VNXVN
R
L' =
and R''
VNNVN
R'
VN-1_
where
Lolio:Jil Jil ifi 0 and io 51
VV-1 E = UNNVN for ..NII
otherwise
0 if < jo or ((io
VN
and (jo
VN
R/ [io jol
2
for j1 0...., VN
Rí Jul otherwise
Lil are all non-zero entries for all 0 S jo: < VN. Further, we require that R/, [io jol, Then for jovN N + i VN N
the basis polynomial (() f) (Z)r jn (7V/]
of M' PLPR has minimum degree JoVN + Ji and maximum degree S n. 1, and for JovN + +j < N, the basis polynomial (2) has minimum degree jo / N + and maximum degree < 1.
Note that in M = M'P the (ji.jo) basis polynomial Ci (2) = T/o.5 (Z) has the required degree bound.
Proof. Let L', R' +- BLOCKY MONARCu(L', R'), and denote their corresponding polynomials as E/ (Z), 1/ (2 respectively for every I Ju < VN. By our definition, for all S jo. 11 NN, 3/ (ZVN) has minimum degree jov N and Eg (Z) has minimum degree j1-
Then it follows that the basis polynomial (( (Z) = Eg (Z) in (ZvR) has minimum degree Ji + JovN.
We now look the degrees of each basis polynomial. From our definition of R all entries
R [oo. = for jo
N and io >
2 This implics that for 0 S jo
and
0 S fi < NN,, we have deg(r') )
<
VN
Since we are only looking at L (Z), note that degree
40

---

deg(E))) S VN -1 Then it follows that
VN
dog (Aiinnto) SVN
1 1+
VN
2
VN
=VN-11 1+
VN
VN
VN
S
V
2
N
1.
2
(Note that in the above JovN + + S VN
by same calculations above as needed.)
For 2
< jo
< N and < j1 < VN, deg(d)) = vN - -1, and Eg (2) degree vN 1.
Then follows that
deg (den S VN 1+ (VN-1) VN
N 1
as desired. (Note that jo N +
NN
VN as needed.)
Finally, we use Lemma 6 and Theorem 8 to conclude the following:
Theorem 11. For any fixed 2 let N [v2n Let L' R NNN with corresponding
coefficient blocks (L.)ugkeexi uSkeeN' R2)) OSkEVN be defined as in Lemma Then for any IR",
BLOCKMONARCHCONN [L','',k' u') w'110: 10: n 11. where k' = (k, ON- u' (u, ON is
causal in u.
Proof. Note that this is the same setting as Theorem 8. If N = [v2n|] then | 2 Then by Lemma 6, the basis polynomials Jj (Z) of M PLPR have deg( JivN + jo V
N for 0 <
Then (34) computes BL.CCKMOARRCCCON (,, u') Since Algorithm performs the same operation as (34), the result follows from Lemma 6 and Theorem 8.
D.5 Bivariate Polynomials with Kronecker Substitution Over Reals
Our earlier results pertain complex evaluation points. In this section define causal convolution over real evaluation points. To do this, we redefine our basis polynomials in terms of the Chebyshev polynomials of the first kind.
In Appendix D.5.1 we recover Theorem for univariate polynomials defined over real evaluation points. This identifies class of matrices that we Appendix D.5.2. We show that these matrices can be used to perform (31). Finally, in Appendix D.5.3 give the analog of Theorem 11 over real evaluation points. However, the resulting matrices not Monarch matrices but they are close. In Appendix D.5.4 and Appendix D.5.5 we discuss how we can exploit this closeness, and compute (31) more efficiently.
define to define structured subelass in
'Reaal that Lemma 6 is stated for basis polynomials do (2) for M' = PLPR where Un (Z)
Looin (Z).
41

---

D.5.1 Univariate Evaluation over Real Numbers
For any integer 0, the Chebyshev polynomial of the first kind with degree is denoted (2) and is defined as
Ta (coos)) 0) def cos(aB),
(46)
and has the property
To cos 0) (-1)" cos(a0).
(47)
To parallel Appendix D.3.2, we consider the class of basis polynomials with the form
N-j-1
(2)
Σ
(48)
evaluated over (WN.))ssieN where
Tí + 3)
def
WN
COS
(49)
N
Let us consider the class of matrices defined over (48) and (49).
Lemma 7. For sequence of points A {WN,........ WN,N J. let M be defined as Mii.j] - 0f (wN,))
(50)
where 13 (Z) and WN, are defined as in (48) and (49), respectively. Then for any vector u, M
equivalent to evaluating the polynomial
N Σ 1=0
u (2)
Uj 9j (Z)
(51)
at each point in A.
Proof. By our definition of M, the column Mí: jl is exactly the evaluation of the polynomial 9j (Z) at each point in A. The claimed result comes from the definition of matrix vector multiplication and (51).
This leads to the following analog of Theorem 7.
Theorem 12. For matrices Ma, M1, M2, each of form as in Lemma 7, the operation
= Mn¹ -(M, k) ④ (M₂ u))
(52)
equivalent to representing the polynomial
f(Z) k(Z) u(Z) mod Tx(
in terms of the basis polynomials
(Z) 0f (2) for j O, O,..,,
where k(Z), u(Z) are defined in terms of the respective basis polynomials corresponding to M1 and M2 as in Lemma 7, and (a)) (Z)) corresponds to M
Proof. For A (WNNJosseN define
(a(2)
II (Z a). aεA
(53)
Then (52) follows since for
- =())=((Z) u(2) mod ga(2).
we have the following for any A:
f(a) = k(a) u(a).
The claim follows from Lemma 7, (52), the invertibility Mo, and the known fact that Tx
qA (Z).
42

---

We also utilize the following result:
Lemma 8. Let qj $ [() be defined as in (48). Then for any 0 Sij < N,
N-1
d (Z) q3 21 (2) (
Σ
astsss (2).
(54)
for some set of coefficients aj+j',
Proof. From (48), we have
1
2 16 (2) =
()laTTa(Z) La (2).
Recall that within a m dimensional vector space of polynomials, we can define a basis by choosing any set of polynomials with degrees ,, m- 1. Because deg
s -
(i + J') polynomials where for 0 < 2 [ (j - j') - 2, the a-th polynomial has degree In other words we can choose the a-th polynomial qx -a-1 (Z). Thus we have
it can be written as a linear combination of any set of _
-
x Z qj' (Z)
Σ
NN ajtY'a -N-a-1 (Z)
for some set of coefficients O'+1''a- Then after reindexing i t- N - a --1 we have
N-1
/ (Z)
1 43 (2)=
Σ
0J- (2 (Z).
i=(N - [
The claim follows by setting 03+3"++544 0, and if N is odd, 0345-1 = 0, and 0j+7', Oj+J'NNi-- for other i.
This allows us to prove the following causality result for convolutions over real evaluation points.
Theorem 13. Fix a family of basis polynomials q6 .q." ...,,.., ..".... _ as defined in (48). Let N Let k"
be a perfect square, n S | k, u Rn and MN defined by basis (d)) (2))
k, -n) and n" u, 0 [
m))
) Then the operation
(MW' (MN- (01z).
MN (0)=),-")) 10:
(55)
defines a causal map in u.
Proof. Let
f f- (-(My' (MN- (0/z), MN 1 u²) (0/47,""))
[0:
1).
(56)
In order to prove that (55) is actually causal in the input E R", we must show that for all 0 S < N, fli] is dependent only on q[]].u[1]. ............ uft]. Let k' (0 k") and u' = -0 uu) By
Lemma 7, MN k' and My u' correspond to the evaluations the polynomials
N-1 uly /3 (Z).
M(Z) Σ K3 T (Z).
and nl(2)
Σ =
(57)
43

---

Let us define
n-1
[ (2), (2), and u(Z) = Σ
n-1
k(Z) = Σ Ko 9
U
97
(Z).
(58)
1
Note that for 0 Ssj <
N
the coefficients ky = ul} = 0. Then (57) becomes
N-1
N
(Z)
Σ
K/j 43 (Z), and v'(2) Σ
"
93 (7),
1
1
which is equivalent to
+ 1
| + -1
K(Z)=
Σ
1
(+1r7(2)) (2), and u'(2)
Σ
"
j=0
p'=0
For 0 $ j <
N
deg
q
1 - j - 1, and deg NN
(0*111) =
N 2 -
1 [ j This implies that qj 1 (Z) and 1+1 (2) are both linear combinations of (Z)
2 - - Chebychev polynomials. Then we can set 0 (Z) / प
(Z). Similarly, note ky. Then it follows that k(Z) = (Z).. and by a similar argument,
that for 0 < < n, N zZ) u' (Z). Then by Theorem 12 we have
i+ ++
f(Z) = k(Z) (()) mod Tx(2)
k(()) (()))
(59)
where the last statement follows since deg (k(Z)) deg (u( 2)) S n 1< 1 implying that their product has deg (K(Z) ū()) deg (I'N (2)) = N. We want to write f(Z) in the form
N-1
((Z) = Σ J (Z)
1=0
for some set of coefficients set of coefficients fi. From (58), (59) becomes
n- 1-11
ky - uj q $ (Z) q + (Z). (2))
f(2) = -
--0
Then by Lemma 8 we have
Σ (Z)
Σ ky
Σ E=+
Opts'
2 ( (Z).
i='=0
We show that for all 0 < i < N, fa is a function of (ue) tw)osrse Then note from the above that each kj and uy appears in terms of q. where i i2j+ + j'. Then we have
N l
f =
Σ
asti'a ky Uj.
J jt''si
as desired.
44

---

D.5.2
Structured Causal Matrices
In this section we narrow our scope from the general class of matrices defined in Appendix D.5.1 1 to particular structured subclass. Now let us define the class of structured causal matrices over the real numbers.
Definition 6. Define
Eg (Z)
f [a] TZ(Z),
(Z)
def
Jo [a] Tz
(60)
2=0
a=0
where
j1 [a] 0 if a) is odd,
(61)
We define structured causal (SC) matrix polynomials as
Lor vx (Z) = L ( 5o (T
(62)
A N x N SC matrix is defined over the set of real evaluation points as
Mi (Jo =4 (ii.do) (WN.i) (ooi) (WN,)).
N
(63)
We note that in (63) deviate from the usual ordering of indices (ji,do) (fo: We show that the SC matrix falls under the category of matrices defined by (50). Lemma 9. Let M denote the set of all matrices defined by of (50), and MSC denote the set of all matrices defined by Definition 6. Then MSC CMC..
Proof. To show that MSC C MMS then it is sufficient to show that for j jovN + j1. any IN (N (Z) is equivalent some oi" (7) as in (48). From (63) have
d (od (2) = L-N- (Z) & -1 (TvN((
N
(64)
Note that
deg ( (1(rnd)))))
= VN - ju (VN - 1
N VNji - 10
= N - 1-
From the fact that deg (T.) a, represented as a linear combination of (Z) like so:
Ta polynomial basis. Then of N
(n v
can be
N N-j-1
(1 v (Z)
Σ
Win.do (a)Ta
which is exactly the form of (48).
Lemma 9 allows us apply the causality result from Appendix D.5.1. Corollary 7. Fix family of basis polynomials goy gi 90.1.....
9/N-1/N-1 //
defined in
(62). For any perfect square N 1, | k R and M' defined by basis ( (1 (Z) OSHidRVN as in (63). Then the operation (55) with MN + M'y defines causal map in u.
Proof. Follows from Theorem 13 and Lemma 9,
45

---

D.5.3 Block Operations Structured Causal Matrices
In this section show how build structured causal matrices through block operations.
Constructing M
Recall that in Appendix D,4.1, we defined L, R in terms of coefficient matrices R, L E RNx/N with blocks Rg, La for 0 S h < VN, noting that for each block R ji
R
VN Ryso
and for Le RNVN
L' ='=PFN L.
These matrices are then diagonalized into R. We use Definition 6 to similarly define L, R, R X with blocks
{Ri Sos.ooosssoeev/ where:
Li (io in] = -i--)-I(WN,))
Righ [oo jol + -NN--I- (WN,))
(65)
and all other entries are zero. Let the coefficient matrices L, R NN.N be defined with respect to blocks as follows:
[[a,ji] - Ra, [o, = = -N------N------
where the entries of VV--i-- and i -N----LN--- are defined as in (60) and (61). Now let CN NXN where
Cxli = (w)
(66)
be the Chebyshev transform. Then analogous to Appendix D.4.1, we define
Righ = =C R
R
='=PCN L.
This allows to give algorithm the following construction algorithm for L and R.
Algorithm 4 BL.OcKSC[L, R)
Input: L, R. E NxVV
Output: Block diagonal matrices L, R E NNN
First, compute L from L
L' P. P.C L
2: for + to VN - do
3: I' + L' Jay N a/N + -1-1.
4: L diag(L
Now compute from R
5: for a 0 0
N do
6: R - CV RlavN N. aVN N + VN -1,
7: R diag (Ra,-.. R R
8: return L, R.
use Algorithm to specify another type of matrix M NNN
Lemma 10. Let
M PLPRP
where L and R are oulputs from BLOCKSC. Then each entry in M is defined as
Mii = --N-joo-1 (WNs))
TVN--o-1-VN--- (W/N
(67)
46

---

Proof. Let Mo PLPR, then M M,P. Then we have
Mol(i1.fu) (f1, LvN- ( VNN-I. (WN,so).
Then we get
M j] MoP[((1, (f1,0)))))) = Mo((E1, (fo (fo. i.]l
This gives
M[((i, (11,10) (WN-))-(WN.)) (F TVN-J--1,NN-11- (WN,
as desired.
Next will discuss the relation between matrices that we get from Lemma 10 and SC matrices.
D.5.4 Constructing M' from M
Since CN & MSC 7, Algorithm is not a proper analog of Algorithm 2. In this section, we show how to convert M produced from Algorithm 4 into matrix with block decomposible form. Recall from (63) that
M o N (1rJ0) vx (WN,)) L ..V -NN-J---- -o - (Tvr
(68)
We note the distinction between the above and (67). Specifically, we note that M is evaluated two sets of evaluation points- the - and V/N -N--i-1,/NNde polynomials are evaluated over the N roots and v roots of unity while M' only evaluated the N-th roots of unity. However, they close due to the following property:
Lemma 11. Let Ta be Chebyshev polynomial of the first kind of degree and define WN, in (49). Then
T T/ (WN.) = (-1)" W.NN.0'
Proof.
VN v N + in + 1
T TV (WN,)) Co8
N
π (No + 2)
= c08
11 +
VN
क (io + +
(-1)st cos
VN
(-1)4t WNN.o
In the above the first equality follows from (46).
Analogous to how we utilized roots of unity, Lemma 11 allows us to express our basis polynomials in terms of two related sets of evaluation points, WNN, and wy.s. The following lemma demonstrates how to translate between M' and M. Lemma 12. For all i.
i -- M
(69)
We do not have proof of this claim, but it seems unlikely that Cy e SC
47

---

Proof. By (68) we have
[i = -NNN-j-- PNN VVN-1o-1,
= =(-8-5--(WN.)) -oN- (--))
where the second statement follows from Lemma 11. Then from (60) and (47) we get
VN- ---nl -1
M = = VV--jo-(WNN (WN,))
VVN -V--o-1-V--j--1011 ((-1)
a=
VN- W-
= -N (WW
Σ
A N 1o-----,------ To (u)
a=0
-a) a) is odd. Then equivalently we Note from (61) that v 1 = (Sr(vV) V 11 -1 1 get
Mi = l LNN- (WN (-1j)----)((VN-----1
VN-
RVN WWN- (-))---A-- (WNs.
a=0
VN-h11
LN-- (wN,s) (-))) (VN)-(VVN---1))
Σ
WN- To (w)s
a=0
-(-1))((8~ VN VN- ( ( 1N-- (WINNN)
Then from (67) we have
= -- (-) =
Block Matrix Multiplication for Structured Causal Matrices Lemma 12 shows us how to compute M' from M. However, this does not mean that the matrix vector multiplication problem for M' can be implemented efficiently. In this section we show how to perform matrix vector multiplication of M' with any RR- from matrix-vector multiplications of M (and these operations are indeed efficient).
The key to our approach involves considering the parity of each block index 0 Si < VN. Let us 1 +
define map MIX R 1
R such that MIX(uo, u,) where
u[i]
mod 2(1/2].
We use this map show that M'u u and M' can be computed efficiently.
Lemma 13, For any u E RN and M E RNN,, M'u can be computed via two matrix-vector mulriplications: M MIX (un 0) and M MIX (0, u,).
Proof. For i = in VN + io N.O == jo < N. let D E RNN be the diagonal matrix defined such that D (-1)
Lemma 12 implies that
if i is even
M (i.
DM (DM[i,))
if i is odd, jı is even otherwise.
48

---

We want shift the (-1) Compute zo M (MIx (un, 0) + MIX (0, M.
(Mix (No 0) MIx (0, u1)). Define y' such that
if ji is even
y' b
ji is odd
It can be verified that D y' = M'u, which completes the proof.
We now give the analogous result for u M'. Lemma 14. For any E TeN and M RNX,, u M' can be computed via two matrix-vector multiplications: MIx (uo; 0) M and MIx (0, u1) M.
Proof. For =i= i V + 10 < AN..Sj=AVV + jo < N. D E RNxN be the diagonal
matrix defined such that Díi, (-1) and u' D Mix (u/. u/)). Lemma 12 implies that
J
if i is even
Mii
DM
i odd, ji is even otherwise.
- (DM[i,))
We want shift the (-1) to u'. Compute Za Mix (u/, 0)T MIx ui)) M, (Mix (ub, 0)" -Mx (0, u²)))
M. If
M'u. u, then one can check that
Zulill
if jı is even
y
z
if Ji is odd
which completes the proof.
Lemma 13 implies that computing M and My in (55) can be done efficiently. However, we do not know how to compute My' y for any arbitrary y. We address this partially in the next section.
D.5.5 Inverse of Chebyschev Transform
In this subsection show how computing Cr' can be done efficiently.
Lemma 15. Let C.v E NxN be defined as (66). Then
CN'
C diag( 1/N, 2/N, 2/N).
Proof. Follows since (C
Given the above, the goal then is to compute C or equivalently u CCy for any R,N Note that is sufficient to show that
CN CN SN
for CN,S, NNN such that u CN and SS can be computed with two invocations of Lemma 14. Indeed we have
(i + () VN + jo)
Cx cos
N
T(i })
7(i + 1)100
cos
+
VN
N
Tii + })i
T(i
+ 5 }
( + 4):00
cos
sin
sin
VN
N
VN
N
(70)
49

---

We this define CN and Sv. along with two additional matrices CN, SN E RNXN such that for C2 and Cx:
M( 5)11
Ti 3)500
CN[i. cos
COs
VN
N
(ig + })5i
π(( + 2)000
Tijsi +
008
VN
N
π (io + 2)11
T(i + 1/ )io
(-1)/u cos
def (-1)4.JCCx(i.) C. (71)
cos
VN
N
and similarly for SNx and SN:
Ti + 2)12
xii 2)30
Sx[1,]= = sin
sin
VN
N
π ( + 5)111
T(2 + 2))00
= (-1)6 sin
sin
def (-1) SN (-I) (72)
VN
N
summarize these results into the following lemma.
Lemma 16. Define CN2 CN, Sy and SN as in (71) and (72), respectively. Then
CN CN SN
where ON and NN are of the form (67).
Finally Lemma 14 and Lemma 16 imply the following: Theorem 14, For any u E RRN
= ==uTCN
can be computed from four calls matrix vector multiplication with Monarch matrices.
Proof Sketch for Theorem 14 Lemma 16 tells that T C' and u T S are sufficient compute CN. and (72) shows us that CN Sy are Monarch matrices that allows u CC and S to be computed from two matrix- vector multiplications with Cy and SN, respectively. Then the claim follows from applying Lemma 14 twice.
D.6 Multivariate MONARCH MIXER with Block Size V/N
In this section, generalize Theorem 1, Theorem 2, and Theorem 3 arbitrary 2 2. In Appendix D.6.1, we set up notation. We then generalize Theorem and Theorem 2 to general p in Appendix D.6.2 and then generalize Theorem 3 in Appendix D.6.3. In Appendix D.6.4 we connect definition of p-variate Monarch in Appendix D.6.2 to the one defined Section 3. Finally in Appendix D.6.5 discuss some extensions and generalizations.
D.6.1 Notation
We will use notation from Appendix D.5.1.
Fix an integer p 2 2. We will be working with indices j, where j = (jo...., (p-1) and i (io ..... ......,,pp)) with 0 < ia.ja VN for every
We will denote the set of all sub-indices (0, V/N), and the operator < denote the lexicographical ordering of vectors in [0, vNY).
For any < < M < N such that divides M and M divides N, define the following permutation matrices:
Pk. diag
Pur. Pu. Po.as
times
50

---

Note that PN,N.N is exactly the as PVIN from earlier.
we use ab, M, N) to denote the corresponding permutation then it takes the input (1p-1....; 10) and maps it to (3p-1s..... 2p' 22.-21..... ........ 2o, &y' ), (where p log, N and p' log, M). Le. (b, M, N) does not change the first - p' sub-indices and then does a left rotation on the re- maining p' sub-indices.
For the rest of the section, assume permutation matrix".
/ N and then consider the following 'sub-index reversal'
p-2
R 6,6
II
P P
e=0
If o RB,, N) is the permutation corresponding to the permutation matrix above, then RR N) maps (fp-1....,;0) 10) +
(p--...;
D.6.2 Generalizing Theorem and Theorem 2 For specific 0 S < let
VN-1
(Xa) ji
pla) 10 [m] Tm (Xa)
(73)
be an arbitrary polynomial of degree We will be interested in the evaluations of the above polynomials over the set
VN in the Chebyshev polynomial basis (see (47)).
A def
(74)
where w VN, is defined in (??).
The p-variate version of Monarch matrices M' E NNN as follows. (See Appendix D.6.4 to see how these are exactly related to the definition in (1)). For every row index E [0, VNY and column index j E [0, &N),, we have
p-1 II ( a=u
(WEN..).
M'
[.j
(75)
To express the above in terms of polynomial basis will need the following definition. For any S a < and define the Lagrange basis polynomial A (a) (Xo...., Xa-1) such that for any 0 < mo. Ma-1 < VN, we have
1 (a)
VN VN.m...) W.,
{1 =
if to = i mı...., mi....,ia-11 Ma-1 otherwise.
We use the above to convert the polynomials in (73) to not depend on the sub-indices in (at least for the definition in (75)). For every and 0 I a < p, define:
( (Xo...,,X) Xa)=
Σ
1 (No..., N.-1) gla) (Xa)-
( 1. ba...).........---- VN)
Note that the summation fixes the last sub-indices in i since the definition of elo) (Xa) only depends on (io, da- 1). This implies that for any 0 < io, da < NN, we have
vN = ((4) (WeT
(76)
$6 gives the well known bit reversal permutation.
51

---

We ready define basis p-variate polynomials. For any index E [0, VNN,, define
p-1
of .o......pp-11.
II
gla) XXa)
(77)
a=D
Then (75) can be re-written
Mii.j] = of (We............................. VV....,,,, VW..
(78)
The above leads to the following result, which generalizes Theorem to general
Theorem 15. Let M', A and gj (Xo....,Xp-11 be as defined in (75), (74) and (77). Then for any vector u, M. u is equivalent to evaluating the polynomial
(X Xp-1) = Σ Uj 4j Xp-1))
(79)
at each point AP.
Proof. By (78), the column [M is exactly the evaluation of the polynomial (77) at each point in A". Then the claim follows from the definition of matrix vector multiplication and (79).
This then leads to the following result, (which generalizes Theorem 2):
Theorem 16. For matrices Mg, M1 M2, each of form in Theorem 15, the operation
f = M -(M, k) @ (M2-))).
(80)
is equivalent to representing the polynomial
(X)= . xX)
mod (TuN .o...... Ten (Xp-1))
in terms of the basis polynomials
where k(X), (X) are defined in terms of the respective basis polynomials corresponding to M1 and M2 in Theorem 15, and 0j (X)s corresponds to Mo.
Proof. Define
(a(2) = Z
(Z a).
(81)
asA
Then (80) follows since for
f(X) k: l(X) (X) mod (qa(Xu),. (XX(Xp11),
have the following for any AP we have:
f(a) k(a) u(a))
Using the known fact that T (Z) N- (z (R.e) QA(Z), the claim follows from
c-=0
Theorem 15, (80), and the invertibility of Mp.
D.6.3 Generalizing Theorem for p 2
To convert Theorem 16 into causal map we basically have blow up n 2n.n. Further, paralleling (48) we need to change the definition (73) to have degree V/N ja instead of the earlier N
VN- 3-
Taa ji
(X.)
=
Σ
7a)
(1)
[m]]
Ti (Xa).
(82)
Y02 0=
52

---

Note that now the RHS only depends on and Ja (let us call the RHS Z VN)
(Xa)). the next
definition becomes easier:
githt (Xo..... X.) = c/ VN)
(Xa).
We are now ready define our causal basis polynomials. For any index j. define
p-1
(Xo...., Xp xp-1
II
(83)
These polynomials form a structured subclass of the polynomials defined in (77).
Lemma 17. The class of polynomials qj :X.....; defined in (83) are a special case of (77).
Proof. This follows from the fact that (82) a special case of (73) for every i, < p.
VN I
We show that the product of two Dj (X....., x-- type polynomials can be written has linear combination of qmm (Xu, ................................ Xp-1) with the indices of being lexicographically larger than the original indices.
Lemma 18. Let qj
(Xo ( Xp-1) be defined as in (83). Then for any j.j 0,
2
(Xo: Xp-
=
Σ
agty N (Xo..... Xp-1)
J+J' (m (NN)
(84)
for some set of coefficients Og.j.mm
Proof. From (83) we have,
p-
93 (Xa)
{al
(o......--11 oj
0.
2/2
(Xa)
a=0
(85)
Let us fix 0 a <.p.
Because (82) is of the same form as in (48), we can apply Lemma 8 to each product iy
(Xa).
(al NN
(Xa), which gives us
%/4
{
(Xa)=
X aa.
(Xa) =
VN-1 Σ
(0)
La. VN)
(Xa).
jn
3/2
Ja ta
ma =ja +3
Going back to (85), we get
p-1 N-1
(Xo...,,X, Xp-1) Tj
(Xo
X,-1)
II
Σ
aa. VN) (X₂)-
jatik.ma
0=0 mg =Jatt
Let 0j+J.m
[Pp-]
(a)
Then we get
of
Xp-1 of
Xp-1))
Σ
Ojt5'. m Ahd (No, Xp-1),
i+3' Sme[0, WNY
as desired.
53

---

We now define the following padding scheme, PAD (k).
Algorithm 5 PAD(k)
Input: k R
indexed as kj for [0,
EN
2
Output: k' E
R
VN
0, VN
1: for j
E
do
2:
if ja
for 0 S <
then
3:
Kk ((o..do--))
f k
63
4: else
5:
kj= = 0
6:
return k'
The above basis and padding scheme allows us to extend Theorem to general p.
Theorem 17. Fix family of basis polynomials oj" (X) defined in (83). Let N 2 1 be perfect power of p. n <
NN
u RR" and M'y defined by basis qj (X). Then the operation
MM' (M's PAD (k) M'y PAD(()) I0:n-
(86)
defines a causal map in u.
Proof. Let PAD (k) u'
PAD (u), and
f = (M'' (M'v k' o M'y u')) 10 -1].
(87)
(0, ('N) P
In order to prove that (86) is causal in the input E R", we must show that for all i E dependent only on uy for i' =
f
By Theorem 15, M k' and M'y ut correspond to the evaluations of the polynomials
k' (Xo...., Xx-1))
Σ
Ky of (Xo..... .p.,)),),, and
JE VT))"
n' Xp-..),,p--))> = Σ
(88)
J' ('N)"
respectively, Let us define
..,....),-,)). =
Σ
kj
Oj
and
JE 0.
u Xp-1) =
Σ
Uj
oj
Xp-1)..
(89)
je [0.
Let us define deg the component univariate polynomials Tla)
0j N =
(VN- 100((
ja -
as the vector of length p consisting of the degrees of
a=0
j,i (Xa) defined as in (82). Then we have deg
Oj
54

---

WN
2 -j
VN
Jp-1
for (jo. -p--)) such that 0 S jo
V
for
2
S a < p. Further, for j
j0
VN
VN
we have
2
2
dog (gF))
VN
j/ -11, ,.... VN 8p-1 -1)
=
VN
VN
VN
J0 1,..., VN
3p-1 1111
2
2
V/
P
VN
.0
5p-1
2
2
UN
Since deg
= deg
ff
we can set
((0) 3,1 (Xa) J (Xa). Similarly, note that for j.
7(a)
as above, K2z = KÚ and uzj Uj- Then follows that (Xo...., Xp-1 1) k (Xo..... Xp- 1), and
by a similar argument, u (Xo: .., xp-1) = u (Xo:... : Xp-1). Then by Theorem 16 we have
(Xo Xp-11) = (Xo. ................ Xp-11 u (Xo. Xp-1) mod (TyN (Xo) ..... T TwN (Xp-1)) l C......,ppp111 u (Xo..... Xp-1)
(90)
where the second line follows by observing that each 0 S < p, have degx, (R(X) u(X))
2
N
VN. We want write f (Xo: ..... Xp-1) the form and observing that NN
2
f (Xo. Xp-1)=
fm In (pp-))
mE [0, V'N)"
for a set of coefficients fm. From (89) and (90) we get
f (Xa,...,X,, Xp-1)= =
Σ
(No...., xp-1) . Oj
Xp-1).
JJ' u.
Then by Lemma 18 we have
(Xu, Xp-1)
aj+J'.m Ih Xp-1))
J 0. 2
J+3' mm &NY
Thus, for any m E (o, vNN we have
fm Σ ajtJm m ky
J+J'
implying that fm depends only uj for = m, desired.
D.6.4 p-variate Monarch
Recall that we have fixed VN (and hence N 6P).
Define the p-vrriate Monarch matrix as follows:
p--
M'
P.R. 6,N
I[
B
(P (Pa++ (p0+1,002,00 6a0z (N)
Bo:
(91)
Where each B, is block diagonal with b x blocks for every 0 S a < p. Recall that Equation (1) has a permutation Po the end while the above definition does not have any permutation the end. One trivial way to show the equivalence of above to Equation (1) is to define Po I. In Appendix D.6.5,
55

---

show that exists other non-trivial choices for Po. Further for < p. the P, in Equation (1) connect to the above definition as follows:
(P (o
for 1 S 1 p for p
P
P
Finally we connect the above definition of p-variate Monarch the earlier definition based on polynomial evaluation:
Lemma 19. Equation (75) can be written as Equation (91).
Proof. In (91), B, would correspond to the evaluations j11 (w (WNN..) from earlier. Specif- ically the following holds for any 0 I (io...., fa-1). (pp-1...., and the row and column 'offsets' within each such block are indexed by &a and ja respectively. Connecting back to ( (+) set
p. We index the g0-1 blocks of B. by
B ( 1n [aa Jal ((a) (w: (WUR...)
(92)
We will prove the claim follows. Let ej RV have 1 in the jth location and O's elsewhere. Our goal is to multiply this vector the right of Equation (91) and show that get the j"h column of M' as defined Equation (75). will do so by induction.
Define ej and y1 = Bu Yo. Then for every
S < define
Yot1 Bo (Por are a .11 N) Ya-
Note that we have
M' Cj = P b,N Nv.
(93)
Next, we claim that for every S S p, we have for every (io....,áa-1) E [0, (NNa-1
Ya Illooo................... ................................ in-2), (jp-1. ... ..) in
(WNN.)))
(94)
j.a
b=0
where for = p, we think of ((io, ..., 2a-2), (p ,a)) 20-1) (io, ... 1p-1).
We first note that Equation (94) is enough to prove the claim. Indeed we have that
p-1
= II 1 (W
Yp
b-
where the RHS in the above is the same as RHS in Equation (75). To the claim note that by Equation (93) then
definition of PR b, we have (P& b,N yr) [(ip-l,..., io)l
established the claim.
To complete the proof we prove Equation (94) by induction on a.
We next consider the base case of a 1. Note that y is just the jth column of By (as yo ej) and that case Equation (94) follows from Equation (92).
For the inductive hypothesis. assume that Equation (94) is true for for some a 2 1. We now want argue Equation (94) for Ya+1. Towards that end define
2, = (P bo-at1, Ya = P P
Note that
Yat1 B Za-
(95)
Now by definition of Po.b--a41, N. we have
Za I( 20-1) (pp-1........) Ya [[(io.... ...., ia-2), (1p-1..... .ja) fia-1)]
...
56

---

We claim that Equation (94) is true for a + 1 from the above along with Equa- tion (95) and Equation (92). Indeed, fix any (fo.-...aa-1). Then in Equa- tion (95) the entry Zu
(p--t.....aa)) gets multiplied by the entries B (Ha 1):(ip i---da) [a Ja] for all values of ia E [0, VNN.. The inductive hypothesis and
Equation (92) then proves the inductive step, as desired.
Finally, we note that we can generalize Algorithm 4 for constructing B, for 0 S a < p (since this is the multivariate case some of the steps bit simplified):
Algorithm 6 BLOCKY MULTIVAR MONARCH(B (0). B B N.
Input: (oo)..... B (i--1) R where b = VN Output: Block diagonal matrices Bo.... Bp- E N
b Compose each output matrix from corresponding input matrix
1: for < 0top do
2: for t to do
3:
B B C B.BB
lbb ab + 1
4: By diag(By (0)
B 1)
5:
return Bo..... B Bp
D.6.5 Extensions and Open Questions
In this sub-section we outline certain (fairly) straightforward extension of our theoretical results and conclude with some open questions.
Comparing Equation (91) to Equation (1) We note that we can post multiply M in Equation (91) with large class of permutations for which Theorem 17 still holds. We outline the technical reason why this is true. At the heart of the argument for why Equation (86) gives a causal map is Lemma 18. Specifically note that the sum in RHS in Equation (84), is over all j 3 m. The main observation is that this partial order still holds if we permute the b-variate representation of j.J and m in the same way. In other words, for any permutation : [0. p) + [0. p) if we define (j) (Jo(o))..--Jopp)) and similarly (J'), (m). Then still have a (j) a (J') I (m). This in turn implies the following. Let Pa be a permutation that maps E [0, V/ to (j) [0, VN Then Theorem 17 holds if we replace M' by M Pa with M as in Equation (91).
Our results as presented are for specific classes of evaluation points. natural question to ask is if our results can be extended to more general set of evaluation points. It turns out that our results for p-variate Monarch matrices can be extended to a wider class of evaluation points. Specifically, for each S a < p. let S. C C with |Sa| = VN. Then our results in this sub-section hold if we replace the evaluation points from AP to S= The only thing that changes our proofs that in Theorem 16, we replace mod (TyN (Xo) ,..., T Tan (Xp-1)) by mod (ss (Xo) .... 4sp- (Xp-1)), where (2) is as defined in Equation (81). This result can then be propagated throughout the rest of our proofs.
Evaluation points
On the other hand, our results in Appendix D.3 and Appendix D.5 do exploit specific properties the evaluation points (specifically (wxx) v/N = - for Appendix D.3 and Ty (WN,s = (-1)" vN for Appendix D.5). To generalize these results to other sets of evaluation points, we need the existence of degree VN polynomial that maps (in VN-to-1 fashion) A to a set of VNV elements. Another interesting open question is to avoid the blowup n - 2" n Theorem 17 and ideally only pay a blowup + 2n for every p 2 2 as were able to do in Appendix D.3 and Appendix D.5 (with p == 2).
E Implementation
57

---

from einops import rearrange import torch from torch import nn
def blockdiag.matzul return torch einsum
w):
"bnm, bbnm.... bm->... bn", w, x. (xx. 1:=1], w. shape [0], v.
(*x.
shape (-1]) reshhppe (#x. shape)
10
class MonarchMatrix (nn Module):
11
12
def
-_init__ (self, sqrt_n: int): super Ou_pit__._ self sqrt_n sqrt_n
15
self.L nn. Parameter (torch. randn ((sqrt_n, aqrt.n, sqrt_n))) self.R mn. Parameter (torch.. randn ((sqrt_n, sqrt sqrt_n)))
16
def forward x):
19
x rearrange (x, x blockdlag_natmullt, self.L)
(m n) ->
(n m)",
20
x rearrange (x, (m n) ->
(n m)", neself. sqrt_n)
blockdiag mlockdiag.matmul (x, self.R)
(earrange
(m
(n m)", n=self sqrt_n)
24
25
class MonarchMizerLayer Monule)
26
def -_init_- (self, sqrt_n: int, sqrt_d: int):
super () 1 iniit_nnn
28
self
HonarchMatrix (sqrt_n)
29
self m2 MonarchMatrix (eqrt__))
30
self m3
(mnarchMatriixsqqrt..d (sqrt_d)
self m4 ManarchMatrixssgrt...
32
33
self n_kernel nn Parameter (torch. randn (sqrt_d 2, sqrt_n
** 2))
34
self d_kernel nn Carameter(torde randn (1, sqrt_d ** 2)) self layer_norm Layerloor (sqrt_d ** 2)
36
37
def (self, x: torch Tensor): shape (b, n, d) self ml (x. transpposss (-1, -2))
38
x_tilde self melf.m2 (seff.fsee....eereee
) transpoe (-1, -2) # mix sequence
y self m4 serchhrelu(sel1 kernel self m3xxttilde))
mix features
40
return self layer_norm (y x_tilde) skip connection
Listing 1: A basic implementation of the M2 layer.
58


Reviews:
[{"review_id": "Rr98dVayUb", "reviewer": "Reviewer_xUjR", "summary": "The Monarch Mixer (M2) combines MLP mixer and Conv mixer and yields a new family of mixers that is formalized in terms of Monarch matrices. the approach is novel and reminds me of an extension of [15] in their reference. \n\nThe main advantage of M2 is in its sub-quadratic computation capability. The authors evaluate their method on a set of large language models and prove their framework functions properly on meaningful and challenging tasks. The paper is well-written and easy to follow. ", "questions": "- Table 3 mentions a \"Transformer\" model. What is the configuration of the model?\n- Tables 4 to 7 mention BERT-s. What is BERT-s? What are the details of its architecture?\n- Table 2 seems to be incomplete. The FLOP utilization for dense matmul is missing from Table 2.\n- Caption of Table 4 mentions GPU, however the body of the Table contains no inference results on GPU. Did the authors perform inference on CPU only? Why GPU inference is not reported in Table 4?\n- A comparison with SwinMLP-B [1] and other SWIN-v2 models is more informative than the relatively older ViT model. \n- What is BERT-s? What is ViT-s? please provide the model details.\n\n[1]: Zheng, Hao, Guohui Wang, and Xuchen Li. \"Swin-MLP: a strawberry appearance quality identification method by Swin Transformer and multi-layer perceptron.\" Journal of Food Measurement and Characterization 16.4 (2022): 2789-2800.", "limitations": "While the method clearly shows that the method works well on Transformer models, convolutional models, and speech applications are ignored. \n", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper stands out in its clear and proper presentation. The idea is well motivated from a hardware perspective, explained intuitively, and studied theoretically. Evaluation of the method on causal and non-causal large language models is an asset and demonstrates a clear potential of the method.", "weaknesses": "- The new approach is only benchmarked on Transformer tasks, while speech applications look relevant but are ignored.\n- It would be fair to see a comparison of inference latency and metric performance of the M2-based BERT with a configuration of BERT that has a smaller number of parameters, but the same number of FLOPs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and insightful suggestions. We hope that the additional experiments reported in the common response have improved the paper, and we look forward to further discussion. Here, we answer the specific weaknesses and questions raised in your review.\n\n**W1. Evaluation on speech applications.** Thank you for the suggestion to evaluate on speech applications. Please see the experiment on speech commands in the common response. We hope this result helps provide further evidence for the generality of our method.\n\n**W2. Comparison against parameter-matched BERT models**. Thank you for your suggestion to evaluate M2-BERT against parameter-matched BERT models. We have reported the results of these experiments in our common response. The 80M BERT model underperforms M2-BERT in quality on GLUE. A highly-optimized 80M BERT achieves higher throughput than M2-BERT on short sequences, but underperforms on inputs longer than 1K.\n\n**Q1. Architecture in synthetics.** For these synthetics, all models, including the Transformer, are small two-layer models with hidden dimension 64.\n\n**Q2 & Q6. BERT-s, ViT-s.** “BERT-s” and “ViT-s” in the submission are typos, we mean BERT-base and ViT-B (the -s terminology is the equivalent model for GPT2-s). These are now fixed in the draft.\n\n**Q3. Dense Matmul FLOP Util.** We report FLOP utilization of dense matmul here, and we have added it to our updated manuscript:\n\n|      |   **4K**  |  **16K**  | **65K**  |\n|-----:|:-----:|:-----:|:-----:|\n| **A100** | 63.0% | 78.0% | 80.0% |\n| **4090** | 74.6% | 96.7% | 97.9% |\n\n**Q4. GPU Inference.** We have reported the GPU inference results in the common response, and we will add them to Table 4 in the updated manuscript.\n\n**Q5. Swin Comparisons.** Thank you for the suggestion to compare to Swin-v2 and Swin-MLP. We have run these experiments and reported them in the common response (Table **4** in the accompanying PDF)."}, {"review_id": "rGLSYaG98Y", "reviewer": "Reviewer_qwKK", "summary": "This paper takes a fresh approach by addressing the issue of high complexity in current neural networks. It points out that the computational complexity of Transformers is quadratic with respect to both the sequence length and the feature dimension. Previous papers primarily focused on reducing the complexity related to sequence length, but this paper is the first to propose a method that reduces complexity for both sequence length and feature dimension. The specific approach used is the utilization of second-order Monarch Matrices, with the model structure referred to as the Monarch Mixer. Additionally, the authors introduce a novel initialization method for the Monarch Matrices, enabling them to handle causal l language modeling. The effectiveness of this approach is validated in the areas of non-causal language modeling, causal language modeling, and image classification.", "questions": "1. I'm not very familiar with the Monarch Matrix, but is its core idea to use the product of block-diagonal matrices and permutation matrices as a replacement for the dense matrix?\n\n2. What is the motivation behind the Monarch Matrix? Despite having significantly fewer parameters, it seems to perform comparably to dense matrices in small-scale models. Can this conclusion be extended to models larger than 10 billion parameters?\n\n3. In Line 257, \"We set the Monarch matrices to DFT and inverse DFT matrices, to simulate long convolutions [15, 37], and do not learn them.\" Does this refer to setting the block-diagonal matrices of the Monarch Matrix as DFT matrices?\n\n4. The experiments for Non-Causal Language Modeling and Image Classification should be more comprehensive, such as considering the Monarch matrices as learnable.\n\n5. Could you provide a more intuitive explanation of the initialization for the causal scenario? From examining the code, it seems like the Monarch Matrix is initialized as a lower triangular matrix?", "limitations": "Yes.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Indeed, it is novel to consider the optimization of complexity from both the sequence dimension and the feature dimension. Furthermore, initializing the model for the causal scenario poses a definite challenge, and the authors have successfully accomplished this task.", "weaknesses": "There aren't many weaknesses, for specific questions, please refer to the **Questions** section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and questions. These questions have helped us improve the presentation of our paper. We provide a more detailed explanation of Monarch matrices below, which we plan to add to the paper.\n\n**Q1. Monarch Motivation.** The motivation behind the Monarch matrix is to adapt and generalize the FFT algorithm. The FFT algorithm splits an FFT of size N into smaller FFT’s over portions of the input, interleaved with permutations. More precisely, let $F_N$ denote the Fourier transform of size $N$, and assume N is a perfect square for simplicity. The FFT breaks down $F_N$ as follows:\n\n$F_N = P F_L P D P F_R P,$\n\nwhere $F_L$ and $F_R$ are block-diagonal matrices whose blocks are made of $F_\\sqrt{N}$, $D$ is a diagonal matrix, and $P$ is a permutation that reshapes a 1D input into $\\sqrt{N} \\times \\sqrt{N}$, and takes the transpose.\n\nA Monarch matrix generalizes this computation pattern by “rolling in” the diagonal matrix, and letting the blocks in the block-diagonal matrices be arbitrary instead of fixed to an FFT:\n\n$M = PLPRP$\n\nThis additional flexibility allows Monarch matrices to express a wider class of structured matrices than the FFT (but they are not as completely expressive as a dense matrix). In our paper, we also generalize Monarch matrices past the order 2 phase, so there can be more than two block-diagonal matrices interleaved with permutations (Figure 1, left in our original submission).\n\n**Q2. Scaling Results.** We have seen promising initial scaling results – in the common response, we have reported results for both M2-BERT-Base and M2-BERT-Large. In our original submission, we also saw similar scaling performance on causal language modeling with GPT2-s and GPT2-m equivalent models (Table 9 in the main paper). We look forward to continuing to scale these models and seeing how well the trends hold.\n\n**Q3. Monarch DFT.** A Monarch matrix can exactly express both a DFT and an inverse DFT (see Appendix F, corollary 6 in the original submission for the exact parameterization – they are similar to DFT matrices, but with slight modifications). For these experiments, we set the Monarch matrices to express the DFT and inverse DFT, respectively.\n\n**Q4. Learnable Monarch.** Thank you for the suggestion to extend the experiments for learnable Monarch matrices. Please see the common response for additional experiments along these lines. Making the Monarch matrices learnable yields small benefits in quality.\n\n**Q5. Causal Monarch Interpretation.** One way to interpret a Monarch matrix is to view it as evaluating a polynomial at a set of evaluation points $(a_i)$. When the Monarch matrix is used in a convolution, it is equivalent to multiplying two polynomials, by multiplying their evaluations $h(a_i) = f(a_i)g(a_i)$. The causal parameterization ensures that the resulting polynomial $h(a_i)$ is causal in $f$ – i.e., its coefficients do not depend on coefficients of $f$ that come later in the sequence."}, {"review_id": "1hAwaZlhqc", "reviewer": "Reviewer_zN2M", "summary": "The paper introduces a new neural network layer which runs efficiently on modern GPUs, and exhibits strong performance against state-of-the art on several benchmarks. The layer is based on Monarch matrices, introduced in [7]. Monarch matrices use permutation matrices, and block diagonal matrices to represent dependancies across features and temporal dimensions. Inspired by butterfly matrices in FFT, this matrix parameterization can represent anything from convolutions to fully connected matrices, depending on their order. The Monarch Mixer layer introduced in this paper uses two monarch matrices as well as a matrix (K) that is used in point-wise multiplication. Performance is compared on language and image classification, replacing transformer and fully connected layers with the M2 layer. Finally, the paper presents a theoretical derivation of how this layer can be modified for use in causal language tasks, while maintaining its computational efficiency. ", "questions": "-Does the expressivity of M increase with the order p? Perhaps this is obvious, but it should be stated explicitly in the paper. Can M be used to express any dense matrix? \n\n-Why do you think your model exceeds the performance of other state-of-the art models without using attention? I always thought one of the most important aspects of attention was that it afforded permutation invariance, does M2 serve a similar purpose? ", "limitations": "Limitations were addressed in second paragraph of Section 6. ", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Technically, this paper makes a strong contribution by proposing a computationally efficient layer. The computational performance is benchmarked on modern GPUs. I really liked the discussion on factors affecting runtime performance on modern GPUs -- this is a really valuable introduction to this topic. The proposed layer achieves superior performance, with fewer trainable parameters, in less time. Like MLPMixer, it also does away with the attention mechanism of transformers -- which suggests M2 is also a good architectural inductive bias -- perhaps competitive with attention? Comparing against [7], I would say the largest technical contribution of this paper is on how to modify the M2 layer perform in the causal setting.  ", "weaknesses": "The main weakness of the paper is that it is somewhat hard to read/understand without having to read [7]. Section 4 in particular is quite dense. It seems that the authors struggled to fit the paper within the NeurIPS page limit. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback on the technical contributions of our paper. We are glad that you found the GPU performance discussion insightful, and we appreciate your constructive comments on how to improve the paper.\n\n**W1. More intuitive explanation of Monarch matrices.** Thank you for your suggestions on how to improve the clarity of the paper. We plan to use the extra space in the camera ready to include both a more complete description of Monarch matrices and the motivation behind their definition. We plan to motivate the Monarch matrices via the FFT algorithm, as follows:\n\nThe motivation behind the Monarch matrix is to adapt and generalize the FFT algorithm. The FFT algorithm splits an FFT of size N into smaller FFT’s over portions of the input, interleaved with permutations. More precisely, let $F_N$ denote the Fourier transform of size $N$, and assume N is a perfect square for simplicity. The FFT breaks down $F_N$ as follows:\n\n$F_N = P F_L P D P F_R P,$\n\nwhere $F_L$ and $F_R$ are block-diagonal matrices whose blocks are made of $F_\\sqrt{N}$, $D$ is a diagonal matrix, and $P$ is a permutation that reshapes a 1D input into $\\sqrt{N} \\times \\sqrt{N}$, and takes the transpose.\n\nA Monarch matrix generalizes this computation pattern by “rolling in” the diagonal matrix, and letting the blocks in the block-diagonal matrices be arbitrary instead of fixed to an FFT:\n\n$M = PLPRP$\n\nThis additional flexibility allows Monarch matrices to express a wider class of structured matrices than the FFT (but they are not as completely expressive as a dense matrix). In our paper, we also generalize Monarch matrices past the order 2 phase, so there can be more than two block-diagonal matrices interleaved with permutations (Figure 1, left in our original submission).\n\n**Q1. Expressivity of $M$ with the order $p$.** This is a little subtle - the expressivity actually goes _down_ with increasing $p$, since we decrease the sizes of the blocks (block size $b$ is $\\sqrt{N}$ for order 2, $\\sqrt[3]{N}$ for order 3: for general $p$ the block size is $b=\\sqrt[p]{N}$). This in turn implies that a $p$-variate Monarch has $O\\left(pN^{1+1/p}\\right)$ many parameters i.e. the number of parameters decreases as $p$ increases and hence its expressivity goes down (this follows e.g. from a counting argument).\n\n$p=1$ gives arbitrary dense matrix, but that is because in that case the block size is now $b=N$ and hence is a trivial case. For $p>1$, a single $p$-variate Monarch matrix cannot express an arbitrary matrix.\n\nHowever, one can show that one can express an arbitrary $N\\times N$ matrix as a product of $m=O(N)$  matrices $M_1,\\dots,M_m$ where each $M_j$ is a $p$-variate Monarch matrix (this follows from a known result that an arbitrary matrix can be represented a product of O(N) Toeplitz matrices).\n\nWe will add a discussion of these properties to the main body when we introduce $p$-order Monarch matrices.\n\n**Q2. Why does M2 outperform attention-based models?** There are two pieces – first, we build on prior work to build sub-quadratic replacements for attention. Second, we replace the MLPs with sub-quadratic alternatives, which achieves the same performance with the same height/width but fewer overall parameters.\n\nM2 builds on prior work studying how to replace attention with a sub-quadratic alternative while maintaining high quality. Many of these models use a combination of long and short convolutions with elementwise multiplication, e.g. [1, 2, 3, 4]. These convolutions are often computed with an FFT, which means that they can be expressed using Monarch matrices and elementwise multiplication. We build on the insights from these architectures when using Monarch for sequence mixing in our M2 models with an alternative that scales sub-quadratically in sequence length while maintaining high quality.\n\nIn addition, M2 uses Monarch matrices to scale sub-quadratically in the model dimension by replacing MLPs. Our results – that the dense layers in MLPs can be replaced by sparse(r) matrices without losing quality – may suggest that the current generation of models is overparameterized, and that there exist much more efficient architectures to develop. We are excited by these possibilities, and we look forward to building on these ideas in the future.\n\n[1] Long Range Language Modeling via Gated State Spaces. Harsh Mehta, Ankit Gupta, Ashok Cutkosky, Behnam Neyshabur. ICLR 2022.\n\n[2] Pretraining Without Attention. Junxiong Wang, Jing Nathan Yan, Albert Gu, Alexander M. Rush. ACL 2023.\n\n[3] Hungry Hungry Hippos: Towards Language Modeling with State Space Models. Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, Christopher Ré. ICLR 2023.\n\n[4] Hyena Hierarchy: Towards Larger Convolutional Language Models. Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, Christopher Ré. ICML 2023.\n"}]

Author rebuttal:
# Common Response
We thank all reviewers for their time and valuable comments, which have helped us improve our paper. In this paper, we introduce Monarch Mixer, a new architecture that is hardware-efficient and sub-quadratic in both sequence length and model dimension. We demonstrate that Monarch Mixer can be used as a drop-in replacement for attention and MLP in Transformers in BERT-style, GPT-style, and ViT-style modeling, matching quality with up to 27% parameter reduction and up to 9x faster inference for long sequences. We are excited to take a first step with this work in developing new architectures that are fundamentally more efficient than Transformers while maintaining quality.

We are glad to have received positive feedback on the motivation (reviewers **xUjR**, **zN2M**), clear technical contribution (reviewers **qwKK**, **zN2M**), execution and experiments (reviewers **xUjR**, **zN2M**), and overall clarity of presentation (reviewers **xUjR**, **qwKK**).

In our general response, we are excited to report the results of some additional experiments in pretraining quality, as well as experiments requested by the reviewers:
* Stronger BERT results:
  * M2 achieves GLUE performance that matches BERT-Base from (Devlin et al 2018), with 27% fewer parameters and up to 9X faster GPU inference time for long sequences (requested by reviewer **xUjR**)
  * Scaling up to BERT-Large equivalent – M2-BERT-Large matches BERT-Large in quality with 24% fewer parameters and achieves up to 4X faster GPU inference time (demonstrating scaling performance, as requested by reviewer **qwKK**)
  * Benchmarks against smaller BERT models (requested by reviewer **xUjR**)
* Swin-M2: matching ImageNet accuracy with 32% fewer parameters when replacing attention and MLPs in Swin-V2 with Monarch Mixer, as a drop-in replacement (requested by reviewer **xUjR**)
* Speech: M2 matches state-of-the-art in classification accuracy when classifying raw 16 kHz speech signals on the SpeechCommands task (requested by reviewer **xUjR**)
* Experiments with learnable Monarch matrices in the sequence mixer for CIFAR, achieving 1.5 points in lift from learning the matrices (requested by reviewer **qwKK**)

The results tables for these experiments are in the accompanying PDF for the rebuttal. We plan to include these experiments in our updated manuscript and welcome further discussion on how to improve the paper during the discussion period.

## BERT
Since our initial submission, we have improved our BERT pretraining formula via improvements to pretraining hyperparameters and data. We are happy to report stronger downstream GLUE results, competitive with those from the official BERT-Base model trained by (Devlin et al 2018). We have also included an 80M BERT trained using our formula as a baseline, following the suggestion from reviewer **xUjR**.

Table **1** in the accompanying PDF shows the results. M2-BERT-base is competitive with the Devlin et al BERT-base, with 27% fewer parameters. When scaled up to parameter-match BERT-base, M2-BERT outperforms BERT-base by 1.3 GLUE points on average.

We have also scaled up to a BERT-Large equivalent, which is competitive with the BERT-Large trained by (Devlin et al 2018). Table **2** in the accompanying PDF shows that M2-BERT-large is competitive with BERT-large with 24% fewer parameters. These results suggest that M2 can scale well.  We are excited to scale up to even larger models in future work (as suggested by reviewer **qwKK**).

**BERT GPU Inference**

Next, we report additional results on GPU inference times for different sequence lengths (addressing a question by reviewer **xUjR**). Inference times are reported as throughput, in tokens/ms on a single A100-40GB.

Table **3** (top) in the accompanying PDF shows that M2-BERT-base achieves higher throughput than even highly-optimized BERT models, and achieves 9.0X higher throughput than a standard BERT-base model for long (4K) input sequences.

Table **3** (bottom) in the accompanying PDF reports throughput against the 80M BERT as well. When parameter-matched, M2-BERT is slower than the most highly-optimized attention kernels for short sequences, but still faster for long sequences.

Benchmarks for BERT-large and M2-BERT-large are omitted for space in the rebuttal, but the trends are similar (but the HF model OOM’s earlier, leading to a maximum speedup over HF of 4.34X at sequence length 2K).

## Swin-M2 Experiments
Reviewer **xUjR** suggests comparing ImageNet performance against Swin-V2 and Swin-MLP. Table **4** in the accompanying PDF reports the results of replacing attention and the MLP in Swin-V2, using M2 as a drop-in replacement. Surprisingly, Swin-M2 outperforms Swin-MLP-B, is competitive with Swin-V1-B, and comes within 1 point of Swin-V2-B – even without any hyperparameter tuning or architecture adjustment from our ViT formula.

We expect that performance may improve further with hyperparameter tuning specific to M2. These results provide additional evidence that M2 is a strong drop-in replacement for attention and MLPs in various architectures.

## Speech Applications
Reviewer **xUjR** suggests evaluating M2 on speech applications. Table **5** in the accompanying PDF presents the performance of M2 on Speech Commands-10, a speech classification task over raw 1-second clips sampled at 16 kHz. M2 is competitive with state-of-the-art architectures on this task.

## Learnable Monarch Matrices in Sequence Mixer
Reviewer **qwKK** suggests extending the experiments to include the case where the Monarch matrices are learnable. In our submission, the monarch matrices in the state mixer (i.e. MLP) are already learnable. Table **6** in the accompanying PDF presents an experiment evaluating sequence mixers (i.e. attention) with learnable monarch matrices on sequential CIFAR. Learning the Monarch matrices in the sequence mixer yields 1.5 points of lift. We look forward to further exploring this regime in future work.

Meta-review: This paper proposed a new architecture called Monarch Mixer that is able to achieve sub-quadratic cost in many GEMM-based models including ViT, non-causal and causal language models. 

The reviewers generally agree that this work has made nontrivial contribution to the field, and this paper made a timely contribution to the increased computation need for large models. 

We also appreciate the additional experiments made by the authors, which we would also encourage the authors to include in their camera ready version.
Decision: Accept (oral)