Title: What can a Single Attention Layer Learn? A Study Through the Random Features Lens
Venue: NeurIPS 2023
Accepted: yes
Keywords: transformers

attention

deep learning theory

random features
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

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 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 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 over standard fully connected random-feature models; (2) Concrete and natural classes of functions that can be learned efficiently by a random-feature attention 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.Experiments on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.

1 Introduction

The transformer architecture [86] has achieved remarkable recent successes in many areas 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]-sequence-to-sequence mappings that allow each token within the input sequence to 'attend to' other tokens that are most relevant to the present token, and produce outputs based on those tokens. Attention layers implement this mechanism in a compact way that allows them to handle sequences of arbitrary length using a fixed set of parameters, a crucial reason behind their success in handling long input sequences.

Despite its wide applicability, the theoretical properties of attention layers are less well understood. While multi-layer attention networks (transformers) have been shown to be universal approximators for certain classes of functions, such as equivariant sequence-to-sequence functions [102], their results 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, 35], yet the results are either instantiated on specific target functions such as sparse boolean functions [35], or generic but arguably elusive 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 on generic and natural function classes (e.g., [8, 51, 7, 40] and many other results along the line). This motivates the following open question:

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 at 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 x 0 ∈ R d and N key tokens { x i } i ∈ [ 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 x 0: 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 a 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 an ˜ O ( √ B ( f ⋆ ) /n ) 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 input 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 a 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 zero-mean 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 a 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

Transformers The Transformer architecture, initially proposed by [86], brought about a revolutionary change in natural language processing and has been widely adopted in large language models such as 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].

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 permutationequivariant sequence-to-sequence function, and any function if positional encodings are added. Our

paper instead uses a single attention layer to approximate sequence-to-scalar 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 a 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 as 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 or random-feature 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, a line of work studies the NTK beyond MLPs, including convolution networks [53, 17, 59, 62, 18, 16], residual networks [45, 83, 4], graph networks [97, 47], and transformers [43, 98].

Although the kernel approach is a powerful tool for studying neural networks, it 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 a sequence of N +1 input tokens x 0: N = ( x 0 , { x i } i ∈ [ N ] ) ∈ X = ( R d ) N +1 , where each { x i } i ∈ [ N ] ⊆ R d represents a sequence of key vectors , and x 0 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 i as the query token and all of [ N ] as key tokens. Results obtained in this model can be directly mapped back to full self-attention, by simply concatenating N outputs generated by our model with x 0 ranging over { x i } i ∈ [ N ] . In addition, throughout the paper, we consider scalar-valued attention models, which take the sequence x 0: N as input and give a scalar output in R .

Attention layer We consider a scalar-valued, M -head, multiplicative attention layer that takes x 0: N = ( x 0 , { x i } i ∈ [ N ] ) as the input. The attention layer first applies affine (linear with bias) transformations to the input vectors to obtain {query, keys, values} at each head m ∈ [ M ] :

where Q m , K m ∈ R ( d +1) × d , v m ∈ R d +1 are the parameters of the attention layer, and ˜ x i := [ x i ; 1] for a more compact display. Then, it computes the output value by an attention mechanism

Above, σ : R → R is an activation function applied entry-wisely to each attention score ⟨ q m, 0 , k m,i ⟩ . We choose σ to be the ReLU activation σ ( t ) = max { t, 0 } throughout this paper. Notice that

this choice of the attention non-linearity is different from standard transformers [86] with softmaxattention. We remark that we choose to study the (normalized) ReLU attention for theoretical convenience, and this replacement does not change the essence of the attention mechanism. Such a choice is also recently explored in the literatures such as [80] and [95], which show that transformers with ReLU-attention perform as well as standard softmax-attention 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 a short discussion.

Simplifying the expression, we reparametrize the attention layer (1) and (2) using parameters { ( W m , v m ) } m ∈ [ M ] ⊆ R ( d +1) × ( d +1) × R d +1 :

For technical convenience, we assume all input tokens have unit norm throughout the rest of the paper: ∥ x i ∥ 2 ≡ 1 so that ∥ ˜ x i ∥ 2 ≡ √ 2 , for all i ∈ { 0 } ∪ [ N ] .

Random-feature attention models We consider a random-feature version 2 of the multiplicative attention mechanism (3), where the weight matrices { W m } m ∈ [ M ] have i.i.d. Gaussian entries 3 :

The variance is chosen to be 1 / 4 without loss of generality: this choice of variance is such that ⟨ W m , ˜ x 0 ˜ x ⊤ i ⟩ ∼ N (0 , 1) has a unit variance. The weight matrices { W m } m ∈ [ M ] are then held to be fixed during the entire learning process, whereas the value vectors { v m } m ∈ [ M ] are the learnable parameters. The random-feature attention model with input x 0: N is thus given by

Notice that random-feature attention model is linear in the parameter V , so training this model with a convex loss function gives a convex optimization problem.

Additional notation For any x ∈ R d 1 and y ∈ R d 2 , let x ⊗ y ∈ R d 1 × d 2 denote their tensor product (outer product), and x ⊗ n := x ⊗··· ⊗ x denote the n -fold self tensor product of x . For a tensor A , we use ∥ A ∥ Fr to denote its Frobenius norm. For a function f : X → R , we use ∥ f ∥ ∞ to denote its L ∞ norm. We use O ( · ) (resp. Θ( · ) ) for standard Big-O (resp. Big-Theta) relations. We use ˜ O ( · ) for hiding the multiplicative terms that are logarithmic in problem parameters, including ( M,d,n,N,δ -1 ) . We use Poly( p ) to denote a polynomial of p that is less than p C for some universal constant 0 < C < ∞ .

3 Learning with random-feature attention models

In this section, we study the expressivity and generalization of random-feature 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 a broad class of permutation invariant 4 target functions f ⋆ : X → R that takes form

A different and closely related model is the Neural Tangent Kernel [46, 32], which is however similar in

2

essence to the random feature model in terms of the sample complexity of learning, e.g. [40].

Assume that there exists symmetric tensors { f rs ∈ R d r + s } r,s ≥ 0 such that F : R 2 d → R admits representation

Note that such an expression allows F ( x 0 , x i ) to be any general nonlinear function that admits convergent Taylor expansions. In particular, any polynomials of [ x 0 , x i ] (e.g., β ⊤ x 0 , β ⊤ x i , and ⟨ x 0 , Sx i ⟩ for some β ∈ R d and S ∈ R d 2 ) are within this function class. We will discuss more specific target functions in Section 3.3.

Theorem 1 (Expressivity of RFA model) . Suppose function f ⋆ : X → R takes form (6) . Then for any input distribution P on X , with probability at least 1 -δ (over { W m } m ∈ [ M ] sampled from (4) ), there exists an M -head RFA model (5) with coefficients V = { v m } m ∈ [ M ] ⊆ R d +1 that approximates f ⋆ in L 2 ( P ) up to error

In addition, the norms of the weight of this random-feature attention model are bounded as

Here B ( f ⋆ ) is a complexity measure of f ⋆ defined as

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 1 is contained in Appendix B.1. Our proof relies on standard analyses of infinite-width random feature model with ReLU-Gaussian kernel, combined with a sampling argument to obtain approximation with finite-width.

This theorem is applicable to general functions with a finite B ( f ⋆ ) norm. The 4 k scaling of C k in the summand of equation (10) seemingly confines the target function class to those with exponentially fast decaying ∥ f rs ∥ Fr , which suggests a relatively narrow target function class. However, as 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 { x ( j ) 0: N , y j } j ∈ [ n ] ∼ iid P , where x ( j ) 0: N = { x ( j ) i } 0 ≤ i ≤ N is the j -th token sequence with length N +1 , and y j is the label corresponding to the i -th token sequence. Assume that we are given a loss function ℓ ( ̂ y, y ) that is 1-Lipschitz in ̂ y , and ℓ (0 , y ) ≤ 1 for any y . The population risk is then given by L D ( f ) = E ( x 0: N ,y ) ∼ P [ ℓ ( f ( x 0: N ) , y )] . We consider the empirical risk minimization (ERM) over the RFA model (5),

