Year: 2021

Abstract:
Self-attention, an architectural motif designed to model long-range interactions in sequential data, has driven numerous recent breakthroughs in natural language processing and beyond. This work provides a theoretical analysis of the inductive biases of self-attention modules, where our focus is to rigorously establish which functions and long-range dependencies self-attention blocks prefer to represent. Our main result shows that bounded-norm Transformer layers create sparse variables: they can represent sparse functions of the input sequence, with sample complexity scaling only logarithmically with the context length. Furthermore, we propose new experimental protocols to support this analysis and to guide the practice of training Transformers, built around the large body of work on provably learning sparse Boolean functions.

heading: 1 INTRODUCTION
text: Self-attention mechanisms have comprised a dramatic paradigm shift in deep learning in recent years, appearing ubiquitously in recent empirical breakthroughs in sequence modeling and unsupervised representation learning. Starting with large-scale natural language processing (Vaswani et al., 2017), self-attention has enjoyed surprising empirical successes in numerous and diverse modalities of data. In many of these settings, self-attention has supplanted traditional recurrent and convolutional architectures, which are understood to incorporate inductive biases about temporal and translational invariances in the data. Self-attention models discard these functional forms, in favor of directly and globally modeling long-range interactions within the input context.
The proliferation of self-attention raises countless mysteries for theorists and empiricists. One fundamental question concerns its statistical properties: How should we think about the inductive biases of self-attention models? More specifically, we can ask: Which functions do self-attention blocks prefer to represent? How many (approximately) distinct functions can they represent? To this end, this work initiates an analysis of the statistical foundations of self-attention, as it is used in today’s state-of-the-art models.
Our main technical contribution is a classical norm-based generalization bound for a Transformer network, which can be extended to related and future architectures via a modular abstraction of attention mechanisms. In particular, the capacity (in terms of log covering number) of the function class of bounded weight self-attention heads (and Transformers) grows only logarithmically in the context length, which provides theoretical justification for the empirical observation that attention models can learn long-term dependencies without overfitting. Next, we show that bounded-norm self-attention heads are capable of representing wide classes of sparse functions. This representational capacity result, combined with the generalization results, provides a partial theoretical explanation for the observed sparsity bias of attention models, which we term sparse variable creation.
We accompany this analysis with an experimental study of the sample complexity needed by Transformers to learn sparse Boolean conjunctions, and verify the sample complexity scaling law predicted by the theory in this clean synthetic setting. We discuss how to extend and repurpose this experimental protocol of benchmarking long-context sequence models on synthetic “cryptographic” tasks. We find that Transformers trained with gradient descent can learn sparse parities with noise, which may be of independent interest, exposing the empirical study of Transformers to the rich theory established around this problem.

