Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture

ArXiv: 2310.12109

🎯 Pitch

Monarch Mixer matches Transformer quality across BERT, ViT, and GPT-style models—without using attention or MLPs—by replacing both with a single class of structured matrices that scale sub-quadratically, achieving up to 9.1× higher throughput at long sequences.

---

1. Executive Summary

This paper introduces Monarch Mixer (M2), a new neural architecture that achieves sub-quadratic scaling in both sequence length and model dimension by using Monarch matrices—a structured class of expressive, hardware-efficient matrices that generalize the FFT—as the single mixing primitive along both axes (replacing attention with gated Monarch convolutions and replacing MLP dense layers with block-diagonal Monarch factors). Evaluated on BERT-style masked language modeling (C4/GLUE), ViT-style image classification (ImageNet-1k), and GPT-style causal language modeling (The PILE), M2-BERT matches BERT-base and BERT-large downstream quality with up to 27% fewer parameters while achieving up to 9.1× higher throughput at sequence length 4K, M2-ViT outperforms ViT-b by 1% accuracy with half the parameters, and M2-GPT matches Transformer perplexity at 360M parameters without using either attention or MLPs. The paper also develops a novel theoretical characterization of Monarch matrices as multivariate polynomial evaluation and interpolation, deriving sufficient conditions under which M2 convolutions are causal while remaining sub-quadratic—establishing for the first time that it may be possible to match Transformer quality in auto-regressive settings without attention or MLPs, provided the causal parameterization constraints on polynomial degrees are satisfied.

2. Context and Motivation

The Core Problem: Quadratic Scaling Along Two Dimensions

The fundamental problem this paper addresses is structural: modern machine learning models, particularly Transformers, scale quadratically along both critical axes—sequence length and model dimension. This quadratic scaling creates a compounding bottleneck as practitioners push toward longer contexts (to enable reasoning over larger documents, codebases, or multimodal inputs) and larger model dimensions (to improve representation quality and downstream performance).

The quadratic scaling manifests in two separate but equally constraining ways within a standard Transformer block:

• Attention scales quadratically in sequence length: For an input of $N$ tokens, the attention matrix is $N \times N$, requiring $O(N^2)$ memory and computation. This means doubling the sequence length quadruples the cost. Even with optimized implementations like FlashAttention (Dao et al., 2022), which achieves ~25% FLOP utilization, the $O(N^2)$ FLOP count remains—it is only the constant factor that improves. At very long sequences (8K, 16K, 32K tokens), attention dominates runtime and memory consumption, making long-context deployment expensive or infeasible.

• MLP layers scale quadratically in model dimension: The feed-forward sublayers in Transformers contain two dense matrix multiplications, each requiring $O(d^2)$ computation for model dimension $d$. Since modern models routinely use dimensions of 768 (BERT-base), 4096 (Llama-7B), or larger, these MLP layers consume substantial compute—and that compute grows quadratically as model width increases. MLPs are compute-bound operations that achieve ~95% FLOP utilization on GPUs (Table 1), so their cost is large and scales rapidly.

The combined effect is that Transformers are quadratically expensive in both $N$ and $d$, creating a two-front scaling problem. Much prior work has addressed one front or the other in isolation—sub-quadratic attention for sequence length, sparse MLPs for model dimension—but no single architecture had tackled both simultaneously with a unified primitive.

Why This Problem Matters

The paper identifies several converging trends that make this dual-quadratic bottleneck increasingly urgent:

The push toward longer contexts. Language models are being deployed with dramatically longer context windows—from 512 tokens in early BERT, to 2K in GPT-3, to 32K–100K in Claude and GPT-4, and beyond. In both language and vision, longer sequences enable new capabilities: document-level reasoning, multi-turn dialogue, high-resolution image processing, and video understanding. If architectures scale quadratically in sequence length, these capabilities come at steep and growing computational cost, limiting who can afford to deploy them.

The demand for higher-quality representations through wider models. Model dimension is one of the primary levers for improving representation quality—wider models can learn richer features, support larger vocabularies, and capture more nuanced patterns. The quadratic scaling of MLP layers means that increasing model width by a factor of $2\times$ increases the per-layer FLOP cost by $4\times$, making incremental quality improvements disproportionately expensive.

Hardware utilization is equally important as asymptotic complexity. The paper emphasizes a subtle but critical point (Section 2, Table 1): asymptotic scaling alone does not determine real-world speed. Operations must also achieve high FLOP utilization on modern accelerators to realize their theoretical efficiency. The FFT—used in several recent sub-quadratic sequence mixers—scales as $O(N \log N)$ but achieves only 3–5% FLOP utilization because it is memory-bound, spending most of its time moving data rather than computing. In contrast, dense matrix multiplications achieve 75–98% FLOP utilization. An architecture that is sub-quadratic on paper but achieves poor utilization can be slower in practice than a quadratic architecture with high utilization, at least at moderate sequence lengths.

The training-inference tension. During training, both sequence length and model dimension matter. During inference, particularly for auto-regressive generation that processes one token at a time, the model dimension cost dominates. An architecture that addresses only sequence length (e.g., linear attention) still leaves the MLP bottleneck unaddressed, limiting deployment efficiency for both training and inference workloads.

Prior Approaches and Where They Fall Short

The paper surveys three broad families of prior work that attempt to address parts of this problem, each with characteristic limitations:

Sub-Quadratic Sequence Mixers (Attention Replacements)

Recent work has proposed replacing attention with long convolutions (Fu et al., 2023; Poli et al., 2023; Romero et al., 2022) or state space models (Gu et al., 2021; Hasani et al., 2022). These approaches achieve $O(N \log N)$ or $O(N)$ scaling in sequence length by exploiting the FFT convolution theorem or structured state space recurrences.

Where they fall short: The FFT-based convolutions are memory-bound and achieve poor FLOP utilization (3–5% on GPUs; Table 1, citing Fu et al., 2023), meaning their favorable asymptotics may not translate to wall-clock speedups at moderate sequence lengths. More critically, these methods only address the sequence dimension—they still use standard quadratic MLPs for the model dimension mixer. They solve half the problem.

Sparse and Structured MLPs

Another line of work addresses the model dimension bottleneck by sparsifying MLP layers through pruning (Frankle & Carbin, 2018; Han et al., 2015), low-rank factorization (Zhu et al., 2021), or structured sparse patterns such as butterfly matrices (Chen et al., 2021) and Monarch matrices (Dao et al., 2022). Monarch matrices in particular were shown to be expressive—capable of representing a wide class of linear transforms including convolutions, Toeplitz matrices, and orthogonal polynomials—while scaling sub-quadratically at $O(N^{3/2})$ for sequence length $N$ (with order $p=2$).

Where they fall short: As the paper notes in Section 2, "some of the models can actually be slower than their dense counterparts, due to low hardware utilization." Sparse patterns often introduce irregular memory access, reducing FLOP utilization below what the reduced FLOP count would predict. More importantly, this line of work focuses exclusively on the model dimension—it does not provide a sub-quadratic sequence mixer. The two communities (sub-quadratic attention and sparse MLPs) had largely been developing independently, with no unified architecture.

The FFT-Based Convolution Bottleneck

The paper flags FFT-based convolutions as a particularly instructive case. The FFT convolution theorem states that convolution can be computed as:

$y = K * X = \text{FFT}^{-1}(\text{FFT}(X) \odot \text{FFT}(K))$

This is $O(N \log N)$ in theory but suffers from two practical problems:

1. The FFT is memory-bound—each stage of the Cooley-Tukey algorithm involves data movement across memory hierarchies, and the irregular access patterns of the bit-reversal permutation limit cache effectiveness.
2. It requires complex arithmetic even for real-valued inputs (or careful packing to exploit conjugate symmetry), increasing the constant factor.

The paper's Table 1 quantifies this: at 64K input dimension on an RTX 4090, an MLP achieves 95.5% FLOP utilization, FlashAttention achieves 24%, and an FFT achieves only 3%. The M2 convolution operator (Equation 2) aims to hit a sweet spot: $O(N^{3/2})$ complexity (worse than FFT asymptotically) but 41.4% FLOP utilization (dramatically better than FFT), yielding a net wall-clock speedup.

The Core Insight: A Unified Primitive

The paper's central observation—inspired by work such as MLP-Mixer (Tolstikhin et al., 2021) and ConvMixer (Ng et al., 2022)—is that many modern architectures can be understood as repeatedly mixing information along the sequence axis and the model dimension axis. An MLP-Mixer, for instance, alternates between "token mixing" (an MLP applied across the sequence dimension) and "channel mixing" (an MLP applied across the feature dimension). ConvMixer does the same with convolutions—depthwise convolution for spatial mixing, pointwise ($1\times 1$) convolution for channel mixing.

This "mixer" perspective suggests a natural question: can we find a single primitive that efficiently performs mixing along both axes, with sub-quadratic scaling and high hardware utilization on both?

Monarch matrices are a natural candidate for this unifying role because:

• They are expressive: they capture FFT, convolutions, Toeplitz transforms, orthogonal polynomials, and more (Dao et al., 2022).
• They are parametrically efficient: an order-2 Monarch matrix has $2N^{3/2}$ parameters (from two sets of block-diagonal matrices with $\sqrt{N} \times \sqrt{N}$ blocks), compared to $N^2$ for a dense matrix.
• They are hardware-efficient: the block-diagonal factors can be computed using standard GEMM operations on blocks of size $\sqrt{N}$, which map naturally to tensor core instructions and achieve high utilization.
• They still haven't been systematically applied to the sequence mixing problem, particularly in causal settings—the original Monarch paper (Dao et al., 2022) focused on model dimension compression.

How This Paper Positions Itself

The paper explicitly positions M2 as a proof-of-concept exploration, not a fully optimized production system. It frames its contribution along three axes:

1. Architecture unification. M2 is the first architecture to use the same sub-quadratic primitive (Monarch matrices) for both sequence mixing and model dimension mixing, replacing both attention and MLPs. This is not just an efficiency argument—it is a conceptual simplification. Where Transformers require two qualitatively different mechanisms (attention for cross-token interaction, MLPs for per-token feature transformation), M2 uses structurally identical operations, differing only in which dimension they operate over.

2. Theoretical foundation for causality. The paper develops new theory (Section 4, Theorem 3) showing that Monarch convolutions can be made causal—i.e., the output at position $i$ depends only on input positions $0$ through $i$—while remaining sub-quadratic. This is non-trivial because the natural Monarch parameterization is not causal: the polynomial product $k(X,Y) \cdot u(X,Y)$ undergoes modular reduction (modulo $X^{\sqrt{N}}-1$ and $Y^{\sqrt{N}}-1$) that can "wrap around" higher-order terms to lower indices, violating causality. The paper's solution—interpreting Monarch matrices as multivariate polynomial evaluation and imposing degree constraints on the basis polynomials—enables sub-quadratic causal mixing for the first time, which is essential for auto-regressive language modeling.

3. Empirical breadth across domains and paradigms. Rather than focusing on a single benchmark or architecture variant, the paper tests M2 across three distinct settings—non-causal BERT-style masked LM, ViT-style image classification, and causal GPT-style auto-regressive LM—to demonstrate generality. The results in each domain aim to match or exceed Transformer baselines while using fewer parameters and/or less computation, establishing that the architecture is viable across different training objectives and modalities.

The Hardware-Motivated Design Philosophy

A recurring theme in the paper's motivation is the tension between asymptotic complexity and hardware reality. Section 2's preamble on GPU accelerator cost models (compute-bound vs. memory-bound operations, tensor core throughput vs. general-purpose unit throughput, memory hierarchy bandwidth) is not incidental—it directly motivates the choice of Monarch matrices over alternatives like the FFT.

On modern GPUs (particularly since the V100 introduced tensor cores), matrix multiply operations enjoy a massive throughput advantage: 312 TFLOP/s for half-precision tensor core operations on an A100, versus only 19 TFLOP/s for non-matrix operations. An architecture built around GEMM operations—even if it performs more total FLOPs than an FFT-based alternative—can still be faster because those FLOPs execute at dramatically higher throughput. The M2 convolution at $O(N^{3/2})$ with 41.4% utilization outperforms FFT at $O(N \log N)$ with 3% utilization at realistic sequence lengths (Table 2).

This hardware-aware design philosophy echoes the motivation behind FlashAttention (which the same lab contributed to), but extends it from memory-IO optimization to the choice of mathematical primitive itself. The core insight is that the computational primitive should be chosen for its hardware compatibility, not just its asymptotic complexity, and the architecture should be designed around that primitive.

3. Technical Approach

3.1 Reader Orientation

This paper is primarily an architecture design and theoretical analysis paper whose core idea is that Monarch matrices—a structured class of expressive, hardware-efficient matrices that generalize the FFT—can serve as a single, unified sub-quadratic primitive for mixing information along both the sequence length and model dimension axes in neural networks, replacing both attention and MLPs in Transformers with the same computational building block while achieving $O(N^{3/2})$ complexity (for order-2) with high GPU utilization via GEMM operations. The system being built is a family of architectures—M2-BERT, M2-ViT, and M2-GPT—that demonstrate this unified mixing approach across non-causal, spatial, and causal settings, with the theoretical contribution of characterizing when Monarch convolutions can be made causal while remaining sub-quadratic through a polynomial evaluation and interpolation framework.

3.2 Big-Picture Architecture (Diagram in Words)

The M2 architecture has three major conceptual components that operate identically regardless of whether they are applied to sequences or model dimensions:

1. Monarch matrices (the mixing primitive): Structured matrices parameterized as products of block-diagonal matrices interleaved with permutations. An order-$p$ Monarch matrix $M \in \mathbb{R}^{N \times N}$ is defined as $M = \left(\prod_{i=1}^p P_i B_i\right) P_0$, where each $B_i$ is a block-diagonal matrix with block size $b = \sqrt[p]{N}$ and each $P_i$ is a variant of the bit-reversal permutation. These matrices compute $O(p N^{(p+1)/p})$ FLOPs instead of $O(N^2)$ and are implemented using GEMM operations on the blocks, achieving high hardware utilization.