where the constrained class V M is given by

with K 1 and K 2 being two constants. Theorem 2 below provides the excess risk bound for the empirical risk minimizer.

Theorem 2. Assume M > δ -1 and n > log( dM ) . Let f ⋆ be the minimizer of the population risk L D ( f ) within the target function class (6) (7) . Let ̂ f W M = f W M ( · ; ̂ V ) be the empirical risk minimizer given by (11) , where in (12) we choose K 1 = C √ B ( f ⋆ ) and K 2 = CB ( f ⋆ ) δ -1 , with C being a constant. Then for any joint distribution P , with probability at least 1 -δ over { W m } m ∈ [ M ] sampled according to (4) and { ( x ( j ) 0: N , y j ) } j ∈ [ n ] ∼ iid P , the excess risk is bounded by

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 ∈V M | ̂ L D ( f ) -L D ( f ) | , which cannot be simply controlled due to the unboundedness of the infinity norm of functions in the target function class V M . We dealt with this subtlety by a carefully decomposition of sup f ∈V M | ̂ L D ( f ) -L D ( 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, we call it the random-feature MLP model, in short RFMLP model). In the RFMLP model, we view x 0: N as an input vector instead of a sequence of vectors denoted as vec( x 0: N ) = [ x 0 ; x 1 ; . . . ; x N ; 1] ∈ R d ( N +1)+1 . The RFMLP is given by

We choose the variance of random weights w m to be 1 / ( N +2) to ensure that ⟨ w m , vec( x 0: N ) ⟩ ∼ N (0 , 1) has unit variance. The generalization and approximation properties of the random-feature MLP model have been well-studied in the literature, for example, [7, 8, 60].

We instantiate Theorem 2 on three concrete examples of target functions (calculations of the excess risks in Appendix B.4, where the result for RFMLP are adapted 5 from Arora et al. [7]). In all three cases, the target functions are permutation invariant with respect to { x i } i ∈ [ N ] , by which we naturally expect RFA to achieve better sample complexity than RFMLP in accordance with this structure. Although the comparsion between RFA 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 a special case with a uniform distributional assumption on the input tokens.

Example 1 (Functions of x 0 ) : We consider functions of x 0 (no dependence on x 1: N ) of the form

By Theorem 2, setting M = Θ( d 2 n ) , the excess risk bound gives ˜ O ( √∑ ∞ k =0 k 4 . 5 4 k ∥ A k ∥ 2 Fr /n ) . ✸

̸

As a special case, consider f ⋆ ( x 0: N ) = ( β ⊤ x 0 ) p , which corresponds to taking A k = β ⊗ p for k = p and A k = 0 for k = p . The above excess risk of RFA model and the RFMLP model scales as

Compared to the RFMLP model, the RFA model significantly reduces the necessary sample size by a factor of ( N/ 4) p .

Example 2 (Average of functions of x i ) : We consider average of functions of x i of the form

Theorem 2 then gives an ˜ O ( √∑ ∞ k =0 k 4 . 5 4 k ∥ A k ∥ 2 Fr /n ) excess risk, same as Example 1. ✸

As a specific example, consider f ⋆ = 1 N ∑ N i =1 ψ ( ⟨ β , x i ⟩ ) with ψ ( z ) = z arctan( z/η ) for some η > 2 , ∥ β ∥ 2 = 1 . Using the power series expansion of ψ , the excess risk bound of RFA model and the RFMLP model scale as

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

The latter diverges whenever η ≤ ( N +2) / 2 , in which case the bound is meaningless.

Example 3 (Correlation-weighted functions) : f ⋆ is the following function:

for S ∈ R d × d , { a k } k ≥ 0 ⊆ R , G k ∈ R d k . This target function fully exploits the representation power of the attention layer. Eq. (10) gives B ( f ⋆ ) = O ( ∑ ∞ k =0 C k ( ∑ r + s = k a 2 r ∥ S ∥ 2 r Fr ∥ G s ∥ 2 Fr )) . ✸

As a specific example, consider f 1 ,⋆ = 1 N ∑ N i =1 ⟨ x 0 , x i ⟩ p , corresponding to taking S = I d , F ( t ) = t p , and G ≡ 1 . The excess risk bound of RFA (by Theorem 2) and RFMLP scale as

As another example, consider f 2 ,⋆ = 1 N ∑ N i =1 cos( ⟨ x 0 , x i ⟩ ) ⟨ x ⊗ p i , G ⟩ with ∥ G ∥ Fr = 1 . Then the excess risk bound of RFA and RFMLP scale as

RFA reduces the required sample size by factors of ( N/ 4) p for f 1 ,⋆ and exp( N √ d ) for f 2 ,⋆ .

4 Expressivity of biased random-feature attention model

