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