Year: 2021

Abstract:
Despite the popularity of attention based architectures like Transformers, the geometrical implications of softmax-attention remain largely unexplored. In this work we highlight the limitations of constraining attention weights to the probability simplex and the resulting convex hull of value vectors. We show that Transformers are biased towards local information at initialization and sensitive to hyperparameters, contrast attention to maxand sum-pooling and show the performance implications of different architectures with respect to biases in the data. Finally, we propose to replace the softmax in self-attention with normalization, resulting in a generally applicable architecture that is robust to hyperparameters and biases in the data. We support our insights with empirical results from more than 30,000 trained models. Implementations are in the supplementary material.

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

heading: 2 BACKGROUND AND RELATED WORK
text: 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.

heading: 3 IMPLICATIONS OF SOFTMAX ATTENTION
text: 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 cannot 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.

heading: 4 NORMALIZED ATTENTION POOLING
text: Given the implications of softmax attention, one might seek alternatives. To investigate different architectural choices, we contrast the following architectures in an abstract empirical study:
Transformer Encoder (BERT): As a starting point, we replicate the encoder architecture presented by Vaswani et al. (2017) as described in the code release of Devlin et al. (2019).2 This architecture has been used extensively (Radford et al., 2018; Yang et al., 2019; Raffel et al., 2019; Liu et al., 2019; Clark et al., 2020). Each Transformer-layer consists of two submodules: a multi-head self-attention “layer” and a feed forward network. Both modules are encompassed by residual connections. The multi-head self-attention “layer” consists of a
projection to queries, keys and values, the attention mechanism as well as a mixing layer as described in Section 2. The feed forward network consists of two layers with a GELU (Hendrycks & Gimpel, 2016) non-linearity. Layer normalization (Ba et al., 2016) is applied between incoming and outgoing residual connections. Note that this gives a crucial distinction of this architecture: Embeddings are normalized after they are summed with the residual connection. This yields the implicit dependence on sequence length as discussed in Section 3. We train this architecture with learning rate warm-up and gradient norm clipping (cf. Devlin et al. (2019)).
Modified Transformer Encoder (MTE): To overcome the implicit dependence on sequence length, we modify the architecture by moving the layer normalizations and adding additional normalizations as shown in Figure 3. Note that this is different from the recently studied PreNorm (Parisotto et al., 2019; Nguyen & Salazar, 2019; Liu et al., 2020) that places the normalization before the attention mechanism. To further remove cofounding factors, we remove learning rate warm-up and gradient clipping, but keep a linearly decreasing learning rate schedule, taking Li et al. (2020) as reference. We provide an ablation of all modifications in Appendix D.1. All following architectures apply the same modifications. The resulting MTE architecture still projects attention weights to the probability simplex in the multi-head attention and is thereby limited to convex combinations of value vectors.
Normalized Attention Pooling (NAP): Given the success of online normalization during training - be it through batch- (Ioffe & Szegedy, 2015), layer- (Ba et al., 2016) , group- (Wu & He, 2020), instance- (Ulyanov et al., 2016) or weight-normalization (Salimans & Kingma, 2016) - our main proposal is to simply replace the softmax with a normalization:
aim = normalize([l i,1 m , . . . , l i,N m ]) with normalize(x) j = g · x j − µx σx + b
where µx = 1N ∑ j x j and σx = √ 1 N ∑ j(x
j − µx)2 are the mean and standard deviation of the corresponding input vector x, in our case the logit vector calculated through key-query dot products. Similar to layer normalization (Ba et al., 2016), we introduce gain and bias parameters g and b initialized to 1 and 0, respectively. However, while Ba et al. (2016) introduce gain and bias vectors, we only introduce scalar parameters and broadcast these over the sequence, as we want the architecture to be independent of the sequence length N . Note that while no convex combination can represent the logical XOR, a normalized weighting can – we give a proof in Appendix A.
No Online Logit Normalization (NON): To investigate whether a dynamic normalization of the attention logits is necessary, we also train a model where we use the logits li,jm directly as attention weights, i.e., oim = GELU( 1√ N ∑ j l i,j m · vjm). The factor 1√N is introduced such that the model has in expectation a constant contribution of context at initialization.
Simple Summation of Embeddings (sum): From a theoretical perspective, summation is sufficient for general function approximation (Zaheer et al., 2017; Xu et al., 2019; Segol & Lipman, 2020). Therefore, we investigate to simply replace attention with a sum-reduce-broadcast operation.
2https://github.com/google-research/bert
Max Pooling over Sequence Dimension (max): Similar to sum pooling, we can replace the attention sub-module with a simple max-reduce-broadcast operation over the sequence dimension. Note that max pooling is a powerful operation that yields a direct link to up to d different tokens.
We provide a schematic figure of all architectures in Appendix D. If not varied in a corresponding experiment, we default architecture hyperparameters to L = 2 Transformer-layers (consisting of an attention sub-module and feed forward sub-module each), M = 4 heads to calculate the logits (if applicable), d = 128 as model dimension and train on a total of 3200 batches of 32 example sequences each, using the Adam optimizer (Kingma & Ba, 2014). The hidden dimension of the feed forward sub-modules is 4 · d for the models BERT, MTE, NAP and NON. For the models sum and max we increase the feed forward hidden dimension to approximately match the parameter count.