heading: 1.1 RELATED WORK
text: The direct precursors to modern self-attention architectures were recurrent and convolutional networks augmented with attention mechanisms (Bahdanau et al., 2014; Luong et al., 2015; Xu et al., 2015). Landmark work by Vaswani et al. (2017) demonstrated significantly improved machine translation models via self-attention only; autoregressive language models followed shortly (Liu et al., 2018; Radford et al., 2018; 2019; Brown et al., 2020), as well as self-supervised representation learning (Devlin et al., 2018). More recently, self-attention has demonstrated promise in computer vision (Dosovitskiy et al., 2020), protein folding (Jumper et al., 2021), theorem proving (Polu & Sutskever, 2020), program synthesis (Chen et al., 2021b), and reinforcement learning (Chen et al., 2021a; Janner et al., 2021).
Norm-based generalization bounds. There is a vast body of literature dedicated to establishing statistical guarantees for neural networks, including VC-dimension and shattering bounds (dating back to Anthony et al. (1999)). In recent years, generalization bounds have been established for various architectures under norm bounds including (Bartlett et al., 2017; Neyshabur et al., 2015; 2017; Golowich et al., 2018; Long & Sedghi, 2019; Chen et al., 2019) using covering-based arguments; Jiang et al. (2019) provide an extensive empirical study of how well these bounds predict generalization in practice. Our work complements these bounds by establishing norm-based bounds for attention models. Our main results rely on a novel reduction to the `∞ covering number bound for linear function classes given by Zhang (2002).
Theory for attention models. Our work complements various existing theoretical perspectives on attention-based models. Vuckovic et al. (2020) formulate a dynamical system abstraction of attention layers, arriving at similar Lipschitz constant calculations to ours (which are coarser-grained, since they focus on contractivity and stability rather than finite-sample statistical guarantees). Zhang et al. (2019); Snell et al. (2021) study idealizations of optimization problems for self-attention heads. Wei et al. (2021) propose a definition of statistically meaningful approximation of function classes that ties statistical learnability with expressivity, and show that Boolean circuits can be SM-approximated by Transformers with a sample complexity bound that depends mildly on circuit depth (rather than context size), using a margin amplification procedure.1 See Appendix D for a broader survey of related work, including universal function approximation-style results (which ignore statistical considerations).

heading: 2 BACKGROUND AND NOTATION
text: Notation. We use d to denote the embedding dimension for input tokens. T denotes the number of tokens in an input sequence, a.k.a. the context length or context size of an attention mechanism. And m refers to the number of samples (input sequences) in a data set. ‖ · ‖2 denotes the spectral norm for matrices, and ‖ · ‖p,q denotes the (p, q) matrix norm where the p-norm is over columns and q-norm over rows. ‖ · ‖p denotes the `p norm for vectors. For `2-norm, we drop the subscript. B is generally used to quantify bounds on norms of matrices and L for Lipschitz constants. ∆n−1 denotes the simplex of dimension n, that is, ∆n−1 := {x ∈ Rn : x ≥ 0, ‖x‖1 = 1}. Covering numbers and uniform generalization bounds. Our main technical contribution is a generalization bound arising from carefully counting the number of functions representable by a Transformer. This main complexity notion we use is covering number.
We will use the following definition of∞-norm covering number adapted from Zhang (2002): Definition 2.1 (Covering number). For a given class of vector-valued functions F , the covering number N∞(F ; ε; {z(i)}mi=1; ‖ · ‖) is the smallest size of a collection (a cover) C ⊂ F such that ∀f ∈ F ,∃f̂ ∈ C satisfying maxi ‖f(z(i)) − f̂(z(i))‖ ≤ ε. Further, define N∞(F , ε,m, ‖ · ‖) = supz(1)...z(m) N∞(F ; ε; z(1), . . . , z(m), ‖ · ‖).
If F is real-valued (instead of vector-valued), we drop the norm from the notation. Furthermore for functions parameterized by a set of parameters Θ, we exploit the notation to replace F by Θ.
1Quoting the discussion following Theorem 4.1 of (Wei et al., 2021): “The correct norm-based Rademacher complexity bound to use for Transformers is unclear.”
Recall from Zhang (2002) that for the class of linear functions, Flin = {x 7→ w · x : w ∈ Rd, ‖w‖s ≤ BW }, we have the covering number bound of N∞(F ; ε; {x(i)}mi=1) ≤ O ( B2XB 2 W /ε 2 · log(BXBWm/ε ) , where ‖x(i)‖ ≤ BX for i ∈ [m]. Importantly, note that the covering number has a mild dependence on m, only logarithmic; this logarithmic dependence on m will be helpful when we turn our analysis to the capacity of attention mechanisms.
Generalization bounds. This work focuses on providing log-covering number bounds, which determine the generalization error via standard Rademacher complexity and chaining arguments. The following lemma relates these quantities; we refer the reader to Appendix A.1 for more details.
Lemma 2.2 (Generalization bound via covering number). Let D be a distribution over X × R and let ` : R × R be a b-bounded loss function that is L-Lipschitz in its first argument. For a given function class F and f ∈ F , let risk(f ;D) := E(x,y)∼D[`(f(x), y)] and r̂isk ( f ; (z(i), y(i))mi=1 ) := 1m ∑m i=1 `(f(z
(i)), y(i)). Suppose F satisfies |f | ≤ A for all f ∈ F and logN∞(F ; ε;x(1), . . . , x(m)) ≤ CF/ε2 for all for all x(1), . . . , x(m) ∈ Xm. Then for any δ > 0, with probability at least 1− δ, simultaneously for all f ∈ F and some constant c > 0,∣∣∣risk(f ;D)− r̂isk(f ; (x(i), y(i))mi=1)∣∣∣ ≤ 4cL√CFm (1 + log (A√m/CF))+ 2b √ log(1/δ) 2m .

heading: 3 ABSTRACTIONS OF ATTENTION AND SELF-ATTENTION
text: Attention is not straightforward to define — unlike other architectural components such as convolutions and residual connections, it is a broader model design principle, with numerous variations possible. In this section, we present an abstraction of attention mechanisms, guided by the different manifestations discussed in (Luong et al., 2015), with the goal of enabling a more unified and modular statistical analysis. Subsequently, we show how to represent the Transformer (the predominant attention-based architecture) as a special case of this formulation. Our goal is not necessarily an all encompassing formulation of attention mechanisms, but rather an abstraction that helps to guide general design principles and theoretical analysis.

heading: 3.1 ATTENTION
text: Intuitively, we would like to capture the notion that an output variable selects (“attends to”) a part of the input context on which it will depend, based on a learned function of global interactions (see Figure 1, left). To this end, we define an attention head as a function which maps a context X of T inputs {xt ∈ X}Tt=1 (e.g. the tokens in a sentence, pixels in an image, or intermediate activations in a neural network) and a context z ∈ Z to an output y ∈ Y . In this work, we will exclusively consider X ,Y,Z to be Rd; d is the embedding dimension, and we define the matrix of inputs X = [x1x2 . . . xT ]> ∈ RT×d. An attention head uses z to select the inputs in X to which the output y will “attend”, which we formalize below:
Definition 3.1 (Attention head). An attention head is a function f : X → Y , specified by an “alignment score” function Score : X × Z → R parameterized by θs ∈ Θs, normalization function Norm : RT → ∆T−1, and position-wise transformations φin : X → V, φout : V → Y parameter-
ized by θin ∈ Θin and θout ∈ Θout. The output of an attention head on input X ∈ X T , z ∈ Z is
y = φout ( T∑ t=1 [ Norm ( Score(x1, z; θs), . . . ,Score(xT , z; θs) )] t φin(xt; θin); θout ) = φout ( φin(X; θin) >Norm ( Score(x1, z; θs), . . . ,Score(xT , z; θs) ) ; θout
) where φin(X; θ) = [φin(x1; θ) . . . φin(xT ; θ)]> denotes the row-wise application of φin.
The above definition corresponds to leftmost diagram in Figure 1. Here, V is a vector space of input representations “mixed” by the normalized alignment scores; in this work, we will set V = Rk. A function class of attention heads is induced by specifying parameter classes for {Θs,Θin,Θout}.

heading: 3.2 SELF-ATTENTION AND TRANSFORMERS
text: A self-attention head is a special case of an attention head, in which the context z is one of the inputs xt themselves: interactions between elements in X are used to select the elements of X on which f depends. In this case, we will use the term context to denote X . The focus of this work is to analyze the inductive biases of such a construction. For example, for a self-attention head (see Figure 1 (center)), we would have that the t-th component is:
yt = φout ( φin(X; θin) >Norm(Score(X,xt; θs)); θout ) ,
We now define the Transformer self-attention architecture as a special case of the above. Since a Transformer layer has shared parameters between multiple output heads, it will be convenient to define all T outputs of this layer at once. Definition 3.2 (Transformer self-attention layer). A Transformer attention layer is a collection of T attention heads with outputs y1, . . . , yT , specified by the following choices of function classes (with shared parameters between the heads) where the context for yτ is xτ . • Score(x, xτ ; {WQ,WK}) := x>τ WQW>Kx, WQ,WK ∈ Rd×k(for output yt) • φin(x;WV ) := W>V x, WV ∈ Rd×k • φout(x;WC) := W>C σ(x), WC ∈ Rk×d, Lσ-Lipschitz function σ : R → R applied position-
wise, with σ(0) = 0. • Norm(x) := softmax (x) = exp(x)
1> exp(x)
Defining Y := [y1y2 . . . yT ]> ∈ RT×d and [RowSoftmax(M)]t,: := softmax(Mt,:), we have Y = σ ( RowSoftmax ( XWQ(XWK) >)XWV )WC . Functions from the above class of Transformer layers map RT×d to itself, so that instances from this function class can be composed. Although Definition 3.2 only contains the “self-attention” component, it is not merely a simplified idealization of the full Transformer architecture used in practice. We discuss some remaining discrepancies (positional embeddings, layer normalization, parallel heads, position-wise feedforward networks) in Section 4.3 and the appendix.
Extracting scalar outputs from a Transformer. We introduce one more construction: the canonical way to extract a scalar prediction from the final layer of a Transformer. This is the setup used by the classification modules in BERT (Devlin et al., 2018) and all of its derivatives. For a context of size T , a Transformer layer with T + 1 inputs is constructed, with a special input index [CLS]. The input at this position is a vector x[CLS] ∈ Rd (which can be considered as a constant, a part of the input, or a trainable parameter); the output is a linear function w>y[CLS], for a trainable parameter w ∈ Rd. This defines a class of functions mapping RT×d → R, parameterized by a Transformer layer’s parameters and w, which we call the class of scalar-output Transformers.

heading: 4 CAPACITY MEASURES OF ATTENTION MODULES
text: In this section, we present our main technical results, along with overviews of their proofs. Section 4.1 bounds the capacity of a general attention head. Section 4.2 instantiates this bound for the
case of a single Transformer self-attention head. Section 4.3 generalizes this bound for full depth-L Transformer networks. Our sample complexity guarantees scale only logarithmically in the context length T , providing rigorous grounding for the intuition that the architecture’s inductive bias selects sparse functions of the context. Lastly, in Section 4.4, we complement this capacity analysis by exhibiting classes of functions expressible using low-norm Transformer architectures. Combining these representation results and corresponding capacity bounds, we coin the term sparse variable creation to refer to this inductive bias.
Note: Throughout this section, assume that ‖xt‖2 ≤ BX for all t ∈ [T ]. Note that this allows for the Frobenius norm ‖X‖F to scale with √ T . The key challenge throughout our analysis is to avoid incurring factors of norms which take a sum over the t dimension, by analyzing the attention parameters in appropriately chosen geometries.

heading: 4.1 CAPACITY OF A GENERAL ATTENTION HEAD
text: Recall that the attention head architecture can be represented as a function fhead : RT×d×Rd → Rd parameterized by θs, θin, θout as
fhead(X, z; θs, θin, θout) = φout ( φin(X; θin) >Norm(Score(X, z; θs)); θout ) .
Denote the corresponding function class by Fhead := {(X, z) 7→ fhead(X, z; θs, θin, θout) : θs ∈ Θs, θin ∈ Θin, θout ∈ Θout}. To convert the vector-valued function class to a scalar output function class, we define Fscalar := {(X, z) 7→ w>f(X, z) : f ∈ Fhead, w ∈ Rd, ‖w‖ ≤ Bw}. For simplicity, we will focus only on the attention part and assume that φout is a fixed function (no parameters) and w is fixed. It is not hard to handle these even if allowed to be trainable. For the case of Transformers, we handle this more generally (see Appendix A.6).
Assumption 4.1. We make the following assumptions: 1. φout is Lout-Lipschitz in the `2-norm, that is, ∀a, b ∈ Rk, ‖φout(a)− φout(b)‖ ≤ Lout‖a− b‖. 2. φin is Bin-bounded in `2-norm, that is, ‖φin(a; θin)‖ ≤ Bin‖a‖ for all a ∈ Rd and θin ∈ Θin. 3. Norm is continuously differentiable and its Jacobian satisfies ∀θ ∈ RT , ‖J Norm(θ)‖1,1 ≤ CNorm.
The Jacobian assumption might seem strong. However, softmax (the most commonly used Norm function) satisfies this with Csoftmax = 2 (see Corollary A.7).
We prove the following bound on the covering number of Fhead for m samples, Theorem 4.2 (Attention head capacity). Under Assumptions 4.1, the covering number of Fhead satisfies
logN∞ ( Fhead; ε; { (X(i), z(i)) }m i=1 , ‖ · ‖2 )
≤ inf α∈[0,1]
[ logN∞ ( FScore;
αε
CNormLoutBinBX ; {(x(i)t , z(i))}i∈[m],t∈[T ] ) + logN∞ ( Fin;
(1− α)ε Lout
; {x(i)t }i∈[m],t∈[T ]; ‖ · ‖2 )] ,
where FScore = {(x, z) 7→ Score(x, z; θs) : θs ∈ Θs}, and Fin = {x 7→ φin(x; θin) : θin ∈ Θin}.
Note that the bound is in terms of the N∞ covering number of functions that dependent on dimensions d or k and not T . The effect of T only shows up in the number of samples to cover. The N∞ number for many classes scales only logarithmically with the number of samples (for eg, linear functions Zhang (2002)). This is exactly what allows us to get a log T dependence for Transformers.
Since w is fixed, an ε-covering of Fhead directly gives us an εBw-covering for Fscalar implying, logN∞ ( Fscalar; ε; { (X(i), z(i)) }m i=1 ) ≤ logN∞ ( Fhead; ε/Bw; { (X(i), z(i)) }m i=1 , ‖ · ‖2 ) .
Proof overview. In order to prove the bound, we first show a Lipschitzness property of ftf-head. This property allows us to construct the cover by using the covers for FScore and Fin.
Lemma 4.3 (`∞-Lipschitzness of ftf-head). For any θs, θ̂s ∈ Θs, θin, θ̂in ∈ Θin; for all X ∈ RT×d, such that ∥∥X>∥∥ 2,∞ ≤ BX ,∥∥∥fhead(X, z; θs, θin, w)− fhead(X, z; θ̂s, θ̂in, w)∥∥∥
≤ CNormLoutBinBX ∥∥∥Score(X, z; θs)− Score(X, z; θ̂s)∥∥∥
∞ + Lout ∥∥∥φin(X; θin)− φin(X; θ̂in)∥∥∥ 2,∞ .
The most crucial part of this proof is to ensure that we do not get a spurious T dependence when accounting for the attention mechanism. The key observation here is that the attention part of the network is computed using Norm, whose Jacobian norm is bounded. This allows us to use the meanvalue theorem to move to the maximum (`∞) error over T tokens instead of sum (`1), which could potentially incur a T factor. Furthermore, this allows us to combine all samples and tokens and construct a `∞-cover for mT samples.

heading: 4.2 CAPACITY OF A TRANSFORMER SELF-ATTENTION HEAD
text: Let us now look at the case of a Transformer self-attention head and instantiate the covering bound. For ease of presentation and to focus on the self-attention part, we collapse WQW>K to a single matrix (this does not change the representation), set k = d and remove the linear layer WC2. Then the Transformer self-attention head (for any fixed τ ) can be described as
ftf-head(X;WV ,WQK) := σ ( W>V X >softmax ( XW>QKxτ )) which is obtained from the general formulation by setting the context to be xτ , Score(X,xτ ;WQK) = XW > QKxτ , Norm = softmax and φout = σ.
Let us define the function class of self-attention heads with bounded norm, Ftf-head := {X 7→ ftf-head(X;WV ,WQK) : ‖WTV ‖2,∞ ≤ B∞V , ‖WV ‖ ≤ BV , ‖WQK‖2,∞ ≤ B∞QK}. Since WV ,WQK have dimensions dependent on d and k, bounding their norms does not hide a T dependence. As before, to convert this vector-valued function class to a scalar output function class, we define Ftf-scalar := {X 7→ w>f(X) : f ∈ Ftf-head, w ∈ Rd, ‖w‖ ≤ Bw}. We obtain the following bound on the covering number of Ftf-head as a corollary of Theorem 4.2:
Corollary 4.4. For any ε > 0 and X(1), . . . , X(m) ∈ RT×d such that ∥∥∥X(i)>∥∥∥
2,∞ ≤ BX for all
i ∈ [m], the covering number of Ftf-head satisfies
logN∞(Ftf-head; ε;X(1), . . . , X(m), ‖ · ‖2) . (dLσBX)2 ·
( (B∞V ) 2 3 + (B∞QKBV ) 2 3 )3 ε2 · log(mT )
Here . hides logarithmic dependencies on quantities besides m and T .
Our bounds have a logarithmic dependence on T , highlighting the inductive bias of the transformer towards selecting sparse functions of the context.
Proof overview. The above result follows from bounding the covering numbers of FQK := {z 7→ x>τ WQKz : ‖WQK‖2,∞ ≤ B∞QK} and FV := {z →W>V z : ‖WTV ‖2,∞ ≤ B∞V , ‖WV ‖ ≤ BV }.
Note that |x>τ WQKx−x>τ WQKx| ≤ ‖WQKx−WQKx‖ since ‖xτ‖ ≤ 1, so the covering number of FQK is at most the covering number of the class of functions of the form x 7→ WQKx. Therefore, a bound on the vector-valued linear function class suffices to handle both covering numbers. We derive the following covering bound which gives the desired result.
Lemma 4.5. LetW : {W ∈ Rd1×d2 : ‖W‖2,∞ ≤ B∞}, and consider the function class F : {x 7→ Wx : W ∈ W}. For any ε > 0 and x(1), . . . , x(N) ∈ Rd1 satisfying ∀i ∈ [N ], ∥∥x(i)∥∥ ≤ BX , logN∞(F ; ε;x(1), . . . , x(N); ‖ · ‖2) . (d2B∞BX) 2
ε2 · log(N).
2See Appendix 4.3 for a general analysis.
The proof of Lemma 4.5 actually proves a somewhat stronger bound for the function class given by {W ∈ Rd1×d2 : ‖W‖2,∞ ≤ B∞, ‖W‖2,1 ≤ B1}, with d2B∞B1 in the numerator instead of (d2B∞) 2, but we have kept the latter formulation for simplicity of presentation.
Finally, we discuss how to account for some important architectural modifications.
Positional embeddings. In practice, the permutation-invariant symmetry of a Transformer network is broken by adding a positional embedding matrix P ∈ RT×d to the input X at the first layer. In practice, the embedding matrix is often fixed and non-trainable. Our results extend to this setting in a straightforward way; see Appendix A.4. If these matrices are to be trained from a sufficiently large class (say, ∥∥P>∥∥ 2,∞ ≤ 1), the dependence of the log-covering number on T could become linear.
Multi-head self-attention. In almost all applications of Transformers, multiple parallel selfattention heads are used, and their outputs aggregated, to allow for a richer representation. Our analysis directly extends to this setting; see Appendix A.5 for details. When a single attention head is replaced with the sum of H parallel heads, the log-covering number scales up by a factor of poly(H).
Layer normalization. State-of-the-art Transformer networks are trained with layer normalization modules (Ba et al., 2016), which is generally understood to aid optimization. We keep a variant of layer normalization in the covering number analysis– it proves to be useful in the analysis of full attention blocks (see Appendix A.6), as it keeps the norm of the embedding of each token bounded. Removing these layers would lead to a worse dependence on the spectral norm of the matrices.

heading: 4.3 CAPACITY OF DEEP TRANSFORMER NETWORKS
text: In this section, we will extend our results for L-layer Transformer blocks. Denote the weights of layer i by W (i) := { W
(i) Q ,W (i) K ,W (i) V ,W (i) C
} . Further denote the set of weights up to layer i by
W 1:i = (W (1), . . . ,W i−1). Denote the input representation of layer i by g(i)tf-block(X;W 1:i). We inductively define g(i)tf-block : RT×d → RT×d starting with g (1) tf-block(X;W 1:1) = X (the input):
g (i+1) tf-block
( X;W 1:i+1 ) := Πnorm ( σ ( Πnorm ( f ( g (i) tf-block ( X;W 1:i ) ;W (i) ))) W (i) C ) with f (Z; {WQ,WK ,WV , ·}) := RowSoftmax ( ZWQ (ZWK) > ) ZWV ,
where Πnorm denotes layer normalization3 applied to each row. We use a slightly modified version of LayerNorm where instead of normalizing to norm 1, we project it to the unit ball. Let us denote the class of depth-L transformer blocks by
F (L)tf-block := { X → g(L+1)tf-block(X;W 1:L+1) : ∀ i ∈ [L], ∥∥∥W (i)V ∥∥∥
2 , ∥∥∥∥W (i)K W (i)Q >∥∥∥∥ 2 , ∥∥∥W (i)C ∥∥∥
2 ≤ C2,∥∥∥W (i)V ∥∥∥
2,∞ , ∥∥∥∥W (i)K W (i)Q >∥∥∥∥ 2,∞ , ∥∥∥W (i)C ∥∥∥ 2,∞ ≤ C∞ } .
To obtain a final scalar output, we use a linear function of the [CLS] output, gtf-scalar(X;W 1:L+1, w) = w> [ g ( X;W 1:L+1 )] [CLS],: . Let the scalar output function class be F (L)tf-scalar := {X → w>f(X)[CLS] : f ∈ F (L) tf-block, w ∈ Rd, ‖w‖ ≤ Bw}.
Theorem 4.6 (Theorem A.17 (Simplified)). Suppose ∀i ∈ [m], ∥∥X(i)∥∥
2,∞ ≤ BX , then we have
logN∞(F (L)tf-block; ε;X (1), . . . , X(m)) . (C2Lσ)
O(L) · d 2B2XB 2 wC 2 ∞
ε2 · log(mT ).
Note that the dependence on T is only logarithmic even for deeper networks. The dependence on embedding dimension and (2,∞) norms of the weight matrices is quadratic. As long as the spectral norms of the matrices are bounded by 1 and σ is 1-Lipschitz (which holds for sigmoids and ReLUs), the exponential dependence on L can be avoided.
3Layer normalization allows for the norms of the outputs of each token in each layer to remain bounded by 1. Note that the norm of the entire input can still have a dependence on T . Our results would go through with a worse dependence on the spectral norms if we were to remove layer norm.

heading: 4.4 SPARSE VARIABLE CREATION: AN INDUCTIVE BIAS FOR SELF-ATTENTION
text: The above analysis shows that function classes bottlenecked by self-attention mechanisms are “small” in terms of the context size. In this section, we answer the converse question: which functions of interest can they express? To this end, we show in this section that sparse Boolean functions are realizable by bounded-norm Transformers.
Given a Boolean function f : {0, 1}T → R which only depends on s of its inputs, we represent f using a self-attention head ftf-head composed with a feedforward network fmlp; this is the repeated block in the standard Transformer architecture. Intuitively, ftf-head can select an s-dimensional subset of inputs to “attend to”, while fmlp memorizes an arbitrary function of these s inputs (requiring up to≈ 2s parameters). In the regime of s log T (think of ftf-head ◦fmlp as implementing a single composable Boolean gate), the corresponding statistical guarantees are meaningful. In order to describe this combination of sample-efficient sparsification of rich contexts and subsequent restricted use of universal function approximation, we coin the term sparse variable creation.
Setup. We consider the classes of Boolean functions f : {0, 1}T → R representable by boundednorm scalar-output Transformer heads ftf-scalar : RT×d → R. To do this, we must first fix a mapping from {0, 1}T to RT×d; we discuss several natural choices in Appendix B.1. The simplest of these uses a sum of token and positional embeddings X(b)t,: := ebt + vt, for a set of approximately orthogonal unit vectors {e0, e1, v1, . . . , vT }. After choosing a mappingX(b), the setup of the representation problem is evident: given f(b), find Transformer weights θtf-head and feedforward network weights θmlp such that
ftf+mlp(X(b); θtf-head, θmlp) := fmlp (ftf-head(X(b); θtf-head); θmlp) ≈ f(b), ∀b ∈ {0, 1}T .
Main representational results. We show that Transformer blocks can represent I-sparse Boolean functions, whose values only depend on some subset of indices I ⊆ [T ]. We present informal statements of these approximation results below, and present the precise statements in Appendix B.2.
Proposition 4.7 (Sparse variable creation via Transformers; informal). Under any of the input mappings X(b), we have the following guarantees: • ftf-scalar alone can approximate a particular monotone symmetric s-sparse Boolean function, with
weight norms ‖WQ‖F ≤ O (log(Ts)) , ‖WK‖F , ‖WV ‖F , ‖WC‖F ≤ O(s). • ftf+mlp can exactly represent symmetric s-sparse functions, with the same Transformer weight
norms as above; the feedforward network weights satisfy ‖W1‖F , ‖W2‖F , ‖w‖F ≤ O(poly(s)). • ftf+mlp can exactly represent general s-sparse functions, with the same Transformer weight norms
as above; the feedforward network weights satisfy ‖W1‖F , ‖W2‖F , ‖w‖F ≤ O(2s · poly(s)).
Proof ideas. Each construction uses the same idea: select WQ,WK so that the attention mixture weights approximate the uniform distribution over the relevant positions, then use the ReLU network to memorize all distinct values of f . Full proofs are given in Appendix B.4.

heading: 5 EMPIRICAL SCALING LAWS FOR LEARNING BOOLEAN GATES
text: Our theoretical analysis has shown that Transformers can represent sparse Boolean functions, with sample complexity scaling mildly with the context size. In this section, we present an empirical study of whether Transformer architectures (as they are trained and used in state-of-the-art language modeling) exhibit these scalings in practice.
We introduce a rigorous benchmark for probing the empirical sample complexity of a Transformer: attribute-efficient learning of a planted sparse Boolean function. We choose a family of distinct distributions {D1, . . . ,DN} on {0, 1}T × {0, 1}, corresponding to supervised learning problems, such that the feature distributions are identical, and the uniform mixture 1N ∑N i=1Di is invariant under all permutations of the indices 1, . . . , T . We then select an i∗ ∈ [N ] uniformly at random, train a Transformer binary classifier on m samples from Di∗ , then evaluate the generalization error via cross-validation. Since the architecture, initialization, training algorithm, and training data are permutation-invariant, at least Ω(logN) samples are required to learn this distribution: one sample can only reveal a single bit of information about i∗, via its binary label. We are interested in the empirical scaling of the sufficient sample size m to solve this problem, in terms of N .
Learning sparse AND gates. The simplest instantiation of this experimental framework is a hidden s-way AND under the uniform distribution (i.e. T i.i.d. random bits). The Transformer must learn which of the N = ( T s ) combinations of inputs determines the label, which requires m ≥ Ω(s log T ) samples. The theory predicts the sample complexity of learning a bounded-norm Transformer should match this scaling in T . With hyperparameters typical of Transformer setups used for natural data, we indeed observe that the empirical sample complexity scales sublinearly with T . Figure 2 summarizes our findings; details are provided in Appendix C.1.
Learning sparse parities. We can also replace the sparse AND operation with XOR: the label is the parity of a hidden subset of input bits. This variant emphasizes the “cryptographic” nature of this experimental setup, due to its known computational hardness. Ω(T s) statistical queries (thus, batch gradient descent steps) are necessary (Kearns, 1998); in the presence of noise, the fastest known algorithms for learning parities with noise require TΩ(s) time (Valiant, 2012). Figure 3 (with details in Appendix C.2) show that Transformer models can fit sparse parities. This raises an intriguing question: if theory suggests that “exhaustive search-like” methods are computationally necessary for this problem, why does local search (i.e. gradient-based training) succeed? We leave this computational (as opposed to statistical) mystery as an open direction for future work.

heading: 6 CONCLUSION AND FUTURE WORK
text: We have presented a theoretical analysis of the inductive biases of self-attention models, finding that they can learn sparse Boolean functions with sample complexity scaling logarithmically in the context length. We call this phenomenon sparse variable creation. Our analysis is accompanied by new empirical probes involving training Transformers on sparse Boolean functions, where we corroborate this scaling of the sample complexity in practice. We believe our capacity bounds are improvable (we have only sought to obtain an optimal dependence on T ). Incorporating aspects of computation and depth remains a perennial challenge in this line of inquiry.
Building further upon the principles of attention constitutes a vibrant frontier of empirical research (Tolstikhin et al., 2021; Lee-Thorp et al., 2021; Jaegle et al., 2021b;a; d’Ascoli et al., 2021). We hope that the theoretical foundations and experimental probes presented in this paper will assist in further expanding the breadth of applications of self-attention, and developing more computeefficient, data-efficient, controllable, and reliable algorithmic interventions.

heading: A PROOFS OF CAPACITY BOUNDS
text: In this section we present the full proofs (including the omitted proofs) of our capacity bounds. We also cover relevant background and useful technical lemmas.

heading: A.1 RADEMACHER COMPLEXITY AND GENERALIZATION BOUNDS
text: Here we briefly review Rademacher complexity and its relationship to covering numbers and generalization bounds. We refer the reader to Bartlett & Mendelson (2002) for a more detailed exposition. Definition A.1 (Empirical Rademacher complexity). For a given class of functions F = {f : X → R} and {z(i) ∈ X}mi=1, the empirical Rademacher complexity R̂(F ; z(1), . . . , z(m)) is defined as
R̂(F ; z(1), . . . , z(m)) = 1 m Eε [ sup f∈F m∑ i=1 εif(z (i)) ] ,
where ε is a vector of m i.i.d. Rademacher random variables (Pr[εi = 1] = Pr[εi = −1] = 1/2).
In order to relate the Rademacher complexity and `∞-covering numbers, we use a modified version of Dudley’s metric entropy. Lemma A.2 (Dudley (1967); modified). Consider a real-valued function class F such that |f | ≤ A for all f ∈ F . Then
R̂(F ; z(1), . . . , z(m)) ≤ c · inf δ≥0
( δ + ∫ A δ √ logN∞(F ; ε; z(1), . . . , z(m)) m dε ) for some constant c > 0.
Proof sketch. The original statement is for 2-norm covering number, but the∞-norm case reduces to the 2-norm case because N2(·) ≤ N∞(·). The original statement also fixes δ = 0 rather than taking an infimum. Also, the standard statement has the integral go from 0 to∞, but these are easily replaced with δ and A.
For our paper, we will instantiate the above lemma for log covering numbers scaling as 1/ε2. Corollary A.3 (Rademacher Complexity via covering number). Consider a real-valued function class F such that |f | ≤ A for all f ∈ F . Suppose logN∞(F ; ε; z(1), . . . , z(m)) ≤ CF/ε2, then
R̂(F ; z(1), . . . , z(m)) ≤ c · √ CF m · ( 1 + log ( A √ m/CF )) for some constant c > 0.
Proof. Using Lemma A.2, we have for some constant c > 0,
R̂(F ; z(1), . . . , z(m)) ≤ c inf δ≥0
( δ + ∫ A δ √ logN∞(F ; ε; z(1), . . . , z(m)) m dε )
≤ c inf δ≥0
( δ + ∫ A δ √ CF ε2m dε )
= c inf δ≥0
( δ + √ CF m ∫ A δ 1 ε dε )
= c inf δ≥0
( δ + √ CF m log(A/δ) )
= c √ CF m ( 1 + log ( A √ m/CF )) .
We can now obtain a generalization guarantee from the Rademacher complexity of a function class: Theorem A.4 (Bartlett & Mendelson (2002)). Let D be a distribution over X × R and let ` : R × R be a b-bounded loss function that is L-Lipschitz in its first argument. For a given function class F and f ∈ F , let risk(f ;D) := E(x,y)∼D[`(f(x), y)] and r̂isk ( f ; (z(i), y(i))mi=1 ) := 1 m ∑m i=1 `(f(z
(i)), y(i)). Then for any δ > 0, with probability at least 1− δ, simultaneously for all f ∈ F ,∣∣∣risk(f ;D)− r̂isk(f ; (z(i), y(i))mi=1)∣∣∣ ≤ 4L R̂(F ; z(1), . . . , z(m))+ 2b √ log(1/δ)
2m .
Combining the above, we get: Lemma A.5 (Lemma 2.2 (restated)). Consider a function class F such that |f | ≤ A for all f ∈ F and logN∞(F ; ε;x(1), . . . , x(m)) ≤ CF/ε2 for all x(1), . . . , x(m) ∈ Xm. Then for any δ > 0, with probability at least 1− δ, simultaneously for all f ∈ F ,∣∣∣risk(f ;D)− r̂isk(f ; (x(i), y(i))mi=1)∣∣∣ ≤ 4cL√CFm (1 + log (A√m/CF))+ 2b √ log(1/δ) 2m ,
for some constant c > 0.

heading: A.2 USEFUL LEMMAS
text: Lemma A.6. Consider function f : Rd → ∆d−1 such that the Jacobian of the function satisfies ‖J f(θ)‖1,1 ≤ cf for all θ ∈ Rd, then for any vectors θ1, θ2 ∈ Rp,
‖f(θ1)− f(θ2)‖1 ≤ cf‖θ1 − θ2‖∞.
Proof. By the fundamental theorem of calculus applied to g(t) = f(tθ1 + (1 − t)θ2), followed by a change of variables:
f(θ1)− f(θ2) = (∫ 1
0
J (tθ1 + (1− t)θ2) dt ) (θ1 − θ2),
We have
‖f(θ1)− f(θ2)‖1 = ∥∥∥∥∫ 1
0
J (tθ1 + (1− t)θ2) (θ1 − θ2)dt ∥∥∥∥
1
By Jensen’s inequality:
≤ ∫ 1
0
‖J (tθ1 + (1− t)θ2) (θ1 − θ2)‖1 dt
Using ‖Ax‖1 ≤ ‖A‖1,1 ‖x‖∞:
≤ ∫ 1
0
‖J (tθ1 + (1− t)θ2)‖1,1 ‖θ1 − θ2‖∞ dt
By assumption on the Jacobian: ≤ cf ‖θ1 − θ2‖∞ .
Corollary A.7. For vectors θ1, θ2 ∈ Rp, ‖softmax(θ1)− softmax(θ2)‖1 ≤ 2‖θ1 − θ2‖∞.
Proof. Observe that for softmax, the Jacobian satisfies:
J(θ) = diag(softmax(θ))− softmax(θ)softmax(θ)>).
We have for all θ, h,
‖J(θ)‖1,1 = p∑ i=1 p∑ j=1 |softmax(θ)i(1[i = j]− softmax(θ)j)|
= p∑ i=1 softmax(θ)i 1− softmax(θ)i +∑ j 6=i softmax(θ)j  = 2
p∑ i=1 softmax(θ)i (1− softmax(θ)i)
≤ 2.
Combining the above with Lemma A.6 gives the desired result.
Lemma A.8. For αi, βi ≥ 0, the solution to the following optimization
min x1,...,xn n∑ i=1 αi x2i
subject to n∑ i=1 βixi = C
is γ 3
C2 and is achieved at xi = C γ ( αi βi )1/3 where γ = ∑n i=1 α 1/3 i β 2 3 i .
Proof. The proof follows by a standard Lagrangian analysis.
Lemma A.9 (Contraction of Πnorm). Let Πnorm be the projection operator onto the unit norm ball. For any vectors u, v, we have ‖Πnorm(u)−Πnorm(v)‖ ≤ ‖u− v‖.
Proof. If u, v are both in the unit ball then this follows trivially. Let us assume that ‖u‖ ≥ ‖v‖ and ‖u‖ ≥ 1 WLOG. First suppose ‖v‖ ≤ 1. Let B(1)V = αu be the projection of v in the direction of u, and let B2V = v −B (1) V . Then
‖Πnorm(u)−Πnorm(v)‖2 = ‖u/‖u‖ − v‖2
= ‖u/‖u‖ − (αu+B2V )‖2 = ‖(‖u‖−1 − α)u−B2V ‖2 = (‖u‖−1 − α)2‖u‖2 + ‖B2V ‖2 ≤ (1− α2)‖u‖2 + ‖B2V ‖2 since ‖u‖−1 < α < 1 = ‖u− (αu+B2V )‖2
= ‖u− v‖2
If ‖v‖ > 1, then
‖Πnorm(u)−Πnorm(v)‖ = ‖Πnorm(u/‖v‖)−Πnorm(v/‖v‖)‖ ≤ ‖u/‖v‖ − v/‖v‖‖ < ‖u− v‖.
where the second-to-last inequality follows from the ‖v‖ < 1 case.
Lemma A.10 (Zhang (2002) Theorem 4). Let V : {v : v ∈ Rd1 , ‖v‖ ≤ B1} and Flinear = {x 7→ v>x : v ∈ V}. For any δ > 0 and x(1), . . . , x(N) satisfying ‖x(i)‖ ≤ B2 ∀i,
logN∞(Flinear; ε;x(1), · · · , x(N)) ≤ 36 B21B 2 2
ε2 log(2d4B1B2/ε+ 2eN + 1).

heading: A.3 OMITTED PROOFS
text: Proof of Lemma 4.3. Observe that,∥∥∥fhead(X, z; θs, θin)− fhead(X, z; θ̂s, θ̂in))∥∥∥ = ∥∥∥φout (φin(X; θin)>Norm(Score(X, z; θs)))− φout (φin(X; θ̂in)>Norm(Score(X, z; θ̂s)))∥∥∥
By Lout-Lipschitzness of φout and bound on ‖w‖: ≤ Lout ∥∥∥φin(X; θin)>Norm(Score(X, z; θs))− φin(X; θ̂in)>Norm(Score(X, z; θ̂s))∥∥∥
By triangle inequality: ≤ Lout ∥∥∥φin(X; θin)> (Norm(Score(X, z; θs))− Norm(Score(X, z; θ̂s)))∥∥∥
+ Lout ∥∥∥∥(φin(X; θin)− φin(X; θ̂in))> Norm(Score(X, z; θ̂s))∥∥∥∥ Using ‖Pv‖ ≤ ‖P‖2,∞‖v‖1 and Bin-boundedness of φin:
≤ LoutBin ∥∥∥Norm(Score(X, z; θs))− Norm(Score(X, z; θ̂s))∥∥∥
1
+ Lout ∥∥∥∥(φin(X; θin)− φin(X; θ̂in))>∥∥∥∥ 2,∞ ∥∥∥Norm(Score(X, z; θ̂s))∥∥∥ 1
By Lemma A.6 and the assumption on Norm: ≤ LoutCNorm ∥∥φin(X; θin)>∥∥2,∞ ∥∥∥Score(X, z; θs)− Score(X, z; θ̂s)∥∥∥∞ + Lout ∥∥∥∥(φin(X; θin)− φin(X; θ̂in))>∥∥∥∥ 2,∞
By boundedness of φin and ∥∥X>∥∥
2,∞ ≤ BX : ≤ LoutCNormBinBX ∥∥∥Score(X, z; θs)− Score(X, z; θ̂s)∥∥∥
∞ + Lout ∥∥∥∥(φin(X; θin)− φin(X; θ̂in))>∥∥∥∥ 2,∞ .
Proof of Theorem 4.2. Our goal is to show that for every ε > 0, collection of inputs (X(1), z(1)), . . . , (X(m), z(m)), there is a cover Chead such that for all θs ∈ Θs, θin ∈ Θin, there is some (θ̂s, θ̂in) ∈ Chead such that maxi ∥∥∥fhead(X(i), z(i); θs, θin)− fhead(X(i), z(i); θ̂s, θ̂in)∥∥∥ ≤ ε. Observe that for all θs, θ̂s,
max i∈[m] ‖Score(X(i), z(i); θs)−Score(X(i), z(i); θ̂s)‖∞ = max i∈[m],t∈[T ] ∣∣∣Score(x(i)t , z(i); θs)− Score(x(i)t , z(i); θ̂s)∣∣∣ . Similarly, for all θin, θ̂in,
max i∈[m] ∥∥∥∥(φin(X(i); θin)− φin(X(i); θ̂in))>∥∥∥∥ 2,∞ = max i∈[m],t∈[T ] ∥∥∥φin(x(i)t ; θin)− φin(x(i)t ; θ̂in)∥∥∥ . This crucially allows us to aggregate over the i and t dimensions together.4 Therefore, we can consider N∞ covers for the above to bound the overall covering number.
Let CScore be the εScore-cover (∞) for FScore over inputs { (x (i) t , z (i)) } i∈[m],t∈[T ] of size
N∞ ( FScore; εScore; {(x(i)t , z(i))}i∈[m],t∈[T ] ) .
Also, Let Cin be the εin-cover (∞) for Fin over inputs {x(i)t }i∈[m],t∈[T ] of size N∞ ( Fin; εin; {x(i)t }i∈[m],t∈[T ]; ‖ · ‖2 ) .
We are ready to construct the cover forFhead. Set Chead = {fhead(·; θ̂s, θ̂in))i∈[m] : θ̂s ∈ CScore, θ̂in ∈ Cin}. Then for any θs ∈ Θs, θin ∈ Θin, there exists θ̂s, θ̂in ∈ Chead, such that for all i ∈ [m], using Lemma 4.3:∥∥∥fhead(X(i), z(i); θs, θin)− fhead(X(i), z(i); θ̂s, θ̂in)∥∥∥ ≤ CNormLoutBinBXεScore + Loutεin.
4In the case of the Transformer self-attention mechanism, we will obtain ∞-norm covering numbers for Score and φin that have only logarithmic dependence on the number of examples. Because of this aggregation trick, the resulting covering number for the whole layer will have merely logarithmic dependence on the context length T .
The size of the cover we have constructed is,
log |Chead| = log |CScore|+ log |Cin| = logN∞ ( FScore; εScore; {(x(i)t , z(i))}i∈[m],t∈[T ] ) + logN∞ ( Fin; εin; {x(i)t }i∈[m],t∈[T ]; ‖ · ‖2 ) and we are done.
Proof of Corollary 4.4. By Theorem 4.2, the covering number of Ftf-head satisfies logN∞ ( Ftf-head; ε; { (X(i), z(i)) }m i=1 ) ≤ inf α∈[0,1] [ logN∞ ( FQK ; αε 2LσBVBX ; {(x(i)t , z(i))}i∈[m],t∈[T ]
) + logN∞ ( FV ;
(1− α)ε Lσ
; {x(i)t }i∈[m],t∈[T ]; ‖ · ‖2 )] .
where we have used the fact that for a scalar-output Transformer layer:
• softmax satisfies the Jacobian assumption with Csoftmax = 2 using Corollary A.7.
• Lout is the Lipschitz constant of σ: Lσ .
• Bin is a bound on the norm of W>V x with respect to norm of x: BV .
By Lemma 4.5, for any εQK , εV > 0:
logN∞ ( FQK ; εQK ; {(x(i)t , z(i))}i∈[m],t∈[T ] ) . (dBQK2,infBX) 2 log(mT )
ε2QK
logN∞ ( FV ; εV ; {(x(i)t , z(i))}i∈[m],t∈[T ]; ‖ · ‖2 ) . (dBV2,infBX) 2 log(mT )
ε2V
since WQK ,WV ∈ Rd×d (k = d). We want to choose εQK and εV to minimize the sum of the above two terms, subject to 2LσBVBXεQK + LσεV ≤ ε. By Lemma A.8, the solution to this optimization leads to an optimal bound of:
logN∞(Ftf-head; ε;X(1), . . . , X(M)) . (dLσBX)2 ·
( (B∞V ) 2 3 + (B∞QKBVBX) 2 3 )3 ε2 · log(mT ).
Proof of Lemma 4.5. LetB∞ be an upper bound on ‖W‖2,∞ andB1 be an upper bound on ‖W‖2,1. The approach will be to cover each of the columns of W independently, treating each as specifying a linear function from Rd1 → R. By Lemma A.10, letting V(b) : {v : v ∈ Rd1 , ‖v‖ ≤ b} and Flinear(b) = {x 7→ v>x : v ∈ V(b)}, for any δ > 0
logN∞(Flinear(b); δ;x(1), · · · , x(N)) ≤ cb2B2X log((1 + bBX/δ)N)
δ2 .
given that ‖x(i)‖ ≤ BX for all i.
In fact the cover, which we denote by F̂linear(b; δ), is proper: F̂linear(b; δ) = {x 7→ v̂>x : v̂ ∈ V̂ } for some finite subset V̂ ⊂ V(b). Let
S = { (k1, k2, . . . , kd2) : ki ∈ {0, 1, . . . , d2} for all i and 0 ≤ d2∑ i=1 ki ≤ d2 }
Given any W ∈ W , let vi denote the ith column of W . For each i, let
k∗i = ⌊ d2 B∞ ‖vi‖ ⌋ Let SW = (k∗1 , . . . , k ∗ d2 ). Then SW ∈ S, and
B∞ d2 k∗i ≤ ‖vi‖ ≤ B∞ d2 (k∗i + 1) ∀i.
For every tuple S = (k1, k2, . . . , kd2) ∈ S, let ŴS = { [v̂1v̂2 . . . v̂d2 ] > : v̂i ∈ F̂linear ( B∞(ki + 1)/d2; ε √ (ki + 1)/(2d2) ) for all i } .
Our cover will be F̂ : {x 7→ Ŵx : Ŵ ∈ ⋃ S∈S ŴS}. Note that
|F̂ | ≤ ∑ S∈S d2∏ i=1 N∞ ( Flinear(B∞(ki + 1)/d2); ε √ (ki + 1)/(2d2);x (1), · · · , x(N) )
. ∑ S∈S exp ( d2∑ i=1 B2∞a 2(ki + 1) 2 log(N) d22ε 2(ki + 1)/(2d2) )
= ∑ S∈S exp
( 2B2∞B 2 X log(N)
d2ε2
d2∑ i=1 (ki + 1)
)
≤ |S| exp ( 4B1B∞B 2 X log(N)
ε2 ) where in the last step we used the fact that
d2∑ i=1 ki ≤ d2 B∞ d∑ i=1 ‖vi‖ = d2B1/B∞
Since |S| ≤ (
2d2 d2
) = exp(O(d2 log d2)), we obtain
log |F̂ | . d2B1B∞B 2 X log(N)
ε2
as desired. In particular, we obtain the more concise, but looser, bound from the lemma statement by using the fact that B1 ≤ d2B∞
For a particular W = [v1 . . . vd2 ] ∈ W it is guaranteed that there is a matrix Ŵ ∈ ŴSW , Ŵ = [v̂1 . . . v̂d2 ], such that for each i ∈ [d2],
|v>i xn − v̂>i xn| ≤ ε √ ki + 1
2d2 ∀n ∈ [N ].
where ki is the ith element of SW . We then obtain the desired covering property:
max n∈[N ] ∥∥∥Wxn − Ŵxn∥∥∥ = max n∈[N ] √√√√ d2∑ i=1 (v>i xn − v̂>i xn)2
≤ √√√√ d2∑ i=1 ε2 ( ki + 1 2d2 )
= ε
√ d2 + ∑d2 i=1 ki
2d2 ≤ ε