We now 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 are often more similar to the identity matrix than a mean-zero matrix (Figure 1; see the details in Appendix D.1. This is also observed in 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 W 0 ∈ R ( d +1) × ( d +1) :

Here { W m } m ∈ [ M ] are again Gaussian random matrices sampled according to (4). The biased random-feature attention model is similar to (5) except that a bias weight W 0 is added. Motivated by our observation, we choose W 0 = [ I d × d , 0 d × 1 ; 0 1 × d , 0] ∈ R ( d +1) × ( d +1) , so that the diagonal elements of W 0 + W m will be on average larger than the off-diagonal elements.

4.1 Expressivity of biased random-feature attention

Given the formulation of biased random-feature attention models (thereafter, we call it the biased random-feature attention model, in short BRFA model), a natural conjecture is that this model can better fit functions that are the average of function of ⟨ x 0 , x i ⟩ . We here show that this is indeed the case. In particular, we consider a broad class of target functions g ⋆ : X → R that take forms

Here the scalars { a k } k ≥ 0 ⊆ R and the tensor A ⋆ ∈ R d 5 parameterizes g ⋆ . As we will explain in Section 4.3, confining G to be a degree(3 , 2) polynomial in ( x i , x 0 ) 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 in Theorem 2, when the population risk minimizer gives f ⋆ = g ⋆ , with probability at least 1 -δ , we have

where ε L = 1 / [2 L +1 ( L +1)!] and B ( g ⋆ , L ) = ∥ A ⋆ ∥ 2 Fr · ( ∑ ∞ k =0 | a k | · C k ) 2 , with C k = (2 L + k ) ( k +3) / 2 8 L + k/ 2 . (18)

The proof of Theorem 3 is contained in Appendix C. We provide the intuitions of the result and an overview of 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 x 0 and x i , but are 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 concrete 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 i and x 0 ,

For any η > 0 , if we take n ≥ exp(exp(Θ(1 /η ))) , L = Θ((1 + log log n ) -1 log n ) , and M = Θ( d 2 n ) , the excess risk will scale as ˜ O ( √ ∥ A ∥ 2 Fr /n 1 -η ) . ✸

Compared with the excess risk of the RFA model as detailed in Example 2, the excess risk bound of the BRFA model loses a factor of n -η/ 2 .

Example 5 (Functions of correlations) : Consider a special case of functions of correlations,

For any η > 0 , choosing the same parameters ( n, L, M ) as Example 4, the excess risk bound scales as ˜ O ( √ (log n + p ) ( p +3) 8 p /n 1 -η ) . ✸

Consider the required sample size n ⋆ to reach an accuracy of 0 . 01 . The BRFA model requires n ⋆ = ˜ O ((8 p +48) p +3 ) , whereas the RFA model requires n ⋆ = ˜ O ((4 d ) p ) . Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by a factor of ˜ O ([ d/ (2 p +12)] p ) .

Example 6 (Correlation-weighted functions) : Consider the function

where B ( g ⋆ , L ) is bounded through the Taylor expansion of cos( t ) and (18). For any η > 0 , choosing the same parameters as Example 4, the excess risk bound scales as ˜ O ( √ 1 /n 1 -η ) . ✸

Consider the required sample size n ⋆ to reach an accuracy of 0 . 01 . The BRFA model requires n ⋆ = ˜ O (1) , whereas the RFA model requires n ⋆ = ˜ O (Poly( d ) exp( √ d )) . Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by a factor of ˜ O (Poly( d ) exp( √ d )) .

4.3 Overview of techniques

Here we 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 BRFA ( x 0: N , x ′ 0 ,N ) associated with the infinite-width BRFA model. This kernel has a natural feature map, given by { Ψ k : X → R d 2 k +1 } k ≥ 0 , where

Here ϕ ( t ) = (2 π ) -1 / 2 e -t 2 / 2 is the Gaussian density function, and He k ( z ) denotes the k -th probabilist'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,

whose RKHS norm associated with kernel K BRFA is bounded by B ( ˜ g ⋆ ) = ∑ ∞ k =2 ( k -2)! k 2 4 k ∥ A k ∥ 2 Fr .

Notice that ˜ g ⋆ bears similarities to, but also distinct differences from, g ⋆ as presented in (16). The key difference lies in the ϕ ( ⟨ x 0 , x i ⟩ ) factor in ˜ g ⋆ , which is hard to interpret and analyze. To obtain the excess risk bound for learning g ⋆ , we can use ˜ g ⋆ to approximate g ⋆ in the L ∞ norm. The excess risk for learning g ⋆ can be bounded by the summation of the excess risk for learning ˜ g ⋆ and the approximation error. Acquiring this approximation error bound necessitates a truncation argument of the Taylor expansion of 1 /ϕ ( · ) .

5 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 RFMLP (14). We choose the target functions to be of form

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 set the input dimension as d = 16 and the number of tokens as N = 16 . We fix the width of RFA and BRFA to be M RFA = M BRFA = M = 1000 , whereas the width of RFMLP is set as M RFMLP = 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 0 = 4[ I d × d , 0 d × 1 ; 0 1 × d , 0] ∈ R ( d +1) × ( d +1) , which is four times the matrix investigated in our theory. The input distribution is selected as { x i } 0 ≤ i ≤ N ∼ iid Unif ( S d -1 ) , and we take y = f ⋆ ( x 0: N ) without any noise. We consider three representative target functions: f 1 ,p for p = 2 , 4 , and f 2 ,p for p = 3 , as per (20) and (21). We examine a list of sample sizes n from 2 4 to 2 12 . Prior to training with RF models, we standardize the y i '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.

The left and middle panels of Figure 2 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 our sample complexity analysis detailed in Example 2, where the

10 1

10 2

10 3

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6



f

1, 2

MLP

RFA

BRFA

10 1

10 2

10 3

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

f

1, 4

MLP

RFA

BRFA

10 1

10 2

10 3

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

f

2, 3

MLP

RFA

BRFA







sample complexity bound of RFMLP for learning average of functions of x i is found to be O (( N/ 4) p ) times greater than that of RFA .

The performance comparison between RFA and BRFA depends on the target functions. RFA outperforms BRFA in learning f 1 , 2 and f 1 , 4 , whereas BRFA outperforms RFA in learning f 2 , 3 . The latter phenomenon is as we expected: as demonstrated in Example 3 and 5, BRFA is more powerful than RFA in 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 random-feature attention models, namely RFA (5) and BRFA (15). For general classes of functions that are invariant to the permutation of key tokens x 1: N , the excess risk of RFA (5) can avoid the dependence on sequence length, in contrast to the standard random-feature model RFMLP (14). Moreover, for specific functions that adopt the form of correlation-weighted polynomials (6), the excess risk of BRFA can avoid the polynomial dependence on the dimension. These insights enhance our understanding of the attention mechanism within a simplified context. Finally, our work left open many interesting questions for future work, such as the expressivity of softmax attention, the influence of positional encoding in expressivity, and the expressivity of multi-layer transformers.

Acknowledgment

S. Mei is supported in part by NSF DMS-2210827 and NSF CCF-2315725.

References

E. Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055 , 2023.

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 arXiv:2211.15661 , 2022.

Z. Allen-Zhu and Y. Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems , 32, 2019.

Z. Allen-Zhu, Y. Li, and Y. Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019.

Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

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.

F. Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research , 18(1):629-681, 2017.

F. Bach and L. Chizat. Gradient descent on infinitely wide neural networks: Global convergence and generalization. arXiv preprint arXiv:2110.08084 , 2021.

D. Bahdanau, K. Cho, and Y. Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473 , 2014.

Y. Bai and J. D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:1910.01619 , 2019.

A. R. Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information theory , 39(3):930-945, 1993.

P. L. Bartlett and S. Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research , 3(Nov):463-482, 2002.

S. Bhattamishra, K. Ahuja, and N. Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.11264 , 2020.

S. Bhattamishra, A. Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09286 , 2020.

A. Bietti. Approximation and learning with deep convolutional models: a kernel perspective. stat , 1050:18, 2022.

A. Bietti and J. Mairal. On the inductive bias of neural tangent kernels. Advances in Neural Information Processing Systems , 32, 2019.

A. Bietti, L. Venturi, and J. Bruna. On the sample complexity of learning under geometric stability. Advances in neural information processing systems , 34:18673-18684, 2021.

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-1901, 2020.

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.

A. Caponnetto and E. De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics , 7:331-368, 2007.

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:15084-15097, 2021.

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 in Neural Information Processing Systems , 33:22134-22145, 2020.

L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in neural information processing systems , 32, 2019.

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 as meta optimizers. arXiv preprint arXiv:2212.10559 , 2022.

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.

A. Daniely. Sgd learns the conjugate kernel class of the network. Advances in Neural Information Processing Systems , 30, 2017.

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

L. Dong, S. Xu, and B. Xu. Speech-transformer: a no-recurrence sequence-to-sequence model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP) , pages 5884-5888. IEEE, 2018.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, 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.

S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning , pages 1675-1685. PMLR, 2019.

S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes overparameterized neural networks. arXiv preprint arXiv:1810.02054 , 2018.

E. Dyer and G. Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv:1909.11304 , 2019.

B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases and variable creation in selfattention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

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.

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 , 35:30583-30598, 2022.

R. Ge, J. D. Lee, and T. Ma. Learning one-hidden-layer neural networks with landscape design. arXiv preprint arXiv:1711.00501 , 2017.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems , 33:14820-14830, 2020.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. 2021.

A. Giannou, S. Rajput, J.-y. Sohn, K. Lee, J. D. Lee, and D. Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 , 2023.

M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics , 8:156-171, 2020.

J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

J. Huang and H.-T. Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International conference on machine learning , pages 4542-4551. PMLR, 2020.

K. Huang, Y. Wang, M. Tao, and T. Zhao. Why do deep residual networks generalize better than deep feedforward networks?-a neural tangent kernel perspective. Advances in neural information processing systems , 33:2698-2709, 2020.

A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

S. Jegelka. Theory of graph neural networks: Representation and learning. arXiv preprint arXiv:2204.07697 , 2022.

S. Jelassi, M. Sander, and Y . Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Y. Kim, C. Denton, L. Hoang, and A. M. Rush. Structured attention networks. arXiv preprint arXiv:1702.00887 , 2017.

H. Li, M. Wang, S. Liu, and P.-Y. Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 , 2023.

Y. Li and Y. Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems , 31, 2018.

Y. Li, M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers as algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 , 2023.

Z. Li, R. Wang, D. Yu, S. S. Du, W. Hu, R. Salakhutdinov, and S. Arora. Enhanced convolutional neural tangent kernels. arXiv preprint arXiv:1911.00809 , 2019.

T. Liang and A. Rakhlin. Just interpolate: Kernel 'ridgeless' regression can generalize. 2020.