heading: 5 EXPERIMENTS AND RESULTS
text: Our goal with this work is to provide an insight into the variety of performance implications that the architecture choices entail. We aim to provide these insights independent of any particular application, as these architectures can be applied to a variety of tasks – from NLP (Vaswani et al., 2017; Devlin et al., 2019) to graph neural networks (Velickovic et al., 2018; Yun et al., 2019; Xu et al., 2019; Velikovi et al., 2020) to reinforcement learning agents (Parisotto et al., 2019; Fang et al., 2019; Loynd et al., 2020). We therefore focus on carefully crafted synthetic tasks that (1) are general enough in that we can expect the insights to generalize to a large set of downstream tasks and (2) let us modify key aspects that are hidden in real world data sets, such as a bias towards a certain subtask. The focus on synthetic tasks also allows us to get a better grasp on the learning dynamics – the focus of this work – as we can train thousands of models in diverse hyperparameter combinations. To limit the influence of confounding hyperparameters, we generate new data points for every batch. This allows us to omit regularizations and their confounding factors. See Appendix D.2 for an in depth discussion of this setup as well as regularized ablations.

heading: 5.1 ARGMIN-FIRST-ARGMAX CASE DISTINCTION TASK
text: As a first task, we train the networks to pin-point a specific, input dependent token. Note that the ability to pin-point a specific token is an abstract task relevant to NLP (e.g., question answering or co-reference resolution), graph neural networks (e.g., finding the next hop in a shortest path) as well as reinforcement learning (e.g., action credit assignment). We consider an input pipeline where tokens from a fixed integer-vocabulary are translated to a randomly initialized embedding. To the embedded tokens, a (also randomly initialized) positional embedding is added to provide position-relative information. The sequence of tokens is then processed by several architecture dependent Transformer-layers (cf. Section 4). Finally, each contextualized embedding is projected to a single output. A softmax-crossentropy loss is applied over the sequence dimension. See Figure 4 (middle) for a visualization. To make the task input dependent, we generate the data as given in the pseudo code in Figure 4 (left). Note that the argmin and argmax make this task quite challenging from a learning perspective as the networks start from random embeddings which do not provide any ordering information. Which embeddings correspond to bigger integers and which to smaller
integers has to be inferred during training. Further, the case distinction in this task lets us tweak the data bias towards each sub-task. Specifically, we consider a vocabulary size of S = 100 integers (0-99) and uniformly random sampled sequences of N = 128 tokens in length. This leads to a bias as pargmin = 1 − (1 − 1S )
N ≈ 72.4% of data points require the network to pin-point the minimum in the input sequence, pfirst ≈ 20.1% require to pin-point the first token of the sequence and pargmax ≈ 7.5% require to pin-point the maximum in the input.
5.1.1 VARYING MODEL DIMENSION d
As a first investigation, we are interested in how varying the model dimension d influences the architectures ability to learn the given task. For this, we train each of the architectures for each of the model dimensions d ∈ {8, 16, 32, 64, 128, 256, 512, 1024} using 10 different learning rates and 5 random seeds for each hyperparameter combination. As we want to base our insights on as many results as possible, we derive a novel, human friendly visualization of results. Figure 5 shows the first results as follows: The outcome of each hyperparameter combination is reported as an RGB pixel in the plot, where the R (red) value corresponds to the accuracy of the worst performing random seed, the G (green) value corresponds to the average over the random seeds and the B (blue) value corresponds to the best performing random seed. For each value (R, G and B), the max over the course of training is taken. This assignment roughly translates as follows: The brighter, the better - brighter pixels correspond to higher min-, mean- and max-accuracy. Blue/turquoise pixels highlight a large performance variation across random seeds and black/grey pixels correspond to hyperparameter combinations where none of the random seeds could solve the task. See Appendix L for figures with color channels split. Given that all architectures are applicable to sequences of any length, we investigate how the architectures generalize to sequences of different length. Specifically, we validated each of the models trained on sequences of length N = 128 after every 100 batches on 32 batches with sequences of half the length (N = 64). The condensed results in Figure 5 give rise to the following observations: (1) Most models have some hyper-parameter combinations that learn the task well (white pixels). (2) The optimal learning rate depends on the model size, especially in the BERT architecture. This has profound implications for hyperparameter optimization: Tuning hyperparameters independent of each other might lead to sub-optimal results. (3) Models with probability simplex limitations (BERT and MTE) work for a smaller range of hyperparameters. (4) Our NAP architecture seems to be the most robust to this generalization. We provide case learning curves and additional results in Appendices E and F. Specifically, to be sure, we repeat this experiment with L = 4 and L = 6 Transformer layers, but the results stay the same (see Appendix F.1). We also provide the results of an additional model with learned aggregation weights in Appendix H.