heading: A.4 CAPACITY WITH POSITIONAL EMBEDDINGS
text: Since the Transformer architecture is permutation invariant for all t 6= τ , positional embeddings (fixed or trainable) are typically added to the inputs to distinguish the different positions of the tokens. These positional embeddings are matrices P ∈ RT×d such that P = [p1 . . . pT ]> for pi ∈ Rd. Accounting for the positional embeddings as input, a single Transformer attention head can be expressed as:
ftf-pos(X,P ;WV ,WQK) := σ ( W>V (X + P ) >softmax ( (X + P )W>QK(xτ + pτ ) )) .
For a fixed positional embedding P , let us define
Ftf-pos(P ) := {X → ftf-pos(X,P ;WV ,WQK) : ‖W>V ‖2,∞ ≤ B∞V , ‖WV ‖ ≤ BV , ‖WQK‖2,∞ ≤ B∞QK} . Position embedding just impacts the input into the covering bound argument which effects the bound in terms of the ∥∥P>∥∥ 2,∞ as given below,
Lemma A.11. For all X(1), . . . , X(m) ∈ RT×d such that ∥∥∥X(i)>∥∥∥
2,∞ ≤ BX for all i ∈ [m], and
P ∈ RT×d such that ‖P>‖2,∞ ≤ BP , the covering number of Ftf-pos(P ) satisfies
logN∞(Ftf-pos(P ); ε;X(1), . . . , X(m), ‖·‖2) . (dLσ(BX+BP ))2·
( (B∞V ) 2 3 + (2B∞QKBV (BX +BP )) 2 3 )3 ε2 · log(mT ).
Proof. Observe that ftf-pos(X,P ;WV ,WQK) = ftf-head(X + P ;WV ,WQK). Thus we have, logN∞ ( Ftf-pos(P ); ε; { (X(i)) }m i=1 , ‖ · ‖2 ) = logN∞ ( Ftf-head; ε; { X(i) + P }m i=1 , ‖ · ‖2 ) .
For all i ∈ [m], ∥∥(X(i) + P )>∥∥ 2,∞ ≤ ∥∥∥X(i)>∥∥∥ 2,∞ + ∥∥P>∥∥
2,∞ ≤ BX + BP . Therefore, using Corollary 4.4, we get the desired result.
Therefore our bounds go through for fixed positional embeddings. If we were to train the embeddings, we would need a much finer cover on the embeddings which could incur a T dependence.

