Inductive Biases and Variable Creation in Self-Attention Mechanisms

1 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 unsupervised representation learning . Starting with large-scale natural language processing ( Vaswani et al. , 2017 ) , self-attention has enjoyed surprising empirical successes in numerous and diverse modalities of data . In many of these settings , self-attention has supplanted traditional recurrent and convolutional architectures , which are understood to incorporate inductive biases about temporal and translational invariances in the data . Self-attention models discard these functional forms , in favor of directly and globally modeling long-range interactions within the input context . The proliferation of self-attention raises countless mysteries for theorists and empiricists . One fundamental question concerns its statistical properties : How should we think about the inductive biases of self-attention models ? More specifically , we can ask : Which functions do self-attention blocks prefer to represent ? How many ( approximately ) distinct functions can they represent ? To this end , this work initiates an analysis of the statistical foundations of self-attention , as it is used in today ’ s state-of-the-art models . Our main technical contribution is a classical norm-based generalization bound for a Transformer network , which can be extended to related and future architectures via a modular abstraction of attention mechanisms . In particular , the capacity ( in terms of log covering number ) of the function class of bounded weight self-attention heads ( and Transformers ) grows only logarithmically in the context length , which provides theoretical justification for the empirical observation that attention models can learn long-term dependencies without overfitting . Next , we show that bounded-norm self-attention heads are capable of representing wide classes of sparse functions . This representational capacity result , combined with the generalization results , provides a partial theoretical explanation for the observed sparsity bias of attention models , which we term sparse variable creation . We accompany this analysis with an experimental study of the sample complexity needed by Transformers to learn sparse Boolean conjunctions , and verify the sample complexity scaling law predicted by the theory in this clean synthetic setting . We discuss how to extend and repurpose this experimental protocol of benchmarking long-context sequence models on synthetic “ cryptographic ” tasks . We find that Transformers trained with gradient descent can learn sparse parities with noise , which may be of independent interest , exposing the empirical study of Transformers to the rich theory established around this problem . 1.1 RELATED WORK . The direct precursors to modern self-attention architectures were recurrent and convolutional networks augmented with attention mechanisms ( Bahdanau et al. , 2014 ; Luong et al. , 2015 ; Xu et al. , 2015 ) . Landmark work by Vaswani et al . ( 2017 ) demonstrated significantly improved machine translation models via self-attention only ; autoregressive language models followed shortly ( Liu et al. , 2018 ; Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) , as well as self-supervised representation learning ( Devlin et al. , 2018 ) . More recently , self-attention has demonstrated promise in computer vision ( Dosovitskiy et al. , 2020 ) , protein folding ( Jumper et al. , 2021 ) , theorem proving ( Polu & Sutskever , 2020 ) , program synthesis ( Chen et al. , 2021b ) , and reinforcement learning ( Chen et al. , 2021a ; Janner et al. , 2021 ) . Norm-based generalization bounds . There is a vast body of literature dedicated to establishing statistical guarantees for neural networks , including VC-dimension and shattering bounds ( dating back to Anthony et al . ( 1999 ) ) . In recent years , generalization bounds have been established for various architectures under norm bounds including ( Bartlett et al. , 2017 ; Neyshabur et al. , 2015 ; 2017 ; Golowich et al. , 2018 ; Long & Sedghi , 2019 ; Chen et al. , 2019 ) using covering-based arguments ; Jiang et al . ( 2019 ) provide an extensive empirical study of how well these bounds predict generalization in practice . Our work complements these bounds by establishing norm-based bounds for attention models . Our main results rely on a novel reduction to the ` ∞ covering number bound for linear function classes given by Zhang ( 2002 ) . Theory for attention models . Our work complements various existing theoretical perspectives on attention-based models . Vuckovic et al . ( 2020 ) formulate a dynamical system abstraction of attention layers , arriving at similar Lipschitz constant calculations to ours ( which are coarser-grained , since they focus on contractivity and stability rather than finite-sample statistical guarantees ) . Zhang et al . ( 2019 ) ; Snell et al . ( 2021 ) study idealizations of optimization problems for self-attention heads . Wei et al . ( 2021 ) propose a definition of statistically meaningful approximation of function classes that ties statistical learnability with expressivity , and show that Boolean circuits can be SM-approximated by Transformers with a sample complexity bound that depends mildly on circuit depth ( rather than context size ) , using a margin amplification procedure.1 See Appendix D for a broader survey of related work , including universal function approximation-style results ( which ignore statistical considerations ) . 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 m refers to the number of samples ( input sequences ) in a data set . ‖ · ‖2 denotes the spectral norm for matrices , and ‖ · ‖p , q denotes the ( p , q ) matrix norm where the p-norm is over columns and q-norm over rows . ‖ · ‖p denotes the ` p norm for vectors . For ` 2-norm , we drop the subscript . B is generally used to quantify bounds on norms of matrices and L for Lipschitz constants . ∆n−1 denotes the simplex of dimension n , that is , ∆n−1 : = { x ∈ Rn : x ≥ 0 , ‖x‖1 = 1 } . Covering numbers and uniform generalization bounds . Our main technical contribution is a generalization bound arising from carefully counting the number of functions representable by a Transformer . This main complexity notion we use is covering number . We will use the following definition of∞-norm covering number adapted from Zhang ( 2002 ) : Definition 2.1 ( Covering number ) . For a given class of vector-valued functions F , the covering number N∞ ( F ; ε ; { z ( i ) } mi=1 ; ‖ · ‖ ) is the smallest size of a collection ( a cover ) C ⊂ F such that ∀f ∈ F , ∃f̂ ∈ C satisfying maxi ‖f ( z ( i ) ) − f̂ ( z ( i ) ) ‖ ≤ ε . Further , define N∞ ( F , ε , m , ‖ · ‖ ) = supz ( 1 ) ... z ( m ) N∞ ( F ; ε ; z ( 1 ) , . . . , z ( m ) , ‖ · ‖ ) . If F is real-valued ( instead of vector-valued ) , we drop the norm from the notation . Furthermore for functions parameterized by a set of parameters Θ , we exploit the notation to replace F by Θ . 1Quoting the discussion following Theorem 4.1 of ( Wei et al. , 2021 ) : “ The correct norm-based Rademacher complexity bound to use for Transformers is unclear. ” Recall from Zhang ( 2002 ) that for the class of linear functions , Flin = { x 7→ w · x : w ∈ Rd , ‖w‖s ≤ BW } , we have the covering number bound of N∞ ( F ; ε ; { x ( i ) } mi=1 ) ≤ O ( B2XB 2 W /ε 2 · log ( BXBWm/ε ) , where ‖x ( i ) ‖ ≤ BX for i ∈ [ m ] . Importantly , note that the covering number has a mild dependence on m , only logarithmic ; this logarithmic dependence on m will be helpful when we turn our analysis to the capacity of attention mechanisms . Generalization bounds . This work focuses on providing log-covering number bounds , which determine the generalization error via standard Rademacher complexity and chaining arguments . The following lemma relates these quantities ; we refer the reader to Appendix A.1 for more details . Lemma 2.2 ( Generalization bound via covering number ) . Let D be a distribution over X × R and let ` : R × R be a b-bounded loss function that is L-Lipschitz in its first argument . For a given function class F and f ∈ F , let risk ( f ; D ) : = E ( x , y ) ∼D [ ` ( f ( x ) , y ) ] and r̂isk ( f ; ( z ( i ) , y ( i ) ) mi=1 ) : = 1m ∑m i=1 ` ( f ( z ( i ) ) , y ( i ) ) . Suppose F satisfies |f | ≤ A for all f ∈ F and logN∞ ( F ; ε ; x ( 1 ) , . . . , x ( m ) ) ≤ CF/ε2 for all for all x ( 1 ) , . . . , x ( m ) ∈ Xm . Then for any δ > 0 , with probability at least 1− δ , simultaneously for all f ∈ F and some constant c > 0 , ∣∣∣risk ( f ; D ) − r̂isk ( f ; ( x ( i ) , y ( i ) ) mi=1 ) ∣∣∣ ≤ 4cL√CFm ( 1 + log ( A√m/CF ) ) + 2b √ log ( 1/δ ) 2m . 3 ABSTRACTIONS OF ATTENTION AND SELF-ATTENTION . Attention is not straightforward to define — unlike other architectural components such as convolutions and residual connections , it is a broader model design principle , with numerous variations possible . In this section , we present an abstraction of attention mechanisms , guided by the different manifestations discussed in ( Luong et al. , 2015 ) , with the goal of enabling a more unified and modular statistical analysis . Subsequently , we show how to represent the Transformer ( the predominant attention-based architecture ) as a special case of this formulation . Our goal is not necessarily an all encompassing formulation of attention mechanisms , but rather an abstraction that helps to guide general design principles and theoretical analysis . 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 a learned function of global interactions ( see Figure 1 , left ) . To this end , we define an attention head as a function which maps a context X of T inputs { xt ∈ X } Tt=1 ( e.g . the tokens in a sentence , pixels in an image , or intermediate activations in a neural network ) and a context z ∈ Z to an output y ∈ Y . In this work , we will exclusively consider X , Y , Z to be Rd ; d is the embedding dimension , and we define the matrix of inputs X = [ x1x2 . . . xT ] > ∈ RT×d . An attention head uses z to select the inputs in X to which the output y will “ attend ” , which we formalize below : Definition 3.1 ( Attention head ) . An attention head is a function f : X → Y , specified by an “ alignment score ” function Score : X × Z → R parameterized by θs ∈ Θs , normalization function Norm : RT → ∆T−1 , and position-wise transformations φin : X → V , φout : V → Y parameter- ized by θin ∈ Θin and θout ∈ Θout . The output of an attention head on input X ∈ X T , z ∈ Z is y = φout ( T∑ t=1 [ Norm ( Score ( x1 , z ; θs ) , . . . , Score ( xT , z ; θs ) ) ] t φin ( xt ; θin ) ; θout ) = φout ( φin ( X ; θin ) > Norm ( Score ( x1 , z ; θs ) , . . . , Score ( xT , z ; θs ) ) ; θout ) where φin ( X ; θ ) = [ φin ( x1 ; θ ) . . . φin ( xT ; θ ) ] > denotes the row-wise application of φin . The above definition corresponds to leftmost diagram in Figure 1 . Here , V is a vector space of input representations “ mixed ” by the normalized alignment scores ; in this work , we will set V = Rk . A function class of attention heads is induced by specifying parameter classes for { Θs , Θin , Θout } .

Summary: 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.