heading: 5.1.2 CASE ACCURACY UNDER VARYING DATA BIASES
text: As a next experiment we reset the model dimension to d = 128 and investigate the models under varying data biases. We do so in two distinct ways: First, we vary the sequence length
N ∈ {4, 8, 16, 32, 64, 128, 256, 512}. Note that this implicitly varies the biases pargmin, pfirst and pargmax in the data. Second, we fix the sequence length at N = 128 and vary pargmin by explicitly adding or removing corresponding data points. We report the case specific accuracies in Figure 6 as follows: After every 100 batches, we validate the models on 1000 examples per case. Reported is the best accurracy over the course of training in form of pixel value with R (red) corresponding to the argmin-case accuracy, G (green) corresponding to the first-case accuracy and B (blue) corresponding to the argmax-case accuracy. As a consequence, white pixels correspond to all cases learned and yellow pixels correspond to the argmin- and first-case learned. We make the following observations: (1) If the learning rate is too low, models tend to focus on the majority case (indicated in a shift from blue to red as the bias shifts from the argmax- to the argmin-case with increasing sequence length N and the shift from green to red for fixed N ). (2) If the learning rate is too high, the BERT architecture tends to focus on the first-case. We believe this is due to the architectural bias towards local information as discussed in Section 3. Note that the first-case can be solved by relying on the local positional embedding. (3) Only the NAP and max architecture manage to learn all three cases from highly biased data (see N = 256 and pargmin = 96.5%). This shows that the NAP architecture is not only robust to hyperparameters but also robust to biases in the data. In Appendix F.2 and F.3 we provide further experiments investigating different batch sizes and initialization scales. In Appendix F.4 we investigate the local vs. global aggregation focus further.

heading: 5.2 FIRST TOKEN OUTPUT
text: The task so far requires the architectures to learn an information flow between tokens to distinguish the case and decide per token, whether it is the token that is looked for or not. Now we investigate, whether all this information can also be aggregated into a single token. We therefore modify the architecture output slightly in that we only take the contextualized embedding of the first token and project from it to a vector of size N (see example on the right in Figure 4). Note that this task set-up is harder and can highlight bottlenecks in the information flow across tokens. We fix the sequence lenght to N = 128 and again vary the model dimension d. We report the case specific mean accuracies in Figure 7. Min-, mean- and max-overall-accuracies are given in Appendix F.5. We observe: (1) All architectures learn for (almost) all combinations the now close to trivial firstcase, even though it is not the majority case. (2) The sum pooling architecture does not learn any of the other cases. (3) Only NAP and max learn all three cases in some hyperparameter combinations.
The worse performance of NON highlights the advantage of online normalization of the logits. While the softmax provides some form of online normalization, we hypothesize that the worse performance of MTE and BERT in this task stems from an information bottleneck induced by the probability simplex limitations. To test this hypothesis further, we vary the number of attention heads M . The results are shown in Figure 9. We observe that increasing the number of heads helps the MTE and BERT architecture, supporting our hypothesis. Note however, that MTE and BERT are still outperformed significantly by NAP. In Appendix F.6 we provide a further experiment, varying the depth up to L = 64. Results are complementary.

heading: 5.3 MODE FINDING TASK
text: Given the results so far, one could conclude that max is the best choice due to its simplicity. Note however, that max has an architectural prior that is in line with the underlying task of finding the maximum or minimum of the sequence. To study the effect of architectural priors, we experiment on an additional task: Finding the mode/most common integer in the input sequence. Also this task has ties to NLP (e.g., sentiment analysis), graph neural networks (e.g., consensus/agreement) and reinforcement learning (e.g., count based exploration). Here we remove the positional embeddings, as this task can also be done on sets, and project from the contextualized embedding of the first token to a vector of dimension S (the vocabulary size) over which we apply the softmax-cross-
entropy loss. We keep N = 128 but reduce S to 10 to have meaningful modes. Ties are broken by taking the smallest integer of the ones with maximal occurrence. Results of varying the model dimension d are reported in Figure 8. We observe: (1) sum pooling works well on this task, as it has a suitable architecutral prior. (2) max pooling cannot learn the task, not even with a model dimension d = 1024 = 8 ·N . In Appendix G.1 we provide complementary results varying the vocabulary size. We refer an interested reader to Xu et al. (2020) for more on architecture-task alignment.