T. Liang, A. Rakhlin, and X. Zhai. On the risk of minimum-norm interpolants and restricted lower isometry of kernels. arXiv preprint arXiv:1908.10292 , 27, 2019.

B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 , 2022.

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.

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 , 115(33):E7665-E7671, 2018.

S. Mei, T. Misiakiewicz, and A. Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory , pages 3351-3418. PMLR, 2021.

S. Mei, T. Misiakiewicz, and A. Montanari. Generalization error of random feature and kernel methods: hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis , 59:3-84, 2022.

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, Pittsburgh, PA, USA, June 22-25, 2006. Proceedings 19 , pages 154-168. Springer, 2006.

T. Misiakiewicz and S. Mei. Learning with convolution and pooling operations in kernel methods. arXiv preprint arXiv:2111.08308 , 2021.

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.03688 , 2022.

C. 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.11895 , 2022.

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

S. Oymak and M. Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

S. Oymak, A. S. Rawat, M. Soltanolkotabi, and C. Thrampoulidis. On the role of attention in prompt-tuning. arXiv preprint arXiv:2306.03435 , 2023.

A. P. Parikh, O. Täckström, D. Das, and J. Uszkoreit. A decomposable attention model for natural language inference. arXiv preprint arXiv:1606.01933 , 2016.

K. Y. Patarroyo. A digression on hermite polynomials. arXiv preprint arXiv:1901.01648 , 2019.

J. Pérez, J. Marinkovi´ c, and P. Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv:1901.03429 , 2019.

A. Pinkus. Approximation theory of the mlp model in neural networks. Acta numerica , 8: 143-195, 1999.

A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems , 20, 2007.

A. Rahimi and B. Recht. Uniform approximation of functions with random bases. In 2008 46th annual allerton conference on communication, control, and computing , pages 555-561. IEEE, 2008.

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.

S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175 , 2022.

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.

A. Rudi and L. Rosasco. Generalization properties of learning with random features. Advances in neural information processing systems , 30, 2017.

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.

C. Snell, R. Zhong, D. Klein, and J. Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

D. A. Tarzanagh, Y. Li, X. Zhang, and S. Oymak. Max-Margin Token Selection in Attention Mechanism, Sept. 2023. URL http://arxiv.org/abs/2306.13596 . arXiv:2306.13596 [cs, math].

T. Tirer, J. Bruna, and R. Giryes. Kernel-based smoothness analysis of residual networks. In Mathematical and Scientific Machine Learning , pages 921-954. PMLR, 2022.

A. Trockman and J. Z. Kolter. Mimetic initialization of self-attention layers. arXiv preprint arXiv:2305.09828 , 2023.

J. A. Tropp et al. An introduction to matrix concentration inequalities. Foundations and Trends® in Machine Learning , 8(1-2):1-230, 2015.

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.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

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.

C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 2022.

E. Weinan, C. Ma, and L. Wu. Barron spaces and the compositional function spaces for neural network models. arXiv preprint arXiv:1906.08039 , 2019.

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.

G. Weiss, Y. Goldberg, and E. Yahav. Thinking like transformers. In International Conference on Machine Learning , pages 11080-11090. PMLR, 2021.

M. Wortsman, J. Lee, J. Gilmer, and S. Kornblith. Replacing softmax with relu in vision transformers, 2023.

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.

K. Xu, M. Zhang, J. Li, S. S. Du, K.-i. Kawarabayashi, and S. Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009.11848 , 2020.

G. Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14548 , 2020.

S. Yao, B. Peng, C. Papadimitriou, and K. Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115 , 2021.

G. Yehudai and O. Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems , 32, 2019.

C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems , 34:28877-28888, 2021.

C. Yun, S. Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 , 2022.

D. Zou, Y. Cao, D. Zhou, and Q. Gu. Gradient descent optimizes over-parameterized 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 n be the probabilists' Hermite polynomial of degree n :

Here are some basic properties of Hermite polynomials { He n } n ≥ 0 :

Let ( x, y ) ∼ N ( 0 , [1 , ρ ; ρ, 1]) . Then E [He m ( x )He n ( y )] = n ! ρ n 1 m = n for m,n ≥ 0 .

He n ( -x ) = ( -1) n He n ( x );

He n +1 ( x ) = x He n ( x ) -He ′ n ( x ) .

He ′ n ( x ) = n He n -1 ( x ) .

It can be shown that Hermite polynomials are a complete orthogonal basis of the following Hilbert space with ϕ ( x ) = (2 π ) -1 / 2 e -x 2 / 2 :

For two functions f, g : R → R , we define their inner product ⟨ f, g ⟩ as:

For any function f ∈ L 2 ( R , ϕ ( x )) , we can derive its Hermite expansion:

where a n = ⟨ f, He n ⟩ = ∫ ∞ -∞ f ( x )He n ( x ) ϕ ( x )d x . Then for another function g ( x ) = ∑ ∞ n =0 b n n ! He n ( x ) , the inner product of f and g gives:

Here are some formulae for Hermite expansion of certain functions.

Lemma A.1 (See e.g., [38]) . Let f, g ∈ L 2 ( R , ϕ ( x )) . Then, for any unit vectors u, v ∈ R d , we have that

where f ( x ) = ∑ ∞ n =0 a n n ! He n ( x ) and g ( x ) = ∑ ∞ n =0 b n n ! He n ( x ) .

Lemma A.2 (Inverse explicit expression [69]) . Hermite expansion of x n gives:

Lemma A.3. Let σ c ( x ) := σ ( x + c ) be the shifted ReLU function. Then the Hermite expansion of σ c gives:

where ϕ ( x ) , Φ( x ) are the PDF and CDF of standard Gaussian.

Proof of Lemma A.3. Denote a n,i := ∫ ∞ -c ( x + c ) i He n ( x ) ϕ ( x ) dx . It is easy to obtain that

Then using the two different formulae for a n, 0 ,

we obtain that

Then it is easy to prove by induction that

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 a functional σ ( x ; w ) : X × W → R d . Denote µ as a probability measure on W . We define the (infinite-width) random feature model as:

where H W = { v ( w ) : ∫ W v ( w ) ⊤ v ( w ) µ (d w ) < ∞ } is a Hilbert space with norm ∥ v ∥ 2 H W = ∫ W v ( w ) ⊤ v ( w ) µ (d w ) and inner product ⟨ v , u ⟩ H W = ∫ W v ( w ) ⊤ u ( w ) µ (d w ) . Besides, we endow F with a norm ∥·∥ F and the corresponding inner product ⟨· , ·⟩ F defined as:

We further define the corresponding reproducing kernel K : X × X → R , s.t.

which is positive definite. Define the RKHS induced by this kernel as H K with corresponding norm ∥·∥ H K and the inner product ⟨· , ·⟩ H K . Then we have the following proposition according to [61]:

Proposition A.1. Given the above definition of F and H K , we have that ( F , ∥·∥ F ) = ( H K , ∥·∥ H K ) .

More generally [11], for any feature map ϕ : X → H (where H is a Hilbert space) that induces the kernel K , i.e., K ( x, y ) = ⟨ ϕ ( x ) , ϕ ( y ) ⟩ H , we have for any function f that

which shows the equivalence among different feature maps that generate the same kernel.

A.3 Concentration inequalities

Definition A.1 (Sub-Gaussian and Sub-Exponential random variables [87]) . For a random variable X , its sub-gaussian norm, denoted ∥ X ∥ ψ 2 , is defined as

If σ ≡ ∥ X ∥ ψ 2 < ∞ , we say that X is σ -sub-Gaussian.

For a random variable X , its sub-exponential norm, denoted ∥ X ∥ ψ 1 , is defined as

If σ ≡ ∥ X ∥ ψ 1 < ∞ , we say that X is σ -sub-exponential.

Theorem A.1 (Gaussian concentration inequality (e.g., [89])) . Let ( X 1 , . . . , X n ) be a vector of i.i.d. standard Gaussian variables, and let f : R n → R be L-Lipschitz with respect to the Euclidean norm. Then the random variable f ( X ) -E [ f ( X )] is L -sub-Gaussian, and hence

