What can a Single Attention Layer Learn? A Study Through the Random Features Lens
Authors: Hengyu Fu
Tianyu Guo
Yu Bai
Song Mei
Venue: NeurIPS 2023 poster (2023)
Keywords: transformers
attention
deep learning theory
random features
TL;DR: We study the expressivity of a single random-feature attention layer, and derive concrete learning results from finite samples.

Abstract: Attention layers---which map a sequence of inputs to a sequence of outputs---are core building blocks of the Transformer architecture which has achieved significant breakthroughs in modern artificial intelligence. This paper presents a rigorous theoretical study on the learning and generalization of a single multi-head attention layer, with a sequence of key vectors and a separate query vector as input. We consider the random feature setting where the attention layer has a large number of heads, with randomly sampled frozen query and key matrices, and trainable value matrices. We show that such a random-feature attention layer can express a broad class of target functions that are permutation invariant to the key vectors.  We further provide quantitative excess risk bounds for learning these target functions from finite samples, using random feature attention with finitely many heads.

Our results feature several implications unique to the attention structure compared with existing random features theory for neural networks, such as (1) Advantages in the sample complexity over standard two-layer random-feature networks;  (2) Concrete and natural classes of functions that can be learned efficiently by a random-feature attention layer; and (3) The effect of the sampling distribution of the query-key weight matrix (the product of the query and key matrix), where Gaussian random weights with a non-zero mean result in better sample complexities over the zero-mean counterpart for learning certain natural target functions. Experiments on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.

What can a Single Attention Layer Learn? A Study Through the Random Features Lens
Hengyu Fu" Peking University
Tianyu Guo" UC Berkeley Llanyu.guo@berkeleey. edu
2100010881@stu, pku. edu. on
Yu Bai
Song Mei UC Berkeley songmei@berkeley.ccm edu
Salesforce Al Research yu. baissalesforce. com
Abstract
Attention layers- -which map a sequence inputs to a sequence of outputs are building blocks of the Transformer architecture which has achieved significant breakthroughs in modern artificial intelligence. This paper presents a rigorous theoretical study the learning and generalization of single multi-head attention layer, with activation function replaced by ReLU, which have recently shown comparable performance with the original Softmax activation. We consider the random feature setting where the attention layer has large number of heads, with randomly sampled frozen query and key matrices, and trainable value matrices. We show that such a mandom-fetturg attention layer can express a broad class target functions that are permutation invariant the key vectors. We further provide quantitative excess risk bounds for learning these target functions from finite samples, using random feature attention with finitely many heads.
Our results feature several implications unique to the attention structure compared with existing random features theory for neural networks, such as (1) Advantages over standard fully connected random-feature models; (2) Concrete and natural classes of functions that be leamed efficiently layer. Additionally, we show that the sampling distribution of the query-key matrix (the product of the query and key matrix) matters-A biased Gaussian random matrix results in better sample complexities over the standard zero-mean counterpart for learning certain natural target functions. Eetrrionn..xpemmentts on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.
random-feature attention
Introduction
The transformer architecture [86] has achieved remarkable recent successes in many of artificial intelligence (AI) such as vision, language, speech, graph processing, reinforcement learning, and more recently general AI capabilities [29-31, 19, 73, 101, 22, 77, 65, 20). A central building block in transformers is the attention layers [10] [eo]=ssquenneetoooeee se-quencce mappings that allow each token "attend to" other tokens that most relevant to the present token, and
within the input sequence produce outputs based on those tokens. Attention layers implement this mechanism in compact way that allows them to handle sequences of arbitrary length using a fixed of parameters, crucial reason behind their success in handling long input sequences.
Equul contributions.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Despite its wide applicability, the theoretical properties of attention layers are less well understood. While multi-layer attention networks (transformers) have been shown be universal approximators for certain classes of functions, such as equivariant requence-to-sequeee only focus on the expressive power and do not account for learning from finite samples. Another line of work derives generalization bounds for learning with multi-layer transformers in terms of the number of layers, heads, and weight norms [91, 35J, yet the results are either instantiated on specific target functions such as sparse boolean functions [35], or generic but arguably clusive function classes such as Turing machines [91]. Understandings about the more basic building block-a single attention layer- remain largely open. This is in stark contrast with the situation for fully connected neural networks, where there is by now a decent understanding of the learning and generalization in the important basic case of two-layer neural networks generic and natural function classes (e.g., [8, 51, 7, 40] and many other results along the line). This motivates the following open question:
functions [102], their results
What function classes can be learned by a single attention layer with benign sample complexities?
This work makes progress on this problem by studying the learning and generalization with a single attention layer in the random feature setting [74, 76, 28, 100], in which the query and key matrices are frozen their random initialization, and the value matrices remain to be learnable parameters. Motivated by the attention structure in practical architectures, we consider attention layers that take in a single query token Xy Rd and N key tokens {x, ][E[N] as the input, and produce a scalar-valued output-A simplified setting capturing the essence (the interaction between the query and keys) of attention models. We study the sample complexity of learning certain target functions (of Xo:N) using an attention layer with a large but finite number of heads, and finitely many samples.
Our contributions are summarized as follows.
We show that a Random Feature Attention layer (the RFA model) with sufficiently large number of heads can express a broad class of target functions that are averages over a generic function of two tokens, which are in particular permutation invariant with respect to the key tokens (Section 3.1). We give several natural examples of target functions in this class (Section 3.3) with concrete bounds on the number of heads and weight norms.
We derive D(VB(f.)mm excess risk bound for learning with the RFA model with sufficiently many heads, where B(f-) is an inherent complexity measure of the target function f. and n is the sample size (Section 3.2). When instantiated on the aforementioned examples, the bounds only depend on the imput dimension and not the number of key tokens, improving over a naive two-layer random feature neural network model (RFMLP). Such improvement is expected due to the permutation invariance structure of target functions, aligning with the attention mechanism.
Towards moving beyond standard random feature settings, we study a biased RFA model where the query-key matrices (product of transposed query matrices and key matrices) are drawn from a distribution with non-zero mean (more precisely the identity matrix as the mean), motivated by a similar observation on learned attention layers in practice. We show that this model achieves provably superior sample complexity than the standard mero RFA for learning certain functions of the correlations between the query token and key tokens (Section 4).
Experiments on simulated data verify our theoretical findings in realistic settings of learning from finite samples using a RFA layer with mild number of heads, and characterize the interplay between the complexity of the target function and the sample size (Section 5).
1.1 Related work
The Transformer architecture, initially proposed by [86]. brought about revolu- tionary change in natural language processing and has been widely adopted in large language models such GPT and BERT [72, 29, 19]. At the core of transformers lie the attention layers. which were originally introduced as neural network modules for machine translation tasks [10, 49, 68].
Transformers
A recent line of work investigated the capabilities of transformers by viewing transformers to be function approximators [102], computational models [91, 70, 99, 15, 56), or algorithms [82], and using transformers to perform synthetic reasoning tasks [103]. Among these works, the closest to our work is [102], which shows that multi-layer transformers can approximate any permutation- equivariant pequence-to-sqquena function, and any function if positional encodings are added. Our

---

paper instead uses single attention layer to approximate sequence-to-scalle functions and focuses on the generalization property with quantitative bounds.
In terms of generalization properties of transformers, [35] analyzed the generalization bound of a single attention network through the Rademacher complexity and showed that single self-attention head could efficiently represent a sparse function of the input sequence. Besides, several works studied the sample complexities of vision transformers [48, 50], and prompt-tuning using attention [67] with special target function classes. Our paper also studies the generalization bound of a single attention network, but from the different perspective of kernel methods, and for a more general class of target functions. The kernel limit of transformers was derived in [43, 98], which shows that multi-head attention architectures behave as Gaussian processes the number of heads tends to infinity. However, they do not study the representation power of the limiting kernel.
Besides approximation and generalization capabilities, recent work also studied the limitations [42, 14), internal working mechanisms [36, 81, 94, 64), and in-context learning capabilities [19, 96, 37, 88, 3, 26, 41, 52] of Transformer models.
Theory of random features and neural tangent kernels A recent line of work [28, 51, 33, 32, 6, 7, 104, 66, 25] studied the training dynamics of overparametrized neural networks under certain random initialization, and showed that it converges to a kernel estimator, which corresponds to the "neural tangent kernel" (NTK). These works suggested that one could use kernel uundom-featur models [74] to study the properties of deep neural networks.
For NTK of MLPs and their corresponding random-feature models, there is a vast number of literature that studies their approximation power [12, 71, 8), as well as their generalization properties [13, 21, 93, 92, 54, 55, 76, 79, 100, 57, 60, 40, 39]. More recently, line of work studies the NTK beyond MLPs, including convolution networks [53, 17, 59, 62, 18, 16], residual networks [45, 83, 4), graph networks 197, 47], and transformers [43, 98].
Although the kernel approach is a powerful tool for studying neural networks, received criticism since it does not capture the feature learning of neural networks. Going beyond the kernel regime, a series of works used the mean field method to establish the evolution of the network parameters via
a Wasserstein gradient flow [58, 9, 24, 78]. Several other mechanisms have been proven to obtain superior results over the NTK, including the Quadratic NTK [5, 11, 23, 63], regularization [90], Neural Tangent Hierarchy [34, 44], representation learning [27], and staircase- like mechanisms [1, 2].
2 Preliminaries
We consider sequence of N + input tokens Xp:N (xo, (xIlieN]) E X (Rd)N+1, where each {x; hejai S IRd represents a sequence key vectors, and xo represents the query vector. This model simplifies standard self-attention, which maps N input tokens to N output tokens, where the output i only uses input as the query token and all of [N] as key tokens. Results obtained this model can be directly mapped back to full self-attention, by simply concatenating N outputs generated by model with Xo ranging over {x (x<heeN)) In addition, throughout the paper, we consider scalar- valued attention models, which take the sequence N as input and give a scalar output in R.
Attention layer We consider a scalar-valued, M-head, multiplicative attention layer that takes Np.v = (xp, (etn)) as the input. The attention layer first applies affine (linear with bias) transformations to the input vectors to obtain (query, keys, values) each head E [M]:
9m,0 = Qm Xx; 1] =: Qm Xy E RR,
km. = Km [x;; 1] K X; E Rdd,
(1)
Um,i vn [x;; 1] = v Txx E R [N],
where Qm: Km E (d) Vm E Rd++1 are the parameters the attention layer, and x; == [xxx]] 1] for more compact display. Then, it computes the output value by an attention mechanism
f(xp.N) Σ
Σ fm. x1), f x;) = a(qm,, km,))) Um. E R. (2)
Above, a IR + an activation function applied entry-wisely to each attention score (qm,o; o km.))- We choose to be the ReLU activation ott) max (t, 0} throughout this paper. Notice that
3

---

this choice of the attention non-linearity is different from standard transformers [86] with softmax attention. We remark that we choose study the (normalized) ReLU attention for theoretical convenience, and this replacement does not change the essence of the attention mechanism. Such choice also recently explored in the literatures such as [80] and [95], which show that transformers with Rel.U-attention perform as well as standard softmax-attentin
transformers in certain NLP and CV tasks. Moreover, our results can extend to other activation functions such as the exponential activation (t) exp(t), which is more similar to the standard Softmax activation. We refer to Section B.5 for short discussion.
Simplifying the expression, we reparametrize the attention layer (1) and (2) using parameters {(W my Vm)) vm)]me[[n) C R(d+1)x(d+1) x TRu+1;
= o Qm (.): = (( mm x,x,) (Vm x,)
(3)
For technical convenience, we assume all input tokens have unit norm throughout the rest of the paper: ||x:|| ||2 so that |2 = v2, for all E {0} [N].
Candom-feature attention models We consider a random- version2 of the multiplicative attention mechanism (3), where the weight matrices {
have i.i.d. Gaussian entries*:
(W ma) 1/4), (m, E [M] x [d + 1)².
The variance chosen to be 1/4 without loss of generality: this choice of variance is such that (Wm: Xox; ) ~ N(O, 1) has unit variance. The weight matrices (Wm)me are then held to be fixed during the entire learning process, whereas the value vectors {vm [M][M]] are the learnable parameters. The rantore attention model with input Xo.N is thus given by
MM (x V) =
-M Em= EN o (W... xxxT))
(V),
(5)
Notice that random-feature attention model is linear in the parameter V, so training this model with convex loss function gives a convex optimization problem.
Additional notation For any x E Ted and y E RRd,, let & Te x denote their tensor product (outer product), and S x X ... @ denote the n-fold self tensor product of x. For tensor A, we use || to denote its Frobenius norm. For a function f X - R, we use |l to denote its d norm. We use O--) (resp. 6(-)) for standard Big-O (resp. Big-Theta) relations. We use (-) for hiding the multiplicative terms that are logarithmic in problem parameters, including (M, d, n,, N, 8 1-)). We use (ply(p) to denote polynomial of that is less than p for some universal constant 0 C OC.
Learning with ranuor-feature attention models
In this section, we study the expressivity and generalization of madomm-featree attention models. We will consider a broad class of target functions that can be well approximated and is efficiently learnable by random-feature attention models.
3.1 Expressivity of random-feature attention
Consider broad class of permutation invariant target functions f. : X - IR that takes form
f. fx = IN ((xx,x)))
(6)
AA different and closely related model is the Neural Tangent Kernel [46, 32], which is however similar in essence to the random feature model terms of sample complexity of learning, e.g. [40].
Annother feasible choice for (4) is to sample the key matrix and the query matrix separately with independent Gaussian entries, which however will produce mean zero product matrix similar to (4) in many aspects.
A function [Xo, X1, ..., XN) ) is permutation invariant (with respect to X1.N) ) fXa, Xi,..., xx)
f [Xx, Xo(1), ..., xx))) for any permutation [N] + [N]. We consider permutation invariant target functions since attention layers only fit these functions due the structure of attention models.

---