heading: 5.4 GENERALIZATION TO GRAPH NEURAL NETWORKS AND REINFORCEMENT LEARNING
text: NLP experiments normally require large amounts of data and large models – a setting which would make an investigation of the training dynamics under various hyperparameters prohibitively expensive. We therefore focus in our generalization study on the two other domains where attention based architectures have recently gained traction. Specifically, to see whether the the learning dynamics seen so far generalize, we perform two additional experiments: One investigating the architectures in a graph neural network (GNN) task and one investigating them in a reinforcement learning (RL) setup. Specifically, we train a transformer based GNN in the Protein-Protein-Interaction (PPI) graph node classification task (Zitnik & Leskovec, 2017) and investigate the different architectures as alterations on the working memory graph architecture for RL (Loynd et al., 2020), training on level 3 of the Baby-AI environment (Chevalier-Boisvert et al., 2019). The validation results under varying model dimension d are given in Figure 10. The results replicate several observations made so far. Specifically: (1) All attention based architectures can learn the task for some hyperparameter combinations. (2) There is a correlation between optimal learning rate and model dimension. (3) BERT is hyperparameter sensitive in comparison. See Appendices I and J for experiment details.

heading: 6 CONCLUSION
text: Taking all observations together, we conclude: Many recent works apply some sort of self-attention mechanism involving a softmax that projects the attention weights to the probability simplex. In this work we question the softmax in dot-product self-attention modules. Our theoretical investigation shows that softmax-attention outputs are constrained to the convex hull spanned by the value vectors. In our experiments we show that this can lead to an unwanted hyperparameter sensibility. We show that simpler architectures like max- and sum-pooling perform well when their architectural prior aligns with the underlying task. These architectures however fail in cases where the architectural prior is not suitable. As a solution, we propose to replace the softmax in attention through normalization. Our resulting normalized attention pooling (NAP) architecture is the only architecture of the 6 investigated that is robust in all tasks and setups, even if the data is heavily skewed towards easier sub-tasks. We hope that our work provides a stepping stone to examine architectures under varying biases in the data. Further, we see potential in exploring the correlated effects of hyperparameters. Please also refer to our broader impact statement in Appendix K.