Theorem A.2 (Bounded difference inequality (e.g., [89])) . Consider a function f ( X ) : R n → R . Assume that for any X = ( X 1 , . . . , X n ) and X i, ′ = ( X 1 , . . . , X ′ i , . . . , X n ) , we have difference bound | f ( X ) -f ( X i, ′ ) | ≤ L i . We further assume that the random vector X = ( X 1 , X 2 , . . . , X n ) has independent components. Then

Theorem A.3 (Matrix Bernstein Inequality (e.g., [85])) . Consider a sequence { S k } k ∈ [ n ] of independent random matrices with common dimension d 1 × d 2 . Assume that

Introduce the random matrix

Let v ( Z ) be the matrix variance statistic of the sum

Then we have and for all t ≥ 0 ,

Theorem A.4 (Ledoux-Talagrand contraction inequality (e.g., [89])) . Let { ξ i } i ∈ [ n ] ∼ iid Unif ( {± 1 } ) be independent Rademacher random variables. For any set T ⊂ R n and any family of L -Lipschitz functions { ϕ j } j ∈ [ n ] with ϕ i (0) = 0 , we have

B Proofs for Section 3

Throughout this section, we use the notation ≲ to hide a universal constant C . Also, we use σ ( W , X ) to denote a function of W ∈ R ( d +1) × ( d +1) and X ∈ { X ∈ R ( d +1) × ( d +1) : ∥ X ∥ 2 Fr ≤ 4 } ≡ D that satisfies the following properties:

For any X ∈ D , we have that σ ( W , X ) is L 1 -Lipschitz with respect to W .

For any X ∈ D , the expectation E W [ σ ( W , X )] ≤ L 2 .

For any X ∈ D , we have σ ( 0 , X ) ≤ L 3 .

Here L 1 , L 2 , and L 3 are universal constants. Let { W m } m ∈ [ M ] be sampled from Eq. (4). Let v : R ( d +1) × ( d +1) → R with E W [ ∥ v ( W ) ∥ 2 2 ] ≤ R 2 . Consider a random feature model associated with σ (hereafter, we will refer to it as RF model)

as well as the infinite-width version of the random feature model,

Note that both model RFA and BRFA correspond to special choices of σ . Thus, all lemmas and theorems in this section are applicable to both model RFA and model BRFA .

B.1 Proof of Theorem 1

To prove Theorem 1, we first state two auxilliary lemmas, Lemma B.1 and B.2.

Lemma B.1 (From infinite-width RF model to finite-width RF model) . Consider f v that takes the form as Eq. (26) , with E W [ ∥ v ( W ) ∥ 2 2 ] ≤ R 2 . Let x 0: N ∼ P . Define the ∥·∥ 2 L 2 ( P ) norm by

Then with probability at least 1 -δ , there exists a sequence of vectors { v m } M m =1 ⊆ R d +1 and constant universal C < ∞ that only depends on L 1 , L 2 , and L 3 s.t.

Lemma B.2. Under the setting of Theorem 1, let f ⋆ be a target function of form (6) and (7) . Then there exists an infinite-width RFA model (26) with v : R ( d +1) × ( d +1) → R such that

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

Proof of Theorem 1. For any function that takes form (6) and (7), by Lemma B.2, it admits representation (29) with E W [ ∥ v ( W ) ∥ 2 2 ] ≤ B ( f ⋆ ) . Then by Lemma B.1, since RFA model is a special case of the RF model, there exists { v m } m ∈ [ M ] such that Eq. (27) and Eq. (28) hold with probability larger than 1 -δ . This proves Theorem 1.

B.1.1 Proof of auxiliary lemmas

Proof of Lemma B.1.

Step 1. Proof of Eq. (27) . To prove Eq. (27), we use a truncation argument.

Fix a R W > 0 which we will choose its value later in the proof. Recall that we have W 1 , . . . , W m i.i.d. with W m,i,j ∼ iid N (0 , 1 / 4) . Define v m = v ( W m )1 {∥ W m ∥ Fr ≤ R W / 2 } /M for m = 1 , . . . , M . Consider the truncated infinite-width random feature model f R W v = 1 N ∑ N i =1 E W [ σ ( W , ˜ x 0 ˜ x ⊤ i ) ⟨ v ( W ) , ˜ x i ⟩ 1 {∥ W ∥ Fr ≤ R W / 2 } ] , we have

where ˜ C 1 only depends on L 1 , L 2 , and L 3 . Then using Markov's inequality,

holds with probability at least 1 -δ/ 3 . Next for the difference between f R W v and f v , we have

Here ˜ C 2 is a constant that only depends on L 1 and L 2 . By concentration of functions of Gaussian random vectors (Theorem A.1), ∥ σ ( W , ˜ x 0 ˜ x ⊤ i ) ∥ ψ 2 ≤ L 1 + L 2 for any i . So in the last inequality, we used the bound ( E W [ σ ( W , ˜ x 0 ˜ x ⊤ i ) 4 ]) 1 / 2 by Θ(( L 1 + L 2 ) 2 ) . To bound P ( ∥ W ∥ Fr > R W / 2) , we use concentration of functions of Gaussian random vectors (Theorem A.1) again, and get that

Take C = max ( ˜ C 1 , ˜ C 2 ) . Since E ( ∥ W ∥ Fr ) ≤ ( E ∥ W ∥ 2 Fr ) 1 / 2 ≤ d +1 , by choosing R W = d +1+ C √ log M , the above probability is less than 1 /M 2 . Then

with probability larger than 1 -δ/ 3 . This proves Eq. (27).

Step 2. Proof of Eq. (28) . By Chebyshev's inequality,

Combining with the fact that E [ ∥ v ( W ) ∥ 2 1 {∥ W ∥ Fr ≤ R W / 2 } ] ≤ √ E [ ∥ v ( W ) ∥ 2 2 ] ≤ R , we have

For the second part of (28), Markov inequality gives

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, we have

We first consider a single component in the sum, which is

Let u i,j = 〈 ˜ x 0 ˜ x ⊤ i , ˜ x ′ 0 ( ˜ x ′ j ) ⊤ 〉 / 4 . Let N 2 ( ρ ) denote a bivariate normal distribution with marginals are N (0 , 1) and the correlation is ρ ∈ [ -1 , 1] . Then (30) can be expanded as follows:

Here the coefficients { c ℓ } satisfy

Therefore, the kernel can be expressed as:

Now we reformulate the target function as

where ˜ f rs is a transpose of f rs ⊗ 1 ⊗ (2 ℓ -r -s +1) d +1 with 1 d +1 = (0 , . . . , 0 , 1) , such that 〈 ˜ x ⊗ r 0 ⊗ ˜ x ⊗ s i , f rs 〉 = ⟨ ( ˜ x 0 ⊗ ˜ x i ) ⊗ ℓ ⊗ ˜ x i , ˜ f rs ⟩ for any r ≥ 0 and s ≥ 0 . Then by the feature map equivalence property (23), the RKHS norm of f ⋆ can be bounded as

Thus, using again the property (23) with the original feature map of the random feature model, there exists v : R ( d +1) × ( d +1) → R d +1 such that

Notice that ∥ ˜ f rs ∥ 2 Fr = ∥ f rs ∥ 2 Fr by our construction of ˜ f rs , so that the right-hand-side of the equation above coincides with Eq. (10). This proves Lemma B.2.

B.2 Preliminary proposition for Theorem 2

To prove Theorem 2, 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 lemmas below, we denote X = { x ( j ) 0: N } j ∈ [ n ] and y = { y j } j ∈ [ n ] .

Proposition B.1. Under the setting of Theorem 2. Consider the finite width RF model (25) :

Then with probability at least 1 -δ (w.r.t. W , y , and X ), we have

