Customizing the Inductive Biases of Softmax Attention using Structured Matrices

Thank you for your constructive feedback. We respond to your questions below.

Relationship with Positional Encoding and RoPE.

It 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".

But 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.

We 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".

Do we "aim to be better (quality) or more efficient"?

1. 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.
2. 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.

Grouped Query Attention (GQA). In 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.

Explanation of how Equation 9 gets transformed into S of the form of equation 4. Claim: The $j, j’$ entry of S equals the right hand side of Eq 9. Proof. For simplicity, assume there are just two levels, so Eq 9 reduces to the formula preceding it on line 246. Divide $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}$ Now divide $Q$ into 3 named blocks according to Figure 1b:

• $Q_{11} = XL_1$
• $Q_{21} = X_1 L_2$
• $Q_{22} = X_2 L_2$

Analogously

• $K_{11} = XR_1$
• $K_{21} = X_1 R_2$
• $K_{22} = X_2 R_2$

Finally, plug the above into the definition of S on line 268: $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}$

Consider 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$.

We added this derivation to the appendix.

Is there a maximum "reach" defined by the maximum level?

The 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$.

Does Attention with 1 Head use the full rank D×D matrix?

Yes.

Does our attention always use a single head as well?

We use multiple heads. E.g., in Fig 3, our bilinear models use 8 heads.

How did we choose to report D=128?

In 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.

Thank 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?

Thank you for your thoughtful questions and supportive feedback!

Hyperparameter and Error Bars.

In 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.

The 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.

Response to Review’s Questions.

Q1

You 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)

The 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.

Q2:

We kindly refer you to our response to Reviewer zR4a.

Q3:

Previous 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.

Q4:

• 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.
• 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.

Q5

Let 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.

Q6

The 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.

Q7

We 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.

Q8

We 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.

Writing and Style

Thank you for your suggestions about enlarging figure labels and including citing some standard claims. We have incorporated them in our revised manuscript.

We 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?

We thank you for your supportive feedback. We address your comments below.

Larger Models and Datasets:

We 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.

In 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.

Efficient Implementation.

As 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).

Rethinking Attention with Performers. ICLR 2021

Thank you for pointing out this paper! We have incorporated a discussion of it in the revised manuscript.

We 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.

Thank 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.

Thank you for your thoughtful and positive feedback. We address your questions below.

Contributions in Comparison to Prior Works.

Thank 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:

• 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.
• 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.
• 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.
• 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.

Writing.

Thank you for your helpful feedback regarding the writing and typos. We have updated our manuscript accordingly, alongside additional clarifications inspired by your questions.

Locality.

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.

We 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.

Thank 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.

---

Reference:

[1] Factor Fitting, Rank Allocation, and Partitioning in Multilevel Low Rank Matrices. https://arxiv.org/abs/2310.19214

[2] Monarch: Expressive Structured Matrices for Efficient and Accurate Training. https://arxiv.org/abs/2204.00595

[3] Compute Better Spent: Replacing Dense Layers with Structured Matrices. https://arxiv.org/abs/2406.06248

[4] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. https://arxiv.org/abs/2407.09941

[5] Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices. https://arxiv.org/abs/2410.02117