Quality over Quantity in Attention Layers: When Adding More Heads Hurts
Authors: Noah Amsel
Gilad Yehudai
Joan Bruna
Venue: ICLR 2025 Poster (2025)
Keywords: Learning theory
Expressive capacity
Expressive power
Transformer
Attention
Primary area: learning theory
TL;DR: We prove that attention layers with a small query/key dimension can have weak expressive power. This weakness cannot be overcome without using a massive number of attention heads.

Abstract: Attention-based mechanisms are widely used in machine learning, most prominently in transformers. However, hyperparameters such as the number of attention heads and the attention rank (i.e., the query/key dimension) are set nearly the same way in all realizations of this architecture, without theoretical justification. In this paper, we prove that the rank can have a dramatic effect on the representational capacity of attention. This effect persists even when the number of heads and the parameter count are very large. Specifically, we present a simple and natural target function based on nearest neighbor search that can be represented using a single full-rank attention head for any sequence length, but that cannot be approximated by a low-rank attention layer even on short sequences unless the number of heads is exponential in the embedding dimension. Thus, for this target function, rank is what determines an attention layer's power. We show that, for short sequences, using multiple layers allows the target to be approximated by low-rank attention; for long sequences, we conjecture that full-rank attention is necessary regardless of depth. Finally, we present experiments with standard multilayer transformers that validate our theoretical findings. They demonstrate that, because of how standard transformer implementations set the rank, increasing the number of attention heads can severely decrease accuracy on certain tasks.

Published as a conference paper at ICLR 2025
QUALITY OVER QUANTITY IN ATTENTION LAYERS: WHEN ADDING MORE HEADS HURTS
Noah Amsel', Gilad Yehudai' & Joan Bruna
'Courant Institute of Mathematical Sciences, New York University FFlatiron Institute {noah. (noah.amsel, gy2219, jb44 (444))eenuued
ABSTRACT
Attention-based mechanisms are widely used in machine learning, most promi- nently in transformers. However, hyperparameters such as the number of attention heads and the attention rank (i.e., the query/key dimension) are nearly the same way in all realizations this architecture, without theoretical justification. In this paper, we prove that the rank can have dramatic effect the representational capacity of attention. This effect persists even when the number of heads and the parameter count are very large. Specifically, target function based on nearest neighbor search that can be represented using a single full-rank attention head for any sequence length. We prove that it cannot be approximated by a low-rank attention layer even on short sequences unless the number of heads is exponential in the embedding dimension. Thus, for this tar- get function, rank is what determines an attention layer's power. We show that, for short sequences, using multiple layers allows the target to be approximated by low-rank attention; for long sequences, we conjecture that full-rank attention is necessary regardless of depth. Finally, we present experiments with standard mul- tilayer transformers that validate our theoretical findings. They demonstrate that, because of how all standard transformer implementations set the rank, increasing the number of attention heads can severely decrease accuracy certain tasks.
present simple and natural
INTRODUCTION
Attention brcettoo-bbasd architectures are ubiquitous in contemporary machine learning. The most promi- nent examples are transformers, which are invaluable tools for processing images, audio, time series, PDEs, and biological data in addition natural language. The basic transformer architecture leaves the user free to set several hyperparameters, but few of these have been carefully studied. In fact, in the thousands of papers that use this architecture, certain hyperparameters are almost always kept the same as in the original work of Vaswani et al. (2017) that introduced transformers. (See SceeAAppedixx A for a comparison.) In this paper, we study one of these hyperparameters: the attention rank.
The rank of an attention layer refers to the dimension of the query and key vectors formed by each attention head. The capacity of an attention layer is determined by its rank (r) and its number of heads (H). When the input is sequence of embedding vectors in Rd its total number of param- eters is of order dHr. Most transformers in the literature still use a small rank of between 64 and 128, even though the embedding dimension has grown dramatically from d 512 in the original transformer (Vaswani et al., 2017) to d = 8, 192 in LLaMA (Touvron et al. 2023a). It is not clear whether the expressive power of transformers weakened by maintaining a fixed rank as the di- mension is increased. In addition, nearly all work on transformers sets the number of heads to be H = dd/r (see Appendix A). In fact, this scaling so standard that is hard-coded into libraries like PyTorch (Paszke et al., 2019) and xFormers (Lefaudeux et al., 2022), fact which has likely discouraged experimentation with other scalings. The original motivation for this scaling in Vaswani et al. (001)) was simply to match the parameter count of multi-head attention with that of earlier architectures, which used one full-rank attention head (H = 1, = d). We know no a priori reason or experimental evidence that favors this scaling over any other, as the trade-offs between the rank and the number of heads are still not well uellssnders.odd. This paper raises doubts about whether the universal practice of setting & d and H
d/r is always wise choice.

---

Published as a conference paper at ICLR 2025
To demonstrate how impactful these hyperparameters be, reproduce the influential experi- ment of Garg et al. (2022). which shows that transformers can be trained to perform linear regression. This experiment is widely studied by both empiricists and theoreticians as archetype of in-context learning Von Oswald et al. 2023 Bai et al. 2023) Ahn et al. 2023 Akyürekk et al. 2023 Zhang al. 2024). The original paper set H
In Figure IJ we show that simply reducing H while
maintaining the scaling H = d/r yields significantly better accuracy without changing the number of parameters in the model. (See Appendix B.1 for details.)
The first step towards understanding such hyperparameter scalings transformers is to focus on their expressive power. Our ap-
proach is analogous to a long line of work in the theory of deep learning that has ana-
Mean Squared Error
0.25
lyzed the relative importance of width and depth in determining the expressive power of feedforward neural networks (FNNs), For FNNs, depth two suffices for univer- sal approximation (Cybenko) 1989). but greater depth may be required for efficient approximation. That is, some functions can be efficiently represented by three-
0.20
0.15
0.10
0.05
0.00
layer FNN, but cannot be represented by two-layer FNN unless its width is ex- ponentially wide in the input dimension (Eldan & Shamiri 2016;; [aniely| 2017 Safran & Shamir 2017). This shows that width and depth are not exchangeable; very large width does not compensate for the depth being too small. It is natural ask similar questions about attention- based architectures: Does using low rank (r < d) weaken the expressive power of attention? Does using large number of heads H compensate for this weakness? What is the most parameter-effiieet way to scale up our models?
2
4
8
16
Number of Heads
Figure 1: For the in-context learning task of Garg al. (2022), transformers with fewer heads (that is, higher rank) perform significantly better despite having the same number of parameters. See See[Appendi B. B de- tails.
In this paper, we study precisely these trade-offs the expressive capacity of attention layers. We present simple target function based on nearest neighbor search which can be efficiently approx- imated by a single-head, full-rank attention layer (hence with total of ((d2) parameters), but not by a low-rank attention layer. This separation in expressive power persists even if the total number of parameters satisfies Hdr >> d². This shows that rank, rather than the total number of parameters, can be the main factor influencing an attention layer's capacity. Stacking multiple layers of low-rank attention can yield a moderately-effiicee lengths. We complement these theoretical results with experiments on fully-featured transformer architectures as implemented in PyTorch.
approximation of our target too, at least for short context
1.1 OUR CONTRIBUTIONS
In Section 4 prove rank separation for representing the nearest neighbor function using multi-head attention. This function can be approximated to any accuracy using only a single full-rank head. Yet in the high-dimensional regime, at least Ω ((d/r)//g required to achieve relative squared error Moreover, in the high-
heads of rank accuracy regime (e going to zero with fixed), the required number of heads is exponential: Dexxp(d log(d/r))).
In In(eecion 5)
we establish an exponential separation in both regimes by combining multiple biased nearest neighbor functions. This combined target function can be approximated to any accuracy using polynomially-many full-rank heads, but M(exp(d r)) rank-r heads are required approximate with better than constant error.
In Section 6, we explore ways to circumvent the weakness of low-rank attention. We show that augmenting the attention architecture and adding second, non-linear layer achieve
2

---

Published as a conference paper at ICLR 2025
this using polynomially many heads, but unlike full-rank attention, such constructions may not scale to long sequence lengths.
we support our theoretical results with experiments on standard transformer architectures with multiple layers of attention and MLPs. We show that the full rank models easily learn the target to high accuracy - even recovering our main construction - but the low rank models struggle jet al. (2022), larger rank likewise yields higher accuracy with fewer parameters. Users of standard transformers may not realize that setting !/ could be much worse than !/ 1, but in some cases, it is.
In Section 7
do so. For the well-known in-context learning task of /arg
2 RELATED WORK
Theory of transformers A growing line of work has sought to provide theoretical analysis of transformers and the attention mechanism. Training dynamics, inductive biases, generalization, and in-context learning have all received significant attention. However, papers in these areas nearly
always assume that full-rank attention is used (Bietti et al.
2023 Cabannes et al. 2024 Fu et al. 2023 Zhang et al.] [024]] Jelassi al.,
2023 Sanford et al.] [2224a] Edelman al.
2022]
Bai Bai et al.,
2024 Chen et al. 2024] Deora et al. 2023 Tian et al. 2023). even though many also assume there
multiple heads. Our work provides important context for these results, showing that full-rank
models may not be good proxies for the low-rank transformers used in practice.
Expressive power of transformers Our work belongs to a body of research studying the repre- sentational capacity of transformers. Unlike other topics in transformer theory, results in this area often do apply to low-rank attention. Yun et al. (2019) prove that (exponentially deep) transformers are universal approximators even with rank one. Wei et al. (2022) and Merrill & Sabharwal (2023) show that transformers can simulate Turing machines if their size is allowed to grow with the se- quence length. Kim et al. (2022) and Kajitsuka& & Sato (223) show that transformers are capable of memorizing data. Bhattamishra et al. (2024) show that transformers can efficiently implement a ver- sion of the nearest neighbor algorithm for in-context classification of points on the sphere, but their construction uses attention that is full-rank with respect to the input dimension. Our formulation of the nearest neighbor task is slightly different and can be solved with full-rank attention almost trivially (see Fact 1). Finally, an important line of work analyzes the representational capacity of transformers using classes of formal languages, finite automata, and circuits (Hahn) 2020 Liu et al.)
2022 Hao et al. 2022; Merrill et al. 2022 Strobl et al. 2024). but it does not capture separations in capacity due to rank.
Limitations of low-rank attention Several other studies have investigated the role of the rank of the attention mechanism. [Bhojanapalli et al. (2020) present experiments that challenge the canonical H = d//r scaling. They argue that fixing d and based on the context length N and setting H independently leads to more powerful and efficient models. They also prove that full-rank attention head can produce any attention pattern from any input (for some setting of the weights), but a low- rank attention head cannot; however, this implies nothing about the representational power of low- rank attention or how it relates to the number of heads. In addition, addition,,Lkhooshertt [Lihhohhesttov et al. (2023) show (og suffices to represent any sparse attention pattern. Mahdavi et al. (2024)
that even rank ask how many input-output pairs a low-rank multi-head attention layer can exactly memorize. For their problem, is not worth setting r > N; furthermore the memorization capacity depends on rHl rather than on r or H, supporting the standard scaling. We study the more realistic and practically motivated setting of approximating a natural function over data drawn from a natural distribution.
Unlike Likhosherstov et al. (2023) and Mahdavi et al. (2024), we show that high rank is sometimes essential, irrespective of H.
The paper closest to our own is Sanford et al. (2024b), which proves two separations related to rank. First, they present a function that can be well-approximated by a single attention head if and only if its rank is sufficiently large. This result prompts the following question: can using multiple heads compensate for the weakness of low-rank attention? We answer this question in the negative. Second, they present one dimenssionl function N inputs that is impossible represent exactly unless r rHp > N, where p is the bits of precision. We extend this result in that our lower bounds apply (1) even for N = 2, (2) for infinite finite precision (3) to function approximation over a nat-
3

---

Published as a conference paper at ICLR 2025
ural distribution, not just exact representation. Additionally, our target function en stronger separation: while H 2 (((/r)) suffices in their setting, ours requires H to grow polynomially or even exponentially in d/r to overcome the weakness of low-rank attention. However, their target functions are more closely akin to the kinds of structured reasoning tasks to which transformers are often applied. In partieular, they highlight how attention is naturally suited to capturing pairwise interactions; recurrent architectures struggle to do this efficiently, while transformers struggle to capture third-order interactions.
Finally, an important reference to the dangers of low rank appears in Yang et al. (2022), which introduced the highly influential pP method for hyperparameter transfer. They find that as and r decrease, transformers start to behave badly. Thus, they set r >> d/ H in some cases. See Appendix D.4 and E.2.
Low rank compression and fine-tuning Much recent work in model compression (Lv et al. 2023 Hajimolahoseimi et al. 202T Ben Noach & Goldberg, 2020) and fine-tuning (Hu et al.) 2022) is based on the empirical observation that the weight matrices of pretrained transformers (like those of other neural networks) can be replaced or fine-tuned by diver-dimensional proxies without sacri- ficing performance, and in some cases even helping it (Sharma et al. (002)) Such results contextu- alize our work by showing that full-rank is not always better than low-rank.
Depth-width trade-offs in neural networks Many previous papers have proved separations be- tween neural networks of different depths and between neural networks and kernel methods. Several studies (Eldan & Shamir 2016 Daniely, 2017; Satran & Shamir [007]] Veentrr et al.) (002)) con- structed functions that can be approximated efficiently with 3-layer neural network, but for which 2-layer networks require the width to be exponential in the input dimension. (Telgaskky (0016) and Chatziafratis et al.] (2019) show depth separation for networks with constant input dimension and varying depths. Our lower bounds are also closely related technically to separation results between neural networks and kernel methods. Yehudai & Shamir (2019) prove that random features (or any other kernel method) cannot learn even single neuron unless the number of features magnitude of the weights is exponential in the input dimension. Kamath et al. (2020) improved on their result by removing the dependence on the magnitude of the weights. Ghorbann al. (2021) and Misi- akiewicz & Montanari (2023) study upper and lower bounds in approximating polynomials with kernel methods. They show that essentially, is necessary and sufficient for the number of features
to be exponential in the degree of the approximated polynomial. Our lower bounds are inspired by this work.
SETTING AND NOTATIONS
Attention layers A rank-r attention head is parameterized by the weight matrices Q K, V, e Rdxr. (Some authors call these Wo. WK, Wv. and Wo.) A multi-head attention layer simply the sum of H such attention heads. The input to a multi-head attention layer is a sequence of vectors ..,... ..N E Rdd called the target points and a sequence y ... UM called the source points. (Note that the name "target points" is unrelated to that of the "target function"' we wish to approximate.) If the columns of X E _dxN and Y gdxM are the target and source points, respectively, then a softmax multi head attention layer is function of the form
H
On V. X sm (X K K.Q.Y E MRd
where sm(+) computes the softmax of each column of its input; that is, for each y. it outputs a probability distribution over [N] based on the scores X K, Q,Q y E RNN. A hardmax attention layer is the same, except that the hardmax function hm(-) outputs es,, where is the index of the maximum score. Note that hardmax heads are often considered to be special case of softmax heads, since lim.+.0 sm X Q Y) = hm ( K; ) in pointwise convergence.
Above, we have described so-called cross-attention, which takes both source points and target points as input. The familiar aelf-attention layers are a special case in which the source points and target points are identical: X Y. A given multi-head attention function can be applied to any number of source or target points, since no part of this definition depends on N or M. In addition, it is invariant to permutations of the target points and equivariant to permutations of the source points.
4

---

Published as a conference paper at ICLR 2025
Generalized attention We prove lower bounds against class of functions that generalizes multi-head attention. Rather than computing the attention distribution as sm K,Q,)), we al- low any function depending y and rank-r projection of X that outputs probability distribution over [M]. In addition, we replace O V. with a single matrix Vt. E Raxd Thus, our model is
H
Vi On (K) X,
(2)
where KA Rdxr,, the function Φh R Rd A A is applied column- Y. and A N-1 the simplex. Note that the function oh, may vary between heads. Moreover, we allow V, E Rdxd be full-rank. Note that this class captures, beyond standard transformer architectures, the use of biases, additive positional encodings, and other encoding schemes like RoPE (Su et al., 2024) and ALiBi (Press al., 2022) in the attention layer. We also capture architectures from early works on attention (Bahdanau et al. 2014] Xu et al. 2015). which used feedforward networks to compute the attention scores instead of the "muliplicative" or "dot product" attention scores X T K KQY used in transformers.
Nearest neighbor function The input to the nearest neighbor function consists of a sequence of gd-1 (also denoted by X N) and source
N target points from the unit sphere Zi...., EN EN point y E gd-1
The nearest neighbor function outputs the target point that is closest to the source:
f( bNiy): == arg min 3E)
|| - y//2-
This function is analogous performing semantic search, in which the goal is to retrieve the entry or word in database or context window that most closely matches a query. This function is highly symmetric. Like multi-head attention itself, it defined for any N and is invariant permutations of the target points. It is also invariant to simultaneous orthogonal transformations of X and y. so it has no principal directions, subspaces. or scales.
Data distribution We draw target and source points uniformly from the sphere. For our lower bounds, it is convenient to assume that the target points are orthogonal. For N S d, let (S( denote the uniform distribution over the set of sequences EN E gd-1 for which + j a; J. 2r Such samples can be generated by taking the first columns of a random orthonormal matrix. Note that this is similar in essence drawing the data points independently from the unit sphere, as isotropic random vectors in high dimension are nearly orthogonal. This distribution is invariant to orthogonal transformations of X and of y-
LOW-RANK SEPARATION FOR NEAREST NEIGHBORS
In this section, we study the capacity of multi-head attention to represent the nearest-neighhor func- tion. We show a separation in representational power based on rank. The target can be represented efficiently using full-rank attention, but under the assumptions below, approximating it using low- rank attention requires a much larger model. We begin with the upper bound using a single full-rank attention head:
Fact 1 (Full-rank Efficient Approximation, Equivariant Case), For the target function from Equa-
tion (3) any é >> O, N, d E N there exist K, Q, V E Rddd such that:
[[] v) V KQ v)ll Se.
(4)
y.P........ w................. Unif(3d
The construction is straightforward. Consider for simplicity the hardmax case. Set V = KQ = I so that a.. y|l2 = 2 x._ KQ y. Then hm(X KQ y) =e,- where = arg min ||=>
a/l2 and e; is the ith standard basis vector. Note that this construction using hardmax works for any input distribution on s and any number of points N, as it represents the target function exactly. The softmax case is similar; for the formal statement see appendix AAppendix C.I This construction (or one very similar to it) is easily learned by gradient descent; see Figure 5
5

---

Published as a conference paper at ICLR 2025
We now turn to the lower bound. We show that approximating the target function with rank-r heads requires the number of heads to be large unless r number of target points to two and draw them from the distribution D2(56=1) in which they always orthogonal. Our main result establishes a strong quantitative separation between full-rank and low-rank self-attention layer, even when the total number of parameters is of the same order:
For technical convenience, we set the
Theorem 2 Low-Rank Approximation Lower Bounds, Equivariant Case). There exist universal constants c, e, C and C' such that if either of the following sets of assumptions hold:
1. High-accuracy regime:
S 3d-3, 3, E E
c
and
d+1
11 1 1...2 2220 (og)))
(5)
2. High-dimensional regime:
d2 5,
and
d 2
c/ec
1 1 2e
d
H 3
(6)
2
Then, for any choice of rank-r generalized attention heads or IR + A Vn E Rdxd Raxr the error of approximating the nearest neighbor function is bounded as follows
K₁
H
f(X:y)
Σ h=1
Vn on
KK, x,
2E.
#t 11..2~-D218"" UnUUnif(gd 4)
where f is defined as in Equation (3)
For the proof of Theorem 2 see Appendix C Intuitively, the approximation problem becomes harder as d - oo and as 6 - 0. Theorem 2) combines guarantees in two different regimes. In the first regime, the desired accuracy is small. In this case, the necessary number of heads grows expo- nentially with d - r. In the second regime, the dimension is be large. In this case, the necessary number of heads grows polynomially with d/r. Informally, both regimes show that the error is at least whenever H S (d/r)'/c
We emphasize that the data distribution is V - to the uniform product measure in Wasserstein distance, and we expect our main proof techniques to generalize to this uniform measure, as well as other rotationally invariant distributions. Additionally, while N establishing the separation, we also believe the framework should extend to the general setting of N 2, although this out of the present scope.
sufficient for our goal of
Our proof uses tools from harmonic analysis on the sphere. It is reminiscent of the original depth separation work Eldan Shamir (2016) and (2017), which also exploited the inability of ridge functions to approximate radially-symmen
targets with substantial Due to the rotational symmetry the target function, attention function, and data distribution, we can transform our problem to depend on pair of points Φ1 22 and y drawn uniformly from the sphere, rather than 31, 22 and y. Our target is essentially given by a step function of the form (x, y) + sgn(2 x y), which has a slowly decaying spectrum with respect to the appropriate basis. We construct this basis using spherical harmonics, and like them, our basis functions are organized into orthogonal subspaces based on degree l polynomials. Due to rotational symmetry, the energy of the target function is uniformly spread within each harmonic subspace. In contrast, each attention head is tied to a few principal directions given by the span of K₁, As a result, each head is spanned by only a fraction of the basis functions in each subspace. Thus, with a limited number of heads, it impossible capture substantial fraction of the energy of the target function.
energy.
We now comment on the tightness of this lower bound, focusing on the canonical setting of r = 1. In this case, our lower bound simplifies and strengthens slightly. For fixed and large d, the error of approximation
least e whenever H O (dt)()))). We can construct an upper bound for our problem by considering rank-1 heads to be random features. In Appendix C.8, we argue that we can approximate our target function in the RKHS associated with the feature map (x1 E2, y) + sgn ((m) 2227 kg y). where & and q are drawn uniformly from the unit sphere. The associated

