Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Authors: Yilun Kuang
Noah Amsel
Sanae Lotfi
Shikai Qiu
Andres Potapczynski
Andrew Gordon Wilson
Venue: ICML 2025 poster (2025)
Keywords: Attention
Inductive Biases
Low Rank Bottleneck
Locality Bias
Hierarchical Matrices
Structured Matrices
Language Modeling
In-Context Regression
Primary area: deep_learning->attention_mechanisms
TL;DR: We use specialized structured matrices instead of low rank ones to compute the attention scores. This helps improve two shortcomings of standard attention for certain tasks.

Abstract: The core component of attention is the scoring function, which transforms the inputs into low-dimensional queries and keys and takes the dot product of each pair. While the low-dimensional projection improves efficiency, it causes information loss for certain tasks that have intrinsically high-dimensional inputs. Additionally, attention uses the same scoring function for all input pairs, without imposing a distance-dependent compute bias for neighboring tokens in the sequence. In this work, we address these shortcomings by proposing new scoring functions based on computationally efficient structured matrices with high ranks, including Block Tensor-Train (BTT) and Multi-Level Low Rank (MLR) matrices. On in-context regression tasks with high-dimensional inputs, our proposed scoring functions outperform standard attention for any fixed compute budget. On language modeling, a task that exhibits locality patterns, our MLR-based attention method achieves improved scaling laws compared to both standard attention and variants of sliding window attention.
Additionally, we show that both BTT and MLR fall under a broader family of efficient structured matrices capable of encoding either full-rank or distance-dependent compute biases, thereby addressing significant shortcomings of standard attention.

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Yilun Kuang! Noah Amsel Sanae Lotfi
Shikai Qiu
Andres Potapczynski! Andrew Gordon Wilson'
Abstract
different architectures and domains, it is common to think of attention general purpose tool (Bommasani et al., 2022). However, attention endows models with a certain set of inductive biases that suit some tasks better than others (Lavie et al., 2024). Moreover, attention is computationally expensive in both runtime and memory. A large ongoing research effort aims devise more efficient alternatives to standard attention, especially for long sequences and big models (Katharopoulos et al., 2020; Beltagy et al., 2020; Gu & Dao, 2024; Dao & Gu, 2024; Yuan et al., 2025). Archi- tectures like linear attention based Transformers come with their own sets of inductive biases that can be too limiting for real data and often reduce accuracy (Arora et al., 2024; Jelassi al., 2024). In this work, we show how structured matrices be used to customize attention, changing its in- ductive biases suit particular tasks and thereby improving efficiency.
The core component of attention is the scoring function, which transforms the inputs into low- dimensional queries and keys and takes the dot product of each pair. While the low-dimensional projection improves efficiency, it causes infor- mation loss for certain tasks that have intrinsi- cally high-dimensional inputs. Additionally, at- tention uses the same scoring function for all in- put pairs, without imposing distance-dependent compute bias for neighboring tokens in the quence. In this work, we address these shortcom- ings by proposing new scoring functions based on computationally efficient structured matrices with high ranks, including Block Tensor- Train (BTT) and Multi-Level Low Rank (MLR) matrices. On in-context regression tasks with high-ding sehh-imenniooal inputs, our proposed scoring functions outperform standard attention for any fixed compute budget. On language modeling, task that exhibits lo- cality patterns, our MLR-based attention method achieves improved scaling laws compared to both standard attention and variants of sliding window attention. Additionally, we show that both BTT and MLR fall under broader family of efficient structured matrices capable of encoding either full-rank or distance-dependent compute biases, thereby addressing significant shortcomings of standard attention.
We focus on the attention scoring function, which deter- mines how much each token attends to every other token by taking a dot product of their corresponding query and key vectors. We identify two properties of the scoring function that are suboptimal in certain settings. First, it can suffer from low-rank bottleneck (Amsel et al., 2024). Since the head dimension, which defines the size of queries and keys, is significantly smaller than the embedding dimension, transforming inputs into queries and keys can cause infor- mation about each token to be lost. This bottleneck can limit the power of the attention layer. In fact, there are simple functions that are easy to express with attention if the head dimension is large enough, but provably hard to express otherwise (Amsel et al., 2024).
1. Introduction
Second, the form of the scoring function lacks an inductive bias for encoding distance-dependent compute. standard attention, every pair of tokens is scored using the same function and the same weights, and each token can atttend to the entire context. This feature makes attention power- ful, but for long contexts, it is also expensive. Real-world datasets often exhibit some amount of locality, meaning that the semantics of a token depends most strongly on nearby tokens. Previous efforts have been made to make attention more efficient by promoting locality while still preserving some global communication. For instance, the Longformer architecture combines sliding window attention with few
The attention mechanism (Bahdanau et al., 2016; Vaswani al., 2017) crucial much of contemporary deep leamn ing, providing a powerful and flexible way to process quences. Because of its widespread application to so many
'New York University. Correspondence to: Yilun Kuang <yilum.kuang@@ www.euuuuuuuuuuueo Andrew Gordon Wilson
<andrewgw @ cims.m hylns.yy..ddu>..
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Low Rank (Standard)
Bilinear MLR
Bilinear BTT
R4
La
Rai
Penmutation
Ponnutation
+
(a) Bilinear Forms with Structured Matrices
Standard Attention
MLR Attention
Ka.
Bixx
Q2
Osi
Ng,1
O..1
022
02
Global Interactions
Global Interactions
Local Interactions
Local Interactions
(b) Encoding Distance- Deaennee-D Compute Bias with Multi-Leevel Low Rank (MLR) Attention
Figure 1. Overview of two ways to customize the inductive biases of softmax attention with structured matrices. (a) Standard attention computes dot product between a query and a key via a low-rank bilinear transformation. this work, we replace the low rank product with other structured matrices such as Multi-Level Low Rank (MLR) and Block Tensor Train (BTT) introduced in Section 3.2. (b) Standard attention captures pair- wise token interactions the global context without imposing a Distance-dependent compute bias. We introduce inductive bias for more ficxible interactions between nearby tokens by changing the attention matrix from low rank to Multi-Level Low Rank (MLR), as in Section 3.4.
pair-wise
tokens of standard attention (Beltagy et al., 2020). However, such constructions are often brittle and gain efficiency at the expense of accuracy (Zhang et al., 2024).
way that prioritizes local interactions over long-range de- pendencies. This leads to mild improvements for language modeling compared to standard and sliding window atten- tion and for time series forecasting with long horizons.
To address these limitations, we introduce rich classes of structured matrices into the attention scoring function. Se- lecting appropriate structure allows to refine the in- ductive bias introduced by attention, tailoring it to specific tasks. Prior work has demonstrated the advantages of re- placing dense linear layers with structured matrices (Dao et al., 2022; Qiu et al., 2024; Potapezynski al., 2024), but these approaches treat each linear layer separately. Atten- tion scores built from combination of multiple linear and bilinear transformations. By replacing all of these trans- formations with a single structured matrix, we expand the design space for more efficient and expressive attention mechanisms.
Our contributions are:
We introduce conceptual framework for analyzing and modifying the inductive biases of attention through the structure underlying linear and (bi-)linear transformations.
Section 3.2 and Section 4, we apply this framework eliminate the low-rank bottleneck of standard atten- tion using high-rank BTT and MLR matrices, improv- ing performance task from the literature (Garg et al., 2022).
inherently high-dimensional
We explore two applications of this approach. address the low-rank bottleneck, we incorporate Block Tensor Train (BTT) and Multi-Level Low Rank (MLR) matrices. struc- tured families that have high rank while remaining efficient in terms of parameters and computation. Using this tech- nique, we demonstrate improvements for in-context linear regression, a task where standard attention is constrained by the low-rank bottleneck. To better allocate compute and memory resources, apply MLR matrices in different
In Section 3.3, show that both BTT and MLR matrices-including Monarch (Dao al., 2022), But- terfly (Dao et al., 2020), Kronecker, and Low Rank matrices-can be united under a broader structured family which we call Multi-Level Block Tensor Con- traction (MLBTC).
In Sections 3.4 and 5, we use MLR matrices to in- troduce distance-dependent compute bias, which
2

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
slightly outperforms previous methods in language modeling and time series forecasting.
rewritten follows:
(S) (Wy. Wo.)
This work advances our understanding of attention's is induc- tive biases, exploring its structural limitations and offering principled approach to design more efficient and expressive architectures.
In this paper, we explore alternatives to the standard scoring function given in Equation (2), We now describe two of its drawbacks.
Our codes are available at the following github repository https://github.com http://ggitubboomolluuuuuunnnnn Vilunnkuagg/srruuurr--tttttttt
2.2. Low-Rank Bottleneck
Any D x D matrix defines bilinear form, so in principle
any D x matrix could be used in Equation (2) to define an attention scoring function. Some earlier forms of attention used single D x D weight matrix (Luong et al., 2015), but Vaswani al. (2017) chose use the rank-r matrix W W K, instead. Since r < D, this low-rank scoring function is more efficient than a full-rank one, but it also has weaker expressive power. For example, a natural scoring function for many tasks is s(x, x') = x x', but low-rank matrix is not capable of representing it. This phenomenon is called the low-rank bottleneck of multi-head attention.
2. Inductive Biases and Limitations of Standard Attention
In the following section, we discuss how the built-in in- ductive biases of attention can be inappropriate in certain settings. In Section 2.1, we introduce our notations and review standard multi-head attention (Vaswani et al., 2017). We identify the scoring function for sequence mixing bilinear transformation. In Section 2.2, we show that the scoring function in multi-head attention relies on low rank matrix, which creates information bottleneck for tasks with intrinsically high dimensional inputs, On in-context regres- sion tasks, we show that multi-head attention requires a large enough head dimension to achieve good regression performance. In Section 2.3, we point out that standard attention shares the same scoring function for all tokens within the context without imposing distance-dependen compute bias. Real data often exhibits locality patterns, and thus attention fails to exploit the structure of data for improved efficiency.
Recent work has shown that the low-rank bottleneck can seriously weaken srvention-based models in certain settings. Amsel et al. (2024) demonstrated that for in-context lin- ear regression, task popularized by Garg et al. (2022), one full-rank attention head can significantly outperform low-rank multi- head attention after controlling for model size. They also proved that an attention layer cannot solve the nearest-neighhor problem for points on the dinput" dimensional sphere, even approximately, unless 2 diupuu In both cases, the inputs are intrinsically high-dimensional and cannot be compressed without losing important infor- mation. In such settings, the efficiency of low-rank attention comes at the price of accuracy.
2.1. The Attention Scoring Function
Let H be the number of attention heads, D be the em- bedding dimension, and r be the head dimension, where D Hr. The input to a multi-head self-attention layer is T x D matrix X, where T is the sequence length. The output is a T x D matrix given by
In Figure 2, we demonstrate this unfavorable trade-off for in-context regression, replicating the setting from Garg et al. (2022). Our results show that, across range of model sizes, transformers cannot solve this task unless the head dimen- sion r is nearly as large as the input dimension. Once this threshold reached, transformers achieve high accuracy. In Section 3.2, we explore alternatives to the low-rank scoring function of Equation (2) that are both full-rank and efficient.
H Σ (XWo,
(1)
where a is the row-wise softmax function and Wo., W K, Wv, Wo, Dxr are weight matrices. Attention is usu- ally deseribed with reference queries, keys and values (given by XWQ.. XWK, k., and XWv,) but in this paper we take different perspective. Consider any head and define its attention scoring function s RD x B + IR as follows:
2.3. Distance-Dependent Compute Bias
An important feature of attention-based models is that they allow any pair of tokens to exchange information. stan- dard attention layer completely agnostic to the order of its inputs. In particular, the same scoring function s(-:)) is used for every pair of tokens, regardless of where they appear in the sequence.
s(x,x') x We, W X
(2)
This is simply the bilinear form given by the matrix Wo, WK. We define the score matrix S E RTxT by Sy = xy), Then the output of this head can be
Many real-world tasks exhibit locality patterns, meaning that tokens appearing near each other in a sequence are more likely to be connected. In natural language, for instance,
3

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
1.00
3. Customizing the Scoring Function with Structured Matrices
Input Dimensions
Squared Error
0.75
Ampu. 16
dupps 32
0.50
dugu 64
In this section, we show how to modify the scoring func- tion so as to change its inductive biases and alleviate the issues identified in the previous section. We first introduce our main mathematical tool, structured matrix families, in Section 3.1. We then show two ways to incorporate these structured matrices into the scoring function, addressing the issues identified in Sections 2.2 and 2.3. In addition, in Section 3.3 we identify generalization of the existing structured matrix families. Finally, we cover practical con- siderations regarding efficient implementations of structured matrices and techniques for stable feature learning in Sec- tion 3.5.
0.25
0.00
10"
10"
Rank Ratio r/ imput (Log Scale)
Figure 2. Multi-head attention canmot solve in-context regres- sion unless the head dimension is close to the input dimension duput- We train 6-layer transformers with heads and varying embedding and head dimensions perform in-context regres- sion: given the prompt X1, w Xi...., -N-1, w XN-1, XN for x, E Tupe and unknown w, predict w xN. Plot shows error at N = 2dmput. For details see Section 4.
3.1. Structured Matrix Families
Consider D x D matrices. A structured matrix is one that can be represented using fewer than D parameters. Structured matrices often admit fast algorithmss for matrix- vector multiplication, making them both computationally- and parameter-efficient For instance, low rank matrices have the form LR where L, R E Dxr and & D. They require only 2rD parameters to represent, and the corresponding bilinear form yx H x LR y can be computed with only O(rD) FLOPs. We consider several other families of structured matrices in this paper. Key prop- erties of each class are summarized in Table 1. We also present the extension to structured matrices with rectangular shapes in Appendix A.
words appearing in the same sentence or paragraph must be understood together, while words appearing hundreds of pages apart usually do not influence each other directly. There are exceptions, such as when a particular fact must be recalled across a long context. Attention does not natu- rally exploit the local patterns in the data, so cannot take advantage of this simplifying structure. For autoregressive models, the global nature of attention also requires the keys and values of all previous tokens to be saved in cache, which is expensive.
Due to these drawbacks, previous work has altered attention by imposing a sparsity pattern on the attention score matrix S. For instance, sliding window attention (SWA) masks pairs of tokens that more than some fixed distance apart (Child et al., 2019). In effect, this changes the score ma- trix to use different scoring functions for different pairs of tokens:
Block diagonal matrices contain all zeros except for blocks along the diagonal, which are each dense. We notate this as D WR. where each We is a D x 2 diagonal block and + denotes the direct sum. In a low-rank block diagonal matrix, each of these diagonal blocks is itself low rank. They have the form D L,R,,. where L3 and R4 E R 2 xr Intuitively, this family interpolates between low-rank matrices (p 1) and diagonal matrices (p D).
x Wo
i otherwise.
$35 (x):
Multi-Level Low Rank (MLR) matrices (Parshakova et al., 2023; 2024) are summations of several low-rank block di- agonal matrices with different block sizes. Formally, an L-level MLR matrix is given by
0
where T" is the window size and we drop the attention head index for ease of notation. Some models use a strided version of SWA or give a global receptive field to a few tokens (Beltagy et al., 2020). However, these sparse versions of attention are brittle and underperform standard attention (Arora et al., 2024), they are commonly combined with some form of non-sparse attention (Riviere et al., 2024; Warner et al., 2024; Behrouz et al., 2024b). In Section 3.4, we introduce a flexible and hierarchical distance dependdnt compute bias by making S a structured (but not sparse) matrix.
L P
Σ
L
(4)
where L1,: and R1, are D x ri low-rank factors. Each level represents interactions at a particular scale. MLR matrices with a range of block sizes inherit the abilities of both low-rank (pr = 1) and block diagonal matrices (pr > 1), simultancously capturing coarse- coarse-grained global interactions and fine- grained local interactions. Thus, they
Pt

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Table Properties of Structured Matrix Families. For all listed structures, the number of FLOPs needed to compute the bilinear form x My is (#pparas)). For BTT, we assume = 6 d = D. MLR and BTT both attain high rank relative to their parameter counts.
Structure
Definition
Parameters
Rank
Dense
D
D
Low Rank
LR Σ R L P
2Dr
Multi-Levvl Low Rank (MLR)
R P
2D En
E rirs
Block Tensor Train (BTT)
P V= Lg
P
R2
2D28
D
natural and efficient way to represent hierarchical structure.
visualized in Figure la:
L 2
Block Tensor Train (BTT) matrices were introduced in Qiu et al. (2024) as generalization of Monarch (Dao et al., 2022) and Butterfly (Dao et al., 2020) matrices. They are designed to be efficient, expressive, and full-rank. Qiu al. (2024) and Potapczynski al. (2024) found that they outperform other structured matrix families replacements for linear layers in neural networks. Given hyperparameters b, d and where ab cd = D, a BTT matrix has the form
SMLR (xL((x).3)))== =x]
Σ 4
L R,,x 2531
(6)
1=1 k=1
SBTT(X;, xy) =x]
Pr.
L PR
D R
Xj' (7)
k==1
k=1
As shown in Table 1, a BTT matrix with a a= b======- VD requires only O(D3/2) parameters and FLOPs, but is full rank. For the MLR version, we set the numbers of blocks be powers of two: pt 2 Setting Σ, ri be the "head dimension" r, match the efficiency of standard attention and achieve high or even full rank. In Section 4, we train transformers with these structured scoring functions on in-context regression.
e PR
+ ==1
L
R
It is also possible to interpret these scoring functions as con- structing (higher-dimensional) queries and keys per atten- tion head. See Appendix B for details. Thus our approach also fully compatible with grouped-query attention (GQA) that shares the same KV transformations across multiple query heads (Ainslie et al., 2023).
where Ly E Raxce R4 E Rddxe and PL E DDD PR E Rebexche are fixed permutation matrices. PL permutes the rows by rearranging the dimension b into b. This permutation equivalent reshaping vector into matrix Z E T transposing it to Z E R and vectorizing it to get z' E Rab, Likewise, Pa reshapes vector of dimension e b. into tensor with shape (c, h, s), swaps the first two dimensions, then flattens it back into vector of dimension b. c. s. Qiu et al. (2024) proved that when a = = c = 8 DD, BTT can express any D D matrix. For efficiency, we s 1 2.
3.3. MLBTC: Multi-Level Block Tensor Contraction
In this section, we show that both MLR and BTT are spe- cial cases of a novel family of structured matrices which we call Multi-Level Block Tensor Contraction (MLBTC), MLBTC naturally encodes either the full rank constraints or the distance-dependent compute bias.
Definition 3.1. The Multi-Level Block Tensor Contraction is defined as
3.2. Resolving the Low-Rank Bottleneck with Structured Bilinear Forms
Pí MLBTC R) = Σ PL L
R A=1
R
(8)
As described in Section 2.2. parameterizing the attention scoring function by low-rank matrix Wq W K creates information bottleneck. However, using dense D x D matrix is prohibitively expensive; evaluating the scoring function would require O(D²) operations per attention head. Instead, we propose using a structured matrix that is both high-rank and allows efficient evaluation the scoring function. particular, we use MLR and BTT matrices as
1=1 ==1
where or E R, Py and n are fixed permutation matrices, L [ Ri x EB mi m. k= Mi, n, and r/pl riPt- For square matrices, D.
assume
We show in Appendix C that both MLR and BTT are spe- cial cases of the MLBTC matrices. Since BTT generalizes
5

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Monarch (Dao et al., 2020), Butterfly (Dao et al., 2020), Kronecker and MIR generalizes Low Rank, MLBTC natu- rally contains these matrices as special cases.
of the ranks equals the head dimension: Σ, = r. In this paper, we use I S 8, and we set pt = 2 to create hierarchically nesting levels.' We notate the rank allocation as follows: T rı|r2| ... Irz. When L = 1. MLR attention reduces to standard attention. When 7. > 1, more FLOPPs are allocated to computing each local interaction than each global one, which saves compute overall. Standard attention needs T FLOPs form S, but MLR attention uses only 2 2
3.4. Encoding Distance-Dependent Compute Bias with MLR Attention
As described in Section 2.3, standard attention lacks distance-dependee compute bias, and previous work has sought to create one by imposing a sparse structure onto the score matrix S. We impose an MLR structure instead. MLR matrices inherit the advantages of both low rank and block diagonal matrices. Likewise, our construction inherits the advantages of both standard attention (global recep- tive field) and sliding window attention (distance-dependent compute bias) by organizing the score matrix into hierarchi- cally nested levels.
Another benefit of MLR attention is that it reduces the key cache size during auto-regressive generation. Since K2,1 only pertains to pairs of tokens in the first half of the se- quence, the final token will not attend to it; it only needs 22 for level attention. For level 1, we must only retain kip. Thus, the total key cache size is Σ/= pize(Kr.p)) T TE z"., which is smaller than the size(K) = O(Tr) cache size needed for standard attention. For example, using 8-level MLR with = r/8 would yield 4x savings. Our proposed MLR attention is also compatible with grouped- query attention (Ainslie et al., 2023) if we apply the above transformations to different attention heads, which can lead to further savings in the KV cache size.
We first describe our construction in terms of the scoring function. Divide the sequence into blocks. Define j'j to be 2 when tokens j and j' belong to the same block and otherwise. Then we define two-level low-rank scoring function as follows:
RL xxy RL L,R.)x
d(j,5') = d(j.3') 2
8j (j,((,, (x), =
In addition, our MLR attention is also compatible with rel- ative positional embedding methods like RoPE (Su et al., 2023). RoPE encodes positional information by rotating query and key vectors by an angle proportional to the to- ken's index in the sequence. Thus RoPE imposes a form of locality bias in the scoring function. In contrast, our MLR attention is not trying to encode relative positional information boost the attention score between neighbor- ing tokens. Instead, we change the computational cost of the scoring function based on the tokens' positions. Put differently, standard attention (with or without RoPE) uses the same query/key dimension for all pairs of tokens. We effectively use a smaller query/key dimension for tokens that are far apart. This saves FLOPs compared to standard attention, while still allowing high quality attention scores for neighboring tokens.
where L1, R1 E Pxri and L2, R2 E R Dxrs are low rank factors. The term L2R increases the rank and power of the scoring function, but only for pairs of tokens in the same block. To generalize this construction, we can further divide each block into subblocks. For pair of tokens in the same subblock, let d'j = 3. Define
diis") ELR L.R.R
Sjo (x,, x
Xy'
(9)
This definition extends to any number of subdivisions. The expression in Equation (9) is also equivalent to the matrix form presented in Equation (4). See Appendix D.
We can compute the corresponding score matrix efficiently using batched matrix multiplications. As in standard at- tention, we form query and key matrices Q = XW Q and K XW K- We divide Q into blocks to obtain the left MLR factors and divide K to obtain the right ones as shown in Figure 1b. Finally, we combine the fac- tors QL, and Kr, as in the definition of MLR matrices (Equation (4)): S = ΣL_ Di_ K This is equiv- alent to the score function defined in Equation (9) when RR..
3.5. Practical Considerations
Our attention variants are implemented using batch ma- trix multiplications. We present the most efficient tensor contraction order in Appendix E. Maximum Update Param- eterization (μIP) is a recipe to set the learning rates and initializations of each weight matrix in a neural network for stable feature learning as the width of the network grows (Yang et al., 2022). We describe how to adapt jPP to our structured attention variants, along with our use of normal-
We L
[L.] and Wx- = [R.]
In our construction, the number of levels L. the number of blocks in each level PI..... Pr. and the rank of each level ri...., T' are hyperparameters chosen such that the sequence length is always longer than max Pt and the sum
It is also possible to select the size of each block in a level dynamically according to the input. For instance, one level could divide the tokens into blocks corresponding to paragraphs, while another could divide it into documents.
6

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Sctucture
Squared Erro (log scale)
Amtention (8 Head)
Squared Error
Attention (1 Head)
10-1
Billneardi TT
0.6
BillineurM..
0.4
Structure
10--
Attenion (8 (lad)
Attetion (1 Head)
BilincarlITT
0.0
10 Non-Fmbeediing Compute (log scale)
Model Width (Log Scale)
(a) Input Dimension (digus) 128.
(b) Input Dimension (dingue) = 64.
Figure 3. Both Bilinear BTT and Bilinear MLR outperform standard attention for in-context regression tasks. (a) Bilinear BTT and 1 1-head attention are full rank, so they can learn the task using smaller model width (x-axis) than 8-head attention. Here model width refers to the embedding dimension D. All points on this figure correspond to models that was trained for the same number of steps. (b) For fixed model width (D 256), both Bilinear MLR and Bilinear BTT have smaller regression error compared attention with 8 heads when controlling for compute. Here our bilincar models use 8 heads.
ization layers, in Appendix F. We also include figure which shows stable learning rate transfer across width for our MLR attention model in Figure in Appendix G.
attention suffers from a low-rank bottleneck that prevents it from learning this task even with D = 512. 1-head attention, for which r = D, has a full-rank scoring function. It learns the task at D = 512. Bilinear BTT is also full rank, even with 8 heads, so performs well too. this case, needs only D) 256. We defer a full set of results across multiple input dimensions dapuu to Figure 11 in the appendix.
4. Experiments with the Low Rank Bottleneck
Since both Bilinear BTT and Bilinear MLR in Equations (6) and (7) increase the rank with fewer FLOPs compared to dense matrix, we test if our new scoring functions based on these structured matrices can resolve rank bottleneck on the in-context regression task based on Garg et al. (2022).
Figure 3b shows that both BilinearBTT and BilinearMI.R attention attain significantly higher accuracy when control- ling for training compute. We measure compute in FLOPs, excluding the input and output projection layers (which have the same cost across all these models). We present the full set of sweeps across different model widths D and input dimensions dmput in Figures 7 to 10 in the appendix.
Following Garg et al. (2022), we train transformers on prompts of the form X1, f(11,, X2, :x. to out- put ((xN), where x; E Ridup and fxxx) w T The (O, Id.p.) is firshly sampled for each prompt, so the model must use in-context learning to infer w and apply it to XN- We draw x; ~ (O, Jupee Ianna) and set N 2dhmuu- We train the model causally the loss N EN (f(x.) - f(x))² where f(x,) is the model prediction given the first
linear functional
5. Experiments with Distance-Dependent Compute Bias
5.1. Language Modeling
pairs as context.
We train 6-layer transformers with both standard attention and MLR attention on the OpenWeb'Ter dataset with a batch size of 4, sequence length T = 1024. head dimension
Unlike language models, which have token embedding and unembedding layers, we have an input linear layer that maps input points from dhe to embeddings in IR D and an output linear layer that maps embeddings to R. Following μuP (Yang et al., 2022), the last layer is zero-initialized so that [x)((x,) - fxx)) = at initialization.
= 64, and model width D E {256, 384, 512, 768}. The model is trained with Adam (Loshchilov & Hutter, 2019) and we tune hyperparameters based on juP (Yang et al., 2022). We use character-level tokenization, which improves In Figure we show that the structured attention variants introduced in Section 3.2 outperform standard attention this task. In Figure 3a we plot the error of Bilinear BTT attention with 8 heads and standard attention with and 8 heads across range of model widths D. Standard 8-head
"We did not optimize the implementations of our methods, so they are somewhat slower in wall-clock time than standard attention. However Figure 12 shows that our methods are still superior controlling for wall-clock time and hardware. Since they use parallelizable tensor operations, we believe better speed attainable.
7

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
14x 10
Atterminn
13x
Width 512
Wiah 394
Wiath 256
Valdeiton (Dog sole!
Model Witn - 256
2
Model Widt - 384
768
Levels
Valichation Loss (Log Scale)
1.2 10
Wiald 256
Wiath 394
10"
10
Widls 768
1014
10
1015
1.1x
Non Embeed
Non Eibbed scale)
Model Wath 912
Moder Waths 7014
Valiteeio Love (lng sole)
2x
2x 10
10
9xx
1013
9x: 10
Non: Eibbed Compute (log scale)
Nox Eimbed Compute (lng scale)
10
10"
107
Strusure
Mlon-Embed Compute (Log Scale)
Attention
SWA
Global + SWA
MLR Levels
(a) Scaling Law on OpenWebText
(b) Different Forms of Distance- Deeennene Compute
Figure 4. MLR Attention achieves lower validation loss on Open WebnWe compared to standard multi-head attention and variants of sliding window attention when controlling for compute. (a) We plot the validation loss of both MLR attention with 8 levels and standard attention against compute. We vary over 4 values of model width D E {256, 384,512, 768). ML. attention outperforms standard attention across model widths. (b) We compare MLR attention with variants of sliding window attention (SWA) that also encodes distance-dependeet compute bias. Given the training sequence length T = 1024 and a 6-layers transformer model, SWA refers to 6 layers of sliding window attention with window size T' 128. Global + SWA refers to a combination of standard attention and sliding window attention with T' = 128. The first and fourth layers of the model are standard attention and the rest are sliding window attention. We observe that across the model width D. MLR attention outperforms standard attention, SWA, and Global + SWA when controlling for compute.
our ability study scaling with model size with limited compute budget (Potapczynski et al., 2024).
additional details to Appendix J.
In Figure 4a, we compare the language modeling perfor- mance of standard attention to 8-level MIR attention across model widths D. The rank allocation is [2]]8//4444441]] for a total of r = 64. MLR attention outperforms standard attention across all model widths throughout training, and thus exhibits better scaling law.
6. Related Work
Previous work has replaced dense neural network weight matrices by structured matrices- low displace- ment rank (Thomas et al., 2018), low-rank plus sparse or diagonal (Han al., 2024; Wei et al., 2024), Monarch (Dao et al., 2022), BTT (Qiu et al., 2024), and others (Potapezynski et al., for greater efficiency, as well as for fine-tuning (Hu et al., 2021; Schanobish et al., 2024) and compression- based generalization bounds (Lotfi et al., 2024a:b).
In Figure 4b, we compare MLR attention to standard (global) attention, sliding window attention (SWA) (Child et al., 2019; Beltagy et al., 2020), and an architecture with alternating layers of global and sliding window attention. Like MLR, these architectures aim to take advantage of locality. Across model widths and training time, MLR at- tention performs best when controlling for compute.
Structured matrices have also been used in place of the attention matrix. Chen et al. (2021) approximates the atten- tion matrix by a sparse plus low-rank matrix. Hwang et al. (2024) proposes the matrix mixer framework that identifies sequence mixers like Attention, S4 (Gu et al., 2022), H3 (Fu et al., 2023), Hyena (Poli et al., 2023), and Mamba (Gu & Dao, 2024; Dao & Gu, 2024) with particular structured families. As in our work, they show that structured matrices offer systematic and principled way to explore the design space.
5.2. Time Series Forecasting
We also compare standard attention and MLR attention on time series prediction. The Electricity Transformer Temper- ature (ETT) dataset (Zhou et al., 2021) tracks fluctuations of oil temperature along with six additional power load fea- tures across time. As shown in Figure 14, MLR attention gradually outperforms standard attention oil temperature prediction accuracy as the time horizon grows. We defer
Following sliding window attention (Child et al., 2019; Belt- agy et al., 2020), many methods have been proposed to
8

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
combine local attention with some global component (Belt- agy et al., 2020; Hatamizadeh al., 2023; Arora et al., 2024; Behrouz et al., 2024a; Huang et al., 2023, $3.1). Like our MLR attention, other work has constructed hierarchically nested levels of progressively longer-range attention (Ye et al., 2019; Zhu & Soricut, 2021; Ren et al., 2021; Huang et al., 2023, $3.2). Unlike us, their goal is to make attention sublinear in the sequence length. Thus, they must aggregate tokens in each block together (e.g. using max-pool), depart- ing significantly from the form of standard attention. Our study of the low rank bottleneck in attention follows that of Bhojanapalli et al. (2020); Sanford et al. (2023) and Amsel et al. (2024), though past work has not tried modifying atten- tion to fix it. In addition, many works have tried to reduce the quadratic dependence on the sequence length of stan- dard attention by approximating the softmax computations with random features (Choromanski et al., 2022), removing softmax (Katharopoulos et al., 2020), or in general reducing KV caches in the model architecture (Ainslie et al., 2023;
fine-tuning model compression rather than pre-training. We defer further exploration of these aspects future work.
Acknowledgements
We thank Alan Amin and Hoang Phan for helpful discus- sions and anonymous reviewers for helpful feedback. This work was supported in part by NSF CAREER IIS-2145492. NSF CDS&E-MST 2134216, NSF HDDR-2118110 NSF Award 1922658, BigHat Biosciences, Capital One, and an Amazon Research Award.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Yuan et al., 2025).
References
Ainslic. J., Lee-Thorp. J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head check- points, 2023. URL https: /ttp:///aavvv/..//eeeeeeeeeae..a org/ abs/ 2305. 13245.
7. Discussion
In this paper, we use structured matrices to design vari- ants of attention with inductive biases that are beneficial for certain tasks. Bilinear BTT and Bilinear MLR attention resolve the low-rank bottleneck for in-context regression. Both bilinear structures allow the scoring function to be more expressive compared to low rank matrices while being efficient. While some tasks including language modeling do not seem to suffer from the low rank bottleneck, believe our technique can improve transformers' performance on variety of datasets that are intrinsically high-dimensional, especially scientific and PDE data. In addition, we use MLR attention to create a flexible distance-dependent compute bias that improves performance on language modeling and time series prediction with less FLOPs. Here too, we believe our methods can have further impact on point cloud data and tasks like code generation, which has not only a strong posi- tional structure, but multiple levels of nested structures such as functions, classes, files, and packages. Furthermore, we define the Multi-Level Block Tensor Contraction (MLBTC) class in Section 3.3 to generalize BTT and MLR. We hope future work will explore the benefits and inductive biases of MLBTC's greater flexibility in attention and elsewhere.
Amsel, N., Yehudai, G., and Bruna, J. On the benefits of rank in attention layers, 2024. https: /arxiv. org szg/abs// /aabs//400 /4407 16153.
Ansari, A. F., Stella, L.., Turkmen, C., Zhang, X., Mercado,
P., Shen, H., Shchur, O., Rangapuram, S. S., Arango, S. P., Kapoor, S., Zschiegner, J., Maddix, D. C., Wang, H., Mahoney, M. W., Torkkola, K., Wilson, A. G., Bohlke- Schneider, M., and Wang, Y. Chronos: Learning the Language of Time Series. arXiv:2403.07815 2024.
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., Zou, J., Rudra, A., and Ré, C. Sim- ple linear attention language models balance the recall- throughput tradeoff, 2024. URL https: org/ abs 24402 18668.
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate, 2016. URL https: / htta/aiivvvorr/ aba / /49..499..00 0473.
In this work, we study the attention score matrix XWQ. WT, K, T and the low-rank scoring function at its center. Similar techniques could be applied to the other part of each attention head, defined by the low-rank matrix Wv; WT, Rather than customizing each head a trans- former the same way, different structures could be used for different heads, as in Xu et al. (2025), to promote better di- vision of labor. Finally, our techniques could be adopted for
Behrouz, A., Zhong, P., and Mirrokni, V. Titans: Learning to memorize at test time, arXiv preprint arXiv:2501.000. 00666, 2024a.
Behrouz, A., Zhong, P., and Mirrokni, V. Titans: Learning to memorize at test time, 2024b. URL https: / /arxiv. org/ hagg/ab// /2501 00663.

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer, 2020. URL https: //aaxxvv..oo//////eee/aaaallaaaa org/ hasp//ab///// abaa /204. 05150.
A. Rethinking attention with performers, 2022. URL httpp::/arrtvv.orr/bbbb////
14794.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024. URL https://arxiv.orrt abs 24../4455..0 21060.
Bhojanapalli, S., Yun, C., Rawat, A. S., Reddi, S., and Kumar, S. Low-rank bottleneck in muultihhead attention models. In Proceedings of the 37th International Confer- ence on Machine Learning, ICML'20. JMLR.org. 2020.
Dao, T., Gu, A., Eichhorn, M., Rudra, A., and Ré, C. Learn- ing fast algorithms for linear transforms using butterfly factorizations, 2020. URL https: htt///axvvvvoro/// abs/ /bb33 05895.
Bommasani, R., Hudson, D. A.. Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosse- lut, A., Brunskill, E., Brynjolfsson, E., Buch, S., Card, D., Castellon, R., Chatterji, N., Chen, A., Creel, K., Davis, J. Q. Demszky, D., Donahue, C., Doumbouya, M., Durmus, E., Ermon, S., Etchemendy, J., Ethayarajh, K. Fei-Fei, L.., Finn, C., Gale, T., Gillespie, I.., Goel, K. Goodman, N., Grossman, S., Guha, N., Hashimoto, T., Henderson, P., Hewitt, J., Ho, D. E., Hong, J., Hsu, K., Huang. J., Icard, T., Jain, S., Jurafsky, D., Kalluri, P., Karamcheti, S., Keeling, G., Khani, F., Khattab, O., Koh, P. W., Krass, M., Krishna, R., Kuditipudi, R., Kumar, A., Ladhak, F., Lee, M., Lee, T., Leskovec, Levent, Li, X. L.. Li, X., Ma, T., Malik, A., Manning, C. D., Mirchan- dani, S., Mitchelll E., Munyikwa, Z., Nair, S., Narayan, A. Narayanan, D., Newman, B., Nie, A., Niebles, J. C., Nilforoshan, H., Nyarko, J., Ogut, G., Orr, L., Papadim- itriou, I., Park, J. S., Piech, C., Portelance, E., Potts, C., Raghunathan, A., Reich, R., Ren, H., Rong. F., Roohani, Y., Ruiz, C., Ryan, J., Ré, C., Sadigh, D., Sagawa, S., Santhanam, K., Shih, A., Srinivasan, K., Tamkin, A., Taori, R., Thomas, A. W., Tramer, F., Wang, R. E., Wang. W., Wu, B., WU, J., Wu, Y., Xie, S. M., Yasunaga, M., You, J., Zaharia, M., Zhang, M., Zhang, T., Zhang, X., Zhang, Y., Zheng, L., Zhou, K., and Liang, P. On the opportunities and risks foundation models, 2022. URL ftttss://axxiie www.vv.or//ee..ooo abs///b///// 220021101.10 07258.
Dao, T., Chen, B.. Sohoni, N., Desai, A., Poli, M., Grogan, J., Liu, A., Rao, A., Rudra, A., and Ré, Monarch: Expressive structured matrices for efficient and accurate training, 2022. URL https: http:///xxvvv.ooo /abs/ 2204. 00595.
Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Ré, C. Hungry hungry hippos: Towards language modeling with state space models, 2023. URL https: //arxiv.org// hawg///b/////// 221..22222..11... 22
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 30583- 30598. Curran Associates, Inc., 2022. URL https: arwia.aiivv..... wwwpwwwwwwww.//////eeee.// 01066
Gu, A. and Dao, T. Mamba: Linear-time sequence mod- eling with selective state spaces, 2024. URL https:
/
00752.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces, 2022. URL https:/ paaps://arxiinn /ttp///axvvvv//a///////eeeeaaava httg//bs/22111
Chen, B., Dao, T., Winsor, E., Song, Z., Rudra, A., and Ré, C. Scatterbrain: Unifying sparse and low- Ranzato, M., Beygelzimer, A.,
rank attention. Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17413-17426. Curran Associates, Inc., 2021. URL https: htttp::://proceedingg. meursiiriiinnnnnnnn...ee..... cc ccc/paper_filespppprr//eee ALBffeee5Lcc544777784444488 pdf.
Han, A., J., Huang, W., Hong, M., Takeda, A., Jawan- puria, P., and Mishra, B. SLTrain: a sparse plus low rank approach for parameter and memory efficient pretrain- ing. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / //openreiiww net/ www.morruuiioooo (oorum?iddWx2e44H70ea
Hatamizadeh, A., Yin, H., Heinrich, G., Kautz, J., and Molchanov, P. Global context vision transformers. In In- ternational Conference on Machine Learning, pp. 12633- 12646. PMLR, 2023.
Child, R., Gray, S., Radford, A., and Sutskever, I. Generat- ing long sequences with sparse transformers, 2019. URL https:/ / /atp/aivvvv////// /ttpw/awwwgg/aaaaaae////a httg/abbs.. abs 904. 10509.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang. S., Wang, L.,, and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 096855 2021.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J.. Mohiuddin, A., Kaiser, I.,, Belanger, D., Colwell, L., and Weller,
10

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Huang, Y., Xu, J., Lai, J., Jiang, Z., Chen, T., Li, Z., Yao, Y., Ma, X., Yang. L.., Chen, H., al. Advancing transformer architecture in long-context large language models: A comprehensive survey. arXiv preprint arXiv:2311.23511, 2023.
Parshakova, T., Hastie, T., Darve, E., and Boyd, S. Factor fitting, rank allocation, and partitioning in multilevel low rank matrices, 2023. URL https: /ttpiv.ox.../////...... abs/2310.19211
Parshakova, T., Hastie. T., and Boyd, S. Fitting multilevel factor models, 2024. URL
Hwang, S., Lahoti, A., Dao, T., and Gu, A. Hydra: Bidirec- tional state space models through generalized matrix mix- ers, 2024. URL https: arxiv. org abb///b///// aaaabb//244 2407. 09941.
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Bac- cus, S., Bengio, Y., Ermon, S., and Ré, C. Hyena hicrar- chy: Towards larger convolutional language models, 2023. URL https: //t./aiivvv//////////////////aaa org.iiioorrroo..oooo..ooo .og. / 2302 10866.
Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, E. Repeat after me: transformers are better than state space models at copying, In Proceedings of the 41st In- ternational Conference on Machine Learning, ICML'24. JogR.org, 2024.
Potapczynski, A., Qiu, S., Finzi, M., Ferri, C., Chen, Z.,, Goldblum, M., Bruss, B., Sa, C. D., and Wilson, A. G. Searching for efficient linear layers over continuous space of structured matrices, 2024. URL https:// arxi v. org/ abagg//bb/// 2410 10. 02117.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention, 2020. URL https://arrae
org/abs/ abaaabbs//550 16236.
Qiu, S., Potapczynski, A., Finzi, M., Goldblum, M., and Wilson, A. G. Compute better spent: Replacing dense layers with structured matrices, 2024. URL https:
Lavie, I,,
Gur-Ari, G., and Ringel, Z. Towards understanding inductive bias in transformers: A view from infinity, 2024.
URL httph::h////mm /tt///aaxxx aaaaiiiiivv....///aaaa/aaaaavvva org/ Abs/2402 05173.
//arxiv. org/ abs///b//// 2400 06248.
Liaw, R., Liang, E., Nishihara, R., Moritz, P., Gonzalez, E., and Stoica, I. Tune: A research platform for distributed model selection and training, 2018. URL https: //arxiv.// org/ h/s///b///// /11//110.0000 05118.
Ren, H., Dai, H., Dai, Z., Yang, M., Leskovec, J., Schu- urmans, D., and Dai, B. Combiner: full attention trans- former with sparse computation cost. In Proceedings of the 35th International Conference on Neural Informa- tion Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845399
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- org/ abs/
larization, 2019. URL
1711.05101.
Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhu- patiraju, S., Hussenot, L., Mesnard, T., Shahriari, B.,
Lotfi, S., Finzi, M., Kuang, Y., Rudner, T. G. J., Gold- blum, M., and Wilson, A. G. Non-vacuous generaliza- tion bounds for large language models, 2024a. URL https: /tas:///rxxiii wtr.v..rgg///aoaaaaooooooooooorr abs/2312..... 17173.
Ramé, A., et al. Gemma 2: Improving open language models practical size, 2024. URL https://arxiv. org/abs/2408.0011 18, 1(3), 2024.
Lotfi, S., Kuang, Y., Amos, B., Goldblum, M., Finzi, M., and Wilson, A. G. Unlocking tokens as data points for
Sanford, C., Hsu, D., and Telgarsky, M. Representational strengths and limitations of transformers. In Thirly- seventh Conference Neural Information Processing Systems, 2023, URL https: : optnreview.eeeeee// net/ forum 1imummmdd=3xxxx222
generalization bounds on larger language models, 2024b. URL https::/p/// /aaxivivvvv.///ea//eea/eaaaaavva org abs/2407 18158.
Luong, T., Pham, H., and Manning, C. D. Effective ap- proaches to aateetion-based neural machine translation. In Mårquez, L., Callison- C., and Su, J. (eds.), Proceedings of the 2015 Conference on Empirical Meth- ods in Natural Language Processing, pp. 1412-1421, Lisbon, Portugal, September 2015. Association for Com- putational Linguistics. doi: 10. 10.18653///DD15-111666 URL https: / //a//aaaaathhlcce wtwp:/ggg.rrrooooo wttp/or/////////// 1166/.
Sehanobish, A., Dubey, K. A., Choromanski, K. M., Chowd- hury. S. B. R., Jain, D., Sindhwani, V., and Chaturvedi, S. Structured unrestricted-ramt matrices for parame- ter efficient finetuning. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https:////// openreviewwwweeeeeeeeeeeeereewwe net/ forum? id Ad-MxozgiiWDF.
Olivares, K. G., Challú, C., Garza, A., Canseco, M. M., and Dubrawski, A. NeuralForecast: User friendly state-of- the-art neural forecasting models. PyCon Salt Lake City, Utah, US 2022, 2022. URL https: /ithhb..uub..g/g// com/ Mixtlannent wwweerraaforrcccm forecast.
Su, J.. Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding, 2023. URL htips: /awxivivvvvo www.vi..rrr//.o. org/abs/ /bs/ 2104. 09864.

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Thomas, A. T., Gu, A., Dao, T., Rudra, A., and Ré, C. Learn- ing compressed transforms with low displacement rank. In Proceedings of the 32nd International Conference Neural Information Processing Systems, NIPS' 18, pp. 9066-9078, Red Hook, NY, USA, 2018. Curran Asso- ciates Inc.
Zhang, M., Bhatia, K., Kumbong, H., and Re, C. The hedge- hog & the porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //opperevieee ew net. forum? 4
Zhou, H., Zhang, S., Peng, J., Zhang, S., Li, Xiong, H., and Zhang, W. Informer: Beyond efficient transformer for long sequence time-series forecasting, 2021. URL https: /aaaaaaxvvvvrrr..a/aaa org/
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N.. Kaiser, u., and Polosukhin, Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H.,, Fergus, R., Vishwanathan, S., and Garneut, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://procceddnggs. nwwrriisinnnnnnnns.i.o....m....s ce cc/ cc/paper wwtpppppe_f.lleeeeeeeeeeeeeeeeee [iaaer__filess/pppe///////// 2h0e44455444eee111e000001aalaae
Zhu, Z. and Soricut, R. H-transformer-ID: Fast one- dimensional hierarchical attention for sequences. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Pro- ceedings of the 59th Annual Meeting of the Associa- tion for Computational Linguistics and the 11th Inter- national Joint Conference on Natural Language Process- ing (Volume Long Papers), pp. 3801-3815 Online, August 2021. Association for Computational Linguis- tics. doi: 10.18663/w1222221accoooooo 294. URL https: / waclaatthloo...rrgg///o/o http://2222..111ongaeeeee.a 294/.
pdf.
Warner, B., Chaffin, A., Clavié, B., Weller, O., Hallström, O., Taghadouini, S., Gallagher, A., Biswas, R., Ladhak, F., Aarsen, T., et al. Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference. arXiv preprint arXiv:2412. 13663, 2024.
Wei, X., Moalla, S., Pascanu, R., and Gulcehre, Build- ing on efficient foundations: Effective training of LLMs with structured feedforward layers. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: / http////ppeeeeieeeeeeeeeeee net. forum? www.iidddddiiLiiieeeee.e
Xu, Y., Nag, S., Li, D., Tian, L.,, and Barsoum, E. Mswa: Refining local attention with sulli-scalewindow attention. arXiv preprint arXiv:2501.01039, 2025.
Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer, 2022. URL htt.h:tp/// / arxiv. htg//bb/2200 03466.
Yang, G., Simon, J. B., and Bernstein, J. A spectral condition for feature learning, 2024. URL https: // //arxiv.. www.v.org//.....o.........oooott abs/22110 17813.
Ye, Z., Guo, Q., Gan, Q., Qiu, X., and Zhang, Z. Bp- transformer: Modelling long-range context via binary partitioning. arXiv preprint arXiv: 1XX1::911.04077 2019.
Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z.,, Wei, Y. X., Wang, L.,, Xiao, Z.,, Wang, Y., Ruan, C., Zhang, M., Liang, W., and Zeng, W. Native sparse attention: Hardware-alignee and natively trainable sparse attention, 2025. URL https: // /arxiv..... org/ abs/ 1109.
12

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
A. Rectangular Structured Matrices
In this section, we discuss the extension of Multi-Level Low Rank (MLR) matrix and Block Tensor Train (BTT) matrix to rectangular shape m n. The remaining appendix sections also follow the notations defined for rectangular MLR and rectangular BTT.
A.1. Rectangular MLR
MLR matrix with shape x n is given by
① k=1
L R
(10)
where Di-i 1 R, E R= is a block diagonal matrix with pl blocks at level L. Each block L, is a low rank product for L E R xri and R E x By construction, we have Σk=1 Mi, m and E 71, = n. The
MLR rank of the matrix W is defined TMLR Et= ri.
A.2. Rectangular BTT
BTT matrix of shape x is given by
PL L PR + R k==1
(II)
k=
where R₂ E gdxbs Ly E _axes ab= = m, ed = n. The extra dimension s is the BTT rank. Pt permutes the rows by rearranging dimension ba into ab. This permutation is equivalent to reshaping a vector Re into a matrix Z E R tranposing to get Z Baxb and reshaping it back to vector z' RRab, PR permutes the rows by rearranging dimension cbs into bes.
B. Bilinear Attention Scoring Functions with Structured Matrices
In this section, we describe how the bilinear BTT and bilinear MLR attention scoring functions can be interpreted in terms of queries and keys with projection matrices that are structured instead of dense.
B.1. Notations
The appendix section adopts the following notation. Let B be the batch size, T be the sequence length during training, We, WK, Wv Dxr be projection matrices, and H be the number of heads such that = Hr. For any x; E Dx1 we obtain q k2 V I by the projections qe W k; W Wx,,, V; WJx,. Since matrices in attention usually have square shapes, assume
D.
B.2. Bilinear MLR
Let XW Q WTX T be the form of attention matrix per attention head. We can replace the low rank structure We WT. with MLR(W Q, W x) and obtain the Bilinear MLR projections:
PL WQ..
Q
MLR
P
W
Q
(12)
P1
K MLR
Ps
W
XT
(13)
Ki,s
where W Wo.. E Hn WK,, RR and thus D the shape of the product X D WQ.. is (BT,p, and the shape of the product The permutation matrix P1 rearranges the shape (BT, HrHr) into (B, H, T, r.p:). Pg rearranges the shape (p,Hr, BT) into (B, H, riPt, T).
Wo. Rm
W Hu Thus
W k, X (pillri BT).
13

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
The final output of the BilinearMLL sequence mixing operation O BB is given by
Oaf (N
MLR )
(14)
where Q E TR and K E RR In practice we assume = 2/-1, and mI,K M1, = D/2
B.3. Bilinear BTT
We can replace the low rank structure Wo W
& with QT.. W k and obtain the Bilinear BTT projections:
Q BIT
P3
X
W
QN
(15)
k=1
KBIT
PA
W K. X
(16)
where assume X BTXxD W Q+ R Hes W K2
BR and thus Di.1 Wo. baxbies Oz Wk. k. He Here we assume D = ab = cd. Thus the shape of the product X + WQ. is (BT, bHcs) and the shape of the product + l= WT
K2 X is (cHbs, BT). The permutation matrix P3 rearrange the shape (BT, bHes) into (B, H T, sbc). P4 rearrange the shape (cHibs, BT) into (B, H, sbc, T).
Let a denotes the batch index and 8 denotes the head index such that BTT E Txabe and KBIT Rabbxx T The final output of the bilinear BTT O E RBxHXXXT is compute as
QUTTKBI Oug oM Q K
(17)
where is the softmax function and M is the lower diagonal causal mask matrix. Both B and H are extra batch dimensions that can be computed efficiently using the batch matrix multiplication primitives.
In practice, it's always true that sbe > r where r is the head dimension for standard multi-hear attention.
C. Multi-Level Block Tensor Contraction
To that MLBTC captures all MLR matrices (including low-rank matrices special case), set a₁ P, PR I, k' = k. p/1 Pt. and r/ in Equation (8). To see that it captures all BTT matrices (including Monarch, Butterfly, and Kronecker), set 0 for all but one I and select the dimensions to match Equation (5): P = b, = c, r = Cs, = bs. & = d, and my,k' = a.
Now consider square matrices such that m = n = D. Since BTT can approximate arbitrary dense matrices of shape x D based on Qiu et al. (2024), MLBTC can also express any D matrices with large enough r and r/. Like MLR and BTT, MLBTC matrices be used to define efficient full-rank scoring function. We leave this generalization to future work.
D. MLR Attention Derivation
In this section, we present Lemma D.1 that connects the matrix form of MLR and its scoring function formula.
Lemma D.I. The (j.j') entry of the scoring matrix S = Et= Pc shown in the right hand side of Equation (9).
Kr. has the form x]
(Et L,R.J
d(s''')
Xj as
Proof. For simplicity, assume there are just two levels. So Equation (9) reduces to the following expression:
85 (x, x,) = x3 RT + L2RJ)xx
(18)
X X
Divide X =
into blocks. X2 is the first half of the sequence and X, corresponds to the second half. Divide
14

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
WQ Q
L
L2
Thus
L L2 X₂L1 X
(19)
Now divide Q into named blocks according to Figure lb:
Q XL,
(20)
Q21 X,
(21)
Qz X2L2
(22)
Analogously
K1 XR1
(23)
2 R2
(24)
K X2R2
(25)
Finally, plug the above into the definition of S:
= Qu Qu
Q21 Q21K2 K 0
X
X
0 X x
(26)
K
0
Consider the (j.j") entry of S. If they're in different blocks, then it's x, L,R. xj. If they're 're in the same block, it's x] L,L,, + xx] L2 x,. The same reasoning applies when we have more than two levels.
E. Optimal Tensor Contraction Order
E.1. Bilinear MLR
In this section, we compute the costs to do Bilinear MLR per attention head assuming batch size equal to 1. We write bilinear MLR as
XMLR(WQ, W = X
Σ
θ
W Wo
Q
WT
KIN
X
(27)
Im
km
Σ
!
W War WhIt
X
(28)
Woln
(#WWi.
WK X
(29)
k=1
W
QuA
W KiN
X
(30)
k=
P1
Pr
x
W
Q
W
KrN
X
(31)
We summarize the FLOPs for each of these tensor contraction orders in Table 2. We find that
(@ ki X
X Wa...
(32)
k=
K=1
has minimal compute costs of Dr 2 E 21-1 1r, This the form of bilinear MLR we used all of our experiments.
15

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Table 2. FI.OPs for Bilinear MLR with Different Tensor Contraction Ordering
TENSOR CONTRACTION ORDER
FLOPs
(Low Rank) XWoWkX
1 + 2TDr
T +
X
E
WOl
WKL
X
D' अ + -i==
EL
X
W
Q WK
X
T + Ei F T
D'm TD²
1
E
X
W 4
DEL WK X
2T + T Ei= 2
EL X
+
W
LLE
WT WK, X
LT + 2r
X
WKI X
T + 2TDr
E,2. Bilinear BTT
In this section, we analyze the FL.OPs required for bilinear BTT assuming batch size equal and attention head. We write bilinear BTT
(Pn θ
XBTT(WQ: WK)X
XP ① Wq.
q.
W Kv
x
(33)
k=1
k=1
x Pr k=1
W Q
P
(34)
k=1
where E TXD WQ. E R eex WK. Kr E Rdd with D ab cd. Table 3, we list out the FLOP counts with the
simplifying assumption of a = b = = d VD. Thus we should be using the form
X PL + W Q
PR + W x )
(35)
k=1
k=1
with FLOPs D + D3/2
Table 3. FLOPs for Bilinear BTT with Different Tensor Contraction Ordering
TENSOR CONTRACTION ORDER
FLOPS
XP Wo: Px PR W
sT 2
X P O
Wo, PR W
T + 2/T
F. Practical Considerations.
F.1. Maximal Update Parameterization (nP)
For a dense matrix W E Xdd and an input vector x E Reda, the output hidden state E IR is computed as Wx. uP requires both h and the gradient update Ah have norm ||h|/ = ||hh|22 ((vdm) for stable feature learning under width scaling. This is equivalent imposing the spectral norm constraints on the weight matrix such that (( /oout// d)p) (Yang al., 2024). Thus the weight matrix W should be initialized from (O0 a for ow = ( / p min(1, / don/ /dm)) with learning rate 1/w = O(dout/da) /m). For AdamW (Loshchilov & Hutter, 2019), the learning rate scales as 1/w (11dmm) (Yang et al., 2022). Qiu et (2024) develops μP for structured linear layers that
[ ||]
16

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
expressible compositions of dense batched matrix multiplications and reshapes. Since structured attention variants fit this form, they amenable ytP.
For bilinear MLR with the optimal tensor contraction order shown in Equation (32), we apply μpP on the dense matrices W WQ..& E R Hei and W Ki,, E R where we assume Mi, = R1,k = D/pt. Since the fan-in dimension for both W Q and WK,, are D/pt, we initialize them as W Q WK,, ~ (0, IT. For learning rate, we first select a base value Thase E {0.001, 0.0005, 0.0001, 0.00005, 0. J.00001} embedding dimension of the transformer model. Then for any model width D2 > D1. we the learning rate 7War, and
certain model width Dr. Here the model width the
K,,, to be
D1 Toase D2/pt
(36)
for optimal learning transfer across width.
Bilinear BTT follows the optimal tensor contraction order in Equation (35) where WQ, gaxes and W K, E gdxbs Thus we initialize them as W Q= ~ ~O,, I) and W WK. ~ (N(0, IT). The optimal learning rates are given by
D1
TWa = TJoase
(37)
C5
D1
VWW K2 Thase
(38)
a
Additionally, we apply RMS normalizations the weight matrices WQ. Qe and W K2 for training stability as shown in Qiu et al. (2024). The rest of the parameters in the transformer models follow the recipe in Yang et al. (2022). For MLR attention, we note that the fan-in dimensions are unchanged and thus follows the jPP recipe for a standard transformer architecture.
F.2 QK LayerNorm
Let LN(-) denote layer normalization. We apply normalizations to bilinear MLR over our optimal tensor contraction order to get
Σ
P1 LN
W
Wo..
Qu.k
LN
+ W
W.W. KL.
)
C
(39)
rIPI
for some tunable constant C. We also apply layer normalization to bilinear BTT and get
C
LN X LN P
W QN
P
W K, X
(40)
ab
1=1
k=
for some tunable constant C*, In standard attention, the score matrix is normalized by 1/ before applying softmax. uP recommends normalizing by 1/r instead (Yang et al., 2022). We incorporate the jPP version of normalization into our architectures. That is, for Bilinear MLR and Bilinear BTT, we use C/ript and C* /ab respectively instead of C/ /ript and
C* c+/Vab.
G. μP Additional Figures
this section we show that under the maximum update parametrization (μPP), both our MLR attention and standard attention has stable learning rate across width as shown in Figure 5.
H. ICL Additional Figures
In this section we provide additional figures of in-context regression performance as function of compute and training steps across different input dimensions, model widths, and types of scoring functions shown in Figures 7 to 11. We also provide error bars for Figure 2 as shown in Figure Due to compute budget constraints, we only plotted the error bars when the input dimension 16.
17

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
MLE Anorice Midel Wath 512
Attrace Alxid Wadh 269
Stankel Aftesions
Madal W136 512
Seaubrt Atretion Madel Wiln 768
Vaindation Lass OpeaWWdToo
1.22
1.000
1.25
Learning
Figure show the validation loss an 8-Level MLR attention standard attention OpenWebText across a variety of learning rates. Here we sweep over model widths D {512, 768, 1024) with reduced context length of 256 due to compute constraints. As figure shows, our MLR attention shares the same optimal learning rate across model width, and it's s also consistently better than standard attention when both are properly tuned.
1.0
Input Dimensions
Squared Error
0.8
dipput 16
0.6
0.4
108
Rank Ratio r/ diput (Log Scale)
Figure 6. The performance of standard multi-head attention as function of the rank ratio. We take the runs from Figure 2 and plot error bars when the input dimension is 16, i.c. dinon 16.
I. Language Modeling Additional Figures
In this section we present additional figures for language modeling performances between MLR attention and variants of sliding window attention across model widths shown Figure 13.
J. Time Series Forecasting Additional Figures
We train transformer model with 2 encoder layers, embedding dimension D = 512, and 8 attention heads on the ETTh subset of the ETT dataset. The ETThl dataset records the oil temperature and six power load features collected each hour from different electric power transformer stations (Zhou et al., 2021). The dataset contains records with a variety of time horizons. We select horizons of T E {96, 332) hours. Thus, the time series given to the transformer models has sequence length 7 with feature dimension
18

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
The transformer model used has input linear projection layer that maps points in R² to RP and output linear projection layer that maps points from R D to RR7. Following Olivares al. (2022), sweep over optimal learning rates in the range of {0.00001, 0.0002, 0.005) and SGD iterations in {200, 1000). In Figure 14, we plot the relative improvement of MLR attention over standard attention in terms of Mean Absolute Error (MAE) for three different time horizons. The relative improvement obtained by subtracting the final MAE value of MLR attention with the final MAE value of standard attention normalized by the final MAE value of standard attention. We observe that as sequence length grows, MLR attention outperforms standard attention oil temperature prediction accuracy.
Ray Tune (Liaw et al., 2018)
Finally, we substitute the attention mechanism in the foundational model of Chronos Ansari al. (2024). Similar to our text experiments, MLR attention reduces the computational cost as seen in Figure 15. For this experiment, we train a Chronos model on the data mixture from Ansari et al. (2024) using combination of synthetic kernel data and real data. We use an MLR rank split of 32) |
19

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Model Width
Model Width
Model Width 64
Square Exror (log scale)
10-4
10-1
1014
1013
10
1014
1015
Nom EanbEenbediing Compute (log scale)
Nom Embeddding Compute (log scale)
Non: Embeddiing Compute (log scale)
Model Width
Model Width 256
Model Width - 512
Squared Ertor (log scale)
10-1
1014
10
10
10
10
Non-Embeddiig Compute (log scale)
Con-Embedding Compute (log scale)
Non-Embeddiing Compute (log scale)
Model Width
Model Width = 32
Model Width 64
102
Squared Error (log scale)
10-1
Steps (log scale)
Steps (log scale)
Steps (log scale)
256
Model
Squaee Ervor (log scale)
10'
10"
10" Steps (log scale)
Steps (log scale)
Steps (log scale)
Structure
Atteation (1 Head)
BiliscuBTT
Attention (8 Head)
Figure show in-context regression performances as functions of training steps and training compute (excluding the input and output linear projection layers) across model width D E (16, 32, 64, 128, 256, 512) for diugot = 16. In general, Bilinear BTT outperforms attention with varying number heads.
20

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Model Width
Model Width
Model Width
Square Etror (log scale)
10-1
10-3
100
1014
1014
10F
1014
Non Embedding Compute (log scale)
Nom Embeddding Compute (log scale)
Non Embeedding Compute (log scale)
Model Width 128
Model Width 256
Midel Width = 512
Squared Ertor (log scale)
10-1
10
10
10
10
Non-Embedding Compute (log scalo)
Non-Embedding Compute (log scale)
Non- Embedding Compute (log
Model Width 16
Model Width = 32
Model Width 64
10P
Squared Error (log scale)
10-3
Steps (log scale)
Steps (log scale)
Steps (log scale)
Model Width
256
Model Width
Squired Ervor (log scale)
10
III
10'
Steps scale)
Steps (log scale)
Steps scule)
Structure
Atteation (1 Head)
BiliscuBTT
Attention (8 Head)
Figure show in-context regression performances as functions of training steps and training compute (excluding the input and output linear projection layers) across model width D E (16, 32, 64, 128, 256, 512) for diugot = 32. In general, Bilinear BTT outperforms attention with varying number heads.
21

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Model Width
Model Width
Width
Squared Earor (log scale))
10-4
10+
10
10
1015
16-
10ls
101
Non Embeedinng Compute scale)
Nom Embeddiing Compute (log scale)
Non Embedding Compute (log scale)
Model Width 128
Model Width = 256
Model Width 512
Square Enoo (log scall)
10-1-
10
Non-Embedding Compute (log scale)
Non-Eam Non-Embedding Compute (log scale)
Non: Embedding Compute (log scale)
Model Width
Model Width
Model Width 64
Squared Enor (log (log scale)
1tP
IIP
10 Steps (log scale)
Steps (log scale)
Steps (log scale)
Model Width = 128
Model Width 256
Model Width 512
Squared Ervor (log scale)
10-1
10-4
1P
10
Steps (log scale)
Steps (log scale)
Steps (log scale)
Structure
Attention Head)
Attention (8 Head)
BilineuBBT
BilineaMMLRR
Figure 9. We show in-context regression performances as functions of training steps and training compute (excluding the input
and output linear projection layers) across model width D E {16,33 32, 64, 128, 256, 512) for diggot 64. In general, Bilinear BTT outperforms attention with varying number of heads. Bilinear MLR also outperforms Bilinear BTT when D = 256 or 512. The different Bilinear MI. lines correspond different levels Bilinear MLR based (ri,".",rs)
22

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Model Width
Model Width
Model Width 64
10"
Squared Eiror (lgg scale)
10'
10-1
10th
1013
10
10ah
100
Non Embedding Compute (log scale)
Nom Embeddiing Compute (log scule)
Non Embedding Compute (log scale)
Model Width
Model Width 256
Model Width - 512
Squared Extor (log scale)
t0-1
10-4
10
10
10
10
10" Non-Embedding Compute (log scalo)
Non-Embedding Compute (log scale)
Non-En Bob-EEmeddiing Compute (log scale)
Model Width =
Model Width = 32
Model Width 64
10"
Squared Emror dlog scale)
10
Steps (log scale)
Steps (log scale)
Steps (log scale)
256
Model
Squaee Ervor (log scale)
10-1
10-4
10
1II
10a Steps (log scale)
Steps (log scale)
Steps (log scale)
Structure
Atteation (1 Head)
Attention (8 Head)
BilincaBTT
Figure We show in-context regression performances as functions training steps training compute (excluding the input and output linear projection layers) across model width D E {16, 32, 64, 128, 256, 512) for diupat 128. In general, Bilinear BTT outperforms attention with varying number heads.
23

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
1.0
Stracture
Structure
0.8
Attenionn Head)
0.8
Amenntoon (8 Head)
Attenton Head)
Attention (1 Head)
Squared lirre
BilinarBTT
Error Error
BilinorBTT
0.6
Squared
0.4
0.2
0.0
0.0
10
10
Model Width (Log Scale)
Model Width (Log Scale)
(a) Input Dimension (dinput) = 16.
(b) Input Dimension (clagu) = 32.
1.0
Stracture
Attemion Head))
Attenion Head)
0.8
BilincarBTT
Squared Error
0.6
Squared
0.
Siructure
Attention Head)
0.2
0.2
Head)
DilincurBIT
0.0
0.0
10² Model (Log Scale)
10" Model Width (Log Scale)
(c) Input Dimension (dimpu) = 64.
(d) Input Dimension (dimau) = 128.
Figure 11. Bilinear BTT and 1-head attention full rank structures and thus achieve lower regression errors (y-axis) compared to S-head attention at smaller model width (x-axis). Each one of the four figures correspond to one particular input dimension daput.
200
Strattue
175
BiiseuBTT
Squared Eroor (log scale)
Wall Time (minutes)
142 112484E7 DES
125
100
75
10
0
50
100
125
150
175
200
BilinarBTT
(8 Heal)
Wall Time (minutes)
Structure
(a) Loss as function wall time
(b) Total wall time for in total 500, 000 training steps
Figure 12. Wall time comparisons between bilinear BTT and standard attention for input dimension at dupput 16 and model width D 64 and number of attention heads H 8. We omit the I head attention since has wall time very close to 8 head attention. (a) we show the squared regression error in log scale (y-axis) against the wall time between Bilinear BTT and standard attention. We observe that Bilinear BTT still outperforms attention when controlling for wall time. (b) We plot the total wall time for Bilinear BTT and standard attention for 500, 000 training steps. Although Bilinear BTT 1.35x slower, still has much better performance compared to standard attention.
24

---

Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Glubd 590 musa
SPA
Walli
Valilaim (Lo Scal))
12x11
Malle112 2.1x
12x
Wils 700
13x10
Compute (Log Scale)
Compuse (Log
(a) MLR Attention Global SWA
(b) ML. Attention SWA
Figure 13. MLR attention outperforms both purely sliding window attention (SWA) and combinations of sliding window attention and standard attention with global context across model widths. SWA refers all layers being sliding window attention. Global SWA means sequentially stacking standard attention and SWA.
Relutive Improvement MATE (3))
Relative Improvement in MAE (%)
0.8
0.6-
-2
0.4
0.2
-3
0.0
-4
336
Hotizon
Horizon
(a) 2 Levels MLR Attention ETTh1
(b) 4 Levels MLR Attention ETTh1
Figure the time horizon (i.c. sequence length) becomes larger, MLR attention outperforms standard attention in oil temperature prediction accuracy by around 1%. (a) We show the relative improvement in Mcan Absolute Error (MAB) of 2-levels MLR attention with rank distributions 48/16. For both horizons 96 and 336, MLR attention achieves better temperature forecasting accuracy. (b) We show that 4-levels MLR attention with rank distributions 40]16 4 outperforms standard attention in relative MAE the sequence length becomes larger. Across the horizon, we observe a positive correlation between relative MAE and the time horizon of the ail temperature data in the ETThl datasets. See Appendix for additional details.
(128)
GA (256)
x 10"
MLR (128)
MUR
MLR
Validation
x 10"
Compute (log)
Figure 15. MLR attention reduces computational cost when compared to standard attention. We plot the validation loss of three
distinct model sizes when substituting the T5 atlention mechanism in Chronos (Ansari al., 2024) from standard attention to MLR attention. As shown in the figure, MLR attention achieves the same validation loss compared to standard attention with smaller FLOPs.
25


Reviews:
[{"review_id": "vutl6etSBL", "reviewer": "Reviewer_g9Ct", "summary": "The paper address two issues of the attention computation.\n\nThe first is the bottleneck caused by the low rank computation of the Key and Query matrices in the attention computation. Instead of doing the standard low rank decomposition, they proposed to use structure matrices to represent the attention score, which is full rank but still more efficient in computation that the full DxD matrix (O(D^(2/3)) v.s. O(D^2)).\n\nThe second is not paying more attention to locality. Most problems have locality properties where the closer tokens are often more relevant than the further tokens. By introducing structure matrices, the proposed method introduce locality into attention computation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and positive feedback. We address your questions below.\n\n**Contributions in Comparison to Prior Works.** \n\nThank you for pointing out the connection of our work to the broader literature of structured matrices. We would like to highlight further connections that we did not mention explicitly in the paper: \n- Several prior works [2, 3, 5] replace linear layers in the Transformer model for improved efficiency and a better scaling law. Our work goes beyond them by 1) considering structured matrices in bilinear transformations, like the one that defines the attention score, and 2) using structured matrices to replace data-dependent matrices like QK^T, where the query and key matrices are themselves the outputs of a linear layer. \n- The Multi-Level Low Rank (MLR) matrix [1] was originally proposed in the applied math literature as an extension of low rank matrices to fit multilevel factor models. To our best knowledge, our work is the first to apply MLR matrices to deep learning models.\n- We propose Multi-Level Block Tensor Contraction (MLBTC), a family of structured matrices that generalizes many common structured matrices like Low Rank, Kronecker, Monarch [2], Block Tensor Train [3], Multi-Level Low Rank [1]. We believe this is a novel perspective. We conduct experiments to explore the various inductive biases that these structured matrices encode when applied to a Transformer architecture.  \n- The Hydra paper proposed the matrix mixer framework that unifies common sequence mixer modules (e.g. CNN <-> Toeplitz, Linear Attention <-> Low Rank, Mamba <-> Semi-Separable, etc.) with their underlying structured matrices. We believe our efforts add to this line of work of exploring novel structured matrices for sequence mixing.\n\n**Writing.** \n\nThank you for your helpful feedback regarding the writing and typos. We have updated our manuscript accordingly, alongside additional clarifications inspired by your questions.\n\n\n**Locality.** \n\n We believe that problems with long context lengths and a hierarchical structure are most amenable to our approach. The computational savings of MLR attention compared to standard attention increases with the sequence length. And the multi-level nature of MLR matrices make them well-suited to hierarchically structured data, where tokens become gradually more and more related as their distance decreases. One application we are excited about is code models. Code repositories tend to be large, pushing the limits of standard transformers' context windows. They are organized hierarchically into folders and files, and each code file is further organized by an Abstract Syntax Tree, with hierarchically nested classes, methods, control structures (loops), lines, and expressions. Our method may help transformers read repositories more efficiently, like humans do, focusing most of their effort on the immediate context, but keeping the global structure in mind too.\n\nWe are also intrigued by potential applications to non-text data, such as DNA sequences, phylogenetic data, time series data, and graph data. We plan to apply our method to some of these settings in future work. \n\n\nThank you again for your detailed review. We hope we were able to address all of your questions and that you would consider raising your score. Please let us know if you have any additional questions or comments that we can address or discuss.\n\n________\n**Reference:** \n\n[1] Factor Fitting, Rank Allocation, and Partitioning in Multilevel Low Rank Matrices. https://arxiv.org/abs/2310.19214\n\n[2] Monarch: Expressive Structured Matrices for Efficient and Accurate Training. https://arxiv.org/abs/2204.00595\n\n[3] Compute Better Spent: Replacing Dense Layers with Structured Matrices. https://arxiv.org/abs/2406.06248\n\n[4] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. https://arxiv.org/abs/2407.09941\n\n[5] Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices. https://arxiv.org/abs/2410.02117", "claims_and_evidence": "* Better at the efficiency and performance trade-off.\n  - Figure 10 in the Appendix shows the loss as a function of wall time.\n\n* Locality.\n  - The paper includes experiment on regression problems where the standard attention performs poorly and show that the proposed method can improve the performance (Figure 3).", "methods_and_evaluation_criteria": "Yes. The proposed methods introduce structure into the attention computation that was originally represented as low-rank matrix multiplication. The method is evaluated on in-context regression, which highlight the low-rank bottleneck issue.", "theoretical_claims": "N/A.", "experimental_designs_or_analyses": "The paper uses experiments measuring squared error on in-context regression tasks and shows that the proposed method outperforms the standard method.", "relation_to_broader_scientific_literature": "The paper is proposing structured matrices in attention computation. There were other previous works using different design of structured matrices.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "The paper is clear in the target problem and proposed methods to solve the problem.\n\nThe writing of the paper can be improved.", "other_comments_or_suggestions": "Typos:\n  - line 033. Section 1, paragraph 2, third last line. \"express with attention\".\n  - line 304. Section 4, paragraph 3, last line. The equation is missing a \"(\".", "questions_for_authors": "I'm wondering if there are textual problems (non-regression problems) where the locality can have an impact.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FsG4SQX736", "reviewer": "Reviewer_1ojm", "summary": "The authors propose to address a common limitation of existing softmax attention layers: the information bottleneck when using small head dimension. To do this, they propose to bake in a locality bias into the structured parameterisation of the attention weights.\nThe attention mechanism is introduced cleanly as a billinear form, from which they can introduce structured and parameter efficient bilinear forms. This work appears to build upon the work of Parshakova et al but in the context of designing an efficient attention layer.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your supportive feedback. We address your comments below. \n\n**Larger Models and Datasets:**\n\nWe are hopeful about the prospect of scaling up these experiments to even larger models and datasets. As a first step in that direction, we now present a new experiment on hyperparameter transfer for our architectures. The maximum update parameterization (muP) has become a crucial tool for scaling models up because it provides a way to transfer the results of hyperparameter tuning from small models to large ones. This is far less expensive than tuning hyperparameters for large models directly. For architectures that use structured matrices, muP does not work out of the box, but in Appendix E of our paper, we provide a recipe for adapting muP to our architectures. Our latest experiment validates this recipe.\n\nIn the figure provided in the link https://drive.google.com/file/d/18pI--DmWWqB3PqFd-dEaziOypc5rRIP8/view?usp=sharing, we show the validation loss of an 8 Level MLR attention and standard attention on OpenWebText across a variety of learning rates. In our paper, the maximum model width we used is 768 with a context length of 1024. Here we sweep over model widths  512, 768, and 1024 with a reduced context length of 256 due to compute constraints. As the figure shows, our MLR attention shares the same optimal learning rate across model width, and it’s also consistently better than standard attention when both are properly tuned. This result suggests MLR will continue to perform well on larger models trained with more data.\n\n**Efficient Implementation.** \n\nAs the review notes, a thoughtful IO-aware implementation is needed to ensure that savings in FLOPs translate to savings in wall-clock time on a GPU. Like standard attention, our proposed methods rely on a few batch matrix multiplications followed by softmax. The block diagonal structure of MLR matrices is easy to parallelize, and we think the highly structured permutation matrix in BTT can be fused with the surrounding multiplications. Thus, we are optimistic that techniques similar to Flash Attention can be applied to our attention variants (and indeed, to general MLBTC matrices). \n\n**Rethinking Attention with Performers. ICLR 2021**\n\nThank you for pointing out this paper! We have incorporated a discussion of it in the revised manuscript. \n\nWe would like to note that while the Performers paper proposes an algorithm for computing attention more efficiently achieving a linear instead of quadratic dependence on the sequence length, the function they are (approximately) computing is the standard attention function. As a result, they still have a low rank bottleneck and they lack a locality bias in the allocation of compute. In contrast, our approach to improving transformer models is to replace the standard attention function with something else. Our proposed methods have different inductive biases from standard attention, which we show leads to better performance in several cases. However, we retain the quadratic dependence on the sequence length. We leave it to future work to find efficient algorithms for (approximately) implementing our proposed architectures.\n\nThank you again for your detailed and supportive review. We hope we addressed your questions. Please let us know if you have any additional questions or comments that we can discuss.", "claims_and_evidence": "The main claim is that current attention layers have an information bottleneck, which is more significant when the head dimension is reduced. Naturally, this depends on the data, but the intuition is valid, the results confirm this, and the authors provide significant references to existing works that highlight both the theoretical and practical cases where this can be a problem.\n\nSecondly, they propose that a locality bias is a good way to introduce more parameter efficiency, without degrading performance. The authors confirm this with emperical results.", "methods_and_evaluation_criteria": "The authors evaluate on the OpenWebText dataset and the ETT time series forecasting dataset, where they observe better performance as the sequence length grows. Although the first dataset is moderate size, a large scale dataset and larger models would give this paper a much bigger impact. Although understandably, there are new difficulties when scaling up these methods to ensure theoretical flops does translate to real wall clock performance reductions.", "theoretical_claims": "I looked through the derivations and they appear to be correct.", "experimental_designs_or_analyses": "The wall clock time is great to see and strengthens the practical impact of this paper. The implementation using batch matrix multiplication is simple and effective for these structured block matrices.", "relation_to_broader_scientific_literature": "This work naturally builds upon existing and recent development of multi-level low-rank matrices. These fit very nicely into the attention matrix formulation and are well motivated. This work has a very broad impact, and further practical developments on efficient implementations could enable its widespread adoption among practitioners.", "essential_references_not_discussed": "Other than [1], the discussions are relatively complete.\n[1] Rethinking Attention with Performers. ICLR 2021", "other_strengths_and_weaknesses": "One of the main limitations is not seeing more quantitative results, larger models etc.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_flag": "Flag this paper for an ethics review."}}, {"review_id": "ixBhCzc186", "reviewer": "Reviewer_z3b6", "summary": "--- increased score from 3 to 4 after comment from authors ---\n\nThis work proposes a new way to parameterize the query-key operation in attention matrices. When the key and query matrices are low-rank, as it is the case in the vanilla transformer architecture, high-dimensional input data might suffer from a low-rank bottleneck leading to low performance. While introducing full-rank matrices would solve this problem, this increases the number of flops needed to run the model. Therefore the authors propose a new version of parameterizing this operation via highly structured high-rank matrices. This not only increases efficiency compared to the dense matrices and allows representing high-rank operations, but also allows taking in distance between tokens as a parameter to model locality.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful questions and supportive feedback!\n\n**Hyperparameter and Error Bars.** \n\nIn all experiments, for both our methods and baselines, we sweep over the learning rate while keeping other hyperparameters fixed. In all figures, including Fig 3 and Fig 4, we plot only the best learning rate for visual clarity. Since our code derives from the NanoGPT codebase, the other hyperparameters (weight decay, Adam betas, etc.) are well-tuned for standard attention already. We did not have the compute budget to sweep multiple random seeds for our larger experiments, though we can add error bars for the smaller-scale experiments to our revision.\n\nThe choice of learning rate is further discussed in Appendix E.1. Our implementation adopts muP, a method for stable hyperparameter transfer across model width. Please refer to the “Larger Models and Datasets” section in our response to Reviewer 1ojm for a figure that plots loss against various learning rates. We hope that this helps address your concern that we indeed pick the best learning rate for both the baseline and our method.\n\n**Response to Review’s Questions.** \n\n### Q1\nYou can refer to the `compute_model_FLOPs` function in `src/utils.py` provided in our codebase. We use FlopCountAnalysis from the fvcore library. This function traces the forward pass of the model and records the total number of FLOPs for all the operations. To compute the non-embedding FLOPs, we subtract the FLOP count of the token embedding and language modeling head (the first and last linear layer)\n\nThe MLP layers have weight matrices of size 4D x D, but the attention layers have many heads. It would be too expensive for *each* head to have a D x D weight matrix. For example, Llama 3, 70B has H = 64 heads, D = 8192, and 80 layers. If each head used an unstructured DxD weight matrix instead of the D x (D/H) matrices $W_Q$ and $W_K$, it would have 330 billion more parameters.\n\n### Q2:\nWe kindly refer you to our response to Reviewer zR4a. \n\n### Q3:\nPrevious methods were for linear layers of the form $X \\mapsto AX$. We use it for $x \\mapsto X W_Q W_K^\\top X^\\top$. That is, previous work on replacing linear layers with structured BTT could not account for the fact that a structured matrix is *already* being used in standard attention, since $W_Q$ and $W_K$ function not as two separate linear transformations, but a single, low-rank (bi-)linear transformation.\n\n### Q4:\n- Like standard attention, our proposed methods use multiple attention heads in each attention layer. Our modifications to standard attention are applied separately to each head. In Fig 3, we compare our method to standard attention with 1 head as a baseline since 1-head attention is full rank like our version, but that isn’t our proposed method. We don’t explicitly test for induction heads, but we think our methods can implement them as well as standard attention, especially since they succeed at language modeling and in-context linear regression.\n- There is some work on the tradeoff between number of heads and rank. Theoretical: Amsel (https://iclr.cc/virtual/2025/poster/27747) and Sanford, “Representational Strengths and Limitations of Transformers” (https://openreview.net/pdf?id=36DxONZ9bA). Empirical: Bhojanapalli (https://dl.acm.org/doi/10.5555/3524938.3525019) and Appendices D.4 and E.2 of the muP paper https://arxiv.org/pdf/2203.03466.\n\n### Q5\nLet us fix $a=b=c=d=\\sqrt{D}$, as in our experiments. We are now free to set s to anything between 1 and $\\sqrt{D}$ to trade off efficiency for expressivity. (Because of the cited result, setting $s > \\sqrt{D}$ would be pointless.) We chose to set s = 1 or 2 to maximize speed, and we find that the model is still expressive enough to perform well (even better than low rank attention). In fact, when s=1, BTT matrix is exactly a Monarch matrix (https://arxiv.org/abs/2204.00595), which is an expressive matrix class already.\n\n### Q6\nThe 8-head and 1-head models we test have exactly the same compute intensity. This is because the 1-head model has an 8x larger head dimension, due to the Hr=d rule.\n\n### Q7\nWe mean to say that each point on this graph corresponds to a model that was trained for the same number of steps. We have now clarified this point in the caption.\n\n### Q8\nWe are currently using multiple heads of BTT or MLR attention. We generally compare our attention to standard attention with the same number of heads. In Fig 3, we additionally compare to standard attention with one head. We explain this more clearly in the revision.\n\n**Writing and Style**\n\nThank you for your suggestions about enlarging figure labels and including citing some standard claims. We have incorporated them in our revised manuscript.\n\nWe made a significant effort to address your questions, including paper edits which we feel improve the clarity of our paper. We would appreciate it if you would consider raising your score in light of our response. Do you have any additional questions we can address?", "claims_and_evidence": "The claims on the parameterizations themselves are valid, i.e. how they relate to one another and their parameter counts.\n\nThe evidence that the new parameterizations are more efficient is shown in experiments on language, in context learning and regression. However, it is unclear from the manuscript how the training time, hyperparameter search and final performance, as well as the number of heads influence one another. For practicioners, it would be useful to obtain a deeper understanding of there links, to be able to decide in which cases to use this architectural component specifically.", "methods_and_evaluation_criteria": "see above.", "theoretical_claims": "There is no strong theoretical claims except the parameter efficency that are to check, the inclusion of the related matrix families seem ok.", "experimental_designs_or_analyses": "It would be useful to include error bars in all experimental plots over several runs to estimate the variance and significance of the differences.", "relation_to_broader_scientific_literature": "The work examines attention mechanism closely, which is itself being studied widely due to the high frcation time and compute constraints it ensues. Next to the other models of decreasing computational efficiency this is another valid and interesting approach.", "essential_references_not_discussed": "Two factors that might be discussed more in depth are the following:\n\n(Q1) Since the authors mainly advocate that their approach is more efficient, it would be useful to know how exactly the flops and \"non-embedding compute\" are calculated. Also it would be useful to know how this approach impacts parallelization. Similarly, when the authors mention on page 5, \"DxD matrix is prohibitively expansive\", I would like to understand better why. After all the rule of thumb is that the hidden layers have size 4D - so the computation there is even more expensive. What is the difference between the two settings that motivates your statement?\n\n(Q2) The statements about locality never consider the fact that most input data is decorated with positional information via positional embedddings. It seems that this already adds a sort of positional bias to the matrix when the positional embeddings are relational and their outer product represents a structured matrix. How does this relate to your approach?\n\n(Q3) You mention on page 4 that the BTT matrix has been used as a replacement for linear layers in neural networks. This is elaborated on in section 6, first paragraph. Can you make the special properties of the query-key-matrix more explict, that warrant a special examination, rather than simply treating it as a linear layer and replacing it with those previous methods?\n\n(Q4) Is there work that investigates the trade-off between the number of heads and the attention rank?This would seem related to the efficiency discussion, as well as your experiments where you compare 1- and 8-head models. Some known mechanisms, such as induction heads [1] expllicitly use several heads to execute different functions, it is unclear how your proposal would be able to implement them.\n\n[1]Olsson, Catherine, et al. \"In-context learning and induction heads.\" arXiv preprint arXiv:2209.11895 (2022).", "other_strengths_and_weaknesses": ".", "other_comments_or_suggestions": "The figure lables are quite small compared to the text.\n\nEven though at this point in time some elements of the introduction seem like universal truths, it would be correct to cite the appropriate related work when mentioned, i.e.\n- L.43: the attention mechanism/transformer \n- L.47: transformers are being used as general purpose tools\n- L.48: Transformers have specific inductive biases\n- L.14: a large ongoing research effort ... for long and big models\n- L.34: lacks a bias for prioritising local interactions - how is this related to positional encodings? - why is the locality bias from them not enough?", "questions_for_authors": "See all questions (QX) before in the relevant sections.\n\n(Q5) On page 5, left first paragraph you choose $s=1$ or $s=2$. I might have misunderstood the previous section, but does this imply that the largest matrix you can fully represent for $s=1$ is a 1x1 matrix? Could you elaborate on this together with the practical values you choose in your experiments? Or is there a typo in the $a=b=c=d=s=\\sqrt{d}$ constraint?\n\n(Q6) Why do the 8- and 1-head attention structures in Fig.3b) have the same points on the x-axis? Should the 8-head one not be more compute intense? Maybe it woud be helpful to clarify which settings you are using here in the caption.\n\n(Q7) In Figure 3a), what do you mean by \"controlling for the number of training steps\"? \n\n(Q8) This might be me, but can I use several heads of MLR and BTT, or combine them? Would it make sense to compare the 8-head attention to the 8-head setting of the structured matrices? \n\nThe responses to the questions here would improve the presentation and understanding of the paper, specifically to make the efficiency improvements more clear empirically - this would strengthen the main claims of the paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cmODsg5lZo", "reviewer": "Reviewer_zR4a", "summary": "--- score update to 3 (weak accept) from 2 (weak reject) after the rebuttal ---\n\nThe paper consider the Score function in attention computation. This is one core component of modern LLMs. The common setting for multi-head attention is that each head has a \"low-rank\" bottleneck, as the product of key and query projection matrices ($W_Q W_K^\\top$) is a low-rank matrix. Prior work has shown that this precludes learning some functions when doing in-context regression.\nFurthermore, regular attention uses the full context which is computationally expensive, relative positions are encoded via positional encodings and are not explicitly treated in the attention itself anymore. An existing proposal to make the attention more efficient is sliding-window attention.\n\nThe paper proposes families of structured matrices to parametrize $W_Q W_K^\\top$ in ways that are still computationally efficient but lead to full rank matrices.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We respond to your questions below. \n\n**Relationship with Positional Encoding and RoPE.**\n\nIt is certainly true that many modern transformers use positional encoding schemes that capture relative position. For instance, RoPE encodes positional information by rotating query and key vectors by an angle proportional to the token’s index in the sequence. This can help the scoring function treat neighboring tokens alike. So in a way, RoPE does have a “locality bias”, or at least a \"locality awareness\".\n\nBut the goal of our MLR attention is different. We are not trying to simply encode relative positional information or boost the attention score between neighboring tokens. Instead, we change the computational cost of the scoring function based on the tokens’ positions. Put differently, standard attention (with or without RoPE) uses the same query/key dimension for all pairs of tokens. We effectively use a smaller query/key dimension for tokens that are far apart. This saves FLOPs compared to standard attention, while still allowing high quality attention scores for neighboring tokens. Thus our MLR attention is compatible and complementary with existing positional embedding schemes, including RoPE.\n\nWe acknowledge that our choice of the term “locality bias” in the context of MLR attention may have obscured our meaning. In our revised manuscript, we instead call it \"distance-dependent compute allocation\".\n\n**Do we \"aim to be better (quality) or more efficient\"?**\n1) Bilinear BTT improves expressive power by increasing the effective query/key dimension of each head (“the rank”). However, it costs more parameters and compute than standard attention (see Table 1). That is, it increases “quality” at the possible expense of efficiency. Figure 3b shows that this trade is worth it.\n2) MLR attention improves efficiency by decreasing the effective query/key dimension for pairs of tokens that are far apart in the sequence. That is, it saves FLOPs at the expense of expressiveness. Because real-world sequence data has local structure, this reduction in expressive power probably does not matter.\n\n**Grouped Query Attention (GQA).** \nIn GQA, the same KV transformation is shared across several heads. Our methods can be thought of in terms of query and key transformations too; see Appendix B.2 and the discussion under Eq 9. Thus, our approach is fully compatible with GQA. We include a detailed discussion of GQA in our revision.\n\n**Explanation of how Equation 9 gets transformed into S of the form of equation 4.**\nClaim: The $j, j’$ entry of S equals the right hand side of Eq 9.\nProof. For simplicity, assume there are just two levels, so Eq 9 reduces to the formula preceding it on line 246.\nDivide $X = \\begin{bmatrix}X_1 \\\\\\\\ X_2\\end{bmatrix}$ into blocks. $X_1$ is the first half of the sequence and $X_2$ corresponds to the second half. Divide $W_Q = \\begin{bmatrix}L_1 & L_2\\end{bmatrix}$ (line 270). Thus $$Q = X W_Q = \\begin{bmatrix}X_1 L_1 & X_1 L_2 \\\\\\\\ X_2 L_1 & X_2 L_2\\end{bmatrix}$$\nNow divide $Q$ into 3 named blocks according to Figure 1b:\n- $Q_{11} = XL_1$\n- $Q_{21} = X_1 L_2$\n- $Q_{22} = X_2 L_2$\n\nAnalogously\n- $K_{11} = XR_1$\n- $K_{21} = X_1 R_2$\n- $K_{22} = X_2 R_2$\n\nFinally, plug the above into the definition of S on line 268:\n$$S = Q_{11}K_{11}^\\top + \\begin{bmatrix}Q_{21}K_{21}^\\top & 0 \\\\\\\\ 0 & Q_{22} K_{22}^\\top\\end{bmatrix} = XL_1R_1^\\top X^\\top + \\begin{bmatrix}X_1 L_2 R_2 ^\\top X_1^\\top & 0 \\\\\\\\ 0 &X_2 L_2 R_2^\\top X_2^\\top\\end{bmatrix}$$\n\n\nConsider the $j, j'$ entry of $S$. If they’re in different blocks, then it’s $x_j^\\top L_1 R_1^\\top x_j$. If they're in the same block, it's $x_j^\\top L_1 R_1^\\top x_j + x_j^\\top L_2 R_2^\\top x_j$.\n\nWe added this derivation to the appendix.\n\n**Is there a maximum \"reach\" defined by the maximum level?**\n \nThe first level is global, so all tokens interact. In the notation of Eq 9, this is because *all* pairs of tokens $j, j’$ have $d(j, j’) \\geq 1$.\n\n**Does Attention with 1 Head use the full rank D×D matrix?** \n\nYes.\n\n**Does our attention always use a single head as well?** \n\nWe use multiple heads. E.g., in Fig 3, our bilinear models use 8 heads.\n\n**How did we choose to report D=128?** \n\nIn all four subplots of Figure 9, standard 1-head attention and BilinearBTT (both full rank) converge well before standard H=8 attention (which is low rank). This is the only claim we make with this figure or with Fig 3a. We picked d=128 because it was the largest and most realistic. In Fig 3b, we further show that Bilinear BTT is better than 1-head attention because it trains faster in terms of FLOPs.\n\nThank you again for your detailed feedback and your questions. We made a significant effort to address your questions, and we would appreciate it if you would consider raising your score in light of our response. Do you have any additional questions we can address?", "claims_and_evidence": "The paper claims to a/ solve the problem encountered with low-rank approximations for problems that intrinsically require high-rank. b/ improve language modelling tasks by including a locality bias. \n\nThe claims are not made very precise throughout the paper (no formal statements) and the empirical evidence seems a bit weak.\n\nAlso the claim \"attention uses the same scoring function for all input pairs without imposing a locality bias\" in the abstract seems misleading to me. Early models used positional encoding with the embedding, whilst Llama models and many others use ROPE in each attention layer. Hence, it is certainly the case that the attention use some relative information about the the tokens.\n\nOverall it is a bit unclear whether they aim to be better (quality) or more efficient. Or pareto-better.", "methods_and_evaluation_criteria": "The benchmark with the in context regression seems a bit limited. \n\nFor the language modelling task, I would expect some discussion of the positional encoding.", "theoretical_claims": "The only theoretical claim is that MLBTC (Definition 1) subsumes MLR and BTT matrices. This seems correct.", "experimental_designs_or_analyses": "see above.", "relation_to_broader_scientific_literature": "As mentioned above I am missing a discussion of positional encodings. Furthermore, modern LLMs often use GroupedQueryAttention. I think for larger impact those should be discussed.", "essential_references_not_discussed": "see above", "other_strengths_and_weaknesses": "--", "other_comments_or_suggestions": "Overall, I am somewhat insecure about the impact of the paper. I don't see a large theoretical contribution and the empirical results are a bit limited.", "questions_for_authors": "### 1. I am struggling a bit to follow the Locality Bias discussion in section 3.4.\n\nCould you give me a bit more explanation of how Equation 9 get's transformed into sth of the form of equation 4?\n\nIn particular (9) requires some explicit handling of the relative indices, whereas 4 does not. \n\nAlso in (9) my understanding is that there is a maximum \"reach\" defined by the maximum level. This would mean that it cannot consume the full context and is actually more like a sliding window.\n\n---\n### Figure 3:\nAm I understanding correctly that Attention with 1Head uses the full rank $D \\times D$ matrix? And yours is always using a single head as well, right? \nFrom an expressivity perspective, 1 Head Attention should be strictly better that BTT. Are you hence making a claim about training dynamics?? Also in Appendix Figure 9, the ordering seems not consistent. How did you choose to report D=128 in the main document?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: The paper proposes a method for inducing different inductive biases in attention-based models by designing attention alternatives with structured families matrices. It targets two perceived weaknesses of standard attention - the "low rank bottleneck" caused by low head dimension used as inner dimension, and the uniform allocation of computational resources to pairs of tokens irrespective of their distance in the sequences. It addresses them with two families of structured matrices, block tensor train (BTT) and multi-level low rank (MLR). 

Reviewers found the problem and approach interesting and well-motivated. The empirical results were a focus of some criticism, as the experiments were deemed only moderate in scale, and the reported performance measures were only combinatorial (FLOP counts) rather than wall-clock times. This leaves some lingering doubts about the viability of the method in practice and its potential impact. However, the innovation of the paper and the preliminary empirical results were deemed sufficiently strong by the reviewers to accept the paper. 

Several points were clarified and expanded in the rebuttals, which the authors are encouraged to include in the final version.
Decision: Accept (poster)