heading: A REPRESENTATION POWER WITH RESPECT TO XOR
text: Lemma 1. No convex combination can represent the binary exclusive OR (XOR) function defined on binary inputs x1 ∈ {0, 1} and x2 ∈ {0, 1} by the indicator function as XOR(x1, x2) = 1x1 6=x2 .
Proof. Suppose there exist convex combination weights a1 and a2 with a1 + a2 = 1, such that a1 · x1 + a2 · x2 represents the XOR function. Plugging in x1 = x2 = 1 yields a1 · x1 + a2 · x2 = a1 + a2 = 1, which gives the contradiction.
Lemma 2. Given the two binary inputs x1 ∈ {0, 1} and x2 ∈ {0, 1}, there exists an affine mapping f : {0, 1}2 → R2, such that the normalized weighting given by
f1(x1, x2)− µf(x1,x2) σf(x1,x2) · x1 + f2(x1, x2)− µf(x1,x2) σf(x1,x2) · x2
is equivalent to the logical exclusive OR given by the indicator function asXOR(x1, x2) = 1x1 6=x2 .
Proof. For a vector l ∈ R2, the standard deviation σl can be simplified to
σl = √√√√1 2 ∑ i∈{1,2} (li − µl)2
= √√√√1 2 (( l1 − l1 + l2 2 )2 + ( l2 − l1 + l2 2 )2)
= 1
2 |l1 − l2|
and the normalization function reduces to
normalize(l) = [ l1 − µl σl , l2 − µl σl ]T =
[ l1 − l2 |l1 − l2| , l2 − l1 |l1 − l2| ]T
=  [1,−1]T if l1 > l2 [−1, 1]T if l1 < l2 undef. if l1 = l2
As an example, consider the affine mapping f(x) = l = [3x1 + 1, 2x2]T , which for x1 ∈ {0, 1} and x2 ∈ {0, 1} results in the function
3x1 + 1− 2x2 |3x1 + 1− 2x2| · x1 + −3x1 − 1 + 2x2 |3x1 + 1− 2x2| · x2 = 1x1 6=x2
We note that for a realization of such an affine mapping across tokens given the weight sharing constraints of the discussed architectures we would need x1 and x2 to be distinguishable for the mapping to keys and queries, e.g., through positional embeddings. This however does not invalidate our conclusion that normalized weighting is more expressive than softmax weighting, as we do not require the inputs that are weighted to be distinguishable.
B VANISHING GRADIENT ANALYSIS
For better readability we will use subscripts here to index sequence elements and omit the implicit reference to the head and vector dimension used in the main text. Lemma 3. The partial derivative g = δsoftmax(l)iδli is limited to the range g ∈ (0, 0.25].
Proof.
g = δsoftmax(l)i
δli = exp (li) ∑ j exp (lj)− exp (2 · li)(∑ j exp (lj)
)2 = ai · (1− ai) with ai = softmax(l)i Maximizing this function we get that the maximum of g = 0.25 is achieved for ai = 0.5. Noting that ai > 0, the claim follows.
To expand this further, we introduce the notion of α-bandwidth.
Definition 1. We call the range ∆x in which the absolute value of a gradient factor δf(x)δx exceeds a threshold α, i.e. ∣∣∣ δf(x)δx ∣∣∣ > α, its α-bandwidth BWα ( δf(x)δx ).
Please refer to Figure 11 for a visualization of this definition. Intuitively, the α-bandwidth gives a measure for the range of operation, where at least an α-fraction of the gradient gets backpropagated.
It also provides a measure of how fast a function saturates: Given x with δf(x)δx > α, updating x to x ′ = x ± BWα ( δf(x) δx ) puts x′ into the saturated region where δf(x ′)
δx′ ≤ α. A function with a small αbandwidth can therefore more easily saturate than a function with a larger α-bandwidth.
Lemma 4. The α-bandwidth of g = δsoftmax(l)iδli can be calculated as
BWα(g) = ln
( 1− 2α+ √ 1− 4α
1− 2α− √ 1− 4α
)
Proof. Substituting x = exp(li) and y = ∑ j 6=i exp(lj) we can write g as g(x) = x x+y · ( 1− xx+y ) .
Solving g(x) = α yields x = y2α ( 1− 2α± √ 1− 4α ) . Reversing the substitution of li = ln(x) and subtracting the larger from the smaller value, we get:
BWα(g) = ln ( y
2α
( 1− 2α+ √ 1− 4α )) − ln ( y 2α ( 1− 2α− √ 1− 4α )) = ln ( 1− 2α+ √ 1− 4α
1− 2α− √ 1− 4α
)
Corollary 1. The range in which g = δsoftmax(l)iδli > 0.01 is smaller than 9.2.
Proof. Calculating the α-bandwidth for α = 0.01 we get BW0.01(g) ≈ 9.170. Note that BWα(g) is independent of the length of the vector l as the remainder y = ∑ j 6=i exp(lj) cancels. For the self-attention in the Transformer this translates to the fact that BWα(g) is independent of the sequence length. This means that the Transformer architecture cannot adapt this range of operation according to the sequence length. Aggregated information from long sequences inevitably needs to be averaged or scaled to fall within this range.
We further note that the constant 0.01-bandwidth on its own is not the problem. In fact, a tanh-nonlinearity has a smaller 0.01-bandwidth. However, in the standard Transformer architecture, gradients are scaled dependent on the layer input. More specifically, dropping sequence indices to simplify notation, we have Lemma 5. For k = Wkx + bk, q = Wqx + bq , WTq Wk + WTk Wq invertible and ||x|| ≥ c1 we can find constants c2 > 0 and c3 > 0 such that∣∣∣∣∣∣∣∣ δlδx ∣∣∣∣∣∣∣∣ ≥ c2 · ||x|| − c3
Proof. Writing out l = <q,k>√ dh we get
l = 1√ dh (xTWTq Wkx + (b T kWq + b T qWk)x + b T q bk)
Using δx TAx δx = x T (A+AT ) and substituting gb = bTkWq + b T qWk we get
δl δx = 1 dh (xT (WTq Wk +W T k Wq) + gb)
Setting gx = xT (WTq Wk+W T k Wq) and solving for x we get (W T q Wk+W T k Wq) −1gTx = x. With this, we can bound ||gx|| with the operator norm on (WTq Wk + WTk Wq)−1. Specifically, if we set c2 as
c2 = 1√
dh · ||(WTq Wk +WTk Wq)−1||
we get 1√ dh · ||gx|| ≥ c2 · ||x||. We can then bound ∣∣∣∣ δl δx ∣∣∣∣ as∣∣∣∣∣∣∣∣ δlδx ∣∣∣∣∣∣∣∣ = 1√dh ||gx + gb|| ≥ 1√dh
∣∣∣∣||gx|| − ||gb||∣∣∣∣ Therefore, if we set c1 = 1c2· √ dh · ||gb|| we get ||gx|| ≥ c2 ·
√ dh · ||x|| ≥ ||gb|| and hence∣∣∣∣∣∣∣∣ δlδx
∣∣∣∣∣∣∣∣ ≥ 1√dh (||gx|| − ||gb||) ≥ c2 · ||x|| − c3 for c3 = 1√dh · ||gb||.
Lemma 5 shows that the magnitude of the gradient factor scales with the magnitude of the input. Therefore, a single large input can yield a correspondingly large parameter update. This in turn can put the softmax easily into the saturated region, since the softmax only has a constant α-bandwidth. Once saturated, the softmax leads to vanishing gradients henceforth.

heading: C SEQUENCE LENGTH DEPENDENT FOCUS
text: For Figure 2 we sample 16’384 value, key and query vectors of dimension dh = 128 per sequence length N ∈ {1, 2, 4, 8, 16, 32, 64, 128, 512, 1024, 2048} from a normal Gaussian N (0, Idh) - Idh being the dh-dimensional identity matrix. We split the samples to form the sequences and calculate the corresponding output vectors oi for i ∈ {1, . . . , N}. Here, the softmax attention outputs are calculated as described in Section 2, while the mean-, sum- and max-outputs are calculated as mean, sum- and max-reduce of the value vectors over the sequence dimension. For the normalized results we take the sum-output vectors and normalize them (over the dh-dimensional vector dimension). Note that such a normalization can be applied to any of the aggregation methods to get qualitatively similar results. The plots in Figure 2 are generated by reporting the standard deviation over all output values and the mean norm of the output values, respectively.
Given the numerous successes of Transformers in natural language processing, we conjecture that a bias towards local information might be beneficial in language modeling. However, the implicit dependence on sequence length in a model that should be oblivious to different input sequence lengths is questionable. We leave an in depth investigation to future work.

