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 looks at one specific aspect of transformer architectures - the normalization step that constrains the attention vectors in a probability simplex. The authors argue that this restriction has some limitations including limited information flow, makes the models biased towards local information at initialization and increases the sensitive to hyperparameters. To mitigate this, the authors propose an architecture that replaces the softmax in the attention mask with a simple normalization. Then, with several experiments on a synthetic dataset, the authors demonstrate that the proposed architectural change leads to models that are robust to hyperparameters, and less sensitive to biases in the data.