heading: A.5 CAPACITY OF MULTIPLE PARALLEL HEADS
text: In virtually all practical applications of Transformers since their inception, instead of using one set of weights for an attention head, there are parallel attention heads, which have separate identicallyshaped parameters; their outputs are concatenated. For the purposes of this analysis, suppose we have
ftf-heads
( X; { W
[h] V ,W [h] QK }H h=1 ) := H∑ h=1 ftf-head ( X;W [h] V ,W [h] QK ) .
Let us define the class of multi-head self-attention with H heads as Ftf-heads := { X 7→ ftf-heads ( X; { W [h] V ,W [h] QK }H h=1 ) :
∀h ∈ [H], ‖W [h]V > ‖2,∞ ≤ B∞V [h], ‖W [h]V ‖ ≤ B [h] V , ‖W [h] QK‖2,∞ ≤ B ∞ QK [h] } .
Lemma A.12. For all X(1), . . . , X(m) ∈ RT×d such that ∥∥∥X(i)>∥∥∥
2,∞ ≤ BX for all i ∈ [m], the
covering number of Ftf-heads satisfies
logN∞(Ftf-heads; ε;X(1), . . . , X(m), ‖·‖2) . (dLσBX)2·
(∑H h=1(B ∞ V [h]) 2 3 + (2B∞QK [h]B [h] V ) 2 3 )3 ε2 ·log(mT ).
Proof. For all h ∈ [H], let Ch be an εh-covering of Ftf-head with weight bounds corresponding
to head h. Since ftf-heads ( X; { W [h] V ,W [h] QK }H h=1 ) = ∑H h=1 ftf-head ( X;W [h] V ,W [h] QK ) , we have
C := C1 × . . .× CH5 is an (∑H h=1 εh ) -covering for Ftf-heads. Using Corollary 4.4 (and optimizing for εh using Lemma A.8, by breaking them into individual errors for each head), we have
log |C| = H∑ h=1 log |Ch| ≤ H∑ h=1 ≤ (dLσBX)2 ·
(∑H h=1(B ∞ V [h]) 2 3 + (2B∞QK [h]B [h] V ) 2 3 )3 ε2 · log(mT ).
To see the dependence on H , consider the setting where the weight bounds are the same for each head (dropping the [h] subscript), then we get,
logN∞(Ftf-heads; ε;X(1), . . . , X(m), ‖·‖2) . (dLσBX)2·H3·
( (B∞V ) 2 3 + (2B∞QKBV ) 2 3 )3 ε2 ·log(mT ).

heading: A.6 CAPACITY OF DEEP TRANSFORMER NETWORKS
text: We will consider an L-layer transformer. Let us denote the weights of layer i by W (i) :={ W
(i) Q ,W (i) K ,W (i) V ,W (i) C
} such that ∥∥∥∥W (i)K W (i)Q >∥∥∥∥ 2 ≤ B(i)QK , ∥∥∥W (i)V ∥∥∥ 2 ≤ B(i)V , ∥∥∥W (i)C ∥∥∥ 2 ≤ B(i)C
and ∥∥∥∥W (i)K W (i)Q >∥∥∥∥
2,∞ ≤ B∞QK
(i), ∥∥∥W (i)V ∥∥∥
2,∞ ≤ B∞V
(i) and ∥∥∥W (i)C ∥∥∥
2,∞ ≤ B∞C (i). Let us further
denote the set of weights up to layer i by W 1:i = (W (1), . . . ,W i−1). Let the input representation of layer i be g(i)tf-head(X;W 1:i). We inductively define g with g(1)tf-head(X;W 1:1) = X
g (i+1) tf-head
( X;W 1:i+1 ) = Πnorm ( σ ( Πnorm ( f ( g (i) tf-head ( X;W 1:i ) ;W (i) ))) W (i) C ) with
f (Z; {WQ,WK ,WV ,WC}) = RowSoftmax ( ZWQ (ZWK) > ) ZWV ,
where Πnorm is applied row-wise. Our final output is gtf-scalar(X;W 1:L+1, w) = w>g
(L) tf-head
( X;W 1:L+1 ) [CLS] for ‖w‖ ≤ Bw.
In order to construct a cover, we will first bound the distance between the function g with different weight parametersW 1:L+1 and Ŵ 1:L+1. This bound will depend on the closeness of the parameters which will allow us to construct a cover of the network in an iterative fashion by constructing covers of each layer.