heading: D ARCHITECTURES
text: We provide a schematic of 1 Transformer-layer of each architecture investigated in Figure 12. Our base architectures consist of 2 such layers followed by a projection to the output dependent on the task as described in the corresponding sections (cf. Section 5.1, 5.2 and 5.3).
D.1 ARCHITECTURE MODIFICATION ABLATIONS
An empirical ablation of the modifications that lead from the BERT architecture to the MTE architecture is given in Figure 13. The plots are generated as described in Sections 5.1.1 and 5.1.2.
The first column in Figure 13 corresponds to the original BERT architecture, trained with gradient norm clipping and learning rate warm up.
The second column (- warm up) corresponds to the same architecture, but trained without learning rate warm up. Here we see that too high learning rates learn even less without learning rate
warm up in the BERT architecture, hinting at a necessity for learning rate warm up for the original architecture.
The third column (- grad. clip) reports the results if we further remove gradient clipping from the training schedule. This does not seem to have a big impact in our setup.
Next, we report in the forth column (+ normalize) the results of moving the layer normalization before the residual addition and introducing an additional layer normalization right after the attention mechanism as well as on the hidden layer of the feed forward network. Note that this change removes the bias towards local information discussed in the end of Section 3. We see that this change leads to a profound shift in focus in regions where the learning rate is high: models with the original normalization focus on the (local) first-case, while models with our normalization focus on the (majority) argmin-case. This is in line with the insights stated in Section 5.1.2.
Finally, we report in the fifth column (+ GELU) the results of adding an additional GELU layer after the attention. These results correspond to the MTE architecture used throughout the paper.
Apart from the performance landscape changes just mentioned, the best hyper-parameter accuracies remain similar throughout all modifications, cf. Table 1.
For completeness, we also provide results for just replacing the softmax in BERT with normalization in Figure 14, that is, we train NAP with learning rate warmup, gradients clipped, no GELU after the attention and layer normalization as in BERT. The results are complementary to those presented in the main text: the model is still more robust to hyperparameter changes, albeit not performing well
on higher learning rates. On these high learning rates, the model focuses (similar to BERT) on the local first task. This is due to the placement of layer normalization, which with the initialization of the mixing layer yields a focus on local information.
D.2 REGULARIZATION EXPERIMENTS
To limit the number of variables which are not accounted for in the experiments, we focus on the infinite data but limited training time regime. In this regime, every batch consists of new data points. We believe that this regime is of paramount interest in future research, as more devices create a constant stream of data and training is more limited by the available training time than the available data. This regime allows us to omit regularization in all architectures as over-fitting is not an issue. In fact, our supplementary experiments below as well as related work (Lan et al., 2020) show that regularization does not help in this regime. We leave a comparison of the architectures in the limited data regime to future work.
Here, we show empirical results supporting the intuition that L2 as well as dropout regularization does not help in our setup. For each of our tasks, we take our default hyper-parameters (d = 128, L = 2, M = 4, N = 128) and train 5 random seeds per learning rate for models with regularization, varying the dropout rate in {0.0625, 0.125, 0.25, 0.5} and the L2 regularization weighting in {0.0001, 0.001, 0.01, 0.1}. Tables 2, 3 and 4 report the best mean accuracy achieved with the small number behind the accuracies indicating the regularization used, 1 referring to the smallest, 4 to the largest. We underline the results where regularization did lead to an improvement in mean accuracy. Note however that these improvements should be taken with a grain of salt, as (1) none of these improvements is significant considering the performance variation across random seeds and (2) the regularized values are likely to be overestimated, as the max is taken over 40 averages (4 regularization values times 10 learning rates) as compared to 10 averages (10 learning rates) in the unregulated case.
Overall we note that none of the architectures consistently benefits from regularization in our setup and regularization often decreases mean performance. Further, we point out that the best performance with regularization is most of the times achieved with the smallest regularization.
Table 4: Regularization results in the mode finding task. The top three rows correspond to the best mean training accuracy, while the bottom rows correspond to the best mean validation accuracy when validating on sequences of twice the length.
BERT MTE NAP NON sum max unregularized 99.6% 99.8% 99.6% 98.7% 99.8% 14.4% with dropout 93.9%1 93.3%1 94.3%1 91.8%1 93.3%1 24.5%1 with L2-regularization 99.5%1 99.9%2 99.7%1 98.8%1 99.9%4 14.4%2 unregularized 95.3% 95.4% 94.9% 91.3% 95.8% 13.5% with dropout 94.8%2 95.4%2 93.8%1 92.6%1 95.7%2 13.4%4 with L2-regularization 94.7%1 96.0%1 94.9%1 94.7%2 95.8%1 13.7%1

