MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention
Authors: Can Yaras
Alec S Xu
Pierre Abillama
Changwoo Lee
Laura Balzano
Venue: NeurIPS 2025 spotlight ()
Keywords: structured matrices
efficient attention
monarch
Primary area: deep_learning

Abstract: Transformers have achieved state-of-the-art performance across various tasks, but suffer from a notable quadratic complexity in sequence length due to the attention mechanism. In this work, we propose MonarchAttention -- a novel approach to sub-quadratic attention approximation via Monarch matrices, an expressive class of structured matrices. Based on the variational form of softmax, we describe an efficient optimization-based algorithm to compute an approximate projection of softmax attention onto the class of Monarch matrices with $\Theta(N\sqrt{N} d)$ computational complexity and $\Theta(Nd)$ memory/IO complexity. Unlike previous approaches, MonarchAttention is both (1) transferable, yielding minimal performance loss with no additional training, even when replacing every attention layer of the transformer, and (2) hardware-efficient, utilizing the highest-throughput tensor core units on modern GPUs. With optimized kernels, MonarchAttention achieves substantial speed-ups in wall-time over FlashAttention-2: $1.4\times$ for shorter sequences $(N=256)$, $4.5\times$ for medium-length sequences $(N=4K)$, and $8.2\times$ for longer sequences $(N=16K)$. We demonstrate the quality of MonarchAttention on diverse tasks and architectures in vision and language problems, showing that it flexibly and accurately approximates softmax attention in a variety of contexts.

MonarchAttention. Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention
Can Yaras University of Michigan cjyarasCumi ch. edu
Alec S. Xu
Pierre Abillama
University of Michigan alecx@umi ch. edu
University of Michigan @bbillam@uii ch, edu
Changwoo Lee University of Michigan cwooleeGumic ch. edu
Laura Balzano
University of Michigan g girasoleGumich. edu
Abstract
Transformers have achieved state-of-the-ar performance across various tasks, but suffer from a notable quadratic complexity in sequence length due the attention mechanism. In this work, we propose MonarchAttention - novel approach to sub-quadratic attention approximation via Monarch matrices, an expressive class of structured matrices. Based on the variational form of softmax, we describe an algorithm compute an approximate projection of softmax attention onto the class of Monarch matrices with O v dd computa- tional complexity and ( memory/IO complexity. Unlike previous approaches, Monar chAttention is both (1) transferable, yielding minimal performance loss with no additional training, even when replacing every attention layer of the Trans- former, and (2) hardware-efficiee
efficient optimization-bant
utilizing the throughhtt tensor core units modern GPUs. With optimized kernels, Monar chwarentttttttii tial speed-ups in wall-time over FlashAttentions 2: 1.4x shorter sequences (N = 256), 4.5x for medium-length sequences (N = 4K), and 8.2x for longer sequences (N = 16K). We demonstrate the quality of MonarchAttention on diverse tasks and architectures in vision and language problems, showing that it flexibly and accurately approximates softmax attention in variety of contexts. Our code is available at https: github. com/ womm/iiyaras/moarcch-aaaa
achieves substan-
Introduction
have become the dominant architecture for generating and processing various data modalities, such as text (Brown et al.) 2020). images (Dosovitskiy et al., (202T) and speech (Radford et al.) (003). Central the Transtormer" success is attention, the mechanism through which complex interactions within sequential data are captured through weighted combinations of embeddings at every position in the sequence. Famously, the attention mechanism has a quadratic-time complexity ((N2d) in the length of the sequence N, where d the head dimension, which is key bottleneck for both training and inference, particularly in long sequence problems. To address this, numerous works have proposed sub-quadratic substitutes for attention. Yet, such approaches are either (1) not transferable, requiring training from scratch or fine-tuning of existing models, or (2) do not yield speed-ups in practice (except on extremely long sequences) due to a gap between theoretical complexity and practical considera- tions for modern GPUs, especially compared to highly optimized implementations (Dao et al. (0226).
Over the past decade, Transformers (Vaswani et al. 2017)
In this work, we propose Monarchattention: novel sub-quadratic attention substitute
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

cosformer
Leearr attettion
perto neer
rost forner
monerch attention
sal teax
Figure 1: Approximation softmax attention via Monar chAttenionn By directly optimizing the softmax variational objective constrained to Monarch matrices, Monar chattentioo yields accurate zero-shot approximation softmax attention compared to other hardware-friente friendly, efficient attention baselines, Attention maps extracted from RoBERTa on the SQuAD dataset in Section 4
based on approximating the attention matrix via Monarch matrices (Dao et al. 2022a). a class of expressive structured matrices. At first glance, this is computationally infeasible for sequence length N, computing an exact projection onto the set of Monarch matrices has a super-quadratic O(N2/N) / -time complexity, not to mention that we need to form the entire N x N attention matrix. Instead, we reframe the computation of the attention matrix an optimization problem terms of the variational form of softmax, and exploit low-dimensional structure in the variational objective when constrained to the set of Monarch matrices - this yields a sub-quadratic O (.NNNN))tmm approximation, where is the head dimension. This approach combines two optimization based improvements for its success: the variational form of the softmax nonlinearity with structured factorization approach analogous to those for low-rank approximation of matrix (Chi et al.) 2019), where rather than computing a full SVD and truncating to the desired rank, one can more efficiently minimize a Frobenius norm objective constrained to the set of low-rank matrices. We briefly review prior work on structured matrices, including Monarch matrices, well as existing approaches to efficient attention.
Structured & Monarch Matrices. We use the phrase "structured matrices" to mean those that admit sub-quadratic storage and matrix-vector multiplication, such as low-rank or sparse matrices. There are many useful classes of structured matrices, such as those with low displacement rank (Kailath et al. 1979), which includes Toeplitz. Hankel, Vandermonde, Cauchy matrices (Pan (001); ; orthogonal polynomial transforms (Chihara, 2014). which includes discrete Fourier/cosine and Hadamard transforms; butterfly factorizations (Da et af. 2019), which implement fast matrix-vector multiplication via recursive divide-and-come con-quuer algorithm similar to that of fast Fourier transforms (FFTs); and Monarch matrices, an expressive family of structured matrices (generalizing butterfly matrices and thereby many fast transforms) that overcome unfavorable memory access patterns typical to FFT-like algorithms by implementing matrix products via batched dense matrix multiplications (also called matmuls) on fast tensor cores found in modern GPUs.
Sub-Quadratic Attention. Nearly all approaches to sub-quadratic attention approximate the attention matrix by structured matrix, specifically low-rank and/or sparse.
Low-Rank. Motivated by
embeddings, Wang et al. (020)) propose sketching the key and value matrices along the sequence dimension via learnable projections. Katharopoulos al. (2020) introduce linear attention, where the exponential kernel is approximated via inner products of queries and keys lifted via some feature map. Several follow-up works proposed various feature maps, such as the exponential linear unit (ELU) (Katharopoulos et al. 2020), random positive features (Choromanski et al. 202T), rectified linear unit (RcLU) with cosine reweighting (Qin et al. 2022). and learnable single-layer multi-layer perceptrons (MLPs) (Zhang al. 2024g. Xiong et al. (2021) use the Nyström method for computing low-rank approximations by sampling rows and columns.
Sparse. Child et al. (2019) introduce sparsity by applying fixed, structured sparse masks on the attention matrix. In particular, [Cenn et al. (2022) propose particular block butterfly matrix for the sparse mask, which is more hardware-friende
at the cost of reduced expres- siveness. Those that do not enforce structure on the sparsity pattern include Kitaev et al. (2020); Daras et al. (222)) where they utilize locality-sensitive hashing (I-S)) on shared query/key vectors to only compute attention within clusters of similar tokens.

---