2. Sequence mixer (replaces attention): Uses Monarch matrices to implement a gated long convolution. For input $X \in \mathbb{R}^{N \times d}$, the operation is $\tilde{X} = M_2(K_1 \odot M_1 X)$, where $M_1, M_2 \in \mathbb{R}^{N \times N}$ are Monarch matrices (typically fixed to DFT and inverse DFT to recover convolution), $K_1 \in \mathbb{R}^{N \times d}$ is a learned kernel in frequency space, and $\odot$ is elementwise multiplication. This is equivalent to a convolution with kernel $K_1$ when $M_1$ and $M_2$ are DFT matrices. More expressive variants add gating: $V \odot M_2(K_1 \odot M_1(Q \odot K))$, where $Q, K, V$ are linear projections of $X$, reproducing the gated convolution blocks from H3 and Hyena.

3. Dimension mixer (replaces MLP): Uses Monarch matrices to replace the two dense matrices in a standard MLP. For input $\tilde{X} \in \mathbb{R}^{N \times d}$, the operation is $Y^\top = M_4 \sigma(M_3 \tilde{X}^\top)$, where $M_3, M_4 \in \mathbb{R}^{d \times d}$ are learned Monarch matrices (typically order-1, i.e., block-diagonal matrices with block size 4, reducing to a sparse MLP) and $\sigma$ is an optional nonlinearity like ReLU or GeLU.

Information flows as follows: a sequence of embeddings $X \in \mathbb{R}^{N \times d}$ enters a layer → the sequence mixer convolves along the $N$ dimension using Monarch matrices $M_1, M_2$ and a learned kernel $K_1$ → the output $\tilde{X}$ (same shape $N \times d$) is passed through the dimension mixer, which applies learned Monarch matrices $M_3, M_4$ along the $d$ dimension → the final output $Y$ is produced with a residual connection and layer normalization. This layer is stacked repeatedly with standard Transformer-style residual connections and normalization.

3.3 Roadmap for the Deep Dive

• First, the mathematical definition of Monarch matrices and why their parameterization yields sub-quadratic complexity with hardware efficiency — this is the foundational primitive that everything else builds on.
• Second, the two mixing equations (sequence mixer and dimension mixer) that define an M2 layer, establishing the architectural template and showing how Monarch matrices are deployed as drop-in replacements for attention and MLPs.
• Third, the hardware benchmarking of the M2 convolution operator, which establishes the practical motivation for choosing $O(N^{3/2})$ Monarch convolutions over $O(N \log N)$ FFT convolutions — this explains why asymptotic complexity alone is insufficient.
• Fourth, the theoretical analysis that characterizes Monarch matrix multiplication as multivariate polynomial evaluation and interpolation, which provides the mathematical machinery needed to reason about causality.
• Fifth, the causal parameterization theorem that derives sufficient conditions on the basis polynomials for an M2 convolution to be causal while remaining sub-quadratic — this is the critical theoretical result that enables GPT-style auto-regressive models.
• Sixth, the three architecture instantiations (M2-BERT, M2-ViT, M2-GPT), detailing the exact design choices, hyperparameters, and architectural variations that map the general M2 template to specific training paradigms and domains.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architecture and theory paper whose core idea is that a single structured matrix class—Monarch matrices—can efficiently mix information along both the sequence and model dimensions with sub-quadratic complexity and high hardware utilization, replacing both attention and MLPs, and that greedy causality constraints can be enforced through careful control of polynomial degrees in the Monarch parameterization.

---

Monarch Matrices: The Fundamental Primitive

A Monarch matrix $M \in \mathbb{R}^{N \times N}$ of order $p$ is defined by the following factorization:

$M = \left(\prod_{i=1}^p P_i B_i\right) P_0$

where each $P_i$ is a permutation matrix related to the "base $\sqrt[p]{N}$" variant of the bit-reversal permutation, and each $B_i$ is a block-diagonal matrix with block size $b = \sqrt[p]{N}$.

What this factorization means structurally: The matrix $M$ is built as a sequence of $p$ block-diagonal matrices, each interleaved with a permutation that rearranges the input vector. The product is read right-to-left: the input vector is first permuted by $P_0$, then multiplied by block-diagonal $B_1$, then permuted by $P_1$, then multiplied by block-diagonal $B_2$, and so on through $p$ stages. Each $B_i$ has $N/b = N^{1-1/p}$ blocks, each of size $b \times b = N^{1/p} \times N^{1/p}$, arranged along the diagonal.

Why this yields sub-quadratic complexity: A dense matrix-vector multiplication costs $O(N^2)$ FLOPs. The Monarch factorization computes $p$ block-diagonal matrix multiplications. Each block-diagonal multiplication involves $N/b$ independent blocks of size $b \times b$, each costing $O(b^2)$ FLOPs, for a total per stage of $(N/b) \cdot O(b^2) = O(Nb) = O(N \cdot N^{1/p}) = O(N^{(p+1)/p})$. With $p$ stages, the total is $O(p N^{(p+1)/p})$. For $p=2$, this is $O(2 N^{3/2})$; for $p = \log N$ (the theoretical limit), this approaches $O(N \log N)$, recovering the FFT's asymptotic complexity. The key practical insight is that intermediate values of $p$ — particularly $p=2$ — provide a tunable tradeoff: $O(N^{3/2})$ complexity, worse than FFT asymptotically but with dramatically better constant factors due to hardware-friendly block structure.

Why this is hardware-efficient: Each $B_i$ is block-diagonal, meaning the matrix-vector multiplication decomposes into $N/b$ independent smaller matrix multiplications, each of size $b \times b$. These sub-multiplications are exactly the kind of dense GEMM operations that modern GPUs (specifically tensor cores) are optimized for. The permutation operations $P_i$ are simple data rearrangement (reshape, transpose, flatten) that can be implemented as memory operations without arithmetic. This means the entire Monarch matrix-vector product can be implemented using only matrix multiplication, transpose, reshape, and elementwise operations — all of which have optimized implementations in standard libraries like cuBLAS. The paper reports that a simple CUDA implementation of the M2 operator achieves 41.4% FLOP utilization at input size 64K on an RTX 4090, compared to 3% for FFT (Table 1).

The paper focuses primarily on order-2 Monarch matrices (the $p=2$ case), written as $M = PLPRP$, where $L$ and $R$ are block-diagonal matrices (mnemonically, "left" and "right" factors), and $P = P_2 = P_1 = P_0$ is a permutation that reshapes the input from 1D to 2D, transposes it, and flattens it back to 1D. In this case, the block size is $b = \sqrt{N}$. A common instantiation sets $L = R = (I_{\sqrt{N}} \otimes F_{\sqrt{N}})$, where $F_{\sqrt{N}}$ is a $\sqrt{N} \times \sqrt{N}$ DFT matrix and $\otimes$ is the Kronecker product — this choice recovers the standard FFT decomposition within the Monarch framework.

The parameter count: A dense $N \times N$ matrix has $N^2$ parameters. An order-2 Monarch matrix has two block-diagonal matrices, each with $\sqrt{N}$ blocks of size $\sqrt{N} \times \sqrt{N}$, totaling $2 \cdot \sqrt{N} \cdot (\sqrt{N})^2 = 2 N^{3/2}$ parameters. At $N = 4096$, this is a reduction from $\sim 16.8$ million to $\sim 524,000$ parameters (a 32× reduction). This parametric efficiency is important for the dimension mixer, where learned Monarch matrices replace dense MLP matrices and reduce the total parameter count of the model.

Expressivity: The original Monarch paper (Dao et al., 2022) showed that Monarch matrices can exactly represent (with appropriate parameter settings) a wide class of structured linear transforms: the FFT and inverse FFT (by setting blocks to DFT matrices), convolutions (via the FFT convolution theorem), Toeplitz matrices, low-displacement rank transforms, and orthogonal polynomial transforms. This expressivity means that fixing the Monarch matrices to specific transforms (e.g., DFT) recovers known operations (e.g., convolution), while learning the blocks from data allows the architecture to discover task-appropriate mixing patterns.

---

The M2 Layer: Architecting Mixing Along Two Axes

The M2 layer defines how Monarch matrices are deployed to mix information along both the sequence dimension ($N$) and the model dimension ($d$). Given an input $X \in \mathbb{R}^{N \times d}$, the layer computes:

Step 1: Sequence mixing (Equation 2 in the paper).

$\tilde{X} = M_2(K_1 \odot M_1 X)$

where $M_1, M_2 \in \mathbb{R}^{N \times N}$ are order-2 Monarch matrices, $K_1 \in \mathbb{R}^{N \times d}$ is a learned kernel matrix, and $\odot$ denotes elementwise (Hadamard) multiplication.

What this computes operationally: The input $X$ is multiplied on the left by $M_1$ — this is a Monarch matrix-vector product applied independently to each of the $d$ columns (features) of $X$, mixing information across the $N$ sequence positions. The result $M_1 X \in \mathbb{R}^{N \times d}$ is then elementwise-multiplied by the learned kernel $K_1$, which applies a position-dependent scaling to each feature. Finally, the result is multiplied on the left by $M_2$, mixing the scaled values again across sequence positions. The output $\tilde{X}$ has the same shape as $X$ ($N \times d$) but with information redistributed across positions.

Why this equals convolution when $M_1, M_2$ are DFT matrices: When $M_1$ is set to the DFT matrix $F_N$ and $M_2$ is set to the inverse DFT $F_N^{-1}$, we have:

$\tilde{X} = F_N^{-1} (K_1 \odot F_N X)$

This is exactly the FFT convolution theorem: $F_N X$ transforms each feature column to the frequency domain, $K_1$ provides frequency-domain filter coefficients (one per frequency per feature), the elementwise product applies the filter, and $F_N^{-1}$ transforms back to the time/sequence domain. The result is a circular convolution of each feature column with the filter whose frequency response is $K_1$. This is the fundamental operation underlying long convolutions: rather than explicitly constructing and applying an $N \times N$ Toeplitz convolution matrix (which would be $O(N^2)$), we represent the filter in the frequency domain and compute the convolution via FFT in $O(N \log N)$.

Beyond convolution — the gated variant: The paper notes that more expressive sequence mixers can be constructed by adding gating, following the pattern from H3 (Fu et al., 2023) and Hyena (Poli et al., 2023). The full gated sequence mixer is:

$V \odot M_2(K_1 \odot M_1(Q \odot K))$

where $Q = X W_Q$, $K = X W_K$, $V = X W_V$ are linear projections of the input (with weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$), and the innermost operation $M_1(Q \odot K)$ mixes the elementwise-gated query-key interaction product across sequence positions. $V \odot (\cdot)$ then gates the mixed result with the value projection. This structure mirrors the attention mechanism (query, key, value) but replaces the softmax attention matrix with a convolution kernel $K_1$ and Monarch transforms.

Step 2: Dimension mixing (Equation 3 in the paper).

$Y^\top = M_4 \sigma(M_3 \tilde{X}^\top)$

where $M_3, M_4 \in \mathbb{R}^{d \times d}$ are learned Monarch matrices (operating on the model dimension $d$), $\sigma$ is an optional pointwise nonlinearity such as ReLU or GeLU, and $^\top$ denotes matrix transpose.

What this computes operationally: The sequence-mixed output $\tilde{X} \in \mathbb{R}^{N \times d}$ is transposed to $\tilde{X}^\top \in \mathbb{R}^{d \times N}$, so that the $d$ dimension becomes the rows. The Monarch matrix $M_3$ is applied to each of the $N$ columns (sequence positions) independently, mixing information across the $d$ feature dimensions — this is a per-position feature transformation analogous to the first dense layer in an MLP. A nonlinearity $\sigma$ is applied elementwise. Then $M_4$ performs a second feature mixing. The result is transposed back to $Y \in \mathbb{R}^{N \times d}$.

Why this replaces the MLP: A standard Transformer MLP computes $Y = \sigma(\tilde{X} W_1^\top + b_1) W_2^\top + b_2$, where $W_1 \in \mathbb{R}^{d_{\text{ff}} \times d}$ and $W_2 \in \mathbb{R}^{d \times d_{\text{ff}}}$ are dense matrices (with $d_{\text{ff}} = 4d$ typically), costing $O(d \cdot d_{\text{ff}}) = O(d^2)$ FLOPs per token. The M2 dimension mixer replaces these two dense matrices with two Monarch matrices $M_3, M_4 \in \mathbb{R}^{d \times d}$, which cost $O(d^{3/2})$ FLOPs per token (for order-2). This eliminates the need for a separate expansion dimension $d_{\text{ff}}$ — the Monarch matrices operate directly at the model dimension size, mixing features within the same $d$-dimensional space. The nonlinearity is applied between the two Monarch multiplications, providing the necessary nonlinear transformation capacity.

The complete M2 layer thus alternates:

1. Mix along sequence ($\mathbb{R}^{N \times d} \to \mathbb{R}^{N \times d}$ via $M_1, M_2$ Monarch matrices and kernel $K_1$)
2. Mix along features ($\mathbb{R}^{N \times d} \to \mathbb{R}^{N \times d}$ via $M_3, M_4$ Monarch matrices and nonlinearity $\sigma$)

with a residual connection and layer normalization (as shown in the pseudocode in Figure 1).

Simplifications and special cases: In many of the paper's experiments, specific choices are made:

• For M2-BERT and M2-ViT, the sequence mixer Monarch matrices $M_1$ and $M_2$ are fixed to DFT and inverse DFT, making the sequence mixer a standard long convolution. Only the kernel $K_1$ is learned.
• For the dimension mixer, the Monarch matrices are set to order-1 (simply block-diagonal matrices with no permutation between blocks) with block size $b = 4$, effectively making the dimension mixer a sparse MLP where each block of 4 features interacts only among themselves. This reduces the parameter count of the MLP replacement substantially (each dense matrix is replaced by a block-diagonal matrix with blocks of size 4, reducing parameters from $d^2$ to $d^2 / 4$).
• For M2-GPT, the MLP layers are removed entirely (replaced with the identity) and the model is made wider to compensate, so the entire model consists only of sequence mixing layers — making it the "attention-free, MLP-free" variant.