heading: A.6.1 LIPSCHITZNESS OF THE NETWORK
text: To bound the Lipschitzness of the network, we will first bound the distance between f with different weights and inputs.
Lemma A.13 (Instantiation of Lemma 4.3). For anyWK , ŴK ,WV , ŴV ,WQ, ŴQ ∈ Rd×k, for all Z ∈ RT×d such that ∥∥Z>∥∥ 2,∞ ≤ 1,∥∥∥∥(f (Z; {WQ,WK ,WV , ·})− f (Z; {ŴQ, ŴK , ŴV , ·}))>∥∥∥∥
2,∞ ≤ 2 ‖WV ‖2 ∥∥∥(WQW>K − ŴQŴ>K)Z>∥∥∥ 2,∞ + ∥∥∥(WV − ŴV )>Z>∥∥∥ 2,∞
Proof. Consider a fixed row τ of the output of the functions,∥∥∥f (Z; {WQ,WK ,WV , ·}) [τ ]− f (Z; {ŴQ, ŴK , ŴV , ·}) [τ ]∥∥∥ = ∥∥∥W>V Z>softmax (ZWKW>Q zτ)− Ŵ>V Z>softmax(ZŴKŴ>Q zτ)∥∥∥
5Here, × denotes the Cartesian product: the functions obtained by using the every combination of parameters of each individual cover.
By triangle inequality: ≤ ∥∥∥W>V Z> (softmax (ZWKW>Q zτ)− softmax(ZŴKŴ>Q zτ))∥∥∥ + ∥∥∥(WV − ŴV )>Z>softmax(ZŴKŴ>Q zτ)∥∥∥
Using ‖Pv‖ ≤ ‖P‖2,∞ ‖v‖1: ≤ ∥∥W>V Z>∥∥2,∞ ∥∥∥softmax (ZWKW>Q zτ)− softmax(ZŴKŴ>Q zτ)∥∥∥1 + ∥∥∥(WV − ŴV )>Z>∥∥∥
2,∞ ∥∥∥softmax(ZŴKŴ>Q zτ)∥∥∥ 1
By Corollary A.7, ∥∥Z>∥∥
2,∞ ≤ 1, ‖PQ‖2,∞ ≤ ‖P‖2‖Q‖2,∞, and ‖P >‖2 = ‖P‖2: ≤ 2 ‖WV ‖2 ∥∥∥ZWKW>Q zτ − ZŴKŴ>Q zτ∥∥∥∞ + ∥∥∥(WV − ŴV )>Z>∥∥∥2,∞
≤ 2 ‖WV ‖2 ∥∥∥(WQW>K − ŴQŴ>K)Z>∥∥∥ 2,∞ + ∥∥∥(WV − ŴV )>Z>∥∥∥ 2,∞ .
Lemma A.14. For any WK ,WV ,WQ ∈ Rd×k, for all Z, Ẑ ∈ RT×d such that ∥∥Z>∥∥
2,∞ ≤ 1, ‖Ẑ>‖2,∞ ≤ 1,∥∥∥∥(f (Z; {WQ,WK ,WV , ·})− f (Ẑ; {WQ,WK ,WV , ·}))>∥∥∥∥
2,∞ ≤ ‖WV ‖2 ( 1 + 4 ∥∥WKW>Q ∥∥2)∥∥∥(Z − Ẑ)>∥∥∥2,∞ .
Proof. Consider a fixed row τ of the output of the functions,∥∥∥f (Z; {WQ,WK ,WV , ·}) [τ ]− f (Ẑ; {WQ,WK ,WV , ·}) [τ ]∥∥∥ = ∥∥∥W>V Z>softmax (ZWKW>Q zτ)−W>V Ẑ>softmax(ẐWKW>Q ẑτ)∥∥∥
By triangle inequality: ≤ ∥∥∥∥W>V (Z − Ẑ)> softmax (ZWKW>Q zτ)∥∥∥∥+ ∥∥∥W>V Ẑ> (softmax (ZWKW>Q zτ)− softmax(ẐWKW>Q ẑτ))∥∥∥
Using ‖Pv‖ ≤ ‖P‖2,∞‖v‖1: ≤ ∥∥∥W>V (Z − Ẑ)∥∥∥
2,∞ ∥∥softmax (ZWKW>Q zτ)∥∥1 + ∥∥∥W>V Ẑ>∥∥∥
2,∞ ∥∥∥softmax (ZWKW>Q zτ)− softmax(ẐWKW>Q ẑτ)∥∥∥ 1
By Corollary A.7, ∥∥∥Ẑ>∥∥∥
2,∞ ≤ 1 and ‖PQ‖2,∞ ≤ ‖P‖2‖Q‖2,∞: ≤ ‖WV ‖2 ∥∥∥(Z − Ẑ)>∥∥∥
2,∞ + 2 ‖WV ‖2 ∥∥∥ZWKW>Q zτ − ẐWKW>Q ẑτ∥∥∥∞ By triangle inequality:
≤ ‖WV ‖2 ∥∥∥(Z − Ẑ)>∥∥∥
2,∞ + 2 ‖WV ‖2 (∥∥∥(Z − Ẑ)WKW>Q zτ∥∥∥∞ + ∥∥∥ẐWKW>Q (zτ − ẑτ )∥∥∥∞) Since ∥∥∥Ẑ>∥∥∥ 2,∞ ≤ 1 and ‖Pv‖∞ ≤ ‖P>‖2,∞‖v‖:
≤ ‖WV ‖2 ( 1 + 4 ∥∥WKW>Q ∥∥2)∥∥∥(Z − Ẑ)>∥∥∥2,∞ .
With the above lemmas, we are ready to prove the effect of change of weights on g.
Lemma A.15. For any W i+11 , Ŵ i+1 1 satisfying the norm constraints,∥∥∥∥(g(i+1)tf-block(X;W 1:i+1)− g(i+1)tf-block(X; Ŵ 1:i+1))>∥∥∥∥
2,∞ ≤ ∥∥∥∥(W (i)C − Ŵ (i)C )> σ (Πnorm (f ((X; Ŵ 1:i) ; Ŵ (i))))>∥∥∥∥
2,∞
+ LσB (i) C B (i) V
( 1 + 4B
(i) QK )∥∥∥∥(g(i)tf-block (X;W 1:i)− g(i)tf-block (X; Ŵ 1:i))>∥∥∥∥ 2,∞
+ 2LσB (i) C B (i) V ∥∥∥∥(W (i)Q W (i)K > − Ŵ (i)Q Ŵ (i)K >) g(i)tf-block (X; Ŵ 1:i)>∥∥∥∥ 2,∞
+ LσB (i) C ∥∥∥∥(WV − ŴV )>g(i)tf-block (X; Ŵ 1:i)>∥∥∥∥ 2,∞ .
Proof. Unrolling one layer, we have∥∥∥∥(g(i+1)tf-head (X;W 1:i+1)− g(i+1)tf-head (X; Ŵ 1:i+1))>∥∥∥∥ 2,∞
= ∥∥∥(Πnorm (σ (Πnorm (f (g(i)tf-head (X;W 1:i) ;W (i))))W (i)C )
−Πnorm ( σ ( Πnorm ( f ( g (i) tf-head ( X; Ŵ 1:i ) ; Ŵ (i) ))) Ŵ (i) C ))>∥∥∥∥ 2,∞
Using Lemma A.9 for each row: ≤ ∥∥∥∥W (i)C >σ (Πnorm (f (g(i)tf-head (X;W 1:i) ;W (i))))> − Ŵ (i)C >σ (Πnorm (f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i))))∥∥∥∥
2,∞
By triangle inequality for each row: ≤ ∥∥∥∥W (i)C > (σ (Πnorm (f (g(i)tf-head (X;W 1:i) ;W (i))))− σ (Πnorm (f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i)))))>∥∥∥∥
2,∞︸ ︷︷ ︸ (A)
+ ∥∥∥∥(W (i)C − Ŵ (i)C )> σ (Πnorm (f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i))))>∥∥∥∥ 2,∞ .
Let us focus on term (A).
Bounding the norm per row: (A) ≤ ∥∥∥W (i)C ∥∥∥
2 ∥∥∥∥σ (Πnorm (f (g(i)tf-head (X;W 1:i) ;W (i))))> − σ (Πnorm (f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i))))>∥∥∥∥ 2,∞
Since σ is Lσ-Lipschitz and ∥∥∥W (i)C ∥∥∥
2 ≤ B(i)C , for each row:
≤ LσB(i)C ∥∥∥∥Πnorm (f (g(i)tf-head (X;W 1:i) ;W (i)))> −Πnorm (f (g (X; Ŵ 1:i) ; Ŵ (i)))>∥∥∥∥ 2,∞
Using Lemma A.9 for each row:
≤ LσB(i)C ∥∥∥∥f (g(i)tf-head (X;W 1:i) ;W (i))> − f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i))>∥∥∥∥ 2,∞
By triangle inequality:
≤ LσB(i)C ∥∥∥∥f (g(i)tf-head (X;W 1:i) ;W (i))> − f (g(i)tf-head (X; Ŵ 1:i) ;W (i))>∥∥∥∥ 2,∞
+ LσB (i) C ∥∥∥∥f (g(i)tf-head (X; Ŵ 1:i) ;W (i))> − f (g(i)tf-head (X; Ŵ 1:i) ; Ŵ (i))>∥∥∥∥ 2,∞
By Lemma A.13 and A.14 and norm bounds on the matrices:
≤ LσB(i)C B (i) V
( 1 + 4B
(i) QK )∥∥∥∥g(i)tf-head (X;W 1:i)> − g (X; Ŵ 1:i)>∥∥∥∥ 2,∞
+ 2LσB (i) C B (i) V ∥∥∥∥(W (i)Q W (i)K > − Ŵ (i)Q Ŵ (i)K >) g(i)tf-head (X; Ŵ 1:i)>∥∥∥∥ 2,∞
+ LσB (i) C ∥∥∥∥(WV − ŴV )>g(i)tf-head (X; Ŵ 1:i)>∥∥∥∥ 2,∞ .
Combining the above gives us the desired result.
Lastly, we take account of the last linear weight and observe that,
Lemma A.16. For any W 1:L+1, Ŵ 1:L+1 and w, ŵ,∣∣∣gtf-scalar (X;W 1:L+1, w)− gtf-scalar (X; Ŵ 1:L+1, ŵ)∣∣∣ ≤ ‖w‖ ∥∥∥∥g(L+1)tf-block (X;W 1:L+1)[CLS] − g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] ∥∥∥∥+ ∣∣∣∣(w − ŵ)>g(L+1)tf-block (X; Ŵ 1:L+1)[CLS]
∣∣∣∣ . Proof. Observe that,∣∣∣gtf-scalar (X;W 1:L+1, w)− gtf-scalar (X; Ŵ 1:L+1, ŵ)∣∣∣ = ∣∣∣∣w>g(L+1)tf-block (X;W 1:L+1)[CLS] − ŵ>g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] ∣∣∣∣
By triangle inequality: ≤ ∣∣∣∣w>(g(L+1)tf-block (X;W 1:L+1)[CLS] − g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] )∣∣∣∣+ ∣∣∣∣(w − ŵ)>g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] ∣∣∣∣
Bounding the inner product by norms: ≤ ‖w‖ ∥∥∥∥g(L+1)tf-block (X;W 1:L+1)[CLS] − g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] ∥∥∥∥+ ∣∣∣∣(w − ŵ)>g(L+1)tf-block (X; Ŵ 1:L+1)[CLS] ∣∣∣∣ .