The main difficulty of the proof of Proposition B.1 comes from that ℓ ( f W M , y ) might be unbounded and that ℓ ( f W M ( x ( j ) 0: N ) , y j ) are not independent across j (since they share the same { W m } m ∈ [ M ] ). So we begin with several lemmas below.

Lemma B.3. Let { ξ j } j ∈ [ n ] be a set of i.i.d. Rademacher random variables. Under the setting of Proposition B.1,

Furthermore, any fixed X and y ,

Lemma B.4. Under the setting of Proposition B.1. With probability at least 1 -δ/ 3 over X , y , and W ,

Lemma B.5. Under the setting of Proposition B.1. With probability at least 1 -δ/ 3 over X and y ,

Lemma B.6. Under the setting of Proposition B.1. With probability at least 1 -δ/ 3 over W ,

The proofs of Lemma B.3, B.4, B.5, and B.6 are contained in section B.2.1. Now assuming they hold, we proceed to prove Proposition B.1.

Proof of Proposition B.1.

Split the left-hand side of inequality (31), we have

with probability at least 1 -δ . Here the last inequality uses Lemma B.4, B.5, and B.6. This proves Proposition B.1.

B.2.1 Proof of auxiliary lemmas

Proof of Lemma B.3.

First using Rademacher contraction inequality, since ℓ (0 , y ) ≤ 1 , we can center it and only pay an extra term 1 / √ n in the Rademacher complexity. Then by the Rademacher contraction property (Theorem A.4), the problem boils down to bounding the Rademacher complexity of f W M , which is

Fix X , y , and W

By matrix Bernstein inequality (Theorem A.3), for any fixed m ,

where

Using the union bound, we have

Therefore, we can bound its expectation with

Now take expectation over X , y , and W . Since n ≥ log( dM ) , we have

Combine this with Eq. (37) and (38), we prove (32).

Fixing any X and y , only taking expectation over W , we get

Combine this with Eq. (37) and (38), we prove (33). This finishes the proof of Lemma B.3.

Proof for Lemma B.4.

Denote X = { x ( j ) 0: N } j ∈ [ n ] and y = { y j } j ∈ [ n ] and denote

Since W m has independent standard Gaussian entries, by Gaussian concentration inequality (Theorem A.1), we have that

For the conditional expectation, by symmetrization, we have

for any X and y , where the last inequality uses Lemma B.3. Combining (39) and (40) and taking ε = 2 √ K 2 log (6 /δ ) /M , we have

for any X and y . Since the right-hand side is irrelevant to X and y , we get (34). This proves Lemma B.4.

Proof of Lemma B.5.

Let

For each i ∈ [ n ] , let { X ′ , y ′ } differs with { X , y } only on i -th data point. We have

Therefore, h ( x 0: N , y ) satisfies the bounded difference property with the parameter { L i } n i =1 uniformly bounded by Θ( K 1 /n ) . By bounded difference inequality (Theorem A.2), there's a constant ˜ C such that

Combining with Lemma B.3, we have

Therefore, by taking ε = 2 K 1 [log (6 /δ ) / ( n ˜ C )] 1 / 2 , we have

with probability at least 1 -δ . This proves Lemma B.5.

Denote

Since W m has independent standard Gaussian entries, there is a constant ˜ C s.t.

For the expectation, we have

where the last inequality is by Lemma B.3. Combining (43) and (44) and taking ε = [ K 2 log(6 /δ ) / ( M ˜ C )] 1 / 2 , we get

with probability at least 1 -δ . This proves Lemma B.6.

B.3 Proof of Theorem 2

Proof. Given any target function f ⋆ , by Lemma B.1, with probability larger than 1 -δ/ 2 over W 1: M , there exists ˜ V = { ˜ v m } M m =1 such that

with

By our choice of K 1 and K 2 , let f W ̂ v ,M = f W M ( · ; ̂ V ) denote the model trained by (11), and let

Then with probability at least 1 -δ over W , X , and y , we have