Assume that there exists symmetric tensors {fre E R Tr,200 such that F R2d 7 R admits representation
or F(xo: x:) Lr,000 x @ frs).
Note that such an expression allows F(xo, x,) to be any general nonlinear function that admits convergent Taylor expansions. In particular, any polynomials of [Xo. x,] (e-g.. B Xo, B x; and Sx;) for some 3 E IRd and E IRd are within this function class. We will discuss more specific target functions in Section 3.3.
Theorem (Expressivity of REA model). Suppose function f+ : X + IR takes form (6). Then for any imput distribution P X, with probability least 8 (over { {W m mee| sampled from (4)), there exists an M M-head RFA A model (5) with coefficients V = f. (P)) up to error
C Rd+1 that approximates
Ex (fo - SM (XoiN;
so v
(d log log Mog M
(8)
In addition, the norms of the weight of this random-featur attention model are bounded as
Σ Ivmll2 I =
M
B(f+)6-
M
|/ 11 20
B(I+)6-1)
(9)
M
M
m=1
m==
Here B(f.) is a complexity measure of fa defined as
B(J.) = Ci Em l[Coollir:
C2 ===kh kk-54 V
(10)
In case where f. admits multiple representations of the form (7), B(f.) is the infimum of the right-hand-side over all such representations.
The proof of Theorem is contained in Appendix B.1. Our proof relies on standard analyses of infinite-widhe random feature model with LeaU-Gaun kernel, combined with sampling argument to obtain approximation with finite-width.
This theorem is applicable to general functions with finite B(f+) norm. The 4 scaling of Cje in the summand of equation (10) seemingly confines the target function class those with exponentially fast decaying ||f///Fr. which suggests relatively narrow target function class. However, we will demonstrate in the forthcoming examples, this class includes a diverse range of functions.
3.2 Generalization and sample complexity of learning
Given n samples {xo:N= Mj } je(n] P. where x() = {xx" JosisN is the j-th token sequence with length + 1, and Vy is the label corresponding to the i-th token sequence. Assume that we are given loss function E(T, 3) that is 1-Lipschitz in y. and E(0, s 1 for any y. The population risk is then given by LD(f) = (x [[(P (f((f(()))))))))))))))))()))(()) v)]. We consider the empirical risk minimization (ERM)
over the RFA model (5),
arg minvevy TD (s V)), where the constrained class VM is given by
in(s)= Σ (((()))))))
(11)
VN = Mr ΣM |/v/l22 S Ki, EM |/ K2/M
(12)
with K2 and K2 being two constants. Theorem below provides the excess risk bound for the empirical risk minimizer.
Theorem 2. Assume M 5-1 and log(dM). Let fa be the minimizer of the population risk (f( within the target function class (6) (7). Let FW = fM (:: V) be the empirical risk minimizer given by (11), where in (12) we choose K1 = C B(f) and K₂ CB(f=)) with C being constant. Then for any joint distribution P, with probability at least
over {W m [Wm)meja) melad] sampled
according to (4) and (x Us)} j)(j)))))) Midd P. the excess risk is bounded by
L Lo(fa) oō (BBL.)
d88 d'8-1
1)
(13)
M
5

---

The proof of Theorem 2 is contained in Appendix B.2. The proof mostly uses the Rademacher complexity bound for the supremum of empirical process. The main non-trivial technical challenge lies in showing the concentration of sup, ( - (f))), which cannot be simply controlled due the unboundedness of the infinity norm of functions in the target function class Vas- We dealt with this subtlety by carefully decomposition of SuP [Tp(f) [.(()]- The seemingly unnatural constraint set (12) is used in bounding different terms in this decomposition.
3.3 Examples and comparison
We next give the sample complexity for learning several examples of target functions using the random-feature attention model. We will compare its sample complexity for learning these functions with that of the standard random-feature model [74] (thereafter, model, in short RFMLP model). In the RFMLP model, we view X as an input vector instead of a sequence of vectors denoted as vcc(X.:N) = [xo; Xi; 1] E IR [N((+))++ The RFMLP given by
call it the mandom-feature MLP
M
MLP JMI (xo.NiV)
Σ o (Wm, vec(Xa:N))) Um m=1
N I/(N 2)).
(14)
We choose the variance of random weights Wm be 1/(N + 2) ensure that (Wm: vec(Xp:N) N(O, 1) has unit variance. The generalization and approximation properties of the candom-feature MLP model have been well-studied in the literature, for example. [7, 8, 60].
We instantiate Theorem on three concrete examples of target functions (calculations of the excess risks in Appendix B.4, where the result for RFMLP are adapted from Arora al. [7]) In all three cases, the target functions are permutation invariant with respect to {x: by which we naturally expect RFA to achieve better sample complexity than RFMLP in accordance with this structure. Although the comparsion between RFA A and RFMLP is based on comparing upper bounds on the sample complexity of both models, existing work has also derived lower bounds on the sample complexity of RFMLP, which aligns with the upper bound for RFMLP we used. We do not invoke these lower bounds, as they apply to special case with a uniform distributional assumption on the input tokens.
Example 1 (Functions of xo):
We consider functions of X (no dependence X1:N) of the form
fx(x::))
Σ k=
X
Sk Ax), A Rdd with B(S.) C& || by (10).
By Theorem 2, setting M = (ddn), the excess risk bound gives Ö( Ei k- n).
As special case, consider f. (Xo:N) = (B) (8)x))" 'o)",, which corresponds taking A& 0o for p and A₁ for k & p. The above excess risk RFA model and the RFMLP model scales as
RFA : (Pool(l)) 4r 31
RFMLP: Ö (Pooly(p) / (N + 1)11|/1
Compared to the RFMLP model, the RFA model significantly reduces the necessary sample size by factor of (N/4)².
Example 2 (Average of functions of x;): We consider average of functions of x; of the form
(x
1 N
Ab), A E R... with B(f.) - Σ C. |A by (10).
Jmm
Theorem then gives an Ō VV 64 A l2/// excess risk, same as Example 1. O
As specific example, consider fa n Σn (1B, xx)) with (b=)) (arctan(z/1) for some 7) > 2, || 2 = 1. Using the power series expansion v the excess risk bound of RFA model and the RFMLP model scale as
RFA ō o(VE V k45 = Õ(V
REMLP: : 5 O K45[(N + 2)/(22))kk/m /m).
By deriving the correspouding results for Random Features instead of Neural Tangent Kernels.

---

as
Figure 1: Visualization of weight matrices of the 2nd, 5th, 8th, and 11th layers of the BERT-Base model. Each row contains weight matrices of layer. All matrices are clipped to the top-left 32 x 32 block. Lighter color indicates a larger absolute value.
The latter diverges whenever 7) S (N 2)/2, in which case the bound is meaningless.
Example 3 (Correlation Correlatton-weighted functions):
f. is the following function:
[x(x.:)==
F((xo. G(x)))((x), F(() Σ a G(x.)
@k G
K=U
R=U
for S E Raxd, {ax)a20 C IR, G2 Rdd This target function fully exploits the representation power of the attention layer. Eq. (10) gives B(fa) = O Ca
G, (l.)).
N As specific example, consider f1,+ Σ i= (Xo, x,)",, corresponding to taking S La, F(I) t", and G = 1. The excess risk bound of RFA (by Theorem 2) and RFMLP scale as
RFA : (PPol(P))/(ddy))
RFMLP δ
2)dp/n).
&p G) with | 1. Then the As another example, consider f2,- = N Σ cos( (x)) excess risk bound RFA and RFMLP scale
RFA: δ (Poly(dd) / /etvasp/n)
REMLP Ō (Poly(nNaal /
C2N42waND//
RFA reduces the required sample size by factors of (N/4)² for fi. and exp(N Vd) for f2...
Expressivity of biased random-feature attention model
We move beyond the Gaussian weight assumption by exploring alternative possibilities for the weight distribution in the attention heads. We observe empirically that the weight matrices in transformer architectures learned in practice often similar the identity matrix than mean-zero matrix (Figure 1; see the details Appendix D. This is also observed a recent and concurrent work [84]).
Towards understanding this effect, we consider an alternative attention model with biased random weights, where the bias is a fixed matrix Wo E R(d+1)x(d+1),
M
N
W. Wo (Xo:N; V) A
(( + W XoXT)) (Vm, x;)
(15)
m=
1=1
Here (W... m [ MM again Gaussian random matrices sampled according (4). The biased rundom-feature attention model is similar to (5) except that a bias weight Wo added. Motivated by our observation, we choose Wo = [axd 04x1 0 x 0] E [Rd++))((++1) so that the diagonal elements of Wo W will be average larger than the off-diagonal elements.
7

---

4.1 Expressivity of biased random-feature attention
Given the formulation of biased candom-feature attention models (thereafter, we call the biased mundom-feature attention model, in short BRFA A model), a natural conjecture is that this model can better fit functions that are the average of function of (x), We here show that this is indeed the case. In particular, we consider a broad class of target functions X - R that take forms
g. (xoxNN) = 1 Ei=l F((x0, x.) (xxo, x;)
F(t) -x agth, G(xo, x;) = (x) X A.).
(16)
Here the scalars {ax 3200 S and the tensor A, E RR parameterizes gr. As we will explain in Section 4.3, confining G to be degree-(3, 2) polynomial in (x,, Xo) is essential to our theoretical results. Our next theorem provides the excess risk of learning target function g+ using the BRFA model (15).
Theorem 3. Given the same setting and assumptions as Theorem 2, when the population risk minimizer gives J- = 9x, with probability at least 1 5, we have
LD(JM WW. - (D( - ö inf VB 1i (V
P8
+ sn l1smllzo) 1):
(17)
M
where - 11/12++ 1)!! and B L) = ||A.| (Ei=o |aæ| Cx)², with C = (2L + k) k)
(18)
The proof of Theorem 3 is contained in Appendix C. We provide the intuitions of the result and an overview the proof technique in Section 4.3.
4.2 Examples and comparison
Compared to the target functions (7) discussed in Section 3.1, functions in (16) may not express the average of arbitrary functions of Xg and x;, but well-suited to express functions of correlations. Consequently, we anticipate that the BRFA model will outperform the RFA model in learning functions of correlations. We will now present three conerete examples of target functions (16), and compare the excess risk of the BRFA model to that of the RFA model. The proof of excess risk is contained in Appendix C.3.
Example 4 (Low degree polynomials): Consider average of polynomials of x; and Xp.
N
A), with B(g., L) = |, &gg2.. by (18).
N
For any 17 > O, if we take (xp(exp(O(L/n))) L @((1 + log n) log n), and M (d)n), the excess risk will scale as O(VIAA /
Compared with the excess risk of the RFA model as detailed in Example 2, the excess risk bound of the BRFA model loses factor of n n/2
Example 5 (Functions of correlations): Consider a special case of functions of correlations,
1 N
Σ
Xo. x;) ( x1)
B Bes--1 with B(g., L.) (2I + p) by (18).
9+
1=1
For any 17 O, choosing the same parameters (n, L, M) as Example 4, the excess risk bound scales as (Vog n + p)(++3)87 /p/n1--1)
Consider the required sample size n, to reach an accuracy of 0.01. The BRFA model requires n. = Ö(8p + 48)++,) whereas the RFA model requires Ō ((4d))). Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by factor of ((d/((p + 12)]
Example (Correlation- weighted functions): Consider the function
cos( (xo, x:) (x 83 G.
with |G|| and B(g., L) (((8)224)
g.
N
8

---

where (gg.. L) is bounded through the Taylor expansion of cos(t) and (18). For any 77 > 0, choosing the same parameters as Example 4, the excess risk bound scales O (Vv1/n!_>) o
Consider the required sample size n, reach an accuracy of 0.01. The BRFA model requires n, = O1), whereas the RFA model requires n. = (d)ly(dd exp(vd)). Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by a factor of ((Poly(d) exp(Vd)).
4.3
Overview of techniques
Here provide the intuition and an overview of the technique of Theorem 3, with the proof details in Appendix C. To show the sample complexity of learning with the BRFA model, the first step is to derive the kernel K BRFFA (Xo,N, x6.v) associated with the infinite-width BRFA model. This kernel has
natural feature map, given by (Vr: X
Ie td
1:20- where
(x(X):)) = Σi (xxx, x.)) Hee-2((Xo))) xx)),
x,
X
Vk 1 2.
Here @(t) (2T) 1/2e - is the Gaussian density function, and He denotes the k-thh proba- bilist's Hermite polynomial, with detailed expression and properties given in Appendix A.1. This feature map implies the learnability of the following target function class by the BRFA model,
g. (Xo.N) N/4 E (xx0 xx)) Ek=2 Hex-2((xo- xx)) (x) SK+
Ak)
(19)
whose RKHS norm associated with kernel KERFa is bounded by B = 2 1)
Notice that 9. bears similarities to, but also distinct differences from, 9. as presented in (16). The key difference lies in the of (xx,,x)) factor g., which is hard to interpret and analyze, To obtain the excess risk bound for learning Gr. we can use to approximate risk for learning 9+ can be bounded by the summation of the excess risk for learning J+ and the approximation error. Acquiring this approximation error bound necessitates a truncation argument of the Taylor expansion of 1/((-).
in the 7 norm. The excess
Numerical experiments
We test our theory by experimentally approximating two types of target functions using the three models under investigation RFA (5), BRFA (15), and REMLP (14). We choose the target functions to be of form
(xoN) = (x.N) N E (Xo, xx) (B,x))
N ΣM x,)
PEN, q REN, N,
BE B
(20)
B BES-11
(21)
The first target function (20) is a specific instance of Example 2, whereas the second target function (21) has been considered in both Example 3 and 5.
In our experimental setup, we the input dimension as d 16 and the number of tokens as N 16. We fix the width of RFA and BRFA to be MRFA MrRFa = M = 1000, whereas the width of RFMLP set as MAFMLP = M (d + 1) 17000. This configuration ensures an equal number of parameters across all three models. To further accentuate the test risk difference between the BRFA and RFA, in BRFA we use a bias matrix of W. Laxa, 02x1; 01xd; 0] E [R(++)x(ddI) which is four times the matrix investigated in our theory. The input distribution selected {x, JuKsss N ~id (Sd-1) and we take = f- (Xp.N) without any noise. We consider three representative target functions: fI. forp= = 2, 4, and f2, for 212, Prior to training with RF models, we standardize the y, 's to have zero mean and unit standard deviation, ensuring that the trivial risk equals 1. We train the RF models using square loss with ridge regularization, selecting the ridge parameter to minimize the test error. The experimental results are displayed in Figure 2.
3. as per (20) and (21). We examine a list of sample sizes n from 2 to
The left and middle panels of Figure demonstrate a noticeable separation between RFMLP and the other two random-feature attention models for learning these target functions. RFMLP can hardly approximate the target function, whereas RFA and BRFA exhibit significantly better performance. This observation is consistent with sample complexity analysis detailed in Example 2, where the

---

2.2
4.4
6.
1.4
1.4
14
1.3
1.
3.
REA
BRTA
03
100
Test error of three RF models for learning f1.2 (left), f1,4 (mid), and f2.3 (right), as per (20) and (21). We set d, N = 16, MaFA MaRFA = 1000, and MArиL 17000. We train the RF models using square loss with ridge regularization, with the ridge parameter selected minimize test error. The test error is calculated using Neeest 1000 fresh samples. The figure reports the mean and normalized standard error of test error, based on 5 independent experimental instances.
Figure 2:
sample complexity bound of RFMLP for learning average of functions of x; is found to be (((N/4))) times greater than that of RFA.
The performance comparison between RFA and BRFA depends on the target functions. RFA outper- forms BRFA in learning f1.2 and fi,4, whereas BRFA outperforms RFA in learning S2.3. The latter phenomenon is we expected: as demonstrated in Example and 5, BRFA is more powerful than RFA approximating the correlation- weighted functions.
We have conducted further experiments with various other target functions, detailed in Appendix D.
6 Conclusion
In this work, we introduced and examined the expressivity of two kundom-feature attention models, namely RFA (5) and BRFA (15). For general classes of functions that are invariant to the permutation of key tokens XI:N. the excess risk of RFA (5) can avoid the dependence on sequence length, in contrast the standard random-feature model RFMLP (14). Moreover, for specific functions that adopt the form of correlation-weighae -eiggted polynomials (6), the excess risk of BRFA can avoid the polynomial dependence the dimension. These insights enhance within a simplified context. Finally, our work left open many interesting questions for future work, such the expressivity of softmax attention, the influence of positional encoding in expressivity, and the expressivity of multi-layer transformers.
understanding of the attention mechanism
Acknowledgment
Mei is supported part by NSF DMS- 2210827 and NSF CCF-2315725.
References
[1] E. Abbe, F. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
[2] E. Abbe, F. Boix-Adsera, and T. Misiakiewiez. Sgd learning neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXi::2302. 2023.
[3] E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arX.v.2211.16611. 2022.
[4] Z. Allen-Zhu and Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
10

---

[5] Z. Alle Y. Li, and Y. Liang. Learning and generalization in overparameterizee neural networks, going beyond two layers. Advances in neural information processing systems, 32. 2019.
[6] Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[7] S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
[8] F. Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-68) 2017.
[9] F. Bach and L. Chizat. Gradient descent on infinitely wide neural networks: Global convergence and generalization. arXiv preprint arXiv:21 10. 10.08084, 2021.
[10] D. Bahdanau, K. Cho, and Y. Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint 1oXiv:1409.0473, 2014.
[1]] Y. Bai and J. D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:1910.01619: 2019.
[12] A. R. Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions Information theory, 39(3):930 945, 1993.
[13] P. L. Bartlett and Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 46N44)4433-48,, 2002.
[14] S. Bhattamishra, K. Ahuja, and N. Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009. /2244 2020.
[15] S. Bhattamishra, Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09288 2020.
[16] A. Bietti. Approximation and learning with deep convolutional models: kernel perspective. stat, 1050:18, 18, 2022.
[17] A. Bietti and J. Mairal. On the inductive bias of neural tangent kernlss. Advances in Neural Information Processing Systems, 32, 2019.
[18] A. Bietti, I.. Venturi, and J. Bruna. On the sample complexity of learning under geometric stability. Advances neural information processing systems, 34:1 1 18684, 2021.
[19] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-19999, 2020.
[20] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303. 12712, 2023.
[21] A. Caponnetto and F. De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7:331-368, 2007.
[22] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 144115044-5099 2021.
[23] M. Chen, Y. Bai, J. D. Lee, T. Zhao, H. Wang, C. Xiong, and R. Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances Neural Information Processing Systems, 14222134-22144 2020.