heading: A.6.2 CONSTRUCTING THE COVER
text: The cover construction follows the standard recipe of composing covers per layer (as in Bartlett et al. (2017)).
Theorem A.17. Let F (L)tf-scalar represent the class of functions of L-layer Transformer blocks satisfying the norm bounds (specified before) followed by linear layer on the [CLS] token. Then, for all X(i)
logN∞(F (L)tf-scalar; ε;X (1), . . . , X(m), ‖ · ‖2) .
log(mT )
ε2 ×
( B 2 3 w +
L∑ i=1 αi 2 3 ( d 2 3B∞C (i) 2 3 + d 2 3 ( 2LσB (i) C B (i) V B ∞ QK (i) ) 2 3 + k 2 3 ( LσB (i) C B ∞ V (i) ) 2 3 ))3
where αi = ∏ j<i LσB (j) C B (j) V (1 + 4B (j) QK).
Proof. Our goal is to show that for every ε > 0, and collection of inputs X(1), . . . , X(m), there is a cover C of vectors in R(m) such that for all W 1:L+1 and w satisfying the norm bounds, there is some v ∈ C such that maxi |gtf-scalar(X(i);W 1:L+1, w)− v| ≤ ε.
In each layer of the transformer, W (i)Q and W (i) K always appear together in the form W (i) K W (i) Q
> .
Therefore, we will overload notation and define W (i)QK : W (i) K W (i) Q
> . Our cover C will be proper,
consisting of vectors of the form (gtf-scalar(X(i); Ŵ 1:L+1, ŵ))i∈[m]. We will build the cover iteratively by finding finite collections of matrices Ŵ1:i for each layer.
First observe that for any collection of Z(1), . . . , Z(m) ∈ RT×d1 , and any W, Ŵ ∈ Rd1×d2 ,
max i∈[m] ∥∥∥W>Z(i)> − Ŵ>Z(i)>∥∥∥ 2,∞ = max i∈[m],t∈[T ] ∥∥∥W>z(i)t − Ŵ>z(i)t ∥∥∥ . This crucially allows us to aggregate over the samples and context length. In particular, we can apply Lemma 4.5 with the input vectors (z(i)t )i∈[m],t∈[T ]; a total of mT input vectors. Specifically, for any
ε andW(d1, d2, α) := {W ∈ Rd1×d2 | ‖W>‖2,∞ ≤ α} with fixed Z(i) satisfying ∥∥∥Z(i)>∥∥∥
2,∞ ≤
1, Lemma 4.5 gives us such a cover.
First let us build a cover for one Transformer layer with inputs Z(1), . . . , Z(m). We will begin with creating an εV -cover ŴV for the function class of linear transformations given by WV : {W ∈ Rd×k, ∥∥W>∥∥ 2,∞ ≤ α, ‖W‖2 ≤ s} and εQK-cover ŴQK forWQK := {W ∈ R
d×d, ‖W‖2,∞ ≤ β, ‖W‖2 ≤ r} and inputs Z(1), . . . , Z(m). For each pair of ŴV ∈ ŴV and ŴQK ∈ ŴQK , we construct an εC-cover ŴC(ŴV , ŴQK) for WC : {W ∈ Rk×d, ‖W‖2,∞ ≤ γ, ‖W‖2 ≤ c} and
inputs { σ (
Πnorm ( f ( Z(i); ŴV , ŴQK )))}m i=1 . Our final cover is
Ŵ := { (ŴV , ŴQK , ŴC) : ŴV ∈ ŴV , ŴV ∈ ŴV , ŴC ∈ ŴC(ŴV , ŴQK) } .
Using Lemma A.15, we can show that Ŵ is an ε-cover for g(·; {WV ,WQK ,WC}) and inputs Z(1), . . . , Z(m) where
ε = εC + 2LσcsεQK + LσcεV .
Using Lemma 4.5, the size of the cover is,
|Ŵ| ≤ |ŴV ||ŴQK | max ŴV ∈ŴV
ŴQK∈ŴQK ∣∣∣ŴC(ŴV , ŴQK)∣∣∣ =⇒ log |Ŵ| . ( kα2
ε2V + kβ2 ε2QK + dγ2 ε2C
) log(mT ).
We are now ready to inductively construct a cover for the deeper network. Suppose we have a ε(i)-cover Ŵ1:i for g(·;W 1:i) on X(1), · · · , X(m). We show how to construct an ε(i+1)-cover for g(·;W 1:i+1). For every element Ŵ 1:i ∈ Ŵ1:i we construct a( ε
(i) C + 2LσB (i) C B (i) V ε (i) QK + LσB (i) C ε (i) V ) -cover Ŵi(Ŵ 1:i) for the transformer layer (as above) on
inputs { g(X(j); Ŵ 1:i) }m j=1 . Consider the cover
Ŵ1:i+1 := { (Ŵ 1:i, Ŵ (i)) : Ŵ 1:i ∈ Ŵ1:i, Ŵ (i) ∈ Ŵi(Ŵ 1:i) } .
By Lemma A.15, this gives,
ε(i+1) = LσB (i) C B (i) V (1 + 4B (i) QK)ε (i) + ε (i) C + 2LσB (i) C B (i) V ε (i) QK + LσB (i) C ε (i) V .
The size of the cover is
|Ŵ1:i+1| ≤ |Ŵ1:i| max Ŵ 1:i∈Ŵ1:i ∣∣∣Ŵi(Ŵ 1:i)∣∣∣ . Inductively applying this, we get
ε(L+1) = L∑ i=1 ∏ j<i LσB (j) C B (j) V (1 + 4B (j) QK) (ε(i)C + 2LσB(i)C B(i)V ε(i)QK + LσB(i)C ε(i)V )
= L∑ i=1 αi ( ε (i) C + 2LσB (i) C B (i) V ε (i) QK + LσB (i) C ε (i) V ) where αi = ∏ j<i LσB (j) C B (j) V (1 + 4B (j) QK).
The size of the cover is
log ( |Ŵ1:L+1| ) ≤ L∑ i=1 k2B∞V (i)2 ε
(i) V
2 + d2B∞QK
(i)2
ε (i) QK
2 + d2B∞C
(i)2
ε (i) C
2  log(mT ). Notice that the layer-norm maintains the norm bound on the inputs. Lastly, we need to cover the linear layer on the [CLS] token and compose it with the cover of g1:L (as before). Using Lemma A.10 and A.16, we can get the final ε-cover C with
ε = Bw L∑ i=1 αi ( ε (i) C + 2LσB (i) C B (i) V ε (i) QK + LσB (i) C ε (i) V ) + εw
and size
log |C| . B 2 w log(m)
ε2w + L∑ i=1 k2B∞V (i)2 ε
(i) V
2 + d2B∞QK
(i)2
ε (i) QK
2 + d2B∞C
(i)2
ε (i) C
2  log(mT ) Using Lemma A.8, the size of the cover for fixed ε gives us the desired result.

heading: B PROOFS FOR SPARSE FUNCTION REPRESENTATION

