Inductive Biases and Variable Creation in Self-Attention Mechanisms
Authors: Benjamin L. Edelman
Surbhi Goel
Sham M. Kakade
Cyril Zhang
Venue: ICLR 2022 Submitted (2022)
Keywords: transformers
attention

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. We show that bounded-norm Transformer layers create sparse variables: they can represent sparse Lipschitz functions of the input sequence, with sample complexity scaling only logarithmically with the context length. We propose new experimental protocols to support the analysis and guide the practice of training Transformers, built around the rich theory of learning sparse Boolean functions.

Under review as a conference paper at ICLR 2022
INDUCTIVE BIASES AND VARIABLE CREATION IN SELF-ATTENTION MECHANISMS
Anonymous authors Paper under double-blind review
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 vari- ables: 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 prac- tice of training Transformers, built around the large body of work provably learning sparse Boolean functions.
INTRODUCTION
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 unsuper- vised representation learning. Starting with large-scale natural language processing (Vaswani et al.) 2017). self-attention has enjoyed surprising empirical successes numerous and diverse modali- ties of data. In many of these settings, self-attention has supplanted traditional recurrent and con- volutional 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 fun- damental 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 classical norm-based generalization bound for Transformer network, which can be extended to related and future architectures via a modular abstraction of at- tention mechanisms. In particular, the capacity (in terms 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 noumded-norm norm self-attention heads are capable of representing wide classes of sparse functions. This representa- tional capacity result, combined with the generalization results, provides partial theoretical expla- nation 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 Trans- formers to learn sparse Boolean conjunctions, and verify the sample complexity scaling law pre- dicted by the theory in this clean synthetic setting. We discuss how to extend and repurpose this experimental protocol of benchmarking 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.
sequence models on synthetic "cryptographic"

---

Under review as a conference paper at ICLR 2022
1.1
RELATED WORK
The direct precursors modern self-attention architectures were recurrent and convolutional net- works augmented with attention mechanisms (Bahdanau et al. 2014) Luong et al. 2015] Xu al, Landmark work by Vaswani et al. (2117 demonstrated significantly improved machine trans-
2015).
lation models via self-attention only; autoregressive language models followed shortly (Liu et al.
2018 Radford et al. 2018 2019 Brown et al. 2020), 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 al. 2021). theorem proving (Polu &
Sutskever|202D) program synthesis (Chen et al.
2021b), and reinforcement learning (Chen et al.,
2021a Janner 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. [017]] Neyshabur al., 2015] 2017 Golowich et al. 2018: Long & Sedghil 2019, Chen et al. 2019) using covering-based argu- ments; 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 fo 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 dynamical system abstraction of at- tention layers, arriving at similar Lipschitz constant calculations to ours (which are coarser- coarser-grained, since they focus contractivity and stability rather than umite-sample statistical guarantees). [Zhang eet al.] (2019); Snell et al. (202)) study idealizations of optimization problems for self-attention heads. Wei al. (2021) propose definition of statistically meaningfal approximation of func- tion 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 cir- cuit depth (rather than context size), using a margin amplification procedure See Appendix D a broader survey of related work, including universal function approximation-styll results (which ignore statistical considerations).
2 BACKGROUND AND NOTATION
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 refers to the number of samples (input sequences) in data set. || |l2 denotes the spectral norm for matrices, and ll.. denotes the (p.g) q) matrix norm where the p-norm is over columns and q-norm over rows. || Il denotes the Ep norm for vectors. For fo-nrmm we drop the subscript. B is generally used to quantify bounds norms of matrices and L for Lipschitz constants. A denotes the simplex of dimension n, that is, A n- = { E R.= 2 0, ||=]|| = 1}.
Covering numbers and uniform generalization bounds. Our main technical contribution is 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 co-orm covering number adapted from Zhang (002):
Definition 2.1 (Covering number). For given class of vector-valued functions F, the covering number N (Fie; (200)m [I) is the smallest size of a collection (a cover) C F such that vs E F. af E C satisfying max; (2())) - [(z(0)] || S e. Further, define No. (F,e, sup,.")... Noc (Feez (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 O we exploit the notation to replace F by e.
Quoting the discussion following Theorem 4.1 of (ei et 2021 "The correct norm-based Rademacher complexity bound to use for Transformers unclear."
2

---

Under review as a conference paper at ICLR 2022
Y1 15
39
3/1
3/2
Victs
irear
31
31
attention head
self-attention layer
scalar wwwf-ettereneeee atteetitenion output
Figure 1: Diagrams of attention modules described in Section 3 alignment scores (grey edges) de- termine normalized attention weights (blue), which are used to mix the inputs ZT.7. Left: Attention with a general context. Center: Stackable self-attention layer, where F1.T is the input well as the context. Right: Auxiliary [CLS] token to extract a single scalar from a self-attention layer, providing real-valued function class for classification or regression tasks.
Recall from Zhang (2022) that for the class of linear functions, Fim fo w W E Rd, |||| I Bw ). we have the covering number bound of (F; ) S O (B3 Biy/ fcw log Bx Bwm/s), where ||x:(0 || S Bx for E [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 de- termine the generalization error via standard Rademacher complexity and chaining arguments. The following lemma relates these quantities; we refer the reader to Appendix [A. for more details.
Lemma 2.2 (Generalization bound via covering number). Let 'D be a distribution over X x R and let & I. x R be b-bounded loss function that is L-Lipschitz in its first argu- ment. For a given function class F and f E F, let iis(f;D) == Blz:g)=D(((((r v)l and risk (z(()))) . (0)) == 1 Em ((f(2())), (.)). Suppose F satisfies Ifl S A for all f E F tn
and
...."" ......!!! ) Cy/e² for all for all x!!), ,(!),..., E X Then for any 5 > O, with probability at least 1 8, simultaneousty for all f and some constant >
risk(f: -risk (fi(E(), (0(20))))))))) )=)
CF (1+ log (AVm/C7)) + +2bV 4cLL m
log(1/0) (1/5)
2m
ABSTRACTIONS OF ATTENTION AND SELF-ATTENTION
Attention not straightforward to define unlike other architectural components such as convo- lutions 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 diseussed in (Luong et al. 2015), with the goal of enabling more unified and mod- ular statistical analysis. Subsequently, we show how to represent the Transformer (the predominant attention- based architecture) as 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.
3.1 ATTENTION
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 learned function of global interactions (see Figure T left). To this end, we deñne an attention head as a function which maps a context X of T' inputs {z, E X)T (e.g. the tokens in a sentence, pixels in an image, or intermediate activations neural network) and a context z E Z to an output y E V. 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 = [z1 ... Fry E Rxxd An attention head uses z to select the inputs in X to which the output y will "attend". which formalize below:
Definition 3.1 (Attention head). An attention head is a function f X + Y, specified by "align- ment score" " function Score X x Z - R parameterized by 8, E O., normalization function Norm RT AT- and position-wise transformations din X V, Pout - y parameter-
3

---