where from (48) to (49) we use the definition of f W ̂ v ,M and f v W ∗ ,M . From (49) to (50), we bound ≪ Yu notes: fix ≫ L D ( f W ̂ v ,M ) -̂ L D ( f W ̂ v ,M ) and L D ( f v W ∗ ,M ) -̂ L D ( f v W ∗ ,M ) by sup f | L D ( f ) -̂ L D ( f ) | and use the Lipschitzness of ℓ ( 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 1 and K 2 into the equation. This proves Theorem 2.

B.4 Proof of Examples in Section 3

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

Consider the RFMLP model

where vec( x 0: N ) = [ x 0 ; x 1 ; . . . ; x N ; 1] ∈ R dN + d +1 . For target functions that take forms in (6) and (7), define B MLP ( f ⋆ ) = ∑ ∞ k =0 ˜ C k ∑ r + s = k ∥ f rs ∥ 2 Fr with ˜ C k = k 3 . 5 ( N + 2) 2 k . In case where f ⋆ admits multiple representations of the form (7), B MLP ( f ⋆ ) is the infimum of the right-hand-side over all such representations.

Then we consider the empirical risk minimizer over the RFMLP model:

where the constrained class V MLP M gives

Proposition B.2 (The sample complexity of RFMLP ) . Let f ⋆ = arg min f ∈F L D ( f ) be the population risk minimizer within the target function class (54) . Assume M > δ -1 and n > log( dM ) . Take K 1 = C √ B MLP ( f ⋆ ) and K 2 = CB MLP ( f ⋆ ) δ -1 in (57) , with C being a constant. Let ̂ f MLP M = f MLP M ( · ; ̂ v ) be empirical risk minimizer of model RFMLP . Then for any joint distribution P , with probability at least 1 -δ over { w m } m ∈ [ M ] sampled according to (55) and { ( x ( j ) 0: N , y j ) } j ∈ [ n ] ∼ iid P , the excess risk is bounded by

Proof of Proposition B.2. The proof is basically the same as that of RFA model. We only give a sketch of the proof. Firstly, with a few modifications of the proof, we can show that Lemma B.1 also holds for RFMLP , with B ( f ⋆ ) replaced with B MLP ( f ⋆ ) . Lemma B.2 is slightly different, we have the kernel expansion:

Therefore we can rewrite f ⋆ as

where ⟨ ˜ f rs,i , [vec( x ′ 0: N )] ⊗ ℓ ⟩ = ⟨ f rs , x ⊗ r 0 ⊗ x ⊗ s i ⟩ . Thus,

The right-hand-side gives the formulation of B MLP ( f ⋆ ) . Then, a similar version of Proposition B.1 holds for RFMLP model f W M , i.e., with probability at least 1 -δ (w.r.t. W , y , and X ),

Then combining all of the above equations and following the proof in Section B.2, we get (58).

Remark 1. The representation in (59) is not unique. With a more careful choice of the representation of the target function f ⋆ , we can get a better bound for B MLP ( f ⋆ ) , which is given by

B.4.2 Proofs of Examples

Proposition B.3 (Restatement of Example 1) . For functions of x 0 of the form

we have that B ( f ⋆ ) = ∑ ∞ k =0 C k ∥ A k ∥ 2 Fr . Setting M = Θ( d 2 n ) , the excess risk bound gives ˜ O ( √∑ ∞ k =0 k 4 . 5 4 k ∥ A k ∥ 2 Fr /n ) . Moreover, consider f ⋆ ( x 0: N ) = ( β ⊤ x 0 ) p . The above excess risk of RFA model and the RFMLP model scales as

Proof of Proposition B.3. This follows by direct calculation.

Proposition B.4 (Restatement of Example 2) . For f ⋆ = 1 N ∑ N i =1 ψ ( ⟨ β , x i ⟩ ) with ψ ( z ) = z arctan( z/η ) with η > 2 and ∥ β ∥ 2 = 1 . The excess risk bound of RFA model and the RFMLP model are

Proof of Proposition B.4. Use the power series expansion of ψ , we have

Plug it into the formula of B ( f ⋆ ) (10) and B MLP ( f ⋆ ) (61), we get

and B MLP ( f ⋆ ) = Θ( ∑ ∞ k =1 k 4 . 5 [( N +2) / 2 η ] 2 k ) . Therefore, by Theorem 2 and Proposition B.2, we get their excess risk. This proves Proposition B.4.

Proposition B.5 (Restatement of Example 3) . For f 1 ,⋆ = 1 N ∑ N i =1 ⟨ x 0 , x i ⟩ p , the excess risk bound of RFA (by Theorem 2) and RFMLP scale as

For f 2 ,⋆ = 1 N ∑ N i =1 cos( ⟨ x 0 , x i ⟩ ) 〈 x ⊗ p i , G 〉 with ∥ G ∥ Fr = 1 . Then the excess risk bound of RFA and RFMLP scale as

Proof of Proposition B.5. For f 1 ,⋆ , direct calculation gives the value of B ( f 1 ,⋆ ) and the excess risk follows. For f 2 ,⋆ , using the Taylor expansion of cos( z ) , we get

Plug it into the formula of B ( f ⋆ ) (10) and B MLP ( f ⋆ ) (61), we get

Note that for any z > 0 ,

Plug z = 4 d into (62) for B ( f 2 ,⋆ ) gives the excess risk for RFA model. As for B MLP ( f 2 ,⋆ ) , we have

√

Using similar argument, we get that B MLP ( f 2 ,⋆ ) is bounded by Θ(Poly( pdN )4 p exp(2( N +2) d )) . 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 σ ′ ( t ) = exp( t ) . The random feature attention model with exponential activation (we call it ERFA ) is given as follows.

We consider deriving the explicit form of the Kernel K ERFA ( x 0: N , x ′ 0: N ) associated with ERFA model

For any i, j ∈ [ N ] , let u i,j = 〈 ˜ x 0 ˜ x ⊤ i , ˜ x ′ 0 ( ˜ x ′ j ) ⊤ 〉 / 4 . For a single component in (64), we have

Here d ℓ = e/ℓ ! for any ℓ ≥ 0 . Therefore, the kernel can be expressed as:

Thus, for any target function f ⋆ that in the form of (6), we can redefine its complexity measure as

and obtain the corresponding results of Theorem 1 and Theorem 2 by repeating the proof of these two theorems.

C Proofs for Section 4

We consider the empirical risk minimizer over the BRFA model (15),

where the constrained class V M gives (12), copied here for reader's convenience,

Denote G to be the set of all functions in the function class (16), i.e.,

We restate Theorem 3 in Theorem C.1 with detailed assumptions.

Theorem C.1 (Restatement of Theorem 3) . Assume M > δ -1 , n > log( dM ) and let L ∈ Z ≥ 1 . Let g ⋆ = arg min g ∈G L D ( g ) be the population risk minimizer within the target function class G (68) . Take K 1 = C √ B ( g ⋆ , L ) and K 2 = CB ( g ⋆ , L ) δ -1 in (67) , with C being a constant. Let ̂ f W , W 0 M = f W , W 0 M ( · ; ̂ V ) be empirical risk minimizer given by (66) . Then for any joint distribution P , with probability at least 1 -δ over { W m } m ∈ [ M ] sampled according to (4) and { ( x ( j ) 0: N , y j ) } j ∈ [ n ] ∼ iid P , the excess risk is bounded by

where ε L = 1 / [2 L +1 ( L +1)!] and

C.1 Auxiliary results for the proof of Theorem 3

In this section, we give some auxiliary results used in the proof of Theorem 3. The proof will be given in Section C.2. We define the biased transformer with infinite width (informally corresponding to Eq. (15) with M →∞ ), given by

and

where h i := ⟨ W 0 , ˜ x 0 ˜ x ⊤ i ⟩ and W is sampled according to (4). Here we set W 0 = [ I d × d , 0 d × 1 ; 0 1 × d , 0] . Then h i = ⟨ x 0 , x i ⟩ .

We consider a class of target functions ˜ g ⋆ : X → R that takes form

for some coefficients { D ℓ ∈ R ( d +1) 2 ℓ +1 } ℓ ≥ 0 and

where ϕ and Φ are the PDF and CDF of the standard Gaussian random variable, respectively. Lemma C.1 below provides a counterpart of Lemma B.2 for the BRFA model and the target function (71).

Lemma C.1. Any function ˜ g ⋆ of form (71) can be expressed exactly as an infinite-head random feature attention model (70)

where the coefficients v ( · ) satisfy

Given Lemma C.1, we can get a counterpart of Theorem 1 for the BRFA model as Proposition C.1 below.

Proposition C.1 (Counterpart of Theorem 1 for BRFA ) . Suppose function ˜ g ⋆ : X → R takes form (71) . Then for any input distribution P on X , with probability at least 1 -δ (over { W m } m ∈ [ M ] sampled from (4) ), there exists an M -head BRFA model (15) with coefficients V = { v m } m ∈ [ M ] ⊆ R d +1 that approximates ˜ g ⋆ in L 2 ( P ) up to error

In addition, the norms of the weight of this random-feature attention model are bounded as

Here B ( ˜ g ⋆ ) is defined alternatively as

Furthermore, we could approximate the target function in the form (16) to any precision by a function that takes form (72) , which is discussed in Lemma C.2.

Lemma C.2. For any target function g ⋆ in the form of (16) , and for any precision ε ℓ := 1 2 ℓ +1 ( ℓ +1)! , there exists a function f W 0 v in the form of (72) such that

where the tensor A ⋆ ∈ R d 5 parameterizes g ⋆ .

C.1.1 Proof of auxiliary lemmas

Proof of Lemma C.1.

The function (72) can be interpreted as a linear function of the feature map

So the kernel w.r.t. the feature map takes the form that

where h i = ⟨ x 0 , x i ⟩ and h ′ j = 〈 x ′ 0 , x ′ j 〉 . Similar to the proof of Lemma B.2, we also consider a single component of the equation above first, which has the form of

We expand σ ( ⟨ W , ˜ x 0 ˜ x ⊤ i ⟩ + h i ) by Hermite polynomials in the space L 2 ( R , ϕ ) using Lemma A.3,

Therefore, we have

Using Lemma A.1, we obtain the expansion:

where

Then by taking summation with respect to i and j , we derive the expansion of the kernel:

Then for the target function that takes the form in (71), we have the RKHS norm of ˜ g ⋆ bounded by:

by the feature equivalence property (23). Thus, there exists coefficients v ( W ) such that and

Proof of Proposition C.1.

Note that Lemma B.1 is also applicable to the model BRFA . Combining it with Lemma C.1 proves Proposition C.1.

Proof of Lemma C.2.

Firstly, we consider approximating x p using the function class { ϕ ( x )He ℓ ( x ) } n ≥ 0 , which is equivalent to approximating x p ϕ -1 ( x ) using Hermite polynomials. We take p = 0 first and compute the Hermite expansion of the 2 ℓ -th-order term in the Taylor expansion of ϕ -1 ( x ) = e x 2 / 2 , i.e. ψ ℓ ( x ) = ∑ ℓ n =0 x 2 n 2 n n ! . Using Lemma A.2, we have that:

where c ℓ,k = 2 k (2 k )! ∑ ℓ -k m =0 (2 m +2 k )! 2 2 m +2 k ( m + k )! m ! . Similarly, for any p ≥ 0 , define p 0 = ⌊ p 2 ⌋ with r = p -2 p 0 . We have that

where c ℓ,k,p 0 = ∑ ℓ + p 0 -k m =max( p 0 -k, 0) (2 m +2 k + r )! 2 m + k -p 0 ( m + k -p 0 )! · 1 2 m m !(2 k + r )! . Then we can bound c ℓ,k,p 0 as follows:

Next, we give an upper bound of c ℓ,k :

Here we use the inequality that

and we will use it again in the following proof. Since the function

can be written as

where e = diag(0 , 0 , ..., 1) . By Lemma C.1, there exists v p,ℓ ( W ) s.t.

with

Notice that

where ℓ p = ℓ + p 0 , and by the inequality that

we obtain an upper bound:

Finally, we consider the target function (16)

By approximating 1 N ∑ N i =1 a p ⟨ x 0 , x i ⟩ p G ( x 0 , x i ) separately and adding the approximation functions together, we obtain a v ℓ s.t. f W 0 v ℓ ( x 0: N ) = ϕ ( ⟨ x 0 , x i ⟩ ) ψ ℓ ( ⟨ x 0 , x i ⟩ ) g ⋆ ( x 0: N ) and

and ∥ ∥ g ⋆ -f W 0 v ℓ ∥ ∥ ∞ = ∥ ( ϕψ ℓ -1) g ⋆ ∥ ∞ ≤ ∥ ϕψ ℓ -1 ∥ ∞ ∥ g ⋆ ∥ ∞ ≤ e 2 ℓ +1 ( ℓ +1)! ∥ g ⋆ ∥ ∞ = ∥ g ⋆ ∥ ∞ ε ℓ . This finishes the proof of Lemma C.2.

C.2 Proof of Theorem C.1

Proof. For any L > 0 , using Lemma C.2, we can find a function f W 0 v L such that ∥ ∥ g ⋆ -f W 0 v L ∥ ∥ ∞ ≤ ∥ g ⋆ ∥ ∞ ε L with B ( f W 0 v L ) ≤ 8 3 B ( g ⋆ , L ) . Follow the same manner as the proof of Theorem 2. We can get that

Therefore, we have that

Taking infimum over L proves Theorem 3.

C.3 Proof of Examples in Section 4

Proposition C.2 (Restatement of Example 4) . Consider the target function

It has norm bound B ( g ⋆ , L ) = ∥ A ∥ 2 Fr L 3 8 2 L . So for any η > 0 , if we take n ≥ exp(exp(Θ(1 /η ))) , L = Θ((1+loglog n ) -1 log n ) , and M = Θ( d 2 n ) , the excess risk will scale as ˜ O ( √ ∥ A ∥ 2 Fr /n 1 -η ) .

Proof of Proposition C.2. The value of B ( g ⋆ , L ) can be obtained by definition and by direct calculation. As for the second part of the proposition, take L = r log n , where r > 0 is a parameter that is to be chosen to minimize the excess risk. Eq. (17) becomes

where ˜ O hides all the logarithm factors and constants of n, d , and M . To minimize the excess risk, we need to make the two terms in (78) have the same scale. So we set 3 r -1 / 2 = -r log(2 /e ) -r log r -r log log n. Denoting the solution as r ⋆ , we get that

So this gives r ⋆ < (1 / 2 + 1 /e )(log(2 e 2 ) + log log n ) -1 . Assume r ⋆ < 1 (otherwise 3 r ⋆ > 1 / 2 and the excess risk is meaningless), then r ⋆ > (2 log(2 e 2 ) + 2 log log n ) -1 . Therefore, we get that r ⋆ = C (1 + log log n ) -1 , with C being a constant. As a result, when choosing L = r ⋆ log n = Θ((1 + log log n ) -1 log n ) , the excess risk Eq. (78) scales as ˜ O ( ∥ A ∥ Fr n 3 r ⋆ -1 2 ) = ˜ O ( ∥ A ∥ Fr n C/ (1+log log n ) -1 2 ) . As a result, let n > exp(exp( C/η -1)) where C is a constant, we get an excess risk that scales as ˜ O ( √ ∥ A ∥ 2 Fr /n 1 -η ) . This finishes the proof of Proposition C.2.

Proposition C.3 (Restatement of Example 5) . Consider the target function

It has B ( g ⋆ , L ) = (2 L + p ) p +3 8 2 L + p . For any η > 0 , choosing the same parameters ( n, L, M ) as Example 4, the excess risk bound scales as ˜ O ( √ (log n + p ) ( p +3) 8 p /n 1 -η ) .

Furthermore, to reach an accuracy of 0 . 01 , the BRFA model requires n ⋆ = ˜ O ((8 p +48) p +3 ) , whereas the RFA model requires n ⋆ = ˜ O ((4 d ) p ) .

Proof of Proposition C.3. The value of B ( g ⋆ , L ) can be obtained by direct calculation, and we use the same method as the proof of Proposition C.2 to get the ˜ O ( √ 1 /n 1 -η ) excess risk. To reach an accuracy of 0 . 01 , we can set L = 3 , note that ε L ∥ g ⋆ ∥ ∞ < 0 . 006 . Therefore we just need to choose n ⋆ s.t. n > ˜ O ((6 + p ) p +3 8 6+ p ) , and this gives n ⋆ = ˜ O ((8 p +48) p +3 ) for the BRFA model. Direct calculation using Theorem 2 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

It has B ( g ⋆ , L ) = Θ((8 e ) 2 L ) . For any η > 0 , choosing the same parameters as Example 4, the excess risk bound scales as ˜ O ( √ 1 /n 1 -η ) .

Furthermore, to reach an accuracy of 0 . 01 . The BRFA model requires n ⋆ = ˜ O (1) , whereas the RFA model requires n ⋆ = ˜ O (Poly( d ) exp( √ d )) .

Proof of Proposition C.4. We use the expansion of the cos function,

Then use the Eq. (17), we get that

This gives the formula for B ( g ⋆ , L ) . To reach an accuracy of 0.01, we set L = 3 . Then B ( g ⋆ , L ) = Θ(1) . So BRFA needs n ⋆ = ˜ O (1) . And Theorem 2 and Example 3 show that RFA model would require n ⋆ = ˜ O (Poly( d ) exp( √ d )) . 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

The target function (79) is a specific instance of Example 1, and the target function (80) is a specific instance of Example 3. In (80), we sample Z ij ∼ iid N (0 , 1 /d ) for ( i, j ) ∈ [ d ] 2 and vary γ 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 f 3 , 2 (79), which solely depends on x 0 . We observe that a larger sequence length N leads to a 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 in Example 1, where the sample complexity bound of RFMLP for learning average of functions of x i is found to be O (( N/ 4) p ) times greater than that of RFA .

Figure 4 demonstrates the performance comparison between RFA and BRFA under different choices of the token dimension d . 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 n as we state in Example 5. In 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 is consistent with our analyses that BRFA saves a Poly(d) factor in the sample complexity for d ≫ p .

Figure 5 demonstrates that BRFA has no advantage in approximating the target function f 4 ,γ for γ = 0 (left panel; i.e., ( S ij ∼ iid N (0 , 1 /d ) , for ( i, j ) ∈ [ d ] 2 ). However, as γ increases, BRFA outperforms RFA and their separation increases with a larger γ (middle and right panel). Notice that lim γ →∞ ⟨ x 0 , Sx i ⟩ /γ = ⟨ x 0 , x i ⟩ . So this result also conforms to our analysis that BRFA is adept at approximating functions of correlations as 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 in 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 Q ij ∈ R 768 × 64 and the key matrix as K ij ∈ R 768 × 64 . We compute W ij = √ 768 · Q ij K ⊤ ij for i, j from 1 to 12, and then take the absolute value for each entry of the weight matrices W ij . Figure 1

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6



N

=4

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

N

=8

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

N

=16

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6



d

=8

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

d

=16

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

d

=32

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6



=0

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

=2

MLP

RFA

BRFA

10 1

10 2

10 3



0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

=8

MLP

RFA

BRFA

shows the heat maps of weight matrices of the 2nd, 5th, 8th, and 11th layers, where all matrices are clipped to the top-left 32 × 32 block. As we can see, a 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nTypo 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.\n\nQUESTIONS:\nYou 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?\n\nLIMITATIONS:\nNo need", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nWhile 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$\n\nQUESTIONS:\nWould 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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. \"\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThere is no negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}