---

Published as a conference paper at ICLR 2025
kernel integral operator diagonalizes in the same basis of tensorized spherical harmonics used to decompose the target function above, and thus the kernel ridge regression approximation can be explicitly analysed by bounding the spectral decay of the kernel. Then, via standard arguments from random feature expansions (Bach) 2017b), one can transfer the approximation guarantees from the RKHS to the random feature model, provided that H $2(d//=). Thus, for = 1 and fixed the approximation lower bound of Theorem caaptures the qualitatively correct behavior, though its precise dependence on d may not be tight.
5 EXPONENTIAL SEPARATION FOR LOW-RANK ATTENTION
In the previous section, we proved polynomial separation between low-rank and full-rank trans- formers in the constant accuracy regime (part 2 of Theorem 2D. That is, to achieve error smaller than e, where is a constant not depending on or r, the number of heads must be at least poly ($). prove stronger, exponential separation in this regime. That is, we find target
In this section, function that cannot be approximated by low-rank transformers to ((1)-rrror unless the number of heads is exp(It(d r)). This new target function is defined
2df
('.1,...,,2 EN. = Σ
(-1)² arg max (d - y11 + bo,5 iE.
Dj
(8)
where bys [t i=1 This function can be viewed as sum of polynomially many nearest
o.w.
neighbors functions, each with different bias term. Using this target function, we can show the following exponentially-strong separation between full-rank and low-rank attention:
Theorem 3. The function f" defined above satisfies the following:
1. For any N 2 and e > O, there exists a full-rank transformer T with 2dl + 2 biased attention heads such that:
(o DN
E
2N. xN.
.N. 12 S
(9)
2. For any rank-r transformer T with < there exists universal constant > such that if d.H. maxxee[] [|V&|| exp(c(d -)) then:
#1 D1 (su-1) [T(0,, y) - (x (21, 9)ll
1
(10)
40
yoA
The full proof can be found in Appendix D The above theorem states that the function f defined above can be approximated by a full-rank transformer with polynomially many heads, but rank- r transformers need exponentially-ne many heads to approximate it up constant. This result is reminiscent of the exponential separation between 2- and 3-layer networks from Eldan & Shamir (2016): [Daielyy (0117), which the target function requires polynomially many neurons to be approximated by 3-yeyrr network, but exponentially many neurons for 2-layer network.
Theorem 3 achieves a stronger separation than Theorem but uses a more complicated target func- tion. In addition, the full rank construction here requires attention to have bias term inside the soft- max. Another difference that y is drawn from a Gaussian instead of the uniformly from sphere, but in sufficiently high dimension, this difference is insignificant. Note also that E [| O(1), so our separation holds for the relative error too (see the discussion in Apppendix D.2
Remark 4 (Bound on the weights). Note that unlike Theorem 2] the bound in Theorem 3 does not apply to the number of heads per se; either H or the norm of the output weight matrix Vn, must be exponential overcome the hardness. A similar bound is found in Yehudai & Shamir (2019) which inspires our proof. In Kamath et al. (2020) the authors were able to remove the dependence on the intricate analysis involving the SQ-dimension: however, in
scale of the weights by applying our case it is not clear how to extend their technique because of the dependence on r. We conjecture that it is still possible to remove this dependence and leave it for future work.
7

---

Published as a conference paper at ICLR 2025
The proof of the lower bound (part 2 of Theorem 3) is inspired by and extends the proof technique of Yehudai & Shamir (219)). which separates kernel methods from 2-layer networks. For the intuition behind the proof, see Appendix D.I] We note that to achieve exponential separation with constant error, we used a target that contains polynomially-many nearest neighbor functions. In fact, pre- vious results (Hsu al., 202T] [Satran al., 2019) imply that such strong separation cannot be achieved with just single nearest neighbor function. For a more in-depth discussion of this issue, see Appendix D.2
6 EFFICIENT APPROXIMATION USING DEPTH
In the previous sections, we showed that a single layer of low-rank attention fails to represent the target unless the number of heads is very large. In this section, we take up the question of whether additional layers of depth can overcome this weakness. We present construction that approximates the target function (with slightly modified inputs) using two layers and only polynomially many rank-1 heads. However, we present constructions only for the case where the context length N = 2, which is also the setting of our lower bounds. We conjuncture that any construction using low-rank heads introduces an unfavorable dependence N, significant weakness compared to full-rank attention.
Our constructions are based on the strategy we call "majority voting", which we briefly describe here. Consider the case of N = 2 target points and hardmax attention. The output of each head, like the target function itself, is either æ1 Z2- A random rank-1 head is weakly correlated with the target; the probability that it outputs the correct answer is 1/2 R(1/vd). Thus, combining many such random heads together, their mode (the output with the most "votes") matches the target function with high probability. We second layer calculate the "majority vote" of the heads in the attention layer.
Standard attention mechanisms make it difficult to count the number of votes each target point received-or even to remember what the target points æ; and 22 were-since the next layer gets only a linear combination of them with unknown coefficients. Therefore, we slightly modify the attention layer to facilitate the majority voting strategy. We concatenate labels to the vectors that allow us to count how many times 21 and Ø2 appear in the sum. We then use a second layer of attention to look up the full vector corresponding to the majority label. This labeling can be implemented by concatenating positional encodings to the input points. That is, instead of inputting
II.... EN E gd-1 to the transformer, now input
x1
ZN for b. E R A linear bN
b
transformation can be used to map the output of this (d + e) dimmenionall transformer back Rdd Note that our target function is permutation-invariann,, so the order of the points is irrelevant to the task hand. Thus, these concatenated "positional encodings" function more like a modification to the architecture. They provide extra input dimensions that serve as scratch space in which the model can perform discrete operations like counting and indexing without corrupting the input data. Also note that, because they change the dimension of the inputs and of the transformer, these concatenated positional encodings are different from the positional encodings used in practice such as RoPE (Su ct al. 2024) and ALiBi (Press et al. 2022), which are included under our framework of generalized attention.
The following theorem describes majority voting construction using a two-layer transformer with rank-1 heads, concatenated positional encodings, and self-attention. Because we use self- attention, the source and target points are now the same. For the precise definition of this model and the proof of this theorem, see (Appendix E. and Appendix F.4
Theorem 5. There exist universal constants Cı, C2 such that for all d > CI. and E (0, and H 2 C2 & there exists a 2-layer, rank-1 transformer T with H heads and 2-dimensional positional encodings (as defined in Defimition 33) for which
202 2/ Unif(S (x [2;; g) T ([æ)
22 2)11 S.
(11)
One might wonder whether the concatenated positional encodings necessary to make this con- struction work, especially since they break permutation invariance in order to represent a permu- tation invariant target. In Appendix E. we present an alternative construction Theorem 32) that is
8

---

Published as a conference paper at ICLR 2025
Layers
Layers
Layers
100
Mean Squared Error
10-1
mm
10-2
Params per layer
d²
d
10-3
d
2
8
16
32
64
2
4
16
32
64
1
4
8
16
32
64
Rank
Figure 2: Standard transformers trained on the farthest neighbor function. The dimension d = 64 and the number of imput points is N 16. Line shows best of five runs (see (Appendix for the spread). Across different numbers of layers, high-rank models significantly outperform low-rank models that have the same number of parameters.
permutation invariant. However, it modifies the architecture by concatenating the outputs of the at- tention heads together instead of summing them, and it passes the result (which has dimension dl)) to an MLP layer instead second attention layer.
Although these constructions assume N = 2 source points, it seems feasible to generalize them to larger N. However, the major drawback of such a generalization that the size of the transformer will depend on N. Even the simple step of calculating the majority between N possible terms does not seem to be possible without at least a linear dependence on N. On the other hand, Fact I shows that the target function can be approximated for any using a single full rank attention. We conjecture that such dependence on N is necessary when using low-rank attention:
Conjecture 6. There is no multi-layer transformer (with fixed size and weight matrices) of rank < d that approximates the target of Equation (3) for all N.
That is, while it may be possible to construct a transformer that approximates the target for a given fixed N (as we do above), conjecture that there is such construction that is independent of N. Proving or refuting the above conjecture would have very different implications. A counterexample would mean that the the weakness of low-rank can be compensated by depth, and thus the rank does not play a decisive role in the expressive power of multi- transformers. A proof would show that, even in the multi-layer case, low-rank attention is fundamentally weaker than high-rank attention.
7 EXPERIMENTS
In this section. we complement our theoretical results with experiments on more realistic archi- tectures []] We train fully-featured transformer encoders-wwiie include multiple layers of self- attention, MLPs, skip connections, and normalization Relaxing the assumptions our theory, we draw N input points uniformly and i.i.d. from dd-1 without constraining them be orthogonal, and make no distinction between source and target points. To test our predictions, we modify the attention layers by allowing the number heads to be different from the standard setting H d/r.
our nearest neighbor target function.
Our first experiment studies the influence of rank on the performance of the transformer for various numbers of heads (H) and layers (L). We fix the dimension d 64 and the number of points N 16. We use three different scaling rules for the number of heads: H = d /r for E {1, 1.5, 2};
in other words, set number of parameters per attention layer be of order rdH d In this
Our code is available wttptttpt/////// https:://itthbb.com
httpaaaeeeaattttttttaoooo formers

---

Published as a conference paper at ICLR 2025
experiment, we use no positional encodings. Figure 2 plots the results, showing the best of five runs for each setting. Each line corresponds to one choice of с (that is, one choice of scaling rule for the number of heads) so the number of parameters per layer is constant along each line. When I = 1, the results suggest that using full-rank (r = d 64) is necessary and sufficient to learn the target function accurately; even d²//2 /2 heads of rank d/2 fails. For I. > 1, the trade-off between rank and accuracy is more favorable, but full-rank attention still significantly outperforms low-rank attention, even using fewer parameters. The best low-rank model we trained (L 5, c 2, r = 32) performs no better than the worst full-rank model (L = 1,c = 1. = 64) despite having 80x fewer parameters in its attention layers. (Here we are excluding the models with L 5, c 1, which seem to suffer from optimization difficulties on this problem.) In short, standard transformer with H = (full-rank) performs much better on this task than one with even slightly larger H (lower rank).
In Appendix B.2 we present additional experiments exploring the solution learned by full rank heads (cf. Fact the use of positional encodings (cf. Theorem 5), and the role of the number of points N. In each case, the empirical results support the assumptions and conclusions our theory. We also give full details of our models and training procedure.
8 CONCLUSIONS AND LIMITATIONS
In this paper, we have investigated the role of rank in attention mechanisms. We question the nearly universal practice of using low-rank attention and setting the number of heads according to H d/r. We show that for a simple and natural target function inspired by semantic search, low-rank attention is fundamentally weaker than full-rank attention, even when H >> d/r. We demonstrate this strict separation between the low-rank and high-rank regimes both theoretically, by proving hardness of approximation in the shallow setting, and empirically, through experiments with standard multilayer transformers. Our results suggest that using a larger rank can improve the expressivity and parameter-efficieney of attention.
That said, our theoretical analysis is inherently limited to the study of shallow transformers, and the results of Section 6jillustrate how adding depth may overcome the limitations of low-rank self- attention in some cases. However, we hope that our findings will motivate theoreticians and prac- titioners to more carefully consider the settings and scalings of transformer hyperparameters. In particular, they show that if theoreticians assume attention full-rank, their analyses may fail to accurately describe the transformers used in practice. Moreover, increasing the rank or number of heads beyond their standard settings may have practical benefits that allow SOTA models to be smaller overall. Such benefits may have been obscured by overreliance on heuristics like H d/r, which prevent us from isolating the effects of any one hyperparameter. Clearly, much remains to be understood about the successes and failure modes of attention- based architectures.
Our findings suggest several promising directions for future work. The basic transformer architec- ture of Vaswani et al. 2017) allows the user to set variety of hyperparameters. However, trans- formers have been scaled almost exclusively by increasing the embedding dimension and number of layers, without significantly changing the other hyperparameters (see [Table I). While considerable prior work has studied scaling laws for the dimension and number of layers, we believe that fu- ture research should also consider the other hyperparameters and seek to understand the trade-offs. dependencies, and scaling laws that govern them. Here, we focus on the query/key rank and its relationship to the number of heads, but the depth and width of the MLPs and the value/output rank are also of interest.
Additionally, the rotational invariance of the input data distribution is instrumental in establishing our lower bounds. Given the inherently discrete nature of text- text-based transformers, natural question is to understand how to generalize our techniques beyond the iaationally-invariinn setting. Another direction for future work is to understand the relationship between the rank and the context length. Focusing on the N = 2 case suffices for us to prove rank separation, but we believe a similar result should hold at least for all N < d: Figure 7 provides preliminary experimental evidence. Understanding the N > 2 case may also help address a final open question: What is the relationship between rank and depth? In particular, does Conjecture holld
10

---

Published as a conference paper at ICLR 2025
Acknowledgements: This work was partially supported by the Alfred P. Sloan Foundation, and awards NSF RI-1816753, NSF CAREER CIF 1845360, NSF CHS-1901091 and NSF MML-MoDL 2134216. We thank Ohad Shamir for useful discussions while this work was being completed.
Reproducibility Statement: Assumptions of all our theoretical results described in the main text, and complete proofs are given in Appendices C to E Details of all experiments are given in Appendix BJ and our source code included in the supplementary material and available at aittpps: //github. com/ wwwa///ohhnseee/tteeee synthetic.
ion- formers All data used in our experiments in
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023.
Curran Associates Inc.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
net forum?
id 40-08
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-66 681, 2017a.
Francis Bach. On the equivalence between kernel quadrature rules and random feature expansions. Journal of Machine Learning Research, 18(2)):1338 2017b.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. CoRR abs/1409.0443,, 2014. doi: lbss/14990073..
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and Levine (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 57125-57211. Curran Associates, Inc., files/paperr2023/ Canerr-ooofeeeeeee
2023. URL
File
Matan Ben Noach and Yoav Goldberg. Compressing pre-trained language models by matrix decom- position. In Kam-Fai Wong, Kevin Knight, and Hua Wu (eds.), Proceedings of the Ist Conference of the Asia- Asia-Pacific Chapter of the Association for Computational Linguistics and the 10th Inter- national Joins Conference on Natural Language Processing, pp. 384-889, Suzhou, China, De- cember 2020. Association for Computational Linguistics. URL orq/20200 2ac2-maaclmall. 88
Satwik Bhattamishra, Michael Hahn, Phil Blunsom, and Varun Kanade. Separations in the rep- resentational capabilities of transformers and recurrent architectures, 2024. URL https:
//
09347
Srinadh Bhojanapalli, Chulhce Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Low- rank bottleneck in multi-head attention models. In International conference machine learning, pp. 864-873 PMLR, 2020.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bot- tou. Birth of a transformer: A memory viewpoint. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural In- formation Processing Systems, volume 36, pp. 1560-1588. Curran Associates, Inc., 2023. URL httpt: htt.a///s////poccc........rao... F1ee33577aaaaa3eeeeettteeeeeepep

---

Published as a conference paper at ICLR 2025
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Hemighan, Rewon Child, Aditya Ramesh, Daniel M. Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Vivien Cabannes, Elvis Dohmatob, and Alberto Bietti. Scaling laws for associative memories. In The Twelfth International Conference on Learning Representations, 2024. URL https: / //openreview.co http://ww.etttooommo2.ccc/22///a
Vaggos Chatziafratis, Sai Ganesh Nagarajan, loannis Panageas, and Xiao Wang. Depth-wiide width trade- offs for relu networks via sharkovsky's 's theorem. arXiv preprint arXiv: 19120439,, 2019.
Siyu Chen, Heejune Sheen, Tianhao Wang. and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint wrX1: 2024.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sashank Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Bar- ret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov. Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Weier-Hellstem Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: scaling language modeling with pathways. J. Mach. Learn. Res., 24(1), mar 2024. ISSN 1532-4435.
George Cybenko. Approximation by superpositions a sigmoidal function. Mahhmmttic of control, signals and systems, 2(4):303-314, 1989.
Amit Daniely. Depth separation for neural networks. In Conference on Learning Theory, pp. 690- 696. PMLR, 2017.
Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-hear attention. arXiv preprint UrXiv:2310.12600, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10. 10.18633/1/N19-111 1423. URL http:://cclanthlllogggoooooo///s
N19-1423
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //ppezeeviww
Benjamin L Edelman, Surbhi Goel, Sham Kakade. and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pp. 5793 18 PMLR, 2022.
Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Conference learning theory, pp. 907-940. PMLR, 2016.
12

---

Published as a conference paper at ICLR 2025
Christopher Frye and Costas J Efthimiou. Spherical harmonies in p dimensions. arXiv preprint 1rXiv:1205.3548, 2012.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. In seirth-eeventh Conference Neural Information Processing net/ forum? wd=wxGGuZDSJR
Systems, 2023. httph::t/pppeeeeeeeeee
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can trans- formers learn in-context? a case study of simple function classes. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 30583-30598. Curran Associates, Inc., 2022. URL htct:://procceddnsssneeerrr
httpaa/peeiieeeeppee//22222 Fi aaleee22880bbellaaaeceeseeeeeee0
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two- layers neural networks in high dimension. The Annals of Statistics, 49(2):1029 1054, 2021. doi: 10.1214420-ASS1999 URL httpt:://///ooo org/ /00 1214 /0 AOS1990
Yuan Gong, Yu-An Chung, and James Glass. AST: Audio Spectrogram Transformer. In Proc. Interspeech 2021, pp. 571-575, 2021. doi: 10.24457/nterspechh.2222 2021-666..
Dirk Groeneveld, Iz Beltagy. Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402. 00838, 2024.
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 1:1,6-171, 2020.
Habib Hajimolahoseini, Mehdi Rezagholizadeh, Vahid Partovinia, Marzieh Tahaei, Omar Mohamed Awad, and Yang Liu. Compressing pre-trained language models using progressive low rank de- composition. Advances in Neural Information Processing Systems, 2021.
Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention trans- formers: Perspectives from circuit complexity. Transactions of the Association for Computational Linguistics, 00:800-810, 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Thomas Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Au- relia Guy, Simon Osindero, Karén Simonyan, Erich Elsen, Oriol Vinyals, Jack Rae, and Lau- rent Sifre. An empirical analysis of compute-optimal large language model training. In S. Koyejo, Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 30016-30030. Curran Associates, Inc., 2022. URL. httt:s:://proceeednggg..oe wwwpeeeewedddddeeeeeeeee..e..eee wwtppppp..pppppppppppppppp..eeee ftteeefffleeepppeeeereeee fdf
Daniel Hsu, Clayton H Sanford, Rocco Servedio, and Emmanouil Vasileios Vlatakis Gaaaakii-Gaarakkuuiii On the approximation power of two-layer networks of random relus. In Conference on Learning Theory, pp. 2463-2461. PMLR, 2021.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022. URL https: pectps://openreviie
net forum?
iddnZev/eeFFff]
Samy Jelassi, David Brandfonbrener,, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying, arXiv preprint arXiv:2402.01032, 2024.
Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators?. 2023.
Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Probabilistic variants of dimensional and margin complexity. In Conference on Learning Theory, pp. 2236- 2262. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2025
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transform- ers. In The Eleventh International Conference on Learning Representations, 2022.
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selec- tion. Annals of statistics, pp. 1302-1338, 2000.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable trans- former modelling library. 2022.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive flexibility of self-attention matrices. Proceedings of the AAAI Conference on Artificial Intelligence, 37(7): 37(())
8773-8781, Jun. 2023. doi: 10.1609/aaii. 23777.2665. URL https: aaai.
index. php AAAI/AAAAAA http///rtttteeieeeeeeeeeeeeeeeee 26055
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamunthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210. 10749, 2022.
Xiuqing Lv, Peng Zhang. Sunzhu Li, Guobing Gan, and Yucheng Sun. LightFormer: Light-weight transformer using SVD-based weight transfer and parameter sharing. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguis- fics: ACL 2023, pp. 10325-10335, Toronto, Canada, July 2023. Association for Computational findings-acl.656. URL
Linguistics, doi: 10. 088836653//2222 htw.:/66///222...nnnnnnaaall...a 2023. w.g.ffnninggs acl. 656
orgg
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. In The Twelfth International Conference on Learning Representations, 2024. URL [RLLLttts///ppereiiieee
net forumpid--Mr
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:23 2023.
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843 856, 2022.
Theodor Misiakiewiez and Andrea Montanari. Six lectures on linearized neural networks. arXiv preprint arXiv:2308. 13431, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lever, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference Learning Representations, 2022. URL
httpss://ppenreviee... http://wwweeeeeetteooooommmddeee foormmiddR&ss@0
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, al. Improving language under- standing by generative pre- mre-training, 2018
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 8748 076888773.... PMLR, mlr, wwwppppessssp....... /v13// radford21a. 18-24 Jul 2021. URL /tttp:://proceeee html
14

---

Published as a conference paper at ICLR 2025
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXfr:2112.11446, 2021.
Itay Safran and Ohad Shamir. Depth-width tradeoffs in approximating natural functions with neural networks. In International conference on machine learning. pp. 2979-2987. PMLR, 2017.
Itay Safran, Ronen Eldan, and Ohad Shamir. Depth separations in neural networks: what is actually being separated? In Conference on Learning Theory, pp. 2664-2666. PMI.R 2019.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and loga- rithmic depth, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024b.
Ohad Shamir. Distribution-specffic hardness of learning neural networks. The Journal of Machine Learning Research, 19(1): 133--11333111,3 2018.
Pratyusha Sharma, Jordan T. Ash, and Dipendra Misra. The truth in there: Improving reasoning in language models with layer-selective rank reduction. In The Twelfth International Confer- ence on Learning Representations, 2024. URL ozX92bu8VA
net/ forum?id
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What Formal Lan- guages Can Transformers Express? A Survey. Transactions of the Association for Computa- tional Linguistics, 12:543-561, 05 2024. ISSN 2307-387X. doi: 10.1 16.1112/ a 00663. URL
https: //t::://oiiio org /10 1162/tac _aeee_____00 00663
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Matus Telgarsky. Benefits of depth in neural networks. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir (eds.). 29th Annual Conference on Learning Theory, volume 49 of Proceed- ings of Machine Learning Research, pp. 1517-1539, Columbia University, New York, New York, USA, 23-26 Jun 2016. PMLR. URL htetps::/proceeedings ml htwp:///.........//a9993 tel gargkr 6 html
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, YaGuang Li, Hongrae Lee, Huaixiu Steven Zheng, Amin Ghafouri, Marcelo Menegali, Yanping Huang, Maxim Krikun, Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Vin- cent Zhao, Yanqi Zhou, Chung-Chie Ching Chang, Igor Krivokon, Will Rusch, Mare Pickett, Pranesh Srinivasan, Laichee Man, Kathleen Meier-Hellstern, Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Soraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Diaz, Ben Hutchinson, Kristen Olson, Alejandra Molina, Erin Hoffman-John, Josh Lec, Lora Aroyo, Ravi Rajakumar, Alena Butryna. Matthew Lamm, Viktoriya Kuzmina, Joe Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Aguera-Arcas, Claire Cui, Marian Croak, Ed Chi, and Quoe Le. Lamda: Language models for dialog applications, 2022.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding train- ing dynamics and token composition in 1-layer transformer. Advances Neural Information Processing Systems, 36:71911177999 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint urXiv:2302. 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint wrX2v:2307.09288 2023b.
15

---

Published as a conference paper at ICLR 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gareett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
meppr/rriiieri.......pppeee httpi//sscccppppeeeieeeee//eeeee (ttpeee__ieeeepppe/20111111 Papear.
Luca Venturi, Samy Jelassi, Tristan Ozuch, and Joan Bruna. Depth separation beyond radial func- tions. Journal of machine learning research, 23(122):1-56, 2022.
Roman Vershynin. High- digeessiil probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 25:11-1-00 2022.
Kelvin Xu, Jimmy Ba. Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual tention. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 Proceedings of Machine Learning Research, pp. 2048 2057, Lille, France, 07-09 Jul 2015. PMLR. URL http:: http::/d///pccceeingg. ml press/v37/
xug15.hte
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer, 2022. URL
/bb//
2203.03461
Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In H. Wal- lach, H. Larochelle, Beygelzimer, F. d'Alché-Bue, F. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpe::/prrocedins.eerrr...... http:///epppeeeieeeeeeeeeeeeeeee files/paper/2018/ File/ Filee//8813222222233333322222333
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Research, 25(49):1-55, 2024. URL http://// http://ce //jrr.
htt/papprrs/25/255
1001//221110...11 html
A HYPERPARAMETERS OF TRANSFORMER
The transformer architecture Vaswani et al. 2017) leaves the user free to set the following hyper-
parameters:
The embedding dimension (d)
The number of layers (L)
The width of the MI.Ps (w)
The depth of the MI.Ps (D)
16

---

Published as a conference paper at ICLR 2025
The rank of the Wq and Wx matrices for each head (r))
The rank of the Wv and Wo matrices for each head (r2)
The number of attention heads in each layer (H)
In this paper, we consider the dimension d to be given by the domain of the target function, rather than being hyperparameter as in language modeling. As Table I shows, only d and have been significantly changed relative to the original model. For all models of which we are aware, w lies within a factor of two from Vaswani al. (2017), lies within factor of four, and and r2 are not changed at all. H has been scaled, but always according to the standard scaling (up to factor of 2).
Table 1: Hyperparameter settings popular transformer models (largest versions reported). Except for d and L, they are strikingly consistent, See text of Appendix A for notation.
Year
Model
d
T
D
r
H
2017
Attention is all you need (Vaswani et al.
2017)
512
6
4d
2
64
d/r d/r
2018
GPT, GPT-2 (Radford et al. 20088 2019
768
12
4d
2
64
r
2019
Bert-Large (Devlin
(2199
1,024
24
4d
64
d/r
2020
GPT-3 (Brown et al. 2020)
12,288
96
4d
2
128
d/r
2021
ViT-Huge (Dosovitskiy al. 2021
1,280
32
4d
80
d/r
CLIP (text encoder) (Radford et al. 2021
768
12
4d
2
64
d/r
CLIP (image encoder) (Radford et al. 2021
1,024
24
4d
64
d/r
Jurassic-1
13,824
76
4d
2
144
d/r
Gopher 280B Rae et al. 2021
16,384
80
4d
128
d/r
AST (audio) GGog et al. 2021
1,024
24
4d
2
64
d/r
LaMDA (Thoppilan et al.
2022
8,192
64
8d
2
128
2d/r
2022
Chinchilla 70B (Hoflman et al.
2022
8,192
80
4d
2
128
d/r
2023
PaLM (howdheey et al. 2024)
18,432
118
4d
2
256
r
2d/3r
(aaiaa2-Toouiei Touvron al. Llama-2
2023a b
8,192
80
8d/3
2
128
d/r
2024
OL (Groeneveld et al. 2024)
8,192
80
8d/3
2
128
d/r
B ADDITIONAL. EXPERIMENTS AND DETAILS
B. IN-CONTEXT LEARNING OF LINEAR FUNCTIONS
We reproduce striking experiment from Garg et al. (2022) that demonstrates the ability of trans- formers to perform in-context learning. In this task, the prompt is sequence of input-output pairs from randomly selected function it has not seen before, plus additional input point. The model is trained to predict the output corresponding to this final point. That is, given the sequence
3t f ....)..... UN f(æN), Equery
(12)
the desired output is (( ). In the linear version of this task, f(x) σ. The input points 21,..., N. the query point Qquery. and the target weight ur are drawn from the isotropic Gaussian distribution in 20 dimensions. The model consists of linear embedding layer mapping from 1g20 to Rd 12-layer transformer, and linear unembedding layer from Rd to IR.
We use this experiment models, specifically the consequences of using rank that is too small for the task. We run the exact code Garg et al. (2022), changing nothing the data, model, or training procedure except for the embedding dimension d and the number of heads H. As is standard, the rank is given by d/ H. 256, which is quite large relative to the intrinsic dimension of the inputs, and // 8, 32. We show that we can still achieve good performance with much smaller models, such as d = 48, but only if the rank is sufficiently large. Figure plots the performance of these models. Figure I showw the same data, but focusing only the d = 48, N = 40 case. Across different embedding dimensions, the high-rank models (H 2) significantly outperform the low-rank ones (H 8 16).
benchmark to demonstrate the importance of rank in attention based
The original paper uses a model with embedding dimension
17

---

Published as a conference paper at ICLR 2025
1.0
Embedding Dimension
Mean Squared Error
-
32
0.8
48
64
. 80
0.4
0.2
0.0
2
4
8
16
Number of Heads
Mean Squared Error
10a
Number of Heads
10"
4
10-1
5
10
20 Context Length
15
25
30
35
40
Figure 3: Performance 12-layer transformers on in-context learning of linear functions 20 dimensions. Top panel shows the enror prompt of N 40 input-output pairs for various embedding dimensions d. Bottom panel shows = 48 given various context-lengths N. Across different choices of embedding dimension and context length, performance is significantly improved by reducing the number of attention heads H. Due to the standard scaling rule = d/H. this is equivalent increasing the attention rank.
18

---

Published as a conference paper at ICLR 2025
Layers
Layers
Layers
Mean Squared Error
10"2
10-2
Params per layer
d²
d
10-3
d
2
8
16
32
64
2
4
16
32
64
1
4
8
16
32
64
Rank
Figure 4: Copy of Figure 2) with the addition shaded region showing the range over five runs. With few exceptions, these ranges quite narrow.
B.2 NEAREST NEIGHBOR TASK
In this section, we provide details about the models and training procedure used for the experiments in Section 7 and present additional experiments
nearest neighbor target function.
Model and training details We use the PyTorch implementation of transformer encoders (Pazzke ict al. 2019) with two modifications. First, we generalize the standard scaling H d/r, allowing try H d=.5 /r and H d²/r.) Second,
H to be any multiple of d/r. (In particular, replace the layer normalization with RMSNorm (Zhang Sennrich, 2019), standard choice in modern transformers (Tourron et al., 2023a (Cowddhery al. (0224) that is also better suited to our target function. We train with biases, but preliminary experiments showed that these make little difference We run each experiment on single Nvidia GPU (usually V100) for
than
few hours.
Since we are using self-attention, there is no distinction between the source and target points. The N input points are drawn uniformly and i.i.d. from gd-1 and they are not constrained to be orthogonal. We change our target function accordingly. For each input point, the target now outputs whichever of the other points is farthest from it. We output the farthest instead of the nearest point because otherwise, each point would map to itself. The loss function the average mean squared error over the N points. We do not any attention mask. In particular, themselves. Our dataset synthetic, so we train that never repeat. We train on 10 batches of size 256 each. For all experiments, with the same learning rate 0.01 and a learning rate schedule of cosine annealing with a lincar warm-up.
allow points to attend to test on a stream freshly generated samples AdamW W
Statistical consistency of main experiment Our focus in this paper is on representational capacit rather than learnability or generalization. However, for completeness, we reproduce Figure 2 with the addition of shaded region showing the range are strikingly consistent between runs. Optimization difficulties were observed for the regimes L = 3, params = d and L 5, params d2 In the first case, mostly overcame these performing three additional runs each (so eight instead of five) for the settings r = 16 and r 32.
five runs (Figure 4). In most cases the results
Full-rank solution In the full-rank case the transformer learns the target, but what representation has it learned? Figure 5 suggests that, in some cases, it is very nearly the construction of Fact
Note that biases in the key, query, and value transformations have a different role from additive positional encodings. These biases differ between heads but are constant across tokens; in contrast, the positional encod- ings differ hetween tokens but not heads. The hiases implemented by PyTorch are also slightly different from those studied in Section
19

---

Published as a conference paper at ICLR 2025
180
3000
175"
Angle btwn. KQ
2000
170"
IKQ
165
160
64
64
Num heads
Num
Figure 5: Properties of learned KQT matrices for full-rank models with layer. Boxplots show distribution over heads from five runs, each on a model which has between and 64 full-rank heads. Left panel plots Frobenius angle with the identity: arccos (( (KQ) ) /[| [[=]]][]=])
Results
show that KQ nearly equals - cf for 1000 in all cases.
10-1
MSE
Positional Encoding
10-2
None
Additive
Concatenated
2
4
8
16
32
64
Rank
Figure 6: Standard 3-layer transformers with positional encodings the the farthest neighbor task (d = 64, H 512/r, Ñ = 16). Line shows best five runs; shaded region shows range over five runs. Concatenating 6-dimensional learned absolute positional encodings to the input points im- proves performance somewhat, but does not ameliorate the low-rank bottleneck. Additive positional encodings do not improve performance because they do not increase the embedding dimension.
Recall that in Fact we use hardmax attention head with K₁ however, the farthest neighbor target function and softmax heads, so the corresponding construction K, = -el for > 1. The first panel shows the median Frobenius angle between the matrices K QT and I learned by the full-rank, single layer models in the previous experiment. This shows that K, very nearly equals -I up to a constant factor. Moreover, as the second panel shows, the norm of this matrix is large, which causes the softmax act like hardmax. Results are similar for three layer networks with a single full-rank head, but when L > I and / > 1, it seems the network learns some other, less interpretable strategy represent the target.
our experiments
Positional encodings Since our target function is sermutationiinvarinn no positional information exists in the data. However, in /Section 6 showed that concatenated positional encodings can help low-rank attention succeed when L > 1 by giving the model extra dimensions scratch space. The positional encoding schemes used in practice, like additive encodings /Vswaan et al.) (017), RoPE (Su et al. 2024) and ALiBi (Press et al.] (002), cannot be used this way, being versions of the generalized attention heads studied in this paper. In Figure 6 we experiment with positional
20

---

Published as a conference paper at ICLR 2025
1.2
Num Points
256
Mean Squared Error
1.0
64
16
0.8
4
0.6
0.4
0.2
0.0
2
16
32
64
Rank
Figure 7: Effect of the number of points (N) the difficulty of learning the farthest neighbor function. Full-rank attention learns an accurate representation across many Ns, but the performance of low-rank attention degrades N grows. Dimension is 64. All models have two layers with H d heads each. Line shows best five runs; shaded region shows range over five runs.
encodings. As expected, additive encodings fail improve performance at all. Concatenated 6- dimensional positional encodings help a little bit, since the transformer is now larger (d = 70 instead of 64) and has more parameters. However, unlike in Theorem 5] these encodings do not mitigate the low-rank bottleneck in practice, even for 3-layer transformer. Thus, the low-rank construction of Theorem is primarily of theoretical interest.
we explore how the number of input points N affects the difficulty of learning the target function. fix 64, H d²/r, and the number of layer 2. The results show that, predicted by Fact I] the full-rank heads leam the target accurately across range of N. However, the low-rank heads suffer declining accuracy as N grows. This accords with [Conjecture 6 which predicts that low-rank transformers of fixed size fail to accurately represent the target for sufficiently large N.
Role of N In Figure 7
C PROOFS FROM SECTION
In this section, prove the upper bound Fact 1 the lower bound Theorem 2 and some important properties relating to the approximation of the target by random heads.
We begin with the proof of Fact in Appendix C. In Appendix C.2 we review the basies of spherical harmonics and describe the corresponding family of ultraspherical orthogonal polynomi- als on the interval. In Appendix C.3, we construct basis for functions of pairs of points on the sphere that we will use to analyze the target and the attention mechanism. In [Appendix C.4 we show how expand the target function in this basis, proving the critical properties slow spectral decay and rotational invariance between basis elements of the same degree. In [Appendix C.5] we expand a single attention head in this basis, showing that the number of basis elements with which it is correlated is limited by the rank of the attention head. In Anppppediix C.6 use these results to obtain a lower bound on the error of approximation that depends only on certain universal constants related to the spherical harmonics, particularly the number of spherical harmonics of given degree and the coefficients of the ultraspherical expansion the sign function. In ln/ppeddi C.7] we ana- lyze this expression derive bound the necessary number of heads that depends only the
21

---

Published as a conference paper at ICLR 2025
dimension d, the rank and the error level Finally, in Appendix C.8 analyze construction that approximates the target function using random rank- heads.
C.I PROOF OF FACT
Let > O. We set V I, KQT = al for > to be chosen later. Since wt y Unif(Sd-1), for every E [I,.... NN. there exists > (which depends on e) such that for the set:
Ax:= = {111. NN. y) E (Sd 1 : + j, (2, xzz) yl 8),
(13)
have that Pr((x11,.... ...... IN.y) & As) S Note that:
arg m y) = arg max ||=; X smmaX y) a0 æ;
y/
(14)
at
where the convergence is uniform on As. and the equality follows since all the vectors are from the unit sphere. In particular, there exists 0 such that:
sup
X sm (ax y)
arg max ai
x y/|
(15)
2
(20 ..... NE EA
Combining both bounds and taking expectation over the vectors finishes the proof.
C.2 SPHERICAL HARMONICS
We begin by reviewing basic results from the theory of spherical harmonies. Let 7(+) denote the uniform distribution over gd-1 and define the inner product (s-), over L²(dd-1) as follows
(f.g), == Sga. f(z)g(z)d7((z)
(16)
polynomial H Rd IR is called harmonic and degree- homogeneous if
H 0, H(ax) a (x)
(17)
spherical harmonic of degree is the restriction harmonic homogeneous polynomial to the sphere Sd-1 That is, a function Y gd-1 + R is spherical harmonic of degree é if and only if the Rd + function defined by
- ||
(18)
|
is harmonic homogeneous polynomial of degree l. The set of spherical harmonics of degree & on gd-1 form function space F, C (Sd(Sd-)). These subspaces have the following dimensions
(Theorem 4.4 of Frye & Efthimiou (2012)):
28 + d (1 ((+d-3) l e-1
(dd, 0) == dim
(19)
The reason spherical harmonics useful is that the Fz linearly independent, and their direct N(d.f) an sum is (gd-1). That is, if (Y;)) an orthonormal basis of Fr. then 0 (Y2)
orthonormal basis of (Sd-)) with respect to (,.
For unit vector e, let denote the distribution of x when æ ~ Then for E [-1,1],
Ad-2
(2) d=3
ua(t)
(20)
Ad-1
where A&-1 is the surface area of 5d-1 (see Lemma 4.17 of Frye & Efthimiou (2012)). Define the following inner product over functions mapping [-1, + R:
(f.9) == S" f(t)g())ua(tddt
(21)
The ultraspherical polynomials P; [-1,1] 1] + IR for E N2o are defined by the following proper- ties:
22

---

Published as a conference paper at ICLR 2025
1. Pt has degree L
2. 8 l' = (Pt. ae) =0
3. =1
These polynomials form an orthogonal basis for 1, uz), which includes all bounded func- tions on [-1,1]. Moreover, they are intimately connected to the spherical harmonics. We exploit three such connections. First (Equation 4.30 Frye Efthimiou (2012)
11,11),,))
1
(22)
(dd,))
Second, the addition formula states that each ultraspherical polynomial can be expressed in terms of the spherical harmonics of the same degree and vice versa (Theorem 4.1 4.11 of Frye & Efthimiou 2012
N(d.))
Pe(x y) N(d,e) E (x)))
(23)
Finally, the Hecke-Funk formula (Theorem 4.24 of Frye Efthimiou (2012) gives the relationship between the ultraspherical expansion of L + f(() and the spherical harmonic expansion of f(x y). For any degree-l spherical harmonic Ye.
(f((z.))))a) : SSan f(x ()) Y((x) Pt)
(24)
We will make use of the ultraspherical expansion of particular functions: Definition 7. Let {ae} be the ultraspherical series for arcsin and let {ne} be the ultraspherical series for sign. That is,
Pt(t)
arcsin(l) =
Of
|| PElluu
(25)
E=0
Pe(t)
sign(t)
ne Pe
At V [-1,1]
(26)
C.3 ORTHONORMAL BASIS FOR TARGET AND ATTENTION HEADS
The goal of this section is define the orthonormal basis that we will use to analyze the (surrogate) target and attention functions. We define the input space for these functions as follows: X gd-1, x sd-1, We denote elements of this set by (2., y) or for short. For any two functions, define their tensorization by
(f 2 ()) = f(x)f(y)
(27)
We let F = @ be the uniform measure on X. We also define a feature space Ω = gd- x dd-1 and denote elements of this space by (q, k) or w. Of course, X, but since they are used in different contexts, we use separate notation for readability.
We define the feature mapping that we will use to analyze the surrogate target and attention func- tions:
Definition 8. Define the "rank-1 head" function p: X x S > {+1} by
p(z,w) == sign (xr kq y)
(28)
and the feature map linear operator T (T) - T.2 by
(Tw) / ((z
(29)
Note that Frye & Efthimiou 2012 has extra factor of A in the theorem statement. This is because they use different normalization for the spherical harmonies.
23

---

Published as a conference paper at ICLR 2025
The intuition is follows. For fixed value of w (k: q), the function p(,, w) acts like hardmax attention head with rank 1. More precisely, æ = 21 2 and V = I. then p(z, w) is the output of the head applied the souree y and targets 21 and 22. projected onto æ. Furthermore, Tu is weighted linear combination of all possible rank-1 hardmax heads.
We will construct basis using functions of the form T(Y & Y") for spherical harmonics Y and Y'. The rationale for choosing this basis is as follows. T defines a positive semidefinite operator TTT ( + (22), which is described by the following formula:
(T"Tu)(w) I. andt
[p(2, (p(z,w) ulw u(u')ddF(w")
(30)
Functions of the form Y @ Y' will turn out to be eigenfunctions of this operator. To why, must first analyze the kernel Esnr [ø(2, w),,,, which we do in the following lemma.
Lemma 9.
E znu
T n2 arcsin(q q') arcsin(k' k')
(31)
Proof. To begin, we compute a closely related property - the probability that the signs are equal:
Pr [)( w²) Pr [(x, k) (q. v) (xx, k') (q' y) > 0]
(32)
2
(33)
Let I be the angle between q and q' and let o be the angle between and k', We have
Pr( (y d¹) 0]
(34)
Pr((xr,)) (ar, k') 0]
(35)
Pr ④) (y, ) 0 (x. (x,
> 0
(36)
0
Pri q) (y. m.y
(a.
5
0
(37)
0
0 Φ
Pr Pr k) (x, (y.
1
(38)
2.9
A bit of algebra now shows
(1
0
6
Pr [p(zww) p(z,w')) =
+
(39)
π
2
2
πT2
2
8
2
(40)
By definition, 0 = arccos( q')) and o = arccos( (k, k')). Using the identity arcsin(z) = m/2 arccos(2), we obtain
1
2
Pr [p(z,w) = p(z,w')] z
aresin(q q²) arcsin(k' Tk))
(41)
2 n2
Finally,
E. jo(z,w)p(z,u) 2N0
Pr jø(z, w) = p(z, w)] Pr [p(z,w) + p(z,w)))
(42)
zuf [p(z, w) = p(z,w))) -1 Pr
(43)
arcsin(q" T q²) arcsin TkT))
(44)
The above lemma gives us handy expression for T-T that allows to show the following:
24

---

Published as a conference paper at ICLR 2025
Lemma 10. Let Y, be spherical harmonics of degrees & and f", respectively, Then Y @ Y' an eigenfinction of the operator T-T:
))) 2² (((.))
Oxg
.Y Y
(45)
Proof. is easily seen that
(T*))()) (00)=/a (100)=/((zz))000) plzz))((=)d7(2
(46)
and thus, substituting and changing the order of integration
[T 2 (')[() = / E ((z,,,)))(,,' (Y
(47)
Applying Lemma 9 and expanding (w) and @ Y
4
arcsin(q" q') arcsin(k (Y @ (')
(48)
π²2 Jn
/ Ige- aresin(k T k (k')(k)
aresin(q d
(49)
n2 ga-1
Applying the Hecke-Funie formula EEquation (24)) to the first integral,
[ arcsin(q (rrcin(("T()) (d))((()) = (g) (arcsin, Pélus
(50)
Pe Pe
= -())) aarssin, ||
(51)
= =Y(g)
at
(52)
(d(
By the same logic, the second integral equals Y'(k') oel //rvv E). Combining these proves the lemma.
The previous lemma immediately implies that the functions T(Y @ Y") form an orthogonal basis: Lemma 11. Let B be of orthonormal spherical harmonics. Then the elements of (T &YY) Y,Y' E B} are also orthogonal. Furthermore, Y and Y' have degrees and E, then
4
[| @ Y T² (Nd.e)
(53)
Proof. Let Y,, Y,, Y2, Yy E B. Let Y/ have degree l and Yú have degree l.. Then
(T(Y, @ T(Y @ (Y,)) = (Y, 2 Y, TT2 @ o,,))
(54)
= (Y 2 y, Yv 8Y) 7 (ddd
(55)
But (Y, ② &Y, Y, ② Y is one if Y, = Yd and Y, Yy, and zero otherwise.
C.4 EXPANSION OF THE TARGET FUNCTION
We define a surrogate target function that will turn out to he the relevant one for our analysis.
Definition 12. The surrogate target function f X + IR is
f(z) = sagn(x
(56)
After change of variables (æ, w) = (x) x2 x + Ø2), our original target function reduces simply to f(x)x + w. We now wish to expand f in the basis (T(Y & Y')). We will first need the following lemma, which describes the correlation of rank-1 head with the surrogate target function.
25

---

Published as a conference paper at ICLR 2025
Lemma 13. Fix w (q. k) S2 Then
20
(f.pl.w)); = 2=
cPPila k)
(57)
where
2
C2 10
(58)
Proof. By definition,
(frp(-w)),
= E siggn( y) sign(æ T kg 2/8/ea
T y)]
(59)
Let 7+ denote the uniform measure the hemisphere {ar E s x k 0), and the uniform measure on the opposite hemisphere. Then we can decompose the expectation as follows:
1
E siig(xx T y) sign(æ kg y)] 2/
E [sign(x: y) sign(q T yy]
(60)
2 avor- y-T
E
[sign( y) sign(q T y)]
(61)
2
Given any fixed unit vectors æ, q we have that
= arccos(x: 4)
Pr(sign(x y) sign(q' y)] [y)]
(62)
y
Therefore,
E[sign(æ y) sign(q y)] Pr Prisign(x* y) sign(q T y)] Pr(sign(x v) * sign(q y)] (63)
y
Pr Prisiig( (x y) si y)]
(64)
y
2 arccos(x: q))
(65)
Plugging this into the expression above,
1
2 arccos(xx q)
1
arccos(x T 4))
E
E
(66)
2 a~T+
m
2 x--t
n
1
1 E
E
aarcoos x q)]
[arccos(x T 9)]
(67)
π
2
2
mm
1
+ 1
E
[sign( k)) arccos( T 4))
E
[sign( T k)) arccos( T q)]
(68)
2 xou7T
2 mmat_
R axy
[ii T k)) arccos( T 4)]
(69)
(70)
Using the identity arccos(t) _ arcsin(t) and the fact that Ez [sign(<> [k]] 0.
2
E sign T k)) arcsin( q))
(71)
π
2
sign(( k)), (r.sin((:,9))).
(72)
We now expand sign((- ((, k))) and arcsin((:, q)) in a basis of spherical harmonics. By Hecke-Funk,
(iign((: k)), Y2 Y/ (k) (sign, Pt) Y/ (k))eelPPlll
(73)
aresin((; ()), Y2) -Y((g) (aresin, Pi) (eee
(74)
(75)
26

---

Published as a conference paper at ICLR 2025
Thus, writing the inner product in the basis of spherical harmnoics,
N(d.))
2 <ssign((; ar)ak)))))))))))))))))) (r))
2
Σ
(P()))/)PPlllu) (76)
π
N(d.e)
Σ (((1)/1(1) j=1
neae| Pelles
(77)
8=
Applying the addition formula
Equation (22)
2
Naael q)
(78)
2 naaele(k g)
(79)
t=
(80)
We now expand our surrogate target function f our basis (T(Y @ Y')}. The following lemma shows that f is orthogonal to any basis element for which Y * Y', and that the coefficient of T(Y & Y') only depends only on the degree of Y. That is, the energy of f is evenly spread across all elements of (T(Y, Yi) Ye Fr).
Lemma 14. Let Y, Y be spherical harmonics of odd degree. Let E be the degree of Then
f T(Y &Y) IITYY & Y') (Y) ), VN(d,))
mE Sv.Y
(81)
where dy,y" 1[Y Y'J. That is, if the basis element is built from two identical spherical harmonics of degree É, then its correlation with the target function depends only on E; otherwise is zero.
Proof. Expanding, switching the order of the integrals, and applying
Lemma 13
(I.((Y
( ( (())((z)))(> & (")())((WW)F(((
(82)
(F:pl:w)), (Y & (F
(83)
ce / (r(qTT)(2 & (')(((((
(84)
P=
Expanding the integral over Ω and applying Hecke-Funk (Equation (24)).
80
[ Pe Pelq ()
ce
(85)
Jgd
F
20
/
Ce
(Pe Pe) )
(86)
gga-1
Σ ==2
ce Per 11 (Y,Y",
(87)
(88)
N(d,))
27

---

Published as a conference paper at ICLR 2025
Finally, applying the formula for 0ı from Lemma 13 and the formula for IT(Y Y from Lemma 11.
(i
T(Y &Y) [T(Y & 1IFF/= (((,e) IT(Y & 1)1F= N(d,e) v /
C
- nece
(N(d,
(89)
Up now, we have constructed a basis without showing that its span includes our target function.
Lemma 25 (in Appendix C.8) verifies that, fact, f lies in this span. This lemma is not needed for the proof of Theorem 2 but is used in the kernel approximation Appendix C.8) It also shows that this step of the proot is tight. We do not lose anything by lower bounding the error only on the part of f that lies in the span of our basis functions.
C.5 EXPANSION OF THE HEAD FUNCTIONS
In this section, expand the low-rank attention head function in our basis (T(Y @ Y'). Unlike the target function, the energy of an attention head is not spread out, but concentrated on a few basis elements in each harmonic. We first need the following lemma, which will to bound the number of these special basis elements.
Lemma 15. Let Ae be the span of the harmonics of degree i on g that zero after marginal- izing onto the first r coordinates. Then
dim(Fe/At) = :(r,)))))))))) S (T+)
(90)
where Fe/Ae is the orthogonal complement of A in Fe. Furthermore, (1 =
Proof. Let L Fr (B,) be the linear operator which marginalizes a degree l spherical har- monic function on the first r coordinates. (Here, B, is the unit r-ball.) That is,
y-8d / ([DVi="
(cf)())=
x
(91)
By definition, Ar the null space of L. We will show below that the range of contains only polynomials of the first coordinates of degree at most f. The dimension of the space of polynomials in dimension of degree at most & is (*tt). Thus, by the rank-nullity theorem,
dim(Fe) dim(AA) (r+))
(92)
and therefore
dim(Fe/At) dim(F) - dim(Ae) < ( (r+e)
(93)
We will now show that the range of L contains only polynomials in the first r coordinates of degree most Each spherical harmonic the restriction to sd of harmonic homogeneous polynomial on Rd, so it suffices to show that L maps monomials of degree exactly É in Rd to polynomials of degree at most in the first r coordinates. Let
d
Y
x y
II
Pi
II imr++
y.
(94)
be one such monomial. If any of pr Pa is odd, then L[Y] = O. If all are even, then
II (v.V Duv1-1las 1 P.
(95)
y-ga-r-1 1=7+1
d
II (=++1
(1 - 131113
E
V
(96)
i=1+1
28

---

Published as a conference paper at ICLR 2025
a polynomial in a whose highest degree term has degree (Ei=1 pi) + (En (Ein-1 which equals the degree of the original monomial.
For the special case of r = 1, suffices to show that has rank one, or equivalently that its nullspace has dimension N (dd, E) - 1. Let Y1 = Pr((e é. 1),, where é E Rd is the first standard basis vector. By Theorem 4. of Prye & Efthimiou (012), Yi is spherical harmonic of degree l. Complete an orthonormal basis (Yi,... (N(a,)) Fe. Our goal is show that LY; for all j E [2, })} (with equality in the weak sense).
To do this, it suffices to show that (Y,) = for all
(Pe, = E
[P(())(CY)(<=))
(97)
wout
(r(x)
(
x
D1
(98)
govud
DER
= E (P(x)) 2
(99)
where
S But by definition, Pr(x) Yi
(<
x
for all
y
y gd-2 Continuing from above,
E (M(=)) (,), =0
(100)
for all j + 1.
Lemma 16. Let X a square matrix. Let D the uniform distribution over orthogonal matrices. Then,
E Q xQ] tx)))
(101)
Q~D
Proof. Let qki denote the entry in the kth row and ith column of Q. Then the (i,j) entry of the expectation is
E Q XQ Q~
ΣΣ oue glonsss K
(102)
So long (k:, i) * (L,j), then conditional distribution of qrj given qui is symmetric, since negating the fth row (or jth column) of Q would produce another orthonormal matrix. Thus, if (k; # (f,j), then the expectation is zero. The only non-zero terms are
E Q
XQ
Σ
Fek B Bl
(103)
Q~D
k
Since the marginal distribution of each row (or column) is uniform on the unit sphere, the variance of each entry is 1.
Lemma 17. Define (r.,)) in Lemma 15 Assume the rank d and consider the functions (x) x Vaa tr K, xr, y) for dn Ie' x 3d-1 - and K₁ E Rdxr for h [H). Then there exists a subspace A S. Ft of dimension at least N(d, E) - M M(r, {) such that T(Y 8 Yo) is orthogonal to 9h
for for any Yt E Ae and any H.
Proof. The first part of the proof gives construction for Ac. Fix y. q and h and define
h (k)
E ((lz.wggntee axe sign(( kg y) Vx Tn (K x.,))
(104)
200
Define K K
k). As first show that this function only depends particular projec- tion of V. not on itself. Choose basis such that the column span of K is span( ((e), er)).
where 1 < r' min(r + 1, d). Then
rewrite V
A
where A B²'''', The D
B
C
distribution of x is isotropic and independent of y. Therefore, we can rotate it without affect- ing the expectation. In fact, draw random orthogonal matrix from any distribution, and
29

---

Published as a conference paper at ICLR 2025
Eu, Q (Qx)) will equal Es [f))]. We draw random orthogonal matrices that fix the column span
of K, that is, matrices of the form Q I
Q where Q E R(d-r")x(d-r) is uniformly
distributed orthogonal matrix. Then,
(k(())
E a.Q
sign(x Q kgq y) Q Vi (n(K (KTQx,,1
Qx, y)
(105)
A QTC
BQ òn (K a. y)
E
sign(æ kg y)
(106)
Q
Moving the expectation over Q inside, the off off-diagonal blocks are both O. Applying Lemma 16 the bottom right block becomes tr D) I. Thus, letting A' A tr(D) I,
A
BQ D
E
tr(D) I+UA'U
(107)
QTC
Q
I
where
defined to be the column span of K. In all,
(k(k)
IE 2t
sign(x kq y)æ (tr(D) + + -UAUU x "n(K' x. x.
(108)
Now that we have reduced V, we can more clearly see the implications of the rotational invariance of the distribution of æ. Let O be an arbitrary orthonormal matrix. Then
hk (k) =
E si O
To kqq y)æ TO (tr(D) I UAAU Oz oh ( Ox, v)
(109)
E --T
sign(a TO ""kqqs ") T (Ir(D) 0 OUAAU 0) x on (K Oz, y)
(110)
hoT K (O
(111)
where the last step follows because O is precisely the column span of OT K. Thus by Weyl's fundamental theorem of invariant functions, there exists h R" - such that
hx (k) M(K TT)
(112)
Let TK denote the marginal distribution of T on the column space of K and let TK- denote its marginal distribution on the orthogonal complement of the column space of K. Then the random vector v + v /1- ||l.] where v ~ TK and v ~ -K is distributed uniformly on the sphere. Let Y be spherical harmonic that is zero after marginalizing the onto the column space of K. (For example, if K = K Opxd-r]. then marginalizing onto the column space means taking the average of the function over the final r coordinates.) Then
/ bx (k))))((())d
(hK, Y)
(113)
3a-1
/ l
-I--tii--ii ((ol(//)) (orrrki(o=tttrr (o=)dTk(())
(114)
/ (ix(()) (fv(o=o-vi-1o)))) |o] dT<(1)
(115)
0
(116)
Let Ah C Fr be the space of spherical harmonics of degree that have this marginalization property with respect to K,. Let A. = NxAA. Recall that Ndd, E) is the dimension of Fe, and M () the dimension of the orthogonal complement of A in Fe, denoted Fe/Ag. Thus,
dim()A) =
= (F 2 ((a,)))H-M(rr,) (117)
remains to show that for all E Ar, T(y @ Yx) is orthogonal to gh.
(
(T(Y Y), ())
(118)
f E (f
Elgoa..g..)nnn v(g)T())
(119)
Jga-1
30

---

Published as a conference paper at ICLR 2025
But for any fixed y and q.
ga- Elp(a, y. (2) = (hKn, 1)
(120)
by the calculation above, where the final step follows because Y E C A
C.6 PROOF OF THEOREM 2
Theorem 2 (Low-Rank Approximation Lower Bounds, Equivariant Case), There exist universal constants e, e, C and C' such that if either of the following sets of assumptions hold:
1. High-accuracy regime: r S -3.
c d+1
and
H SSC. 2 (r81)
2. High-dimensional regime: d d 5,
and
T
C/e
1 1 d H H
(6)
2 2e r r C'
Then, for any choice of rank-r generalized attention heads dr, 1RTx2 + A Vn E Rdxd Kn Rdxr the error of approximating the nearest neighbor function is bounded as follows
((X:9)
Σ V.Xon (K, x,9) h=1
2e,
g--Unif(sd-1)
2
where f is defined as Equation (3)
Proof. We lower bound the error by projecting it onto the unit vector (x) x2)/(v2). For conve-
nience, define basis
T + 22
(121)
2
V2
The joint distribution of x and w is the same as that of æ and 22- They are each uniformly distributed the sphere, and they are always orthogonal. The projection of the target function onto yields the surrogate target function Definnitoo 12
D2 V2
f(X:8)
sign ((z, x2, x2, v)) =:
1 f(z y)
(122)
V2
/2"
Let the attention weights produced by softmax head be t and to = head before multiplication with V is
t1.
Then the output of the
t
1
tx + (1 - 1)2²
x
/2 w
(123)
/2
Letting T(KT (x,y) y) = (41 - 122/ the inner product of the head with æ is
x Vx O(K y) + x V
(124)
Notice that, since the conditional distribution of w given is symmetric, the correlation of the second term above with the surrogate target zero:
(84 1) [(x, 3) x Vw
(125)
31

---

Published as a conference paper at ICLR 2025
Thus, have the following lower bound:
H
E
((X:y)
Σ
Vn din (Ky) (a1 - 222); v)
(126)
21 gruUnif(Sd 4
h=1
2
E
00 s(X:y)
1n Xon (K) 12), yy)
(127)
an (2 g(-Uni((3(
f(x,y) - Σ
VA (K,,,, x. x,
Σ
Viw
(128)
t. 18
hml
hm
y~Unnf((d-1)
H
E
f(x,y) - Σ
Vv on (K, x; y))
(129)
2. Uuir
2
h=1
gh
(130)
2
hm
where (2) = x. VA (K,, x. 9). Construct the space Ae S Fg according to Lemma 17 and
let (Y;) dim Ae be orthonormal basis of Ap. Then each element in the following set is orthogonal ]==]
to each ga (z):
T(Y & Yi [T(Y & Y
dim(Ae)
(131)
Furthermore, by Lemma II
this set orthonormal. Thus
dim(Ar)
T(Y & YY) gh: [I & YII
Σ gh
Σ
Σ
(132)
dim(Ac)
T(Y & YJ [T & Y2Il
Σ
(133)
V N(d, E)
dim(Ae)
(134)
E odd
where the final step follows from Lemma 14 By the construction of A. Lemma 17
dim(Ar) 2 N(d, E) H M(r E)
(135)
and thus
1 1- H. M(r E) N(d,
Σ
n
(136)
8 odd
Appealing either to Lemma 22 or Lemma 23 finishes the proof.
C.7 ASYMPTOTICS
Lemma 18. Let m > l and E odd. Then
[( (
(4) dt
m (1 2) --1)t=(-1))(
( 1)!
(137)
{-
32

---

Published as a conference paper at ICLR 2025
Proof. We have
J (1))) di (1)) (1
d
8-1
())
ym
=
(1
(138)
dt
2
d
m
k 2k
(139)
dl
k
(-1)2+((-1)2
(e 1)1
(140)
(-1 2
Lemma 19. Define ny as in Definition 7
For odd E, mn v pाetaa
Proof. From the definition, we have
nid = (N(,))AAu2
Ad-1 J
-1)(0)(1-
(141)
From the Rodrigues formula for Pt,d
(Frye & Efthimiou
2012
Proposition 4.19), we have
Mhd = 2 Add
(-1)²
d (((4)
(1 (d+t(-33/22
(142)
Ad-1
2 + (d (01-3)] 3)/2)/
di
Now, using Lemma 18 we obtain
Bid = 2 V ((d,
2(1 + (1- (d-3)2)) (-1) (1 (d (33/)/22)
(-1)²
(-1)
(1 - 1)!
(143)
Ad-1
1-1
and thus, using
Ao-2 As-1
C² Vd, we have
1)!((d ()((((3))2)) (1+ (d -3)/2)
mnal ~
(
(144)
(I+ (d- (d-323)/2)
(H14_2112)
Vd (Nd,1222
4
(145)
1
(+((-33/2)
Vā (dd,1)2-
11 (1((=3)1
(146)
(T))) (1t
Using Stirling's approximation, we obtain
1+d (1 (1+d) 1/2 (1 (1+d) d)
N(d,t) 1 ld
(147)
1(1 (d--2)
~ (1 dd) +d-- 332 1/2 d+33/2
(148)
well
11 d-3
ld
3d-3 (1+))dd--44)))) 1)-44222)/22
(149)
(5))) (1112))
[(d +1]
(1 + d) d
(150)
leading to
Im.al ~ ||0 ~ + d) 1/4/ 1 3/4
(151)
(152)
claimed.
33

---

Published as a conference paper at ICLR 2025
Lemma 19 shows that n decays slowly with f. Using / v 2 1/E8 and including by a fudge factor that is slightly smaller than 1, we get form that is better suited the proof of our lower bounds:
Corollary 20. There exists a universal constant e' such that ni 2 c for all sufficiently large d and & (say, for all d, > 4).
Lemma 21 (Decay of ag). For odd, we have at = /
Proof. We start from aresin a arccos and the kernel representation (Bach 2017a Section 3.1)
(TT arccos(x y)) Eogge-1 [1 001ly O]]
(153)
2x
Now, from the Hecke-Funk formula, we have, up to zeroth-harmonic terms, the following corre- spondence between the Gegenbauer expansion of aresin and that of of sign, given precisely by me- Fix any x E Sd-1 Then
(aressn Pt)
aresin(w Pr ydy(dy)
ʃ sign(æ 0) sign(y P( (dd
4
(sign, Pt) / sign(æ 0) Px (d))
1 Pr (1)ggn P1)
(154)
Since (aresin, Pt) aellPell and (sign, Pt) mel Pell so at 1 102 (d.
Lemma 22. There are universal constants e and C such that the following hold: Assume r S d-3, S J and H S C 2d (r+11 log, Then
( 1/ 2c
Σ 1-H ouid
(155)
((d.
Proof.
28 + 2
d 2(1
(f+
(dd
(156)
l
(-1 1
Applying Stirling's approximation,
2(d,)) f+ d -3 f 1113000 2)3 (L + d - (+)dd3))td-11 (d.0.b 2)
(e+ d 3) 2-2-2
(157)
(158)
Meanwhile, Lemma 15 and Stirling's approximation give
r 1+
(r ty+t0us
M(r
(159)
pr+0.5pe+0.3
By assumption, r < d- 3, s0
M(r,)) ((d,t)
r+l ++6+0.5
(d - 2-dd--1
(160)
1 d -3
p + d - 8-
(d (d-22 -22d--155
(161)
pr+0.5 + d 2
34

---

Published as a conference paper at ICLR 2025
The above expression is decreasing in É. Thus for all l 2 pd + 1,
M(r,))
(d - 2)
2
(162)
E)
(1 u)(d 2))d-r-2
+0.5
(163)
(1 + 1
(1+ 1)=+ -+
(r+0. (
(164)
By assumption, c/c d + so the above holds with μr 1 for all / c/c:
M(r,))
- ( log,
(165)
E(
Also by assumption, S C' 24 (r+1) (2g2))))))) Setting C' appropriately, (1 (1- H
M(r.))
N(d,8)
2 2
for all & c/c Finally, applying Corollary 20
M(r,)) ((d,
1 d"
Σ odd
1 H
m/
2
(166)
2
E00 f odd
"
(167)
8
F20/c
d²
(168)
c
Setting c = /4 completes the proof.
Lemma 23. There is universal constant such that the following holds. If 2 5.
2c
d
(169)
6 2ea
and
1
H
S
(170)
2
2e
+
then
Σ 1- M
me 2
(171)
N(d,))
ouid
(172)
Proof. Recall the formula for N(d, E) from Equation (19)
Lower bounding, for e 2 and d 2 5,
(d 28 + d 2
(E+d d 3)
(173)
1-1
8 + d - 3
(E+ d -3 3
8-1
[( d - 3 8
>
(174)
1-1 1
d++"
(175)
28
(176)
Meanwhile, Lemma 15 gives
(r+()
(ctr+t
c((r E))
M(r,)) (1) 8
<
(177)
Thus
((r ((d,
r + E
r r+l
<
2e
2e
(178)
d+t
d
35

---

Published as a conference paper at ICLR 2025
The above is decreasing function of & for all & < following holds for all & [ 2g] M(r,8) (dd,))
d
Assume that
Then the
2n
2n2
r
r+
2e
(179)
d
Assume H S 2 28 r+t
d
Then for all E
[ 20]]
M(r,
1 1 H
2
(180)
Ndd,
Finally, applying Corollary 20
2c/r
M 8) N(d,
c"
c"
Σ
1 1-
th
(181)
2
4
1
4
2c
É odd
e odd
Emoje
Setting c ==="" d//8 completes the proof.
C.8 KERNEL RIDGE REGRESSION AND RANDOM FEATURE APPROXIMATION
In this section, we analyze a simple approximation the nearest neighbor function by standard rank- attention heads. We show that (((=a(//e) heads suffice to achieve squared approximation error of c, nearly matching the lower bound of Theorem 2| First, we reduce this problem to approximating the surrogate target function f by rank-1 I hardmax heads. Then we approximate f in the RKHS generated by rank-1 hardmax attention heads (that is, generated by the feature map T). Finally, we appeal to standard arguments to conclude that we can approximate by a finite linear combination of random rank-1 hardmax heads.
Recall that a standard rank-1 attention layer has the form En On on Xmm (x ki q, y) for qr. kn. Uh On E Rdd For simplicity, in this section we use rank-1 heads without value/output transform, that is En ar Xsmm (X k, y) for IR. Any such head can be constructed out of d standard rank-1 heads by setting 2n, Ci Vo ae, for [d], this simplification does not meaningfully change our result.
Lemma 24. For any u E L (2)), there exists a rank-1 attention layer that approximates the nearest neighbor function f up to expected squared error 1 f where T is defined as in Definition and f is the surrogate target function of Definition 12
Proof. As in the proof of Theorem 2 define
E2
I + D2
w
(182)
V2
V2
We can rewrite the target function in terms of the surrogate target function as follows:
f(11, m2:
f(x +
(183)
V2
V2
Likewise, we can write rank-1 hardmax attention head as
x hm (x kg 3)
\ p(x, /2
where p( y: g, k) := sgn(æ kg y) is defined as in Equation (28) An "averaging head" is an attention head that always returns the average of the target points, regardless the source point. can be implemented by rank-1 softmax head by setting q = k 0:
X sm X Toy)
2 + 82
w
=
2
/2
36

---

Published as a conference paper at ICLR 2025
We construct our approximation to by taking a linear combination of hardmax heads with coeffi- cients given by u plus a single averaging head with coefficient 1 - Sn u(q, k) (q,
/ u(q. k) X hm [ kg y) dF k) + [I / u(q. (), (X, +
x + 22
2
(184)
To analyze its error, we use the Pythagorean theorem. Due to the averaging head, the projection of the error onto w zero. What remains the projection of the error onto x::
((X:v)
1/8 u(g k.) X hm (x kq T y) di k)) + (1-// ulq, k)) k:)
w
E
v2
(S u~Unni( 1)
(185)
x ((::)) / uuq, kb) xTx x hm (xTaq kq v) dF(q.
=: 1 | Tul
E 2.y=Unir(($a-1)
(186)
[(x,8) / pla, k)u(q, k)dFq.k)) =: m.y=-Unif($aa)) 1) 2
1
1 1 Tu (187)
E
By the above lemma, our task is to find a finitely supported signed measure u for which f ~ Tu. We next show that it is possible to exactly represent f using measure that is not finitely supported.
Lemma 25. The surrogate target function ff lies in the span of (T(Y) Furthermore, f = T where u. Ω TR. is defined as follows:
u(w)
N E) Pela" k).
(188)
2
odd
N(d.e) be an orthonormal basis for Fe. Applying Lemma 14 the norm of the projection of f onto the span of (T(Y? @ Yf)) is
Proof. For each odd f, let (Y:)
N(de)
T(Y & YT) ) IT(Y & YOIL = Σ
2
N(d.)
Σ
J.
(189)
(dd,))
Summing across all (odd) degrees, the energy equals that of f itself.
Σ M21+1 = sign || = 1- = 11f112
(190)
Thus, the projection of / onto this basis equals f. In addition, this implies that f is in the range of
T:
N(dE)
T(Y & YA) IT(2 2 Y
TO & YY) IT(?) @ Y/II=
j ズ=Σx
Σ
(191)
odd
N(d.)
Σ
Σ
me
T(Y Y2)
(192)
2 N(I))(d,)) (N(
odd
N(d.0)
T
π
Σ YY
(193)
2
at
odd
T(u).
(194)
where, by the addition formula,
u(w)
Σ
Me N(d, E) Pelq k)
(195)
2 odd
at
37

---

Published as a conference paper at ICLR 2025
x10
Dimension =
ulq.k)
10
Dimension = 6
u k.
0-
110
Dimension 8
5
u(q, k)
0
-5
x10
Dimension = 16
a(q. k)
0-
-5
77
arccos(q k)
Figure 8: Approximation u(-) of Equation (188) for several dimensions, using degree-50 ultras- pherical expansion. Heads with L(q, are equivalent to those with angles + π up to sign flip. For large dimension, the distribution over L(q, k) induced by u approaches Gaussian with mean O.
Thus, it is possible exactly represent the surrogate target with infinite number of rank-1 heads, each weighted according to u( (-F). See Figure 8 for an illustration this function. We can think signed measure
of u(())d(-) hardmax head function is odd in each of its arguments q and &. Since u(-) is also an odd function, get the same results by restricting this measure to | /. 51- Figure 8 shows that for large d, the (restricted) measure u(-) approaches a Gaussian distribution centered at angle 0.
rank-1 heads that depends only Z(q. k)) Notice that the
We have now shown how to represent f using T. This representation gives us great deal of insight into the structure of f for the following reason. Implicit in the discussion above is the reproducing kernel Hilbert structure induced by the map T. as the following lemma shows:
Lemma 26. Let H C (2)) be the image T. Then H is a reproducing kernel Hilbert space with norm:
Iffll = inf([lu]|= a G. =Tu}
(196)
and kernel:
(2,2') E (p(z,w)o(z whn
(197)
38

---

Published as a conference paper at ICLR 2025
The proof is given in Bach (2017a), Appendix A. Also note that kernel of this RKHS directly corre- sponds to the operator TT by the following formula:
/
())
E wop
wi) ())
(198)
If our target function f were an element of this Hilbert space, we would immediately be able to approximate using random features. Unfortunately, / & H because
|l = || =
Σ odd
M1 ag
((d, v(d,E)=00 00.
(199)
However, we can approximate f by an element of H obtained from solving a ridge regression prob- lem. For any À > 0, let fa be the solution to the following optimization problem:
min
|f f11 + A fEH
(200)
By tuning X, we can find an function that accurately approximates f and that is smooth enough to be approximated using random features. The following lemma constructs this fx. Though we obtained this construction by solving Equation (200) for brevity we do not prove that it is the solution since it is not necessary for our construction.
Lemma 27. For any regularization parameter À > O, there exists a function fx E H. for which
AN(d, E)
II
fAlI
(201)
22a)22 + AN
odd
Proof. Define
Bx Tgx
(202)
9x = Σ N(d.)
n -Yi 8 Y?)
(203)
oodd
i=1
LasT
78 (2))3 + Ad(
(204)
Then by Lemma 26
| H S |1
Σ E(d,
(205)
odd
4
Mdd,0tt +
EN(d,
(206)
(d
66
1
N(d, 002 +
(207)
X²
8=1
00.
(208)
Thus f E H. Furthermore, by the representation j j= T of Lemma 25
N(d.))
- fAll IT (u
1f=
Σ odld
7 T(Y? 8 &Y.)
(209)
2 az
f=1
39

---

Published as a conference paper at ICLR 2025
By Lemma 11
this is equal to
N(d.
71
|
(210)
2 oz
Fodd
==1
2 of
e)
T T
(211)
2 ai.
72 e(d,))
Eodi
2
aLYe
(212)
π
odid
(N(( E) (2)) + AN( )²
2
(213)
We now derive an informal expression for the kernel ridge regression approximation using tuned regularization and describe its implications for random feature approximation in the high- dimensional regime. From Lemma 19 and Lemma 21 we have n S 5-3/2 and a ~ (d/Nd,, E). By Lemma 27] for the kernel ridge regression approximation fA to attain squared error we should set À so that XdV E*) ~ aZ-, where e 1/e². This roughly ensures that only degrees f 2 l are kept, while L $ E are shrunk, and hence
II fAll S Σ
S 2
1 Σ
8-3/2 ~cc
(214)
E25 E odd
125-2
We thus obtain o2. (dd 3)-2
Now that we have a sufficiently accurate kernel ridge regression approximation fx E H. can approximate it using random features. The key quantity controlling the number of random features needed is the degrees of freedom of the kernel integral operator, defined as The eigenvalues of TT the same those of T*T By
D(A) tr (TT + --
with the (l, eigenvalue having multi- x" VN(d.)N(d..)) le, 2 z0)
Lemma 10, these are
plicity N(d, (d,
Hence
D(A) Σ (dd.))
((d.))
S EN(d.EN((aE)) (id.))
(215)
+
A
E
E.8"
By Lemma 21
((d.)N(d.)))
M(d.))
2
222
D(A)
Σ Ti E.E'
Σ
2
5 (ede
(216)
c
In the high-dimensional regime (where is fixed and d goes to infinity), D(A) = θ 00: (8/1²)
By standard arguments about random feature expansions (Bach 2017b). if the number of random features H is of the order II 2 D(X) (og(()) - d2/c" then with high probability the random features achieve the same approximation accuracy as the associated kernel ridge regression solution fa. It is likely that a better rate can be obtained by drawing the random features from problem-specific distribution instead of uniformly at random. Observe that the condition required by our lower bound in the rank-1 case has the same form, though somewhat weaker dependence on d. It is H < 1 (d,, + or 0
for sufficiently large d
'To see this from Equation (136) recall that M(1, E) 1, follow the final steps of Lemma 22 and use the fact that can replace c" for large d.
40

---

Published as a conference paper at ICLR 2025
D PROOFS FROM SECTION 5
D.1 PROOF INTUITION
Part (i) of Theorem 3 is easy prove. We approximate each summand of the target (Equation (8)p by a single attention head with a large enough temperature. We note that this construction doesn't t depend on the input distribution, and can be readily generalized any other continuous distribution. The error parameter affects the necessary temperature (i.e., the magnitude of the weights), but not the number of heads.
The proof of part (ii) of Theorem 3/ is more intricate. The crux of the proof is to construct linear combination of threshold functions that behaves like periodic function with high frequency. Our proof is inspired and extends the proof method Yehudai & Shamir (2019) for separation between kernel methods and 2-layer neural networks. Note that each summand of the target in [Equation (8) can be written as
arg max (zi; y) bt.; = 1 (( - 22 v) + by > 0)æı + [((11 22 y) + bj < 0)æ²
(217)
where b7 = b,.1 2r.2. Let a == 2d + 1, and recall that b,, = for = 1 and 0 for i + 1. Rewriting the target function using the above substitution, we can show that it is periodic in the interval i-a,00 a). considered as function of (x1 - 22 22. y). Denote this function by da We extend
result from Shamir (2018) to show that:
[V - 22, y)) g (Kx,, K x2, y)l S ||9|| exp( -))
(218)
for any function g with bounded outputs if K has rank r. Finally, a rank-r transformer with H heads can be written linear combination of such functions with bounded outputs, where the weights of the linear combination are Vn.
We draw y from Gaussian instead of uniformly from the unit sphere in order use the result from Shamir (0118) There, the bound on the correlation is expressed in the Fourier domain. Since the Fourier transform of Gaussian function is also a Gaussian function, the calculations become significantly easier.
D.2 FROM POLYNOMIAL TO EXPONENTIAL SEPARATION
As can be seen from Theorem to get an exponential separation between low- and high-rank trans- formers, we used a more complex function than the nearest neighbor as used in TTheoemm 2] It is natural to ask whether this more complex function is necessary to achieve this strong separation.
We claim that it is not possible to achieve exponential separation using a target that is just one nearest. neighbor function, even with a bias term. Fix some r < d, and consider some target function f that is hard to approximate using rank-r attention heads. There are three important parameters that play a role in the separation result: (1) The desired approximation error e; (2) The number of heads Il required to achieve this desired error; and (3) The Lipschitz constant L of the function f.
In Hsu et al. (202I) is proven that both necessary and sufficient to approximate any target function by polynomially many random ReLU neurons if both O(1) and L = O(1), that is, if they are both independent d. However, if either e = polytaj or L = poly(d). then exponentially many random ReLU neurons are required. Our results resemble approximation by random neurons r. Using this analogy, Theorem 2 covers two cases: (1) E = O(1) and L = (11) and thus the separation is polynomial; and (2) 6 = (1/dd) and 0(1) where the separation exponential.
(with ReLU or other activations), where the "effective" dimension
Thus, to achieve exponential separation for 0(1) we need to construct hard to approximate target function with Lipschitz constant I, poly(d). One method is to "cheat" by multiplying an
L O(1) target function by a large number that depends on d. However, this will not result in a small relative error; effectively, by multiplying the target by some constant, we also multiplied construct target function f with || < O
by the same constant. In Theorem 3
"This can be seen by equation 224 and I Lemma 28 where the target for N 2 can he written as sum of
two periodic functions, both bounded
41

---

Published as a conference paper at ICLR 2025
meaning that hardness result applies the relative error too. However, this target function is highly oscillatory, resulting in a Lipschitz constant of L = poly(d), and thus we get exponential separation for a relative error of O(1). We believe that the result from Hssu al. (2021) can be extended beyond random ReLU neurons. If so, then achieving an exponential separation with "simple" function (i.e. Lipschitz constant independent of d) would be impossible.
We also note that a similar impossibility result is known for separation between 2- and 3-layer neural networks, Safran al. (2019) shows that such an exponential separation cannot be achieved if both = (()) and L = ((1).
D.3 PROOF OF ITEM (I) IN THEOREM 3
The proof is similar the proof Fact where it is applied each head separately. Namely, let > O, and for each = I,.... 2d +I 1, set V22 I, K.Q al for a > to be chosen later. There
exists > (which depends on e) such that for the set:
As.t { ( NN, y) Vi + j. 1( - 2T) "oyyy> > -6}
(219)
we have that Pr([(21,..... NN y) & As,t) < 2 Again, note that in As. softmax converges to hardmax uniformly as + cc. In particular, there exists > such that:
sup
X sm (a (x's y + b.)) arg max ||= yll + +bt.))
bb.s)
(220)
4d +
(#) WN 2) EAA.,
We define additional head with V = I. KQ 0. Since the attention matrix is zero, for any target vector this head will attend to uniformly to the source vectors. Hence, the output of this head will be Σ Ez. Summing over all the above heads approximates the target in expectation up to an error of e, which finishes the proof.
D.4 PROOF OF ITEM (II) IN THEOREM 3
For algebraic simplicity, we first rescale the problem by a factor of vd (see Lemma rescaled target function:
Define the
21 Σ
("' UN y)
(-1)*
arg max
(1æ:; - 2 be,s
Σ ai
(221)
Vd
(E
2va i=1
We also draw æ1, 22 and y from scaled versions of the uniform and Gaussian distributions, that we now wish to lower bound the following:
[T(1,, 22, y) - (x)) 22; y) )²
(222)
351 1.- Da V
((.
In the following proof, when taking norms inner products of functions, we always consider them to be expectation over N (0, I), i... expectation over y. Whenever we consider the expectation over 1 and x2, explicitly write the expectation symbol. To normalize the expectation over y,
we introduce the constant ea :
v25
We will first construct a periodic function on the real line using a linear combination of thresholds.
Let E Ny2 and denote (x) = 1(x + 2 0). We define the following functions:
(x) (-1)" 2'
(x)
Σ HH(a-n)(z)) (-1)" 2
(223)
1=
n=1
Note that since N = and | =
||اجع|| we can write the target for in the following sense:
(* (z 32, y) + ((x) x2, y/)
Vd th ((x2 31, 9))
(224)
For of notation, we will denote vha w and we will prove the following for dh. However, note that similar proofs can be shown for eh in the exact same manner, thus we omit them for brevity. First, th has the following properties:
42

---

Published as a conference paper at ICLR 2025
Lemma 28. The function (x) defined in Equation (223) satisfies that:
1. It is periodic function in the interval a-a, a), and odd if is an odd number.
2. For every w with ||||| 2 d, if > ||//|| then [[w, 1) 2 40
Proof. Ly a-aaa-2]. There is no [l,..., 2a such that [~0] ], [ 22 E [a-12, a-no+ For every < Mu or n > no + 2 we have that Ha-n (zo) H (In + 2), since the bump in the
threshold either left 2o or right of Hy 2. We also have that Ha-no (za) -a-no+1 (oo)
Ha-no (xi + 2) Ha-- (2o + + 1. Hence va (w0 2), which means it is periodic
function with a period of
If a is an odd number, then for every E [-1,0] we have vha (rn) = and for every za E [0, 1] have tha (zo) 2-2 Since it is periodic with period of 2, it is odd in the interval [-a, a].
For the second item, since æ has a spherically symmetric distribution, we can assume w.l.o.g that ||||| e1. We now have that:
(ow -)>11 = {
(al(w 2))1 |xll
dæ
(225)
DERd
# dxl
[o e
00
00
Ca
[ [||l =1)12e
dx ...
e
da
(226)
000
-cc.
10
1 (2)( शापा. de
(227)
V2x
/2]lw|l
(aa=))
(228)
|wi| ev - w
where we used that if 2 S 2 || then e " S c-l. Since a > ||w|], then in the interval [-v2 wwl, v2 [|l]] there are at least [|]]]]
intervals the form [n, n+22 for E {-a.... a a-2)
where fn pm+ (2))2 2 1. In total, we can bound the norm by:
1
[1va(( (1((,,))))
(229)
4ev/2m 2m 40
Next, we now show that the correlation of this function with any other function that depends only on Wi...., Wy is small:
Theorem 29. Let g(wi,..., Wr, 4) be some function that depends on the first r coordinates of w
with sup |g(x)) < 1, and take 2dt + 1. Then, we have that:
E
[[ w)) g(wi...., Wr +>]l S Sexp(-c(d -r)
(230)
w-Lu(Vagd-)) 1) N(
for some universal constant c > 0.
Proof. For a vector denote by its last r coordinates. Using the law of total expectation, can rewrite the expectation in the following way:
E woll(( aoed
[Iraa((u) y)) -g(w...... Wr, 9)|l
1) I)
Σ ways (u 8) g(w,. ..... wr 4) V1.... Ur ]
E
E
E
[ wr
10 M.
11
EE
1/a
Σ
wiys + (10, 8) g(wi.... Wry wn 11 11...... Dr Wi, Wis.... Wr
W1 We M Mi....
(231)
43

---

Published as a conference paper at ICLR 2025
Namely, we consider the expectation conditioned on drawing the first coordinates of both and y. Note that we could change the order of expectations since all the expectations are bounded and finite.
Let & be a continuation of th from [-a, a] to IR such that it is periodic. Fix wi.... Wr, .r....... and denote by Ei=1 way/ and | 2p. Using Claim 30 we have that:
(2 (00) (ss (m, (n,)))i] ca (exp(-ca(d-)) r)) + Σ exp(( exp((nnp)) (232)
E
Note that in the above equation, is independent of ip (although it does depend on Wi, ...,, wy)) ) and also that ||g|l < 1 since sup= [g(x)) (recall that the norm is w.r.t a Gaussian measure).
We now have that:
E wn-ld(2pBd [I()()), va (x..))))l
((s) (z) w-2((2084
(Kot), ea (20, - 25s + (20, -1)
E (n
E
(233)
The first term in Equation (233) can be bounded by c1 (exp(-oe(d co(c r)) + In= exp( -np?)) by
Equation (232)
For the second term, by Sauchy-Schwartt we have that:
E [| (((())
(s + (w, +)) ( + (Io,
(234)
||g||
via + :) (ss (w
(235)
$ E Pr[[s + (w. y) a)]
(236)
where we used that || S and it independent of wn and that v=) = (²) for every | S a. We have that s + (w, 7) = (w. y), and (w. y) ~ ((. for every ur of norm vd. In particular, for 2 2d there is some constant c3 such that Pr(]s (w, ÿ) a) s exp(-cade Combining the above we have that:
(K(()), va(s + (w0, +)XI s.. (exp(-c2(d- (d)d))r))) r)) +
Σ exp(-np²)) + exp(-
E w-1((2208----)
(237)
We now go back to Equation (231) and consider the conditional probability over F1...., Vr and Wi.... W,. Note that when taking expectation over .i.....B4 Br we either have that (w, y) which happens w.p > 1 exp(- (ycdd or (w, y) 2 a in which case, since, [9(2)]. [2)(2) S 1 for every 2 E R also their product bounded by
Finally, we consider the expectation over wi..... Wy. We need to show that with high prob- ability. p = | large. Instead, we will consider the probability over Wy+1...., We (note that since || = vd, if we lower bound ||l||| it will also upper bound VE w Since is sampled uniformly from u(vas==1), we instead consider sampling 2; from (O 1) and setting (w); vd. it I By standard concentration bound on the norm of Gaus- sian random variables (see Section 3.1 in /eeshynin (2018) there some constant 64 such that Pr[[|æ]|| & [0.9(d - r), 1.1 - r)]) I exp(-ca(d r)). Also, Ei 2? has a x² distribu- tion with d r degrees of freedom. From Lemma in Laurent & Massart (2000) have that
Ei w 2 (d r)) I exp(-cs (d r)) for some constant Cs. Together, there is some constant Cg such that |r] 2 1 (d r)) S exp(-co(d
Pr
Note that p > c'vd-rth then Eie cxp(-np²) S cxp(-c'dd-)) Combining all the above and changing the constant terms appropriately, there is some universal constant > 0 such that:
E
E
[lrb. ((w,y)) g(w......................... wr, x)ll s <ep(-eedd r))
(238)
w-24( 1) (1
44

---

Published as a conference paper at ICLR 2025
Claim 30. For any f (0, Ia)). odd periodic function vb E - R and have that:
R. if > d we
E (0 (K (w, ->0 ci || exp((cod) Σ exp((-naa)
(239)
here d, c,11,02 C2 > are some universal constants.
Proof. The proof is similar to the proof of Claim 1 from Ychudai & Shamir (2019) (which is directly derived from Lemma 5 in Shamir (2018)), except for two changes:
Here we have an absolute value over the inner product, instead of square as in Claim 1.
2. We consider translation of th namely our periodic function is rb(s + for a fixed
For the first item, this is a direct application of Jensen's 's lemma:
VIROO v( (w, .>) S
[KS(), (ss + (w, ->,>)>
E w=
E
(240)
w
where now we can apply Claim from Yehudai & Shamir (019). For the second item, note that (b + ) is also a periodic function, and Lemma from Shamir (018) applies to it in the same way as it does on (-)).
We now ready to prove the main theorem:
Proof of Theorem 3 part (ii). Let T be some transformer with H/ heads, and output matrix Vn for each head E [H]. Then T can be written T(x1, #2; y) Σ V, X dh (K,X, y), where Φn 1grx2 A 1 is some generalized attention function on the simplex, and K₁, Rdxr, We have that:
E 22.
p (I11 T2, y) T(x1 F2, y)]l
(241)
E
Mi ((01 x2. v))
w/ ( ((x2 - x1,9)) 42
T(x,, 22, 8/)
(242)
E
wb ((x1 x2, w))
21
vo ((22 x1,9)) 22
(243)
2t W2.y
Wtt ((2)
32 4)) 21 ,T,,T / -1 ((22 2τ 9)) 32 T(E,, v)
d
(244)
We will split the expectation in Equation (244) using the linearity of expectation and bound each term separately. By Lemma 28 item (ii), and since ||x; 1 d² for 1,2 we have that:
E
((x1 - x2 y) vo
IE
b ((x2 - æ. v))
22
2
1
(245)
40
01,03.00
45

---

Published as a conference paper at ICLR 2025
For the two last terms in Equation (244) we have:
KK vo ((x1 22 1))
21
E 2:
T 22, y)
(246)
d
H
E
wht ((x) 22, 31)
21
V4 X 0% K y)
(247)
d
h=1
( 22, 0))
Vn X
Kn X, Kx,,,))11
(248)
A=1
E a22.y
[21 Vix) .115 ((01 x2 3)
o (K 8)]l
(249)
recall that the output of dn is in I2 (since N = 2). We can bound 5 Vn s d max, independent of since | = |22||| d. By Theorem 29 we can bound:
E W
((01 22 y) on X y)ll & exp(-c(d-r)
(250)
for some universal constant e > 0. similar way we can bound the correlation w.r.f v/- Combin- ing all the above, and plugging back to Equation (244) we get that there is some universal constant >> 0 such that:
[115 (21, y) T(x1, 22, y)l
2
2dH max V.ll exp(-c'(d r)). (251)
48 y22.y
40
In particular, 2d max; | 1 exp(- (''' r)) S
40 then:
E #t 20..3/
(11" 22, y) T( 22 m)|l
1 2
(252)
40
Rearranging the terms, we get that the above inequality happens if dH max, || 11
80 exp( (d Changing the constant c' and rescaling everything by factor of vd (according
to Lemma 31) finishes the proof.
The following lemma justifies our rescaling the target function and input distributions by a factor of Vd.
Lemma 31. Define fo as in Eqqation (8) and f as in Equation (221) For each attention layer T, there is an attention layer T for which
E
[TT,,, - (xx 22, y) [?]
(253)
11. Da (Vi'' M~NN (0.I)
1
[T 12, y) (x=1, v) 1?]
(254)
$1.02-D.1S (0
Proof. Let
H V, x,y)
((X,))= :=
(255)
h=1
Then construct the following attention layer:
H
T'X,y) = (VVVV (X
(256)
h=1
46

---

Published as a conference paper at ICLR 2025
where (A) Φn (n(A/vd. v/vd). For instance, if Φn y) sm(A sm(ATQny) then 6 (A, 3) sm (A AAATTQQAA (Qn Clearly, T(x / Vd. y/ Vd) (X, y). Likewise, by construc- tion, (011 ... NN y) = (x11//// 2N Vd, y Vd). Thus,
| 22, y) - f 22. m)|
E
D1.00r~D (0
(1
|| | x2/ 22/ -f tel/Vd |
E
1ae 02 Dal-DDaavaa (v
y~
[1), | ,x 2) y)
E
201 022 a0g-Da(vas* (0,1)
E PROOFS FROM SECTION AND AN ADDITIONAL CONSTRUCTION
In Section 6 present construction Theorem 5) that uses concatenated positional encodings to facilitate the majority voting strategy. This construction has the strange property that it breaks the permutation invariance of standard attention layers in order approximate function that is permutation invariant. It also increases the dimension of the transformer. This begs the question of whether these properties are necessary to allow low-rank attention to represent the target. Below, we presenting an alternative construction that does not have these properties. Instead, it modifies the attention mechanism by concatenating the outputs of the heads together rather than summing them. It then passes the concatenated outputs an MLP layer that computes the mode.
Theorem 32 (Majority Voting Approximation Upper Bound), There exist universal constants & C1, C2 Ca. Ca > O such that for all d C1, (0, 1), and H 2 C2 2x, there exist vectors T1..... qu and a f feedforward network g R Rd of width czd H such that
X x q1 9)
E Unit(3))
((01,
Se - exxp((cccdd) (257)
X (X (4t
This construction shows that using constant-depth MLP to combine the heads can overcome the weakness of low rank attention. The full proof can be found in Appenddi E.3) The idea behind the construction of the MLP g(-) is to perform inner product between the outputs of the heads,
allowing us to compare which one of the outputs T₁ or 22 received more votes. The inner products can be approximated by ReL.U network, long the input vectors not too close each other, which happens with exponentially large probability. This is the cause of the extra exponentially small term in the loss.
E. TWO-LAYER TRANSFORMERS WITH MASKED SELF ATTENTION
In Theorem we use the following definition of a multi-layer transformer with self-attention for our majority voting construction. Self-attention simply means that the source and target points are the same, in contrast to the cross-attention used in the rest of this paper. We modify the attention mechanism by adding self-excluding mask so that each input point cannot attend to itself (see below, where we form X; by deleting the ith column of X). Following standard practice, we also use a skip connection. We do not need MLP or normalization layer, though our construction can easily be extended to include them.
47

---

Published as a conference paper at ICLR 2025
Definition 33. A rank-r self-masked transformer layer with H heads is a function T RAXN RdxN parameterized by rank-k attention heads {[M,, Vi) and defined as follows:
X; = 21
24+1
(258)
H
Xx T M1
T:(X) = 2; +
Vn X sm
(259)
h=1
(260)
Here, T. denotes the ith output (or ith column of the output) [Ti
(x(
2-layer, rank-r transformer with concatenated positional encodings is a function T I Rdx parameterized by positional encoding matrix E BA.xN and two (d + dim-ennnno self-masked transformer layers, T() and T(2), and output-layer matrix A E Rdx(d+d)) and defined as follows:
A (T())
X E
(261)
Note that in the above definition, we consider the Nth output point of the transformer to be the output of the model as whole.
E.2
LEMMAS
To prove Theorems 5 and 32 we will need several lemmas.
The first shows that for fixed set of inputs, drawing rank-1 head randomly will have the same output as the target f with probability slightly larger than 1.2 This lemma justifies our majority voting strategy.
Lemma 34. Fix 11, 22, y E Sd-1 with (x) x2, 3) have that:
2 a for some a 0. Then for > Ci we
Pr (0-1((84-1)
arg max (x,, q) (y. q) = arg max (a.,, y) 2 +
1
a
(262)
Vd
for some universal constants C12 C2 >
Proof. In the proof, all probabilities are for q ~ U(d-1),, thus we omit this notation. Denote w. = #1 22 and assume w.l.o.g that (we y) 0, the other direction similar. We can write:
Pr arg max (x, q) (y. q) arg mrg (x. vi
= Pr (sgn( (wo y)) sgn( (w. q) (y.)))
=Pr ((w, q) (y. q) > 0)
Since the above probability is rotation invariant w.r.t q, can assume w.l.o.g that w ei. Hence
we can write
where ÿ E IR--1 and ā = (w, v). Thus, the above probability is equal to:
y
Pr (ā) + (q. 7)) 0)
(263)
Pr Pr(q (āq + (q. 7) > 0/q > 0) (q1 (āq (q. 9)) 0lgi < <0)
(264)
2
Pr + (q. 9) 0/g 0) | Pr (āqı + (q. < 0/gi 0)
(265)
-Pr (āg + (q. Ø) 0/q1 0)
(266)
48

---

Published as a conference paper at ICLR 2025
where the last equality is by the symmetry of the distribution of q. Note that if (q. y) 0 which happens w.p 2. then the term inside the above probability is positive. Hence, we can write:
Pr (āqI
(q. v)
0lq1 > 0)
1
+
Pr (āqı
+ (9. 5) 0/4 > 0, (q. F) <0)
2 2
1
28
(q. B) I 0) Pr (q.
à Vi (q.3) <0)
22 Pr 091
91
(267)
vd
We will now lower bound each probability separately. First, note that if we sample u ~ N (0, / then Tull has the same distribution as q1- By the concentration of the norm of Gaussian random vari-
ables (see Vershynin 2018)) Section 3.1), there is a constant > 0 such that w.p > 1 - exp(-cid)
we have ||/||| E [0.9, 1.1]. There is also a constant C2 E (0, {) such that Pr #1 > 3 C2. This
bounds the first probability term in Equation (267). For the second term, note that ||]| By the same reasoning above can write:
I || = 1.
Pr (1(a.))
8)
à
(q.
ÿ)
< 0 2 Pr
(q. Ø)
I Vd
(q. g) <0)
(268)
vd
=Pr
(1)
U/
2 (1 exp(-cid) Pr #2
|12) S
a 0.9
(269)
|
vd
The above probability is bounded by erf
4-0.9
a-0.9 vd
where this inequality is since erf(z) >
vd
for = E [0, 2] In total, we can bound this probability by
2 (1 - exp((c(()) 0.9
Pr (q. (q.9))
(q g <
(270)
vd
We take d ē so that exp(-ci S 2. Combining the two bounds, and changing the universal constant finishes the proof.
The following lemma shows that a random draw of inputs will satisfy certain condition which allows the use of the previous lemma.
Lemma 35. Let c 0, then:
Pr ayn.aayy-UUnittt- (1 x2, y)
c) S (1 exp(-cid))- 2eVd,
(271)
where ej > 0 is some universal constant.
Proof. By the symmetry of the distribution, we can assume w.l.o.g that y e1 Also, note that for u, v N (0, I), we can view the distribution of (æ1)1 and (æ22) as an and us Combining the above, we get that:
( [1
21
Pr 22. unnit (1 (z) 22, y)
S c) = Pr N(..0.
S
(272)
|ull
|0
By the concentration of the norm of normal random vectors (see Vershynin (2018) section 3.1)
have w.p > 1 1- --xepp-ccdde that [|w]] |||||| < 1.1 for some universal constant >> Also 2 = 41
U1 ~ N (0, 2)- Hence, the above probability can be upper bounded by Pr (o. (1æl < 1.1c) S
erf <v)) Note that erf(x) < 2r for every x > 0, hence the above probability can be bounded by (1 exp( exp(-c1d)) 2d
Li
The following lemma shows a construction of the majority function over H input vectors. This construction uses an approximation of the inner product of two inputs using ReL.U network.
49

---

Published as a conference paper at ICLR 2025
Lemma 36. Let U1,... Hu E {-+, C R where (x_, < 0.1. Let v² be the mode of
V..... vH. Then there exists a 4-layer feedforward network g [R(H+2) + I with width df for some universal constant e and weights bounded by 2 such that
11
11
U
(273)
UH
UH
x
x
x
Proof. Let a be the finally d coordinate of v := (0]
L- x=1 E Rd((++2), and let be the second last block of d coordinates of v. Note that either 2.4 and a_ or the other way around. We construct network that calculates the inner product between x and each e; up to accuracy of 101r By Lemma 37 there is such 2-layer network M1 Rd( - 2d+ with width cd2 H for some universal constant >> O and weights bounded by 2. We add 2d more neurons which act as two identity matrices to keep the last 2d coordinates of v. We add an additional output layer to M1 which sums all the outputs of the inner products.
We now construct another network M₂ Rdd+1 + Rdd which either output if the sums of the inner product is larger than 0.2 H or á otherwise. Note that by assumption that (æ,, 22) S 0.1, M2 will output the mode of the v; ss This is because M1 calculates inner products up to an error of 1077 summing over such inner products returns the exact sum plus an error which is bounded by TO-- Composing M1 and M2 provides an MLP which will output either x_ x_ depending who the mode.
The total width of the network cad2 H, since calculate inner products up to error of 1O and the depth of the network is 4.
We next show that shallow neural networks can approximately compute the inner product of two vectors.
Lemma 37. Let > O. There exists 2-layer network N (sd-1)2 R with width ed and weights bounded by 2 that calculates (x,x') to accuracy e. Here > O is some universal constant.
Proof. By Lemma 6 in Daniely (0177) there exists a depth 2 network Nequore R that calculates x²
in [-2, 2] with an error of 5. width of at most 32d and weights bounded by 2. For each coordinate
2
E [d] we compose the linear function (z); + (ax), with N Naguare get depth 2 network that
calculates ((m))+ up to error of 5. Summing over these networks for every index i and
subtracting results in network that calculates x²) with an error of and width 32d
L
Finally, the following lemma shows that if we draw random rank-1 attention heads, taking their "majority vote" will approximate the target function f. The rate of approximation depends on the number of sampled heads and on the input dimension.
Lemma 38. Let M : (Rd) + IRd be the majority function over vectors in Rd. Namely, given set of H vectors, M outputs the vector which appears the most times in the set, and breaks ties randomly. For a vector qu define gb. (2 x2; = arg max,, (æt; qu)- (ny- qa.). There exist universal constants C1 C2 such that H sad", then with probability at least exp(c2d) samples q1..... q8 Unif(Sd-1) )), we have that:
E 28 10 Umi [1][(21, 22: M (f9(2),22;0) S6,
(274)
Here, f is defined as in Equation (3),
Proof. Fix 22, y with (x) x2, 3) 2 e. Denote by As the event over sampling q ~ Uni which output if arg max; (Ris9k)) (y,qn) arg max; (a,, y) and 0 otherwise.
50

---

Published as a conference paper at ICLR 2025
=1) 2 va if d c1 for some universal constants C1 C2 > O. Note that the events {AA are independent when 21, 22 are fixed. Hence, we can use Hoeffding's inequality:
By Lemma 34 have that Pr(AA, =
H
1
21 S 2exp(-2H1 ! C2 Val
Pr Pr
An
(275)
H
1=1
By setting & <le and H va
d²
we get that:
1
Pr
An
S 2 exxepp-22p)
(276)
2
h=1
From now on, we condition on the event that qi..... qp are sampled such that 1 Σ AA 2 }. which happens w.p > 1 2 functions OA (21, 22, y) will output the same vector as f21 22, y)
Note that if this event happens, then the majority of the
By Lemma 35 we have that (r (ar - 22, y) S e) S (1
26 for some universal
constant cg > 0. Hence, we get that:
21 21.02 U [[]( 22, 8) M (19(2), 22, y) 1
(277)
=Pr (1 (x1 - x2, 8) ()) F
| [[ x2. - M (1g(x1, w/m= 11|
(x1 - x2, 1 1 S
+
(278)
+Pr (1 (xt - -02, 3) 2 2) E [ 22, y) - M (1g(01, 22 y) )) (20 x2, y)
(279)
<(1 exp-x(((-) 2evd 1+1 exp(-202d) S evd
(280)
where we choose d large enough such that 1 exp( ccdd) 2 1 and the constant > 0 accordingly. Replacing e with ova finishes the proof.
I 1/2 and changed
E.3 PROOF OF THEFREM 32
Proof. By Lemma 38 there exist q ....... .H. E gd-1 such that if H 2
cgd
U [|] x2, y) - M (1g(11, x2, / |1 Sc
E
(281)
where gh ( 3) arg maxx, (æ; qn)) (x.qr)) and M is the majority function. We can take H 2 instead of H by increasing the constant by a factor of at most 2.
which will be defined later. We also pick some qo E gd-1 and define M aqood and M -aqod- Note that if go & (211, 222, y) and arg max,, a 2 then arg max,, x, MH H = 2. and vice versa.
We define M. aq.q. for 1,..., H 2 for some
Let g Rdd + Rd the 4-yyer network with width of H as defined in Lemma 36 which simu- X am(X" Miy)
arg maxe, (x. Miy)
lates the majority. Denote by :
and by Umax
X sm(X MHV)
arg maxe; (z.) MHY)
We have that:
E m2 Uniif((dd)
[If( x2: g [+)1
E 21 -niis--)))((
[1(( v) (Umax))
E [llo (Vmax) 9 (w))1?
an an..mayy-Unii amnnit-
(282)
51

---

Published as a conference paper at ICLR 2025
We will bound each term separately. For the first term in Equation (282) we can write:
(2 022 y) g (Umax) 11
(283)
E
(2 22; y) (Umax) ) || (2), x22) - 0.1
(r((2), 22) S 0.1)+
(284)
+ E [|] 225 y) 9 (Umax)
)
(2) x2)) > 0.1 Pr((21, 222) ) > 0.1)
(285)
By Lemma 38 the first term is bounded by c. For the second term, note that (a 222 y) g (Umax)ll < 2 since the output of each function is unit vector. Also, by stan- dard concentration of random vectors on the unit sphere (see Section Vershynin (2018)), (201))) there is universal constant C3 0 such that Pr((x1, x2) > 0.1) I exp(-c3d). Hence, we can bound
y) (Pmas))l )1( I E + 2 expp(ccd))
We will bound the second term Equation (282) uniformly for any 21, 22. y. Note that g ReLU neural network with 4 layers, width H and weights bounded by 2. Hence, we can bound its Lipschitz constant by the multiplication of the Frobenius norm of weights matrices, which bounded by (4(c1d +) Hence:
1/g (Umax) (v)ll² < (4(cIe ||mmax wll
(286)
s (((c)d(1)) H max x xmm(X Myy) arg max(x, May)
(287)
There is > which depends on such that for the set:
As = {x1, 22 y E gd-1, Vgh. (x) xx2) quqTTv y 8}
(288)
have that Pr((01, 22, v) $ As) (M(1)a Note that X sm(a qh y)
arg max (o.. qu y) uniformly on As for every qh- Hence, we can find a > large enough such that:
2 $ (H(4))HH)))) max X sm(X (X M,y) arg max May)
sup 21 ME84-1
(289)
This bounds Em, v..if-- [lg (Mmax) g (w))l Sc
Combining both bounds from Equation (282) we have:
E 21 (4 [a) #2: v) و (w)]l Se exp(-cad)
(290)
where we changed the constant 63 accordingly.
E.4 PROOF THEOREM
Proof. We first define the construction. Let qi..... Qrr be such that the conclusions of Lemma 38 are satisfied (e.g. by drawing them uniformly from the unit sphere). Let E
We call
the second dimension of the positional encodings the "scratch space". We construct the heads of the first layer as follows: For each h, let
qh M 1) a
0 0
0
qr
0
O 1
0]
(291)
0
The number of heads in the first layer is H. The weights of the second layer of the transformer are defined as:
0
M
(2)
1
0
TOTO
0
v
e
(292)
0
52

---

Published as a conference paper at ICLR 2025
for the standard basis vectors c,, and | > 0 will be defined later. The number of heads in the second layer is d. Finally, we set the output layer as = La 0j.
We will prove the correctness of the construction. For the following argument, assume that each head uses hardmax instead of softmax. Note that by a similar argument used the proof of Theorem 32 this incurs extra loss of for any 0 the cost of increasing a.
When the first layer is applied to the input y. the scratch space of the output of each head aj 9%.9 y 22 qh 4n y and -1 otherwise. Let Sy, Smy Sa; be the sum of the scratch spaces of all the H heads (we will in fact only use 8y). Note that 8y > 0 if the majority of the heads outputted at and Sy if the majority outputted for 2t2. All other dimensions of the output are O. Thus, after the skip connection, the output of the first layer is
(1
22
y 0
)
32
y
T(1
1
-1
0
(293)
Sw
Sw2
Sy.
For the second layer of attention, note that each head attends æ1 sy > and æ2 otherwise. By summing d such heads, where each head corresponds to some standard basis vector, the output of the second layer is
22 - 1014
)
y
((2)
T())
- 3
(294)
Sy O, and the same with 02 if sy > Finally, the after the output layer, the output of the entire transformer is By æ1 if 8y > 0, B y 22 otherwise.
By taking first take 3 > L we get that the output of the transformer' is the same as the output of the majority of the rank-1 attention heads of the first layer of the transformer, up extra error of By Lemma and taking the number of heads H to be large enough, we get that the majority of the heads in the first layer approximates the target up to an error of Scaling appropriately finishes the proof.
53


Reviews:
[{"review_id": "KG3dgGP1RX", "reviewer": "Reviewer_zXCj", "summary": "This paper presents a theoretical analysis of the role of rank within attention mechanisms. It challenges the prevailing practice of employing low-rank attention and discusses the implications related to the selection of the number of heads. The author establishes that low-rank attention exhibits inferior performance compared to full-rank attention, indicating that the adoption of a higher rank has the potential to enhance attention performance. Preliminary experiments are conducted utilizing toy examples with synthetic data.", "questions": "N/A", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The theoretical analysis seems correct.", "weaknesses": "1. The rank of attention is a significant hyperparameter in the design of transformers. A common convention involves the utilization of low-rank attention, typically establishing the number of heads as ( H = d/r ). This paper, however, contests this design choice, proposing that a higher rank can enhance performance. It is crucial to note that the paper does not address the speed-accuracy trade-off associated with this adjustment. It is widely recognized that high-rank attention may yield superior performance at the expense of increased computational costs. When evaluating overall performance, particularly in terms of accuracy within a predetermined computational budget, prevailing practices may ultimately provide more favorable outcomes.\n\n2. The experiments presented in this study lack robustness, as they are primarily limited to toy experiments. I would appreciate observing performance metrics derived from real-world data applied to standard transformer sizes. It is well established that theoretical performance often diverges from practical outcomes in deep learning; thus, empirical experimentation is essential.\n\n3. This work indicates that shallow transformers may experience limitations due to low-rank attention. However, it is imperative to ascertain how these limitations manifest in deep transformers, as shallow transformers are not commonly employed in practice. If this limitation has been substantially mitigated in deep transformers, it may render further examination of this issue unnecessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "nTimlQ7OFJ", "reviewer": "Reviewer_8W3p", "summary": "This article examines the limitations and potential of low-rank attention mechanisms in transformer models, demonstrating that while low-rank attention heads require significantly more heads to match the performance of full-rank models in approximating functions like nearest neighbors, and these limitations can be mitigated by increasing the depth of the model. Through theoretical analysis and empirical validation, the study highlights that full-rank models inherently possess superior representational capabilities, especially with fewer heads, and suggests that adding more layers could partly overcome the deficiencies of low-rank models, though at the cost of increased computational complexity.", "questions": "1. Could simple experiments or additional references to other studies and conclusions be designed to intuitively show the impact of the low-rank problem on the performance of mainstream Transformer models?\n2. Could you further elaborate on how the proposed “majority voting” method for improving low-rank models enhances mainstream Transformer models and validate this with relevant experiments? For the experiments, model selection could refer to those in Appendix B.1 and the models used in [1][2], while the datasets could refer to those in [1][2] or other widely recognized benchmark datasets.\n[1] Bhojanapalli S, Yun C, Rawat A S, et al. Low-rank bottleneck in multi-head attention models[C]//International conference on machine learning. PMLR, 2020: 864-873.\n[2] Shazeer N, Lan Z, Cheng Y, et al. Talking-heads attention[J]. arXiv preprint arXiv:2003.02436, 2020.\n3. Also, please refer to weaknesses for other concerns.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A deeper exploration of the low-rank problem in Transformer models.\n2. The paper is well written and easy to follow.\n3. Authors provide ample mathematical proofs to support their conclusions.", "weaknesses": "1. The authors mentioned after Theorem 2 that the theoretical framework should be extendable to cases where N>2. Could you provide more specific explanations for the reasoning behind this inference? This would help further understand the applicability of your theory to specific problems.\n\n2. Although the authors have demonstrated theoretically and experimentally that low-rank attention models are insufficient for fitting certain functions in various scenarios and are significantly weaker than full-rank attention models, further clarification is needed on how these issues impact current mainstream Transformer models (such as the new models shown in Table in Appendix B.1), how the proposed methods in the paper apply to these models, and how performance improvements are achieved. I believe that related experimental results and methodological extensions would greatly help illustrate the contribution of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "4lPvENqrqh", "reviewer": "Reviewer_ugwf", "summary": "The paper \"THE LOW-RANK BOTTLENECK IN ATTENTION\" investigates the impact of the rank of attention matrices on the representational capacity of attention-based mechanisms, particularly in transformers. It challenges the common practice of using low-rank attention and proposes that the rank can significantly influence the model's ability to approximate certain functions. Specifically, the authors present a simple and natural target function based on nearest neighbor search that can be represented using a single full-rank attention head for any context length. The paper presents theoretical analysis and empirical experiments to support its claims, suggesting that increasing the rank or the number of attention heads may lead to more expressive and parameter-efficient models.", "questions": "Please refer to the Weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel perspective on attention mechanisms: The paper offers a fresh perspective on the role of rank in attention mechanisms by using a simple and natural target function based on nearest neighbor search that can be represented using a single full-rank attention head for any context length, which is an interesting aspect of transformer architectures.\n\n2. Theoretical and Empirical Rigor: It combines theoretical proofs with empirical experiments, providing a robust exploration of the implications of low-rank attention on model capacity and efficiency.", "weaknesses": "1. The results may rely heavily on the assumption of rotational invariance in the data distribution, which may not hold in all real-world applications.\n\n2. To make it easier for readers to understand, I kindly suggest that the authors explain in more detail the differences between this paper and previous work [1].\n[1] Low-Rank Bottleneck in Multi-head Attention Models. ICML 2020\n\n3. Can the proposed method demonstrate its effectiveness on more attention-based models?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "AW1E0s76dR", "reviewer": "Reviewer_aLKP", "summary": "This paper questions the conventional number that community use for the rank of attention matrices and the number of attention heads. Authors provide in-depth theoretical explanations and experiments to support their arguments. In the high-accuracy regime, the required number of heads is growing exponentially to remain the performance.", "questions": "1. I am wondering the reason why authors choose to analysis nearest neighbor functions and are there any other choices of functions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Transformer have dominated many areas but there have been few studies on the choices of numbers of attention heads and dimensions used in attention mechanism. This paper raises doubts about this which is valuable for community to pay attention.\n\n1. The paper is well written and easy to follow. In-depth theoretical explanations are provided.\n\n2. For a simple and natural target function -- nearest neighbor function, authors show low-rank attention is fundamentally weaker than full-rank attention even when choosing very large head numbers.\n\n3. Also, this paper explores the solutions to mitigate the weakness of low-rank attention", "weaknesses": "1. The paper studies only is  limited to shallow transformers which are not practical to large model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": {}}]

Author rebuttal:

Decision: Accept (Poster)