heading: E CASE LEARNING CURVES
text: Figures 15, 16 and 17 show the case accuracies over the course of training. The corresponding color plot is given in Figure 18. Besides the observations made in the main text, a few additional insights can be noted: (1) Cases are mostly learned in the order of their occurrences (recall that 72.37% of the examples are from the argmin case, 20.09% are from the first case and 7.53% are from the argmax case). This is to be expected when training with gradient descent, cf. Chatterjee (2020). (2) This order is not always given in the BERT architecture.
Besides the focus on the first case if the learning rate is too high - discussed in the main text - we also highlight a curiosity that occurs when the model dimension is too small (see plot highlighted in with red in Figure 15): The first case is learned and then unlearned in favor of the argmin case. Note that all 5 random seeds follow this pattern. Note also that for a different learning rate, the opposite holds as seen in the plot just below the highlighted plot. We highly encourage an interested reader to check out our code release,3 which includes these results as well as visualization scripts to inspect them further.
3See supplementary material.

heading: F CASE DISTINCTION TASK - ADDITIONAL RESULTS
text: Figure 18 shows the accuracy in the case distinction task when evaluating with the same sequence length as used during training, i.e., N = 128.
F.1 VARYING MODEL DIMENSION WITH 4 AND 6 LAYERS
Figures 19 and 20 show the results when varying the model dimension for deeper models, specifically 4 and 6 Transformer-layer deep models. The results do not differ which is why we choose to set the default to L = 2 layers, which in turn allows us to run more experiments. Note that we also present an experiment in F.6 where we explicitly vary the depth up to L = 64 Transformer-layers.
F.2 VARYING BATCH SIZE
In Figure 21 we provide the case accuracy results of an additional experiment, varying the batch size. In this experiment we train the models using different batch sizes, adjusting the number of training steps accordingly to keep the total number of training points seen constant. With this experiment we aim to show the training behaviour of the different architectures if we go from single example batches (many, potentially noisier updates) to batches of size 128 - a batch size in which each batch contains in expectation several examples per case, but fewer updates are made to the network parameters. Besides replicating several insights made in the main text, this experiment shows: (1) smaller batches require a smaller learning rate, supporting our argument that hyper-parameters should not be optimized independent of each other. (2) The focus of BERT on the first-case when the learning rate is too high is amplified in smaller batches.
F.3 VARYING INITIALIZATION SCALE
In Figure 22 we provide the case accuracy results of an additional experiment, varying the initialization standard deviation of the truncated normal distribution with which we initialize the weight matrices. The results show additionally to the observations made previously that the softmax based models BERT and MTE struggle to learn if the initialization is too large. This is likely due to the fact that the softmax is already saturated if the initialization too large. Also, larger initializations seem to require a larger learning rate.
F.4 LOCAL VS. GLOBAL FOCUS UNDER VARYING DATA BIAS
To investigate the local vs. global focus further, we conduct another experiment under varying data bias. Specifically we look at a task where each token has to output its own position (identity-case), if there is a 64 in the input sequence. If there is no 64, each token has to output the argmin position. The results when varying the percentage of identity-case data points are given in Figure 23. Red represents that only the identity-case is learned while turquoise represents that only the argmincase is learned. We can see that BERT struggles to learn the argmin-case (which requires global information) as soon as it is not the majority case anymore.
F.5 FIRST TOKEN OUTPUT - VARYING MODEL DIMENSION
Section 5.2 discusses the case accuracies when training on the case distinction task with outputs taken from the first token. In Figure 24 we addtionally give best the min-, mean- and max-accuracies over the course of training. The top row corresponds to in-distribution/training accuracy (N = 128) while the bottom row corresponds to out-of-distribution generalization accrucay when validating on sequences of half the length (N = 64). Again we note a correlation between optimal learning rate and model dimension, especially in the BERT and MTE architecture. We also note that BERT and MTE have a large performance variation across random seeds in this setup.
F.6 FIRST TOKEN OUTPUT - VARYING DEPTH
In this section we investigate whether our results are tied to the shallow architecture of L = 2 Transformer layers. We therefore vary the number of Transformer-layers L and report the results on the case distinction task with outputs taken from the first token in Figure 25. The results lead us to the following observations: (1) The BERT architecture does seem to perform better when the number of Transformer-layers is increased to L = 4. However, the performance degrades if we further increase the depth. (2) The NAP architecture achieves a higher best mean accuracy and performs well on a wide range of depths. (3) The max architecture performs well on the biggest range of hyperparameters. This is due to the beneficial architectural prior as discussed earlier.