---

[24] L. Chizat and F. Bach. On the global convergence of gradient descent for over parameterizee models using optimal transport. Advances neural information processing systems, 31, 2018.
[25] L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
[26] D. Dai, Y. Sun, L. Dong, Y. Hao, Z. Sui, and F. Wei. Why can gpt learn in-context? language models secretly perform gradient descent meta optimizers. arXiv preprint arXiv:2212.10559 2022.
[27] A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
[28] A. Daniely. Sgd learns the conjugate kernel class of the network. Advances in Neural Information Processing Systems, 30, 2017.
[29] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
[30] L. Dong, S. Xu, and B. Xu. Speech- tranch-transformer: a no-recurrence sequence-to- coanccett-seqquee -e-sueecce model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), pages 5884-5888. IEEE, 2018.
[31] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. De- hghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
[32] S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International conference machine learning, pages 1675-1685 PMLR, 2019.
[33] S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over- parameterized neural networks. arXiv preprint urXiv:1810.0200 2018
[34] E. Dyer and G. Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv: 1330, 2019.
[35] B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases and variable creation in self- attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[36] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
[37] S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 15:35583-35995 2022.
[38] R. Ge, J. D. Lee, and T. Ma. Learning me-hidden-llayr neural networks with landscape design. arXiv preprint wrXiv:1711.0055 2017.
[39] B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 33:14820-14833 2020.
[40] B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. 2021.
[41] A. Giannou, S. Rajput, J.-y. Sohn, K. Lee, D. Lee, and D. Papailiopoulos. Looped trans- formers as programmable computers. arXiv preprint arXiv:2301. /3/6, 2023.
[42] M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguisties, 8:156-171, 2020.
12

---

[43] J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nugp and ntk for deep attention networks. International Conference on Machine Learning. pages 4376 -3366 PMLR, 2020.
[44] J. Huang and H.-T. Yau. Dynamies of deep neural networks and neural tangent hierarchy, In International conference on machine learning, pages 4542-4551. PMLR, 2020.
[45] K. Huang, Y. Wang. M. Tao, and T. Zhao. Why do deep residual networks generalize better than deep feedforward networks? neural tangent kernel perspective. Advances in neural information processing systems, 33:2698-2709, 2020.
[46] A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel:
Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[47] S. Jegelka. Theory graph neural networks: Representation and learning. arXiv preprint 0D0::2204.06697, 2022.
[48] S. Jelassi, M. Sander, and Y. Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37833 2022.
[49] Y. Kim, C. Denton, L. Hoang, and A. M. Rush. Structured attention networks. arXiv preprint 1rXiv:1702.0000 2017.
[50] H. Li, M. Wang. S. Liu, and P.-Y. Chen. A theoretical understanding of shallow vision trans- formers: Learning, generalization, and sample complexity. arXiv preprint (xXiv:2302.06015 2023.
[51] Y. Li and Liang. Learning overparameterized neural networks via stochastic gradient descent structured data. Advances in neural information processing systems, 31, 2018.
[52] Y. Li, M. E. Ildiz, D. Papailiopoulos, and Oymak. Transformers as algorithms: General- ization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067, 2023.
[53] Z. Li, R. Wang, D. Yu, S. S. Du, W. Hu, R. Salakhutdinov, and S. Arora. Enhanced convolu- tional neural tangent kernels. arXiv preprint arXiv:1911.00 2019.
[54] T. Liang and A. Rakhlin. Just interpolate: Kernel "ridgeless" regression can generalize. 2020.
[55] T. Liang, A. Rakhlin, and X. Zhai. On the risk of minimum norm interpolants and restricted lower isometry kernels. arXiv preprint arXiv:1908, 10292, 27, 2019.
[56] B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:22 10749, 2022.
[57] S. Mei and A. Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667-766, 2022.
[58] S. Mei, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, L15633):E7665-77771 2018.
[59] S. Mei, T. Misiakiewicz, and A. Montanari. Learning with invariances in random features and kernel models. In Conference Learning Theory, pages 3351-3418. PMLR, 2021.
[60] S. Mei. T. Misiakiewicz, and A. Montanari. Generalization error of random feature and kernel methods: hypereontractivity and kemel matrix concentration. Applied and Computational Harmonic Analysis, 59:3-84, 2022.
[61] H. Q Minh, P. Niyogi, and Y. Yao. Mercer's theorem, feature maps, and smoothing. In Learning Theory: 19th Annual Conference on Learning Theory, COLT 2006, Piusburgh, PA, USA, June 22-25, 2006. Proceedings 19, pages 154-168. Springer, 2006.
13

---

[62] T. Misiakiewicz and S. Mei. Learning with convolution and pooling operations in kernel methods. arXiv preprint arXiv:2111.08306 2021.
[63] E. Nichani, Y. Bai, and J. D. Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. arXiv preprint arXiv:2206.00 0000222200.0688, 2022.
[64] Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[65] OpenAL. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[66] S. Oymak and M. Soltanolkotabi. Toward moderate overparameterizatioo:: Global conver- gence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84-105, 2020.
[67] S. Oymak, A. S. Rawat, M. Soltanolkotabi, and C. Thrampoulidis, On the role of attention in prompt-tuning. arXiv preprint arXiv:2306.03435 2023.
[68] A. P. Parikh, O. Täckström, D. Das, and J. Uszkoreit. A decomposable attention model for natural language inference. arXiv preprint arXiv:1606.0199 00663, 2016.
[69] K. Patarroyo. digression on hermite polynomials. arXiv preprint 19Xii:1901.01. 2019.
[70] J. Pérez, Marinković, and P. Barceló. On the turing completeness of modern neural network architectures. arXiv preprint 1eXiv::901.03444 2019.
[71] A. Pinkus. Approximation theory of the mlp model in neural networks. Acta numerica, 8:
143-195, 1999.
[72] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[73] Radford,
Kim, C. Hallacy, A. Ramesh, Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervi- sion. In International conference on machine learning, pages 8748-8763. PMLR, 2021.
[74] A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
[75] Rahimi and B. Recht. Uniform approximation of functions with random bases. 2008 46th annual allerton conference on communicution, control, and computing, pages 555-561. IEEE, 2008
[76] A. Rahimi and B. Recht. Weighted sums of random kitchen sinks: Replacing minimization with randomization in learning. Advances in neural information processing systems, 21, 2008
[77] S. Reed, K. Zolna, E. Parisotto, G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, et al. A generalist agent. arXiv preprint 2756:72205.061 2022.
[78] G. M. Rotskoff and E. Vanden-Eijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stat, 1050:22, 2018.
[79] A. Rudi and L. Rosasco. Generalization properties of learning with random features. Advances in neural information processing systems, 30, 2017.
[80] K. Shen, J. Guo, X. Tan, S. Tang, R. Wang, and J. Bian. A study on relu and softmax in transformer. arXiv preprint arXiv:2302.06461. 2023.
[81] Snell, R. Zhong, D. Klein, and Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.0700 2021.
14

---

[82] D. A. Tarzanagh, Y. Li, X. Zhang, and S. Oymak. Token Selection in Attention Mechanism, Sept. 2023. URL http:://axriiii httpaaaawi...raaa/////aa//// 13596. #rXiv:2306.13596 [cs, math].
[83] T. Tirer, J. Bruna, and Giryes. Kernel-based smoothness analysis of residual networks, In Mathematical and Scientific Machine Learning, puges 921-954, PMLR, 2022.
[84] A. Trockman and J. Z. Kolter. Mimetic initialization of self-attention layers. arXiv preprint urXiv23305.09828. 2023.
[85] J. A. Tropp al. An introduction to matrix concentration inequalities. Foundations and Trends@ in Machine Learning, 8(1-2):1-30,, 2015.
[86] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[87] Vershynin. High-dimensional probability: An introduction with applications data science, volume 47. Cambridge university press, 2018.
[88] J. von Oswald, E. Niklasson, E. Randazzo, Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677, 2022.
[89] M. J. Wainwright. High-dimensiona statistics: non-asymptotic viewpoint, volume 48. Cambridge university press, 2019.
[90] C. Wei, J. D. Lee, Q. Liu, and T. Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
[91] C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: a case study on ap- proximating turing machines with transformers. Advances in Neural Information Processing Systems, 35: 15:12071-1208 2022.
[92] E. Weinan, C. Ma, and L. Wu. Barron spaces and the compositional function spaces for neural network models. arXiv preprint arXiv: /906. 08039, 2019.
[93] E. Weinan, C. Ma, and L. Wu. A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics. Sci. China Math, 2019.
[94] G. Weiss, Y. Goldberg, and E. Yahav. Thinking like transformers. In International Conference Machine Learning, pages 11080-11090 PMLR, 2021.
[95] M. Wortsman, J. Lee, J. Gilmer, and S. Kornblith. Replacing softmax with relu in vision transformers, 2023.
[96] S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 2021.
[97] K. Xu, M. Zhang, Li. S. Du, K.-i. Kawarabayashi, and S. Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009. 11848, 2020.
[98] G. Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
[99] S. Yao, B. Peng. C. Papadimitriou, and K. Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint ar 105.:2005..111.... 2021.
[100] Yehudai and O. Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 32, 2019.
15

---