Low-Rank + Sparse. Inspired by robust PCA, Chen et al. (2021) decompose the attention matrix into sum of two matrices: an unstructured sparse component using LSH and low-rank component that is constructed via linear attention. HHan et al. (2024) propose to subsample columns of the non-normalized attention matrix based on row norms of the value matrix, while estimating the softmax normalization factors from few large elements via LSH.
We note that there are significant drawbacks to the approaches described above. Pure low-rank methods are often fast and hardware-friendly, but are typically not suitable as drop-in replacements for attention in pre-trained Transformers due to the prevalence of "strongly diagonal". high-rank attention matrices where attention weights are concentrated locally in sequence. Making up for this with a fixed sparsity pattern does not allow for data-dependent support of the attention matrix, necessary for zero-shot conversion. Finally, sparsity pattern improve on accuracy over low-rank approximations but suffer from significant overhead due to GPU incompatibility. MonarchAttentiors both worlds: it is fast and hardware-friendly due to utilization of tensor cores for batched matmuls, while computing highly accurate approximations to the extent that directly replace softmax attention with no additional training.
Sasiityy/SH--beeeee approaches that do not have fixed
the other hand, achieves the best of
We conclude this section by discussing closely related works. Dao et al. (202bb) propose FlashAttention on, an IO-aware streaming algorithm for computing exact softmax attention. We show in Section 3 that each step of MonarchAttention can be written as a FlashAttention-/e computation, allowing for similar IO savings FlashAttention - in fact, we demon- strate that MonarchAttention achieves strictly better worst-case IO complexity compared FlashAttention. We also note that Dao et al. (2022b) propose further accelerate FlashAttention using block butterfly attention masks, so MonarchAttention can be viewed as a generalization of block-sparse FlashAttention to more general Monarch matrices. Finally, MonarchAttent is closely related to Monarch Mixer (Fu et al. 2023), , mixer-type architecture (Tolstikhin et al. 202T) that utilizes Monarch instead of dense matrices for token and channel mixing. MonarchAttent also uses Monarch matrices for mixing tokens however, it is based on the attention operation which is data-dependent, unlike Monarch Mixer.
Organization. In Section we discuss preliminaries on (softmax) attention, and Monarch matrices. In Section 3 describe the MonarchAttention algorithm and implementation. In Section { we evaluate MonarchAttention in variety of settings for zero-shot conversion to sub-quadratic atten- tion and benchmark its implementation. In Section 5 we discuss limitations and future directions.
2 Preliminaries
Notation. We use [N] to denote the index set f1. .,,.,..... We use A to denote the (N 1) dimensional unit simplex, given by A {a IR a ł 0, (IN, a) = 1}. We denote the m m identity matrix by In. We use the notation A,jk to denote an element of a 3-way tensor, and A to denote a slice. We use ōu to denote the Kronecker delta that is 1 if 1- = 1 and otherwise O.
Softmax.
The softmax function IR: simplex, and is defined as
A maps N real numbers to the (N dimeensiio unit
exp(2:) Vi [N]] Σ, exp(z,)
[softmax(z)]
An alternative definition
Blondel et al. 2019) is given by the following variational form: softmax = arg max (a, H()),
a
where H(a) = -. a, log a; is Shannon entropy. See Appendix for equivalence of and
Attention. Given query, key, value matrices K, V E Nxd where N is the sequence length and d is the head dimension, single head of standard softmax attention computes cotepuisrScompuees
O softmax (QKT) V.
(3)
'Typically, the QK matrix is scaled by factor of 1/2 but this be absorbed into Q
3

---