The implementation is simple: Listing 1 in the appendix provides a complete PyTorch implementation in under 40 lines. The Monarch matrix multiplication is decomposed into two blockdiag_matmul calls (each multiplying a block-diagonal matrix by the input) interleaved with rearrange (reshape + transpose) operations that implement the permutation $P$. The full layer combines two Monarch matrices for sequence mixing, two for dimension mixing, a learned kernel, an activation function, and a residual connection with layer norm — all in $\sim 15$ lines of forward pass code.

---

Hardware Efficiency: Why $O(N^{3/2})$ GEMMs Beat $O(N \log N)$ FFTs

The paper includes hardware benchmarking of the core M2 convolution operator $M_2(K_1 \odot M_1 X)$ (Equation 2) to motivate the choice of Monarch matrices over FFT-based convolutions.

GPU cost model context (Section 2): Operations on GPUs are classified as either compute-bound (limited by arithmetic throughput) or memory-bound (limited by data movement between memory and compute units). Modern GPUs have a massive asymmetry: on the A100, tensor cores achieve 312 TFLOP/s for half-precision matrix multiply, while non-matrix operations are limited to 19 TFLOP/s. Memory bandwidth is 2 TB/s for the 80 GB HBM2, but only 19 TB/s for the much smaller L2 cache (20 MB). This means that operations that can be structured as dense matrix multiplications (GEMMs) and that keep data in cache will dramatically outperform operations that involve irregular memory access patterns or non-matrix arithmetic.

Why FFT is memory-bound: The FFT algorithm (specifically Cooley-Tukey) recursively decomposes the transform into smaller DFTs, interleaved with "butterfly" operations and bit-reversal permutations. At each recursion level, the data must be read, combined, and written back, creating many passes over the data with irregular access patterns. The bit-reversal permutation at the end requires non-sequential memory access. The arithmetic operations are simple (complex multiply-add) and don't saturate tensor cores. The result is that FFT implementations spend most of their time waiting for memory rather than computing, achieving only 3% FLOP utilization (Table 1).

Why Monarch matrices achieve higher utilization: The block-diagonal structure of Monarch matrices means that the arithmetic is concentrated in dense matrix multiplications of $b \times b$ blocks (with $b = \sqrt{N}$ for order-2). These are exactly the operations that tensor cores are designed for. The permutations between blocks are pure data rearrangement (reshaping a 1D vector to 2D, transposing, flattening), which can be implemented as pointer/index manipulation without explicit data movement in optimized implementations. The paper's simple CUDA implementation (calling standard cuBLAS subroutines for the block GEMMs) achieves 25.6% FLOP utilization on A100 and 41.4% on RTX 4090 at input size 64K (Table 2).

The critical Table 2 (hardware benchmarks): The paper benchmarks the M2 operator $M(K_1 \odot M X)$ at varying input sizes $N$, with block size fixed to $\sqrt{N}$. For each $N$, the table reports:

What this table shows: At small input sizes ($N = 4\text{K}$), the M2 operator is only marginally faster than dense matmul because the overhead of the permutation operations dominates the relatively small FLOP savings — M2 achieves only 4.78% utilization on A100 at this size. As $N$ grows, the sub-quadratic FLOP reduction kicks in: at $N = 16\text{K}$, M2 is already 5.1× faster than dense; at $N = 64\text{K}$, 20.6× faster; at $N = 256\text{K}$, over 55× faster (the dense matrix doesn't fit in GPU memory at this size, making the speedup effectively infinite for practical purposes). The RTX 4090 achieves better utilization than the A100 because it has a larger and faster L2 cache, which better amortizes the data movement costs of the permutation operations. The appendix includes a roofline analysis (Figure 4) confirming that the M2 operator's bottleneck is data movement rather than compute, and that the RTX 4090's improved cache architecture directly improves performance.

The key comparison to FFT: While the paper doesn't directly benchmark FFT in Table 2, Table 1 provides the context: at 64K input dimension, FFT achieves 3.0% utilization on an RTX 4090, while M2 achieves 41.4%. M2 performs more total FLOPs ($O(N^{3/2})$ vs. $O(N \log N)$) but those FLOPs execute at over 13× higher throughput. The net result is that M2 convolution is faster than FFT convolution at the sequence lengths relevant to modern language models (4K–8K tokens).

The A100 vs. RTX 4090 comparison reveals the memory bottleneck: On the A100, which has 80 GB of HBM2 but a relatively small L2 cache, the M2 operator achieves only 25.6% utilization at $N = 64\text{K}$ — the permutations cause data movement that the L2 cache cannot fully absorb. On the RTX 4090, which has a larger and faster L2 cache, utilization jumps to 41.4% at the same size. This suggests that M2 performance will continue to improve on future GPU architectures with larger caches, and that the FLOP utilization ceiling is primarily a function of memory hierarchy rather than arithmetic capability.

---

Theoretical Analysis: Monarch Matrices as Polynomial Evaluation and Interpolation

The theoretical contribution of the paper is a characterization of Monarch matrix multiplication as multivariate polynomial evaluation and interpolation. This framework provides the mathematical language needed to analyze when M2 convolutions are causal, and to derive sufficient conditions for causality that can be satisfied by constraining the degrees of the basis polynomials.

Connection between Monarch matrices and polynomials: The starting point is the observation (generalizing results from Dao et al., 2022) that an order-2 Monarch matrix $M = PLPRP \in \mathbb{R}^{N \times N}$ has columns that are evaluations of specific bivariate polynomials over a 2D grid of evaluation points.

Fix evaluation points $A = \{\omega^0, \omega^1, \ldots, \omega^{\sqrt{N}-1}\}$, where $\omega = e^{2\pi i / \sqrt{N}}$ are the $\sqrt{N}$-th roots of unity. Let $\ell_{j_1}(X, Y)$ for $0 \leq j_1 < \sqrt{N}$ be arbitrary bivariate polynomials with $\deg_X(\ell_{j_1}), \deg_Y(\ell_{j_1}) < \sqrt{N}$, and let $r_{j_1, j_0}(Y)$ for $0 \leq j_1, j_0 < \sqrt{N}$ be arbitrary univariate polynomials with $\deg(r_{j_1, j_0}) < \sqrt{N}$. These polynomials are not fixed a priori — they are defined by the entries of the block-diagonal matrices $L$ and $R$ that parameterize the Monarch matrix.

The entries of $L$ and $R$ are set by evaluating these polynomials:

• $L_{i_1, j_1}[i_0, i_0] = \ell_{j_1}(\omega^{i_1}, \omega^{i_0})$: the $(i_1, i_0)$-th "row" of block $(j_1, j_1)$ of $L$ is the evaluation of polynomial $\ell_{j_1}$ at the grid point $(\omega^{i_1}, \omega^{i_0})$.
• $R_{j_1, j_1}[i_0, j_0] = r_{j_1, j_0}(\omega^{i_0})$: the $(i_0, j_0)$-th entry of block $(j_1, j_1)$ of $R$ is the evaluation of polynomial $r_{j_1, j_0}$ at $\omega^{i_0}$.

Theorem 6 (in the appendix) proves that this mapping from polynomial coefficients to Monarch matrices is a bijection — every Monarch matrix corresponds to a unique set of polynomials and vice versa.

Theorem 1 (Monarch multiplication as polynomial evaluation): For any vector $u \in \mathbb{R}^N$ (indexed by pairs $(j_1, j_0)$ with $0 \leq j_1, j_0 < \sqrt{N}$), the product $M u$ is equivalent to evaluating a bivariate polynomial $u(X, Y)$ at all points in $A^2 = A \times A$, where:

$u(X, Y) = \sum_{j=0}^{N-1} u_j f_j(X, Y)$

and $f_j(X, Y) = \ell_{m(j)}(X, Y) \cdot r_{m(j)}(Y)$, where $m(j) = j \bmod \sqrt{N}$ extracts the "column" index of $j$ in the 2D grid representation.

What this means operationally: Each column of $M$ (corresponding to a specific $(j_1, j_0)$ pair) is the evaluation of a specific basis polynomial $\ell_{j_1}(X, Y) \cdot r_{j_1, j_0}(Y)$ over the $\sqrt{N} \times \sqrt{N}$ grid $A^2$. When we multiply $M$ by a vector $u$, we are forming the linear combination $u(X, Y) = \sum_j u_j f_j(X, Y)$ in the basis defined by these polynomials, and then evaluating that combined polynomial at each grid point. The $i$-th entry of $M u$ is $u(\omega^{i_1}, \omega^{i_0})$.

Why this framing is useful: It converts linear algebra (matrix-vector products) into polynomial algebra (evaluation and interpolation). Properties of the linear operation — most importantly, whether it is causal — can be analyzed in terms of the degrees and structure of the basis polynomials $f_j(X, Y)$.

Theorem 2 (M2 convolution as polynomial multiplication): Consider the M2 convolution operation:

$f = M_0^{-1}((M_1 k) \odot (M_2 u))$

where $M_0, M_1, M_2$ are Monarch matrices, $k$ is a kernel vector, and $u$ is the input. When the evaluation points $A$ are the $\sqrt{N}$-th roots of unity, this operation is equivalent to:

1. Representing the kernel $k$ as a bivariate polynomial $k(X, Y)$ in the basis corresponding to $M_1$.
2. Representing the input $u$ as a bivariate polynomial $u(X, Y)$ in the basis corresponding to $M_2$.
3. Multiplying the two polynomials and taking the result modulo $(X^{\sqrt{N}}-1, Y^{\sqrt{N}}-1)$: $h(X, Y) = k(X, Y) \cdot u(X, Y) \bmod (X^{\sqrt{N}}-1, Y^{\sqrt{N}}-1)$
4. Interpolating $h(X, Y)$ back to coefficient form in the basis corresponding to $M_0$.

What this means operationally: The elementwise product in the frequency (evaluation) domain corresponds to polynomial multiplication in the coefficient domain. The modular reduction $\bmod (X^{\sqrt{N}}-1, Y^{\sqrt{N}}-1)$ arises because evaluating at the $\sqrt{N}$-th roots of unity means that $X^{\sqrt{N}} = 1$ and $Y^{\sqrt{N}} = 1$ at all evaluation points — any term $X^a Y^b$ where $a \geq \sqrt{N}$ or $b \geq \sqrt{N}$ "wraps around" to $X^{a \bmod \sqrt{N}} Y^{b \bmod \sqrt{N}}$, which is exactly the effect of polynomial multiplication modulo $(X^{\sqrt{N}}-1, Y^{\sqrt{N}}-1)$.

Why this matters for causality: The modular reduction is the source of non-causality in the general M2 convolution. Under modular reduction, a high-degree term $X^{a} Y^{b}$ with $a \geq \sqrt{N}$ would contribute to a lower-degree coefficient $X^{a \bmod \sqrt{N}}$, potentially making the output at a low index depend on input at a high index — which violates causality (the output at position $i$ should depend only on inputs at positions $\leq i$). To make the convolution causal, we need to prevent this wrap-around.

---

Causal Parameterization: Controlling Polynomial Degrees

The paper's main theoretical result is Theorem 3, which gives sufficient conditions on the basis polynomials of the Monarch matrices such that the M2 convolution is causal — i.e., the output at each position depends only on inputs at earlier positions — while remaining sub-quadratic.

Step 1: Converting bivariate to univariate polynomials. The first step is to convert the bivariate basis polynomials $f_{j_1, j_0}(X, Y) = \ell_{j_1}(X, Y) \cdot r_{j_1, j_0}(Y)$ into univariate polynomials via the Kronecker substitution:

$q_j(Z) = \ell_{m(j)}(Z, Z^{\sqrt{N}}) \cdot r_{m(j)}(Z^{\sqrt{N}})$

where $j = j_1 \sqrt{N} + j_0$ is the flattened 1D index, $m(j) = j \bmod \sqrt{N}$ (the column index in the 2D representation), and the substitution $X \leftarrow Z, Y \leftarrow Z^{\sqrt{N}}$ maps the 2D grid to a 1D sequence. This substitution "diagonalizes" the 2D grid: the point $(X, Y) = (\omega^{i_1}, \omega^{i_0})$ in the original bivariate evaluation corresponds to $Z = \omega^{i_1 \sqrt{N} + i_0}_N$, where $\omega_N = e^{2\pi i / N}$ is the $N$-th root of unity. This works because $\omega^{i_1}_N = (\omega^{i_1 \sqrt{N} + i_0}_N)$ and $\omega^{i_0}_{\sqrt{N}} = (\omega^{i_1 \sqrt{N} + i_0}_N)^{\sqrt{N}}$.

Why this conversion is necessary: Causality is fundamentally a 1D property — "output at position $i$ depends on input up to position $i$" — that is most naturally analyzed on a univariate polynomial basis over the $N$-th roots of unity. The Kronecker substitution converts the bivariate structure of Monarch matrices into a univariate form where degree ordering corresponds to temporal ordering.

Step 2: Defining minimum degree constraints. The paper introduces the concept of a polynomial with a minimum degree: a polynomial $\bar{q}_j(Z)$ has minimum degree $j$ if its lowest-degree non-zero term is at $Z^j$. This means $\bar{q}_j(Z) = \sum_{a=j}^{N-1} \bar{q}_j[a] Z^a$ — all coefficients for degrees below $j$ are zero.

The causal parameterization imposes two constraints on the univariate basis polynomials $q_j(Z)$:

1. For $j < N/2$: $q_j(Z)$ must have minimum degree $j$ and maximum degree $< N/2$. This means the polynomial's terms range from $Z^j$ to $Z^{N/2 - 1}$. The low-degree terms (below $j$) are zero, and the high-degree terms (above $N/2 - 1$) are zero.

2. For $N/2 \leq j < N$: $q_j(Z)$ must have minimum degree $j$ and maximum degree $\leq N-1$. These polynomials can extend up to degree $N-1$, but their minimum degree still matches their index $j$.

In the bivariate polynomial parameterization, these constraints translate to requiring that the component polynomials $\ell_{m(j)}(X, Y)$ and $r_{m(j)}(Y)$ have minimum degrees $m(j)$ and $\lfloor j / \sqrt{N} \rfloor$ respectively, and that the combined polynomial $q_j(Z)$ has maximum degree $< N/2$ for indices $j < N/2$.

Step 3: The causal map. Given Monarch matrices $M_N$ parameterized with basis polynomials satisfying the above degree constraints, and given an input vector $u \in \mathbb{R}^n$ with $n < N/2$, the operation:

$\text{output} = \left(M_N^{-1} \left(M_N (k, 0^{N-n}) \odot M_N (u, 0^{N-n})\right)\right)[0:n-1]$

is causal. Here $(u, 0^{N-n})$ denotes padding the input with $N-n$ zeros, and $[0:n-1]$ denotes taking the first $n$ elements of the result.

Why this is causal: The proof (Theorem 8 in the appendix) shows that under the degree constraints:

1. The polynomial product $k(Z) \cdot u(Z)$ has maximum degree $\leq 2n - 2 \leq N - 2$, which is strictly less than $N$. This means the product does NOT trigger modular reduction modulo $Z^N - 1$ — there is no wrap-around. The convolution is equivalent to ordinary polynomial multiplication (no modular reduction).

2. For any indices $j, j'$ with $0 \leq j, j' < n$, the product of basis polynomials $q_j(Z) \cdot q_{j'}(Z)$ can be expressed as a linear combination of basis polynomials $q_i(Z)$ for $i \geq j + j'$ (Lemma 2). This means that the input coefficient $u_j$ multiplied by $k_{j'}$ only contributes to output coefficients $f_i$ where $i \geq j + j' \geq j$. In particular, the output at index $i$ depends only on input coefficients $u_j$ for $j \leq i$, which is exactly the definition of causality.

What the "zero-padding" and "truncation" do: The input $u$ of length $n$ is zero-padded to length $N$ (with $N = \lceil \sqrt{2n} \rceil^2$) before the Monarch transforms. This ensures that the polynomial multiplication $k(Z) \cdot u(Z)$ stays within the non-wrapping degree regime ($< N/2$). After the inverse Monarch transform, only the first $n$ outputs are kept — the rest are discarded because they correspond to the zero-padded region.

The "blowup" factor: This parameterization requires padding the input from length $n$ to length $N = \lceil \sqrt{2n} \rceil^2 \approx 2n$. This $2\times$ expansion is the cost of ensuring causality — the computation operates on vectors roughly twice as long as the actual input/output. The resulting convolution is still sub-quadratic: $O(N^{3/2}) = O((2n)^{3/2}) = O(n^{3/2})$, which is sub-quadratic in the original input length $n$.

Practical implications for the GPT setting: In M2-GPT, this causal parameterization is used for the sequence mixer. The Monarch matrices $M_1$ and $M_2$ in Equation 2 are parameterized with basis polynomials satisfying the degree constraints from Theorem 3, and the input vectors are appropriately zero-padded. The convolution kernel $K_1$ is similarly padded. The result is a causal long convolution — each output token depends only on the current and previous input tokens — that can be computed in $O(N^{3/2})$ time using only GEMM operations (via the block-diagonal structure of Monarch matrices).

The block algorithm implementation (Appendix D.4): The paper provides pseudocode (Algorithm 3: BlockMonarchConv) for computing the causal Monarch convolution using only GEMM operations. The key steps are:

1. Construct the block-diagonal matrices $L$ and $R$ from coefficient matrices $\tilde{L}, \tilde{R}$ using the Fourier transform (for the complex case) — Algorithm 2 (BlockyMonarch).
2. Form $M = PLPRP$.
3. Compute $M k$ and $M u$ (the frequency-domain representations).
4. Elementwise multiply the results.
5. Compute $M^{-1} (\cdot)$ (the inverse transform).

For the causal variant, the coefficient matrices $\tilde{L}, \tilde{R}$ are constrained using the degree conditions from Lemma 6, ensuring that the resulting Monarch matrices produce the basis polynomials required by Theorem 3. The full algorithm (Algorithm 3) and its causal specialization (Theorem 11) show that the entire computation requires $O(N^{3/2})$ FLOPs and uses only GEMMs on blocks of size $\sqrt{N}$, matching the non-causal M2 convolution in complexity while preserving causality.

---

Architecture Instantiations: M2-BERT, M2-ViT, M2-GPT

The paper instantiates the general M2 layer in three specific architectures, each tailored to a different training paradigm and domain. The key design choices for each are:

M2-BERT (non-causal masked language modeling):

The architecture replaces BERT's attention and MLP with M2 layers as follows:

• Sequence mixer: A bidirectional gated convolution with a residual convolution. The Monarch matrices are set to DFT and inverse DFT (fixed, not learned), making this a standard long convolution in frequency space. The gating follows the H3/Hyena pattern: short depthwise convolutions are applied after the linear projections $Q, K, V$, then the Hyena-style gated long convolution computes the sequence mixing. Short depthwise convolutions (with small kernel sizes) are added after the query, key, and value projections to capture local patterns before the global long convolution mixes across the full sequence — this is a common design choice in long convolution architectures that improves quality.
• Dimension mixer: A sparse MLP where the two dense matrices are replaced with order-1 Monarch matrices (block-diagonal matrices) with block size 4. This means features are partitioned into groups of 4 that interact within each group, with no cross-group interaction. This is a simple form of structured sparsity that reduces parameter count while maintaining representational capacity.
• Expansion factor: 4, matching the standard Transformer MLP expansion ratio (the hidden dimension is 4× the model dimension in standard BERT; this is preserved via the block-diagonal parameterization).
• Model sizes:
  • M2-BERT-base (80M): width 768, 12 layers
  • M2-BERT-base (110M): width 960, 12 layers
  • M2-BERT-large (260M): width 1536, 12 layers
  • M2-BERT-large (341M): width 1792, 12 layers
• Training: Pretrained on C4 with the bert-base-uncased tokenizer, sequence length 128, global batch size 4096, for 70,000 steps using decoupled AdamW with learning rate $8 \times 10^{-4}$ and weight decay $1 \times 10^{-5}$, linear decay with 6% warmup, 30% MLM masking rate.

M2-ViT (image classification):

The architecture is based on HyenaViT-b (which replaces ViT attention with gated long convolutions), with further Monarch substitutions:

• Sequence mixer: Replaces the HyenaViT long convolution with the M2 operator (DFT and inverse DFT Monarch matrices), maintaining the same gating structure. The image is split into patches (16×16 pixels), flattened to a sequence, and processed through M2 layers that mix across patch positions.
• Dimension mixer: Replaces the MLP blocks with order-1 block-diagonal matrices (block size 4), same as M2-BERT.
• Model sizes: ViT-b configuration (12 layers, hidden size 768, MLP intermediate size 3072, 12 attention heads) yields M2-ViT-b at 45M parameters (compared to 87M for ViT-b and 88M for HyenaViT-b).
• Training: Sinusoidal position embeddings, global average pooling (no class token), RandAugment, Mixup ($\alpha=0.8$), CutMix ($\alpha=1.0$), Random erasing, AugMix, AdamW with cosine decay, base learning rate $1 \times 10^{-3}$, batch size 1024, 300 epochs, image size 224×224.

M2-GPT (causal auto-regressive language modeling):

This is the most radical architecture variant — it uses neither attention nor MLPs:

• Sequence mixer: Combines the Hyena convolutional filter (the state-of-the-art attention-free LM at the time) with the head structure from H3 (tying long convolution parameters across heads and computing outer products between different input projections). The FFT operations in Hyena are replaced with causal Monarch matrices parameterized according to Theorem 3 — this is the key use of the theoretical results. The Hyena architecture uses two convolutions: a short convolution (implemented as a standard local convolution) and a long convolution (originally FFT-based, now Monarch-based with causal constraints).
• Dimension mixer: Removed entirely — replaced with the identity function. The model is made wider to compensate for the lost capacity.
• Head structure: The H3 head extension groups the model dimension into heads (each of dimension $d_h = 16$), ties long convolution parameters within each head, and computes the outer product between $X_2$ and $V$ projections before applying the long convolution — this creates a richer interaction pattern without requiring a separate MLP.
• Model sizes:
  • Small: width 1160, 18 layers, 145M parameters, learning rate $6 \times 10^{-4}$
  • Medium: width 1344, 40 layers, 360M parameters, learning rate $8 \times 10^{-4}$
• Training: The PILE dataset, matching Hyena's training protocol with varying amounts of tokens (5B, 10B, 15B), learning rate decayed appropriately for each budget.

Why the MLP is removed in M2-GPT: The paper's dimension mixer (Equation 3) uses Monarch matrices to replace the MLP. In M2-GPT, this is taken to the extreme — the dimension mixer is eliminated entirely, and the sequence mixer is made expressive enough (via the H3 head structure with outer products) to handle both sequence mixing and feature transformation. This is a proof-of-concept that an architecture can function without any MLP-like component, relying solely on sequence mixing operations. The wider model compensates for the lost parameters and nonlinearity: removing the MLP layers (which typically have expansion factor 4) roughly halves the parameter count at a given width, so the model is made wider to maintain the desired total parameter count.

4. Key Insights and Innovations

Innovation 1: A Unified Sub-Quadratic Primitive for Two Historically Separate Scaling Problems

The paper's most conceptually distinctive move is not inventing a new structured matrix, but recognizing that the same matrix class can solve two problems the field had treated as independent: the quadratic cost of attention in sequence length and the quadratic cost of MLPs in model dimension.

Prior to this work, the research communities addressing these two bottlenecks operated largely in isolation. The efficient attention community developed linear and sub-quadratic sequence mixers—long convolutions (Poli et al., 2023; Romero et al., 2022), state space models (Gu et al., 2021), linearized attention (Katharopoulos et al., 2020)—but left the MLP untouched, accepting its quadratic scaling in model dimension as a fixed cost. The sparse MLP community developed butterfly matrices (Chen et al., 2021), Monarch matrices (Dao et al., 2022), and pruning methods (Frankle & Carbin, 2018) to reduce the cost of feature mixing, but did not address sequence mixing. These were two separate research threads with different mathematical toolkits, different benchmarks, and different communities.

The "mixer" perspective—that architectures like MLP-Mixer and ConvMixer operate by repeatedly mixing along sequence and model dimension axes using a single operator type—was present in prior work (Tolstikhin et al., 2021; Ng et al., 2022), but those architectures used quadratic operators (dense MLPs or standard convolutions) for both axes. M2 is the first architecture to use a genuinely sub-quadratic operator along both axes simultaneously, achieving an integrated solution where the same primitive handles both mixing directions with the same computational characteristics (sub-quadratic FLOP count, GEMM-based high utilization).

This unification is more than an engineering convenience—it changes the conceptual framing of architecture design. Rather than thinking of "the attention module" and "the MLP module" as distinct components requiring different optimization strategies, the M2 perspective treats them as two instances of the same abstract operation (structured matrix mixing) applied to different tensor dimensions. The question shifts from "how do we optimize attention?" or "how do we sparsify MLPs?" to "what is the right structured matrix class for mixing, and how do its properties trade off across both axes?"

The evidence that this unification is genuinely beneficial—not just aesthetically pleasing—comes from the parameter efficiency results. M2-BERT matches BERT-base downstream quality with 27% fewer parameters (Table 3), even though the sequence mixer is not learning its Monarch matrices (they're fixed to DFT). The parameter savings come primarily from the dimension mixer, where learned block-diagonal matrices replace dense MLP weights. This demonstrates that the structural efficiency of Monarch matrices translates directly to reduced model size without quality degradation—a practical benefit of the unified approach.

Innovation 2: Viewing Monarch Convolutions Through the Lens of Multivariate Polynomials to Characterize Causality

The paper's theoretical contribution is not just proving that causal M2 convolutions exist, but providing a diagnostic framework that reveals why general Monarch convolutions are non-causal and what structural constraints restore causality while preserving sub-quadratic complexity.

Prior work on structured matrix primitives (Dao et al., 2022; Chen et al., 2021) treated them as static building blocks—pre-defined linear operators with certain mathematical properties and computational costs. The question of how to use these matrices in auto-regressive architectures (where each output must depend only on past inputs) was not addressed, because these matrices were primarily studied in bidirectional or fixed-size-input settings.

The polynomial evaluation framework (Theorems 1 and 2, Section 4) reframes the problem in a way that makes the non-causality source visible and tractable. The key insight is that the elementwise product in the frequency domain—which computes convolution—corresponds to polynomial multiplication in the coefficient domain, and that the evaluation at roots of unity introduces a modular reduction $(X^{\sqrt{N}} - 1, Y^{\sqrt{N}} - 1)$ that causes degree wrap-around. A high-degree term $X^a Y^b$ with $a \geq \sqrt{N}$ contributes to the coefficient of $X^{a \bmod \sqrt{N}}$ after modular reduction, meaning a late-time input can influence an early-time output. This is the precise mathematical mechanism behind non-causality.

This diagnostic framing matters because it converts a binary question ("is this causal?") into a tunable design space. The degree constraints in Theorem 3—requiring basis polynomials to have minimum degree $j$ and maximum degree $< N/2$ for indices below $N/2$—are sufficient conditions that prevent wrap-around by ensuring the polynomial product never exceeds degree $N-2$. This is not a black-box fix; it is a constructive characterization that tells the designer how to parameterize the Monarch factors (specifically, which entries of the coefficient matrices $\tilde{L}$ and $\tilde{R}$ can be non-zero) to guarantee causality.

The theoretical depth is reinforced by the generalization to arbitrary order $p$ in Appendix D.6, which shows that the same polynomial framework extends naturally: $p$-variate polynomials with degree constraints yield causal maps, with the required padding scaling as $2^p \cdot n$ (a limitation the authors explicitly flag for future work). This generality suggests the framework is fundamental to understanding structured matrix primitives in sequence modeling, not a one-off trick for order-2 Monarch matrices.

The practical significance is clear: without this characterization, M2 would be restricted to non-causal applications (masked language modeling, image classification). The ability to derive causal M2 convolutions is what enables M2-GPT—the attention-free, MLP-free auto-regressive language model that matches Transformer perplexity on The PILE (Table 8). This is the first demonstration that a model with neither attention nor MLPs can be competitive on causal language modeling, and the theoretical analysis is what makes it possible.

Innovation 3: Asymptotic Complexity Is Not Destiny—Hardware Utilization as a First-Class Design Constraint

The paper makes a forceful empirical and conceptual argument that FLOP count alone is an insufficient metric for architecture design on modern accelerators, and that the choice of mathematical primitive should be driven by hardware compatibility, not just asymptotic scaling.

This is not a new observation in computer systems generally—the gap between theoretical FLOPs and wall-clock time is well-known in HPC—but the paper operationalizes it as a primary architecture design criterion in a way that is distinctive in the ML architecture literature. Most architecture papers report parameter counts, FLOP counts, and sometimes wall-clock throughput as separate metrics. M2's argument is that the choice of Monarch matrices over FFT-based convolutions is justified by the interaction between complexity class and hardware utilization, not by either alone.

The critical evidence is Table 1 and Table 2, taken together:

• FFT convolutions: $O(N \log N)$ FLOPs, but 3% FLOP utilization (memory-bound, irregular access patterns)
• M2 convolutions: $O(N^{3/2})$ FLOPs, but 41.4% FLOP utilization (GEMM-based, tensor-core-friendly)
• Result: M2 is faster in wall-clock time at sequence lengths $\geq 4\text{K}$, despite performing ~13× more FLOPs than an asymptotically equivalent FFT at $N = 64\text{K}$

This inverts the standard narrative. The usual story in efficient ML is: "we found a way to reduce FLOPs/parameters by X× while maintaining quality." M2's story is: "we slightly increased FLOPs relative to the best-known asymptotic algorithm, but structured those FLOPs to run at dramatically higher throughput, yielding net wall-clock speedups." The sub-quadratic scaling matters (it enables 20.6× speedup over dense at 64K on A100, Table 2), but the hardware-aware choice of primitive is what makes the sub-quadratic scaling translate to real speed rather than paper speed.

The RTX 4090 vs. A100 comparison (Table 2) reinforces this point and adds a forward-looking dimension. The RTX 4090's larger L2 cache enables 41.4% utilization vs. 25.6% on the A100 at the same problem size, because the data movement costs of the Monarch permutations are better amortized. This suggests that M2's performance will improve with future GPU architectures that have larger caches, while FFT's utilization (limited by its fundamentally irregular access patterns) may not benefit as much. The architecture's efficiency is aligned with hardware trends in a way that FFT-based approaches are not.

This insight has broader implications for architecture design beyond M2. It suggests that researchers should:

1. Evaluate new primitives on FLOP utilization and roofline analysis, not just FLOP count.
2. Consider whether a slightly "worse" complexity class with better hardware mapping can outperform an "optimal" complexity class with poor mapping at realistic problem sizes.
3. Design primitives around tensor core capabilities (dense GEMMs, regular block structure) rather than around mathematical elegance alone.

The paper is careful not to overclaim—the M2 operator's utilization is still well below dense matmul (~80%), leaving room for optimization—but the principle of hardware-aware primitive selection is established as a core design philosophy.

Innovation 4: The First Empirical Demonstration That Attention and MLPs Are Both Dispensable in Auto-Regressive Language Modeling

M2-GPT's results (Table 8) represent a significant existence proof: it is possible to build a competitive causal language model that uses neither attention nor MLPs, relying entirely on Monarch-based convolutions for sequence mixing and removing the dimension mixer entirely. At 360M parameters trained on 15B tokens of The PILE, M2-GPT achieves 9.0 perplexity vs. 9.1 for a Transformer and 9.2 for Hyena.

This is distinctive because most prior work on attention-free language models (Hyena, H3, S4, RWKV) removed attention but retained MLPs. The MLP was considered essential for per-token nonlinear feature transformation—attention handled cross-token interaction, but the MLP provided depth and representational capacity within each token's representation. Hyena (Poli et al., 2023), the state-of-the-art attention-free LM at the time of this paper, used gated long convolutions for sequence mixing but kept standard MLP blocks for the dimension mixer.

M2-GPT challenges this assumption by showing that the H3 head structure—which computes outer products between different input projections before applying the long convolution—provides sufficient feature mixing capacity to replace the MLP entirely, provided the model is made wider to compensate for the lost parameters. This is a architectural minimality result: the model is stripped down to its essential sequence-mixing core, demonstrating that the MLP was not strictly necessary for competitive quality, at least at this scale.

The significance is partly philosophical (what is the minimal set of operations needed for language modeling?) and partly practical (removing MLPs eliminates the $O(d^2)$ dimension-mixing cost entirely, making the architecture genuinely sub-quadratic in both axes for training and inference). It also connects to a broader question in architecture design: can we collapse the distinction between "sequence mixing" and "feature mixing" into a single operation type applied to different dimensions, as the mixer perspective suggests? M2-GPT takes this to its logical extreme—the feature mixing is absorbed into the sequence mixing through the head structure's outer products, so there is literally only one type of operation in the model.

The negative result with ReST$^{EM}$ (Appendix K, Figure 16)—where an RL-style training procedure caused the revision-like M2 model to degrade substantially—highlights that this minimal architecture is sensitive to training methodology, and that the positive results depend on careful design choices (offline data construction, specific head structure, causal parameterization). This fragility is itself an important finding: the simpler the architecture, the more the training dynamics and data matter, because there are fewer redundant pathways to compensate for optimization failures.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates across three domains: (1) C4 (Raffel et al., 2019) for BERT-style masked language modeling, pretrained with the bert-base-uncased tokenizer at sequence length 128; downstream evaluation uses the GLUE benchmark (Wang et al., 2018), with fine-tuning following the procedure of Izsak et al. (2021) using standard per-task metrics (F1 for QQP/MRPC, Matthew's correlation for CoLA, Spearman's correlation for STS-B, accuracy for others). (2) ImageNet-1k (Deng et al., 2009) for image classification, with standard 224×224 resolution, following the training recipe of T2T-ViT (Yuan et al., 2021) including RandAugment, Mixup, CutMix, Random Erasing, and AugMix. (3) The PILE (Gao et al., 2020) for causal language modeling, the standard benchmark used by Hyena and H3, evaluated at three training budgets (5B, 10B, 15B tokens) with appropriate learning rate decay.

Base model(s). For BERT experiments: M2-BERT replaces the attention with bidirectional gated Monarch convolutions and replaces the MLP dense matrices with order-1 block-diagonal matrices (block size 4). Model scales: M2-BERT-base (80M and 110M parameter variants) and M2-BERT-large (260M and 341M parameter variants), compared against BERT-base (110M) and BERT-large (340M) from Devlin et al. (2019). For vision: M2-ViT-b (45M) is compared to ViT-b (87M; Dosovitskiy et al., 2020) and HyenaViT-b (88M; Poli et al., 2023). For causal LM: M2-GPT at 145M and 360M parameters, compared against Transformers (125M, 355M) and Hyena (155M, 360M).

Metrics. (1) GLUE score: the average across 8 GLUE tasks (MNLI, RTE, QNLI, QQP, SST-2, STS-B, CoLA, MRPC), using the standard per-task metrics as specified above; (2) ImageNet accuracy: Top-1% and Top-5% classification accuracy on the ImageNet-1k validation set; (3) Perplexity: standard language modeling perplexity on The PILE test set; (4) Throughput: tokens per millisecond for forward passes on A100-40GB GPUs, measured at varying sequence lengths (512 to 8192); (5) CPU inference latency: milliseconds per example with batch size 1 on a 48 vCPU, 96 GB RAM GCP n2-standard-48 instance (Intel Cascade Lake), averaged over 10 examples following the protocol of Funtowicz (2021).

Baselines. For GLUE: BERT-base and BERT-large (Devlin et al., 2019) as reported in the original paper. For throughput: HuggingFace BERT implemented in standard PyTorch, and BERT optimized with FlashAttention (Dao et al., 2022). For ImageNet: ViT-b (Dosovitskiy et al., 2020), HyenaViT-b (Poli et al., 2023), ViT-b + Monarch (which replaces only the MLP blocks with Monarch matrices, keeping attention), and ResNet-152 as an external reference. For causal LM: standard Transformers with attention and MLPs, and Hyena (Poli et al., 2023) which uses FFT-based gated long convolutions with standard MLPs.

Generation budget / compute accounting. For BERT experiments, all models are pretrained for a fixed number of steps (70,000) with the same batch size (4096 sequences) and sequence length (128), making FLOP comparisons implicit through parameter count and throughput measurements. For throughput benchmarks, the comparison is wall-clock tokens per millisecond at identical sequence lengths, with all models running on the same A100-40GB hardware. For causal LM, all models are trained for the same number of tokens (5B, 10B, or 15B) with model sizes approximately matched, enabling direct perplexity comparison at equivalent training compute budgets.

Cross-validation / statistical protocol. For GLUE fine-tuning, a small hyperparameter search is performed over learning rate, weight decay, and number of epochs per task, following the established protocol of Izsak et al. (2021). For ImageNet, hyperparameters (weight decay, stochastic depth rate, base learning rate) are tuned on a held-out validation set (10% of training data). For BERT pretraining, the masked language modeling validation accuracy is reported on a held-out split of C4, but no explicit cross-validation is needed since the pretraining aims to match classical BERT recipes. No complex cross-validation is performed for throughput or perplexity comparisons since those are single-configuration evaluations.

Main Quantitative Results

Non-Causal Language Modeling (M2-BERT)

Headline result: M2-BERT-base matches or exceeds BERT-base on GLUE downstream quality with substantially fewer parameters, while achieving dramatically higher throughput at long sequences.

GLUE quality (Tables 3, 4, 9): The 80M-parameter M2-BERT-base achieves an average GLUE score of 79.9, matching BERT-base's 79.6 with 27% fewer parameters (Table 3). When parameter-matched at 110M, M2-BERT-base achieves 80.9, surpassing BERT-base by 1.3 points. The larger models show similar trends: M2-BERT-large (260M) achieves 82.2 GLUE score, matching BERT-large's 82.1 with 24% fewer parameters, while the parameter-matched 341M variant reaches 82.8—0.7 points above BERT-large (Table 4).

Per-task breakdowns (Table 9, Appendix B.1) reveal that M2-BERT-base (80M) performs competitively across all tasks: MNLI 78.4/78.6, QQP 86.7, SST-2 92.0, STS-B 86.3, with CoLA (53.0) showing the largest gap to published BERT-base numbers. The 110M variant closes or reverses most gaps, achieving MNLI 79.6/80.5, RTE 69.3, QNLI 86.0, QQP 87.0, SST-2 92.3, STS-B 86.9, CoLA 56.0, and MRPC 89.2.

GPU throughput scaling (Table 5): At sequence length 512, M2-BERT-base (80M) achieves 386.3 tokens/ms—similar to FlashAttention BERT-base (367.4) and 1.9× faster than standard HuggingFace BERT-base (206.1). The advantage grows dramatically with sequence length: at sequence length 4096, M2-BERT-base reaches 353.9 tokens/ms vs. 179.1 for FlashAttention BERT (2.0× speedup) and 39.0 for standard BERT (9.1× speedup). At sequence length 8192, M2-BERT-base sustains 320.1 tokens/ms while FlashAttention BERT drops to 102.4 (3.1× speedup) and standard BERT exhausts GPU memory (OOM). This trend confirms the sub-quadratic scaling advantage: as sequence length doubles from 2048 to 4096, M2 throughput drops only 7% (378.9 → 353.9), compared to 30% for FlashAttention BERT (257.2 → 179.1).

Parameter-matched comparison (Table 10, Appendix B.2): Against a BERT model sized to 79M parameters (rather than the standard 110M BERT-base), M2-BERT-base (80M) shows 1.6× speedup over HuggingFace BERT at sequence length 512, growing to 7.5× at sequence length 4096. FlashAttention achieves higher throughput than M2 at short sequences (433.3 vs. 386.3 at 512) but is overtaken starting at sequence length 2048 (335.2 vs. 378.9), indicating that M2's FLOP advantage dominates FlashAttention's IO optimization at moderate-to-long sequences.

BERT-large throughput (Table 11, Appendix B.2): M2-BERT-large (260M) nearly matches FlashAttention BERT-large (340M) at sequence 512 (122.5 vs. 125.0 tokens/ms) and outperforms it at all longer lengths: 118.6 vs. 111.9 at 1K, 109.4 vs. 91.6 at 2K, and 94.5 vs. 54.5 at 4K—a 1.7× speedup at 4K. Compared to standard HuggingFace BERT-large, speedups reach 4.3× at sequence length 2048.

CPU inference latency (Table 6): At short sequences (512), M2-BERT-base is slower than BERT-base on CPU (289 ms vs. 182 ms, or 0.6× speedup), attributed to the data movement overhead of Monarch permutations dominating the FLOP savings at small sizes. The trend reverses at sequence length 1024 (361 vs. 389 ms, 1.1× speedup), and M2's advantage grows to 1.4× at 2048, 2.8× at 4096, and 6.5× at 8192 (1820 ms vs. 11820 ms). The paper notes that further IO-aware optimization could improve short-sequence CPU performance.

Image Classification (M2-ViT)

Headline result: M2-ViT-b achieves higher ImageNet accuracy than ViT-b and HyenaViT-b with roughly half the parameters, despite using neither attention nor MLPs.

ImageNet-1k accuracy (Table 7): M2-ViT-b (45M parameters) achieves 79.5% Top-1 accuracy and 94.5% Top-5 accuracy. This outperforms ViT-b (87M parameters; 78.5% Top-1, 93.6% Top-5) by 1.0 percentage point, HyenaViT-b (88M; 78.5%, 93.6%) by the same margin, and ViT-b + Monarch (33M; 78.9%, 94.2%) by 0.6 points. M2-ViT-b also outperforms ResNet-152 (60M; 78.6%, 94.3%), a dedicated ConvNet architecture extensively tuned for ImageNet. The 45M parameter count is roughly half that of ViT-b (87M), demonstrating significant architectural efficiency—the model is both smaller and more accurate.

Swin comparison (Table 12, Appendix B.3): When M2 is applied as a drop-in replacement for attention and MLP in the Swin-V2 architecture (Swin-M2-B), it achieves 83.5% Top-1 accuracy. This is competitive with Swin-V1-B (83.5%), outperforms Swin-MLP-B (81.3%), and comes within 0.7 points of the fully optimized Swin-V2-B (84.2%). Notably, this result was achieved without any hyperparameter tuning or architecture adjustment specific to the Swin variant, suggesting further gains are possible.

CIFAR-10 (Table 14, Appendix B.5): On CIFAR-10, M2-based architectures show consistent advantages: ViT + Monarch (607K parameters) achieves 79.0% vs. 78.6% for standard ViT (1.2M parameters), HyenaViT-M2 (741K) achieves 80.8% vs. 80.6% for HyenaViT (1.3M) and vs. 78.6% for the standard ViT. The pattern replicates ImageNet—fewer parameters, comparable or better accuracy.

Causal Language Modeling (M2-GPT)

Headline result: M2-GPT—using neither attention nor MLPs—matches or slightly outperforms Transformers and Hyena in pretraining perplexity on The PILE across two model scales and three training budgets.

Perplexity scaling (Table 8): At the ~150M parameter scale:

• 5B tokens: M2-GPT (145M) achieves 12.9 PPL vs. 13.3 for Transformer (125M) and 13.1 for Hyena (155M).
• 10B tokens: M2-GPT 11.6 vs. Transformer 11.9 and Hyena 11.8.
• 15B tokens: M2-GPT 10.9 vs. Transformer 11.2 and Hyena 11.1.

At the ~360M parameter scale:

• 5B tokens: M2-GPT (360M) achieves 11.0 PPL vs. 11.4 for Transformer (355M) and 11.3 for Hyena (360M).
• 10B tokens: M2-GPT 9.6 vs. Transformer 9.8 and Hyena 9.8.
• 15B tokens: M2-GPT 9.0 vs. Transformer 9.1 and Hyena 9.2.

M2-GPT holds a consistent 0.2–0.4 PPL advantage over the Transformer baseline and a 0.1–0.3 PPL advantage over Hyena across all data budgets and both model sizes. The gap does not widen with more training data (it remains ~0.2 PPL at both 5B and 15B tokens for the 360M models), suggesting that the architectural advantage is stable but modest.

Architecture implications: This is the first result showing an attention-free, MLP-free architecture matching Transformer quality on causal language modeling—a strictly stronger result than Hyena, which maintained MLPs. The M2-GPT architecture entirely removes the dimension mixer and compensates by making the model wider (the model depths match equivalent Hyena models, but widths are 1160 for the 145M variant and 1344 for the 360M variant).

Speech Classification (Appendix B.4)

Speech Commands-10 (Table 13): M2 achieves 97.9% accuracy on this raw audio classification task (1-second clips at 16 kHz), outperforming S4 (97.5%), WaveGan-D (96.3%), and CKConv (71.7%). The Transformer and Performer baselines did not fit in GPU memory (marked "x"), preventing direct comparison. This result demonstrates that the M2 architecture generalizes beyond text and images to raw audio processing, where long sequences (16,000 time steps for 1 second at 16 kHz) make sub-quadratic scaling particularly relevant.

Associative Recall (Appendix B.8)

In-context learning at extreme lengths (Table 16): On the associative recall synthetic task (vocab size 20), the causally-parameterized M2 model demonstrates in-context learning across sequence lengths from 512 to 128K tokens, achieving 98.7% at 0.5K, 99.4% at 2K–32K, and 99.4% at 128K. The Transformer baseline achieves 100.0% at 0.5K–8K but does not finish training within a week at 32K and 128K (marked ✗). This demonstrates that the causal M2 parameterization successfully enables sub-quadratic processing at sequence lengths far beyond what is practical for quadratic attention, while maintaining near-perfect in-context learning accuracy.

Ablation Studies and Robustness Checks

Learnable vs. fixed Monarch matrices in the sequence mixer (Table 15, Appendix B.6): On sequential CIFAR (a non-gated convolutional architecture based on long convolutions), using learnable Monarch matrices in the sequence mixer improves accuracy from 91.0% to 92.5% compared to fixed Monarch matrices (set to DFT). This 1.5-point lift demonstrates that learning the mixing patterns can be beneficial beyond pre-specified transforms, though the paper uses fixed DFT matrices for most main experiments for simplicity and to isolate the effect of the Monarch structure itself.

Older architecture variant (M1-BERT, Appendix B.9): The paper reports results for an earlier, non-gated version of the M2-BERT architecture, termed M1-BERT, trained on English Wikipedia and Bookcorpus (rather than C4). This model could match Transformers on masked language modeling accuracy but underperformed on downstream GLUE fine-tuning (Table 19: 69.9/70.5 MNLI vs. 74.5/74.7 for BERT-base). The paper attributes this gap to sub-optimal training hyperparameters (optimized for throughput using NVIDIA MLPerf settings) and the absence of gating, which motivated the switch to the gated architecture described in the main text. This negative result serves as an implicit ablation: gating is important for downstream transfer quality, not just for pretraining loss.

Throughput sensitivity to sequence length (Tables 5, 10, 11): The throughput benchmarks across three different model sizes (M2-BERT-base at 80M, BERT at 79M, M2-BERT-large at 260M) all show the same qualitative pattern: M2 is comparable to or slightly slower than FlashAttention at short sequences (512) but gains an increasing advantage as sequence length grows. The consistency across model scales strengthens the claim that the throughput advantage is a structural property of the sub-quadratic scaling, not an artifact of a particular model size or implementation.

CPU vs. GPU performance tradeoff (Table 6): The CPU inference results reveal that M2's performance advantage is not universal across hardware platforms—at sequence length 512, M2 is actually 1.6× slower than BERT on CPU. This cross-hardware comparison serves as an implicit ablation on the importance of the GPU memory hierarchy: the Monarch permutations' data movement costs are better absorbed by GPU caches than by CPU memory subsystems, making the architecture's efficiency somewhat platform-dependent.

Alternative training data for BERT (Appendix B.9): The M1-BERT results (trained on Wikipedia/Bookcorpus vs. the C4-trained M2-BERT) show that the architecture's effectiveness is somewhat sensitive to training data and hyperparameters—the gated architecture trained on C4 substantially outperforms the non-gated architecture on Wikipedia/Bookcorpus, though the confound of different training recipes prevents isolating the effect of data alone.

Critical Assessment

Claim 1: M2 matches BERT-base and BERT-large on GLUE with 27% and 24% fewer parameters. Supported, but with important caveats about what "matches" means. M2-BERT-base (80M) achieves 79.9 average GLUE vs. BERT-base's 79.6—a 0.3-point advantage with 27% fewer parameters (Table 3). M2-BERT-large (260M) achieves 82.2 vs. BERT-large's 82.1—a 0.1-point advantage with 24% fewer parameters (Table 4). These are genuine parameter efficiency gains. However, the paper does not compare against more recent BERT training recipes (e.g., RoBERTa-style longer training, which typically improves GLUE scores substantially) or against BERT models trained on the same C4 dataset with the same hyperparameters—the BERT baselines are taken from published papers with their own training recipes, creating a confound between architecture and training procedure. The BERT baseline numbers (79.6 for base, 82.1 for large) are from Devlin et al. (2019), trained on Wikipedia and Bookcorpus, while M2-BERT is trained on C4—a different and arguably larger/more diverse dataset. This makes the parameter-efficiency claim somewhat confounded by training data differences.

Claim 2: M2 achieves up to 9.1× higher throughput at sequence length 4K. Supported cleanly for the HuggingFace baseline, but the more relevant comparison is against FlashAttention—which is what any practitioner would use for long-sequence BERT. Against FlashAttention, M2-BERT-base achieves 2.0× speedup at 4K and 3.1× at 8K (Table 5). These are meaningful improvements, but the 9.1× headline number should be understood as relative to an unoptimized baseline, not the best available alternative. Importantly, the paper does not report training throughput—only inference (forward pass). Since training involves both forward and backward passes with optimizer states, the throughput advantage might differ for full training workloads, and this is not measured.

Claim 3: M2 outperforms ViT-b by 1% on ImageNet with half the parameters. Supported by the reported numbers (79.5% vs. 78.5%, Table 7), but this comparison merits scrutiny. The ViT-b baseline (78.5%) is from the original Dosovitskiy et al. (2020) paper; more recent ViT training recipes (e.g., Steiner et al., 2021; Beyer et al., 2022) achieve substantially higher accuracy (often >82% Top-1 for ViT-b) through improved augmentation and regularization. M2-ViT uses the T2T-ViT training recipe (Yuan et al., 2021) which itself improves over the original ViT training—so the 1% gain may partially reflect training recipe improvements rather than architectural superiority. The paper would be strengthened by a comparison against a ViT-b trained with the same recipe (the ViT-b + Monarch entry at 78.9% uses the same recipe but still falls short of M2-ViT-b's 79.5%, providing a more controlled comparison that shows a 0.6-point gain from replacing attention with Monarch convolutions).

Claim 4: M2-GPT matches Transformer perplexity without attention or MLPs. Supported at the tested scales (145M–360M parameters, up to 15B tokens), but the scale is limited. These are relatively small language models by modern standards—the largest M2-GPT is 360M parameters, which is roughly BERT-large scale and far below GPT-3 (175B) or Llama (7B+). The consistent 0.2–0.3 PPL advantage over Hyena (Table 8) is encouraging, but it remains unknown whether this advantage persists, disappears, or reverses at billion-parameter scales and trillion-token training budgets. The paper's claim is appropriately hedged ("showing for the first time that it may be possible to match Transformer quality without attention or MLPs"), and the results support this possibility but do not confirm it at production scale.

Potential weaknesses in the experimental design:

• No training throughput comparison for M2-GPT. While BERT throughput is measured extensively (Tables 5, 6, 10, 11), the GPT-scale experiments report only perplexity, not wall-clock training time. Since one of M2's main claimed advantages is hardware efficiency, the absence of training throughput for the causal LM setting is a notable gap. The causal parameterization requires 2× padding (from Theorem 3), which means the effective sequence length is doubled—potentially eating into throughput gains.

• Single tokenizer for language experiments. All BERT and GPT experiments use the bert-base-uncased tokenizer. The interaction between Monarch matrix block sizes (which depend on sequence length $N$) and tokenizer vocabulary size is not explored; a different tokenizer producing different typical sequence lengths could shift the throughput tradeoff.

• No long-context evaluation for M2-BERT. The paper demonstrates that M2-BERT has higher throughput at long sequences (4K, 8K), but never evaluates whether the model actually benefits from longer context on downstream tasks. The GLUE fine-tuning uses sequence length 128 regardless of the pretraining sequence length, so the quality comparison does not measure long-context capability. This is an important missing experiment: if M2-BERT cannot effectively use longer contexts due to the convolutional architecture's inductive biases, the throughput advantage is less meaningful.

• No learning rate or hyperparameter sweep for causal LM. M2-GPT uses learning rates ($6\times 10^{-4}$ and $8\times 10^{-4}$) that are described as matching the Hyena models, but no sweep is reported to confirm these are optimal for the M2 architecture specifically. Given the architectural differences (no MLP, different mixing primitive), the optimal hyperparameters could differ.

• The throughput benchmarks don't account for the filter generation cost. In M2-BERT, the sequence mixer uses fixed Monarch matrices (DFT), but the convolution kernel $K_1$ must still be applied. The paper's throughput measurements include the full forward pass, so this cost is captured, but the breakdown between convolution compute and other operations (projections, normalizations, activations) is not provided, making it hard to assess how much of the throughput advantage comes from the Monarch convolutions specifically versus parameter reduction in the dimension mixer.

Experiments that would have strengthened the paper:

• A FLOP-matched comparison across architectures at the same total training FLOPs, rather than same step count or same token count, to isolate architectural efficiency from training FLOP differences.
• Direct measurement of FLOP utilization for the full M2-BERT forward pass (not just the Monarch operator in isolation, as in Table 2), to quantify how much of the theoretical Monarch advantage translates to end-to-end model throughput versus being diluted by other operations.
• Scaling M2-GPT to larger sizes (1B+ parameters) with at least one data point, to test whether the attention-free, MLP-free architecture remains competitive or hits a quality ceiling.
• Evaluation of M2-BERT on a long-document task (e.g., long-text classification, document QA) at the same long sequence lengths where it shows throughput advantages, to demonstrate that the architecture both scales efficiently and uses the additional context effectively.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted For and Dominates the Headline Efficiency Gains

The assumption: The entire compute-optimal framework—selecting which test-time strategy (search algorithm, revision ratio) to use for each prompt—depends on first estimating the prompt's difficulty. The paper's method for doing so requires generating 2048 samples per question and averaging either ground-truth correctness (oracle bins) or the PRM's final-answer score (predicted bins). The authors are transparent about this in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference for every prompt. For simplicity, our experiments do not account for this cost in the compute budget of our compute-optimal scaling strategy."

The consequence: The headline 4× efficiency gains (Figures 4 and 8) over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. Generating and scoring 2048 samples per prompt costs at least 2048 generations—an order of magnitude more than any of the test-time budgets studied (max 512 generations). In a realistic deployment, the total cost would be difficulty_estimation + strategy_execution, and the former would dominate the latter, potentially erasing or reversing the reported gains. The "predicted" difficulty bins (using PRM scores instead of ground-truth labels) remove the need for oracle access but do NOT reduce the sampling cost—2048 generations per prompt is still required. This makes the compute-optimal strategy, as presented, impractical for deployment without additional innovations.

What evidence exists in the paper: The paper explicitly flags this gap in Section 3.2 and Section 8, but provides no measurement of what the amortized cost would be. Figures 4 and 8 show that predicted difficulty bins closely track oracle bins, confirming that ground-truth labels are not necessary—but neither figure accounts for the cost of obtaining those predicted bins. The generation budgets on the x-axis of these figures (4, 16, 64, 256, 512) should be compared against a difficulty estimation cost of 2048 generations, making the "savings" relative to a best-of-N baseline that does NOT pay this upfront cost.

Mitigation status: The paper explicitly flags this as a key avenue for future work in Section 8, suggesting training models to predict difficulty directly from the question text, or developing adaptive difficulty estimation that amortizes into the solution process. Neither approach is developed or evaluated, so this remains an unresolved practical barrier. The paper's $4\times$ figure should therefore be understood as an upper bound on achievable efficiency rather than a realized deployment gain.

---

The $14\times$ Larger Model Baseline Is Weakened by Non-Compute-Optimal Pretraining and Greedy Decoding

The assumption: The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately $14\times$ more parameters. The larger model is trained by scaling parameters while holding training data fixed—following the LLaMA paradigm (Touvron et al., 2023) rather than Chinchilla-optimal scaling (Hoffmann et al., 2022). The authors acknowledge this in Section 7:

"This departs from the compute-optimal pretraining setting where the data and parameters are both scaled equally, but we choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding with no test-time compute augmentation—no majority voting, no best-of-N, no search.

The consequence: A Chinchilla-optimal model trained with $14\times$ more total FLOPs (scaling both data and parameters appropriately) would likely outperform a parameter-only-scaled model trained on the same data as the smaller baseline. This makes the pretraining baseline systematically weaker than it could be, potentially inflating the reported advantages of test-time compute. Furthermore, giving the larger model even a modest test-time compute budget (e.g., best-of-8 with majority voting) would create a substantially stronger baseline—comparing compute-optimal test-time scaling for the small model against greedy decoding for the large model compares two different inference strategies, not just two different model sizes. The headline result that test-time compute with a smaller model "can outperform a $14\times$ larger model" (Section 1) therefore overstates the advantage relative to what a practitioner would actually deploy—someone with access to a larger model would almost certainly use at least basic test-time augmentation (majority voting, temperature sampling) rather than pure greedy decoding.

What evidence exists in the paper: Figure 9 and the bar charts in Figure 1 show the FLOPs-matched comparison, with the larger model represented as horizontal lines (its greedy performance is constant regardless of test-time budget). The paper does not include any variant where the larger model receives test-time compute, nor does it compare against a Chinchilla-optimal larger model. The results in Table 5 of the referenced example paper show that even simple majority voting provides substantial gains—the larger model with majority voting would shift the horizontal lines upward, reducing or reversing the test-time compute advantage.

Mitigation status: The paper acknowledges the non-Chinchilla-optimal pretraining as a limitation explicitly in the Section 7 text quoted above, but does not address the greedy decoding asymmetry. The constraint to scale only parameters follows the LLaMA paradigm, which the authors argue is "canonical," but this choice makes the comparison most relevant to LLaMA-style training specifically, not to the broader question of optimal compute allocation across pretraining and inference.

---

Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Compensate for Fundamental Capability Gaps

The assumption: The compute-optimal framework assumes that the base model's pass@1 on a given prompt is non-trivially above zero—that is, the model can produce correct solutions at some non-negligible rate when sampled sufficiently. The framework optimizes how to spend test-time compute to find or refine those correct solutions, but does not create new capability.

The consequence: On the hardest questions (difficulty bin 5), no test-time strategy—search, revisions, or their compute-optimal combination—produces meaningful improvement regardless of budget. The paper's Figure 3 (right) shows bin 5 accuracy hovering at 1–3% for all methods across all budgets up to 256 generations. Figure 7 (right) shows bin 5 at roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% while the larger model achieves non-trivial accuracy. This means test-time compute cannot substitute for pretraining on problems outside the base model's capability range—any problem where the base model's pass@1 is near zero will remain unsolved regardless of how the test-time budget is allocated.

What evidence exists in the paper: The difficulty-bin breakdowns consistently show this pattern across all methods. Figure 3 (right): bin 5 flat near 1–3% for beam search, best-of-N, and all budgets. Figure 7 (right): bin 5 flat across all sequential-to-parallel ratios. Figure 9: bin 5 (blue line) flat near 0–5% for revisions while the $14\times$ larger model's greedy performance (star) is substantially higher. The paper is explicitly candid about this in the Section 7 discussion, noting that on hard problems "pretraining is almost always more effective" and that "test-time compute provides essentially zero benefit regardless of budget, meaning that some capabilities can only be acquired through pretraining."

Mitigation status: The paper is transparent about this limitation and characterizes it precisely, which is a strength rather than a weakness of the analysis. However, there is no proposed mitigation—this appears to be a fundamental boundary condition: test-time compute amplifies existing capabilities but does not create them. For deployment, this means practitioners should not expect test-time compute to handle genuinely novel or out-of-distribution problems that exceed the base model's training distribution. The difficulty estimator itself can serve as a detection mechanism—if the estimated difficulty falls into bin 5, the system should route the query to a larger model or a human rather than spending compute on strategies that will fail.

---

Revision Model Suffers from a 38% Correct-to-Incorrect Reversion Rate with Only Partial Mitigation

The assumption: The revision model is fine-tuned exclusively on trajectories where all in-context answers are incorrect, followed by a correct target (Section 6.1). This training data construction teaches the model to generate corrections, but provides no signal for what to do when the current answer is already correct.

The consequence: At inference time, approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step (Section 6.1). This is a direct consequence of the training distribution mismatch: the model never saw examples where the correct action is to preserve the current answer rather than change it. The paper mitigates this with majority voting or verifier-based selection across the entire chain (picking the best answer from any point rather than always taking the last revision), but these are post-hoc patches. They recover performance but at the cost of efficiency—the chain must generate additional revisions past the point where a correct answer was already produced, essentially wasting compute on revisions that will introduce errors before the selection mechanism picks the earlier correct answer.

What evidence exists in the paper: The 38% reversion rate is reported in Section 6.1. Figure 6 (left) shows that pass@1 at each step gradually increases through the chain but with variance—some steps degrade from the previous step. The paper's within-chain selection mechanism (verifier or majority) is described in Section 6.1 as the mitigation, and Figure 6 (right) shows sequential revision with within-chain selection outperforms parallel sampling. However, the ablation showing what performance would be without within-chain selection (just taking the last revision output) is not provided, making it difficult to quantify how much the reversion problem costs.

Mitigation status: The within-chain selection (majority voting or verifier-based selection) partially mitigates the symptom but does not address the root cause—the training data construction creates the reversion tendency. The paper does not explore training the revision model to recognize when no revision is needed (e.g., by including correct-to-correct trajectories in training data) or adding a "stop revising" signal. The ReST$^{EM}$ experiment (Appendix K, Figure 16) further demonstrates the fragility of revision training: attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, suggesting the approach is sensitive to training methodology in ways that are not fully understood.

---

Single Benchmark, Single Model Family—Generality of Difficulty-Dependent Scaling Trends Is Unverified

The assumption: All experiments use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The paper claims in Section 4 that "this model is representative of the capabilities of many contemporary LLMs," but provides no evidence that the main findings generalize beyond this specific combination.

The consequence: Several aspects of the results could be model-specific or benchmark-specific. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution—a model with different calibration, different error patterns, or different reasoning style might exhibit different difficulty-dependent scaling curves. The revision model's ability to learn from in-context incorrect examples depends on the base model's in-context learning capabilities, which vary substantially across model families. The MATH benchmark consists exclusively of competition-level math problems requiring multi-step symbolic reasoning—it is unclear whether the core finding (beam search hurts easy problems due to verifier over-optimization, revisions help easy problems but need balancing on hard problems) generalizes to other reasoning domains like code generation, logical reasoning, or scientific QA, or to tasks requiring factual knowledge rather than inference.

What evidence exists in the paper: None beyond the MATH benchmark and PaLM 2-S*. The paper does not include experiments on other reasoning benchmarks (e.g., GSM8K, MMLU, HumanEval) or with other model families. The test set of 500 questions, split into five quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on roughly 50 questions per fold per bin—a small sample that could introduce variance in the policy. The paper does not report confidence intervals on the compute-optimal scaling curves, making it difficult to assess whether the observed strategy differences across difficulty bins are statistically reliable.

Mitigation status: The paper acknowledges the single-model limitation implicitly in Section 4 ("we believe this model is representative") but it is not explicitly flagged as a limitation in the discussion (Section 7). The difficulty-dependent patterns are replicated across two distinct methods (PRM search in Figure 3, revisions in Figure 7) and across oracle vs. predicted difficulty bins, which provides some internal consistency but does not address the model-family or benchmark confound. This is a standard first-paper limitation—the framework is proposed and demonstrated on a single well-studied benchmark with a single model family—and replication on additional benchmarks and models is left to future work.

---

Causal M2 Convolutions Require 2× Input Padding, Partially Offsetting Sub-Quadratic Gains

The assumption: Theorem 3 establishes that causal M2 convolutions can be constructed by constraining the basis polynomial degrees and padding the input from length $n$ to $N = \lceil\sqrt{2n}\rceil^2 \approx 2n$. This $2\times$ expansion is the price of preventing modular wrap-around in the polynomial product, which would violate causality.

The consequence: The causal M2 convolution operates on vectors roughly twice as long as the actual input/output, meaning the effective sequence length for the Monarch operations is $2n$ rather than $n$. Since the complexity is $O(N^{3/2}) = O((2n)^{3/2}) = O(2^{3/2} n^{3/2}) \approx O(2.83 n^{3/2})$, the constant factor is nearly $3\times$ larger than a non-causal M2 convolution operating directly on length $n$. This partially offsets the sub-quadratic advantage over attention (which is $O(n^2)$ with no padding). In practice, this means the throughput advantage of M2-GPT over Transformer at a given logical sequence length is smaller than what the raw $O(n^{3/2})$ vs. $O(n^2)$ comparison would suggest, because the M2 model is internally processing a longer sequence.

What evidence exists in the paper: The throughput benchmarks in Section 5 (M2-BERT, Tables 5, 10, 11) are for the non-causal architecture, which does not require padding. The causal M2-GPT architecture is only evaluated for perplexity (Table 8), not for training or inference throughput. The paper does not measure wall-clock training time for M2-GPT or compare it to Transformer training time at equivalent sequence lengths. This is a notable gap: one of M2's main claimed advantages is hardware efficiency, but the causal variant—which is the one needed for auto-regressive language modeling, the most prominent LLM use case—has no throughput measurements. The $2\times$ padding cost means that M2-GPT's effective FLOPs per logical token are higher than a naive parameter count or complexity class comparison would suggest, and the paper provides no empirical data on whether the architecture still achieves net wall-clock speedups in the causal setting.

Mitigation status: The paper acknowledges the $2\times$ blowup implicitly in the theoretical development (Theorem 3 states the $n < N/2$ constraint) and the Appendix D.3 notes this is a "cost of ensuring causality." The generalization to arbitrary order $p$ (Appendix D.6) reveals that the blowup grows to $2^p \cdot n$ in the general case, which the authors flag as a limitation explicitly: "Whether this is possible to do (i.e. have a blowup of 2 instead of $2^p$) for $p > 2$ is an interesting direction for future work." For $p=2$ (the case used in all experiments), the $2\times$ padding is manageable, but the lack of causal throughput measurements means a practitioner cannot assess whether the net efficiency gain justifies switching from a well-optimized attention implementation (FlashAttention) to causal M2 convolutions, particularly at moderate sequence lengths where the quadratic attention cost has not yet become prohibitive.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper changes the conversation about efficient neural architectures by demonstrating that a single structured matrix class—Monarch matrices—can simultaneously address the two quadratic scaling bottlenecks that have defined Transformer costs: attention's $O(N^2)$ dependence on sequence length and the MLP's $O(d^2)$ dependence on model dimension. This is a conceptual reframing rather than a paradigm shift: the separate research communities studying sub-quadratic attention and sparse MLPs are revealed to be working on two instances of the same underlying problem, and a unified solution exists. Practitioners designing efficient architectures now have a concrete primitive that handles both mixing directions with the same computational characteristics—a significant simplification over hybrid designs that combine, say, FlashAttention for sequence mixing with a separate sparse MLP strategy for dimension mixing.

The reframing operates at the architectural level. Prior to this work, the mixer perspective (from MLP-Mixer and ConvMixer) established that many architectures could be viewed as alternating sequence mixing and channel mixing using the same operator, but the operators used were quadratic (dense MLPs or standard convolutions). M2 shows that the mixer perspective can be realized with a genuinely sub-quadratic operator, collapsing the distinction between "the attention module" and "the MLP module" into two instantiations of MonarchMatrix applied to different tensor dimensions. The pseudocode in Figure 1 makes this tangible: the same MonarchMatrix class, with the same blockdiag_matmul and rearrange operations, handles both axes.

The practical significance is shifted architectural priorities. The paper's Table 1 and Table 2 make a forceful empirical case that asymptotic complexity alone is an insufficient metric for architecture design—FFT convolutions at $O(N \log N)$ with 3% FLOP utilization are slower than Monarch convolutions at $O(N^{3/2})$ with 41.4% utilization at realistic sequence lengths. This inverts the standard narrative (lower FLOP count = faster) and directs research attention toward hardware-compatible primitives rather than asymptotically optimal ones. A research direction that previously seemed promising—using FFT-based convolutions as the default sub-quadratic sequence mixer—becomes less attractive, while the GEMM-based structured matrix approach (Monarch, butterfly, kaleidoscope) becomes a more compelling foundation on which to build. The RTX 4090 vs. A100 comparison further suggests that Monarch matrices' performance will improve with future GPU architectures that have larger caches, making this a forward-compatible choice.

The theoretical contribution—characterizing Monarch convolutions as multivariate polynomial evaluation and deriving causality conditions through degree constraints—provides a diagnostic framework that resolves a key barrier to adoption. Before this work, it was not known whether structured matrix primitives like Monarch could be used in auto-regressive settings at all, because the modular polynomial reduction that makes them efficient also makes them non-causal (higher-degree terms "wrap around" to lower indices). The polynomial framework explains why non-causality occurs and gives constructive how to fix it: constrain the minimum and maximum degrees of the basis polynomials. This converts a binary question ("is this causal?") into a design space with tunable parameters, enabling future work to explore different degree constraint regimes and their quality-efficiency tradeoffs.

The most surprising empirical result is M2-GPT matching Transformer perplexity on The PILE at 360M parameters without using attention or MLPs (Table 8). This is an existence proof that challenges a widely held architectural assumption—that MLPs are essential for per-token nonlinear feature transformation. M2-GPT's dimension mixer is the identity function; all feature mixing is absorbed into the sequence mixer through the H3 head structure's outer products. This does not mean all future language models should drop their MLPs, but it establishes that the minimal set of operations needed for competitive quality is smaller than previously demonstrated, and that the head structure provides sufficient mixing capacity at this scale. The result also validates the theoretical causal parameterization: without Theorem 3, M2-GPT could not exist, so the theory enables the architectural minimality.

The paper does not reconcile major prior contradictions in the literature the way some works do—the structured matrix and efficient attention communities were not making conflicting claims so much as solving different halves of the problem. However, it does reconcile the tension between asymptotic efficiency and hardware reality that has made some sub-quadratic methods (like FFT convolutions) fail to deliver expected wall-clock speedups. The reconciliation is empirical: measure FLOP utilization, not just FLOP count.

Several research directions become less attractive after this work. Pure FFT-based sequence mixers that don't address the MLP bottleneck now face a stronger alternative that handles both axes. Sparse MLP methods that reduce FLOPs but achieve poor utilization due to irregular memory access (the paper cites examples that are "actually slower than their dense counterparts") are similarly challenged by Monarch's demonstrated high utilization. The paper implicitly argues that the right architectural primitive should be chosen for hardware compatibility first, with asymptotic scaling as a secondary consideration—a priority inversion relative to much prior work.

Follow-Up Research This Work Enables

Scaling M2-GPT to billion-parameter scales with throughput benchmarks. The paper's most striking result—attention-free, MLP-free language modeling at 360M parameters—is demonstrated only at modest scale and without throughput measurements for the causal variant. A strong follow-up would train M2-GPT at 1B, 3B, and 7B parameter scales on The PILE or a comparable dataset, measuring both perplexity and wall-clock training throughput against equivalently-sized Transformers (with FlashAttention and standard MLPs) and Hyena (with FlashAttention-style convolutions and MLPs). The key questions: (1) Does the 0.2 PPL advantage over Transformers persist, grow, or shrink with scale? (2) Does the causal M2 convolution's $2\times$ padding cost (from Theorem 3) offset sub-quadratic gains in practice, and at what sequence length does M2-GPT become faster than FlashAttention Transformer for training? (3) Does removing the MLP cause a quality ceiling at larger scales where feature mixing capacity becomes a binding constraint? The paper's Table 8 provides the scaling trend at 145M and 360M—extending this to 1B+ would reveal whether the architecture scales gracefully or hits diminishing returns.

End-to-end training throughput benchmarks for M2-BERT, not just inference. The BERT throughput results (Tables 5, 10, 11) measure forward-pass tokens/ms—inference only. Training involves forward and backward passes with optimizer states, which changes the memory bandwidth and compute balance. A follow-up would measure training throughput (tokens/second end-to-end, including backward pass and optimizer steps) for M2-BERT-base and M2-BERT-large against FlashAttention BERT at sequence lengths 512 through 8192 on identical hardware (A100-80GB), reporting both tokens/second and peak GPU memory usage. This would determine whether M2's inference advantage translates to training—where the backward pass through Monarch permutations may have different data movement costs than the forward pass, and where optimizer state memory (which scales with parameter count) interacts with activation memory (which scales with sequence length).

Long-context downstream evaluation of M2-BERT. The paper demonstrates that M2-BERT has high throughput at sequence lengths up to 8192, but all downstream evaluation (GLUE) uses sequence length 128 regardless of pretraining length. There is no evidence that M2-BERT's convolutional architecture actually benefits from longer context on tasks that require it. A critical follow-up would fine-tune and evaluate M2-BERT on long-document tasks: SCROLLS (which includes long-text QA, summarization, and NLI with inputs up to 100K tokens), the Long Range Arena benchmark, or a document-level NLI task, comparing against BERT with FlashAttention at equivalent sequence lengths. The question is not just whether M2-BERT is faster at long sequences, but whether its inductive biases support effective use of that context. Long convolutions with fixed frequency-domain kernels might struggle with tasks requiring sharp position-specific attention (e.g., coreference resolution across distant mentions), and this would manifest as a quality gap even if throughput is superior.

Learning the Monarch matrices in the sequence mixer at scale. The paper's main experiments use fixed DFT/Inverse DFT Monarch matrices in the sequence mixer, achieving competitive quality with only the kernel $K_1$ learned. Table 15 (Appendix B.6) shows that on sequential CIFAR, learnable Monarch matrices improve accuracy by 1.5 points over fixed DFT—but this is a small-scale vision experiment. A scaling study would compare fixed-DFT vs. fully-learnable Monarch matrices in the M2-BERT sequence mixer, measuring both pretraining MLM accuracy and downstream GLUE quality, to determine whether the expressivity gain from learning the mixing basis provides benefits at language modeling scale or whether the DFT basis is sufficient. The concern is that learning $2N^{3/2}$ additional parameters in the sequence mixer could overfit on pretraining or add optimization difficulty; the fixed-DFT variant has the advantage of being a known good basis. If learnable Monarch matrices provide even a small quality improvement at BERT scale, they become the default choice; if they don't, the architecture can be simplified to always use fixed DFT for the sequence mixer.

Combining M2 convolutions with alternative tokenization and patching strategies. The paper uses standard subword tokenization for language modeling and standard 16×16 patch tokenization for images. However, the Monarch block size $\sqrt{N}$ depends on the sequence length $N$, which is determined by the tokenizer. A byte-level or character-level tokenizer produces much longer sequences for the same text, shifting the throughput tradeoff in M2's favor (where the sub-quadratic advantage grows with sequence length). A follow-up could train M2-BERT using byte-level tokenization (e.g., ByT5-style) with sequence lengths of 2048–4096, measuring both quality and throughput against a byte-level Transformer baseline. The hypothesis is that M2's throughput advantage over attention becomes larger at the longer sequences produced by byte-level tokenization, making it a better match for character-aware or byte-level models. Similarly, for vision, testing smaller patch sizes (8×8 or 4×4) would produce longer sequences and test whether M2-ViT can handle finer-grained spatial mixing efficiently.

Cheap difficulty estimation for compute-optimal test-time scaling with M2 inference. The paper's Section 2 and 3.2 discussion of "compute-optimal test-time scaling" (from the referenced example) applies directly to M2 at inference time: different prompts or inputs may benefit from different amounts of sequence mixing computation. A concrete follow-up would apply the difficulty-estimation framework from the referenced example paper to M2-BERT inference: estimate per-input difficulty using a small number of forward passes (4–8 samples) and allocate Monarch convolution depth or block size adaptively—easier inputs get order-1 Monarch matrices (faster, lower quality), harder inputs get order-2 or higher (slower, higher quality). This would create an adaptive inference system that dynamically trades off compute for quality per input, leveraging the fact that Monarch's order $p$ is a tunable knob controlling the complexity-quality tradeoff. The key measurement would be whether such adaptive allocation recovers efficiency gains comparable to the referenced paper's $4\times$ improvement over a fixed-configuration baseline.

Practical Applications and Downstream Use Cases

On-device and edge deployment of BERT-quality language models with long-context support. M2-BERT-base (80M) achieves 79.9 average GLUE score with 27% fewer parameters than BERT-base, and Table 6 shows it achieves 6.5× lower CPU inference latency than BERT-base at sequence length 8192 on a standard cloud CPU instance (1820 ms vs. 11820 ms). For on-device applications—smartphone keyboards, voice assistants, document scanning on laptops—where GPU access is limited or unavailable and long documents need processing, M2-BERT provides a concrete recipe: train an 80M-parameter M2-BERT variant on the target domain, deploy on device CPU, and process documents up to 8K tokens with latency that would be infeasible for standard BERT. The parameter reduction also reduces storage and memory footprint, which is critical for mobile deployment where app size and RAM are constrained. The CPU advantage only emerges at sequence lengths ≥1024 (Table 6), so this is specifically valuable for long-document applications, not short-text tasks like query classification.

Cost-efficient batch inference pipelines for text classification and embedding at scale. The throughput results in Table 5 are directly actionable for organizations running large-scale inference: at 4K sequence length, M2-BERT-base (80M) processes 353.9 tokens/ms on an A100 vs. 179.1 for FlashAttention BERT—nearly 2× throughput. For a pipeline processing millions of documents (e.g., content moderation, semantic search indexing, document clustering), this translates to roughly halving the GPU-hours required. The recipe is: replace the BERT-based encoder in existing pipelines with an M2-BERT variant at the same accuracy level, keeping the same fine-tuning workflow (the GLUE results in Tables 3 and 4 show M2-BERT fine-tunes with standard procedures), and realize throughput gains that scale with document length. Additionally, because M2-BERT has fewer parameters, the model loads faster from disk and occupies less GPU memory, enabling larger batch sizes and further amortizing the per-batch overhead.

Image classification with reduced model size for resource-constrained vision tasks. M2-ViT-b achieves 79.5% ImageNet Top-1 accuracy with 45M parameters—1% higher than ViT-b at 87M, and competitive with ResNet-152 at 60M (Table 7). For applications such as satellite imagery analysis, medical imaging on edge devices, or embedded vision in autonomous systems, where model size is constrained by hardware (FPGAs, mobile GPUs, microcontrollers with ML accelerators), M2-ViT provides a path to ViT-quality accuracy with roughly half the parameter budget. The architecture uses standard ViT input processing (patch tokenization, positional embeddings), so existing ViT training pipelines can be adapted with minimal changes. The CIFAR-10 results (Table 14) demonstrate this transfers to smaller-scale tasks as well, with HyenaViT-M2 (741K parameters) outperforming the standard ViT (1.2M parameters) by 2.2 percentage points. The key practical advantage over using a smaller ViT (which would also reduce parameter count) is that M2-ViT maintains the full model width, preserving representational capacity per token while reducing the parameter cost of mixing across tokens and features.

When to Prefer Monarch Mixer Over Alternative Architectures

The paper positions M2 against specific named alternatives—Transformers with attention and MLPs, FFT-based long convolutions (Hyena/H3), and sparse MLP variants—making the tradeoffs explicit enough to warrant a decision framework:

• Prefer M2 over Transformers with FlashAttention when: (1) sequences are long enough that the throughput advantage is substantial—at sequence length 2048 or longer on GPU (Table 5) or 1024 or longer on CPU (Table 6); (2) the model dimension $d$ is large enough that the MLP cost matters—M2's dimension mixer reduces the $O(d^2)$ MLP cost to $O(d^{3/2})$, so wider models benefit more; (3) parameter count is a binding constraint—M2-BERT achieves comparable quality with 24–27% fewer parameters (Tables 3, 4); (4) the task is non-causal (BERT-style, image classification) where the padding overhead of Theorem 3 doesn't apply. At short sequences (512 or below) and for standard BERT-base-scale models, FlashAttention BERT is comparable or slightly faster on GPU (Table 5: 386.3 vs. 367.4 tokens/ms), so M2's advantage only materializes at longer contexts.

• Prefer M2 over FFT-based long convolutions (Hyena, H3) when: (1) hardware utilization matters more than asymptotic complexity for your deployment GPU—on the RTX 4090, M2 achieves 41.4% FLOP utilization vs. ~3% for FFT at 64K input dimension (Table 1); (2) you need to address both sequence length and model dimension bottlenecks, not just sequence length—FFT-based methods leave the MLP quadratic cost untouched; (3) you want a single primitive to implement, optimize, and maintain rather than separate attention and MLP replacements. The tradeoff: M2's $O(N^{3/2})$ complexity is asymptotically worse than FFT's $O(N \log N)$, so at extremely long sequences (100K+), FFT-based convolutions may eventually overtake M2 in wall-clock time despite lower utilization—the paper does not characterize this crossover point.

• Prefer scaling pretraining (larger models) over M2 when: (1) the problem requires fundamentally new capabilities that only larger-scale pretraining provides—M2 is an architectural efficiency improvement, not a capability creator; (2) the deployment setting is latency-critical and short-sequence, where FlashAttention's lower constant factor dominates—at sequence length 512 on GPU, FlashAttention BERT is 8% faster than M2-BERT at matched parameter count (Table 10); (3) the ecosystem around Transformers (pre-trained checkpoints, optimized inference servers, compiler integrations) provides value that outweighs architectural efficiency—M2 requires training from scratch or adapting existing checkpoints, with less mature tooling.

• Prefer M2 for training from scratch rather than finetuning pretrained Transformers when: the compute budget for pretraining is the dominant constraint, and the task allows training a domain-specific model. M2's reduced parameter count means lower pretraining cost for a given quality target, and because the architecture is structurally different from Transformers, it cannot benefit from initializing with pretrained Transformer weights. For practitioners who routinely pretrain domain-specific BERT models (biomedical, legal, scientific literature), M2 offers a more parameter-efficient starting point, provided the domain data is sufficient for training from scratch.