[101] C. Ying, Т. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34:28877-28885, 2021.
[102] C. Yun, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-seeeee se-uuence functions? arXiv preprint arXiv: 10077, 2019.
[103] Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and Wagner. Unveiling trans- formers with lego: synthetic reasoning task. arXiv preprint arXiv:2206.04300, 2022.
[104) D. Zou, Y. Cao, D. Zhou, and Q. Gu. Gradient descent optimizes over-parametterieed deep relu networks. Machine learning, 109:467-492, 2020.
A Technical tools
A.1 Basics on Hermite Polynomials
In this section, we briefly review Hermite polynomials and their properties. Let He, be the proba- bilists' Hermite polynomial of degree n:
He (2) == (-1)"
d
e +
dr
Here are some basic properties of Hermite polynomials Inz0:
Let (a. N(O, p; 1)). Then E(Hem (e)ee()) nlp"Im== for m, 2 0.
He,(-r) (-1)" (z):
He (z) .HH,, (z) - (x).
He), -00==:] (=Hee_-(x))
can be shown that Hermite polynomials are a complete orthogonal basis of the following Hilbert space with ((z) = (2m) 1/2 1/2,-1//2,
/2
[ f(z) < co)
L3, o(x)) =
f f:
For two functions f,g: R + R, we define their inner product (J,g) as:
(f.g) = E. = [ f(z)g(z)d(z)dz.
For any function f (RR we can derive its Hermite expansion:
f(x) = Σ
On Hc,
n!
n=0
where On (f, He,) S (z)
Then for another function g(x)
Enzo by He, (I), the inner product of f and gives: n!
ha
(f.g) =
n!
Here are some formulae for Hermite expansion certain functions.
Lemma A.1 (See e.g., [38]). Let f.g L (RR ((x)). Then, for any unit vectors u, U Rd we have
that
2
an bn
B [f(T)))((() (o((TT)) x]]
Σ
(u, w)""
n!
where f(x) Ene "e He, (z) and g(x)
E o
b
(x).
16

---

Lemma A.2 (Inverse explicit expression [69]). Hermite expansion of x² gives:
1
x²
nd!
2 ml(n - 2m)) Hen-2nn 2m 2m()).
m==
Lemma A.3. Let (e(x)) == o(x + c) be the shifted ReLU function. Then the Hermite expansion of Oe gives:
(-1)"
(c(x)=(ii + ((c) + 16 He) +
2())),
m!
where d(x), (xx) the PDF and CDF of standard Gaussian.
Proof of Lemma A.3. Denote an := - (x (e,()))) It easy to obtain that
00,0 = (c),,
40,1 cb(c)) + <(c),
01,0 = (bc), and 01 = +(c).
Then using the two different formulae for an,o.
Hen+ (-c)
Hen+ (x)
an
He, (x) (x)((x))r
Q( ((-c))
(o(x)
n n
Men+i(-c) o(-c)+ n1 n+1 an+11 n2 n+1
an+1.0 n n and
n ++1 1
an,0
(IHen-1 (x) - (n
an-1,1 can 1,0 (n 1)0n-2.0 n 122 2,
we obtain that
0n,1 =
+ + (n - an - (n - 1)can
+ c0n-
and
0n,0 0,11111 1,1
00n-1.0 1.0 (n 1) a)nn-200 2,0; 2 2.
Then easy to prove by induction that
0n,1 =
(c) n n2
0, = He, n n21.
This completes the proof.
A.2 Basics on Random Features
In this section, we give some basic properties of the random feature model considered in our work, which can be seen as an extension of the standard random feature model (e.g. of Rahimi and Recht [74, 75]) to the vector-valued case.
Given functional a (x)w)) w)) X' W + Rdd. Denote y probability measure W. We define the wwithh random feature model as:
F= f(r) (v. H) v Hw}.
(22)
where Hw {v(w): Sw v(w) v(w)p(dw) oo) is Hilbert space with norm || Sw v v(())((dw) and inner product (v, u) u) = Iw v(w) ((w))) Besides, we endow F with a norm || and the corresponding inner product (., )) defined as:
|l gll llf g
|| =
inf
ll (f.g)=
(=(v.o(-)) Nyw
4
We further define corresponding reproducing kernel K X R. s.L.
K (ax v) = / a,;; T (yr
which is positive definite. Define the RKHS induced by this kernel as Hx with corresponding norm Il-Ilay and the inner product ( Hg. Then we have the following proposition according to [61]:
17

---

Proposition A.1. Given the above definition of F and H.g. we have that (F, (111)= = (HK. llllxxx))
More generally [11], for any feature map Φ: X - H (where H. is Hilbert space) that induces the kernel K, i.e., (x, y) = (0(x), d(y)) * we have for any function ff that
inf Il lfll f=(u,0())
||
(23)
which shows the equivalence among different feature maps that generate the same kernel.
A.3 Concentration inequalities
Definition A.1 (Sub-Gause Gaussian and Lub-exponential random variables [87]). For random variable X, its sub-gaussian norm, denoted || ly.. is defined as
|| inf { 0: E exp (x²/t²) 2) a || < oc, we say that is
For random variable X, its sub-exponential norm, denoted [ is defined as
|| inf {t > 0 LE exp (|x)//t) s 2}.
Ifo = || < <c, we say that X is o-sub-exponential.
Theorem A.1 (Gaussian concentration inequality (e.g.. (89)). Let (X X,) be vector of i.i.d. standard Gaussian variables, and let f 12 - IR be L-Lipschitz with respect to the Euclidean norm. Then the random variable f(X) - 1E (X)) is L-sub Gauu-Gaassinn, and hence
22 for all t. 2 0. 2e
Theorem A.2 (Bounded difference inequality (e.g., (89]). Consider a finctionn f(X) R2 + R. Assume that for any X = (XI...., Xx)) and X (XI,..,;,, N....,,.... Xn), we have difference bound [f(X)= - f(xt') S Lz. We further assume that the random vector X (X1 X2...., Xn) has independent components. Then
2x
(((((x)- E[f(X)] -[((x)) 1) < 2e Ei=i for all 1 0.
Theorem A.3 (Matrix Bernstein Incquality (e.g., [85])). Consider a sequence {S* of indepen- dent random matrices with common dimension d1 d2. Assume that
ES₁ 0 and ||Sxllop I L almost surely, for each index k. Introduce the random matrix
Z
Sk-
k=1
Let v( ()) be the matrix variance statistic of the sum
v Z) max
| [ZZ
1lpp
E
Z Z
op
n Σ E (k
Σ
E
(S) Sk)
Then have
E Z lpp S 20 Z log (d) + d2)
+ L log (di + d2),
3
for all
-2212
P(1 7 op 2 t)
(di + d2) exp
((Z) + Lt/3
Theorem A.4 (Ledoux- Talagran contraction inequality (e.g., [89])), Let {Ee Je(s mid Unif(n be independent Rademacher random variables. For any set T C R and any family of L-Lipschirz functions {0, jEin] with (0) = O, we have
E sup
<
2L E sup Σ 5.0.1
(24)
i=1
18

---

B Proofs for Section 3
Throughout this section, to denote function of W E [((+1)x(d+1) and X {X R(d+1)x(++)) satisfies the following properties:
the notation & to hide universal constant C. Also, a (W, X)
|X|| I 4) D that
1. For any X D, we have that ( X) is L1 Li-ischhiizz with respect to W.
2. For any X D, the expectation Ew [ (W, X)] S L2.
3. For any X D. have ( X) L3.
Here L1, L2. and Ly are universal constants. Let { m Ime be sampled from Eq. (4). Let v: R(d++)x(d++1 + R with Ew [w]llv(W)I1 < R². Consider random feature model associated with or (hereafter, we will refer to as RF model)
M
N (W
M (No.ni V)
N
XoX,)
(M,xX.)
(25)
m=1
as well as the infinite-width version of the random feature model,
N
1 N
fv (xx:N)
E Ew (( (WW, xxx,) (v (v(W),
(26)
Note that both model RFA and BRFA correspond special choices of a. Thus, all lemmas and theorems in this section are applicable to both model RFA and model BRFA.
B.1 Proof of Theorem
To prove Theorem 1, we first state two auxilliary lemmas, Lemma B.1 and B.2. Lemma B.1 (From widhh RF model to finite-width model), Consider fv that takes the form as Eq. (26), with Ew [llv([]]]]))))) (W) < R2. Let XO.N ~ P. Define the 11-1[ex(p) norm by
1g = =f x(xx.N)) P(dx..)
Then with probability at least 1 8, there exists sequence of vectors { C RR++1 and constant universal C < ou that only depends on Lr, L2 and Lg s.l.
N
(vm. 1
C(d log M) M)88811
fv
(W m. XxX,
S
(27)
12(P)
M
Σ
vmll, 22
CR M
and
Σ
Vm
CR
(28)
M
Lemma B.2. Under the setting of Theorem 1, let f. be a target function of form (6) and (7). Then there exists an inftnite-width RFA model (26) with v g(d+1)x(d+1) + R such that
N
f.
Ew ((W Xxx, (v(W).5))
(29)
with
Ew Ew(lv(W)ll) < B(f+),
where B(f.) is as defined in (10).
The proofs of Lemma B.1 and B.2 are given in Section B.1.1. Now we assume these two lemmas hold, and use them to prove Theorem
Proof of Theorem J. For any function that takes form (6) and (7), by Lemma B.2, it admits represen- tation (29) with Ew (W) < B(f-). Then by Lemma B. since RFA model is special case of the RF model, there exists (vm me[M] {vm]me[N] such that Eq. (27) and Eq. (28) hold with probability larger than 8. This proves Theorem 1.
19

---

B.1.1 Proof of auxiliary lemmas
Proof of Lemma B.
Step 1. Proof of Eq. (27). To prove Eq. (27), we use truncation argument.
which we will choose its value later in the proof. Recall that we have W1, W...... W m i.i.d. with W WmssmaaN((0 1/4). Define Vm v(W m S S Rw /2) for m = 1,..., M. Consider the truncated infinite-width random feature model fRw
Fix a Rw
ΣE N 2i=1 Ew xxx,) (v(W), X.) |[F S Rw/22], we have
M
Ew fllw
σ [ x (Vm, xz) [
N
i= m=1
(()) | 1 S Rw/22]) /M
Ex
Eww
X
N
SEx FEw
2
|
| 3
3 ([[]] S Rww/2)))
.N
ÖRWW R M
where C1 only depends on L1, L2, and L3. Then using Markov's inequality,
1 N
M
3C Ri2
RR
Σ Σ o o Xox,) (Vm, x:)
N
(22(P)
5M
holds with probability at least 1 - 5/3. Next for the difference between f.w and fv. we have
I1 |Rw 1122(m
Ex
[Bew
aW. (W. xxx (v, (3) ([] Rw / 12
N
Ex
Ew (v. x,)²
V
Rw/2)
S RPP(I Ir, Rw Rw/2)
Here C2 a constant that only depends on Ln and Lg. By concentration of functions of Gaussian random vectors (Theorem A.1), (W X,X,] Ilye, < L1 + L2 for any i. So in the last inequality, we used the bound [o (W, X 1/2 by O((L) + L2)²). To bound P[]] | > Rw /2), we use concentration of functions of Gaussian random vectors (Theorem A. again, and get that
P( E() W 2 2/22 S exp (-t²/2).
Take C max C2). Since B[[ [+)) S (EI (B))))) < d 1, by choosing Rw + 1
Oolog M, the above probability is less than 1/M2, Then
N
M
3CR R2
fv
oW, (Vm, x,)
+ CCRPP(I > Rw/2)*
LP(P)
8M
C(log M RR
8M
with probability larger than 1 - 8833. This proves Eq. (27). Step Proof of Eq. (28). By Chebyshev's inequality.
M
6R²
P
|m ll2 [|v(W)l ( (I........ S Rw/2) 2
5M
20

---

I8
((()) 2
(I FF S Rw/2] 8M
8
S
6R
3
Combining with the fact that E
[v(W)l2
(| S Rw/
E
[I([))]]]]
< R.
have
M
6R
P
V l2 2 R
8M
3
m==1
For the second part of (28), Markov inequality gives
M
6R
[l) (I[ m Fl. S Rw/22
6M
P
Vm
I
8M
6R
This proves Eq. (28) and completes the proof of Lemma B.1. Proof of Lemma B.2.
To get the kernel of the RFA model, have
K K(Xo:N:Xo.N Xg,N
1 N2
Ew
Σ
()) x
1,f=1
We first consider single component in the sum, which is
Ew [o (WW xxx; (W (((,))) (x, xy).
(30)
Let Ua (xxx) (X,) /4. Let (p) denote bivariate normal distribution with marginals
N(O, 1) and the correlation p E 1-1, 1]. Then (30) be expanded as follows:
Ew
0r [ (W. X (x,) (x)
0 (Z))) (z,,za~N((a.))
(,,)
1
[T/
arccos wi.j) + / u
(2i, X,)
2m
1
(2f - 3)!!
1
Ui.j +
(20) - 1) Ui
28
(2,
2m
2
t=1
Σ
(XxX]
[o( T) (2.8))
(E(0.1)((2)).21
Σ
(2 86 3 xx,, (xx() 10 & 8x
(E(0.))u(2k)
Here the coefficients {ce} satisfy
(2f 3)!!
co = 1/(2T), 6 1/4, and cet 2T(2))!!(2E - 1)
O(e-2) for l
Therefore, the kernel can be expressed as:
K K(XoNN; x(,N)
1
Σ
Σ
Jo x 2 2- (x) 0x,
NZ
oo N N
= Σ
Va
Σ
2- 2xxxx
(xxx)
j0e 2 x,,
Σ (56 (x) T) 8x²)
(E(0.1J((2k))5)
1=1
j=1
21

---

Now we reformulate the target function as
N
0
fx
Σ
20 x0 frs))
N
t= max(r;s]=t
N
Σ
(Xo 1 x; xx)x x;
60 x;, Prs
N
1-2 20 max ( (1a)
2
Σ
Va frer
N Σ ==1
2 [xxX; /0 8xx),
N
22
max max(ria)=t =f
where f.s is a transpose of fos & 1 1441
0(22-r-s+1)
with 14+1 = (0..... 0, 1), such that
(xo & x fro) = ((x, @ x,) ol @ X,, fre) for any r 2 and s 2 0. Then by the feature
map equivalence property (23), the RKHS norm of f can be bounded as
VCL N
N
of 8 x: )1
|
Σ
Les
(xxx
Hx
2=
max m
i=1
2
2
Σ
Σ
Fre)
V
Vai
/=2
max(n,)=florts
max(n=)=foot-1
00
Σ
Erg
FrY
/=
1 k4.5
Σ
k=
=a
Thus, using again the property (23) with the original feature map of the random feature model, there exists V [d+1))(d-11 Rd+1 such that
N
Ew (W, (v(W), X,),, with
Ew]lv((W)1
Σ
|E
k=0
max(r,e)=k
Notice that ||E = |[f..e by our construction of fr so that the above coincides with Eq. (10). This proves Lemma B.2.
of the equation
B.2 Preliminary proposition for Theorem 2
To prove Theorem we first present and prove the following proposition that gives a high probability bound for the difference between the empirical risk and the population risk. In the proposition and
(1) ( and y +Ws]yeen
lemmas below, denote X
Proposition B.1. Under the setting of Theorem 2. Consider the finite width RF model (25):
M Σ m=
M (Xo.N; V)
( m. (ox,) (Vm, x.).
Then with probability at leas! 1 - 8 (w.r.t. W, y, and X), we have
sup Σ (SW VEVU J=1
Xo:Ni V). v3))
-xx a ( (Xo.Ni V).y)
log(dM) log (nNM)
(6)
K1
K₂
< K1
log 8
(31)
n
M
22

---

The main difficulty of the proof of Proposition B.1 comes from that ( y) might be unbounded and that E(SM W (Xo:N) ), Mj) are not independent across j (since they share the same {Wm) Wmmee(M)) So we begin with several lemmas below.
Lemma B.3. Let {E sejm] be set of i.i.d. Rademacher random variables. Under the setting of
Proposition B.
Exy.WE VEVM
sup
Σ
E
V..iiV), v).
[1) s Kiv
log(dM) log (n
(32)
n
1=1
Furthermore, any fixed X aad y,
Ew sup 1
L ... x K1
log( (gg(dM) log (nNM)
(33)
VEVM
n2 ;=1
n
Lemma B.4. Under the setting of Proposition B.1. With probability at least 1 - -8/3 over X. y, and W,
n
n (SW (x..i (2..ii))..) )]
Σ
( XONi V). Ew
sup VEVM
j=1
1=1
K2 (6g(()))
log(dM) log (nNM)
+ KiV
(34)
M Lemma B.5. Under the setting of Proposition B. I. With probability at least 1 - 5/3 over X and y,
sup VEVM
Ew
JW
9 V). (.15)] Ex.x.w 5=1
Xo:Ni
(
v),
==1
log(dM) log (MNM)
+ KiV log (6/0)
S SKi
(35)
Lemma B.6. Under the setting of Proposition B. With probability at least 1 5/3 over W.
sup VEVM
Ex,y
Σ ( In
1
1 XoNi V) w))] Ex.y.w
(IW
Xo.Ni v) M)
j=1
j=1
K2 log(6/0) 6/8)
log(dM) log (nNM)
Kiy
(36)
M
The proofs of Lemma B.3, B.4, B.5, and B.6 are contained in section B.2.1. Now assuming they hold, proceed to prove Proposition B.I.
Proof of Proposition B.I.
Split the left-hand side of inequality (31), we have
1
n
v). - Exy (M( (x22xs
(x0((xx2x))))) v).v))
sup VEVU
M
11
j=1
n
H
1 Σ in 3=1
< sup VEVx
(f .y))
Ew
n
j=l
+ sup VEVи
Ew
1 -
Σ ((FW m1)]
Ex.y. (e(FW 2i)
n j=1
Ex.y ((FFW) .2)) Ex,y W ((UHW.))
sup VEVU
(dM log (nNM)
6 T
K
K₂
$ SK1
log
11
M
with probability at least Here the last inequality uses Lemma B.4, B.5, and B.6. This proves Proposition B.
23

---

B.2.1 Proof of auxiliary lemmas
Proof of Lemma B.3.
First using Rademacher contraction inequality, since &(0. Σ 1, center it and only pay extra term 1/vn in the Rademacher complexity. Then by the Rademacher contraction property
(Theorem A.4), the problem boils down bounding the Rademacher complexity of S which
(x E EafM (xO.Ni [
Exy.
sup
VEVM
j=1
Fix X, y, and W, we have
E
sup VEVM
M
W
( X
v
M
N
x )
EE
sup VEVU
V
W m.
m=
==1
|
X x
3 KE
max m
( m.
x
(37)
j=1
By matrix Bernstein inequality (Theorem A.3), for any fixed m,
i||
P
Σ
a (W x
x
x
2 2dx
nE² /2
12 + Ke/3)
n
where
A =
max
nN
o( m. x
)o( (W m. x()x((1) (xx() (x(",x((0))
max
W
X +
),
and
m.k.j
K max 2,12
W
xx(0))
N
Using the union bound, we have
max
1 5
Σ
N
(
x
111 26 s 2dM cxp
(- A2 - KE
ns²
mE[a]]
12
N
j=1
Therefore, we can bound its expectation with
N
log(dM)
log (dg(
Eg
max
Σ
X
A +
K] (38)
N
j=1
Now take expectation over X, y. and W. Since 2 (og( we have
log log(dM) A log(d.M)
Ex,xw
K
n
n
log(dM)
(d.M
((Wm::x) x0)x01T) ||
2A
Ex,y
max f.j.m
n
log(d (dM log (RNM)
2A
7
Combine this with Eq. (37) and (38), we prove (32).
Fixing any X and y. only taking expectation over W, we get
log(dM) A log(dM) K
log(dM) log (nNM)
Ew
n
n
Combine this with Eq. (37) and (38), we prove (33). This finishes the proof of Lemma B.3.
24

---

Proof for Lemma B.4.
Denote X (X0:NIJEjm) and y (Vibeenn and denote
(f) XD:Ns V).
1
n Σ
xx:Ni V). 1) ]]
g( W:M; yy sup VEVM
Rw
JM
n
m ==1
Given WIMM = (
m=1 and W I:M = (w 1M
we define
WWEM - W Wirelerr =
m
(EmE(M) [ m
W )
Then we have
9 O(WI.M L:M: x,y) (iii i::: x.
((MM '((Mi' v1.) .1)]
sup VeVu
Σ (W, x o( x( (v)
sup VEVM
Nn
f.m
sup VEVM
Σ
W
W Wl Mvmll2
m
K₂ M
W I:M
Wim Willlyy
Since W has independent standard Gaussian entries, by Gaussian concentration inequality (Theo- rem A.1), we have that
2xp (-
MM²
P [( LMM x) Ew g( ( i:M; x,y) x, 2 y
(39)
[K2
For the conditional expectation, by symmetrization, we have
Ew
g ([WI:AM I:MMi X, y) X
Ew
sup VEVU
JWW
G) XaiNi
( -Ew [e (xoNs v). [w:)]
x
Σ S (x V). 1x),)
2F &
sup VEEU
1=1
log(dM) log (nNM)
$ KKV
(40)
for any X and y, where the last inequality uses Lemma B.3. Combining (39) and (40) and taking
E = 2/ K2 log (6/8)/M, we have
1 Σe
xoni (edx:v). v). .1:) Ew Σ j=1
(
( Xo:Ni v). ws)]
sup VEVM
72
j=1
K2
6
log(dM) ) log (NNM)
log
+ Ki
(41)
M
for any X and y. Since the right-hand side irrelevant X and y, we get (34). This proves Lemma B.4.
Proof of Lemma B.5.
Let
n E
(xx.N.y)
sup VEVU
Ew
Σ ((W (.))
Ex, w
Σ
(S
93
(42)
n j=l
n 1=1
25

---

For each i E [nd, let y' differs with yX.y} only on i-th data point. We have
1(X,y) h(X',''
Ew ( (/l S V) Bb)
Gy Ew (f
v) 03)
$ sup VEVM
XUNN
2 sup
Ew
-
JM
() Xo.Ni
v) SM X0..N1 T
VEVN
j=1
sup
Ew
E mE(M])
|o { xf + oW m. x6 xf"xx("')
| [/vlll
VEVM
LN
K1
Therefore, (xo.N,,)) y) satisfies the bounded difference property with the parameter (Li uniformly bounded by ((Ki/n) /n). By bounded difference inequality (Theorem A.2), there's constant C such that
Önn² (Ōnne² K2 2 (K,²"
[h(X.y) Ex [xXX
2 S exp
Combining with Lemma B.3, we have
Ex,y 1(x,)) y)
n Ex,..
-x,y
sup VEVM
Ew
e
=
Σ
E (MM ]]
n j=1
1=1
n
$ 2Exy
sup VEVU
Σ Ew
[CI
XUNNi
n
j=l
log (nNM)
$ Kiv
Therefore, by taking E = 2K₃ [log (6/8)/n(O)]// 1/2, we have
( m))] Ex.y W Σ e(SM m5)
sup VEVM
Ew
log(dM) log (nNM)
log (1/6)
< Kiy
+ Ki
n
7
with probability at least - 8. This proves Lemma B.5.
Proof of Lemma B.
Denote
n Σ E(S 3/)) m))] in jml
1
a Σ
( 1:M)
sup VEVU
Ex.y
Ex.y. w n j=l
efW .2))
Similar to the proof of Lemma B.4. Given W I:M (Wm)m and W ( m MM m define
|| 1:M WiLMl Fr V.m [m W
12 We have
<W 1::1) - pp( (PW Wit)
Ex.y ((FM
(LM..)5) .b) ( .1)]
< sup VeVa
Bx.y [ x(1 o W s6
( (())T 3x)
< sup VeVa
Nn s.j.m
26

---

sup VEVM
Σ
2
W Wm.llelll
K2 M
|WiM WiM |F
Since Wn, m has independent standard Gaussian entries, there is constant Ö s.t.
ČM
[W(WIM)) Ew
)] 20] 2eexp
(43)
K₂
For the expectation, we have
Ew
go (W 1.0r)]
Ew
sup VEVU
(J JM
Xo:Ni
(iax:)..0 v).
Ex.y. -
Σ
n
(JW m)
n
n
j=1
j=1
E (FM (X
(.xvi)..u V). |
2x
sup VEVU
n
log (d(d log (MNM)
Kiyg
(44)
where the last inequality is by Lemma B.3. Combining (43) and (44) and taking (K2 log(6/0) we get
sup VEVM
Ex.y
(JM
1 (.M.U)) vi) Ex,y w
e ( mi)
n 1=1
1=1
K2 log(1/6)
log Og(dM)log (nNM)
Ki
M
n
with probability at least 5. This proves Lemma B.6.
B.3 Proof of Theorem 2
Proof. Given any target function f., by Lemma B.I, with probability larger than 1 - 8/2 over WIMM. M
there exists V (Vm m=1 such that
SW (No.Ni VIlLA(P) S v /(B++ log M)B(Ja)6-1
llf
(45)
M
with
M
S B(.)) + B(f.)o-
M
B(f.)8-1
Σ v m=l
M S 22 B(fa) and Σ
Vm
(46)
M
By our choice of K₁ and K2, let firm = J (:: V) denote the model trained by (11), and let
A. M = arg min Lp ( (:: V)) and f.r. = JW (Xo.Ni V).
VEVN
Then with probability at least 1 - over W, X, and y, we have
L (JW - T.p(S*)
(47)
< LD (DWW - Ln (1WWaa) + LD (JWMM
+ -
- -DD(FW. Ms) + +LD(fWW.)) - LD + +LD(f())) Lp((f)
(48)
< < LDW - -DW + +ID(fWW, - LD LDD + LD (D(.) LD(
(49)
22 sup fEVM
LD(f) - [f)
+ +11 f-llea(m))
(50)
27

---

log( (gg(MM) log (nN M)
Ki
K2
(df + log
M B).(
S Ki
+
log
+
+
vn
M
M
(51)
B(J+)
(d2 log M)B(51)6-1
27 24
(VlogddM) log(nN v Vlog (6-1)
+
(52)
n
M
B(f-) [[f=]]log (log(dM) log(nN + log(0 (ef log M)B((1)8-1
(53)
7
M
where from (48) to (49) we use the definition of fr.M and fv .M- From (49) to (50), we bound < notes: tic ( (D)() and (Dv .M) (SD .M) by sup, LD(s) - (D) and
use the Lipschitzness of E(f. y). From (50) to (51), we use Proposition B.1 and Lemma B.1. From
(51) to (52), we insert the value of K₁ and K₂ into the equation. This proves Theorem
B.4 Proof of Examples in Section
B.4.1 Excess risk of RFMLP
Denote F to be the set of all functions in the function class (6) and (7), i.e.
F F= (=((((XONN)
1
xi fre) fre E Ed 'symmetric ,r, 2..200 20. (54)
x5
r,520
Consider the RFMLP model
M
JM
(xuNiiv)= Σ
o (Wm. vec(xp:N))) Um
M + 2)), (55)
where vcc(Xp.N) (xo: ; :NN; E RdN+++1 For target functions that take forms in (6) and (7), define ()L = Ex= Ck Ertesk ||fr with Cs = (N + 2)2*. In case where f- admits multiple representations of the form (7), (f is the infimum of the right-hand-side over all such representations.
Then we consider the empirical risk minimizer over the RFMLP model:
V = arg min LDMMMP (+:)),
Ip(() = +. Ei ((((x..).
(56)
where the constrained class MM gives
V = - v= (Um)
EM |uml S K1, EM U S KK//
(57)
Proposition B.2 (The sample complexity of RFMLP). Let f+= arg miggeer (f) be the population risk minimizer within the target function class (54). Assume M > 8-1 and log(dM). Take
K1 = C (VBLPP(..) and K2 = CBBMLP(f.66-1 in (57), with C being a constant. Let PMPP fM (:; be empirical risk minimizer of model REMLP. Then for any joint distribution P, with probability least 1 - 5 over [Wm]mejad) sampled according (55) and (1) V3) (5))ss(()) Gid P. the exeess risk is bounded by
LD L LD(f)) & (VVBALLLEE)
d'O
(58)
M
Proof of Proposition B.2. The proof is basically the same as that of RFA model. We only give sketch of the proof. Firstly, with a few modifications of the proof, we can show that Lemma B. also holds for RFMLP, with B(f-) replaced with Buar (f.). Lemma B.2 slightly different, we have the kernel expansion:
K (X Xo.N)
Σ (2
(Ve(N+2)-
.No(N+2)) [vec(xoNN) 86
vec(x(xx)) 36
28

---

Therefore we can rewrite f. as
22 N
fre. VQ(N + 2] 2) voc(x( 80 ) (59)
f.= Σ
Nee Σ Σ
2810.18U(28)
where (Fre,s, [vec(x/
(x.x)))/)))) = (fra)
Sr 2 X
Thus,
I|S.
C
Σ r+s=k
with -= k + 2)
k=0
The hight-hand-sid gives the formulation of BMLP (f.). Then, a similar version of Proposition B.1
holds for RFMLP model fM i.e. with probability least 8 (w.r.t W.y, y, and X),
n n
sup VEVU
IM
V): v)..)) X
O.N'
Exx (XoiN: V).y)
j=1
(N + 2) log(dM) log (n
K1 Vn
K₂
S Kiv
V (N + log (1/8)
(60)
M
Then combining all of the above equations and following the proof in Section B.2, we get (58).
O
Remark 1. The representation in (59) is not unique. With more careful choice of the representation of the target function fw. we can get a better bound for (f which given by
20
BELP (f+)
Σ C
Σ r+z=k
with C2 1. 1...((N + 2)/2
(61)
B.4.2 Proofs of Examples
Proposition B.3 (Restatement of Example 1), For functions of xo of the form
[x(x..)) =
0 @k AR). A E Ra
R=O
we have that B(f=) = Ei-o O& |ARlEE Setting M = (d) the excess risk bound gives Ö(V ÖlVE ht | In). Moreover, consider fa (2O.N) (B xo)P. The above excess risk of RFA model and the REMLP model scales
RFA: δ ( 1 /1))).)m)
RFMLP Ō ((N+ (1N
Proof of Proposition B.3. This follows by direct calculation.
Proposition B.4 (Restatement of Example 2). For Sx
Di=1 ((B,xx)) xx)) with v=))
z arrtaan (2 with 1, > 2 and || = The excess risk bound of RFA A model and the REMLP model
RFA: (VE 34/ = 5(V1/m)
REMLP: O 145 2)/2nl2//n)
Proof of Proposition B.4. Use the power series expansion of vh, we have
N
(B,x,) 2k
(2k 24-2
km
Plug it into the formula of B(fa) (10) and BMLP (f.) (61), we get
B(f-) = Σ (2x) || 4 e Σ (2/())))
and (MLr(J+) = + k.4.5 + 2 Therefore, by Theorem and Proposition B.2, we get their excess risk. This proves Proposition B.4.
2k
0
29

---

Proposition B.5 (Restatement of Example 3). For f1,+ = 1. N Ei= (x.,x,)" the excess risk bound RFA (by Theorem 2) and RFMLP scale as
RFA: δ (Poly()))(AAd)/l))
RFMLP:
(N22)N/+/)) 2)d
For f2,, - # >m= cos( x:)) (x,"", G) with [| = 1. Then the excess risk bound of RFA and
RFMLP scale
RFA: σ (Polyyud) eAvá4p/n)
RFMLP: ((Poly(oNa)) ( V 22)/m)
Proof of Proposition B.5. For Si direct calculation gives the value of B(fi,z) and the excess risk follows. For J2,x, using the Taylor expansion of cos(z), we get
x,) 26
(-1)*
f2,-
G)
(xxx, 02k
)o2
6p
1
& G
N
(2k)
N
(2k))
1=1 K=
k=
Plug it into the formula of ((aa) (10) and BMILP (J+) (61), we get
d2k
20
(4d)
B(ʃ2,)= = 4PPoly(p)
22
PPPol((dd)
(2k)1
(2h:)))²
k=0
Note that for any >> O,
0
(2)2k
(2))2k 2k
0 (
(62)
(2k))
(2h))
Plug = 4d into (62) for BJ2,.)) ) gives the excess risk for RFA model. As for (JLL)(2,.). we have
dak
+ ((N
((N
22
ak
BMLL(22..) = Poly(p) Σ
4PPoly(ndN)e ((2))))2
((2k)!)2
Using similar argument, we get that BMI.P (/2,) is bounded by ((Poly(pdN)4P exp(2(N + 2)) Using Theorem 2 and Proposition B.2, we can get the excess risk bound. This proves Proposition B.5.
B.5 Random Feature Attention with Exponential Activation
We give a brief discussion on our results with activation function replaced by the exponential function d'(t) ct)(t). The random feature attention model with exponential activation (we call it ERFA) is given as follows.
( V) = IM N
M
EN o' ((W... xoxTT) (Vm, x,)
(63)
We consider deriving the explicit form of the Kernel KERFA ( xb,N) associated with ERFA model
1
KERFA (XUiN Xx.N)
Ew
exp(( Xox, + ()
(64)
N2
For any i.j E [N], let Uijj
(XxXT
/4. For single component in (64), we have
Ew-N(0.1/4)
eepp( (W, x + xx,)))))) (R xx)
| + (x) 11 (x,
exp
8
exp (1 wi.j) (x,, x))
2
Σ
(x,
1!
1=
30

---

di (x,2, (x())) 4 (x;;
==
2 (x1x;
y 8 x 2 (xx yo0 2 x2²
2=
Here di e/él for any & 2 0. Therefore, the kernel can be expressed as:
de Σ (x)))) Joo @ xx²).
KERFA (Xa:N xp:N)
2 (xox,
i=1
Thus, for any target function f. that in the form of (6), we can redefine its complexity measure as
B(f.) = Σx= Cs Emax(n,))= | I
C₁ = k4 V1
(65)
and obtain the corresponding results of Theorem I and Theorem by repeating the proof of these two theorems.
C Proofs for Section
We consider the empirical risk minimizer over the BRFA model (15), arg min Z. V)),
LD(J) = - ((
(66)
where the constrained class VM gives (12), copied here for reader's 's convenience,
VM = =
{vm)m
EM=n |/ s K1. DM || s Kg/M}
(67)
Denote be the set of all functions in the function class (16), i.c.,
g g= 1-{0-
1 N
F(I)
Σ apt G
( x A.)
(68)
We restate Theorem 3 in Theorem C.1 with detailed assumptions.
Theorem C.1 (Restatement of Theorem 3). Assume M 8-1, > log(dM) and let L E Zz1 Let
9x = arg mingeg (.) be the population risk minimizer within the target function class G (68). Take K2 ((g. L) and K2
L), in (67), with C being a constant. Let fM Wa fW (w; V) be empirical risk minimizer given by (66). Then for any joint distribution P, with probability at least 1 d over {W m m sampled according to (4) and ((x Uz) Je(n) Wiid P. the excess risk is bounded by
LD (W.Wo) -LD(g+)
log(dM) log(nNM) + log(8-1) +
(d2 + log M)8 1
+ e1 1 }
$ inf where EL = - (I 1)!] and
BBgg. L.
n
M
B(g., |AA. (Er lax| C.),,
with C2 (2L + k)
(69)
C.1 Auxiliary results for the proof of Theorem 3
In this section, we give some auxiliary results used in the proof of Theorem The proof will be given in Section C.2. We define the biased transformer with infinite width (informally corresponding Eq. (15) with M - 20), given by
N
Wo (Xo:x)
Ew o ((W + Wo, xox.)) (v(W),
(70)
31

---

N
Σ Bw (WW, + h4) (v(W),
where h2 (Wo: x and W is sampled according (4). Here set Wo
[Lxd 0dx1: 0 1xdd Oj. Then h4 = (Xo, x,)
We consider class of target functions g. X R that takes form
N N
9. (Xp.N)
( De),
(71)
i=1
8=
for some coefficients (D, E (Rd+1)2213
122 and
(h,) + x1+))x (1,, & X
(e 0),
Ue Xx, x,)
(E = 1),
Hee (xx) j & &
(E 2 2),
where @ and & are the PDF and CDF of the standard Gaussian random variable, respectively. Lemma C.I below provides a counterpart of Lemma B.2 for the BRFA model and the target function (71).
Lemma C.1. Any function 9. of form (71) can be expressed exactly an infinite- random feature attention model (70)
1 N
(x(.N)) a
(XF:N)
>
Ew Ew [o + Wo. (v(W), x;)],
(72)
where the coefficients v-) satisfy
Ew (Iv(W))
E20
Given Lemma C.I, we can get a counterpart of Theorem for the BRFA model as Proposition C.1 below.
Proposition C.1 (Counterpart of Theorem for BRFA). Suppose function g+ : - R takes form (71). Then for any input distribution P X, with probability at least - 8 (over { (W W m)mmen sampled
from (4)), there exists an M-head BRFA model (15) with coefficients = [Vmmmm[M) S Rdd+1 that approximates 9. in (P)) up error
Ex (g. - JW (No:N: v))
((d((p log log M) M
(d
(73)
In addition, the norms of the weight of this random-feature attention model are bounded as (
Iimil22 2 (BB..) +
B(G.)6-1
Σ 1i <
so o
B )
(74)
M
M
Here B(g.) defined alternatively as
- CR DD
with
C2 - k!4 V 1.
(75)
Furthermore, we could approximate the target function in the form (16) any precision by function that takes form (72). which is discussed in Lemma C.2. Lemma C.2. For any target function gs in the form of (16), and for any precision (211()+ there exists function fWo in the form of (72) such that
11g.
Wo II S 11g+ Et,
and
Ew [[(()) <
a 213 (26 + p) 442
|AA.1
where the tensor A. E E parameterizes g..
32

---

C.1.1 Proof of auxiliary lemmas
Proof of Lemma C.I. The function (72) can be interpreted as a linear function of the feature map
1 N
V (x.NN)
Σ
(W, + h;
So the kernel w.r.t. the feature map takes the form that
K(Xg:N Xb:N)
EEw (VV (x.:),, V (x66N))) )
Σ ISUJN
Ew [w[o ( + (.). + h2)] (x,
N
where h; = (Xo, x.) and = (x,, Similar to the proof of Lemma B.2, we also consider single component the equation above first, which has the form of
Ew [or ( (W, x + h.) ((W x,,TT) + (x,
(76)
We expand o ((W, xoxx) + h.) by Hermite polynomials in the space L (R, の) using Lemma A.3,
Ew (W xox,) + h:) (W, xxxx))]
EE-Nro.,)lofe + h4) (c))))
h.b(h,,) + o(h,)
(f= 0).
((1)) (1)))(()))Ho
(f=1)
(e2 2).
Therefore, we have
a(W, + h2) = h,;0(h;) + +
+ (he x)x,))
(-1) o(h))eee-((h))ee( xx,x
+
Using Lemma A.1, obtain the expansion:
Ew [a X +
X (() (2,
= [hab + 2 (hha) ( + (11)] (x,
+2 (h,) (h/)
(3 (b))
x (5, x,)
2 28 00 0b (h)) Her Hee-z (he) Hoa (ht) (XoxT X3x; ( (2x
1!
= Σ
(pr(Xo, (r( X)))
2=
where
(1.)) + oh::){&
(e 0),
(x x:)
V / I
ux0 (rx x,)
2 (h, X 2 X
(E = 1),
1 2-2-501 (h;;) 2 (h, [xxx]]] @ X;
(E (e 2).
Then by taking summation with respect to i and j, we derive the expansion of the kernel:
K K(XoNN+ Xb:N)
(x x.),
1
pe(x x²))
N
33

---

Then for the target function that takes the form in (71), we have the RKHS norm of g. bounded by:
30
2 2' "DVEEIID.IE..
ID.II
Σ
4 Deli.
HK
by the feature equivalence property (23). Thus, there exists coefficients v(W) such that
N
[W (Xo:N)
Σ
Ew [o[oo(W- Wo. xxxx)) (v(W), x.)
and
Ew [v(2)
220
which proves Lemma C.I.
Proof of Proposition C.I.
Note that Lemma B.1 is also applicable the model BRFA. Combining it with Lemma C.I proves Proposition C.1.
Proof of Lemma C.2.
Firstly, we consider approximating xZ using the function class (e())) which is equivalent to approximating =²6 0z)) using Hermite polynomials. We take 0 first and compute the Hermite expansion of the 26-th-order term in the Taylor expansion of ø (x) ==2 22//,ic. i.c. (e(())
Σ mn= ndd!
2a
Using Lemma A.2, we have that:
x.
(2n)!
Hezn 2m (2)
(e(x))
n"n!
n
2m (2n 2mmm
(2m 21:)!
(e))(a)
2m+ (m + k)). m(2k)
m=
ce He.(He)(()),
k=
where Ce. = (2k)
2t
Σ t-k
(2m-2k)
Similarly, for any p 2 0, define = Lg with r
m==
22m 12k (m kk)mm!
p 2po. We have that
2po
(2n + 2po + r)! n+
Hegn 280 +r 2m (x)
20n!
ndn!
2m (2n 2po 2mm!
n=
n=0
m==
E-
(o)
E+
Σ
(2m 2k: + r)!
1
2 h (m + k - po)! mm + r)!
m=max(pp k.0)
Pg
Ck Ho2x+r(x),
where Cr follows:
P
(2m-2k4r))
Then we can bound Cr.k.po as
m= k.0) 2m 200 (mm 2
Σ
(2m + 2k + r)!
1
Ct
2m 2m-k (m + -Po! 22k + r)!
m k.0)
E+po
(m + k:)! (2m 2h + r)!
n0
(m + k - ppo) 2m+ + k))
2mk + r)
m=U
34

---

t+po-k k
(2m + 2k)!
I 2e+ (21 + p)"
S (2l + p) cctTook
2m+ (m + k:)!
2=mm!
m=0
Next, we give an upper bound of Cr,bi
2
(2m + 2h)!
(21))
22m 2h (m + k)mm!
m=0
2
(2m- + 2k)!
(m + k)! m!!
(2R:)
22m 2k (m k)).m R:)!
m=0
2
1
(m + k)! m!!
(2k)
m.
2
L-k
(m + k 1)!
Σ
Vmn + k
(2k))
m!! (m + 1)!
m=0
2
f-k
S (2k)!
m!!
m=0
2
1!
Vl
(2k)
kll
k))!
Here we use the inequality that
(2n)
S
vn 2
T (2n 1)
2 n!n!
2n
and we will use it again in the following proof. Since the function
1 N (xxx xx)) x.))
Σ Xo, x
x0 AA)
can be written as
E+po
u2 Xx), pe.k A, s (
K=
where e = diag(0, 0, ..., 1). By Lemma C.I, there exists Vp.e( (W) s.l.
N
WW = (b (xo. xx))) xx))
[xo
x, (0)))((() (x 8 3
N
1=1
with
Ew [[Vxe(W]]1]] < |IA
2
Σ
2k+r+2 (2k + 2)
k=0
Notice that
Po
2k+r+2 22k r +2)1c2
P+
+ (2k + + + poot cc+po..)
k=0
2h
(e+ po)!
< (2l + p) +2r
2b+r+ (2k + r 2)1
V Ve+ Pu k(E + Po k))
(2k)
k=0
35

---

E 42k+2 (2k: + 2)1
2h
(kp)! 1
44(2 +
pp) pp +2 (20 + 2po + tr + 2)
Vep
(2k)
k(tp)!
k=0
where lp = l Po. and by the inequality that
fp
2
(cp)
4
2 2)
V
(2k))
k=0
Pp
(2k: 2) +
2
Ep!
2
24k+4
1p
12
(2k)
(Cp k:)!
k=0
4p
24k+4 pp
(2k: + +2k + 1)
π
(2k: + 1)
2
1-2
2
kl(Cp
1:)1
k=0
Ep!
2
< x 2 24p++
Σ
kllep k:!
(2pp)
4V22 x 240
lp
< <4v2 x e (64)&,
obtain upper bound:
8+po
Ew [IVp.e (WIl
< (2e + p) 220
|A Il:
f2ktr+2(2k +r +2 2)1c
k=
A, 2 4(2 + p) 2po 2r (2e + 2p0 + + 4V2 e- (64)
A 2 (26 + +p)ppp+(g((+2)))0
Finally, we consider the target function (16)
N
1
g. (Xo:N) = N
x:). F(t) = Σ apt,, G(Xo, x,) = 02 p=U
-x
x A.)
By approximating # EeS Og (xo, x,) G(x0, x,) separately and adding the approximation functions together, we obtain V2 s.t. WWo (Xu:N) ol (Xx, x; ())))((((x))))) x:))g. (Xu.N) and
S Ew [1 (W)1
apg/tp/2+3/2 (28 + 1+3
||A. 1F- 83 B(ga, L)
P=
and [|g. - /. 1((ovb 1)g, 1loo S | 111 11 S कहापीेााग 119 11.. Ee.
This finishes the proof of Lemma C.2.
C.2 Proof of Theorem C.1
Proof. For any L > using Lemma C.2, we can find function fWo such that ||g. - S 11 S lIg-1l E. with B(f < 83 Big. L). Follow the same manner as the proof Theorem 2. We can get that
LD [..) -LW
(ogldM)log(nN)) + log(8-1) +
(d2 + log M )6
x B(B,,, LJ
M
Therefore, we have that
LW.F.Wo, LD(g.)
LDD ww. - LD(SWO) + LD( L
log(dM) log(nN + log(0-1)
( + log M)0
S V B(g., L)
+
+ 6r. 11 (77)
n
M
Taking infimum over L proves Theorem 3.
36

---

C.3 Proof of Examples in Section 4
Proposition C.2 (Restatement of Example 4). Consider the target function
N
9. N
A).
has norm bound ((g., L) || L L3821, So for any lj O. if we take n 2 exp(exp(@(L/n)). L - @ + log log n) I.), and M = n)) the excess risk will scale as A
Proof of Proposition C.2. The value of (gg., L) can be obtained by definition and by direct calcula- tion. As for the second part of the proposition, take L = r log n. where >> 0 is a parameter that is to be chosen to minimize the excess risk. Eq. (17) becomes
O(I 3r-12/1112 1/2 +
log
(78)
where Ô hides all the logarithm factors and constants of n, d, and M. To minimize the excess risk,
need to make the two terms in (78) have the same scale. So set 3r - 1/2 -r logl (2/c)
r logr - r log log n. Denoting the solution as T., we get that
1/2 logr,
T.
log(2e2) log log
So this gives r, < (1/2 + 1/eg(2))( + log log Assume r, < (otherwise 3r, > 1/2 and the excess risk is meaningless), then r. (2 + log log n)- Therefore, we get that r. C(1 + log lo - with C being a constant. As result, when choosing L = logn O + log lo log n), the excess risk Eq. (78) scales as (1 (23r-=)
O[I Fr, n C/(++lg log ). As a result, let n > exp(exp(C/n) - 1)) where C a constant, we This finishes the proof of Proposition C.2.
get excess risk that scales as
Proposition C.3 (Restatement of Example 5). Consider the target function
N
Xo. (B,x), B BESd-1
N
It has Bgg., L) -2L + For any 11 >> O, choosing the same parameters (n, L, M) as Example 4, the excess risk bound scales as Õ /llog n + p))++3)8 (op++)88/nnn-> /n!-").
Furthermore, to reach an accuracy 0.01, the BRFA model requires n. = ((Spp+88)) whereas the RFA model requires = ={{dd))).
Proof of Proposition C.3. The value of ( L) can be obtained by direct calculation, and we use the same method as the proof of Proposition C.2 to get the Ō 11/nn-n) excess risk. To reach an accuracy of 0.01, we can set L = 3, note that EL ||g. aa 0.006. Therefore we just need choose n. s.t. n > (( + +p)++tg+++) and this gives n, ((Sp + 48)++))++ for the BRFA model. Direct calculation using Theorem gives the value of n. for the model RFA. This finishes the proof of Proposition C.3.
Proposition C.4 (Restatement of Example 6). Consider the target function that has the form
1 N
Σ
cos cxs( xx)) (x. G).
N
has L) ((8))2L). For any n 0, choosing the same parameters as Example 4, the excess risk bound scules as Õ 1/n
Furthermore, to reach an accuracy of 0.01. The BRFA model requires nt = 5(1), whereas the RFA A model requires n., ((Poly(d) exp( exp(Vd)).
37

---

Proof of Proposition C.4. We use the expansion of the cos function,
1
(-1) (2k)
(Xo, xx) 2k
(x. 08 G).
9+ N
0=
Then use the Eq. (17), we get that
(2L 26)
(2k + +3) (8e)" = -((ee)))
B(g., L)
(2k)
(2k)!
k=
This gives the formula for Bgg., L). To reach accuracy of 0.01, we set L = 3. Then B L) (1). So BRFA needs n, = (1). And Theorem 2 and Example show that RFA model would require n, = ((Pol(dd) (xp((dd))ppppp)) exp(vd)) This finishes the proof of Proposition C.4.
D Further experiments
In addition to Section 5, we perform further simulations to examine our theory upon the approximation power of RFA (5), BRFA (15), and RFMLP (14). Besides the target function (21), we consider two additional target functions of the form
PEN, BES--1, (x.:N) = N. En (xo, Sx,) (8,x;), B BEST S +
(X,n() (3)x))"
(79)
(80)
The target function (79) is specific instance of Example 1, and the target function (80) is a specific instance of Example 3. In (80), we sample Ziy~iaN 1/d) for (i.j) E [d]² and vary y in the experiment. Other experimental settings are the same as in Section 5.
Figure 3 demonstrates the effect of sequence length N on the performance of three RF models when fitting the target function fa.2 (79), which solely depends on Fp- We observe that larger sequence length N leads to larger separation of the test error between RFMLP and the other two random-feature attention models. This result aligns with our sample complexity analysis detailed Example 1, where the sample complexity bound of RFMLP for learning average of functions of is found to be ((N/4))) times greater than that of RFA.
Figure 4 demonstrates the performance comparison between RFA and BRFA under different choices of the token dimension As we can see from the left panel of Figure 4, in the case of d 4, RFA outperforms BRFA for large sample size, which may result from that BRFA has slower convergence rate with respect to the sample size as we state in Example 5. the middle and the right panel (i.e., when d is larger), BRFA exceeds RFA, and the largest separation lies in the case of d = 32. This result consistent with our analyses that BRFA saves a Poly(d) factor the sample complexity for d >>>>> p.
Figure 5 demonstrates that BRFA has no advantage in approximating the target function f4,7 for = 0 (left panel; i.c.. (S 1/d), for (i.j) E [d]3). However, as 7 increases, BRFA outperforms RFA and their separation increases with larger (middle and right panel). Notice that lim, (Xg. Sx,) / (Xo, x)). So this result also conforms to our analysis that BRFA is adept at approximating functions of correlations in Example 5.
In conclusion, RFA and BRFA have similar performance in fitting functions without correlation structure, such as (20) and (79), and RFA may behave even better in some cases. However, BRFA is presumably more powerful than RFA approximating functions of correlations.
D.1 Weight matrices in BERT
As noted in Section 4, we plot the query-key matrices (weight matrices) of the BERT-Base model [29]6 and show that many query-key matrices are diagonally dominated. The BERT-Base model has 12 attention layers with 12 heads in each layer. Denote the query matrix in the i-th head of j-th layer as Qij E T 76 and the key matrix as Kij E 768x61 We compute W,, = v768. Q1, KT for i, from to 12, and then take the absolute value for each entry of the weight matrices Wiy. Figure
DDwnloadddd from https: https::/huuggiggaacocoo wwr.tettteeeeeeeeeeeeee www.bbbrttttteeeeeeeeeeeeeeeeeee uncased.
38

---

N 8
N=
1.4
1.4
1.4
14
24
14
6.4
6.-
6.6
RFA
RFA
RFA
BREA
BRFA
DRYA
Figure 3: Test error of three RF models for learning fa.2 (79). We fix d 16 while varying N = 4 (left), 8 (middle), and 16 (right). The other settings are the same as Figure 2.
d =
d=16 16
d= 32
1.4
1.4
3.4
Etrour
0.2
08
RFA
RFA
RFA
BRFA
6.2
10
8.2
16 257
8.0 16
Figure 4: Test error of three models for learning f2.3 (21). We fix N 16 while vary d (left), 16 (middle), and 32 (right). For a fair comparison among RFA, BRFA, and RFMLP, the number of heads of RFMLP taken be MдTLE 9000 (left), 17000 (middle), and 33000 (right). The other settings are the same as in Figure 2.
12
12
1.8
3.5
0.4
MLP
54 MLP
0.4
MLP
BRFA
BRFA
02
BRFA
Figure 5: Test emror of three RF models for learning f4,y (80). We choose y = 0 (left), 2 (middle), and 8 (right). The other settings are the same in Figure 2.
39

---

shows the heat maps of weight matrices of the 2nd, 5th, 8th, and 11th layers, where all matrices are clipped the top-left 32 32 block. As we can see, lot of weight matrices are diagonally dominated. We remark that a very recent and concurrent paper [84] observed similar phenomena for the ViT-Tiny model. They further show that diagonally dominated weight initialization of self-attention layers allows training transformers faster and obtaining higher final accuracies on CIFAR-10 and ImageNet datasets.
40


Reviews:
[{"review_id": "BB0Wjm01bN", "reviewer": "Reviewer_5oTn", "summary": "This paper first studies the expressive power of a random-feature attention layer and then provides the generalization gap of the attention layer. A sample complexity bound is shown, which indicates a larger number of attention heads help the generalization. This paper also compares the results between attention and MLP layers given several target functions. Numerical experiments support the theory. I am willing to update my score after the rebuttal if my concerns are addressed with revisions in the manuscript.\n\n-------------------------------------------------------------------\nAfter rebuttal, I increased my score to 5. Please see comments below.", "questions": "1. Is the comparison in Sections 3.3 and 4.2 a comparison between upper bounds (sufficient conditions) for RFA and RFMLP? I think so. If so, this should be mentioned. This is not very rigorous, but I am ok with it because it is difficult to compare an upper bound and a lower bound.\n\n2. What does \"permutation invariant of the target function\" mean? Why is it needed? ", "limitations": "There is no negative societal impact of their work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is clear and well-written. The proof seems to be solid.\n2. The random feature analysis of attention layers is new to this community. The analysis combines the expressive power and the generalization gap, which is better than some existing works.\n3. The comparison between attention layers and MLP covers a lot of target functions, which is impressive.", "weaknesses": "1. Equation 3 is important for further derivation. From my understanding, Equation 3 indicates that the attention map is fixed for all data. I think attention layers usually have different attention maps for different data. I believe this is a big limitation of this work. Some clarification about \"fixed attention\" is needed to avoid misunderstanding.\n\n2. An attention map fixed at random initialization makes the attention layer useless and meaningless. I think at least the attention layer should be trainable.\n\n3. About the related works of generalization analysis of Transformers, the references discussed are too old. Here are some recent references. I would like to see a discussion of these works (Some of these works are concurrent works). \n[1] Jelassi et al., 2022, \"Vision Transformers provably learn spatial structure. \"\n[2] Li et al., 2023, \"A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity. \"\n[3] Oymak et al., 2023, \"On the Role of Attention in Prompt-tuning. \"\n[4] Tarzanagh et al., 2023, \"Max-Margin Token Selection in Attention Mechanism. \"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments and suggestions.  \n\n> Equation 3 is important for further derivation. From my understanding, Equation 3 indicates that the attention map is fixed for all data. I think attention layers usually have different attention maps for different data. I believe this is a big limitation of this work. Some clarification about \"fixed attention\" is needed to avoid misunderstanding.\n\nCould the reviewer kindly clarify what \"fixed for all data\" means in the context of Equation 3? We want to make sure we fully understand your perspective to make appropriate clarifications in the paper. As in standard practice, our attention matrices $(Q_m, K_m, V_m)$ are the same for all data points (cf. Equation 1).\n\n> An attention map fixed at random initialization makes the attention layer useless and meaningless. I think at least the attention layer should be trainable.\n\nOur attention layer is still learnable since the value matrices $V_m$ are learnable. We only freeze the inner layer (the query-key matrices). This corresponds to the Random Feature setting, which is widely considered in deep learning theory. \n\n> About the related works of generalization analysis of Transformers, the references discussed are too old. Here are some recent references. I would like to see a discussion of these works (Some of these works are concurrent works). [1] Jelassi et al., 2022, \"Vision Transformers provably learn spatial structure. \" [2] Li et al., 2023, \"A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity. \" [3] Oymak et al., 2023, \"On the Role of Attention in Prompt-tuning. \" [4] Tarzanagh et al., 2023, \"Max-Margin Token Selection in Attention Mechanism. \"\n\nWe appreciate the reviewer for suggesting these works, and we will make sure to incorporate and discuss these studies in our revised version. A key distinction between these works and ours is that they only focus on a small class of target functions with special properties, whereas our work covers a large generic class of target functions (polynomials of input tokens of any degree).\n\n>Is the comparison in Sections 3.3 and 4.2 a comparison between upper bounds (sufficient conditions) for RFA and RFMLP? I think so. If so, this should be mentioned. This is not very rigorous, but I am ok with it because it is difficult to compare an upper bound and a lower bound.\n\nOur comparisons are indeed between upper bounds. Existing work has also derived lower bounds on the sample complexity of RFMLPs (e.g., learning degree $p$ polynomial by RFMLP requires $\\Omega((dN)^p)$ samples, (Ghorbani et al., 2021)), which agrees with the upper bound for RFMLP we used though they apply to a special case with a uniform distributional assumption (input vector uniformly distributed on the sphere). \n\nWe will make sure to emphasize this point and add a discussion of the lower bound in our revision.\n\n>What does \"permutation invariant of the target function\" mean? Why is it needed?\n\nFor any function $f(X_1,\\ldots,X_n)$ of $n$ tokens, it is permutation invariant if $f(X_{\\sigma(1)},\\ldots,X_{\\sigma(n)}) = f(X_1,\\ldots,X_n)$ for any permutation $\\sigma:[n]\\to[n]$. We consider permutation invariant target functions, as attention layers can only fit these functions (due to the structure of attention heads). We will clarify the meaning of \"permutation invariance\" in our revision."}, {"review_id": "Seue7ikOLY", "reviewer": "Reviewer_6d3X", "summary": "The paper examines the capabilities of a single-layer multi-head attention layer in a scenario where the Key and Query matrices are predetermined and randomly selected from a Gaussian distribution. The only modifiable component is the Value matrices, and when provided with a convex loss, the minimization problem becomes convex.   The authors establish expressivity results showing that, whether there is or not there is bias in the Key and Query matrices, the model can effectively learn a class of functions that exhibit permutation invariance to the Key vectors. Furthermore, they demonstrate that the sample complexity of their model is superior to that of two-layer random feature networks for the specific function class. The attention model investigated uses ReLu rather than the more common softmax attention", "questions": "- Why not normalize the model with 1/M rather than including this scaling afterwards on the v constraint set? Isn't that also more revealing of the \"proof\" strategy that you compare to expectation?\n- Eventually, can the authors comment on where they see this study leading to? What specifically revealing does it show about attention and what are the authors thoughts on the use of random features in this setting? \n- How would the results chance if instead of W being Gaussian, K, Q with lower inner dimension are Gaussian?\n- If not mistaken I think I have seen in the literature some works experimentally validating performance of random feature attention (although with softmax). This might be worth taking a look as it might help the narrative (sorry that I don't remember on top of my head)\n- Can you comment on how to derive the last inequality in line 736 of Lemma b3\n- In Eq. (25), is it \\sum_{i=1}^N or \\sum_{j=1}^n?\n- How you guys compare the results in line 263 and 264.\n- Have you considered running your own experiments for softmax function instead of ReLU (to see the possible difference or implications for your results)?\n", "limitations": "n/a", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Results are new and although not surprising require certain efforts to prove and could be a useful addition to the literature of random-feature -type models (more to that literature I think rather than to the attention/transformer literature)\n\n- Paper is well written (but please correct the many typos) and the authors provide comprehensive explanations of limitations ", "weaknesses": "- A major weakness is that instead of the common softmax function, the authors opt to use the ReLU function. \nThe use of ReLu is nonstandard in transformers and should be mentioned explicitly in the abstract and contributions. Relu and softmax actually have rather different properties and this should be clarified\n\n- The fixed Key and Query matrices restrict the learning setting to a linear problem.\n\n- The (as the authors admit) seemingly unnatural constraint set in (12) is rather \"artificial\". Why is such a constraint needed? Besides, knowledge of this K1,K2 requires bounds on B(f_*) making it rather impractical\n\n- Thm 2 only applies for bounded loss. Is it possible to extend to say square-loss?\n\n- Comparing RFA to RFMLP is based on comparing upper bounds to each other. Is it known if the latter bounds are tight?\n\n- As the authors acknowledge, it is rather expected that RFA would beat RFMLP for the specific function class that involves correlations between key and query tokens. While the analysis is non-trivial it is questionable what is that we really learn from those bounds regarding attention?\n\n- Even though the results are new and the proofs require efforts, the techniques are rather standard and is not clear if they are revealing of any special properties of attention? If so, this would be interesting to emphasize.\n\n- There are a lot of typos\n\n- My overall concern with the paper is not about the motivation of the setting (use of ReLu and random features) and also it is a bit unclear what the take home message is (other than a technically solid analysis in a mathematically interesting setting)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback. We respond to the comments as follows. \n\n### Response to questions regarding the setting and message:\n\n>A major weakness is that instead of the common softmax function, the authors opt to use the ReLU function.\n\nA significant portion of our results (such as the generalization bound, and expressing functions within the infinite-head random feature space (cf. Eq(26)) using finitely many heads) would not change if we change the ReLU activation to the softmax. However, the ReLU activation enables a more precise understanding of the infinite-head random feature space, which we show includes natural target functions such as polynomials. This property about ReLU has also been used in many existing work on random feature models, such as Arora et al. 2019 and the many follow-ups after.\n\nWe agree that extending our analyses to random features with the softmax activation (especially studying the function space it can express) is an interesting question, which we would like to leave as future work. \n\n>The fixed Key and Query matrices restrict the learning setting to a linear problem.\n\nOur random feature model is linear in the parameter $(v_m)$ but **non-linear in the input $(x_i)_{i=0}^n$**. Such random feature models capture many non-trivial aspects of fully learnable neural networks, and have been a central topic of study in the deep learning theory literature (e.g. Daniely 2017, Mei & Montanari 2021, and the many references therein). \n\n> what is that we really learn from those bounds regarding attention?\n\n> Even though the results are new and the proofs require efforts, the techniques are rather standard and is not clear if they are revealing of any special properties of attention? If so, this would be interesting to emphasize.\n\n> My overall concern with the paper is not about the motivation of the setting (use of ReLu and random features) and also it is a bit unclear what the take home message is (other than a technically solid analysis in a mathematically interesting setting)\n\n\n\nOur results convey several new messages, including\n* The sample complexity of random-feature attention does not depend on the sequence length, which contrasts with the random-feature MLP model (Section 3); \n* A non zero-mean query-key matrix could further improve the sample complexity for learning certain functions of correlations (Section 4). \n\nBoth results are not present in existing random features theory (for fully-connected neural networks). Please refer to our Additional Response to All Reviewers for more details.\n\n\n### Response to technical questions:\n\n> The seemingly unnatural constraint set in (12) is rather \"artificial.” \n\n(12) has two constraints, the total norm and the total norm squared. The total norm squared constraint is essential and natural (equivalent to an $L_2$ regularization on $(v_m)$). By contrast, the total norm condition is a technical condition merely for a slightly tightened sample complexity. \n\n> Thm 2 only applies for bounded loss. Is it possible to extend to say square-loss?\n\nFor squared loss with unbounded labels, we believe a similar result would still hold by using standard machineries such as truncation. We assumed the loss is Lipschitz and bounded at $0$ for simplicity only.\n\n> Comparing RFA to RFMLP is based on comparing upper bounds to each other.\n\nOur comparisons are indeed between upper bounds. Existing work has also derived lower bounds on the sample complexity of RFMLPs (e.g., learning degree $p$ polynomial by RFMLP requires $\\Omega((dN)^p)$ samples, (Ghorbani et al., 2021)), which agrees with the upper bound for RFMLP we used though they apply to a special case with a uniform distributional assumption (input vector uniformly distributed on the sphere). \n\nWe will make sure to emphasize this point and add a discussion of the lower bound in our revision.\n\n> Why not normalize the model with 1/M\n\nWe did not include any normalization for simplicity of the presentation only, and we agree that normalizing by $1/M$ can make the model more intuitive. We will carefully think about whether to add this normalization this in our revision.\n\n> How would the results chance if instead of W being Gaussian, K, Q with lower inner dimension are Gaussian?\n\nIf $Q,K$ are Gaussian and have large inner dimensions, $W=Q^\\top K$ would be have similar to a Gaussian and thus we don't expect our results to change much. However, if $Q,K$ have small inner dimensions, then $W$ would be low-rank and thus behave very differently.\n\n> How to derive the last inequality in line 736 of Lemma b3?\n\nThe equation in line 736 of Lemma b3: We first drop the $\\log(dM)/n$ term since $n\\geq \\log(dM)$. Then we use the bound for the expectation of the maximum of the sub-gaussian r.v.'s.\n\n>How you guys compare the results in line 263 and 264.\n\nThe arguments can be found in the proof of Proposition C.2 (Line 919-930).\n"}, {"review_id": "zJIaE1PdrL", "reviewer": "Reviewer_3Deb", "summary": "The paper considers the representational and generalization properties single-layer scalar-valued transformer models with random key and query matrices and value vectors that can depend on those random matrices. They draw a comparison to the well-studied random-feature models for two-layer neural networks. Concretely:\n* Theorem 1 proves that functions of the form $f_*(x_{0:N}) = \\frac1N \\sum_i F(x_0, x_i)$ can be efficiently represented under the random feature, with quadratic dependence on the input dimension $d$ and no dependence on the sequence length $N$. \n* Theorem 2 extends this approximation-theoretic result to generalization by proving a generalization bound on the empirical risk-minimizing random feature transformer (among transformers with bounded value vectors) that fits a noiseless dataset. The proof follows from bounds on the Rademacher complexity of the family of functions that approximately represent the dataset.\n* The paper gives several examples in Section 3.3 of functions of the above form and show that they admit much stronger learning rates for random feature attention models than for standard random feature models. Generally, standard random feature models have a substantial dependence on $N$, while random feature attention has no such dependence. \n* Theorem 3 proves a generalization bound similar to Theorem 2, but in the regime where the random feature matrices (e.g. the product of the key and query matrices) is biased in favor of larger elements on the diagonal. (This is empirically motivated in the appendix by a finding that BERT weight matrices frequently concentrate mass on the diagonals.) They prove generalization bounds for a restricted family of target functions, where high-degree polynomials of $\\langle x_0, x_i\\rangle$ and low-degree polynomials of $x_0 \\otimes x_i$ may be averaged together. They provide several examples showing how this model can reduce the error rates over the standard random feature attention model.\n* Numerical experiments in Section 5 validate their theoretical results by comparing the error rates as a function of sample complexity of random feature MLPs, random feature attention, and biased random feature attention.", "questions": "Would you mind explaining if you expect that the $\\delta^{-1}$ dependence in Theorem 1 could be improved? As far as I am aware, other random feature approximation papers can often achieve a $\\sqrt{\\log \\delta^{-1}}$ dependence by employing concentration bounds over Hilbert spaces instead of Markov or Chebyshev, as is employed in the proof of Lemma B.1.\n\nI would be interested to know why in particular you chose to focus on the diagonally-biased weight initializations, when the plots in Figure 5 appear to make similarly strong cases in favor of selecting random feature matrices $W$ from distributions that impose either sparsity or low rank. (I am okay with these regimes not being included in the paper; I am mostly just interested in the choice.)\n\nWould it be possible to offer a more concrete comparison with the Rademacher complexity generalization bounds in this work and single-layer covering numbers-based bound of [Edelman et al](https://arxiv.org/abs/2110.10090)? While the regimes are different, I think the paper would benefit from a brief comparison of the advantages and disadvantages of each. ", "limitations": "The work makes modeling assumptions with substantial differences from standard transformer architectures, notably the random weights, the use of a single layer, the lack of softmax, and scalar rather than sequential outputs. The work is upfront those limitations.  \n\nWhile the random features regime for 2-layer neural networks is interesting for studying the representational and generalization properties of certain training regimes, it's well-known that the method (as well as all other kernel-based approaches) fall short when learning even simple learning problems like single-index models. (Or as was argued in Example 1 of this paper, how standard random feature models fall short on sequential tasks whose inputs depend exclusively on $x_0$.) Are there similar illustrative examples that random feature attention units fail to efficiently approximate? \n\nMoreover, when providing examples that separate random feature MLPs and attention models, it may be helpful to account for a particular examples of a target function where the random feature MLP model offers a better error bound, or to make an argument that RFA models will always have superior rates under certain sequential assumptions.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The work is novel, interesting, and relevant to the growing study of the theoretical properties of transformers. The work formalizes some intuitive advantages that attention models hold over standard MLPs in their ability to compute and aggregate pairwise functions of sequential inputs. Theoretical results are presented cleanly and the proofs that I read appeared correct. The work is creative, and draws inspiration for Theorem 3 from empirical observations. There is interesting follow-up work to be done on understanding the strengths and limitations of this model, especially on the optimization front.", "weaknesses": "While the bounds are interesting in their own right, the comparisons between random feature attention and MLP models focus on a few particular examples, whose generality is unclear. Moreover, the comparisons are between upper bounds for both models; ideally, the results would contrast with _lower_ bounds for random feature MLPs.\n\n### Minor pointers\nl244: \"scaler\" -> \"scalar\"\n\nl657: $W$ in equation block should be $W_m$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and the suggestions on our paper. We would appreciate if you could champion our paper in the discussions!\n\n>While the bounds are interesting in their own right, the comparisons between random feature attention and MLP models focus on a few particular examples, whose generality is unclear. Moreover, the comparisons are between upper bounds for both models; ideally, the results would contrast with lower bounds for random feature MLPs.\n\nOur comparisons are indeed between upper bounds. Existing work has also derived lower bounds on the sample complexity of RFMLPs (e.g., learning degree polynomial by RFMLP requires $\\Omega((dN)^p)$ samples, (Ghorbani et al., 2021)), which agrees with the upper bound for RFMLP we used though they apply to a special case with a uniform distributional assumption (input vector uniformly distributed on the sphere).\n\nWe will make sure to emphasize this point and add a discussion of the lower bound in our revision.\n\n>Would you mind explaining if you expect that the $\\delta^{-1}$\n dependence in Theorem 1 could be improved? As far as I am aware, other random feature approximation papers can often achieve a \n $\\sqrt{\\log \\delta^{-1}}$ dependence by employing concentration bounds over Hilbert spaces instead of Markov or Chebyshev, as is employed in the proof of Lemma B.1.\n\nThe $\\delta^{-1}$ dependence comes from the Chebyshev bound in Lemma B.1, which in turns comes from the precondition $\\mathbb{E}[\\|v(W)\\|_2^2]\\le R^2$. Since we only assume a bounded second moment, the Chebyshev bound is likely the best achievable. Further, bounded second moment is likely the best we can do for expressing polynomials due to RKHS norm related arguments (cf. the proof of Lemma B.2).\n\n>I would be interested to know why in particular you chose to focus on the diagonally-biased weight initializations, when the plots in Figure 5 appear to make similarly strong cases in favor of selecting random feature matrices $W$ from distributions that impose either sparsity or low rank. (I am okay with these regimes not being included in the paper; I am mostly just interested in the choice.)\n\nWe acknowledge the reviewer's point that other patterns such as low rank or sparsity may be present in Figure 5. However, the diagonal pattern appears the most predominant, which motivated our setting here. \n\n>Would it be possible to offer a more concrete comparison with the Rademacher complexity generalization bounds in this work and single-layer covering numbers-based bound of Edelman et al? While the regimes are different, I think the paper would benefit from a brief comparison of the advantages and disadvantages of each.\n\nWe will incorporate a discussion of the work cited by the reviewer. However, we note that the settings differ substantially. The previous work considers full attention where $Q, K,$ and $V$ are all trainable, whereas we analyze random feature attention with only a trainable $V$. Consequently, the results of Edelman et al. bound a larger quantity than our result and thus are not directly comparable.\n\n\n>Limitations\n\nWe appreciate the very detailed suggestions on the limitations of our work, such as target functions that cannot be approximated by RFA or those requiring more samples than RFMLP. We will think carefully about these limitations and make sure to add a discussion of them in our revision. \n"}, {"review_id": "tcp4TQb4ex", "reviewer": "Reviewer_giTQ", "summary": "This paper explores the learning capabilities of a single attention layer, assuming keys and queries to be random and frozen (as in the random features model).", "questions": "You consider single query token models. This semplify the derivation and I agree that is a reasonable simplification. Can you still argue more on this choice? Is this equivalent to consider a model with $N$ queries, after a left inner product with the canonical basis vector $e_1$? Is this method used in practice in single-output attention layers? Has it been considered in previous theoretical work?\n\nAre your results optimization agnostic? There is any possible way to bridge the solutions you mention in Theorem1 with the solutions found by a GD algorithm? In this sense, are the assumptions over the constrained class $\\mathcal V_M$ going in that direction?\n\nIt is not clear to me in Theorem2 if your minimizer $\\hat V$ is unique, or if you don't need it, as any $\\hat V$ would respect your claim.\n\nCan you provide more intuition of why the permutation invariance is a property that is intrinsic of tasks where attention outperforms fully connected models?", "limitations": "No need", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is very well written, and the technical claims look formally supported.\n\nThe paper deals with an important problem, which is to theoretically better characterize the representation power of single attention layers.\n\n", "weaknesses": "Typo in line 95 \"theorey\"\n\nThe family of target functions considered in the discussion doesn't contain terms that consider the interaction between token $i$ and token $j$, with $i, j \\neq 0$. Attention model success also hinges on capturing the relation between different tokens in the context. It looks like this work leverages on something orthogonal, and I wonder if the results are really representative of the attention layer power.\n\nTypo in line 176 \"$i$\"\n\nIt is not clear to me how the permutation invariance of the input tokens represents a valuable propertyy over which attention should perform better than MLPRF. It could be that the toy-model analyzed has this power as a natural difference with respect the standard MLPRF, without really capturing the properties of attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback. We respond to the questions as follows.\n\n> You consider single query token models... Has it been considered in previous theoretical work?\nThe family of target functions considered in the discussion doesn't contain terms that consider the interaction between token $i$ and token $j$ with $i \\neq j$. \n\nOur single-query token model is equivalent to **simply looking at the first output token** of a full self-attention (cf. Line 112-116), and abstracting out the query token to be $x_0$ and the key tokens to be $x_{1:N}$. This does not change the model at all, and merely avoids repeating our arguments $N$ times for the $N$ output tokens. Our results can be easily mapped back to full self-attention by repeating our arguments on all $N$ output tokens. In that case, the attention model would involve interaction between any $x_i$ and $x_j$ (instead of just $x_0$ and $x_j$).\n\nThere exist many literature that constructs sequence-to-sequence transformers to implement a sequence-to-single-token function, by \"reading out\" from a certain output token (such as the last output token), for example Garg et al. (2022), Akyurek et al. (2022). Our model can be seen as a single-layer version of these transformers. We will add a discussion about this point in our revision. \n\n> Is this equivalent to consider a model with $N$ queries, after a left inner product with the canonical basis vector $e_1$? Is this method used in practice in single-output attention layers?\n\nIt is equivalent to using only the first token as the query. It is not equivalent to applying $e_1^\\top$ on the left to all $N$ query tokens. \n\n> Are your results optimization agnostic? There is any possible way to bridge the solutions you mention in Theorem1 with the solutions found by a GD algorithm? In this sense, are the assumptions over the constrained class $\\mathcal{V}_M$ going in that direction?\n\nOur results are optimization agnostic. Any standard convex optimization algorithm (including GD) for the constrained problem on $\\mathcal{V}_M$ (or an equivalent regularized problem) can find an approximate solution efficiently. In our experiments, we used Adam with weight decay, which we found was a good enough optimizer on all our problem instances.\n\n> in Theorem 2 if your minimizer $\\hat{V}$ is unique, or if you don't need it\n\n The minimizer $\\hat{V}$ indeed may not be unique, and we don't need uniqueness in Theorem 2; the statement holds for any minimizer (due to the uniform concentration argument).\n\n> Can you provide more intuition of why the permutation invariance is a property that is intrinsic of tasks where attention outperforms fully connected models?  \nIt is not clear to me how the permutation invariance of the input tokens represents a valuable propertyy over which attention should perform better than MLPRF. It could be that the toy-model analyzed has this power as a natural difference with respect the standard MLPRF, without really capturing the properties of attention.\n\nOur attention models can only fit permutation invariant target functions due to the structure of attention heads. On such target functions, our random-feature attention (RFA) models do achieve better sample complexities than RFMLP as they exploit this structure.\n\nWe remark that we also have results on comparing different weight distributions in RFA, where query-key matrices with non-zero means achieve better sample complexities than zero-mean ones for learning certain functions (Section 4). These results are more intrinstic to the attention structure and do not have analogues in RFMLP models to our best knowledge. See our Additional Response to All Reviewers for more details."}]

Author rebuttal:
**Additional Response to All Reviewers**

We thank all reviewers again for their valuable feedback on our work. 

Here we would like to highlight our contributions again, which we believe were missed by some reviewers:

- The random features (RF) model is an important and widely-studied model in deep learning theory (e.g; Daniely 2017, Mei & Montanari 2021, and the many references therein), widely studied for fully connected neural network. This paper **initiates the study of random feature attention models and provides a set of end-to-end learning results**, which we believe should be of broad interest to the community and could inspire follow-up works.

- Our Section 4 considers the case where the $W$ matrices in random-feature attention are initialized as identity plus standard Gaussian, going beyond the standard Gaussian initialization. This setting was motivated by the empirical weight patterns of *pretrained* BERT models. We show that such initialization has non-trivial advantages (in sample complexity upper bounds) over standard Gaussian initialization, for learning certain functions depending on the correlation of tokens. This is a new message that **did not have an analog in existing random feature theory** to our best knowledge.

- Our paper introduces several new techniques, such as sharp analyses of the Rademacher complexity for random-feature attention models (Line 188-193), as well as approximating functions of correlations using identity-biased random-feature attention (Line 283-288 & concrete statement in Lemma C.2). We believe these techniques **could be useful for future work on analyzing attention**.

We would be grateful if the reviewers could reconsider the evaluation taking into account our contributions. We will highlight these points more clearly in our revision. Additionally, we appreciate the reviewers for catching the typos, which we will fix in the revision.


Meta-review: This paper explores what a single attention layer learns through random feature analysis. The analysis provides valuable insights and is of interest to the community. The paper has received generally positive reviews from the reviewers, and the rebuttal has addressed most of the concerns. One remaining issue is the use of ReLU instead of softmax. I understand that this choice is primarily driven by the challenges in analyzing the softmax function in a general context, and the authors have referenced a study [73] suggesting that ReLU could be comparable to softmax (although this comparison has yet to be validated across a wider range of models and datasets). One potential improvement, as also highlighted by one of the reviewers, is to conduct softmax experiments within the paper's simulated environment and gather additional data points to support the claim.

[73] K. Shen, J. Guo, X. Tan, S. Tang, R. Wang, and J. Bian. A study on relu and softmax in
496 transformer. arXiv preprint arXiv:2302.06461, 2023.
Decision: Accept (poster)