Normalized Attention Without Probability Cage

1 INTRODUCTION The concept of neural attention ( Graves , 2013 ; Bahdanau et al. , 2015 ) has sparked a number of architectural breakthroughs . The Transformer architecture ( Vaswani et al. , 2017 ) successfully deploys multi-headed self-attention in several consecutive layers for natural language processing ( NLP ) – a solution that has become increasingly popular ( Vaswani et al. , 2017 ; Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ; Devlin et al. , 2019 ; Yang et al. , 2019 ; Raffel et al. , 2019 ; Liu et al. , 2019 ) . Apart from NLP , self-attention has shown success in applications ranging from image classification ( Parmar et al. , 2019 ; Dosovitskiy et al. , 2021 ) to generative adversarial networks ( Zhang et al. , 2019 ) to reinforcement learning ( Bram et al. , 2019 ; Parisotto et al. , 2019 ; Loynd et al. , 2020 ) . The attention architecture choice is thereby often based on one , if not both , of the following arguments : ( 1 ) Attention helps with the information flow by providing more direct , dynamic links between inputs and outputs . ( 2 ) Attention is directly interpretable as one can investigate the percentages to which different inputs are “ attended ” to . However , this second argument has been challenged recently , as several works show that attention weights do not directly correlate with predictions ( Jain & Wallace , 2019 ; Wiegreffe & Pinter , 2019 ; Pruthi et al. , 2019 ; Brunner et al. , 2020 ; Pascual et al. , 2020 ) . With interpretability already in dispute , we focus on the other argument . Specifically , we question whether attention is the best way to route information and argue that apart from the proven success , softmax attention might not always be the best option . In this work , we investigate the theoretical implications of constraining the attention weights to the probability simplex , and propose an unconstrained alternative based on normalization . We show that the popular Transformer architecture has an innate bias towards local information at initialization and showcase implications of this architectural bias on biases in the data . We investigate different architectures on an abstract level in numerous experiments and demonstrate the advantage of unconstrained attention . In particular , we improve robustness to hyperparameters and show the general applicability of attention based architectures as compared to sum and max pooling . 2 BACKGROUND AND RELATED WORK . Many data processing tasks can be addressed by representing the input as a set or sequence of discrete tokens , e.g. , words in a sentence , frames in a video or points in a point cloud . As a general formulation , we represent each input token by a vector xi ∈ Rd for i ∈ { 1 , . . . , N } , where N is the sequence length and d is the dimensionality of each token . For ease of notation we use the word “ sequence ” throughout , but note that all architectures discussed are also applicable to unordered sequences , i.e. , sets of tokens . Multi-headed dot-product self-attention is a fundamental building block of the Transformer architecture ( Vaswani et al. , 2017 ) . It allows for information exchange between different tokens of the input sequence . More formally , for each attention head m , the input vectors xi are projected by an affine transformation to a query qim , key k i m and value vector vim . The dimensionality of these vectors is chosen as dh = d M , where M is the number of attention heads . The query and key vectors are used for a pairwise dot product , scaled by the square root of the head dimension dh , to form the attention logits li , jm = < qim , k j m > √ dh and atten- tion vectors aim = softmax ( [ l i,1 m , . . . , l i , N m ] ) , where softmax refers to the normalized exponential function softmax ( x ) j = exp ( x j ) ∑ k exp ( x k ) commonly used to project vectors to the probability simplex SP = { aim|ai , jm ≥ 0 ∀j and ∑ j a i , j m = 1 } . The output oim of each attention head m is then given by a weighted sum of all value vectors oim = ∑ j a i , j m · vjm . These attention head outputs are concatenated and mixed trough an affine transformation to form the attention output in the Transformer architecture ( Vaswani et al. , 2017 ) . In this work , we investigate whether constraining the attention vectors aim into the probability simplex through the softmax function is sensible from a geometric perspective . We contrast the multi-head self-attention architecture to attention-inspired architectures without softmax ( discussed in Section 4 ) as well as simpler aggregation methods commonly used . Specifically , while Yun et al . ( 2020 ) show that Transformers are universal sequence-to-sequence function approximators , we question the practical necessity of an attention architecture , when sum pooling ( Zaheer et al. , 2017 ) already provides general function approximation capabilities ( Zaheer et al. , 2017 ; Xu et al. , 2019 ; Segol & Lipman , 2020 ) . Further , we compare to max pooling , a common aggregator choice that has shown good empirical success ( Nagi et al. , 2011 ; Zaheer et al. , 2017 ; Velikovi et al. , 2020 ) . Several works have proposed architectural changes to the Transformer ( Yun et al. , 2020 ; Wu et al. , 2019 ; Lan et al. , 2020 ; Dehghani et al. , 2019 ; Fan et al. , 2020 ; Wu * et al. , 2020 ; Press et al. , 2019 ; Bachlechner et al. , 2020 ) , however , these do not alter the multi-head self-attention . Another direction of research has focused on improving the computational efficiency of Transformers ( Wang et al. , 2020 ; Lee et al. , 2019 ; Katharopoulos et al. , 2020 ; Zaheer et al. , 2020 ; Choromanski et al. , 2021 ) . However , these either retain the softmax ( Lee et al. , 2019 ; Wang et al. , 2020 ; Zaheer et al. , 2020 ) or try to stay close to it ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2021 ; Lu et al. , 2021 ) . We refer to Tay et al . ( 2020 ) for a more extensive overview on works in this direction . Our focus is different : we ask whether the softmax is necessary at all . The Synthesizer ( Tay et al. , 2021 ) and MLP-Mixer ( Tolstikhin et al. , 2021 ) show that an aggregation with learned aggregation weights can already lead to good results . However , these aggregation schemes make the models dependent on the input sequence length – a trait we explicitly seek to avoid . Recently , a few works provide complementary results that motivate architectures without softmax . In particular , Tay et al . ( 2019 ) introduce an attention mechanism that allows for negative attention values . Schlag et al . ( 2021 ) link the head dimension dh to a corresponding limitation in the memorization capacity of linearized models . Shen et al . ( 2021 ) propose an architecture similar to our NON architecture ( see below ) for improved computational complexity . Cao ( 2021 ) also explore an architecture similar to our NON architecture and links it to a Petrov-Galerkin projection for operator approximation . They empirically find in an ablation that a layer normalization on the keys and queries separately improves convergence . Our work goes beyond this , showing the correlated effects between hyperparameters and architecture choices . In contrast to all mentioned works , we provide a viewpoint from the gradient dynamics perspective and investigate the effect of skewed data distributions . 3 IMPLICATIONS OF SOFTMAX ATTENTION . To start our discussion , we highlight an observation that follows directly from attention vectors ai being constrained to the probability simplex SP : Outputs oim are convex combinations of the vectors vim : This in itself has interesting implications . First and foremost , we note that a convex combination of vectors vim can not yield any vector outside the convex hull spanned by the value vectors vim . An illustration of this output “ cage ” is given in Figure 1 ( left ) . We conjecture that this constraint limits the hypothesis space of the optimization , making it more difficult for gradient descent to find a good solution . This conjecture is supported by our experimental results showing an increased robustness to hyperparameter choices when the constraint is removed . Further , we note the following from a theoretical perspective : No convex combination can in itself represent XOR : A formal proof is given in Appendix A . Note that this implication highlights an inability to represent non-linearity . While the binary exclusive OR can be represented in architectures with multiple heads and layers , the insight further underlines our argument : An aggregation with weights constrained to the probability simplex is restrictive . The convex hull view point also leads to an additional insight : Transformers are not aggregation size independent : Specifically , Transformers focus on local information at initialization ; and this focus depends on the sequence length N . To see this , consider the embeddings ei after the first residual connection given by ei = xi + W [ ∑ j a i , j 1 · v j 1 ∣∣∣ . . . ∣∣∣∑j ai , jM · vjM ] + b , where [ ·|· ] denotes concatenation and W and b represent the parameters of the affine transformation that mixes the M attention head outputs to form the attention layer output . Our aim is to show how much the embedding ei is influenced by the local information xi relative to the context information { xj |j 6= i } . We first note that the contribution of context information depends on the initialization of W , where a typical initialization in language models yields a contractive projection.1 This favors the residual connection , i.e. , local information . However , even if we consider W as scale preserving , we note that the magnitudes of the attention head outputs om are upper bounded by the magnitudes of the value vectors vm as a result of the convex hull . Moreover , attention logits are normally close to 0 at initialization ( to have the softmax in the unsaturated region ) . This yields attention to be close to mean aggregation as oim ≈ v̄m = 1 N ∑ j v j m. Considering v j m to be distributed with zero-mean and variance σ2v at initialization , we note that the expected magnitude and standard deviation of v̄m scales proportional to σv√N . This means that the fraction of context information in ei is dependent on the sequence length N and , counterintuitively , is smaller for longer sequences . Specifically , with the constant contribution of the residual connection , Transformers focus more on local information ; and more so in longer sequences than in shorter sequences . For reference , we visualize the dependence of om on N at initialization for different aggregators in Figure 2 . We note that while an architectural bias towards local information might be beneficial in some applications , the implicit dependence on aggregation size N is questionable . Finally , we note that : A softmax can easily saturate , which leads to vanishing gradients : To show this , we look at the term g = δa i , j m δli , jm which gives a gradient factor on the back-propagation path through the softmax . Specifically , we note that g ∈ ( 0 , 0.25 ] and that the range of li , jm for which g > 0.01 is less than 9.2 . This means that any gradient update that changes li , jm by a constant of c > 9.2 , e.g . by updating the bias terms , will most certainly lead to vanishing gradients henceforth . This is particularly worrisome in the Transformer architecture , where gradients correlate with the input due to the multiplicative interactions in the attention . This means that a single large input can result in a parameter update that yields vanishing gradients on all following data points . See Appendix B for proofs . 1As an example , BERT ( Devlin et al. , 2019 ) initializes W with parameters drawn from a truncated normal distribution with standard deviation set to 0.02 . This would only be roughly scale preserving for d = 2500 .

Summary: This paper starts from the observation that current self-attention modules is sensitive to hyperparameter changing. The authors conjecture that this is due to the softmax operator in self-attention, and give some intuitive examples to support their hypothesis (eg., bias towards local information, gradient vanishing). To solve the problem, the authors propose to replace the softmax with a normalization operator. They show that normalized attention is more robust to hyperparameter in a synthetic dataset, as well as real-world settings.