Under review as a conference paper at ICLR 2022
ized by Oi E in and Pout Oout- The output of attention head input X X ,2 2
T
= Dout
Norm S; (.... (coo:...,,,ooe(z,,22..... (z flu): dout
1=1
= Dout Pin (X; 0mm) "Norm Score(r1, "N 2; (c; Pout
where Pin (X; = Pin (z) () oin (IT: 0)] denotes the row-wise application of Vin-
Here, a vector space of input representations "mixed" by the normalized alignment scores; in this work, we will set V RR. A function class attention heads is induced by specifying parameter classes for {O,, Qin Out
The above definition corresponds to leftmost diagram in Figure |
3.2 SELF-ATTENTION AND TRANSFORMERS
A self-attention head special case of attention head, in which the context one of the inputs 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 denote X. The focus of this work is to analyze the inductive biases of such a construction. For example, for self-attention head (see Figure (center)) we would have that the I-th component is:
V/ dout (cin (X; 0m) Norm x4i 0s)); Oout)
We now define the Transformer elf-attention architecture as a special case the above. Since a Transformer layer has shared parameters between multiple output heads, it will be convenient to define all 7 outputs of this layer at once.
Definition 3.2 (Transformer self-attent eelf-attention layer). Transformer attention layer is a collection T attention heads with outputs y!..... specified by the following choices of function classes (with shared parameters between the heads) where the context for y+ T+-
Score( Fy; ( WkY) x. woWW K., Wo, WK &Rox output yr)
din( Wv) = WTT
Wv Rdak
Dout (I; Wc) = W Woo(x), We Rexd, LL L.-Liscchitz function a R R applied position-
wise, with (0)= = O.
Norm(z) softmax (x)
exp(a)
exp(x)
Defining Y = yy132 yrl
R Txd and [RowSoftmax(M])). := (ofttax(MM.)). we have
Y (RowSoftmax (XW) XWv) Wc.
Functions from the above class of Transformer layers map R Rxxd to itself, so that instances from this function class can be composed. Although Definition 3.2 only contains the "self-attention component, is not merely a simplified idealization the full Transformer architecture used in practice. We discuss some remaining discrepancies (positional embeddings. layer normalization, parallel heads, position-wise feedforward networks) Section and the appendix.
more construction: the canoni- cal way to extract 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 context of size T, a Transformer layer with + inputs constructed, with a special input index [CLS]. The input at this position is vector 2 rruss Rd (which can be considered input, or a trainable parameter); the output a linear function w Victas for a trainable parameter E Rdd This defines class of functions mapping I R. parameterized by Transformer layer's ' parameters and w, which we call the class of scalar-output Transformers.
Extracting scalar outputs from Transformer. We introduce
constant, part of the
4 CAPACITY MEASURES OF ATTENTION MODULES
In this section, we present our main technical results, along with overviews their proofs. Sec- tion 4.1 bounds the capacity of general attention head. Section 4.2 instantiates this bound for the
4

---

Under review as a conference paper at ICLR 2022
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 4444 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 this inductive bias.
Throughout this section, assume that || < Bx for all E [T]. Note that this allows for the Frobenius norm || to scale with vT. The key challenge throughout our analysis is to avoid incurring factors of norms which take sum over the dimension, by analyzing the attention parameters in appropriately chosen geometries.
Note:
4.1 CAPACITY OF GENERAL ATTENTION HEAD
Recall that the attention head architecture can be represented as a function fhead Tx x 12d Rd parameterized by 0, Oiu 6 fout as
Shead X, 2; 0s Oin, Oout) Dout (0ma (X; 0in) (0 Oout)
Denote the corresponding function class by Fhead == {(X. + fhead (X, 2; 0z. Gin Pout) 0, E o, Oi in, Oout E @ To convert the vector- valued function class to a scalar output function class, we define Fscalar == ((x,)) - w =(X,=):
Fhead, w E TR: ||w|| - B.
For simplicity, we will focus only on the attention part and assume that oout is 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. Pout is Lont Lipsccitt in the mo-norm, that is, Va, b b R& |1o (a) our (b)l S Lont l|a bll- 2. Oin is Bin-bounded in E2-norm, that is, ( [I S Bmn for all E Rdd and Om O Norm continuously differentiable and Jacobian satisfies vO E RR. || Norm(O))l S CNorm-
The Jacobian assumption might seem strong. However, softmax (the most commonly used Norm function) satisfies this with Csoftmax (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
log No (Fheadi e;
20)) i -11111)))
< inf aE[0.1]
Noo
(scorei CNormLout Bin Bx
+ Nog (Fmi:
(1-a)= (x( he/m))a |1 112-111) (1a)
Lout
where Fscore = (xx, 2) Score(a, z; 02)= 0g E Os), and Fi (a (xn O.a): On E Oua).
Note that the bound is in terms of the No covering number of functions that dependent dimen- sions d or and not T. The effect of T only shows up in the number of samples to cover. The Nos number for many classes scales only logarithmically with the number of samples (for eg, linear functions Zhang (2022). ). This is exactly what allows us to get a log T dependence for Transformers.
Since is fixed, an e-covering of Fread directly gives us Boveervegiigg for Fscalar implying, T 2
Noo (Fecalarie; {(X(1) [1)) }_,, Log (Fheadi e/ Bwi (x))", .1-112)
Proof overview. In order to prove the bound, we first show a Lipschitzness property of fir.head- This property allows to construct the cover by using the covers for Feoore and Fi..
5

---

Under review as a conference paper at ICLR 2022
Lemma 4.3 (loc Lipschitzzesss of fif-head). For any 8g, 0g E 2e Pm Om Omi for all X E R
such that |X 1,,00 S Bx.
2,00
Sheed (X, 0; Oiun w) - fhead (X, 22 Bin. will
CNarm Lont Bin Bx Score( =; Score X, 2; 0.) + Tout (Xmm(X;0m)) Pia (X; (X;0in)
2,00
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 mean- value theorem to move the maximum (loe) error over T tokens instead of sum (61), which could potentially incur a T' factor. Furthermore, this allows us to combine all samples and tokens and construct lo -covr for mT samples.
4.2 CAPACITY OF A TRANSFORMER SELF-ATTENTION HEAD
Let us now look at the case Transformer self-attention head and instantiate the covering bound. For ease of presentation and to focus the self-attention part, we collapse WoW W K to single matrix (this does not change the representation), set k d and remove the linear layer Wo Then the Transformer self-attention head (for any fixed r) can be described as
Stf-head head (; Wv, Wox) == (W.7 X softmax (XWTK=)))
which is obtained from the general formulation by setting the context to be Zr. Score(X, z,; Wox)= = X WK Norm = softmax and Pout = a.
Let us define the function class of self-attention heads with bounded norm, Fit-head == {X H B0x]. Since
fur-head (X; Wv, Wok) | 12,00 B /|Wll I By, [[Waxll2.> Wr, Wor have dimensions dependent on d and k. bounding their norms does not hide a T de- pendence. As before, to convert this vector-valued function class to a scalar output function class, define Fif-scalar = (X + w f(X):
Fif-head: E Rd [ I Bw).
We obtain the following bound on the covering number of Fir-head as a corollary of Theorem
Corollary 4.4. For any & > 0 and X x
((m) E RR xd such that [x
S Bx for all
2,50
i [m], the covering number of Fit head satisfies
(BV) + (B@x )))
log Noc (Frr-hoadi E; X(). X(m), 11 112) S Bx).
log(mm)
e²
Here & hides logarithmic dependencies on quantities besides 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 Fox == {z + x Worz: : [|Woklla. 12.00 I B3O J and Fv == {z + W 2: ||W 2 S B ||Wv S By).
Note that | WQ WoK [ WQ 22 since ||,, S so the covering number of Fox is at most the covering number of the class of functions of the form x + WoK Therefore, bound the vector-valued linear function class suffices handle both covering numbers. We derive the following covering bound which gives the desired result.
Lemma 4.5. Let W: (W E Rd xd2
|| 12,00 & Bxo), and consider the function class F {ar + Wx : W E W). For any > 0 and (1) x (N) E Rdh satisfying Vi E [N), [ < Bx.
(d Bx)
log Noo (Fie; ( ...., (N), (12) S
log.
e2
See Appendix 4.3 for general analysis.

---

Under review as a conference paper at ICLR 2022
The proof of Lemma 4.5 actually proves somewhat stronger bound for the function class given by {W E Fdixd=
W 2, S Boo | S B1}. with dx B1 in the numerator instead of (d2 Bo0), 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-invariee symmetry of Transformer network is broken by adding a positional embedding matrix P E I to the input X at the first layer. In practice, the embedding matrix often fixed and non-trainable. Our results extend to this setting in straightforward way; see Appendix A.4 If these matrices to be trained from sufficiently large class (say, PT 2,00 S 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 self- attention heads are used, and their outputs aggregated, allow for richer representation. Our anal- ysis directly extends to this setting; see Appendix fAr5 details. When single attention head is re- placed with the sum of H parallel heads, the log-covering number scales up by a factor of poly(H)
Layer normalization. Srate-of-the-and Transformer networks are trained with layer normalization modules (Ba et al.] 2016), which is generally understood to aid optimization. We keep 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 the embedding of each token bounded. Removing these layers would lead to worse dependence on the spectral norm of the matrices.
4.3 CAPACITY OF DEEP TRANSFORMER NETWORKS
In this section, we will extend our results for L-layer Transformer blocks. Denote the weights of layer by W(9 w=(() w w:(f), w(O)] }. Further denote the set of weights up layer by
l (W(1), W Denote the input representation of layer i by bef (X; Wl)). We
(i)
(1) inductively define Otf-block : RTxx xd + RTxx starting with Ptt-block (X; W1:1) = X (the input):
(i+1) Pet black (X; w1:+1) = -norr (o [Inorm (J (X: w(() w())
with f (Z: (Wo, WK, Wv. = RowSoftmax (ZWq (ZWx) ZWv.
where IIomm denotes layer normalization applied each row. We use 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
tf-block = X
(L+1) Per-block (X; W w.l.++) ViE [Z].
12
W
w
(1)
|
S C2.
|
W
|W
S Coo
2,00
To obtain final scalar output, Officcalar (X; W1:++1 w) w [[ (X:W1+++]]] (CLS),"' Let the scalar output function class be [(1))
linear function of the [CLS] output,
tisccaaa (X T f(X) (CLs) : tf-block w E Rd ||[||| S B. Theorem 4.6 (Theorem A.17 (Simplified)). Suppose Vi [m], (|
(1.)
< Bxx then we have
(F log tf-blocki (t......, X(()) $ (CLL,) d² &B3 B2 &BB
log(mT).
g²
Note that the dependence T is only logarithmic embedding dimension and (2, oo) norms of the weight matrices is quadratic. As long as the spectral norms of the matrices are bounded by I and a is I-Lipschitz (which holds for sigmoids and ReLUs), the exponential dependence on L can be avoided.
for deeper networks. The dependence
'Layer 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 dependence on T. Our results would go through with worse dependence on the spectral norms if we remove layer norm.
7

---

Under review as a conference paper at ICLR 2022
4.4 SPARSE VARIABLE CREATION: AN INDUCTIVE BIAS FOR SELF-ATTENTION
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 func- tions of interest can they express? To this end, we show in this section that sparse Boolean functions are realizable by bounded-nome norm Transformers.
Given a Boolean function f (0, + R which only depends on of its inputs, we represent using a self-attention head ftf-head composed with a feedforward network fmip: this the repeated block in the standard Transformer architecture. Intuitively, fit-head can select s-dimensional sub- set of inputs to "attend to"", while Smip memorizes an arbitrary function of these 8 inputs (requiring up to = 2 parameters). In the regime of < log'T T (think of fir-head o fmip as implementing single composable Boolean gate), the corresponding statistical guarantees are meaningful. In order to de- scribe this combination of sample-efficient efficientt 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 1) - R representable by bounded- norm scalar-output Transformer heads fif-scalar : Tx + R. To do this, we must first fix a mapping from {0, 1)Th to R²'xd; we discuss several natural choices in Appendix [B. The simplest of these uses a sum of token and positional embeddings X(br.. == eb, + ur, for a set of approximately orthogonal unit vectors {eo, e1, Vl...., 2T}. After choosing a mapping X (b). the setup of the repre- sentation problem is evident: given f(b), find Transformer weights Oer-head and feedforward network weights mlp such that
mlp X (b); Oct head, mlp) == Smip (fif head ( (x))); Oct head); mlp) R s(b).
4b E (0,
Main representational results. We show that Transformer blocks can represent I-sparse Boolcan functions, whose values only depend on some subset of indices I C [T]. We present informal statements of these approximation results below, and present the precise statements in Appendix
Proposition 4.7 (Sparse variable creation via Transformers; informal). Under any of the input map-
pings (b), we have the following guarantees:
Staaa alone can approximate a particular monotone symmetric s-sparse Boolean function, with weight norms [|W < O (log(Ts)) [|WxllF 1|l'll ||Well
S O(s).
furmlo can exactly represent symmetric s- s-sparse functions. with the same Transformer weight norms above; the feedforward network weights satisfy || ||2ll F | S (poly(ss).
ftt mlp can exactly represent general s-sparse functions, with the same Transformer weight norms as above; the feedforward network weights satisfy ||ill [W2 le || S O 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 Full proofs are given in Appendix B.4
EMPIRICAL SCALING LAWS FOR LEARNING BOOLEAN GATES
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 empirical study of whether Transformer architectures (as they are trained and used in -fate-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 family of distinct distributions {D,...., } on {0,1]" x {0, 1}. corresponding to supervised learning problems, such that the feature distributions are identical, and the uniform mixture N "E"' D, is invariant under all permutations of the indices I,...,T. T. We then select an i* E [N] uniformly at random, samples from D;., then evaluate the generalization error
train a Transformer binary classifier via cross- cross-validee validation. Since the architecture, initialization, training algorithm, and training data are permutation invariiant,, at least Q(log N) samples 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 to solve this problem, in terms of N.
8

---

Under review as a conference paper at ICLR 2022
emprr sample complestty
scarph tultrpetrra cure (T466)
210
220
140
1+2.0
Tx000
6.3
140
340
100
Figure 2: Statistically probing Transformer by training it 3-way AND of hidden subset of i.i.d. random bits. Left, center: Sublinear scaling of the empirical sample complexity. Right: Example training curves in the {overfitting, correct) regimes: T 400, m {100, 2003.
0.7
02
342
6.1
03
100
200
300
400
300
250
500
120
1000
1,50
1500
1750
2003
2003
2000
2000
4000
$000
Figure 3: A curious empirical finding: Transformers online batches) are given for this setup with 3, T' {10, 15, 20). showing abrupt phase transi- tions from random guessing perfect classification. See Appendix C.2 for details.
learn sparse parities. 10 loss curves (with
Learning sparse AND gates. The simplest instantiation of this experimental framework is hidden s-way AND under the uniform distribution (i.c. T i.i.d. random bits). The Transformer must learn which of the N = (") combinations of inputs determines the label, which requires m 2 [1 log T') samples. The theory predicts the sample complexity learning bounded-norns 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 sum- marizes our findings; details are provided in Appendix
Learning sparse parities. We can also replace the sparse AND operation with XOR: the label the parity of hidden subset of input bits. This variant emphasizes the "cryptographic" nature of S2(1"*) 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 7-)(s) time (Valiant 2012). Figure (with details 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-bart based training) succeed? We leave this computational (as opposed to statistical) mystery an open direction for future work.
this experimental setup, due to its known computational hardness.
6 CONCLUSION AND FUTURE WORK
We have presented 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 the sample complexity in practice. We believe our capacity bounds are improvable (we have only sought to obtain optimal dependence on T'. Incorporating aspects of computation and depth remains perennial challenge in this line of inquiry.
Building further upon the principles of attention constitutes a vibrant frontier of empirical research Tolstiihhin et al., 2021: Lee-Thorp et al., 2021 Jaegle et al. 202 a: d'Ascoli et al.) 202I). 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 compute- efficient, data-efficient, controllable, and reliable algorithmic interventions.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Martin Anthony, Peter I. Bartlett, and Peter I. Bartlett. Neural network learning: Theoretical foun- dations, volume 9, cambridge university press Cambridge, 1999.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv.1607.06450, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint 4rXiv:1409.0473, 2014.
Peter Bartlett, Dylan J Foster, and Matus Telgarsky. Spectrally-normalizedd margin bounds for neural networks. arXiv preprint arXiv: / 706.08498 2017.
Peter I.. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds
and structural results. Journal of Machine Learning Research, 3:463-482, 2002. URL http:
// www. jmlr wwtp://pppppp///////p. wttps//////aaaattttteeee html
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009. 11244, 2020a.
Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint M0Xi::2006.09266 2020b.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14160 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345, 2021a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harri Ed- wards, Yura Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374. 2021b.
Minshuo Chen, Xingguo Li, and Tuo Zhao. generalization bounds a family of recurrent neural networks. arXiv preprint arXi::1910.12447. 2019.
2019.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does BERT look at? an analysis of BERT's attention. arXiv preprint 1906.:/900.044 2019.
George Cybenko. Approximation by superpositions sigmoidal function. Mathematics control, signals and systems, (44::303-314, 1989.
Stéphane d' Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. arXiv preprint arXiv:2103.10691 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.0388 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 19929, 2020.
Richard M Dudley. The sizes of compact subsets of hilbert space and continuity of gaussian pro- cesses. Journal of Functional Analysis, 133):290-330, 1967.
Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size- fize-independent sample complexity of neural networks. In Conference On Learning Theory, pp. 297-299. PMLR, 2018.
10

---

Under review as a conference paper at ICLR 2022
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2020.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael C Mozer, and Yoshua Bengio. Neural production systems. Advances in Neural Information Processing Systems, 34, 2021.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural networks, 2(5):359- 366, 1989.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nagp and ntk for deep attention networks. In International Conference on Machine Learning, pp. 4376- 4386. PMLR, 2020.
Andrew Jaegle, Sebastian Borgeaud, Jean- Cean-Baptiste Alayrac, Carl Doersch, Catalin lonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint arXiv:2107. 2021a.
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention. arXiv preprint arXiv:2103 0330006,.66 2021b.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence mod- eling problem. arXiv preprint
2021.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 1912.2212.002778 2019.
William B Johnson, Joram Lindenstrauss, and Gideon Schechtman. Extensions of lipschitz maps into banach spaces. Israel Journal of Mathematics, 54(2): 1986.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588 2021.
Michael Kearns. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 4566)9983-1006, 1998.
Giancarlo Kerg, Bhargav Kanuparthi, Anirudh Goyal, Kyle Goyette, Yoshua Bengio, and Guillaume Lajoie. Untangling tradeoffs between recurrence and self-attention in artificial neural networks. Advances in Neural Information Processing Systems, 33, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
James Lee- Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.0332.. 2021.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint arXiv:2/06.03764 2021.
Peter Liu, Mohammad Salch, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser,
and Noam Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint 180iv: 10198, 2018.
Philip M Long and Hanie Sedghi. Generalization bounds for deep convolutional neural networks. arXiv preprint arXiv: /905. 12600, 2019.
Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Pretrained transformers as universal computation engines. arXiv preprint arXiv:2103.05247, 2021.
Minh-
Minh-Thang Thang Luong, Hieu Pham, and Christopher D Manning. Effective approaches to attention- based neural machine translation. arXiv preprint 308i://508.0022 2015.

---

Under review as a conference paper at ICLR 2022
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Norm-based capacity control in neural networks. In Conference on Learning Theory, pp. 1376-1401. PMLR, 2015.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally-normaalzeed margin bounds for neural networks. arXiv preprint arXiv:1707.09564 2017.
Michael A Nielsen. Neural networks and deep learning, volume 25. Determination press San Francisco, CA, 2015.
Ryan O'Donnell. Analysis of boolean functions. arXiv preprint arXiv:2105.10388 2021.
Stanislas Polu and llya Sutskever. Generative language modeling for automated theorem proving. arXiv preprint (tt::22009003393 2020.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting small algorithmic datasets. In ICLR MATH-AI Workshop, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What we know about how bert works. Transactions of the Association Computational Linguistics, 8:842-866, 2020.
Hava T' Siegelmann and Eduardo D Sontag. On the computational power of neural nets. Journal of computer and system sciences, 50(1):132-155) 1995.
Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint 2rXiv:2/03.07601 2021.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:201 1.00066, 2020.
lan Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint 1905009900 2019.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2105.016/0 2021.
Gregory Valiant. Finding correlations in subquadratic time, with applications learning parities and juntas. In 2012 IEEE 53rd Annual Symposium on Foundations of Computer Science, pp. 11-20. IEEE, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of atten- tion. arXiv preprint arXiv:2007.0255 2020.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. arXiv preprint arXiv:2 2021.
Kelvin Xu, Jimmy Ba. Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International conference on machine learning, pp. 2048-2057, PMLR, 2015.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint
arXiv:2006. 14548, 2020.
12

---

Under review as a conference paper at ICLR 2022
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeee functions? arXiv preprint arXfv:1912.10077, 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? arXiv preprint
arXiv:1912.03194, 2019.
Tong Zhang. Covering number bounds of certain regularized linear function classes. Journal of Machine Learning Research, 2(Mar):527-550, 2002.
13

---

Under review as a conference paper at ICLR 2022
A PROOFS OF CAPACITY BOUNDS
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.
A.1 RADEMACHER COMPLEXITY AND GENERALIZZTION BOUNDS
Here we briefly review Rademacher complexity and its relationship to covering numbers and gener- alization bounds. We refer the reader Bartlett & Mendelson (2002) for more detailed exposition. Definition A.1 (Empirical Rademacher complexity). For given class of functions F = { X
R} and (2()) N= the empirical Rademacher complexity R(F; 2()) x...,,(()) defined as
R(F; (())
EE
sup JEF
where a vector of m i.i.d. Rademacher random variables (Prle, = 1] Pr(e, = -1] 1/2).
In order to relate the Rademacher complexity and -coverin numbers, we use modified version of Dudley's metric entropy.
Lemma A.2 (Dudley (1967); modified). Consider real-valued function class FF such that If s A
for all f E F. Then
2(F;z() (m)) S c. inf 8+ f/1 1
log (F::; ..... 2mm) ds
620
for some constant C O.
Proof sketch. The original statement is for 2-norm covering number, but the co-norm case reduces to the 2-norm case because N2(-) S (.). The original statement also fixes = rather than taking infimum. Also, the standard statement has the integral go from 0 to oc, but these are easily replaced with and A.
For our paper, we will instantiate the above lemma for log covering numbers scaling as 1/E². Corollary A.3 (Rademacher Complexity via covering number), Consider a real-valued functtion class F such that Ifl < for all f E F. Suppose log (F;E; , ()))..,((())) S CF/E², then
2(F;2(1)
Cr 2 (m)) sc. V m
(1 (1+1og log (A (AVm/CT)
for some constant e > 0.
Proof. Using Lemma A.2
we have for some constant > 0,
log (F; (m)
2(F (m)) <
inf
8 +
de
620
6
C5 22m
inf
8 +
de
620
Cf
1
inf
8 +
de
420
5
Cj
c inf 620
8 +
V
log (A/8)
CF
[ + log
A (Avm/C7))
vV m
14

---

Under review as a conference paper at ICLR 2022
We obtain generalization guarantee from the Rademacher complexity of function class:
Theorem A.4 (Bartlett & Mendelson (2002)). Let DD be distribution over X' x R and let e: R be b-bounded loss function that is L-Lipschitz in its first argument. For a given func- tion class FF and f E FF, let risk(f::) = ((x.))=(((((() v)) and risk (f: y())) = D Th (([(2(),,)))) ()), Then for any O, with probability at least 1 - 6, simultaneously for all
FEF.
risk(f:D) risk (f: )0))
4L (F:2 (Fr..)...,?(m))
log
2bp
2m
Combining the above, we get:
Lemma A.5 (Lemma 2.2 (restated)). Consider a function class F such that [f] S A for all f E F and log Noo(F: E; ((1) ,...,, x(/)) S Cz/E² for all x(1), x!!)).... ,.il),,............... x E xm. Then for any > 0, with probability least 8, simaltaneously for all f E F,
Cf (1 (1+1og (AVm/C=)) + 2bV risk(f: D) - risk (x " )/2))=²))) S 44LL
log (1/6)
2m
for some constant c 0.
A.2 USEFUL LEMMAS
Lemma A.6. Consider function f : Rd - d such that the Jacobian of the function satisfies || f(0)]l,, S for all 0 IR,, then for any vectors 0,, 02 Rº.
[1 f(02)ll S c 821100
Proof. By the fundamental theorem of calculus applied to g(t) f(t0) + (1 -t)02), followed by a change of variables:
f(0,) ((02) = - J (t01 + (1- dt (0) - 02),
We have
ll((0) - f(02))l = -11%" J [10, + (1- -1--))02 (0) - 02)dt
By Jensen's 's inequality:
$ l ||J (10, + +(1 +11- -1-17 (1)=))2) (0, -82)11 dt
Using ||Ax|] S ||Al|1. |1=
(' || (t01 (1 (1- --)2))))11111 ||0 0z dt
By assumption on the Jacobian:
cs 110 02 0211oo
Corollary A.7. For vectors 01 02 RP. ||softma((1)) softmax(92)l1 S 2111; 82lloc.
Proof. Observe that for softmax, the Jacobian satisfies:
[(0) diag(softmax(()) softmax(()softmax(()]
We have for all 0, h,
P P
|0(0))
Σ> Σ
(offtmax(0))(IIi softmax(0);)
15

---

Under review as a conference paper at ICLR 2022
softmax(O); (1 sofottaax(0)) +
Σ
softmax(0))
j+
2
softmax(0)) (1 softmax(0);)
2.
Combining the above with Lemma A.6 gives the desired result.
Lemma A.8. For 04, 3g 2 0, the solution to the following optimization
min Flooznn
subject to
B5rz = C
is Eo and is achieved at v; C ( (z)
1/3 where
-x=10
1/3 a, 1/83/
Proof. The proof follows by standard Lagrangian analysis. Lemma A.9 (Contraction of norm). Let naamm be the projection operator onto the unit norm ball. For any vectors u, U, have n[n (u) 1l (v)ll S ||u. uu].
Proof. If u are both in the unit ball then this follows trivially. Let us assume that ||u|| 2 |||||| and ||u]|| WLOG. First suppose ||a||| S
Let B = ar the projection of in the direction of u,
and let B3 B/ Then
||_mrrmm (w) 1loorm (v)l1
||u/ |u|| e112
||u/ |u|| (axu BB-)12
II [||l -1 - -aau u - B& 1|2
(I|ul -1 aa))?llul12 |1
s (1 a3112u11? + ||B
since ||u|| <a <
||u (aru + )1?
||u e||
If ||l||| > 1, then
nnrm (u) Inorm [v) = norm (u/[l=l) In [t/I|u|||I S ||//||=||| w/1//lll < I|u [|u-ul|)
where the second-to-last inequality follows from the || case.
Lemma A.10 (Zhang 2002 Theorem 4). Let (v Rd |||||| B1} and Frinear for
T JTaiUe V}. For any 8 and an) ,...,. (NN) satisfying ||
Bg Vi,
B log(2 (4[2 Bafe 2 1).
log OF (inoarrr 8; (1))
x (N)) < 36
2²
A.3 OMITTED PROOFS
Proof of Lemma 4.3
Observe that,
/head (X, z; Os 0in) Snead X, z; 0., 0 0in))]]
/Boou (X;
Norm(Score(Xx 02))) (X
Norm(Score(X, pout (o &.)))
16

---

Under review as a conference paper at ICLR 2022
By Lout Lipschitzness of Pont and bound ||w]|
Lout
Oin (X; (X;0mn) 0m) Nor 2,),) - dm (X; "Nomm:: 2x)m(Sccoe(Xx,:20.))
By triangle inequality:
Lout
| (Xx; (Norm(Score(X, 2; 02)) 2;
(0in(X (X;n(X;mm) 8m,) oi (X; (X: Norm(Score(Xx22.0.))
Lout
Using
PPil I [1PPl,,x0///11 and B,, boundedness of Φin:
< Lout Bm
Norm(Scoree(X, 0,)) Norm(Score(A,, 2; 2;:))
+ Lout
(Qi((X:0m) oin (X; (X))m))
2x
2,00
By Lemma A.6 and the assumption on Norm:
(<on((;) Hia (X; 0mm))
< Lout CNorm | (X; (x;
112,00
Score(X, 2; 0,) Score(X, + Loout
By boundedness of Φin and |X
2,00 I Bx:
Score (x,e(XX2;;)) 2,), Score X, z; =;:))1 + Lout
( din (x;(xx))))
Lout CNorm Bx
2,00
Proof of Theorem 4.2 Our goal is show that for every
0. collection of inputs
(x(1), =(1)),.... ((m) z z(m)), there is cover Chead such that for all O, Ox, Om Oin there
is some (0s, 0m) E Chead such that max; [nead (X(1), 2(1); On, Omn) Shead (x) (0) 02, 0mn) Se.
Observe that for all 0g, 0s. max ||Score ( ;,,) 0.),;,,,Soomre (x ,(00) ie (e(
Score (6) 2000 20
max
kE(m)
Similarly, for all Pm, .m.
max <E(m) ( (x; Φia (X (X()0
[om((a ( de ( (mm)
max <E[m]
2,00
This crucially allows us aggregate over the and dimensions together Therefore, we can consider No covers for the above to bound the overall covering number.
Let Cscor be the EScore (oo) for FSccore over inputs { (0)) [e(m).rem of size
No (Fscorei Eccorei {xx1 (00
Also, Let Cin be the Ein Cover (oo) for Fin over inputs
{a HE(m))eE(T] of size Noc (Fui Sini (x(() beetmluerryi 1-112) l12))
We are ready to construct the cover for Fhead. Set Chead (-eadd 0s, Fm Oi) 0 E Cscore, Om E C.n). Then for any 0, E +,, On E Oin. there exists 0., Oin Chead, such that for all E [m]] using Lemma 4 fnead (X) 200 0,, 0,..) Shead (x)) 20; 0,, R..) Lour B Bx EScore + LoutEin-
'In the case of the Transformer self-attention mechanism, we will obtain oo-norm covering numbers for Score and din 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.
17

---

Under review as a conference paper at ICLR 2022
The size of the cover we have constructed is,
log [Chead] = log [Cscore| + log [Cial
=
No
(Sccore: EScore {(x] ( )=0m.))m))))) + log (Fui Emi faz he(m)) l12)
and done.
Proof of Corollary 4.4 By Theorem 4.2 the covering number of Fir-head satisfies
log Noo (f
(xr(), 221))
i=1
log.Noc Foqx 2L By ((x)) ((0) [e(m).tt))) aE(0,1]
inf
aE
(2t (2) (1 (1-a)) a)- he(m].rmm)]] -1111)1))) l1a)]
+ loo.gNo No
Fv;
Lo
where we have used the fact that for scalar-output Transformer layer:
softmax satisfies the Jacobian assumption with softmax using Corollary A.7
Lout is the Lipschitz constant of a:: La.
Bi is bound on the norm of WTax with respect norm of a: Bv.
By Lemma 4.5 for any EQK,EV O:
(dBQK Bx)² log(mT)
No
(Forieqki (etmmeeeer
2,inf
(dBy 2,in Bx log(m?)
No (Fvjev: ((x,) )amm, -laalla)) S
Eq
since Wor, Wy E Rddd (k = d). We want to choose EQK and ey to minimize the sum of the
above two terms, subject
2L By BxEQx + Loey Se.
By Lemma A.8 the solution to this optimization leads an optimal bound of:
(BB)) + (B By 3x)
log ( (() (...,.(MM), x(M) $ (dL,
log(mm))
g2
Proof of Lemma 4.5
Let Boc be upper bound on || and B₁ be upper bound on [W]l,,
The approach will be to cover each of the columns of W independently. treating each as specifying a linear function from Rd - R.
By Lemma A.10 letting V(b) (U:VE Rd ||l/|| S 6} and Fiincar (b) = {o + V.viVe V(b)),
for any 8
,(()) ct B log((1 bBxx/))))
log Noc (Fincar (b): (b); 8; (1))
82
given that || < Bx for all
In fact the cover, which we denote by F (b; 6), is proper: Frinarr (b: = (T:002::0 for some finite subset V C V(b).
Let
de
S (B1, kd,): k2 E (0.1.....dd for all and 0 <
k; v d2
i=1
18

---

Under review as a conference paper at ICLR 2022
Given any W E W, let e denote the ith column of W. For each i, let
d2
k =
|
Box
Let Sw = (ki,...,kk, Then Sw E S, and
B k S ||ul||
Boc
S
(kg + Vi.
d2
d2
For every tuple S = (k1,, k2 ................................ kaz) E let
Ws 101 Tdal" Di E Flincar (ks + 1)/d2:c/ (k (+++ + 1)/(2d2)) for all i).
Our cover will be F: fo Wx: W E Uses Ws}. Note that
Noo
(Bro + 1)/d2); e (k: + 1)/(22)); x.())
SeSi=1 i=1
d2
B2 a² (h; 1)² log(N)
Σ exp
dz (k: + (2d()
SES
2B² Bx log(N)
Σ exp
(ki +1
d2
SES
4B, B B¹ log(())
$ S
exp
2
where in the last step used the fact that
k
|u:l| //
Since |S| S (2dz) = cxp(O log dz)). we obtain
log f d2Br B& log(N)
as desired. In particular, we obtain the more concise, but looser, bound from the lemma statement by using the fact that B, S Bx
For particular = [v1... va,1 E W it is guaranteed that there is a matrix W E WWww W (o.... Gagl, such that for each (da).
k +1
I. vn - 7. xnl S
[N].
2d2
where Kū is the ith element of Sw. We then obtain the desired covering property:
da
max
Wx,
Wxn
Σ
(o) In
07 xx)
EE M
nE[N]
1=1
h; + 1
62
2d2
1=1
d2 +
da
h
8
S 8
2d2
19

---

Under review as a conference paper at ICLR 2022
A.4 CAPACITY WITH POSITIONAL EMBEDDINGS
Since the Transformer architecture is permutation invariant for all + positional embeddings (fixed trainable) typically added the inputs distinguish the different positions of the tokens. These positional embeddings are matrices E RT'xd such that P = Ps RR. Accounting for the positional embeddings as input, single Transformer attention head can be expressed as:
for
pot [X, P; Wv, WoK)= == (W) + P) softmax ((X + (e +p))))
For fixed positional embedding P. let define
Fit- (P) (X + ftt-pos P: Wv, Wox) |IW 112.00 19lllWll s Bv, [|eella, s BOR)
Position embedding just impacts the input into the covering bound argument which effects the bound in terms of the [
2,00 as given below,
Lemma A.11. For all x , ((m) E Rxx such that |
S Bx for all i E [m], and
P Txd such that || |P 12,00,00 S Bp, the covering number of Fer pos pos(P) P)) satisfies
(BB)) + (2B (2 Bv Bp))
log
X x(m), 1111112) S (dx
log(mT))
Proof. Observe that fir-pos (X, Wv,, Wor) = fit-head (X + P; Wv, Wok). Thus we have,
log.Noo (out (x(1)) i=1 .1 111-11 (1a) =oog (Fusheadd E; (x(0 + i= |l 11.12))
For all i E [m], (X + P) 1 2,00
$ | 11 - -1|PT S Bx + Bp. Therefore, using
2,00
Corollary 4.4 we get the desired result.
Therefore our bounds go through for fixed positional embeddings. we were to train the embed- dings, we would need a much finer cover on the embeddings which could incur a dependence.
A.5 CAPACITY OF MULTIPLE PARALLEL HEADS
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 identically- shaped parameters; their outputs are concatenated. For the purposes of this analysis, suppose we have
H
fii-heads (x: WW/h
(x:(w<)
[N] W
[b] 11 QK
Σ
fir-head (x; wrik)
k=1
A=1
Let us define the class of multi-head self-attention with H heads as
Fis-heads = { {x + Stf-heads
x: (x:fw(h Wilh
willi [h] H QK!
h=1
T vh. E [17), [W]l
112,00 I Boll
||
- BlA
||-W QK
1,,0 S BQK s
[M] {
Lemma A.12. For all
((m) RTxd such that [X(OT
S Bx for all i E [m], the
2.00
covering number F&t heads satisfies
(BE 11)) + (2B (M) B(H
log. No (Fit hfaaheesiiiee i;; (1) (m)) 11:112) S (dL,Bx)2
log(ml).
22
Proof. For all h E [H], let Cn be en-covering of Fer.had with weight bounds corresponding
to head Since SIF-heads xx: (
Wlll
= El Jirthead (x: will| wiik), ) we have
QK
R=1
20

---

Under review as a conference paper at ICLR 2022
= C₁ x .... CB
is an Eh== En) -ooering for Fir-heads- Using Corollary 4.4 (and optimizing by breaking them into individual errors for each head), we have
for En using Lemmal A.8
3 + (2Bōk [M] )
H
H
log [C] =
Σ h=1
log CN v
Bx)e
log(m7).
6²
h=1
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.
heads: : E; x y..... x(m), 11-112) x (dT, (B77) 3 + (2B Bv) 3
log.N No (Fit
log(mm))
8²
A.6 CAPACITY OF DEEP TRANSFORMER NETWORKS
We will consider an L-layer transformer.
Let us denote the weights of layer by W
O w
w(') } such that W
w()
T
S
BOK
W
S B
w S B
2
and
W K
BQK
(i) QK
1
S Bp(
and [
Let us further I BB
2,00 denote the set of weights up to layer by Wi = (W(),.. W
2,00
2,00
Let the input representation
of layer be Dif head W2)), We inductively define with let (X;W1:)) X
(i)
(1)
get het (X; W wla+1)
= Innrm (a
((Inarm
(1(6 head (x; wh:) ; :wr())) w())
(5)
with
(Z; (WQ, WK, Wy, Wcl) = RowSoftmax ZWQ (ZWk) ZWv,
where nanmm applied row-wise.
Our final output is scaaac W1++1 w)
w T (L)
9it (X:WIL+1) [CLS] for ||l||| < B.--
In order to construct cover, we will first bound the distance between the function g with different weight parameters WI++++1 and W++++ This bound will depend on the closeness of the parameters which will allow us construct cover of the network in an iterative fashion by constructing covers of each layer.
A.6.1 LIPSCHITZNESS OF THE NETWORK
To bound the Lipschitzness the network, we will first bound the distance between f with different weights and inputs.
Lemma A.13 (Instantiation of Lemma 4.3 For any WK, WK, Wv, Wv, Wo, We E Rddkkfoo all E IR' Xxd such that || 11,.00 1,
(f (Z: (Z; WK. Wv. f (z; (WQ WK. Wv. wv. .t))
WK.
2,00
s 22 (Wo - Wo x 2T
+ +|(Wer W TTT
2,
2,00
Proof. Consider fixed row of the output of the functions,
f (Z: (Wq, WK. Wv, [) (Z; (Wo. WK. Wv. -)) |7||
= WT. ZT softmax (ZWx W - WP z softmax ZWxWT
'Here x denotes the Cartesian product: the functions obtained by using the every combination of parame- of each individual cover.
21

---

Under review as a conference paper at ICLR 2022
By triangle inequality:
< w Tº (softmax (ZWx W z-) -softmax
(W )
+ +[( Wv) z softmax (WWx 2τ
Using |Puu|| - ||||2, ||
< W z T
softmax WWWWWQ= 2-)
softmax
WW
2,00
+
|(VV - Wv) "T 7
softmax
ZWWW 2, 2,) )11
2,00
By Corollary
A.7 1 11
2,00 1, PQ]12,00 1 || Q 2,20, and || ||P /PT1/2- = ||
1 22
ZWxxo2z - ZWkW ž
+ +WWv- (WW) z
2,00
s 221|Wrl12
(WoWT WOWK) nT
+ + Wv) T
TT
2,00
2,,00
Lemma A.14. For any WK, Wv, We E Rdxk for all 2,Σ E IT such that ||
2,00 S
1, 1, 112.00 S 1.
| (F Wx Wv. (2: Wx, wx.:)))))
2,00
$ 11 [1+4|| (2- -1) 1120000
Proof. Consider a fixed row of the output of the functions,
(2: (2: (Wo, WK. Wv, :)) [r]] f (2; (W; WK. Wv, - -v,
|w [WJ softmax (2WxW]2+) W. 2T softmax (W WT 2.
By triangle inequality:
WV Z 2)
softmax (ZW W. 2+) | W7.27
softmax (ZWx 2+) softmax (ZwxW3 >)
Using
|Pul| S ||
||u|1
WT Z 2
|sofmmax (ZWx W 2
2,00
W ZT
||
softmax (ZW W q2 x+) softmax (2 W
12000
By Corollary
A.7
2
< 1 and
12. & ||P1|a||@|2..x
2,00
< ||Wvll2
(2
2) T
+ 2 Wrilz
ZWxW - ZWxWT
2,00
By triangle inequality:
< <|WW l1
(z 2) T
+ 22
(/ W/ + +1/2 Wo (er 2x) (2)
2,00
Since
|
< 1 and
Prollx < || |
2,00
< <|Wvll2
(1+4 Wx W
(z - 2) T
2,00
With the above lemmas, are ready to prove the effect of change of weights on g-
22

---

Under review as a conference paper at ICLR 2022
Lemma A.15. For any Wi+!, Wi+1 satisfying the norm constraints,
(++1) (ir-bloc((X; (X;Wrltt))>
Def-block (X; W (pltttt)
(i+1)
2,00
- -
[nnrr ((xx: -
2,00
+ Lo B8 B (+44b(k) ||
( (ft-block (x; Wl:t)
(X; Wrlt
Def-block
2,00
2L B 18) B | 1 1
| -
( belkblock
(x; (2:
2,00
+ Lo | (
- T
(x;Wi:)
block
2,00
Proof.
Unrolling one layer, we have
Iefthead (X; writtt - gef 13 (x:W1s+1)
12,00
= =I[(unom (o (Hnorm (f (gifhhadd (x: m()))) )
(8) -Inorr (a (Innorm (5 (fifheead (x: (ll)) ((()
)) w(
2,,00
Using Lemma A.9 for each row:
w a [Inorm (s (gifhead (x;Wl:t), ; W a [Inor (5 (JSthhaad (x; w1
2,00
By triangle inequality for each row:
W (TInorr (f (X:Wla). :w(x)) (Tharm (f (x: ( ))
(A)
+ (I - We! a [Inarm [feffhead (X:Wi:) (7()))
2,00
Let us focus on term (A).
Bounding the norm per row:
(A) W
(Inomm
(s (frrheea (X:Wl:t; w())
[Inorm (I (6)
(x:Wi:;) (7)))
2,00
Since is Lo Lipschitz and [
| I B(O for each row:
Lg B6 Inoom (5
(f
gef-nead (x: :w)) 1Inrm (f(9(x: w(()
(8)
Using Lemma
A.9 for each row:
< Lo (
T f (gifheadd (X;W1::); - Itf (x:Wit): W((1) T
(
(8)
2,00
By triangle incquality:
< To B( |s (Iifhead (x; WX:t: :w())) f
T (if (x:WPlt) ; w(())
(1)
2,00
(0) Log B [Ittieedd (x:W"::): mm))))
(a)
(frthead (x: 17(1)
2.00
23

---

Under review as a conference paper at ICLR 2022
By Lemma A.13 and A.14 and norm bounds on the matrices:
< Lo B B (1 (1+4B/61
(X;W2:) T (x; (X:W)
2,00
21 c
(1)
B (105 -Wo
gef-head
(e)
(x:
)7|
2,,00
L. (WV - W
(6) lef-head (x:
2,,00
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 W1:++1, W1:++1 and w, T.
9tf-scalar
(X; W1L++1 w) stal (x: 11:+1
7
$ |||||
(L+1) Att-block
(x; (.W:++1))
(ots] Def.block (x: WLLt1)
+ (w - w) gef
(x (X:WL))))
[CLs]
(CLS)
Proof. Observe that,
gef scalar (X; W11++1 ur) scalar
(x: (x; W1:L+11 18
(L+1) w
beccblock (X; Wlibtl)
(++1)
X; Wlint))
(cIs)
Def
[CLS]
By triangle incquality:
w
btokblock (X; (1;W1) [cts]
(1+1) let block
X; W11::+1
+ (w (w--w) w)
T (I+1)
lef black
X: W11L++1
[cts]
(CtS)
Bounding the inner product by norms:
< ||w||
(L+1) 9et block (X; (X:12:111
9ef block
(L+1)
(X:W1:L+11
+ (w (w- (I+1)
gef block
[x; X;WWi:b+1
[CLS]
(CTS)
(CTS)
A.6.2 CONSTRUCTING THE COVER
The cover construction follows the standard recipe of composing covers per layer (as in Bartlett et al. 2017
FLL) Theorem A.17. Let sf-scalar represent the class of functions of L-layer Transformer blocks satisfy- ing the norm bounds (specified before) followed by linear layer on the [CLS] token. Then, for all X
log.Noo (FF(L)
if-scalari e; (1)
X(m)
.1)
3
log(mT))
d (2LO BbS ByR (0) + (LO (LOD)1990()) 3))()))))
x
B3
+
2²
where 0; M Lig B B(0) (1 4B
Our goal to show that for every > 0, and collection of inputs X , there is cover C of vectors in (m)) such that for all W1i++ and w satisfying the norm bounds, there is some such that max; [gt (X) WI:++ w) E.
Proof.
In each layer of the transformer, w-( and w-(() always appear together in the form w w-( ()) Therefore, will overload notation and define w w Our cover C will be proper,
24

---

Under review as a conference paper at ICLR 2022
consisting of vectors of the form (Otf-scalar (x WIL++ ()))eim)- We will build the cover itera- tively by finding finite collections of matrices Wil for each layer.
First observe that for any collection of , 7(m) RTxd, and any W. W E B
W WTz
WT (0)
W
max m
max kE
2,20
This crucially allows us to aggregate over the samples and context length. In particular, we can apply Lemma 4.5 with the input vectors (2( (e( total of mT input vectors. Specifically, for any
and W(d1, a2 (W E Rdixda xda || 12,00 a} with fixed ( satisfying 2
2,00
1, Lemma 4.3 gives us such cover.
First let us build cover for one Transformer layer with inputs 2 ....., Z(m). We will begin with creating an Ey-cover Wy for the function class of linear transformations given by Wy : {W Bdxk
| 2, a. ||ll2 I s) and EQK-cover Wox for Wox == (W E Rdxd | 2,00 S 8, Wl S r and inputs z())..., ...., 2(m),
For each pair of Wy E Wy and Wox E Wox,
construct an EC-cover WWv, WoR) for We {W E Rexd 1 || S 7, W/l2 I cf and
inputs for (Inorm (Z() Wy, Wox) ))) 1= Our final cover is
W = (Wr, WOR. We): : Wy E Wr, Wy E Wv, Wo WoWW,, Wox)
Using Lemma A.15
show that W is E-cover for g(: (Wv, WoK, Wcl) and inputs
z(1). 7 (m) where
8 sc + CSE + Leye
Using Lemma
the size of the cover is,
[W s [WvllWok
Wo(Wy, Wax)
max WrEWy Wex EWQx
ka 2
ki
dy
= log W
S
+
log(mI).
Et
EQK
et
We are now ready to inductively construct a cover for the deeper network. Suppose we have a cover Wil for g(: o(;;Wl:j on x()))...,x(m)) X(m) We show how to con- struct an (i+t))cover for g(: g(;: Wls+1). For every element Win E Wa we construct (EC + +2, Blt
(i) -cover W,(Wl:s) for the transformer layer (as above) on &OK + Lb SV
(1)
inputs {x(x(); Whn) j=1
m
Consider the cover
W = (wii wil): Wia E Wis wii e W.
By Lemma 1. this gives,
e(+1) = Ly B B (1 ( )e(k) +el0 + 2L, BL, BlQ (i)
EQK + Lo B (0)
ey.
The size of the cover is
W
Wi(Wi::)
max W
ewili
Inductively applying this, we get
(I++))= = Σ L8
BP
(1+ 4B
(E0 + 2L. BL B.0 BV + L, Bl0 (0) Ev
11 ja
25

---

Under review as a conference paper at ICLR 2022
= Σ ai (c + 21. B
B (1) ok
+ Lo B (
where -i Lo Lo B (1 (Bk
The size of the cover is
k2t BBp(00
d B (1)
de BBx
log (nolutt))
QK (i) 2
log(mT).
(1)
QK
E
Notice that the layer-norm maintains the norm bound on the imputs. Lastly, we need to cover the linear layer on the [CLS] token and compose it with the cover of gh (as before). Using Lemma
A.10 and A.16 we can get the final e-cover C with
E = B Σ
(1) (6e + 2L. ( B FOK + La L.
+ Ew
and size
B log(m)
1.
12
d B (1) QK
d
LBz(())
log CI x
log(mTT
E
(0)
(i) 2
(i)2
EKK
E
Using Lemma A.8 the size of the cover for fixed e gives us the desired result.
B PROOFS FOR SPARSE FUNCTION REPRESENTATION
B. SETUP
Reductions from Boolean functions to Transformers. In order to establish our function ap- proximation results, must first define canonical mapping between length-7 Boolean strings E {0, 1)" and Transformer inputs E Rxx The key point (which has also been considered since the inception of the Transformer (Vaswani et al., 2017). and continues to be crucial consider- ation in practice (Dosovitskiy et al. 2020)) is that the network's 's pernutation-cquivarian symmetry needs to be broken by assigning different embeddings different indices of b. There several possible natural choices here, which are all of practical interest:
Deterministic positional embeddings. Fix positional embedding matrices P E Rxxd E E IR xd, and a special direction ViCLS] E Rd such that the + + vectors { U {Eg:)e beo.1 U {MrcLss] are an approximately orthonormal basis for Id (see below). The input to the Transformer is then X = Es + P. where E, E 7x such that [Eo]r,- = Eb.;: for each t E [T]. In the ftf-scalar formulation, choose the auxiliary input æ (CLs] to be the constant vector Vicis]- This closely matches applications of Transformers in NLP (Vaswani et al. (2017).
Trainable positional embeddings. Like the above, but P trainable parameter; we still require approximate orthogonality of (io, U (Viois) J. It is also possible to consider the case where E and U Uccus] 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, we limit our consideration to binary tokens. This simplifies our constructions and improves statistical rates; additionally, is a popular and stllestudied
alternative (Vaswani et al. 2017 Devlin et al. 2018 ct al. 2020).
Radford et al. 2018 2019 Brown
Bag of vectors. Fix a matrix V E R with approximately orthogonal rows (like the deterministic P), but choose the Transformer input
X Vdiag(b).
26

---

Under review as a conference paper at ICLR 2022
This construction replaces positional embeddings with positional "indicator vectors" which can be swapped between any of the Transformer's ' imput positions. It has the advantage of being symmetric with respect to permutation of the Transformer's input positions: it turns out that
Jet scalar (V Vdiag(b))= ftt scalar (VTidiag()) (b))
for any T T permutation matrix II. It is also the most natural construction when consider- ing the composition of sparse Boolean functions across multiple layers: a layer can output combinations of the basis rows vy for further function composition, like Boolean gates.
Approximately orthonormal basis. Each the Boolean function approximation constructions will rely 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 {uz ||v=|| = 1 in Rd:: for each + j. we have l0 vy < A. When A = O, the maximal T Lindenstauss lemma (Johnson et al. 1986) implies that the maximal set of is of size exp(O(dA2)), For given choices of d, A and maximal ('1,...,,r') J, our construction is valid for contexts of length T S T". For the special vectors c VICLsI, we will assume that these are exactly orthogonal to the 2 and each other, so that the v, must be basis in dimension d- 1 or d.3. 3. This for clarity only- it reduces the number of error terms to propagate through the analysis.
for which such set exists is d; for E (0, 2), the
Self-attention block. In each construction (which specifies an input X E R''xd, will specify the parameters Wo. WK, Wv, We, w 61 of a scalar-outpur Transformer fif-scalar, which takes an input X E R(T+1)xd, the auxiliary token input will be the constant vector 2 (cLs]] V(cts) E Rd. The internal activation function a is chosen to be the identity. Summarizing, the functional form of fif-scalar Fif-scalar in these constructions is
fif-scalar (X; Wo. WK. Wy, Wc, e1) softmax (Uici)) WoW XWy Wee1. In the intermediate lemmas, it will also be useful to consider the corresponding attention head output fir-head (X; Wo, WK, Wv, Wc) softmax (vTeus: WoW x) XWrWo. and its projections fir-head [Idpro onto the first dproj coordinates.
Feedforward networks. We establish some notation for feedforward networks. An L--laer feed- forward network, with activation function R - IR and dimensions V1....,df+. parame- terized by weight matrices W, E gd xd. and maps x Rdi to y Rdi++, by the iterative equations
o == (a Wi), : := Wt). i= 1= L-1,
fmip(;; WI,..., WE) T = = =vL W..
When d = 1, we will use the notation w instead of Wt. will convenient to incorporate bias weights by introducing an extra input coordinate W; 1 and augmenting the linear function accordingly:
W + IV. 1]W,
The full definition of the Transformer layer composes a self-attention layer (ftt-ayer Rdxx + Txxd with a position-wise feedforward network (fmlp [R + Rd). We will use this combination of modules to establish our function approximation results: ftr-aaeer acts as sparse bouteneck, while fmip approximates function of the selected coordinates. For our single-layer constructions, it is most convenient to establish notation for ocalar-output Transformer with a feedforward network. To this end, define Fir to be the function class with the same Score, Norm, éin functions as in Fef-scalar (thus, the same parameters Wo, WK, Wv), with identity activation function, but feedforward neural network replacing the linear Pout and w. Concretely, with T. = 3 and the ReLU activation function (-)+, Fir-mlp contains functions of the form
Self-attention composed with feedforward network.
arbitrary
fut 0) = ((v) W1) Wi)+ w,
softmax (T WoWFX'Y xT) XWyWow, with parameters 0 = (Wq: WK. Wv. We, Wi, W2,
27

---

Under review as a conference paper at ICLR 2022
Multiple self-attention heads. setup is multi-headed self-attention. We will extend the definition of the single-headed ftf-head to
The final component will need for the function approximation
fer-heads
x: (
w
thi wilh
Σ
fir-head
(x: W W lM] w
1==1
A=1
and substitute this definition into fit+mip when discussing a multi-headed construction.
Classes and properties of Boolean functions. We will call Boolean function f (0. y I-sparse if it only depends on fixed subset I - [T] of its inputs:
b. 1/4 VIEI - ((b) S(V).
Overloading notation, I s, will also call s-sparse. We will call an I-sparse Boolean function f symmetric if its value is invariant under permutation of the indices in I:
ii E Z h; = |} = 1 E I 8/ = =1]1 = f(b) = ((b).
Further, we will call an Z-sparse valudd symmetric Boolean function f {0, 13" + y mono- tone if f(b) is monotonically increasing in r := {ii I b, = 1}/- If, for some 7 > 0, holds that +1) 2 f(r) + for each r 0, - 1, we call f y-strictly monotone. A vector-valued
I-sparse Boolean function f: {0, 1}" + Rd is r-injective if
(b) - f(b')ll 21
for each b that differ at some position i E I: f is called B-bounded [l(())l I B for all b E {0, 1)².
Uniform approximation. For some 2 0 and function f: 10,1 + Rd, we say that J E F e-uniformly approximates /
/nder the mapping b + ((b) if
f(x(()) - f(6) SE,
Ab 10.1)*.
B.2 RESULTS
We give an overview the function approximation results under each input mapping -> X(b), as multi-part proposition:
Proposition B.1 (Sparse variable creation with Transformers). The function classes Fifscalar Fip contain the following classes of sparse Boolean functions:
Deterministic positional embeddings: For any I, Fif-scalar approximate particular monotone symmetric I-sparse f. with Transformer weight norm bounds from the real- valued construction in Lemma B.2 Fet+mip with head can exactly represent any sym- metric s-sparse f. with the same bounds on Transformer weight norms, and feedforward network weight norms scaling (ppoly(s)). Fir+ with 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 (ppoly(s) 22).
Trainable positional embeddings: For any I, Fer can approximate particular mono- tone symmetric I-sparse f. with positional embedding and Transformer weight norm bounds from the real-valued construction in Lemma B33] Fir+mlp with 1 head can actly represent any symmetric s-sparse f. with the same bounds on and Transformer weight norms, and feedforward network weight norms scaling as ((poly(s)) Fet+mip with s heads can exactly represent any sparse f. with P and Transformer weight norm bounds from the vector-valued construction Lemma B.3 and feedforward network weight norms scaling as (s) 2*).
Bag of vectors: For any I, Fif-scalar can approximate a particular monotone symmetric I-sparse f. with Transformer weight norms from Lemma B.4 Fer+ with 1 head can represent any symmetric s-sparse f. with the same Transformer weight norm bounds, and feedforward network weight norms scaling O(poly(s)). Fir+ 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 (spoly(s) 2*).
28

---

Under review as a conference paper at ICLR 2022
The formal statements obtained by (r/4)-niffomly approximating y-strictly monotone or y-injective function with self-attention alone (Lemmas B.2, B.3 B.4), then applying robust uni- versal function representation construction (Lemmas B.5 B.6 appropriately. They organized as follows:
Lemma B.2 (Deterministic P. MLP). Suppose (() = P+E, with deterministic P. Let S [T] such that [I] = S d, k, and A < 1/s. Then, for all 0 < Y S 1, there exists a 1-bounded, (2/s)- strictly monotone symmetric I-sparse Boolean function gr {0, 1) -> TK and Transformer head parameters such that J&t scalar (X (X);; WQ. WX, Wv, Wc, w e1) (n/4))uuifyrmly approximates 9T. The norms satisfy
log
ST
Wall
W
[Wvllp S2,
|Wollp S L.
Also, there exists a 1-bounded, 2-injective I-sparse Boolean function & : =00117 7 R and s-headed Transformer parameters such that fet-head (x(b); { wIl w!N w w M o II, uniformly approximates 1z- The norms of each head satisfy
K
h=1
log 1 8
W K
|
1,
22
1.
Lemma B.3 (Trainable P. no MLP). Suppose X(b) P E. with trainable P. Let I C [T] such that [I] d. Then, for any 0 < S 1, and with the same gI in Lemma B.2 there exists P and Transformer head parameters such that fti (X);; WQ, W Wv, We. w = e1)
(nifoounyy approximates gr. The norms satisfy
88
| 12,1 S,
I[Wllle S Slog
[WKllp 1
[[vll 22,
|Wll S 1.
Also, for the same S' as in Lemma
B.
there exists and s-headed Transformer parameters such o II, uniformly approximates d'z- The norms of n=1
that fit head (X()); Will [h] w each head satisfy
w w
M
|| 12,1 S.
| S log
87
WiNh 1,,
ww 2,
S
Lemma B.4 (Bag of vectors, no MLP). Suppose X X(b) = V + diag(b). I [T] such that [ = < d, ki and A < 1/s. Then, for all sÂ <? < 1, there exists s-bounded, (1/s)- strictly monotone symmetric I-sparse Boolean function 9z 10.1)T + IR and Transformer head parameters such that ful-calar (X(b): WQ. WK, Wv, Wo, e1) (r/4))unifonme approximates
gT. The norms satisfy
log
(8T((++A) 7-96
[Woll: S
[WKllF 11
[[ S 2s,
|| S
8A
Also, there exists a 1-bounded, (1/s) iinjectvve I-sparse Boolean function d/z (0, 10. 1 + and Transformer head parameters such that fir head (X(b); WQ. W K, Wv, Wc) o II, uniformly approx- imates g/T- The norms satisfy the same bounds as above.
Lemma B.5 (Monotone to symmetric functions via MLP). Let f: :0001 10.13T - R be any real-valued symmetric s-sparse Boolean function with index set Let Wo, K, Wy, We be the parameters of function
Stt-boad ( X; Wo, WK, Wv, Wc) == softmax (VTCLS: WoWIX' XWwWe.
and let IT : .R. + R be the projection onto the first coordinate. Suppose that under some mapping b + X(b). ftt-head II, (nnisoonyy approximates a B-bounded 7-strictly monotone sym- h metric I-sparse Boolean function g 1),
1] + R. for some 7,. Then, there exists a function fur+mpp Fir+mip with the same weights WQ. WK. Wy. We. and 3-layer feedforward network weights W1, W2, w, such that
futmip X (x(b) f(b).
vb E (0,
29

---

Under review as a conference paper at ICLR 2022
with dimensions (d2, ds) (4(s + 1), 1)) and weight norms satisfying
8 8max(1, B) & S
8s
||2/ S max If(b).
6E(0.1)7
Lemma B.6 (Injective to arbitrary functions via MLP). Let f {0, 1] - R be any real-valued
s-sparse Boolean function with index set such that [I] = s I d. Let We, WK. Wv, Wc be the parameters of a function
fie-head (X; Wo, WK, Wy, Wc) == = softmax (UTCLs: WoW XWvWe.
and let 11, Ra Rº be the projection onto the first coordinates. Suppose that under some mapping b + X (()), ftt-head oll, (7/4)-uniformly approximates a 7-injective f-unjicttivfuuncionn
{0, 1 + I2 satisfying (lg(b)lloc I B. Then, there exists a function fif+mlp E Fif+ with the same weights Wo. WK. Wv, Wc. and 3-layer feedforward network weights W1, W2, w. such that
fit (X())((()))))))))))))))) = [(b),
vb E {0, 1)"
with dimensions (d2, d3) = -4s22, , 2.2. 2n) and weight norms satisfying
8 8max(1, B)
W22lo S 88
|| I
|| S max If(b).. be (E(1,1)7
B.3 USEFUL LEMMAS
We will use construction which approximates "hard selection" of 8 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 2 E (R ( (-oc)) such that 2 2 R for each 1 s, and < 0 for each s + < t T. Define z' E (IR U (--o])) so that 2 21 for St S and 2t -00 for + 1 SEST. Then, letting e = in the definition of softmax(-), we have
T s
2T
||sotmma(2z)) - -sootmax(22slll S
sexp(R)
exp(Rt
Proof. We have
T
exp(21)
||softmax(=') softmax(=)]l
exp(2t) (Texp(z²) 1 (Tep(2))
Σ
1 exp(22)
1=++1
The first summation is equal to
1 exp(22) 1 exp(2)
T S eex((R)
while the same upper bound holds for the second summation, since each term is at most exx(
Our results on approximating arbitrary sparse Boolean functions will depend on a generic con- struction for robustly approximating arbitrary function f Rd R with a feedforward neural network. For simplicity of presentation, we use a standard 3-layer ReLU network construction, which exactly represents piecewise constant function in specified regions.
Lemma B.8 (Exact function representation with a 3-layer ReLU net). Les f IRd, + R. and let 01,-..,Ia Rdd, such that ||~i|1,o < B for each i E [m]. | - -yyloo 2 48 for each i + j E [n]. Then, there 3-layer feedforward network with ReLU activations, with parameters W, E R(d,-1ixdd W2 E 1dd++1)xds w E dd such that
Smip (zi + 2) f(x;) for all i E [n] and || s 6, where ReLU(x)) == I+ = x0) is applied entrywise, with
d2 4ndy, d3 2n.
| max(1, (,
dj
W2
|| max [f(x)))
8
"For example, this follows from the discussion in Chapter 4 of Nielsen 2015) Here, W1, W2 have bias terms; does not.
30

---

Under review as a conference paper at ICLR 2022
Proof. First, construct one-dimensi "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 -
pz(x) = (x/6 (x/0)=
x/0
- 551050
x 220
Next, we construct the bump function
(x)(2) == 1z(2) - v2(20 - x).
By this construction, we have (1) = 1 for 0 S 20 and (1) 0 for x: S - and x 2 36, interpolating linearly on [- 0] and [26, 36]. Next, define
48(1;;)0)
== 188 I0 0)
To 5
3o + +1) +11
+1)
+
+2
8
+
+
so that (5 20) for xol
8, 1 zo) for |z zol 2 28.
We construct the first layer Wi E R(d+1)x (4)(d)) using these bump functions: indexing the 4nd
dimension by (/ E [4], [n], E [d], we construct
11
[Wilaas))
[Wil.e2, ==
- +
I + 1
11
-
(1).(3. = 5g +2.1
[Wila(axa) =
5/2
+1
so that
(|x IJ_ [Wils.a,..) (ill.... (Wilg,3...) (Wibsaax))) -1 -1 11 (1 (zrl))-
The second layer is used to construct activations which are indicators of whether is in the neigh- borhood of each v;. For each 2,, we will simply average the d, mie-dimensional indicators for each coordinate, and implement a threshold function 15/d,(1 (1 w). We choose W₂ E TR + with the 4ndy + 1 dimension indexed by (h, i, and an extra bias dimension 1, and the 2n dimen- sion indexed by (h' {1,2), E [n]] so that
[W2] [Wal0.ss,00e))) [1 -1 -1 1/a. 1li r'] 5 1,
[We]..u.e) := dj g
d, (W2L1(2.1) = -
Finally, the third (output) layer w E I22", with dimensions indexed by (h E {1,2), e [n]), multiplies the indicators of each v; by the desired f(x.):
w(L)) f(x)). W(2,x) == ff(x)).
For any zo E TRd,, let (xo) be the set of such that ||. - volloc I 8. By this construction, for each x E B3;);), we have f(x) Fú as required.
Note that use 3-layer ReLU networks for function approximation in order minimize the intro- duction 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 ReI.U) which typically used by Transformers in practice.
It is possible embed this construction fir with a 2-layer ReLU network, by using (We, Wi, W2) and introducing a nonlinearity after Wc. without changing the results.
When dy = 1, W2 is unnecessary (one can represent ff directly using the bump function basis).
31

---

Under review as a conference paper at ICLR 2022
B.4 PROOFS
Throughout the constructions in each case, we will refer to standard coordinate bases in several spaces:
Ea, E1 E Id denote the embeddings of the 0, 1 tokens Eo,:, E
& denotes the standard basis in R
(d) denotes the standard basis in Rd
T (Ct.)) denotes the standard basis in IR R++ with the special [CLS] index.
(T) (T)
Recall that the e2 form A approximmaee orthonormal basis for Rd Vrcus] co, 61 are exactly orthogonal to each of them well as each other, and d is chosen such that these conditions can be met.
Let n(i) be unique bijection between I and [s]. Let UI
E-
Approximate vector equality. We will use u Re to denote that two vectors u, v E Rd satisfy ||u olloo c.
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
Σ X
where X; = = if = and - b; = Sct
Wq= Ru (N)T
Wx vzei TK)T
Wv = :Fa E1)e (K)
(k) (d) We ej 6l
where R will be chosen later. Then, by approximate orthogonality,
T Vccss WoW =Tecss WQWW(P + Ex) = UTLLe WoW PT ReA R Σ (T)
iEi
By Lemma
B.7
softmax (u(ens) WoW
(T)T
<
2T exp R 2Rs
TET
Finally, have
Xw E,W. =
Σ (Te(T) (d)T
(E(T]
so that by Hölder's 's inequality,
firhhad (X) o II softmax (UT WQW XW Woe
27
Xi
To get (n/4)-uniform approximation, we choose
log 87'
R
Multiple heads, deterministic P. For h = I..... s, and the same R as above:
W!ll Ru(cus]e(())
(K)T wilh [k]
=n-1(1e (K)T
will (Fo E
(K) W!l
(b) (d)
e cj
32

---

Under review as a conference paper at ICLR 2022
This is the same construction as above, but each head only selects one of the coordinates in I. Thus, by the same analysis.
fti-head (X) II,
Xien(s)
exp( TiEa
VEZ
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, set the trainable positional embeddings be
Pa: = f0r
0
otherwise
Set
We == Rete (
(K)T
WK == 01 (K)T
Wy = (Eo E
We = 9 c
(k) (d)
Now, we have (with equality)
[CLS] WOW
R
6,
(T)T
IE
so that Lemma B.7 gives
softmax (V(css) WoWxx]]
(T)
2T exp(R)
iET
Like before, we have
fir-head (X) o II1 = softmax (v[ma) W Wee
Σ Xi-
27
EET
To get (niff)u approximation, choose
87
R = log
Multiple heads, trainable P. For each i E I, set the trainable positional embeddings to be
i I otherwise
Pa. n(
For = I,...,,,, s, and the same as above:
with This is the same construction as above, but each head only selects one of the coordinates in Thus, by the same analysis.
w Rutesste (K)T
whii c 6i
[h]
(d) (K)T
W
[h]
(Eg
Ej)e
(K)T
(k) (d)
cp
f head (X) IT,
(d)
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 E where b; =
Bag of vectors, scalar output. The function will approximate is defined follows:
(r(r)=
where Σ by
$ [Z]
r ++1
VEr
33

---

Under review as a conference paper at ICLR 2022
Note that this function is (1/s)-stricn monotone, and has absolute value bounded by Set
Wo= RV (cns) e (K)T
WK
(UI ((crss) )e (A)T
Wv
Σ (K)T (11)
Vicss
Σ
m(1))
We: = Σ
c
(k) (d)
KET
EE
iET
where R will be chosen later. Then, by approximate orthogonality,
ViCLs (elssWoWis WoW X ERAA R
Vicus] + Σ bic,
(T)
so that by Lemma B.7
softmax (u (vccss) We WoW K x
(T)
Σ
bye (T)T
2T exp(R R - 2Rs
Cioss]
r+1
SEI
Finally, we have
XWvWoe (d) = -88 (T)
[CL8]
Σ
bi
T UI
(2)
2.A
8e (cLo]
+ Σ bie (2)
(E(T]
1EI
so that
[ftr-ead o 111 -g((1)) 5
Σ )
softmax (U[CLs) Wq WQW
1
((cLs]
(T)
+
b,e;
(T)
(/xwowocf))
+
r+1
IEI
+ ||oftmax (UTOTs) W ( (sA)
27 A
exp(R R ( +A.
To get (7/4) (rn4)-uniform approximation, choose
log 87=(1+A) 7 5- 5
R
1-sÀ
Bag of vectors, s-dimensional output. We use the same construction as above, except
Wc = Σe (k) (d)T
This will allow us approximate the function
(b) -1-11111 Σ (b: - 1)
iEI
which is (1/s) -ijje-ineective and has absolute value is bounded by 1. Then, for each i I. we have (T) c(cis] + ujJuz Vz (T) - E. - ectos bic; (T)
XWv Woe!" (d)
Repeating the above analysis for each coordinate, we have
Seat (X) o ITg Re (r),,,
where a slightly tighter bound
21(1+
E exp(R - ( + +A
comes from the fact that XWw (e of R suffices for (n/f)-uniform approximation.
(d) is now bounded by 1 instead of s. The previous choice
Proof of Lemma B.5 This follows by instantiating Lemma B.8 with == 7/8, 1,, = s + 1.
Notice that a (7/ uniifor approximation of y-strictly monotone function satisfies the conditions needed for Lemma B.8
Proof of Lemma B.6 This follows by instantiating Lemma B.8 with 5 7/8, d, s,n 2
Notice that a (7 4)-uniform approximation r-iijecttvve function satisfies the conditions needed
for Lemma B.8
34

---

Under review as a conference paper at ICLR 2022
C EXPERIMENT DETAILS
C.I 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 I dependence of the sample complexity arising from the analysis.
Experimental setup. Synthetic tasks, parameterized by sample size m and context T, were gen- erated by the protocol described in the main paper: in each trial, one of the (3) subsets of indices was selected uniformly random, under i..d. random inputs X Unif 177). 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 E {100, 15.. ,50, 400} U (500, 600, 700, 800} and B E (50, 60, 70...., (0,....,00).. 50 inde- pendent trials were performed (re-randomizing the dataset generation, random initialization, and dropout masks). Cross-validation was performed on a holdout sample of size 10 every 10 itera- tions. At the end of 1000 training iterations, the trial was counted as success if the maximum validation accuracy throughout training was greater than 0.99. (In 100% of runs, the training loss was driven 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
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.
com/ Soon//aaraath//iinneee
Hyperparameters. A fixed architecture was used (d = 64, 16 parallel heads), with trainable positional embeddings initialized with Gaussian entries (N0, o²), 0.02, input token embed- dings (corresponding to 0, 1, [CLS]), and 2 output embeddings (corresponding to the possible la- bels 1). For regularization mechanisms, typical choices were used: 0.1 for fattention, embedding, output} dropout; 10 weight decay. The Adam optimizer was instantiated with typical parameters 1 = 10 B1 0.9, B2 = 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 {100, 200, 400, 800} with standard emrors 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 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 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- settings (such as BERT pretraining), they require optimized GPU implementations to run a reasonable timeframe (days, as opposed to months).
C.2 LEARNING PARITIES
A natural question for further exploration is whether Transformers can learn other Boolean func- tions. 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' O'Dnnell (02T); furthermore, there are computational hardness barriers; see the works cited in the main paper for further context.
35

---

Under review as a conference paper at ICLR 2022
success probabilities
empirical sample complexity
200
1.0
180
Pr[val acc > 0.99)
0.8
sample size
160
0.6
120
100
T-100
80
0.2
T-200
T-400
60
0.0
T-800
x
3 x 10 context length
100
160
180
sample size
Figure 4: Enlarged plots from Figure
2 the main experimental validation of our theory.
Experimental setup. These experiments match the setting of the main AND experiments, except for few differences, which we enumerate here. Gradient-bard based training done on streaming online losses (so that there is no vaaiing/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 trend), pointing 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 unable to scale this phenomenon (i.e. learn any parities) context sizes above 50.
(rather than to measure any particular
D ADDITIONAL RELATED WORK
In this section, we discuss some additional related work.
Attention-based architectures. Building upon the success of modern attention-based architec- tures, 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 rele- vant to any architecture that 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 the representational power of self-attention architectures in regimes where the statistical guarantees are necessarily weak or vacuous (i.e. there Bhattamishra et al. 22200a b) establish universal function approximation and Kuring-completenes.. which have been known for previous architectures (Siegelmann Sontag 1995). Our work is significantly finer -her-egainedd 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 infinitewidth limit.
many functions the class). Dehghani et al. (2018); Yun al. (2019);
Likhosherstov et al. (2021) analyze the sparsity patterns representable by a self-attention head, with results superficially similar to ours: when the embedding dimension at least logarithmic in the context length, all sparse matrices can be approximately realized by attention head. However, their analysis is not about the capacity of the function class: quantifies over the input X, and holds the parameters (Wq. Wx....) 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 emror scales log(T).
36

---

Under review as a conference paper at ICLR 2022
Interpreting attention mixtures. A line of empirical work ("BERTology") has made progress understanding and interpreting state-of-the-arr Transformer language models by examining the acti- vations of their attention mechanisms (Clark et al., 2019 Tenney al. 2019 Rogers et al. (020). In some cases, these works have found instances which Trannformers 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 statistically meaningful way.
Other theoretical work on self-attention. Kerg et al. (220)) analyze self-attention and its ben- efit 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 results that analyze the statistical and representational properties of attention-based architectures, this work focuses the computational aspects of gradient-based methods 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 maximun-rewaar path from random walk observations) when using Transformers for offline reinforcement learning. [u et al. (2021 I) probe the transfer learning capabilities of pretrained Transformers, and consider some simple Boolean tasks. Our experimental protocol of learning sparse Boolean functions provides simple and fundamental setting for elucidating computational and statistical properties of sequence modeling architectures.
37


Reviews:
[{"review_id": "apDyPhTvMDZ", "reviewer": "Reviewer_k53Q", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:  \n1 This paper asks an essential question about the expressive power of attention and derives low-norm-based bounds for attention head, self-attention head, and the transformer model.\n\n2 The bound in this paper has logarithm dependence on sequence length. These results indicate a finding that Transformer selects sparse functions of context and validates them with empirical analysis.\n\n3 The theoretical analysis is rigorous with assumptions and detailed proofs.\n \nWeaknesses: \n\n1 The paper is not clear enough. The formulation is a little complicated, although the authors describe the well-known methods. \nOn page 5, the authors use the term tf-head before defining them; It would be better if the authors describe the two experimental tasks in detail;\n\n2 This paper could do more practical experiments to show the usage of the bound or the finding of sub-linear scaling law.\n\n3 The Related Work subsection has discussed many previous studies, but the paper does not compare or discuss previous methods in the theoretical analysis or in the empirical comparison.\n\nTypos:\n\" (for eg, linear functions\"-> e.g. linear functions, page 5\n\n", "summary_of_the_paper": "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. \n", "summary_of_the_review": "I like the contribution of generalization bounds for Attention and Transformer models. But there are some weaknesses. Therefore, my recommendation is marginal above the borderline.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZzeNBJYzQ4o", "reviewer": "Reviewer_YFzD", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n- The paper is very well written.\n- The paper tackles an important problem of studying the inductive bias of attention module.\n- The idea of sparse variable creating is interesting and relevant to the larger deep learning community in general. \n\nCons:\n\n- The experimental results are only on a synthetic task. It would be useful to study it on more complex tasks (as mentioned in the future work section of the paper). \n\nRelated Work:\n\n- The basic idea involves the  sparsity of functions learned by attention systems, and more precisely attention heads as found in Transformers. The analogy to \"attention heads\" as functions in typed argument programming languages and representation of different \"positions\" as different variables has been made before such that different heads only need to attend to sparse subset of positions [1, 2].\n\n- [1] RIMs, https://arxiv.org/abs/1909.10893\n- [2] NPS, https://arxiv.org/abs/2103.01937 ", "summary_of_the_paper": "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. \n", "summary_of_the_review": "The paper presents a theoretical analysis of the inductive bias of self-attention models, showing that the transformer architecture can learn sparse functions with some modifications. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RJL5UKo1rY0", "reviewer": "Reviewer_yc2n", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. The writing might need some improvements. The notation m and T were initially very hard to understand, until to the final page, I realize m denotes sample size while T denotes context length.\n\n2. At Corollary 4.4, the paper mentioned \"...highlighting the inductive bias of the transformer towards selecting sparse functions of the context\" and then later in section 5, \"The theory predicts the sample complexity of learning a bounded-norm transformer should match this scaling in T\". This is very interesting to me, but not sure how it is verified in Fig 2. There are some log-scaled correlation in the middle figure, but don't know what it says. For instance, the sample complexity should match the scaling of T, but for what goal?\n\n3. Based on the main takeaway I learnt from this paper, I wonder if it means, imagine in a coreference resolution task, the mention coreference is super dense while the context length is short. This could be challenging for a transformer to handle. But what if keeping the same amount of coreference links but synthetically increase the context size, would the task become easier? I imagine this can be a more real task (even though it might still be synthetic, but more realistic than a dataset of boolean functions) to experiment with. Furthermore, having a realistic data is good to show if the main conclusion of this paper is valid and the bound is a tight one in real application. Otherwise, the logarithmic dependency might not be established.\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "I think this paper proposes an interesting conclusion from its derivation of the theoretical bound on self-attention modules. The amount of contribution on the theoretical part seems substantial. But the verification part is weak. I don't feel I am confident to eat the takeaway from what I see in the experiment section.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2I6AsqACjah", "reviewer": "Reviewer_8Th6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper adresses an important gap in theoretical understanding of attention-based systems that are widely used in state of the art models. It is very clearly written, despite the complexity of the subject matter and the difficulties inherent to the notation of attention methods. The motivation of the main findings is clear and timely, and I believe the general community will be interested and benefit from this contribution. Although I am not an expect in all the tools used in the derivations of the main results, I believe the work is correct, novel, and significantly contributes to our understanding of Transformer-like systems. \n\nI believe there are two main points that could be improved to further strengthen this contribution.\n\n1- Sparcity:\nOne of the central findings of this paper involves the sparsity of functions learned by attention systems, and more precisely attention heads as found in Transformers. Indeed, on page 6 immediately following Corollary 4.4, the authors write \n\"Our bounds have a logarithmic dependence on T highlighting the inductive bias of the transformer towards selecting sparse functions of the context\"\nI think the link between the bounded covering numbers and representation sparsity is not clearly fleshed out and could be unpacked a little further to appeal to a broader audience. Two things could be done. \nFirst, for the interested non-expert, why does logarithmic dependence of bounds governing covering numbers lead to *sparse* functions ? Some context and details here would be welcome. \nSecond, most would already know informally that the softmax operation in attention weight computation leads to sparsity. This proof formalizes this in some way, and a discussion about the important role of the \"Norm\" component in the bounds, which is crucial, could be especially enlightening. Given these novel bounds, what would be the impact of using different Norm operations ? (e.g. Gumbel softmax v.s. softmax). Furthermore, not only is the softmax operation key in this result, it is also the major bottleneck for computational complexity, and a lot of work has been devoted to circumvent this problem (e.g. approximation of softmax by random kernels). How do the proposed bounds shed light on these issues ?\n\n2- Experiments:\nI understand that his is a theoretical contribution and that experiments play an supporting role in this paper. Nevertheless, the boolean functions experiments are Taylor-made to fit the bounds presented, and do not necessarily address the generalization properties associated with these bounds. A serious discussion about the type of experiments that could be carried out at scale by STOA systems, that would confirm the theoretical results, would strengthen the message. Moreover, two minor issues with regards to the experiment section remain. (1) While the finding of \"sub-linear\" scaling of sample complexity is consistent with scaling rates found analytically, the experimental results are far from confirming the same scaling. Sublinear can mean many things and is somewhat loose, and the experiments are noisy in that regard. A discussion about this limitation and what could be done to improve them would be welcome. (2) It appears this section might have been shortened or was added hastily as the main figure (Fig 2) has an incomplete caption (no description of the Left or center panel), and it appeals to the \"{overfitting, correct} regimes\" which are not discussed in the text (unless I missed it).\n\nSome minor issues to correct:\n\n* In definition 3.1, the Norm function definition involves a \\Delta, I'm not sure what this represents. Is this a typo ? I might be mistaken but I don't recognize this as standard notation.\n\n* In the proof overview of Thm 4, the second sentence is misswritten: \"This property allows us to construct the cover by into covering...\"\n\n* The paragraph following Lemma 4.3 also suffers the same faith: \"The key observation here is that the attention part of the network is computes using Norm...\"\n\n* In the description of the first numerical experiment in Section 5, it may be be good to explicitly refer to the appendix section containing more details as it is not clear how the task is implemented from the short description.\n\n*In the Related work section, Theory for attention models, the authors might consider citing recent work directions establishing gradient norm bounds for (self-)attentive systems. e.g. \"Untangling tradeoffs between recurrence and self-attention in artificial neural networks. Kerg et al., NeurIPS 2020\"", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper is well-written, correct, and provides results that are significant and advance our understanding of attention mechanisms. I recommend publication. The reason I do not give it the highest score is because of two slight weaknesses: (1) lack of interpretation of the obtained bounds with respect to choices of Norm operation (2) Simplistic and only mildly convincing experiments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a theoretical analysis of self-attention modules, using Lipschitz conditions.

It suffers from two main weaknesses: the clarity of the presentation, and the weak experimental section.
Decision: Reject