heading: B.1 SETUP
text: Reductions from Boolean functions to Transformers. In order to establish our function approximation results, we must first define a canonical mapping between length-T Boolean strings b ∈ {0, 1}T and Transformer inputs X ∈ RT×d. The key point (which has also been considered since the inception of the Transformer (Vaswani et al., 2017), and continues to be a crucial consideration in practice (Dosovitskiy et al., 2020)) is that the network’s permutation-equivariant symmetry needs to be broken by assigning different embeddings to different indices of b. There are several possible natural choices here, which are all of practical interest:
• Deterministic positional embeddings. Fix positional embedding matrices P ∈ RT×d, E ∈ R{0,1}×d, and a special direction v[CLS] ∈ Rd, such that the T + 3 vectors {Pt,:}Tt=1 ∪ {Ej,:}j∈0,1 ∪ {v[CLS]} are an approximately orthonormal basis for Rd (see below). The input to the Transformer is then X = Eb + P , where Eb ∈ RT×d such that [Eb]t,: = Ebt,: for each t ∈ [T ]. In the ftf-scalar formulation, we choose the auxiliary input x[CLS] to be the constant vector v[CLS]. This closely matches applications of Transformers in NLP (Vaswani et al., 2017).
• Trainable positional embeddings. Like the above, but P is a trainable parameter; we still require approximate orthogonality of {Ej,:}j∈0,1∪{v[CLS]}. It is also possible to consider the case where E and v[CLS] are trainable (matching the way token embeddings are trained in practice). This becomes important in the regime of large vocabulary sizes that require embeddings to capture shared information between tokens; however, this is not necessary for our constructions, as we limit our consideration to binary tokens. This simplifies our constructions and improves statistical rates; additionally, it is a popular and well-studied alternative (Vaswani et al., 2017; Devlin et al., 2018; Radford et al., 2018; 2019; Brown et al., 2020).
• Bag of vectors. Fix a matrix V ∈ RT×d with approximately orthogonal rows (like the deterministic P ), but choose the Transformer input
X := V diag(b).
This construction replaces positional embeddings with positional “indicator vectors” which can be swapped between any of the Transformer’s input positions. It has the advantage of being symmetric with respect to permutation of the Transformer’s input positions: it turns out that
ftf-scalar(V diag(b)) = ftf-scalar(VΠdiag(b)), for any T×T permutation matrix Π. It is also the most natural construction when considering the composition of sparse Boolean functions across multiple layers: a layer can output combinations of the basis rows vi for further function composition, like Boolean gates.
Approximately orthonormal basis. Each of the Boolean function approximation constructions will rely on a basis set of vectors, which will be used as positional embeddings (or the variable indices in the bag-of-vectors construction). We will fix a set of approximately orthonormal vectors {vi : ‖vi‖ = 1}T ′
i=1 in Rd: for each i 6= j, we have |v>i vj | ≤ ∆. When ∆ = 0, the maximal T ′ for which such a set exists is d; for ∆ ∈ (0, 12 ), the Johnson-Lindenstrauss lemma (Johnson et al., 1986) implies that the maximal set of is of size exp(Θ(d∆2)). For given choices of d,∆ and a maximal {v1, . . . , vT ′}, our construction is valid for contexts of length T ≤ T ′. For the special vectors e0, e1, v[CLS], we will assume that these are exactly orthogonal to the vi and each other, so that the vi must be a basis in dimension d−1 or d−3. This is for clarity only– it reduces the number of error terms to propagate through the analysis.
Self-attention block. In each construction (which specifies an input X ∈ RT×d, we will specify the parameters WQ,WK ,WV ,WC , w = e1 of a scalar-output Transformer ftf-scalar, which takes an input X ∈ R(T+1)×d; the auxiliary token input will be the constant vector x[CLS] := v[CLS] ∈ Rd. The internal activation function σ is chosen to be the identity. Summarizing, the functional form of ftf-scalar ∈ Ftf-scalar in these constructions is
ftf-scalar(X;WQ,WK ,WV ,WC , e1) = softmax ( v>[CLS]WQW > KX >)XWVWCe1. In the intermediate lemmas, it will also be useful to consider the corresponding attention head output
ftf-head(X;WQ,WK ,WV ,WC) = softmax ( v>[CLS]WQW > KX >)XWVWC , and its projections ftf-head ◦Πdproj onto the first dproj coordinates.
Feedforward networks. We establish some notation for feedforward networks. An L-layer feedforward network, with activation function σ : R → R and dimensions d1, . . . , dL+1, is parameterized by weight matrices Wi ∈ Rdi+1×di , and maps x ∈ Rd1 to y ∈ RdL+1 , by the iterative equations
y>1 := σ(x >W1),
y>i+1 := σ(y > i Wi), i = 1, . . . , L− 1,
fmlp(x;W1, . . . ,WL) > = y> := y>LWi.
When dL+1 = 1, we will use the notation w instead of WL. It will be convenient to incorporate bias weights by introducing an extra input coordinate Wi ∈ R(di+1+1)×di , and augmenting the linear function accordingly:
y>i Wi 7→ [y>i 1]Wi.
Self-attention composed with a feedforward network. The full definition of the Transformer layer composes a self-attention layer (ftf-layer : RT×d → RT×d) with a position-wise feedforward network (fmlp : Rd → Rd). We will use this combination of modules to establish our function approximation results: ftf-layer acts as a sparse bottleneck, while fmlp approximates an arbitrary function of the selected coordinates. For our single-layer constructions, it is most convenient to establish notation for a scalar-output Transformer with a feedforward network. To this end, define Ftf+mlp to be the function class with the same Score,Norm, φin functions as in Ftf-scalar (thus, the same parametersWQ,WK ,WV ), with identity activation function, but a feedforward neural network replacing the linear φout and w. Concretely, with L = 3 and the ReLU activation function (·)+, Ftf+mlp contains functions of the form
ftf+mlp(X; θ) = ( (y>W1)+W2 ) + w,
y = softmax ( v>[CLS]WQW > KX >)XWVWCw, with parameters θ := (WQ,WK ,WV ,WC ,W1,W2, w).
Multiple self-attention heads. The final component we will need for the function approximation setup is multi-headed self-attention. We will extend the definition of the single-headed ftf-head to
ftf-heads
( X; { W
[h] Q ,W [h] K ,W [h] V ,W [h] C }H h=1 ) := H∑ h=1 ftf-head ( X;W [h] Q ,W [h] K ,W [h] V ,W [h] C ) ,
and substitute this definition into ftf+mlp when discussing a multi-headed construction.
Classes and properties of Boolean functions. We will call a Boolean function f : {0, 1}T → Y I-sparse if it only depends on a fixed subset I ⊆ [T ] of its inputs:
bi = b ′ i ∀i ∈ I =⇒ f(b) = f(b′).
Overloading notation, if I = s, we will also call f s-sparse. We will call an I-sparse Boolean function f symmetric if its value is invariant under permutation of the indices in I:
|{i ∈ I : bi = 1}| = |{i ∈ I : b′i = 1}| =⇒ f(b) = f(b′). Further, we will call an I-sparse real-valued symmetric Boolean function f : {0, 1}T → Y monotone if f(b) is monotonically increasing in r := |{i ∈ I : bi = 1}|. If, for some γ > 0, it holds that f(r+ 1) ≥ f(r) + γ for each r = 0, . . . , s− 1, we call f γ-strictly monotone. A vector-valued I-sparse Boolean function f : {0, 1}T → Rdf is γ-injective if
‖f(b)− f(b′)‖∞ ≥ γ for each b, b′ that differ at some position i ∈ I; f is called B-bounded if ‖f(b)‖∞ ≤ B for all b ∈ {0, 1}T .
Uniform approximation. For some ε ≥ 0 and a function f : {0, 1}T → Rd, we say that f̂ ∈ F ε-uniformly approximates f under the mapping b 7→ X(b) if∥∥∥f̂(X(b))− f(b)∥∥∥
∞ ≤ ε, ∀b ∈ {0, 1}T .

heading: B.2 RESULTS
text: We give an overview of the function approximation results under each input mapping b 7→ X(b), as a multi-part proposition: Proposition B.1 (Sparse variable creation with Transformers). The function classes Ftf-scalar,Ftf+mlp contain the following classes of sparse Boolean functions:
• Deterministic positional embeddings: For any I, Ftf-scalar can approximate a particular monotone symmetric I-sparse f , with Transformer weight norm bounds from the realvalued construction in Lemma B.2. Ftf+mlp with 1 head can exactly represent any symmetric s-sparse f , with the same bounds on Transformer weight norms, and feedforward network weight norms scaling as O(poly(s)). Ftf+mlp with s heads can exactly represent any s-sparse f , with Transformer weight norm bounds from the vector-valued construction in Lemma B.2, and feedforward network weight norms scaling as O(poly(s) · 2s).
• Trainable positional embeddings: For any I, Ftf-scalar can approximate a particular monotone symmetric I-sparse f , with positional embedding and Transformer weight norm bounds from the real-valued construction in Lemma B.3. Ftf+mlp with 1 head can exactly represent any symmetric s-sparse f , with the same bounds on P and Transformer weight norms, and feedforward network weight norms scaling asO(poly(s)). Ftf+mlp with s heads can exactly represent any sparse f , with P and Transformer weight norm bounds from the vector-valued construction in Lemma B.3, and feedforward network weight norms scaling as O(poly(s) · 2s).
• Bag of vectors: For any I, Ftf-scalar can approximate a particular monotone symmetric I-sparse f , with Transformer weight norms from Lemma B.4. Ftf+mlp with 1 head can represent any symmetric s-sparse f , with the same Transformer weight norm bounds, and feedforward network weight norms scaling as O(poly(s)). Ftf+mlp with 1 head can also exactly represent any s-sparse f , with the same bounds on Transformer weight norms, and feedforward network weight norms scaling as O(poly(s) · 2s).
The formal statements are obtained by (γ/4)-uniformly approximating a γ-strictly monotone or γ-injective function with self-attention alone (Lemmas B.2, B.3, B.4), then applying a robust universal function representation construction (Lemmas B.5, B.6) appropriately. They are organized as follows: Lemma B.2 (DeterministicP , no MLP). SupposeX(b) = P+Eb with deterministicP . Let I ⊆ [T ] such that |I| = s ≤ d, k, and ∆ < 1/s. Then, for all 0 < γ ≤ 1, there exists a 1-bounded, (2/s)strictly monotone symmetric I-sparse Boolean function gI : {0, 1}T → R and Transformer head parameters such that ftf-scalar(X(b);WQ,WK ,WV ,WC , w = e1) (γ/4)-uniformly approximates gI . The norms satisfy
‖WQ‖F ≤ log ( 8T γ ) 1− s∆ , ‖WK‖F ≤ s, ‖WV ‖F ≤ 2, ‖WC‖F ≤ 1.
Also, there exists a 1-bounded, 2-injective I-sparse Boolean function g′I : {0, 1}T → Rs and s-headed Transformer parameters such that ftf-head ( X(b); { W [h] Q ,W [h] K ,W [h] V ,W [h] C }s h=1 ) ◦ Πs uniformly approximates g′I . The norms of each head satisfy∥∥∥W [h]Q ∥∥∥ F ≤ log ( 8T γ ) 1− s∆ , ∥∥∥W [h]K ∥∥∥ F ≤ 1, ∥∥∥W [h]V ∥∥∥ F ≤ 2, ∥∥∥W [h]C ∥∥∥ F ≤ 1.
Lemma B.3 (Trainable P , no MLP). Suppose X(b) = P + Eb with trainable P . Let I ⊆ [T ] such that |I| = s ≤ d, k. Then, for any 0 < γ ≤ 1, and with the same gI as in Lemma B.2, there exists P and Transformer head parameters such that ftf-scalar(X(b);WQ,WK ,WV ,WC , w = e1) (γ/4)-uniformly approximates gI . The norms satisfy∥∥P>∥∥
2,1 ≤ s, ‖WQ‖F ≤ log
( 8T
γ
) , ‖WK‖F ≤ 1, ‖WV ‖F ≤ 2, ‖WC‖F ≤ 1.
Also, for the same g′I as in Lemma B.2, there exists P and s-headed Transformer parameters such that ftf-head ( X(b); { W [h] Q ,W [h] K ,W [h] V ,W [h] C }s h=1 ) ◦Πs uniformly approximates g′I . The norms of each head satisfy∥∥P>∥∥ 2,1 ≤ s, ∥∥∥W [h]Q ∥∥∥ F ≤ log ( 8T γ ) , ∥∥∥W [h]K ∥∥∥ F ≤ 1, ∥∥∥W [h]V ∥∥∥ F ≤ 2, ∥∥∥W [h]C ∥∥∥ F ≤ 1.
Lemma B.4 (Bag of vectors, no MLP). Suppose X(b) = V + diag(b). Let I ⊆ [T ] such that |I| = s ≤ d, k, and ∆ < 1/s. Then, for all s∆ < γ < 1, there exists an s-bounded, (1/s)strictly monotone symmetric I-sparse Boolean function gI : {0, 1}T → R and Transformer head parameters such that ftf-scalar(X(b);WQ,WK ,WV ,WC , w = e1) (γ/4)-uniformly approximates gI . The norms satisfy
‖WQ‖F ≤ log ( 8Ts(1+∆) γ−s∆ ) 1− s∆ , ‖WK‖F ≤ s+ 1, ‖WV ‖F ≤ 2s, ‖WC‖F ≤ s.
Also, there exists a 1-bounded, (1/s)-injective I-sparse Boolean function g′I : {0, 1}T → Rs and Transformer head parameters such that ftf-head(X(b);WQ,WK ,WV ,WC) ◦Πs uniformly approximates g′I . The norms satisfy the same bounds as above.
Lemma B.5 (Monotone to symmetric functions via MLP). Let f : {0, 1}T → R be any real-valued symmetric s-sparse Boolean function with index set I. Let WQ,WK ,WV ,WC be the parameters of a function
ftf-head(X;WQ,WK ,WV ,WC) := softmax ( v>[CLS]WQW > KX >)XWVWC , and let Π1 : Rd → R be the projection onto the first coordinate. Suppose that under some mapping b 7→ X(b), ftf-head ◦ Πs (γ/4)-uniformly approximates a B-bounded γ-strictly monotone symmetric I-sparse Boolean function g : {0, 1}T → R, for some γ. Then, there exists a function ftf+mlp ∈ Ftf+mlp with the same weights WQ,WK ,WV ,WC , and 3-layer feedforward network weights W1,W2, w, such that
ftf+mlp(X(b)) = f(b), ∀b ∈ {0, 1}T ,
with dimensions (d2, d3) = (4(s+ 1), 2(s+ 1)) and weight norms satisfying
‖W1‖∞ ≤ 8 max(1, B)
γ , ‖W2‖∞ ≤
8s γ , ‖w‖∞ ≤ max b∈{0,1}T |f(b)|.
Lemma B.6 (Injective to arbitrary functions via MLP). Let f : {0, 1}T → R be any real-valued s-sparse Boolean function with index set I such that |I| = s ≤ d. Let WQ,WK ,WV ,WC be the parameters of a function
ftf-head(X;WQ,WK ,WV ,WC) := softmax ( v>[CLS]WQW > KX >)XWVWC , and let Πs : Rd → Rs be the projection onto the first s coordinates. Suppose that under some mapping b 7→ X(b), ftf-head◦Πs (γ/4)-uniformly approximates a γ-injective function g : {0, 1}T → Rs satisfying ‖g(b)‖∞ ≤ B. Then, there exists a function ftf+mlp ∈ Ftf+mlp with the same weights WQ,WK ,WV ,WC , and 3-layer feedforward network weights W1,W2, w, such that
ftf+mlp(X(b)) = f(b), ∀b ∈ {0, 1}T , with dimensions (d2, d3) = (4s2s, 2 · 2s) and weight norms satisfying
‖W1‖∞ ≤ 8 max(1, B)
γ , ‖W2‖∞ ≤
8s γ , ‖w‖∞ ≤ max b∈{0,1}T |f(b)|.

heading: B.3 USEFUL LEMMAS
text: We will use a construction which approximates a “hard selection” of s indices using the softmax mixture; for this, we will need to quantify the approximation error when the inputs to the softmax function are bounded. Lemma B.7 (Softmax truncation). Let z ∈ (R ∪ {−∞})T such that zt ≥ R for each 1 ≤ t ≤ s, and zt ≤ 0 for each s+ 1 ≤ t ≤ T . Define z′ ∈ (R ∪ {−∞})T so that z′t = zt for 1 ≤ t ≤ s, and zt = −∞ for s+ 1 ≤ t ≤ T . Then, letting e−∞ = 0 in the definition of softmax(·), we have
‖softmax(z′)− softmax(z)‖1 ≤ 2 T − s s exp(R) < 2T exp(R) .
Proof. We have
‖softmax(z′)− softmax(z)‖1 = s∑ t=1 exp(zt) ( 1 1> exp(z′) − 1 1> exp(z) ) + T∑ t=s+1 exp(zt) 1> exp(z) .
The first summation is equal to
1− 1 > exp(z′) 1> exp(z) ≤ T − s s exp(R) ,
while the same upper bound holds for the second summation, since each term is at most 1s exp(R) .
Our results on approximating arbitrary sparse Boolean functions will depend on a generic construction for robustly approximating an arbitrary function f : Rd → R with a feedforward neural network. For simplicity of presentation, we use a standard6 3-layer ReLU network construction, which exactly represents a piecewise constant function in specified regions. Lemma B.8 (Exact function representation with a 3-layer ReLU net). Let f : Rdf → R, and let x1, . . . , xn ∈ Rdf such that ‖xi‖∞ ≤ B for each i ∈ [n], ‖xi − xj‖∞ ≥ 4δ for each i 6= j ∈ [n]. Then, there is a 3-layer feedforward network with ReLU activations, with parameters W1 ∈ R(df+1)×d2 ,W2 ∈ R(d2+1)×d3 , w ∈ Rd3 7, such that
fmlp(xi + z) = f(xi)
for all i ∈ [n] and ‖z‖∞ ≤ δ, where ReLU(x) := x+ = max(0, x) is applied entrywise, with d2 = 4ndf , d3 = 2n,
‖W1‖∞ ≤ max(1, B)
δ , ‖W2‖∞ ≤ df δ , ‖w‖∞ ≤ max i∈[n] |f(xi)|.
6For example, this follows from the discussion in Chapter 4 of (Nielsen, 2015). 7Here, W1,W2 have bias terms; w does not.
Proof. First, we construct a one-dimensional “bump” function basis, and propagate the Lipschitz constants. A threshold function with a linear “ramp” of width δ can be obtained from a linear combination of 2 ReLU functions:
νδ(x) := (x/δ + 1)+ − (x/δ)+ =  0 x ≤ −δ x/δ + 1 −δ ≤ x ≤ 0 1 x ≥ 0 .
Next, we construct the bump function
ψδ(x) := νδ(x)− νδ(2δ − x). By this construction, we have ψδ(x) = 1 for 0 ≤ x ≤ 2δ and ψδ(x) = 0 for x ≤ −δ and x ≥ 3δ, interpolating linearly on [−δ, 0] and [2δ, 3δ]. Next, define
ψδ(x;x0) := ψδ(x− x0 + δ)
= ( x− x0 δ + 2 ) + − ( x− x0 δ + 1 ) + − ( x0 − x δ + 2 ) + + ( x0 − x δ + 1 ) +
so that ψδ(x;x0) = 1 for |x− x0| ≤ δ, ψδ(x;x0) = 0 for |x− x0| ≥ 2δ.
We construct the first layer W1 ∈ R(d+1)×(4nd) using these bump functions: indexing the 4nd dimension by (h ∈ [4], i ∈ [n], j ∈ [d]), we construct
[W1]:,(1,i,:) :=
[ 1 δ I
−xiδ + 2 · 1 >
] , [W1]:,(2,i,:) := [ 1 δ I
−xiδ + 1 >
] ,
[W1]:,(3,i,:) :=
[ − 1δ I
xi δ + 2 · 1 >
] , [W1]:,(4,i,:) := [ − 1δ I
xi δ + 1 >
] ,
so that( [x 1]> [ [W1]j,(1,i,:) [W1]j,(2,i,:) [W1]j,(3,i,:) [W1]j,(4,i,:) ]) + [1 −1 −1 1]> = ψδ(x; [xi]j).
The second layer is used to construct n activations which are indicators of whether x is in the neighborhood of each xi. For each xi, we will simply average the df one-dimensional indicators for each coordinate, and implement a threshold function νδ/df (1 − x). We choose W2 ∈ R(4ndf+1)×(2n), with the 4ndf + 1 dimension indexed by (h, i, j) and an extra bias dimension ⊥, and the 2n dimension indexed by (h′ ∈ {1, 2}, i′ ∈ [n]) so that
[W2](h,i,:),(h′,i′) := [1 −1 −1 1]h · 1[i = i′] · 1
δ · 1,
[W2]⊥,(1,i′) := 1− df δ , [W2]⊥,(2,i′) := − df δ .
Finally, the third (output) layer w ∈ R2n, with dimensions indexed by (h ∈ {1, 2}, i ∈ [n]), multiplies the indicators of each xi by the desired f(xi):
w(1,i) := f(xi), w(2,i) := −f(xi).
For any x0 ∈ Rdf , let Bδ(x0) be the set of x such that ‖x− x0‖∞ ≤ δ. By this construction, for each x ∈ Bδ(xi), we have f(x) = xi, as required.
Note that we use 3-layer ReLU networks for function approximation in order to minimize the introduction of unnecessary notation. Some remarks:
• It would be routine to replace this construction with any architecture which can represent an arbitrary function approximately (Hornik et al., 1989; Cybenko, 1989); this includes the 2-layer feedforward networks (and nonlinear activations other than the ReLU) which are typically used by Transformers in practice.
• It is possible to embed this construction in ftf+mlp with a 2-layer ReLU network, by using (WC ,W1,W2) and introducing a nonlinearity after WC , without changing the results.
• When df = 1, W2 is unnecessary (one can represent f directly using the bump function basis).

heading: B.4 PROOFS
text: Throughout the constructions in each case, we will refer to standard coordinate bases in several spaces:
• E0, E1 ∈ Rd denote the embeddings of the 0, 1 tokens E0,:, E1,:.
• e(k)1 , . . . , e (k) k denotes the standard basis in Rk.
• e(d)i denotes the standard basis in Rd.
• e(T )1 , . . . , e (T ) T , e (T ) [CLS] denotes the standard basis in RT+1 with the special [CLS] index.
• Recall that the vi form a ∆-approximate orthonormal basis for Rd, v[CLS], e0, e1 are exactly orthogonal to each of them as well as each other, and d is chosen such that these conditions can be met.
Let n(i) be a unique bijection between I and [s]. Let vI := ∑ i∈I vi.
Approximate vector equality. We will use u ≈ε v to denote that two vectors u, v ∈ Rdu satisfy ‖u− v‖∞ ≤ ε.
Proof of Lemma B.2. We construct attention heads such that the softmax mixture always selects the indices in I. Single head, deterministic P . We seek to approximate the 1-bounded, (2/s)-strictly monotone function
1
s ∑ i∈I χi,
where χi = +1 if bi = 0 and −1 if bi = 1. Set
WQ := Rv[CLS]e (k)> 1 , WK := vIe (k)> 1 , WV := (E0 − E1)e (k)> 1 , WC := e (k) 1 e (d)> 1 ,
where R will be chosen later. Then, by approximate orthogonality,
v>[CLS]WQW > KX > = v>[CLS]WQW > K (P + Eb) > = v>[CLS]WQW > KP > ≈Rs∆ R ∑ i∈I e (T )> i .
By Lemma B.7,∥∥∥∥∥softmax (v>[CLS]WQW>KX>)− 1s∑ i∈I e (T )> i ∥∥∥∥∥ 1 ≤ 2T exp(R− 2Rs∆) .
Finally, we have
XWVWC = EbWVWC = ∑ i∈[T ] χie (T ) i  e(d)>1 , so that by Hölder’s inequality,
ftf-head(X) ◦Π1 = softmax ( v>[CLS]WQW > KX >)XWVWCe(d)1 ≈ 2T exp(R−2Rs∆) 1
s ∑ i∈I χi.
To get (γ/4)-uniform approximation, we choose R = log ( 8T γ ) 1− s∆ .
Multiple heads, deterministic P . For h = 1, . . . , s, and the same R as above:
W [h] Q := Rv[CLS]e (k)> 1 , W [h] K := vn−1(h)e (k)> 1 , W [h] V := (E0−E1)e (k)> 2 , W [h] C := e (k) 1 e (d)> h .
This is the same construction as above, but each head only selects one of the coordinates in I. Thus, by the same analysis,
ftf-head(X) ◦Πs ≈ 2T exp(R−2Rs∆) ∑ i∈I χie (d) n(i).
This function is clearly 1-bounded and 2-injective.
Proof of Lemma B.3. The constructions closely follow Lemma B.2, but are simpler.
Single head, trainable P . For each i ∈ I, set the trainable positional embeddings to be
Pi,: := { v1 i ∈ I 0 otherwise .
Set
WQ := Rv[CLS]e (k)> 1 , WK := v1e (k)> 1 , WV := (E0 − E1)e (k)> 1 , WC := e (k) 1 e (d)> 1 .
Now, we have (with equality)
v>[CLS]WQW > KX > = R ∑ i∈I e (T )> i ,
so that Lemma B.7 gives∥∥∥∥∥softmax (v>[CLS]WQW>KX>)− 1s∑ i∈I e (T )> i ∥∥∥∥∥ 1 ≤ 2T exp(R) .
Like before, we have
ftf-head(X) ◦Π1 = softmax ( v>[CLS]WQW > KX >)XWVWCe(d)1 ≈ 2T exp(R) 1
s ∑ i∈I χi.
To get (γ/4)-uniform approximation, we choose
R = log
( 8T
γ
) .
Multiple heads, trainable P . For each i ∈ I, set the trainable positional embeddings to be
Pi,: :=
{ e
(d) n(i) i ∈ I
0 otherwise .
For h = 1, . . . , s, and the same R as above:
W [h] Q := Rv[CLS]e (k)> 1 , W [h] K := e (d) h e (k)> 1 , W [h] V := (E0 − E1)e (k)> 1 , W [h] C := e (k) 1 e (d)> h .
This is the same construction as above, but each head only selects one of the coordinates in I. Thus, by the same analysis,
ftf-head(X) ◦Πs ≈ 2T exp(R) ∑ i∈I χie (d) n(i).
Proof of Lemma B.4. This input mapping does not use position embeddings, and does not need multiple heads to implement arbitrary (non-symmetric) functions. The constructed monotone and injective functions are slightly different, but the proof strategy is very similar. The key difference is that the softmax mixture is uniform only on the positions i ∈ I where bi = 1. Bag of vectors, scalar output. The function we will approximate is defined as follows:
gI(r) := r − s r + 1 , where r = ∑ i∈I bi, s = |I|.
Note that this function is (1/s)-strictly monotone, and has absolute value bounded by s. Set
WQ := Rv[CLS]e (k)> 1 , WK := (vI + v[CLS])e (k)> 1 ,
WV := ∑ i∈I vie (k)> n(i) − v[CLS] (∑ i∈I e (k) n(i) )> , WC := ∑ i∈I e (k) n(i)e (d)> 1 ,
where R will be chosen later. Then, by approximate orthogonality,
v>[CLS]WQW > KX > ≈Rs∆ R ( v[CLS] + ∑ i∈I bie (T )> i ) ,
so that by Lemma B.7,∥∥∥∥∥softmax (v>[CLS]WQW>KX>)− 1r + 1 ( e (T )> [CLS] + ∑ i∈I bie (T )> i )∥∥∥∥∥ 1 ≤ 2T exp(R− 2Rs∆) . Finally, we have
XWVWCe (d) 1 = −se (T ) [CLS] + ∑ i∈[T ] bi · v>i vI · e (T ) i ≈s∆ −se (T ) [CLS] + ∑ i∈I bie (T ) i ,
so that |ftf-head(X) ◦Π1 − gI(r)| ≤∥∥∥∥∥softmax (v>[CLS]WQW>KX>)− 1r + 1 ( e (T )> [CLS] + ∑ i∈I bie (T )> i )∥∥∥∥∥ 1
(∥∥∥XWVWCe(d)1 ∥∥∥∞ + s∆) + ∥∥softmax (v>[CLS]WQW>KX>)∥∥1 (s∆)
≤ 2Ts(1 + ∆) exp(R− 2Rs∆) + s∆.
To get (γ/4)-uniform approximation, we choose R = log ( 8Ts(1+∆) γ−s∆ ) 1− s∆ .
Bag of vectors, s-dimensional output. We use the same construction as above, except WC := ∑ i∈I e (k) n(i)e (d)> n(i) . This will allow us to approximate the function
g′I(b) = 1
r + 1 ∑ i∈I (bi − 1)e(d)n(i),
which is (1/s)-injective and has absolute value is bounded by 1. Then, for each i ∈ I, we have XWVWCe (d) i = −e (T ) [CLS] + v > i vI · e (T ) i ≈s∆ −e (T ) [CLS] + bie (T ) i . Repeating the above analysis for each coordinate, we have ftf-head(X) ◦Πs ≈ε g′I(r), where a slightly tighter bound
ε = 2T (1 + s∆)
exp(R− 2Rs∆) + s∆ comes from the fact that ∥∥∥XWVWCe(d)i ∥∥∥∞ is now bounded by 1 instead of s. The previous choice of R suffices for (γ/4)-uniform approximation.
Proof of Lemma B.5. This follows by instantiating Lemma B.8 with δ = γ/8, df = 1, n = s + 1. Notice that a (γ/4)-uniform approximation of a γ-strictly monotone function satisfies the conditions needed for Lemma B.8.
Proof of Lemma B.6. This follows by instantiating Lemma B.8 with δ = γ/8, df = s, n = 2s. Notice that a (γ/4)-uniform approximation of a γ-injective function satisfies the conditions needed for Lemma B.8.

heading: C EXPERIMENT DETAILS
text: C.1 EMPIRICAL SCALING LAWS FOR LEARNING SPARSE AND GATES
In this section, we provide details for the empirical sample complexity scaling law experiments, which are the main empirical verification of the log T dependence of the sample complexity arising from the analysis.
Experimental setup. Synthetic tasks, parameterized by sample size m and context T , were generated by the protocol described in the main paper: in each trial, one of the ( T 3 ) subsets of indices was selected uniformly at random, under i.i.d. random inputs X ∼ Unif({0, 1}T ). A 1-layer Transformer network (with the scalar output convention) was trained with Adam (Kingma & Ba, 2014) and batch gradients of the cross entropy loss for binary classification. For each choice of T ∈ {100, 150, . . . , 350, 400} ∪ {500, 600, 700, 800} and B ∈ {50, 60, 70, . . . , 200}, 50 independent trials were performed (re-randomizing the dataset generation, random initialization, and dropout masks). Cross-validation was performed on a holdout sample of size 104 every 10 iterations. At the end of 1000 training iterations, the trial was counted as a success if the maximum validation accuracy throughout training was greater than 0.99. (In 100% of runs, the training loss was driven to < 10−4, with 100% training accuracy, within 1000 iterations.)
Architecture. Like (Chen et al., 2021a), our experimental setup is based on a popular PyTorch implementation (https://github.com/karpathy/minGPT), with some optimizations for faster 1-layer training and inference. This implementation includes widely-used architectural details which deviate slightly from the theoretical presentation; refer to the referenced repository for details. All hyperparameter settings left undiscussed are taken from the defaults in this codebase.
Hyperparameters. A fixed architecture was used (d = 64, 16 parallel heads), with trainable positional embeddings initialized with Gaussian entries N (0, σ2), σ = 0.02, 3 input token embeddings (corresponding to 0, 1,[CLS]), and 2 output embeddings (corresponding to the possible labels 0, 1). For regularization mechanisms, typical choices were used: 0.1 for {attention, embedding, output} dropout; 10−4 weight decay. The Adam optimizer was instantiated with typical parameters η = 10−3, β1 = 0.9, β2 = 0.999.
Results and discussion. Our findings are summarized by Figure 4, enlarged from the main paper. On the left, the fraction of successful trials is plotted for T = {100, 200, 400, 800} with standard errors derived from the normal approximation. Notice that in this “low-data” regime, Transformer training is quite sensitive to the stochasticity in the experiments, including the training data sample (which needs to be sublinear in the context size), so there is no sharp “phase transition”. On the right, critical sample sizes are shown, defined as the smallest m for which the fraction of successes was greater than 0.5, with standard errors derived from 50 bootstrap samples. We observe that this architecture is able to solve this “planted sparse Boolean function” task with a sublinear scaling in the sample size.
Infrastructure and computational costs. Each training run took at most 20 minutes on an NVIDIA RTX A6000 GPU (with most of computation time spent on cross-validation). Although these experiments are not in the same regime as state-of-the-art settings (such as BERT pretraining), they require optimized GPU implementations to run in a reasonable timeframe (days, as opposed to months).

heading: C.2 LEARNING PARITIES
text: A natural question for further exploration is whether Transformers can learn other Boolean functions. In particular, the statistical probe presented for the s-way AND function is equally valid for XOR. The latter is arguably more intriguing: parities are the basis elements in the monomial (i.e. Fourier) expansion of a Boolean function; (O’Donnell, 2021); furthermore, there are computational hardness barriers; see the works cited in the main paper for further context.
Experimental setup. These experiments match the setting of the main AND experiments, except for a few differences, which we enumerate here. Gradient-based training is done on streaming online losses (so that there is no training/validation split), with batch size 2048, for 10000 iterations (since parities take significantly longer to fit).
Discussion. The key observation of this exploratory experiment is to point out the phenomenon that Transformer architectures can fit sparse parities at all (rather than to measure any particular trend), pointing to computational mysteries discussed in the main paper. We do not present empirical scaling laws for this set of experiments, as the experiments with small batch sizes are higher in variance, and we are unable to scale this phenomenon (i.e. learn any parities) at context sizes above T = 50.

heading: D ADDITIONAL RELATED WORK
text: In this section, we discuss some additional related work.
Attention-based architectures. Building upon the success of modern attention-based architectures, a large body of work (e.g. Goyal et al. (2020; 2021), and the works cited within) has sought to explicitly induce model sparsity and modularity in the architecture design. Our analysis is relevant to any architecture that uses a softmax (or similar) bottleneck for statistical capacity, and could inform design principles for norm-based capacity control of these architectures.
Expressive power of Transformers. Several works establish results on the representational power of self-attention architectures in regimes where the statistical guarantees are necessarily weak or vacuous (i.e. there are too many functions in the class). Dehghani et al. (2018); Yun et al. (2019); Bhattamishra et al. (2020a;b) establish universal function approximation and Turing-completeness, which have been known for previous architectures (Siegelmann & Sontag, 1995). Our work is a significantly finer-grained analysis, in which we establish a hierarchy of function classes (indexed by sparsity s) representable by these architectures, with tight (in terms of T ) statistical guarantees. Hron et al. (2020); Yang (2020) analyze properties of the kernels induced by Transformers at the infinite-width limit.
Likhosherstov et al. (2021) analyze the sparsity patterns representable by a self-attention head, with results superficially similar to ours: when the embedding dimension is at least logarithmic in the context length, all sparse matrices can be approximately realized by an attention head. However, their analysis is not about the capacity of the function class: it quantifies over the input X , and holds the parameters (WQ,WK , . . .) to be constant (rather than vice versa). This finding serves as an interesting complement to our result: even though the attention mixture weights can take on exponentially many sparsity patterns for distinct inputs, the generalization error scales as log(T ).
Interpreting attention mixtures. A line of empirical work (“BERTology”) has made progress on understanding and interpreting state-of-the-art Transformer language models by examining the activations of their attention mechanisms (Clark et al., 2019; Tenney et al., 2019; Rogers et al., 2020). In some cases, these works have found instances in which Transformers seem to have learned features that are reminiscent of (sparse) hand-crafted features used in natural language processing, without explicit supervision. Our analysis formalizes the intuition that self-attention heads can represent sparse interactions within the context in a statistically meaningful way.
Other theoretical work on self-attention. Kerg et al. (2020) analyze self-attention and its benefit for learning long-term dependencies by establishing gradient norms bounds and showing how attention helps address the problem of gradient vanishing in recurrent networks. In contrast to our results that analyze the statistical and representational properties of attention-based architectures, this work focuses on the computational aspects of gradient-based methods on recurrent networks with self-attention.
Synthetic experiments with Transformers. Power et al. (2021) train small Transformer networks on synthetic algebraic tasks, and discover an abrupt phase transition from overfitting to correct generalization similar to ours. Tay et al. (2020) propose some synthetic tasks for benchmarking the ability of Transformer variants to capture long-range dependences. Chen et al. (2021a) present a synthetic demonstration of extrapolation (inferring a maximum-reward path from random walk observations) when using Transformers for offline reinforcement learning. Lu et al. (2021) probe the transfer learning capabilities of pretrained Transformers, and consider some simple Boolean tasks. Our experimental protocol of learning sparse Boolean functions provides a simple and fundamental setting for elucidating computational and statistical properties of sequence modeling architectures.

Review summaries:
This paper claim that they contribute to the norm-based generalization bound of the Transformer model. The paper uses its Lipschitzness to get the bound of self-attention model. With this bound, the paper analyses the inductive biases of self-attention modules; it investigates which function and the long-range dependencies the self-attention modules represent. This paper aims to answer this new question. It presents a theoretical finding that bounded-norm TransformerTransformer can fit sparse functions of the input sequence with sample complexity scaling only logarithmically with the context size. This paper also empirically validates this finding and the scaling law.  

This paper proposes a theoretical bound on inductive biases self-attention module can represent. It gives detailed walkthrough about how such bound can be derived and the result complements conclusion from prior works. The conclusion of this paper is very interesting that the complexity of sparse boolean function that a self-attention module can capture is logarithmically dependent of the context size. It could cast insights on future probing works and dataset designs. A downside, however, is the experiment part which only focused on a synthetic task. Analysis is largely absent. 

The paper provides a theoretical analysis of the functioning of the self-attention modules. The paper shows that the Transformer architectures with some modifications (for ex. bounded norm) can create sparse variables. The paper proposes an experimental protocol to support the analysis.  

This paper provides rigorous norm-based bounds on the expressivity and inductive biases attention and self-attention provides when learning complex functions. The authors leverage established analytical tools involving complexity measurements, together with novel derivations for standard attention heads transformations,  to provide novel into generalization properties of Transformers, and into the types of functions they prefer to learn. Finally, the authors provide a set of simple numerical experiments using Boolean functions that serve as basic checks on their theory.