where the softmax function is applied across rows. The computational complexity of attention is ((N2d) for each forward pass, because the matrices Q, K, V are data-dependent.
Monarch Matrices.
Given N m X b for integers m, b, we define a block rank-one matrix
B E NxN.
B
B im
B
where Bys Rik mamb
B
Bbm
for some Likk 1m R&s I8 for j E [b] and ls E [m]] follows that
L
R2
B =
P
L
R
where Lj Im for j E [b] and Rs E Babbb for h E [m], and P E RNN is a "transpose" permutation matrix whose (i + 1)-th row given by Ca(()+1 where
i1 = b. (i mod m) +
ie (O....,N-11,
m
Given the above, a Monarch matrix M E NXN is given by M = P other words, a row-permuted block rank-one matrix. When m = VN, storing such matrix requires only (NNvN) space, while matrix multiplication (matmul) with a matrix E Nxx can be computed efficiently in ((NvNd) operations (as opposed to @(N2d) for dense matrices) with batched matmuls and transposes:
MV O, where 0b(1-1)45.0 = L Yjku Σ Riji Vo 11+i
(4)
for i,j E [b], k E [m]. useful characterization of M is in block form:
M11
Mim im
M
(5)
Mm
M M
where Mi. E B [MR]a L Lysa Raji, Vi,j E [6],k, k, E [m].
3
MonarchAttention
The main goal of MonarchAttentionn is to find a Monarch matrix E NxN in (( time such that M softmax(QK Then, we can approximately compute the output O MV using efficient matmul. We can do this by viewing the softmax operation as optimization problem via its variational form P whose objective can be efficiently maximized with exact alternating steps when constrained to Monarch matrices. As shown in Figure this yields highly accurate approximations to the softmax attention matrix.
Softmax Objective.
First, from
2 we can write
(QK
arg max f(A:Q. = (A, HH()) AEANXN
H(A),
(6)
where A N denotes matrix whose rows lie A and H(A)= - Eis Aig log Aij. Here
demonstrate how handle this objective function when A has Monarch structure, and defer discus- sion of the simplex constraint to the next paragraph. For dense matrix A, computing {(A; Q K)
2 Pæ corresponds to row-major reshaping E I2 to mmb transposing to R. 6xm, then row-major flattening back to N See Appendix for example.

---

InageNet
SQuAD
186
108
00
80
Top 5 accuracy
monarch- at tention
68
60
per
cos forner
Linear-a aattrrttnnn ryntoroee softrax
48
40
20
8
8
1.00
1.25
1.50
1.75
2,00
2.25
2.
2.75
Total Attention FLOPs
168
Total Attention FLOPS
1e9
Figure 2:
Zero-shot conversion of attention layers for image classification and question answer- ing. We vary hyperparameters for various baselines to evaluate model quality compute tradeoff. Left. Top-5 accuracy vs. total attention FLOPs across all layers for ViT on ImageNet. Right. Fl score total attention FLOPs across all layers for RoBERTa on SQuAD.
requires d(V2d) operations, which is the same as computing (QQKT) directly. However, when A Monarch matrix M- B, we have
[(P B - (P B,QK H B) =(B H(B) E( Qj, Kx),
where @ PQ, and 2, E Rmxdd K& Roxd are the j-th and k-th block of rows of Q K respectively. Then, for each [b], e [m] we evaluate f on the rank-one matrix Byx Lyx
f(Byk: Qy. K2) (Lzr R ²(Ke) - E L5
1,1
[ Like K2 Raj) Lik log Ljkt
Σ
Litt Raiis log Raji
1,
1,
(Q] Like K Rag) + (1RR)) H(Lgx) (1J H(Rxs)-
Thus, for each j E [b), E [m] we only need O((m + b)d) operations to compute [(Bzk; Qj, K&) due to Q Lyk and K, Raj- We emphasize that the rank-one strueture implies separability of the entropy term, meaning we can compute the entropy on Lyk and Raj individually and avoid the need materialize Byk, which would ineur ((mb) cost as opposed O(m + b). Since there are b many By matrices, we have in total @((m² 62m)dd operations to compute f(M; Q, K), which for b VN is ONN) improving on the dense computation factor of /N.
Alternating Maximization with Constraints. We will now explain the alternating maximization approach for optimizing f. When L is fixed, the objective is concave in R, and vice-versa- therefore. we can derive closed form expressions via KKT conditions for and that maximize f with one of L or R fixed, which will constitute a single update step. Evaluating (and therefore differentiating) f w.r.t L and R can be done in O(N VNd) time, which will be the same complexity one of these steps. For T steps, this will require ( /VN) computation: provided that T o( /N), this will still sub-quadratic. However, the constraint M E A NxN presents a challenge in current form, since this requires materializing M to check that each entry non-negative. Instead, use the fact that
L E A Rej E A Vj E [), vk, [m] = M E NXN.
i.e., slices of L, individually lying on the unit simplex is sufficient to enforce the constraint on M. This is casily seen from (5) obbioossl if LiKt, Raji 0, then [Mulji 2 0. Moreover, this also
5

---

16.4
monarch attenntion
16.8
nystronforme
ROUSE-1
ROUGE
15.8
teaa
32
Sequence Length
+ 1824
15.4
2848
4806
15.2
8192
16
10
16th
10
Attention FLOPs
Attention FLOPs
Figure 3: Zero-shot conversion of attention layers for long sequence summarization, We vary the sequence length of the text to be summarized to evaluate model quality vs compute tradeoff. We report recall-based ROUGE-1 and ROUGE-L scores total attention FLOPs across all layers for BART on BookSum-chapters.
enforces the sum-to-one constraint, as rows M sum as
Σ [Muli
E
Σ L Σ Rajs
1.
We now present the updates for L, R. Initializing L (0) dд1 identity, have
R(O) softmax; (2
z (1)
Rkji 3 CR,kj*
(7)
(4) softmax: (z(1)
( Lijht B Cljke
(1)
(8)
for E [T], where softmax,, softmaxx are applied along is and index dimensions respectively, and
((t) R.kii
Σ a R.kju
(t) K
(t)
Σ
(t-1)
kin
jkl
ilus
jkt
((0)
Σ
a (t)
Lijku flas
(t) Lijko
Σ
R
K kios
CLih
(8)
Σ
R log Rijo
(10)
where Q Kei Rdd are the (b (1- 1) + 5- and (b (k - 1) + i)-th rows of Q and K respectively. The full derivation is provided in Appendix C.T After T steps, we obtain the final Monarch approximation (T)) ({QKT) with factors ((r) and R(T), from which we output (T)V using 4 A naïve implementation of the full algorithm is provided in Appendix C.2 We
discuss in Appendix Cow padding can be incorporated into MonarchAttenti for when N
not divisible by b. We empirically demonstrate the fast convergence of MonarchAttention in Appendix E.
Implementation.
To minimize data movement and memory usage on GPU, we do not materialize L or R in binh-bandwidth memory (HBM). In addition to Q, K, V, O, we only need to maintain s a
other intermediate values are only materialized in on-chip SRAM, fusing all operations between the a (and (t) similarly e) variables. For instance, from the above update equations, the computation of a from OR is given by
(t) az c et from p and (10), resulting ((Nd) additional memory.
Σn a K
a
(0) Lijkv
softmax,
R.kjv (0)
K kios
CR.kj
"The a and e variables can share the same memory location as those corresponding to R can be derived from L (and versa)

---

def ftl.cl_keeeellee cR, Kb) # Computes at, cL from al, cB softmax (bmm (aR, Kb. transpose(1, 2)) / cR(: i, None], dim-2) 3 cL sum (Rg(R), dim=2). transpose(O, 1)
2
4 aL = (Rm Kb) transpose(O, 1) return al, cL
def ar_ kernel cL, Qb) Computes al, from al, cl sof tmax (bmm (aL, Qb transpose (1, 2)) cLC:, i, None], dim-1) cR - sum (L, dim-2) transpose (0, 1)
aR = bam (L,, Qb) transpose 1) return aR, cR
13
def ceeeeeeeeeeaaaeeeeeeeeeeeaa_aa cR, Kb, Vb): Puse al_ ct_ kernel Monarch matmul 1st step sof tmax (bmm (aR, Kb. transpose 2)) / cR[:, i, Nonel, dim=2) cL sum 1Rg(R), dim=2) transpose 1)
aL bem (R, Kb) transpose (0, 1) y bam (R, Vb). transpose(O, 1)
17
return al, cL
def kerkerellle (al, y, cL, Qb) Nonarch matmul 2nd step softmax (bmm (Qb, aL. (ranspose 2)) cLC: None, :,, dim=2)
22
z bmm(L, y) transpose 1)
return 2
def monarch monarch_attentiinnn attentionnnn (Q, K, T): $ 80, K, (M, d), I: number of steps Qb Q r.reshape (m, b, d) transpoae 1)
26
Kb K. reshape (m, b, d)
Vb = V reshape (m, b, d)
aR Q reehhape (m, d)
cR ones(m, b)
for in rT-ge
al, cL all_c_ (aR, cR, Kb) aR, ar_cr_kernel (aL, cL, Qb)
36
aL, y, cL 2 kerkeellllee (aL, y, cL, Qb) reshape d)
cR, Kb, Vb)
return o
Figure 4: Python-like code for MonarchAttention.. Each kernel materializes all intermediate
arrays in SRAM reduce data movement.
which can be seen as a batched attention computation, meaning we can implement FlashAttention-Iie kernel to reduce IO between HBM and SRAM. However, several aspects of this computation make it particularly 10-efficient. Besides the fact that K acts both the K and V matrices in 3 the effective sequence length is VN. This eliminates the need for tiling along the sequence length, except for very long sequences having O @(vNd) S, where S is the size of on-chip SRAM. This means that we have an optimal IO complexity
(dd)) for a single call, as opposed to
(t) the worst-case O( complexity FlashAttention. The computation of oR from a!!) as well as the Monarch matmul, can be written in similar fashion. Based on this, MonarchAttents only achieves significant speed-up up over FlashAttention for longer sequences, but also for shorter ones. Python-like code for Monar ConarchAttention is given in Figure
4 Experiments
this section, we evaluate the zero-shot performance (no additional training) MonarchAttention for converting pre-trained/fine-tneel Transformer attention layers to sub-quadratic attention in four different model/task settings. compare with previous low-rank attention methods (Katharopoulos
et al. 2020; Choromanski al. 2021 Xiong et al., 2021 Qin et al. 2022); see Appendix D. for

---

(a) Softmax
(b) XonarchAttention
(c) Nyströmformer
Figure Visual quality of generated images for zero-shot conversion of attention layers. Example images generated by with softmax (left), MonarchAttention (middle), and Nyströmformer (right).
Only the first half the attention layers of DiT are replaced.
more details on the baselines. We specifically exclude low-rank methods with learnable components (Wang al.] 2020) ZZhang al., (0224) since
focused on the zero-shot setting, well as sparsity/LSH--aat approaches (Kitaev et al.) 20200 Daras et al. 2020 Chen et al. [02 Han et al.
2024), since these do admit efficient implementations on current GPUs. We also benchmark our
Tast implementation of MonarchAttention, comparing with F FlashAttention
(Dao 2024
Image Classification with Vision Transformer. We convert all 12 attention layers, each having 12 heads with sequence length N 197 and head dimension d' = 64, of the 87M parameter ViT- B (Dosovitskiy et al., 2021 that has been pre- traitrained on ImageNet-21K (Deng et al. 2009) and fine-tuned on ImageNet- IK (Russakovsky et al., (015) for image classification. mance at different FLOP counts, we vary the number of steps T E {1, 3} for evaluate the perfor- and vary the rank for Performer and Nyströnformer; see Appendix D.2 for more details on set-up. The results are shown in the left panel of Figure Monar chAttentionn achieves significant improvement over other baselines compared to the original softmax attention, MonarchAttent loses only 5% accuracy to reduce attention FLOPs by 80%, matches the performance to reduce attention FLOPs by 50%.
Question Answering with Encoder-Only Transformer. We convert the initial 4 and final 4 layers of the 12 attention layers, each having 12 heads with sequence length N 384 and head dimen- sion d = 64, of the 125M parameter RoBERTa-B (Liu et al. 2019) that has been pre-trained on large English corpus and fine-tuned on SQuADI.1 RRajuurkar et al. 2016) for question answering. and vary the block size Appendix [D.3]
To evaluate the performance at different FLOP counts, we fix for MonarchAttenti on, and vary the rank for Performer and Nyströmformer; for more details on the set-up. The results are shown in the right panel of Figure 2 Once again, MomarchAttentice achieves significant improvement over other baselines compared to the orig- inal softmax attention, MonarchAt tention loses only 10 points in F1 score to reduce attention FLOPS by 60%, matches the performance reduce attention FLOPS by 35%. We additionally provide a per-layer ablation study in Appendix E.2 to understand the impact of converting different layers in this setting.
Summarization with Siccode--eecoder Transformer. We convert all attention layers. each having 12 heads with head dimension d = 64, in the encoder of the 139M parameter BART-B (Lewis) al. 2020) that has been pre-trained on a large English corpus and fine-tuned on BookSum-chapter (Kryscinski et al. 2022) for summarization. We only convert the encoder model and leave the decoder intact. To evaluate the benefits of sub-quadratic attention for processing longer sequences, we truncate the text to be summarized to various sequence lengths N for each method, with T E {2, 3} for MonarchAttention on depending on the sequence length; see Appendix [D.4 for more details on achieves a strictly better ROUGE score (Lin 2004) vs. FL.OPs tradeoff than even softmax attention, due to accurate and efficient processing of longer sequences. In particular, the N 8192 MonarchAttent model
set-up. The results are shown in Figure g We see that MonarchAttent
8

---

Table 1: Quantitative results for zero-shot conversion of attention layers for image generation. We report FID and sFID (using the original softmax attention model as reference) of DiT when replacing all half of the attention layers.
Layers Replaced
Method
Total Attention FLOPs (10%)
FID (L)
sFID (1)
Softmax
8.46
Nyströmformer
3.30
5.97
13.47
All
MonarchAttention
3.44
2.82
5.09
Nyströmformer
5.88
8.17
19.01
First Half
MonarchAttention
5.95
0.39
0.66
Nyströmformer
5.88
6.76
13.58
Second Half
MonarchAttention
5.95
1.98
3.36
improves on the N 2048 softmax attention model by 0.75 on ROUGE-1 I and 0.5 on ROUGE-L with slightly fewer FLOPs, while the N = 8192 Nyströmformer model with similar FLOPs does strictly worse than softmax. Although we are mainly focused on the zero-shot setting, we also investigate the post-swap fine-tuning performance and stability of MonarchAttention on this task in Appendix
Image Generation with Diffusion Transformer. We convert a subset of the 28 attention lay- ers, each having 16 heads with sequence length N 256 and head dimension d 72, of the 675M parameter DiT-XL (Peebles and Xie (023) that has been trained on ImageNet (Deng et al.) 2009). We consider replacing either all layers, the first 14 layers, or the last layers, and fix T for MonarchAttention see Appendix D.5 for more details on the set-up. Examples of generated images with each method for replacing the first 14 layers are shown in Figure 5 where MonarchAttentice the Nyströmformer ones are extremely noisy. We also quantitatively evaluate the (s)FID scores (Heusel et al., 2017) of MonarchAttentioo compared with Nyströmformer, using images gener- ated with the original softmax attention model as reference the results are reported in Table [ MonarchAttention noticeably outperforms Nyströmformer with similar FLOP count. In partieular, using MonarchAttention in the first half the DiT layers results in extremely small FID and sFID from the softmax attention model's 's images, while reducing FLOPs by nearly 30%. We provide more fine-gra -en-igri version of Table [ for MonarchAttent ion in Appendix E.4 to investigate the impact of replacing different quarters of the DiT model.
produces clear images resembling those of softmax attention, while
Node Classification with Graph Transformer. Beyond image and language tasks, we also evaluate the performance of Monar chAttentionn for converting attention in Graph Transformers graph node classification task in Appendix E.5
Benchmarking MonarchAttention. Finally, validate that the computational/IO complexity reduction achieved by MonarchAttention translates into actual speed-ups on the NVIDIA A40, a modern GPU. We implement the pseudo-code described in Section B as four separate Triton kernels and compare against the fastest available implementations of FlashAttention
either the which calls the CUDA
Triton implementation PyTorch's 's scaled_ _eweedddd___ppeeeeeeeeeeeeeeeeed product_attents backend for FlashAttention-2 Using a fixed batch size E, number of heads H, and head dimension d, we sweep the input sequence length N. generate random Q, K, V matrices, and compare the run-times of Flashhttention and Monarchattention (with vN and = 1) in Figure [6] (left). As sequence length increases, MonarchAttention consistently outperforms notably achieving up 8.2x speed-up with N 16384. To highlight gains for shorter sequences, we implement MonarchAttention as a single fully-fused Triton kemnel, with single thread block computing a single head. For fixed sequence length N 256, number of heads H, and head dimension d. we sweep the batch size E and compare the run-time of the fully-fused MonarchAttente
Fl PlashAttention
Figure 6 (right). With smaller batch sizes, we have low utilization of hardware, since we compute a single head with single thread block. However, as we increase the batch size, MonarchAttention achieves up to 1.4x speed-up over Fl FlashAttention-2
kernel against FlashAttention-2

---

Norralized Buntime ve. Sequence Length (E=1)
Nermalized Buntine vs. Batch Size - 2561
1.8
6.2
0.9
1801
2011
1596
$192
1038
2.0
35
32
01
326
356
515
1800
7808 4066 8152
Sc5seric Length
Figure 6: Run-times of MonarchAttention and Lassatteniion-222 across various se- quence lengths. Normalized runtime FlashAttention an-aninin22n- on NVIDIA A40 GPU. Left: sweep of sequence length N with E = 1, H 12. and d 64. Right: sweep of batch size E with N 256, H 12, and d 64.
slowest, 0 fastest) of MonarchAttentior and
5 Conclusion
To conclude, we discuss several limitations and future directions for this work.
Autoregressive Transformers. In this paper, we have primarily focused on zero-shot approxima- tion of attention in encoder-based Transformers, yet most generative language models of interest today are decoder-based autoregressive Transformers. One limitation MonarchAttention is that it fundamentally cannot be applied to autoregressive generation, since there is "attention matrix" to approximate instead, single queries are passed through the attention mechanism in streaming fashion, and discarded after each decode step. Instead, MonarchAttention could be used during prefill to efficiently compute keys/values for a provided prefix (e.g., long prompt or context). Moreover, MonarchA chAttenioon can be applied emerging aon-autoregressive paradigms such as diffusion language models (DLMs) to accelerate language generative models.
Accelerating Training. While we have presented MonarchAttention as a direct replacement for softmax attention with additional training, it can also accelerate pre-training fine-tuning well. cannot be applied at inference time for autoregressive Transformers,
Although MonarchAttents it can be applied training such models. However is not obvious how to efficiently incorporate the causal mask into MonarchAttention.
Beyond Monarch Softmax. Finally, although MonarchAttention specifically designed for approximating softmax attention via Monarch, this idea can be generalized to other (possibly dynamic) structured matrix classes, as well as other nonlinear mappings. Monarch, while efficient, allocates the same number of parameters to each uniformly sized block. Yet, in practice, different regions of the attention matrix may require more fine- geee-graindd approximation. On the other hand, a-entmax mappings (Peters et al. 2019), which generalize softmax, produce sparser attention matrices, thereby resulting in better approximations via block rank-one type matrices.
Acknowledgement
LB and CY were supported in part by NSF CAREER award CCF-1845076 NSF award CCF-2331590 and an Intel Early Career award. LB was also supported by the University of Michigan Crosby award. CY and AX were supported by NSF CCF 312842. PA and CL were supported in part by COGNISENSE, one of seven centers in JUMP 2.0. Semiconductor Research Corporation (SRC) program sponsored by DARPA. We thank Samet Oymak (University of Michigan) for discussion and use computational resources provided by an Amazon Research Award on Foundation Model Development.
10

---

References
Mathieu Blondel, Andre Martins, and Vlad Niculae. Learning classifiers with lenchel-young losses: Generalized entropies, margins, and algorithms. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 606-615. PMLR, 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Beidi Chen, Tri Dao. Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34: 17413-17426, 2021.
Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. International Conference on Learning Representations, 2022.
Yuejie Chi, Yue M Lu, and Yuxin Chen. Nonconvex optimization meets low-rank matrix factorization: An overview. IEEE Transactions on Signal Processing, 67(20)55239-3) -269,, 2019.
Theodore Chihara. An Introduction to Orthogonal Polynomials. Courier Corporation, 2014.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint S0KI::1904.100000 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference machine learning. pages 1517-1527. PMLR, 2019.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning, pages 4690- 4721. PMLR, 2022a.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35: 16344-16359 2022b.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficicn attention using asymmetric clustering. Advances in Neural Information Processing Systems, 33:6476-6488 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255 IEEE, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Dan Fu, Simran Arora, Jessica Grogan, Isys Johnson, Evan Sabri Eyuboglu, Armin Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm- architecture. Advances in Neural Information Processing Systems, 36777446-77663 2023.

---

Insu Han, R Jayaram, A Karbasi, V Mirrokno, D Woodruff, and A Zandieh. Hyperattention: Long- context attention near-linear time. In International Conference on Learning Representations. 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Thomas Kailath, Sun-Yuan Kung, and Martin Morf. Displacement ranks of matrices and linear equations. Journal of Mathematical Analysis and Applications, 68(2):395 407, 1979.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference for Learning Representations, 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In
International Conference on Learning Representations, 2020.
Wojciech Kryšciński, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: collection of datasets for long-form narrative summarization. In Findings of the Association for Compulational Linguistics: EMNLP 2022, pages 6536-6558, 2022.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. Bart: Denoising scquence-to-sequence pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7871-7880, 2020.
Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, et al. Datasets: A community library for natural language processing. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 175-184, 2021.
Chin- Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pages 74-81, 2004.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
Victor Pan. Structured matrices and polynomials: unified superfast algorithms. Springer Science Business Media, 2001.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF international conference on computer vision, pages 4195-4205 2023.
Hongbin Pei, Bingzhe Wei, Kevin Chen Chuan Chang, Yu Lei, and Bo Yang. Geom-gen: Geometric graph convolutional networks. In 8th International Conference on Learning Representations, ICLR 2020, 2020.
Ben Peters, Vlad Niculae, and André FT Martins. Sparse sequence-to-sequene models. In Pro- ceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1504-1519, 2019.
Zhen Qin, Weixuan Sun. Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan. Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022.
Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision. International conference on machine learning, pages 28492-28518 PMLR, 2023.
12

---

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2383-2392, 2016.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 10111411-145151 2022.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115 1-252, 2015. doi: 10.1007/s1 1263-015-0816-y;
Ilya O Tolstikhin, Neil Houlsby. Alexander Kolesmikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261- 24272, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04770 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-an natural language processing. arXiv preprint 1rXiv:1910.03771, 2019.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 14138-14148, 2021.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive lincar attentions with softmax mimicry. In International Conference on Learning Representations, 2024.
13

---

Appendix
A Equivalence of Softmax Definitions
Consider the optimization problem in 2
f(a) = aiz -a log a, s.t. a:20, Σa. =
max
From the KKT stationarity condition, have
1 1- (10)+x(1->0)
+ [(a)+
((a)
+ Σ uia = 0
2; 2-(1+ log a:) + μu; 0,
da:
where X e R. E R-V are dual variables. From complementary slackness a:pt; = and the fact that log a, is not defined for a; = 0, must have μ4 = 0, which gives
log a; = zi 2
a; /x)p) exp(A + 1).
Finally, from the constraint E; 1, we must have exp(A + 1) = E, exp(zj), which gives the form of softmax (T).
B Monarch Background
We provide example of the transpose permutation P in Section 2 Recall that applying P NxN to a vector a E corresponds to row-major reshaping x to m transposing to Rhxm then row-major flattening back to N This is equivalent to applying permutation matrix whose (i+ 1)th row is given by (z(+)+1 where
o (i) = b. mod m) + i (O...., 11..
As an illustrative example, let N 6, 3, and m 2. The action of P is given by the following steps:
3
reshape 2x3
2
3
transpoae
2
5
flatten
2
6
3
6
5
3
In matrix form, we have
0
0
0
0
0
0
1
0
0
0
0
C
Details for MonarchAttentiinn
C.1 Updates
Derivatives.
We evaluate ff with A M as a Monarch matrix. Using (5) we have
([MQQ.K) = Σ LiuRRaji 2
K Kiiv
Σ Likt Raji log( Like Rajs)
Σ PijaBaji Ruiji K Σ
Lyial Pajii log Raji Σ
Ragi Lojaa log Lzat-
14

---

The derivatives of f w.r.t. each factor are given by
Df(M;Q. Q, K) = BI.jjkt Cr.jk (1+ + log jiht) YL.sk:
(11)
0L
Q.Q) K) BR.ji MR.kj - (1 log Cn,jj,
(12)
dRgs
where
BL,jkl = EQ,t. jiv Σ (RASK
kioo, CLisk =
Σ Rkji log Riji,
Σ Rkjs:
Yhok
BRAKI = Σ K lév 3(n))
YR,Kj =
Σ
L log Lyst
CR..j
Σ Likl-
We derive updates for each factor based on maximizing / with the other factor fixed.
L update.
First, fix R Amxbbb and consider
max
K) s.t. Likl 2 0,
Egn =1
L
k
From the KKT stationarity condition, we have
i
f(M:Q,K)
Σ AER ( Lyyt + + MeL.eeLaa =
aLi
where d E Rx
Mr E bxxmx are dual variables. Along with 11 we have
BL.jH! - CL,jk - (1 log L - ALJI + MLight =
Now, from complementary slackness fLIRht = and the fact that log Lyæt is not defined for
L = 0, we must have ML.jki O. Moreover, since E A mxxbbb we have Yr.jk = Altogether, we have
exp(BL.jkt CL.jk)
log Liks BL.skt CL,jk AL
Liu
exp(AL.f1* +1)
Finally, from the constraint Σx L,Kl = 1, we must have exp (L + 1) E. exp(BLLjRt - -L. which gives the final closed form update:
L softmax; (ZL), ZL.skt BL.jHt CL.fkr
where softmax, is applied along the k index dimension.
R update. Similarly, we fix L A 6xmxm and consider
max f(M:Q.K) s.t. Raji 0, Σ Rkji 1.
From the KKT stationarity condition, have
a J(M:Q,K) + Σ ( ( Rkjs
Ma Raji
RRgji
where AR E jpmxb
uR E Ruuxbbb are dual variables. Along with 12). we have
Bn,iii Yn,k - (1 + log Reji)en.ss - ARAI + PR,kjs = 0.
As before, from complementary slackness PR.kji Rajj = 0 and the fact that log Rays is not defined for Raji 0, we must have fAikje Altogether, we have
log Rays BRAKji
7n,jj
AR,kj -1 Rajs eXp(TR.M CR.Kj + AR.jj/CR.RR.. CR, + 1)
exp(BR,tji/Cn..
CR.Kj
Finally, from the constraint E; Reys = 1, we must have sp(7/AA5//R.A + An,ss/Cn.ts + 1)
E, exp((Bnkgji/cmm) Cn,kj) which gives the final closed form update:
R softmax, (ZR), ZRA BR,
where softmax, is applied along the index dimension.
15

---

C.2 Naïve Algorithm
We provide pseudo-code for the naive version of MonarchAttention in Figure This direct implementation of alternating maximization for finding the Monarch factors, where we highlight the choices to (1) initialize L to be block identity, and (2) update R before include this here for completeness, that readers can directly compare this code with the algorithmic steps. The implementation in Figure reorganization of Figure 7 where L, are not matcrialized in HBM, resulting in FlashAttention-lkk kernels.
each iteration. We
Q: array of size (M, d)
2
K: array of size (N, d)
V: array of size (N, d)
T: number steps
5
def monarch_ard monarch_attention attention (Q, K, V, T):
L stack(b (m)])
Qb = einshape (1j)v~>jlv", q, j=b)
Kb einshape (ki)v->kiv", i=b)
10
Alternating mazimization for L, R for in rTnge(T):
update
aR einaum ("jkl,jlv->kjr"
Qb)
15
bR einsum ("kjv,kiv->Xj] aR, Kb)
16
cR = einaum ("jk1->kj", L)
sof omft (bR cR[: i, Nonel, axis=2)
update
20
aL einsum
R, Kb)
bL = einsum ("jkv,jlv->jkl", aL, Qb)
cL einsum ("kji->jk", 1og(R))
sof tmax (bL cL[:, i, None], axis=1)
Nonarch multiply
Vb = einshape (" (ki))->kiv", V, i=b)
einsum (" (ej R, Vb)
einsum ("kmmm"jj,,vvlll jkv.jk->>jjv" L,
29
einshape "1jv~>(1jvv* 2)
return 0
Figure Naive implemention MonarchAttention
C.3 Padding and Masking
In practice, the sequence length N may not be divisible by the desired block size b. In such cases, we round the number of blocks m to m'' = [N/b], and set the new sequence length N' m'b. post-padding Q,K. K. V to have N' rows. However, columns of the padded Monarch attention matrix M N''N'' x' are zero, since these correspond to padded rows of V. This is also an issue when batched sequences of different lengths are padded to a maximum length avoid dynamic resizing.
need to take special care that the final N' N
From is clear that set all columns of M beyond the N-th column zero, it is sufficient to set Rej to zero whenever b(k 1) + N. Thus, we simply form the mask E 1gmm xbxb given by
0
- 1)-1iiisS
Wkji
-o
otherwise,
which we then add Zn before softmax in 7
We can also pre-pad the sequence, which would be
change the above condition to b(k + 2 N' N.
16

---

D Experimental Details
D.1 Baselines
We describe the baselines used in Section
Linear-attentil (Katharopoulos et al. 2020
approximates exp(q )(q) o(k)
where Φ R2 - Ie' is a kernel feature map, resulting in rank approximation to softmax attention:
(Q))((K)
softmax QK
<( 0 K)
1x1
Katharopoulos et al. (2020) propose the map <(x) = + clu(x) with r = where clu is
the exponential linear unit applied element-wise.
perf ormer (Choromanski et al. 2021 a linear attention method using the fact that
|| w
||:]1 )
exp(q' k) = RE (o..a
exp
ToT
q
exp w Tu
2
2
to construct random kernel feature map
|
Φ(x)
exp
[exp(w]<))
exp(w,, xx)]
Vr
sdd (O La).
where Wi,..., w
cosf ormer (Qin et al.
2022) is linear attention method utilizing position-dependee kernel
feature maps of the form
(x) = sin
3d 2N
relu(a;)
cos
2N
relu(x;)) E [N],
which produces a rank 2d approximation.
nystr ormer (iiong al. 2021) computes landmark Q K E gr from Q, K by averaging N/r consecutive spans rows, which are used to approximate softmax attention via the quadrature method:
F = softmax QKT). )) B = softmax (QKT), = softmax QK
voftmax(QK]) FA+B.
where A+ denotes the pseudoinverse of Á, producing rank approximation.
D.2 Image Classification with VIT
The VIT-B model fine-tuned on I is retrieved from the Hugging Face transforment
library (Wolf et al. 2019) as google/vit- seatt-bae---pcchhhe patchh6--22 224. The ImageNet-1K KK evaluation dataset is retrieved from the Hugging Face datasets library (Lhoest et al. 2021) as imagenet using the validation split. We vary the following hyperparameters:
conarch-attentioon
14 and T E (1,2,3}
performer: {16, 32, 48, 64, 80, 96}
nystromformer: E {16, 24, 32, 40}
The choice of b 14 for wwonarch-attenteeo is due to the fact that images patched in 14 14 grid. We also apply post-padding as described in Appendix C.3 for the [CLS] token, since is appended at the end of the sequence.
D.3 Question Answering with RoBERTa
The RoBERTa-B model fine-tuned on SQuAD1. retrieved from the Hugging Face transformert library as
The SQuAD1.1 evaluation dataset is retrieved from the Hugging Face datasets library as squad using the validation split. For evaluation, we truncate and pad to sequence length of 384. We vary the following hyperparameters:
17

---

Table 2: Hyperparameters used for BART summarization.
Sequence Length
Method (b, T)
Total Attention FLOPs (10")
Softmax
9.66
1024
Nyströmformer
64
1.93
conteetntton
(32, 3)
1.96
Softmax
38.7
2048
Nyströmformer
80
4.41
MonarchAttent
(32, 2)
3,93
Softmax
155.
4096
Nyströmformer
112
10.6
MonarchAAteettion (64, 2)
10.9
Softmax
619.
8192
Nyströmffrmer
160
35.0
MonarchAAttention (64, 2)
31.4
nonarch-attentiin T= 1 and {24, 48, 96, 128] performer: E {32, 64, 96, 128, 160, 192)
nystromformer: {16, 32, 48, 64)
D.4 Summarization with BART
The pre-trained BART-B model is retrieved from the Hugging Face gransformery library as beeebook/bart-bbee The BookSum BookSum-chapters training/evaluation dataset is retrieved from the Hugging Face datasets library kmfoda/bookssu using the train and validation splits respectively. BART employs learned positional embeddings up to 1024 sequence length, and since we are interested in long-sequence summarization up to 8192 tokens, we linearly interpolate the encoder positional embeddings up to 8192 tokens, before fine-tuning leave the decoder positional embeddings intact. fine-tune for 5 epochs with batch size of and learning rate of 10 using the Adam optimizer (Kingma and Ba 2014) without weight decay, with the input and summary sequences truncated and padded to 8192 and 512 tokens respectively, For evaluation, truncate the input sequence the corresponding sequence length in Figure The hyperparameters for each method across sequence lengths are shown in Table
BookSSm-chhateers we
D.5 Image Generation with DiT
The pre-trained model DiT-XL is retrieved from the Hugging Face transformers library as Facebook/DiT-XI-22- Following Peebles and Xie (2e33)))3) we generate images using 32 sam- pling steps, 2 x 2 patch size, and a classifier-free guidance scale of 1.5. We use the following hyperparameters:
monarch-attentin
= 16 and T' =
amftrrmffomer:: 32
To create the images in Figure 5 we used random seed of ( input the same 36 random Gaussian
samples into all three models. To obtain the results Table [ we again used a random seed 0, and generated 50K images from each type of model, again using the same 50K random samples across all models.
E
Additional Experimental Results
E.1
Convergence of MonarchAttention
We empirically show that MonarchAttentioon converges very quickly the ViT image classification setting described in Appendix D.2 We measure of the value of the softmax variational objective
18

---

; Q. ) as defined in (6) in the O-th layer and 5-th head for the first few T the results are shown in Table We see that by 2, the variational objective has already reached a stationary 1 closing most of the gap. This demonstrates how MonarchAttention is able
point, with even reach an accurate solution with even just one step.
Table 3: Convergence of f(M(") Q, K)
T
0
1
2
F(M 418
937
940
940
E.2 Layer Ablation for RoBERTa Question Answering
We investigate the impact of replacing different layers with conacchAtenttion in the RoBERTa question answering setting described in Appendix D.3 The results are shown in Table
Table 4: Fl score when replacing given layer of RoBERTa with MonarchAttention
Layer
0
3
5
8
10
11
FI (t)
92.7
92.5 92.6
92.0
91.4
91.9
86.3
91.1
92.7
93.2
93.0
92.3
The initial and final layers are more amenable replacement compared to the middle layers.
E.3
Post-Swap Tine-Tuning for BART Summarization
We demonstrate the improved performance of summarization rost-replacement. We follow the setting and hyperparameters of Appendix but fine-tune with MonarchAttention or Nyströmformer instead of the standard softmax attention for sequence length N = 8192. The results are shown in Table 5 where we also copy the zero-shot results from the main body for reference.
when fine-tuning on BART
Table 5: Post-swap fine-tuning vs zero-shot performance on BART summarization.
Type
Method
Total Attention FLOPs (10")
ROUGE-1 (t)
Softmax
618
35.57
MonarchAttentioon
31
35.38
Fine-tuned
Nyströmformer
35
33.94
MonarchAttent.
31
34.60
sero
Nyströmformer
35
32.86
We see that MonarchAttention stable under fine-tuning, and is indeed strong approxima- tion even beyond the zero-shot setting. In particular, almost fully close the gap with to the original softmax model, while benefiting from more efficient fine-tuning with roughly 5% of the attention FLOPs.
E.4 Layer Ablation for DiT Image Generation
We investigate the impact of replacing different quarters with MonarchAttention in the DiT image generation setting described in Appendix D.5 The results shown in Table We that replacing earlier layers is more favorable than later layers, with the most degradation resulting from replacing the final layers of the DIT.
E.5 Graph Node Classification with GraphGPS
We evaluate the efficacy of MonarchAttention beyond language and image modalities by consid- ering zero-shot conversion of attention in the GPS graph Transformer (GPS-GT) (Rampäšek et al.
19

---

and vary Table 6: (s)FID scores when replacing different quarters DiT with MonarchAttention
Quarter
First
Second
Third
Fourth
FID (L)
0.07
0.21
0.34
1.35
sFID (4)
0.14
0.39
0.52
2.23
2022) on the Actor graph dataset (Pei et al. 2020) for node classification. particular, we convert the attention layer in single-layer 170K parameter GPS-GF GT with 2 heads, sequence length N = 7600, and head dimension d 48. We train for 1500 steps with learning rate of 5 x 10-4 using the Adam optimizer, and also evaluate on the train splif] To evaluate the performance at different FLOOP counts, we vary the number steps E {1, 2, while fixing b = 128 MonarchAttent the rank {128, 192, 256, 320} for Nyströmformer. The results are shown in Figure g
GPS Actor
100.0
97.5
95.0
Train Accuracy
92.5
monarch-attention nystronformer
90.0
so tmax
87.5
85.0
82.5
80.0
16"
16
Total Attention FLOPs
Figure Zero-shot conversion of attention layers for graph node classification. We vary hyperparameters to evaluate training accuracy vs. total attention FLOPs across all layers for GPS on the Actor dataset.
MonarchAttention achieves favorable quality-efficiene tradeoff, almost matching the softmax model quality with an order magnitude less FLOPs. On the other hand, Nyströmformer achieves roughly 18% less accuracy with 2-3x more FLOPs than Monar MonarchAttention
"We evaluate on the train split due to low performance on the validation/test sets due to limited data/training from scratch. Since we are primarily interested in the post-train zero-shot approximation capabilities of MonarchAttention this does not introduce bias in the evaluation.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We support the claims made made in the abstract and introduction through extensive experiments in Section 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss several limitations of the proposed algorithm in Section whose addressing we pose as future work.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
21

---

Answer: [NA]
Justification:
No theoretical results in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide exact Python-like pseudocode for the proposed algorithm, along with precise set-ups for each experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
22

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will release code for all experiments after the review process. We also specify publicly-available model checkpoints used for each experiment.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/guides/ wttaeeee//ooeeSummaaiiiiiiiiioie
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Coaee//ddeeeebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide full details for hyperparameters and settings used in Section 4 and the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: There is little randomness in the experiments carried out in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
23

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specify the GPU used for benchmarking run-times in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:///eeeeeeeeeeci/eeee///e/ee
Answer:
[Yes]
Justification:
All research conducted conforms to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We believe this work be foundational and not tied to any specific application or societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any new data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite and provide publicly available links to all data and models used.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
25

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification:
No new assets are developed.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
26

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy buctps://neure wwwparipppsc.coooeeeeeeeee cc/Conferencees///////e for what should or should be described.
27


Reviews:
[{"review_id": "5O0xUafQxT", "reviewer": "Reviewer_DpNT", "summary": "This paper proposes MonarchAttention, which approximates each attention layer by fitting a structured Monarch matrix, enabling sub-quadratic attention without any model retraining. In zero-shot swaps on ViT, RoBERTa, BART, and DiT-XL, it cuts attention FLOPs by a large margin while incurring minor accuracy drops.", "questions": "**Questions**:\n\n1. Per-Layer Error Analysis? The paper reports only end-to-end task metrics when swapping *all* layers or a fixed prefix. Some layers (e.g., early vs. late) may tolerate approximation much better. I was wondering if a layer-based ablation might offer more insights into where MonarchAttention is safe and where it breaks. \n2. How do you select the block dimensions m and b for arbitrary sequence lengths? Have you evaluated the impact of padding or uneven block sizes on approximation quality at sequence ends?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**: \n\nThe paper tests MonarchAttention on pretrained models without any fine-tuning or weight updates and achieves sub-quadratic complexity. Experiments are done across various model architectures and benchmarks.\n\n**Weaknesses**:\n\n1. The paper doesn’t show any experiments showing that MonarchAttention remains stable after fine-tuning, which might be critical for any downstream applications. I wonder if the authors have attempted any post-swap fine-tuning with MonarchAttention in place, and if so, how stable is training compared to the original softmax?\n2. The same rank-one block approximation is applied uniformly to every row of the attention matrix. In some tasks, however, some tokens require more fine-grained attention while others need only broad context. The rigid Monarch factorization cannot adaptively allocate more capacity to certain rows, potentially harming accuracy on critical positions.", "was_revised": true, "final_justification": "Post-rebuttal final justification: the reviewer thanks the author for the in-depth rebuttal and addressing some of the reviewer's concerns. I will maintain the original score.", "rebuttal": "We thank the reviewer for the positive feedback and thoughtful remarks. We address each of the reviewer’s comments and questions in the following.\n\n*The paper doesn’t show any experiments showing that MonarchAttention remains stable after fine-tuning, which might be critical for any downstream applications. I wonder if the authors have attempted any post-swap fine-tuning with MonarchAttention in place, and if so, how stable is training compared to the original softmax?*\n\nAlthough our paper highlights the zero-shot capabilities of MonarchAttention, the reviewer brings up an interesting point regarding training with MonarchAttention. We repeat the fine-tuning procedure for obtaining the summarization model with N=8192 reported in Figure 3, except we replace the softmax attention layers with MonarchAttention prior to fine-tuning. We show the ROUGE-1 results compared with zero-shot replacement:\n\n|     Softmax     | MA (fine-tuned) | MA (zero-shot) | Nyst. (fine-tuned) | Nyst. (zero-shot) |\n|-----------------|-----------------|----------------|--------------------|-------------------|\n|      35.57      |      35.38      |      33.94     |       34.60        |       32.86       |\n\nIn summary, training with MonarchAttention is stable and in fact improves on zero-shot substantially. We will include these results in the camera ready version.\n\n*The same rank-one block approximation is applied uniformly to every row of the attention matrix. In some tasks, however, some tokens require more fine-grained attention while others need only broad context. The rigid Monarch factorization cannot adaptively allocate more capacity to certain rows, potentially harming accuracy on critical positions.*\n\nWe certainly agree that some tokens require more degrees of freedom compared to others, and we believe that our approach can be generalized to structured matrices with variably-sized rank-one blocks – we will include a discussion of this in the camera ready version. However, an adaptive method that identifies important tokens would likely introduce additional overhead (as with LSH-based approaches), whereas our method is a generic yet fast approximation that works well in practice. We note that Monarch is an expressive class of matrices that can already represent many interesting structured matrices (Proposition 3.2 in [1]). \n\n[1] Dao et al. Monarch: Expressive Structured Matrices for Efficient and Accurate Training.\n\n*Per-Layer Error Analysis? The paper reports only end-to-end task metrics when swapping all layers or a fixed prefix. Some layers (e.g., early vs. late) may tolerate approximation much better. I was wondering if a layer-based ablation might offer more insights into where MonarchAttention is safe and where it breaks.*\n\nWe appreciate the reviewer’s suggestion on this – in addition to the results shown in Table 1, we conduct some more experiments here to gauge the sensitivity to choice of layer for replacing attention. For the RoBERTa SQuAD experiment, we consider replacing different layers with MonarchAttention (keeping the rest intact):\n\n| Layer | 0    | 1    | 2    | 3    | 4    | 5    | 6    | 7    | 8    | 9    | 10   | 11   |\n|-------|------|------|------|------|------|------|------|------|------|------|------|------|\n| F1    | 92.7 | 92.5 | 92.6 | 92.0 | 91.4 | 91.9 | 86.3 | 91.1 | 92.7 | 93.2 | 93.0 | 92.3 |\n\n\nAccording to this, the initial and final layers are more amenable to replacement compared to the middle layers for BART. On the other hand, we also provide a more fine-grained version of Table 1 by replacing quarters (rather than halves) of the DiT model, where replacing earlier layers is more favorable:\n\n| Quarter |  FID   |  sFID  |  FLOPs (10^9) |\n|---------|--------|--------|---------------|\n| First   |  0.07  |  0.14  |      7.20     |\n| Second  |  0.21  |  0.39  |      7.20     |\n| Third   |  0.34  |  0.52  |      7.20     |\n| Fourth  |  1.35  |  2.23  |      7.20     |\n\nWe will include these additional experiments in the camera ready version.\n\n*How do you select the block dimensions m and b for arbitrary sequence lengths? Have you evaluated the impact of padding or uneven block sizes on approximation quality at sequence ends?*\n\nFirst, we note that varying block size b yields a tradeoff between FLOPs and model quality (b=1 or N gives back dense representation, m=b=sqrt(N) is least compute) – this is demonstrated in the RoBERTa SQuAD experiment shown in the right panel of Figure 2. In this sense, we simply choose b (m is just the ceiling of N/b) according to our computational budget. In the case of ViT with N=197=14*14+1, there is an obvious choice of b=14 for the number of patches along the width of the image, and pre-padding since the cls token is at the beginning of the sequence – in fact, if we use post-padding as opposed to pre-padding, accuracy drops from 93.26% to 92.59% for T=2. We will include discussion for these hyperparameter choices in the camera ready version."}, {"review_id": "iLukBISo57", "reviewer": "Reviewer_jzk9", "summary": "Attention requires quadratic complexity, sub-quadratic approximation via Monarch matrices (structured) may improve complexity. Transferable, with no additional training and hardware-efficient make it a potentially good way of reducing attention overhead.", "questions": "I think this is a good paper, and it just ‘works’, which is great. However, I would really love to see the ‘best possible’ sub-quadratic variant of _any one_ of these models (even monarch-attention!). This is optional, but can help. I am very option to raise my rating to an accept once I hear the author's views on this.", "limitations": "None", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- Flashattn style kernels on quadratic-size blocks is very useful and can have big impact on performance. \n- Zero-shot implies faster adoption. \n- Very well written paper, easy to follow.\nWeaknesses:\n- Is it possible to have Figure 2, but with train-from-scratch variants as well on performer, cosformer, linear attention etc? Even a single variant (just to upper bound sub-quadratic attention) would be very insightful, not to serve as a ‘negative’ signal but to ~upper bound the performance of sub-quadratic attention.", "was_revised": true, "final_justification": "I agree with the rebuttal by authors, and think this paper should definitely be accepted + highlighted.", "rebuttal": "We thank the reviewer for the positive reception. \n\nThe suggestion to upper bound performance of sub-quadratic attention is very interesting, but unfortunately training any of the models considered in the paper from scratch is currently infeasible due to time and resource constraints. As an alternative, we repeat the BART summarization experiment but fine-tuning *with* Nystromformer and MonarchAttention (instead of zero-shot replacement), which we believe is a reasonable proxy for the “best possible” sub-quadratic model:\n\n|     Softmax     | MA (fine-tuned) | MA (zero-shot) | Nyst. (fine-tuned) | Nyst. (zero-shot) |\n|-----------------|-----------------|----------------|--------------------|-------------------|\n|      35.57      |      35.38      |      33.94     |       34.60        |       32.86       |\n\nAccording to this, fine-tuned MonarchAttention is very close to the original softmax model and may act as an upper bound on sub-quadratic attention performance. We will include these results in the camera ready version. Overall, this paper focuses on zero-shot approximation, but we believe that MonarchAttention is a strong alternative to attention when training from scratch, to be explored in future work."}, {"review_id": "2QylnxHW5C", "reviewer": "Reviewer_nTcP", "summary": "The authors propose to approximate the softmax attention matrix with a Monarch matrix in order to decrease the computational cost from $\\Theta(N^2 d)$ to $\\Theta(N \\sqrt{N} d)$, where $N$ is the sequence length and $d$ the embedding dimensionality. To that end, they cast attention as an optimization problem over candidate \"attention\" matrices, and derive an alternating maximization scheme that monotonously improves a Monarch approximation to the attention matrix in $\\Theta(N \\sqrt{N} d)$ without ever materializing the attention matrix. Empirically, the authors demonstrate runtime improvements over FlashAttention and much better attention approximation than linear transformers in image generation and language encoding. They demonstrate further that MonarchAttention achieves strictly better performance / FLOPs in summarization than truncated softmax.", "questions": "For well-definedness of the entropy terms, L and R must be constrained to\nbe positive. I don’t see this mentioned anywhere in Section 3 before subsection “Alternating Maximization with Constraints”.\n\nThe authors state: “When L is fixed, the objective is concave in R”. This seems to be a typo, since the negative entropy is convex.\n\nThe chosen $T$ is not always indicated in the experimental section, which initially gives the impression of it being quite large. I would suggest to explicitly state this for every setup to assure the reader no information is being hidden.\n\nA discussion of how Monarch attention integrates with GQA and MLA could be a noteworthy addition.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- The paper is very clearly written.\n- The related work is well-explored and contextualizes the author’s contribution. \n- The experimental evaluation is extensive and convincing. \n- The paper may have wide applicability, in particular for graph transformers and large reasoning models. \n\nWeaknesses:\n- The applications most impacted by a $\\sqrt{N}$ acceleration of attention (e.g., graph attention and very long-context settings) are only explored to a limited extent. \n- MonarchAttention only achieves improvements over vanilla attention if the number of optimization steps $T$ is in $o(\\sqrt{N})$. Although in the experimental evaluations $T$ only reaches the modest size of $3$, the paper lacks insights on the convergence rate of the optimization problem as a function of $T$.", "was_revised": true, "final_justification": "The authors reaffirmed the strength of the paper and in particular elaborated on the empirical convergence rate of their optimization scheme. The paper is a strong contribution to the scientific community, I thus still recommend acceptance.", "rebuttal": "We thank the reviewer for the positive reception and insightful points. We address each of the reviewer’s comments and questions in the following.\n\n*The applications most impacted by a sqrt(N) acceleration of attention (e.g., graph attention and very long-context settings) are only explored to a limited extent.*\n\nBeyond the summarization task shown in Figure 3, we agree with the reviewer’s point that the acceleration brought by MonarchAttention would be most obvious in very long sequence problems. Due to time and resource constraints, we cannot include experimental results in such settings for the rebuttal, but we plan to conduct experiments on a task from the Open Graph Benchmark to be included in the camera ready version.\n\nFor future work, we also hope to team with researchers working on important long-context applications (who may hopefully also have sufficient computational resources) in order to test MonarchAttention on very large scale problems.\n\n*MonarchAttention only achieves improvements over vanilla attention if the number of optimization steps is in o(sqrt(T)). Although in the experimental evaluations T only reaches the modest size of 3, the paper lacks insights on the convergence rate of the optimization problem as a function of T.*\n\nWe empirically find that alternating maximization usually converges in a couple of steps – to quantify this, we can measure the softmax variational objective f(M) in equation (6) after each step of alternating maximization in MonarchAttention for a particular ViT layer:\n\n| T     | 0    | 1    | 2    | 3    | \n|-------|------|------|------|------|\n| f(M)  | 418  | 937  | 940  | 940  |\n\n\nWith T=2 we have already reached a stationary point in the objective. We will include these results in the camera ready version to emphasize the rapid convergence speed. \n\n*For well-definedness of the entropy terms, L and R must be constrained to be positive. I don’t see this mentioned anywhere in Section 3 before subsection “Alternating Maximization with Constraints”.*\n\nThe reviewer is correct that we constrain L and R to be positive for the entropy term. We intended to delay discussion of this to the “Alternating Maximization with Constraints” subsection to solely focus the first half of Section 3 on computing the objective in equation (6). In the camera ready version, we will explicitly point out that we do not worry about the constraint on L, R until later in the section.\n\n*The authors state: “When L is fixed, the objective is concave in R”. This seems to be a typo, since the negative entropy is convex.*\n\nSince we are maximizing entropy (not negative entropy) in equation (6), the objective is indeed concave.\n\n*The chosen T is not always indicated in the experimental section, which initially gives the impression of it being quite large. I would suggest to explicitly state this for every setup to assure the reader no information is being hidden.*\n\nWe thank the reviewer for this valuable suggestion – for the camera ready version, we will emphasize that T is very small throughout the paper, and provide explicit values in the main paper rather than the appendix.\n\n*A discussion of how Monarch attention integrates with GQA and MLA could be a noteworthy addition.*\n\nWe believe that MonarchAttention should integrate with GQA/MLA quite seamlessly. From the modeling side, MonarchAttention is largely orthogonal to GQA/MLA, since they operate on the KV pairs prior to the attention operation. On the other hand, since the intermediate steps of MonarchAttention are FlashAttention-like calls, any optimizations of FlashAttention afforded by GQA/MLA would directly transfer to MonarchAttention."}]

Author rebuttal:

Meta-review: **Paper summary**

To speed up the inference time for transformer models authors propose to approximate attention (quadratic in sequence length) with Monarch matrices introducing MonarchAttention which is sub-quadratic in sequence length. Authors propose this approximation in a zero-shot manner: model is used as is without any fine-tuning and approximation is done via variational form of softmax within 2-3 optimization steps (empirical results) thus preserving sub-quadratic complexity. The paper also provides and discusses the hardware-efficient with the proper memory/IO implementation. As a result authors get 1.4x speed up compared to FlashAttention2.0, better approximation than linear attention for several tasks (4 different models / tasks in vision and language were considered).

**Arguments taken for decision making**

*Strengths*
- **Reviewers jzk9, nTcP** pointed out that the paper is clearly and well written, easy to follow
- **All reviewers** pointed out on a wide applicability and potential (e.g. graph transformers and large reasoning models) as the method is zero-shot without any fine-tuning and reduces compute from quadratic to sub-quadratic complexity
- **Reviewer DpNT, nTcP** pointed out extensive empirical results (including across various model architectures and benchmarks)

*Weaknesses*
- Resolved
	- **Reviewer jzk9** pointed out absence of upper bound performance of sub-quadratic attention. During rebuttal due to time constraints authors provided a variant of upper bound by fine-tuning a model with MonarchAttention after softmax zero-shot approximation and the reviewer was satisfied.
    	- I also agree with the authors points that work primarily focused on zero-shot, so this is not a critical issue for paper acceptance in my view.
	- **Reviewer nTcP** pointed out that the paper lacks insights on the convergence rate of the optimization problem for softmax approximation and this could be in the end $sqrt(N)$ steps making it again quadratic. During rebuttal authors provided empirical justification showing that in 2-3 steps only optimization saturates and reaches optimal point. The reviewer is satisfied with empirical support though theoretical guarantees will be great to have.
	- **Reviewer DpNT** asked for per layer analysis and some clarifications which authors provided during rebuttal and simply can be added into appendix.
	- **Reviewer DpNT** raised issue that some positions may need fine-grained attention and approximation can hurt. After rebuttal and discussion the reviewer agreed that there are trade offs between quality and practicality.
- Partially resolved / not resolved
	- **Reviewer nTcP** pointed out on absence of empirical results in the most impactful areas e.g., graph attention and very long-context settings. During rebuttal authors mentioned that they are working on Open Graph Benchmark and will provide results in camera ready due to time constraints. Authors also mentioned that ideal test will be very long sequence problems.
    	- I agree that these tasks could strengthen the paper even more  and speed up adoption, but authors listed them as limitations in the paper, and I believe results presented are strong to encourage the community to try broader experiments with MonarchAttention.
	- **Reviewer DpNT** pointed out that while the method works in zero-shot, it is unclear if it will be stable in fine-tuning regime as for downstream tasks we often need to fine-tune models. During rebuttal authors provided results of first approximating with MonarchAttention and then fine-tuning that model on summarization task showing that it is stable and given close to softmax attention results. However the reviewer agreed only partially with results as the training had the same loss / data as pre-training.
    	- I think the solution could be to just fine-tune the model with softmax attention for any new downstream task and new objective and after that do zero-shot approximation with MonarchAttention. For this angle as zero-shot inference speeding up, I think, the paper has strong empirical justification. Anything related to training with MonarchAttention speeding up the training itself, I believe, is worth separate paper and was out of authors' focus.

**Recommendation summary**

All reviewers (with different degrees) support the paper acceptance after discussion acknowledging novelty, well written text, extensive empirical justification and potential high impact of the work. Reviewers also asked to extend results to training from scratch and fine-tuning for downstream tasks: overall authors listed these as limitations and I agree that the work potentially can grow into bigger applications widespreading across different domains and tasks, pre-training and post-training stages, and all pieces hard to cover and include in one paper. Current results authors provided I think are enough per NeurIPS standards showing applicability and speed up for inference only and I am looking forward to seeing bigger impacts on the community opening space for a variety of works to make the paradigm shift (applying MonarchAttention to speed up training). Given all this, **I recommend strong acceptance of the paper.**

**Recommendation to authors**

Dear authors, please include important pieces and clarifications from the discussion into the final version of the paper as you promised. If possible, it will be great to have Open Graph Benchmark added as requested by Reviewer nTcP (this can speed up adoption), proper fine-tuning results as requested by Reviewer DpNT (this can speed up adoption), and proper upper bound as requested by Reviewer jzk9 (this can motivate a lot for the work extension showing potential even for training from scratch). But the later is optional to me as your current results already show high potential impact for inference and leaving training applicability as a future work.
Decision: Accept (spotlight)