heading: G MODE FINDING TASK - ADDITIONAL RESULTS
text: G.1 VARYING VOCABULARY SIZE
Figure 26 shows the results of an additional experiment, varying the vocabulary size S while keeping the sequence length N = 128 constant during training. For this experiment, we also vary the total number of training steps and set it to 400 ·S, to keep the number of examples seen per vocabulary token approximately constant. We also include zero-shot generalization results when testing on
sequences of twice the length (N = 256). Compared to the case distinction task we can do such a generalization evaluation here as we do not learn any positional embeddings in this setup. We make the following observations: (1) max completely fails to learn in any of the vocabulary sizes. Note that the shading to the left merely corresponds to the majority class base rate. (2) NAP struggles when the vocabulary consists of only 2 tokens. This is expected, as the mean subtraction in the normalization effectively removes the task relevant information (the mode) in this case. Note however, that for a high enough learning rate, the model learns to use the bias parameter b introduced in the normalization - effectively reverting to sum pooling. (3) While all models learn the task well on small vocabularies, NAP outperforms all other approaches significantly when S gets larger then the training sequence length, cf. Table 5.

heading: H LEARNED AGGREGATION WEIGHTS
text: Here we present the results of an additional model for fixed aggregation size. In particular, we ask how the results look if we learn the aggregation weights as parameters of the network. For this model, we take the sum model and replace the sum-broadcast operation with a fully connected linear layer across the sequence dimension. This is akin to the Random Synthesizer presented by Tay et al. (2021), but without the softmax. The results presented in Figure 27 confirm our conjecture that the softmax limits the information flow, as this new model is more robust to hyperparameters than BERT and MTE.

heading: I PPI EXPERIMENT
text: For the PPI node classification task we train a GNN based on our architectures with aggregation over neighbors instead of the whole sequence. Here we use L = 3 transformer layers, as the number of layers specifies the k-hop neighborhood from which information is aggregated in a GNN (k = L). The remaining settings are left at the default (M = 4 attention heads, no regularization). We train with one graph per update for 32 epochs.
All runs are repeated with 5 different random seeds. We track the validation performance over the course of training. Figure 10 shows the highest validation performance in each hyperparameter combination, with the results of all models given in Figure 28. For the test performance reported
we take the hyperparameter combination with the best average validation performance and use the checkpoints with highest validation performance over the course of training.

heading: J BABY AI EXPERIMENT
text: To test our architecutres in an RL environment, we base ourself on the open source implementation of the working memory graph agent presented by Loynd et al. (2020). Specifically, we take a factored representation of the environment and let the agent learn, what it needs to remember. Please refer to Loynd et al. (2020) for the core ideas here. In our experiments, we slightly modify their code, replacing the multi layer perceptron for the policy and value outputs with a simple linear layer each. Also, we replace the ReLU non-linearities used with GeLUs. These modifications are done to reflect a similar setup as investigated in the rest of the paper. Similarly we fix M = 4, L = 2 and set the “WMG Memo size” equal to the model dimension d. We chose level 3 of the Baby-AI environment (Chevalier-Boisvert et al., 2019) based on the results presented by Loynd et al. (2020). That is, we chose a task which trains decently fast and which should be solvable for a variety of hyperparameter combinations given the allocated training time. For the hyperparameters which do not alter the architecture, we used those which Loynd et al. (2020) reported to work best for their architecture. For completeness, we list them in Table 6. The complete result of all architectures investigated are given in Figure 29. Also here we ran every hyperparameter combination with 5 different random seeds.

heading: K BROADER IMPACT STATEMENT
text: Our work contrast different architectures on an abstract level. Due to the abstract nature of the study, there is no direct risk associated with system failure. However, we would like to mention potential implications of the findings. Our main proposal is an architecture which is more robust to hyperparameters. This has a potential positive benefits on the environment as less computation needs to be invested to reach a decent performance. Further, our architecture shows an increased robustness to skewed data distributions. This can have negative and positive aspects. On one hand, outliers inconsistent with ethical norms, such as hate speech and extreme views, might get picked up by the model more easily. On the other hand, the model might be able to represent minority groups more accurately, even though they are underrepresented in the data. In general, we see it as vital to get a better understanding of how different architectures learn to represent different biases in the data and hope that our work can provide a stepping stone in this direction.

heading: L FIGURE REPLICATION
text: Since our way to visualize the results uses the full color spectrum, it might be difficult for people with color blindness to verify and learn from our findings. We therefore replicate the main figures here with the R, G and B channels split into three separate plots. We also encourage readers to get in contact with us if they face any other difficulty reading our plots. We replicate Figures 5-10 in

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

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. 

The authors investigate the self-attention module architecture and provide a theoretical analysis for the limitation of using softmax function for normalization in the module. In this regard, an unconstrained normalization function is proposed as alternative in the transformer. The authors compare it with baseline transformer and several variants to validate the effect of the proposed normalization function in transformer.

The authors explore the implications of using softmax to implement attention mechanisms in NLP models, particularly transformers. They highlight the theoretical limitations of using softmax-attention and propose a normalization-based attention mechanism (NAP) that overcomes these constraints.. The authors conduct a set of synthetic experiments to compare different attention mechanisms’ performance and demonstrate that the proposed mechanism performs well compared to the traditional softmax attention while being  less sensitive to hyperparameters change.