Normalized Attention Without Probability Cage
Authors: Oliver Paul Richter
Roger Wattenhofer
Venue: ICLR 2022 Submitted (2022)
Keywords: Attention
Transformers
Neural Architecture
Aggregators

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 max- and 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.

Under review as a conference paper at ICLR 2022
NORMALIZED ATTENTION WITHOUT PROBABILITY CAGE
Anonymous authors Paper under double-blind review
ABSTRACT
Despite the popularity of attention based architectures like Transformers, the ge- ometrical implications of softmax- remain largely unexplored. In this work we highlight the limitations of constraining attention weights the prob- ability simplex and the resulting convex hull of value vectors. We show that Transformers are biased towards local information initialization and sensitive to hyperparameters, contrast attention to max- and sum-pooling and show the per- formance implications of different architectures with respect to biases in the data. Finally, propose to replace the softmax in self-attention with normalization, resulting generally applicable architecture that is robust to hyperparameters and biases in the data. We support insights with empirical results from more than 30,000 trained models. Implementations are in the supplementary material.
INTRODUCTION
The concept of neural attention (Graves 2013) Bah- danau et al.) 2015) has sparked number of ar- chitectural breakthroughs.
The Transformer ar- chitecture (Vaswani et al. 2017) successfully de- ploys multi-headed self-attention in several consec- utive layers for natural language processing (NLP) a solution that has become increasingly popu-
Accuracy
NAP BERT
lar /Vaswani al. 2017 Radford et al. 2018; 2019
Hyperparamter Setting
Brown et al. 2020
Devlin et al. 2019 Yang et al.
2019 Raffel et al.
2019 Liu et al., 2019). Apart
Figure 1: Softmax attention outputs can only lie within the convex hull spanned by the value vectors o' (blue region). Removing this constraint with our normalized attention pooling (NAP) leads to an increased robust- ness to hyperparameters (See Section 5.1.1
from NLP, self-attention has shown success in ap-
plications ranging from image classification (Par- mar et al. 2019 Dosovitskiy al. 2021) to gen-
erative adversarial networks (Zhang el al. 2019) to reinforcement learning (Bram et al. 2019 Parisotto cct al.] 2019; Loynd et al. 2020). The attention ar- chitecture choice is thereby often based 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 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 directly correlate with predictions (Jain & Wallace 2019, Wiegreffe & Pinter 2019; Pruthi et al., 2019 Brunner et al., 2020 Pascual et al. (020). 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 innate bias towards local information at initialization and showcase implications this architectural bias on biases in the data. We investigate different architectures an abstract level in numerous experiments and demonstrate the advantage of uncon- strained attention. In particular, we improve robustness to hyperparameters and show the general applicability of attention based architectures compared sum and max pooling.

---

Under review as a conference paper at ICLR 2022
2 BACKGROUND AND RELATED WORK
Many data processing tasks can be addressed by representing the input as a set or sequence of dis- crete tokens, c.g., words in sentence, frames in a video points in a point cloud. As general formulation, we represent each input token by a vector x' E Rd for E (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 un- fundamental
ordered sequences, i.e., sets of tokens. Multi-headed dot-product self-attention 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 x' are projected by an affine transformation to query Iin- key kmm and value vector un The dimensionality of these vectors is chosen as d, M
d 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 du, to form the attention logits 1 <qi,. Z and atten- vdl
tion vectors a'n = (l ( (I,N), where softmax refers to the normalized exponential exp(xr)
function softmax(x)) = Σx [x)) commonly used to project vectors to the probability simplex Sp = {at 1 2 V/j and Ej an = =). The output o' of each attention head m is then given by weighted sum of all value vectors o Σ, and vi, These attention head outputs are con- catenated and mixed trough an affine transformation to form the attention output in the Transformer architecture (Vaswani al. 2017).
In this work, we investigate whether constraining the attention vectors a;, into the probability simplex through the softmax function is sensible from a geometric perspective. We contrast the multi-head self-attention architecture to attention-inspires architectures without softmax (discussed in Section 4 as well as simpler aggregation methods commonly used. Specifically, while (unn et al. (2020) show that Transformers are universal sequence-to-sequenc: 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.) |017] 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. [2200] Wu et al.
2019 Lan et al. [2200] Dehghani et al. 2019; Fan et al., 2020 Wu* al. [020)] Press et al. [019
Bachlechner et al.] 2020). however, these do not alter the multi-head self-attention. Another di- rection of research has focused on improving the computational efficiency of Transformers (Wang)
ct al.
2020 Lee et al. 2019 Katharopoulos et al.) 22200] Zaheer et al., 2020: Choromanski et al. However, these either retain the softmax (Lee al. 2019]]Wang et al. 2020; Zaheer al.
2021
2020
or try to stay close to it (Katharopoulos et al. 2020, Choromanski et al.] 2021 Lu et al., We refer to Tay et al. (2020) for more extensive overview on works in this direction. Our Tocus is different: we ask whether the softmax is necessary at all. The Synthesizer (Tay et al.) (02T) and MLP-Mixer (Tolstikhin et al., 202T) 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)
2021.
introduce an attention mechanism that allows for negative attention values. [Schlag et al. 2021) link the head dimension d, corresponding limitation in the memorization capacity of linearized models. [Shen et al.) (202T) propose an architecture similar to our NON architecture (see below) for improved computational complexity. Cao (202I) also explore an architecture similar to our NON ar- chitecture and links it to a Petrov- Gaaeekin projection for operator approximation. They empirically find in an ablation that a layer normalization on the keys and queries separately improves conver- gence. Our work goes beyond this, showing the correlated effects between hyperparameters and architecture choices. 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 a' being constrained to the probability simplex Sp:
2

---

Under review as a conference paper at ICLR 2022
Outputs on are convex combinations of the vectors vin This in itself has interesting implica- tions. First and foremost, we note that a convex combination of vectors U cannot yield any vector outside the convex hull spanned by the value vectors D An illustration of this output "cage" is given in Figure | (left). We conjecture that this constraint limits the hypothesis space of the opti- mization, making it more difficult for gradient descent to find good solution. This conjecture is supported by our experimental results showing an increased robustness to hyperparameter choices when the constraint 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 Note that this implication highlights an inability to represent non- linearitt. 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 fo- cus on local information at initialization; and this focus depends on the sequence length N. To see this, consider the embeddings et after the first residual connec- x' x'+ W [E E vf E3 a 1.j v'y + b. where [+|-] de- notes concatenation and W and b represent the parameters of the affine transfor- mation that mixes the M attention head outputs to form the attention layer output. Our aim is to show how much the embed- ding e' influenced by the local infor- mation x' relative to the context informa- tion { * i}. We first note that the contribution of context information depends on the initialization of W, where a typi- cal initialization in language models yields a contractive projection [ 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 vec- tors Um as result of the convex hull. More- over, attention logits are normally close to at initialization (to have the softmax in the unsaturated region). This yields attention to be close to mean aggregation as of, U. N >; vi Considering vi,, to be distributed with zero-mean and variance of at initialization, we note that the expected magnitude and standard deviation of im scales proportional to VN This means that the fraction of context information in
tion given by e'
attention normalized
meam
Stm
max
Norm
10
10a
10
10
10-1
10
11
11
10
10a
102
103
10
10a
102
10
The standard deviation (o) and norm of a pooling output depends on the sequence length N (x- axis) if the output is not normalized. Softmax atten- tion outputs (blue) scale similar to mean pooling at initialization, i.e., Transformers focus more on local information in longer sequences. Experiment details are given in Appendix C
Figure 2:
1
e 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 We note that while an architectural bias towards local information might be beneficial in some applications, the implicit dependence aggregation size N is questionable. Finally, we note that:
visualize the dependence of om on N at initialization for different aggregators in Figure [
softmax can easily saturate, which leads to vanishing gradients: To show this, we look at the term g a0t.
Se which gives gradient factor the back pack-propggation path through the softmax.
Specifically, we note that E (0, 0.25) and that the range of H for which y > 0.01 is less than 9.2. This means that any gradient update that changes fi.j 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 single large input can result in parameter update that yields vanishing gradients on all following data points. See Appendix for proofs.
As an example, BERT (Delinn et al. 2019 initializes W with parameters drawn from truncated normal distribution with standard deviation set 0.02. This would only be roughly scale preserving for d = 2500.
3

---

Under review as a conference paper at ICLR 2022
4 NORMALIZED ATTENTION POOLING
Given the implications of softmax attention, one might seek alternatives. To investigate differ- ent architectural choices, we contrast the following architectures in an abstract empirical study:
Transformer Encoder (BERT): As starting point, we replicate the encoder architecture pre- sented by Vaswani al. (2017) described in the code release of Devlin et al. (2019) This architecture has been used extensively (Radford
BERT
MTE
NAP
Feodlorward Layer
LayerNorm
GELU
Softmax Attention
Normalized Attetion
iet al. 2018 Yang et al.
2019; Raffel et al.
Q Query Projection
2019
Liu et al., 2019, Clark al., (200)).
Key Projection
consists of two sub- modules: multi-head self-attention "layer"
Each Transformer-lae
OKW
V- Válue Projection
and feed forward network. Both modules are encompassed by residual connections. The multi-head self-attention "layer" consists of
Figure 3: Difference in Transformer layer.
projection to queries, keys and values, the attention mechanism as well as a mixing layer as de- scribed in Section 2 The feed forward network consists of two layers with a GELU (Hendrycks & Gimpel, 2016) non- non-rinearit.. Layer normalization (Ba et al. 2016) is applied between incoming and outgoing residual connections. Note that this gives crucial distinction of this architecture: Embed- dings 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 up and gradient norm clipping (cf. Devlin et al. 2019
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 Note that this different from the recently studied PreNorm (Parisotto al.,
Modified Transformer Encoder (MTE):
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 linearly decreasing learning rate schedule, taking Li et al. (2020) reference. We provide an ablation of all modifications Appendix [.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 through batch- (Toffe & Szegedy, 2013), layer- (Ba et al. 2016) group- (Wu & He 2020),
instance- (Ulyanov al.) (016) or normaliiomaiiioiionn (Salimans & Kingma 2016) our main
proposal is to simply replace the softmax with normalization:
23 Fx +
a = normalize((P?. ....y 1)
with normalize(ae) =
dx
where Pa N Σ od and om V E pa) are the mean and standard deviation of the corresponding input vector 2. in our case the logit vector calculated through key-query dot products. Similar to layer normalization (Ba et al. (0166), we introduce gain and bias parameters and b initialized to and 0, respectively. However, while Ba et al. (0016) 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 represent the logical XOR, normalized weighting
give proof in Appendix A
No Online Logit Normalization (NON): To investigate whether a dynamic normalization of the attention logits is necessary, also train model where we use the logits 15,1 directly as attention weights, i.e., of = GELU( /
Σ, Pi.i m v. The factor V/ 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-brooe broaccas operation.
https:://github..cmm
resea rcb/bert
4

---

Under review as a conference paper at ICLR 2022
// Cose distinction task data generator
inputs - random integer sequence
0.11
0.03
0.05
0.81
Predictions
0.92 0.03 0.04 0.01
if 64 in inputs
// argmin case
label argmin(inputs)
else if 50 inputs
//first
Pooling
label- 0
else
// aragmax case
label argmax(inputs) return (inputs, labe)
X
Embeddings
N.50
997.
42.
64.
33]
Inputs
52.
50.
67.
33]
Figure 4: Left: Pseudo code of the data generation. The case distinction points 64 and 50 are chosen arbitrarily. Right: Task setup for outputs across all tokens (cf. Section and outputs from the first token (cf. Section 5.2). Green boxes represent the trainable network layers (shared across tokens) while red boxes represent the pooling across tokens. The targets of the displayed examples would be [0. 0, O, 1] and [1, 0, O, OJ, respectively.
Max Pooling over Sequence Dimension (max): Similar to sum pooling, we can replace the attention sub-module with a simple max-reduce-broodd operation over the sequence dimension. Note that max pooling is a powerful operation that yields direct link to up to d different tokens.
We provide a schematic figure of all architectures in Appendix D not varied a corresponding experiment, we default architecture hyperparameters to T, 2 Transformer-lyyer (consisting of an attention sub-module and feed forward sub-module each), .M = heads to calculate the logits (if applicable), d 128 as model dimension and train sequences each, using the Adam optimizer (Kingma & Ba 2014). The hidden dimension of the feed forward sub-mot mobulls 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.
total of 3200 batches of 32 example
5 EXPERIMENTS AND RESULTS
Our goal with this work is to provide insight into the variety of performance implications that the architecture choices entail. We aim to provide these insights independent of any particular application, these architectures
be applied to variety of tasks from NLP (Vaswani al.,
Devlin et al. 2019) to graph neural networks (Velickovic et al.] 2018 Yun et al.] 2019] 9 Xu al. 2019 Velikovi al. 2020) reinforcement learning agents (Parisotto et | 2019 Fang al. Loynd et al. 2020). We therefore focus on carefully crafted synthetic tasks that (T) are general enough in that we can expect the insights to generalize us modify key aspects that are hidden in real world data sets, such as a bias towards certain sub- task. 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 in depth discussion of this setup as well as regularized ablations.
2017
2019
large set of downstream tasks and (2) let
5.1 ARGMIN-FIRST-ARGNN
CASE DISTINCTION TASK
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 (c.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, (also randomly initialized) positional embedding is added to provide position-relatioo information. The sequence of tokens is then processed by several architecture dependent Transformer-layers (cf. Section 4 Finally, each contextualized embedding is projected to single output. A softmax-croot coftmaxx-crosentrroyy loss is applied over the sequence dimension. See Figure (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 learning perspective as the networks start from random embeddings which do not provide any ordering information. Which embeddings correspond to bigger integers and which smaller
5

---

Under review as a conference paper at ICLR 2022
BERT (97.740448)
MTE (7744038)
NAP (9990488)
NON (9754068)
Sum (84.5+6.88))
max (97.810.58)
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
6.9c-6
MONTWONT
WONT ONt -MONin MiN
SONTXONT
Figure 5: Learning rate (y-axis) vs. model dimension d (x-axis) on the argmin Seemin-fisst-aree -irst-argmax case distinction task (with output across all tokens). The pixels' R (red), G (green) and B (blue) values correspond to min-, mean- and max-accuracy. respectively, of the corresponding hyperparameter combination see main text for details. The plot shows the validation accuracy when validating on sequences of length N = 64. Crosses indicate the combination for best validation accuracy, which report with standard deviation behind the model name.
integers has to be inferred during training. Further, the case distinction this task lets us tweak the data bias towards each sub-task. Specifically, we consider vocabulary size of S = 100 integers (0-99) and uniformly random sampled sequences of = 128 tokens length. This leads to a bias as Pargmin = 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.
}) se 72.4% of data points require the network to pin-point the
5.1.1 VARYING MODEL DIMENSION d
As a first investigation, we are interested in how varying the model dimension d influences the ar- chitectures ability learn the given task. For this, we train each of the architectures for each of the model dimensions d E {8, 16, 32, 64, 128, 256, 512, 1024) using 10 different learning rates and random seeds for each hyperparameter combination. As we sults as possible, we derive a novel, human friendly visualization of results. Figure 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 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 hyperparam- eter combinations where none of the random seeds could solve the task. See Appendix [ for figures with color channels split. Given that all architectures are applicable to sequences of any length, investigate how the architectures generalize sequences of different length. Specifically, 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 give rise to (1) Most models have some
base our insights on as many re-
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 hyper- parameters independent of each other might lead to sub-optimal results. (3) Models with probability simplex limitations (BERT and MTE) work for 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, repeat this experiment with = and L = 6 Transformer layers, but the results stay the same (see Appendix F.T). We also provide the results of an additional model with learned aggregation weights in Appendix
the following observations:
5.1.2 CASE ACCURACY UNDER VARYING DATA BIASES
As a next experiment we reset the model dimension to = 128 and investigate the models un- der varying data biases. We do in two distinct ways: First, vary the sequence length

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
Swm
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.60-5
2.0c-5
5.9e-6
-MON
16
64
128
250
512
1
1
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
6.9e-6
9.0%
18
36.2%
72.4 96
86.2
96 56.5%
98.3
9.0%
18.1
36.2
72.4
86.2
06.5
98.3%
9.0%
18,
36.2
72.4%
86.2%
93. 93.1%
96.5% 56.5%
98.3%
9.0%
18 %6
36.2
12.4
86
93.1
96
98.3%
9.0%
18.1 1%
36.2%
7244%
86.2%
03:1 %
96.3%%
98.3
9.0%
18.1
36.2
72.4%
86.2%
03.1
06 8.3%
96.3
Figure 6: Biased data results on the case distinction task (with output across all tokens). RGB
pixel values correspond argmin-,fire and
respectively. Top row:
Learning rate (y-axis) vs. sequence length N (x-axis). Bottom row: Learning rate (y-axis) vs.
percentage argmin-case the data (x-axis) with fixed 128.
N E (4,8, 16, 64, 128, 256, 512). Note that this implicitly varies the biases Pargmin- Pfirst se
and Pargmaz in the data. Second, fix the sequence length 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) cor- accuracy, G (green) corresponding to the first-case accuracy and B (blue) corresponding to the argmax-cast accuracy. As 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 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 too high, the BERT architecture tends to focus on the first-case. We believe this is due to the be
responding to the argmin-
architectural bias towards local information discussed in Section Note that the first-case solved by relying on the local positional embedding. (3) Only the NAP and max architecture manage 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 F2 and F3 we provide further experiments investigating different batch sizes and initialization scales. In Appendix F.4 we investigate the local vs. global aggregation focus further.
5.2 FIRST TOKEN OUTPUT
The task so far requires the architectures to learn information flow between tokens to distinguish the case and decide per token, whether it is the token that looked for or not. Now we investigate, whether all this information architecture output slightly in that we only take the contextualized embedding of the first token and project from it vector of size N (see example the right in Figure 4). Note that this task set-up harder and can highlight bottlenecks
also be aggregated into single token. We therefore modify the
the information flow across tokens. We fix the sequence lenght N 128 and again vary the model dimension d. We report the case specific are given in Appendix F.5 We observe: (1) All architectures learn for (almost) all combinations the now close trivial first- case, even though it is not the majority case. (2) The sum pooling architecture does not learn any of other cases. (3) Only NAP and max learn all three cases some hyperparameter combinations.
mean accuracies in Figure 7 Min- mean- and muxx-overall-accuaade

---

Under review as a conference paper at ICLR 2022
BERT (92.240.3)) MTE (7584425.38)
NAP (988:00225)
NON (88.9=1.1%)
sum (23.7600.5))
max (98.6-20.18)
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2c-4
6.6e-5
2.0c-5
5.9c-6
MONTA 1024
102
1023
Figure 7: Learning rate (y-axis) vs. model dimension d (x-axis) on the case distinction task with output from the first token. RGB pixel values correspond to the accuracies. Crosses indicate the best accuracy, reported behind the model name.
BERT (999-40)))
MTE (100+
NAP (9.40.19)
NON (9940.8)
sum (1000
max (15.3+0051))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0e-5
5.9e-6
TMONING
1024
Figure 8: Learning rate (y-axis) vs. model dimension (x-axis) on the mode finding task. RGB pixel values correspond to min, mean and max accuracy. Crosses indicate the combination for best accuracy, reported behind the model name
The worse performance of NON highlights the advantage of online normalization of the
BERT
MTE
NAP
0.3
logits. While the softmax provides some
0.09
0.027
form of online normalization, we hypothe-
0.008
size that the worse performance of MTE and
0.002
7.3e-4
BERT in this task stems from an information
2.2e-4
bottleneck induced by the probability sim-
6.6e-5
2.0e55
plex limitations. To test this hypothesis fur- ther, 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, sup- porting our hypothesis. Note however, that MTE and BERT are still outperformed sig- nificantly by NAP. In Appendix [F6) we pro- vide further experiment, varying the depth up to L = 64. Results are complementary.
5,9e-6
mon
Figure 9: Learning rate (y-axis) vs. attention heads M (x-axis) on the case distinction task (first token output). RGB pixel values correspond to min, mean and max accuracy. Crosses indicate combination for best accuracy, reported behind the model name.
5.3 MODE FINDING TASK
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//grreement and reinforcement learning (e.g., count based cxploration). Here remove the positional embeddings, as this task can also be done on sets, and project from the contextualized embedding of the first token to vector of dimension S (the vocabulary size) over which apply the softmax-cross-
8

---

Under review as a conference paper at ICLR 2022
BERT 0.9.)))00%
MTE (9.1+008)
NAP 199.1-0.18)
BERT (99.3-00226)
MTE (99.40002%)
NAP (9944019))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.00-5
2.0c-5
5.9e-6
WONT -OnN Ninn
Figure 10: Learning rate (y-axis) vs. model dimension d (x-axis). RGB pixel values correspond to min, mean and max validation performance. Crosses indicate the best combination (validation), from which report the test performance. For NON, and see Appendices | and Left: Protein-prow seeein-protein-inten interaaiiontnnnnnn task. Shown is the node classification F1 Fl-score. Right: Altered working memory graph agent in the Baby-Al level reinforcement learning task. Shown is the success rate.
entropy loss. We keep N 128 but reduce S 10 to have meaningful modes. Ties broken by taking the smallest integer of the ones with maximal occurrence. Results of varying the model dimension are reported in Figure 8 We observe: (1) sum pooling works well on this task, has suitable architecutral prior. (2) max pooling cannot learn the task, not even with a model dimension 1024 N. Appendix G. [. we provide complementary results varying the vocabulary size. We refer an interested reader to Xu al. 2020) for more architecture-task alignment.
5.4 GENERALIZATION TO GRAPH NEURAL NETWORKS AND REINFORCEMENT LEARNING
NLP experiments normally require large amounts of data and large models setting which would make an investigation of the training dynamics under various hyperparameters prohibitively expen- sive. We therefore focus in our generalization study the other domains where attention based architectures have recently gained traction. Specifically, to see whether the the learning dynamics seen far generalize, we perform two additional experiments: One investigating the architectures in graph neural network (GNN) task and one investigating them in a reinforcement learning (RL) setup. Specifically, we train transformer based GNN the Protein-Protted Inttraitinttaaoiion (PPI) graph node classification task (Zimik & Leskoved 2017) and investigate the different architectures as al- terations on the working memory graph architecture for RL (Loynd et al.) 2020), training on level 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 com- binations. (2) There is a correlation between optimal learning rate and model dimension. (3) BERT hyperparameter sensitive in comparison. See Appendices and [ experiment details.
of the Baby-Al environment (Chevalier-Boisse
CONCLUSION
Taking all observations together, we conclude: Many recent works apply some sort of self-attention mechanism involving a softmax that projects the attention weights the probability simplex. In this work we question the softmax in dot-product self-attention modules. Our theoretical investiga- tion shows that anttetx-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 num-pooling perform well when their architectural prior aligns with the underlying task. These architectures however fail in cases where the architec- tural prior is not suitable. As a solution, we propose to replace the softmax normalization. Our resulting normalized attention pooling (NAP) architecture is the only architec- ture 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 work provides 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 our broader impact statement in Appendix K
attention through

---

Under review as a conference paper at ICLR 2022
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 16070066504 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. arXiv preprint arXiv:2003.04887, 2020.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio, Neural machine translation by jointly learning to align and translate. Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Con- ference Track Proceedings, 2015. URL
a409/ 0473
Timo Bram, Gino Brunner, Oliver Richter, and Roger Wattenhofer. Attentive multi-task deep rein- forcement learning. 07 2019.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. BeeglerJJfffree Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Floring Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wat- tenhofer. On identifiability in transformers. In International Conference on Learning Represen- tations, 2020. URL ht htts: /ttpt:::s///ppeeeeeee net B
Shuhao Cao. Choose transformer: Fourier or galerkin. CoRR, 0bbs/2105.14995. 2021. URL
Satrajit Chatterjee. Coherent gradients: An approach to understanding generalization in gradient
descent-based optimization. In International Conference on Learning Representations, 2020.
URL http:httss///peeeeeeeet wwt. fnret/forum' drtooruiiirrreeer rd=rrreyorrrFFsss
Maxime Boiiseter-Boisvert,, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. BabyAl: First steps towards grounded language learning with a human in the loop. In International Conference on Learning Representations, 2019. URL
dattps://openrevier
net/ www.foorrridiieree dederummocdr
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In International Conference on Learning Representations, 2021. URL https: //openreview. weeee//orrmmidd=aaaaaa
Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning. Electra: Pre-training text encoders as discriminators rather than generators. In International Conference on Learning Representations, 2020. URL
net wtt.et//orumimieeeeeee xMH1BtvM
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019. URL https:
//openzeview. wwt. wtw.ee//oorommiallydaaa........a
Jacob Devlin, Ming- Win Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT 2019, Minneapolis, MN, USA. June 2-7, 2019, Volume 1 (Long and Short Papers), pp. 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.18653/v 1/n19 1423. URL
10

---

Under review as a conference paper at ICLR 2022
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openreview. wtr./ooumt/tteeeeeoooo i= cbFdNTTy
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. In International Conference on Learning Representations, 2020. URL
htttps://openrevier net http://wwww.ntttooumdd---d----e
Kuan Fang, Alexander Toshev, Fei-Fei Li, and Silvio Savarese. Scene memory transformer for embodied agents in long-horizon tasks. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA. USA, June 16-20, 2019, pp. 538 547. Computer Vi- htLhttp: /tL/tttp://pppaccessss openaccess. sion Foundation/ / IEEE, 2019. doi: 10.1 ht019//tmm/aaggScceeee Scene Beeeeeeeeeeeeeeeeeeeeeeeeeee Taanners__Troooooooaaa cceevffccm/conttei cont CWWR_CVVRC22_0 2eteR_2019hhmmmm/e for
Long-_Logg-___ Ho r zon Tasks_assssaaa CVPR________________________ 00i__20155 2019_papeere html
Alex Graves. Generating sequences with recurrent neural networks. CoRR, abs/1308.0850, 2013.
hwt.htttp//ariiv.rrr/////// wtt.ooggg//////....// 0850
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
arXiv:1606.084/5 2016.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 448-456. JMLR.org, 2015. URL html
http://prroceddings mwr wwwpss/v/3/ioss.
Sarthak Jain and Byron C. Wallace. Attention is not explanation. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the
Association for Computational Linguisties: Human Language Technologies, NAACL-HLI 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 3543-3556. Association for Computational Linguistics, 2019. doi: 16663//1/119-1357. URL https: wtt.:/ooo/11....3/.... 1337
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th Inter- national Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, vol- ume 119 of Proceedings of Machine Learning Research, pp. 1165-1165. PMLR, 2020. URL mllr
20mmmmll
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.:4412.6999 69800 2014.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020.
net forum?
id=HleA7AEtvs
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for basedd cermuuationninvarian neural networks. In Kama- lika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Confer- ence on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, vol- ume 97 of Proceedings of Machine Learning Research, pp. 3753-3753. PMLR, 2019. URL /ttp:///www....p........////////
mlr. wwwps//////eeeell/eeeeeee
Mengtian Li, Ersin Yumer, and Deva Ramanan. Budgeted training: Rethinking deep neural network training under resource constraints. In International Conference Learning Representations, 2020. /ttttp:://openrevieee www.:/ee/fooummmeeemmmm wtt.::/oorrunntlllyy....v.EE
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. arXiv preprint 088249, 2020.

---

Under review as a conference paper at ICLR 2022
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer. and Veselin Stoyanov, Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692 1692, 2019. http:///aax
abs 19t//997....0 191.1.17772.112
Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs. In Hal Daum III and Aarti Singh (cds.), Proceedings of the 37th In- ternational Conference on Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 6404-6414, PMLR, 13-18 Jul 2020. URL http:ht:pp//ceeeeeeeeeeeeee .ma.eddiig.mm....
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. SOFT: softmax- free transformer with linear complexity. CoRR, abs/21 19.421,,. 2021. UtLhtt:::://ar:
11945
Jawad Nagi, Frederick Ducatelle, Gianni A Di Caro, Dan Cireşan, Ueli Meier, Alessandro Giusti, Farrukh Nagi, Jürgen Schmidhuber, and Luca Maria Gambardella. Max-pooling convolutional neural networks for vision- based hand gesture recognition. In 2011 IEEE International Confer- on Signal and Image Processing Applications (ICSIPA), pp. 342 347. IEEE, 2011.
Toan Q Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention. arXiv preprint 11X.::9100005995 2019.
Emilio Parisotto, H Francis Song. Jack W Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. arXiv preprint arXiv: 2019.
Niki Parmar, Prajit Ramachandran, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens. Stand-alone self-attention in vision models. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Bucc Emily B. Fox, and Roman Garnett (eds.), Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 68 80, 2019.
Damian Pascual, Gino Brunner, and Roger Wattenhofer. Telling bert's full story: from local attention to global aggregation. arXiv preprint arXiv:2004.059166 2020.
Ofir Press, Noah A Smith, and Omer Levy. Improving transformer models by reordering their sublayers. arXiv preprint arXiv:1911.080: 1.08644, 2019.
Danish Pruthi, Mansi Gupta, Bhuwan Dhingra, Graham Neubig, and Zachary C Lipton. Learning to deceive with attention-based explanations. arXiv preprint (tt::/999007933. 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019.
Tim Salimans and Diederik P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5- 10, 2016, Barcelona, Spain, pp. 901, 2016.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th Interna- tional Conference Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, vol- ume 139 of Proceedings of Machine Learning Research, pp. 2355-9336 PMLR, 2021. URL http::/proceediin ngs. www.rrppsssvti ptress/v139/ aagllallllllaaac .htll
12

---

Under review as a conference paper at ICLR 2022
Nimrod Segol and Yaron Lipman. On universal equivariant set networks. In International Confer- ence on Learning Representations, 2020. URL https:://pppneeeeeew net/ formm??ormm?mmm d HkxTwwrKDB
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In IEEE Winter Conference Applications of Computer Vision, WACV 2021, Waikoloa, HI. USA, Jamuary 3-8, 2021, pp. 3530-3538, IEEE, 2021. doi: 10. 1109/NACV4663022001000 00:57
URLpttss:://aatt org 10. 1109 WWW./ 2021.
Yi Tay, Anh Tuan Luu, Aston Zhang. Shuohang Wang, and Siu Cheung Hui. Compositional de- attention networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bue, Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 6135-6145 Curran Associates, Inc., 2019. URL BtLLhtps:: http://prceeeeec..coocccooooo.se / neurips cc/
Yi Tay, Mostafa Dehghami. Dara Bahri, and Donald Metzler. Efficient transformers: A survey. CoRR, abs/2009.06732. 2020. httpt:ttp///xiiii wttp/org////ooo//oo http:///2099.663333
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models, 2021.
Ilya O. Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mip-mixer: An all-mlp architecture for vision. CoRR, abs/2105.01601, 2021. URL
http//bb//21055 05. 01
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing gredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, UI- rike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: An- nual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 5998-6008, 2017. URL http://papers. nips. cc/paper/ 1esl-attenioonissalllll
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April May 3, 2018, Conference Truck Proceedings. OpenReview.net, 2018. URL https::t/////cc http::///ppeeevvieeeeeee / net. wwt.:nneteot/em... wtw.ammmdddrxxxiiieoom
Petar Velikovi, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. In International Conference Learning Representations, 2020. URL
htttss://openreve ew. net.
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, 2442200604768, 2020. URL https::/arxxv.orgg///oooo 2006.04765
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Meth- ods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pp. 11-20. Association for Computational Linguistics, 2019. doi: 10.18653/11DD19-1002 URL https://dda wttt://///...1//////11/13113 /
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In International Conference on Learning Representations, 2019. URL
ew. http://wwwweeeeeeeeeeee/ooeeeean
13

---

Under review as a conference paper at ICLR 2022
Yuxin Wu and Kaiming He. Group normalization. Int. J. Comput. Vis., 128(3):742- 755. 2020. doi: 10.1007/1122--011-0000666 URL https: http://.......1.110..1100 s 111263-019-01192
Zhanghao Wu* Zhijian Liu*, Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. In International Conference on Learning Representations, 2020. URL https:
//openzeview. net/ Lesee//ormmmid-BBeee.. HHPP
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In 7th International Conference on Learning Representations, ICLR 2019, New Or- leans, LA, USA, May 6-9, 2019. OpenReview.net, 2019, URL httl:tttt////pereeee//
net/
Eorum? femmmidddryye 6m5km
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S. Du, Ken ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? In International Conference on Learning Representa- fions, 2020. URL httttttt:://peeeee/e/e
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 5754-5764, 2019.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Con- ference on Learning Representations, 2020. URL https://apenzeviiet entttt/eeeeeeeeeeeeeeeee nnt/ forum?
id (d-ByxRMONEVE
Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, and Hyunwoo J. Kim. Graph trans- former networks. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché- Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Sys- tems 32: Anmual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 11960-11970, 2019.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander J. Smola. Deep sets. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Pro- cessing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 3391-3401, 2017. URL http://papers nips wwtpcc//ppppp..e.. 6931 deep-set.s
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Trans- formers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, De- cember 6-1 2020, virtual, 2020. URL 2020 2ttp://2....11212222222225333333
cc/paper/ abb Abst htm_
Han Zhang, Ian J. Goodfellow, Dimitris N. Metaxas, and Augustus Odena. Self-attention generative adversarial networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 7354- 7363. PMLR, 2019. URL http:uttt//po///oooo/te httt:::/pproceeednggg...t mir wttp://wwww.pppe..reeeeeeeme wwwps/ere/11 waazghaagg9 19d.. html
Marinka Zitnik and Jure Leskovec. Predicting multicellular function through multi-layer tissue networks. Bioinform., 33(14))i190--198, 2017. doi: 10. 10..003/hioinffmmaaiisss222222 URL
14

---

Under review as a conference paper at ICLR 2022
A REPRESENTATION POWER WITH RESPECT TO XOR
Lemma 1, No convex combination can represent the binary exclusive OR (XOR) function defined on binary inputs I1 10,1) and 22 E [0, by the indicator function as X XOR(21, (22) = =21
Proof. Suppose there exist convex combination weights and d2 with di + d2 = 1, such that 01 #1 + 02 22 represents the XOR function. Plugging in #1 = d2 yields a1 31 + 02 D2
a1 d2 1, which gives the contradiction.
Lemma 2. Given the two binary inputs E {0,1) 1 and x2 E {0, 1), there exists an affine mapping f: ::{0,1}² + R2, such that the normalized weighting given by
f1(11, 22) (f(11122)
f2 x2) (f(
21
T2
(((122))
3()
is equivalent to the logical exclusive OR given by the indicator function as XOR(E1,22) = 121
Proof. For a vector I E R2, the standard deviation o can be simplified to
2 (1e
Σ
12)
(1(1,2)
11 +
+ (11 1 +2
2
l2
and the normalization function reduces to
11 P4 la MM
normalize(()
1
1 la
1/ 1
111 12
1/2 12
[1, 11.-112
iʃ >
[-1, 1jr
if
unde f-
is 11 = l2
As example, consider the affine mapping ((x) = = -3xr 1, 202]2 which for æ {0,1)
and Ø2 E }} results in the function
3x + 1 2z2
3x 1 1+ 200
11 +
x2 1
3r + 202
|3
+ 2x2
We note that for a realization of such an afline mapping across tokens given the weight sharing constraints of the discussed architectures would need æ and x2 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 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 the head and vector dimension used the main text.
Lemma 3. The partial derivative 6eoftmax(1); is limited to the range g E (0, 0.25).
81
15

---

Under review as a conference paper at ICLR 2022
Proof.
6softmax(4)
exp (4) E, exp (1/) exp (2.4)
0; (1 (1-a;) a,) with softmax(();
84g
Exexp exp (1))
Maximizing this function we get that the maximum of g 0.25 achieved for a; = 0.5. Noting that > 0, the claim follows.
To expand this further, we introduce the notion of a-bandwidth.
Definition 1. We call the range Ar in which the absolute value of gradient factor 6f(a) exceeds threshold a, i.e.
5f(x) 52
a, its a-bandwidth BW.,
(f(w))
Please refer to Figure 11 for visualization of this definition. Intu- itively, the a-bandwidth gives a measure for the range of operation, where at least an a-fraction of the gradient gets backpropagated.
8f(x)
also provides measure of how fast a function saturates: Given with ff(z) a, updating to x'
BW
BWa
(f(())
puts x' into
the saturated region where 8f(x²) I A function with small a- dx²
X
bandwidth can therefore more easily saturate than a function with a larger e-bandwidth.
Figure 11: Visual definition
Lemma 4. The a-bandwidth of 6suftmax(L)) can be calculated as 84,
- 2a + vI 4a
BW In
2a V1 40
Proof. Substituting exp(l:) and y Eizi exp(li) can write g as g(x) 24z (1- xty
Solving g(x) yields x y. (1 2a + /1 4a). and subtracting the larger from the smaller value, we get:
Reversing the substitution of l In(x)
2a
BW) = In y (1
2a + vI 44) In1
1/ (1 2a vI 4aa
2xx
2x
1 - 2nx + v1 4a
= In
20x VI 4a
Corollary 1. The range in which g 6softmax((), 0.01 is smaller than 9.2.
Proof. Calculating the a-bandwidth for a = 0.01 we get BWo.o1 9.170.
Note that BW)(g) is independent of the length of the vector as the remainder y E fet exp(ly) cancels. For the self-attention the Transformer this translates to the fact that BW. is indepen- dent 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.0 D1-bbndwidthh on its own not the problem. In fact, tanh-non- linearity has 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 WA + by. q = W + by W Wr. + WT W, invertible and ||>|| 2 Ci we can find constants Co and 03 such that
2 02
x - C3
16

---

Under review as a conference paper at ICLR 2022
Proof.
Writing out |
Sq.k> vdz
get
(x W W (bg W, + W
WR W.) bx)
v dh
Using
5x Az bm
x² x= + A²) and substituting gr. - b/. Wg by Wi we get
81
(xr (W Wr W.W,) gb)
Setting 9n = xx ("' WT + W W;) and solving for a we get (W W + WW = a. With this, we can bound || with the operator norm on (W W. + W WTW,)" W Specifically, if we set 22 as
1 1 We W
C2
Vd,,
get vih |19mm
2 C2
||||||
We can then bound
8l
Ix + goll
2 dd
1 ||
gb
dæ
ddh
Therefore, if set 61
|19b|1
we get ||
2 C2 vdn ||x
and hence
51
8æx
2
ddh
(l
|9
l) 2
||
C3
for 1_ -1.b|1. Jan
Lemma 5 shows that the magnitude of the gradient factor scales with the magnitude of the input. Therefore, single large input can yield correspondingly large parameter update. This in turn can put the softmax easily into the saturated region, since the softmax only has a constant a- a-bandwidth. Once saturated, the softmax leads vanishing gradients henceforth.
C SEQUENCE LENGTH DEPENDENT FOCUS
For Figure 2 we sample 16'384 value, key and query vectors of dimension dn 128 per sequence length N E (1, 4, 8, 16, 32, 64, 128, 512, 1024, 2048} from normal Gaussian (0, Ia.. Ia, being the d, -immennsioa identity matrix. We split the samples form the sequences and calculate the corresponding output vectors o² for i (1....,Ny N). Here, the softmax attention outputs are calculated 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 dy -ddmenniional vector dimension). Note that such normalization can be applied to any of the aggregation methods to get qualitatively similar results. The plots in Figure are generated by reporting the standard deviation over all output values and the mean norm of the output values, respectively.
Given the numerous successes Transformers in natural language processing, we conjecture that 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 in depth investigation future work.
D ARCHITECTURES
We provide a schematic of Transformer-lant of each architecture investigated in Figure 12 Our base architectures consist of such layers followed by a projection to the output dependent on the task described in the corresponding sections (cf. Section 5.1 5.2 and 5.3)
17

---

Under review as a conference paper at ICLR 2022
Layer normalization
normalization
norpalization
Food forward layer
Food forward layer
Feed forward layer
GELU non-linearity
GELU non cinn-lineritty
Layer normalization
Layer normulization
GELU non-lincarity
Feed forward
Feed forward
Feed forward layer
Laver normalization
Layer normal ization
Layer normalization
Mixing
Mixing
Mixing layer
non- lineriicarity
non- non-lincarity
Multi-head soltmax atlention
Layer normalization
Layer normalization
Multi-lead softmax attenitin
Mulii-bead normallaed weighting
Queries
Keys
Values
Querics
Keys
Values
Queries
Keys
Values
(a) BERT
(b) MTE
(c) NAP
Layer normalization
Food forward layer
GELU non- lineurity
Laver nocmalization
Laver normalization
Layer normalization
Feed forward layer
Feed forward layer
Foed forward layer
GELU nan- lineurity
GEL U cin-liieurity
nommalization
normlliaaiion
Foed forward layer
Foed forward layer
Layer normalization
Mixing layer
GELU non- lincariiyr
Layer nermalization
Layer nommalizatonn
scale by 1/v Sequence lenghi
Multi head dot product weighting
Sum-redut remt-reduce over sequence
Max-reduuz over sequence
Queries
Keys
Values
(d) NON
(c) sum
(f) max
Schematics of Transformer- block of the different architectures investigated. Green layers correspond to the main weight matrices that are trained. Note that displayed dimen- sions are to scale the hidden dimension of the feed forward layer larger than the model dimension and the hidden layer size in the feed forward network of "max" and "sum" are adjusted approximately match the parameter count of the other architectures.
Figure 12:
18

---

Under review as a conference paper at ICLR 2022
BERT
warm up
grad. clip
normalize
GELU
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9c-6
16
16
128
1024
1024
1024
1022
00
1
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
16
[024
1022
150
312
1024
1022
0.3
0.09
0.027
0.008
0,002
7.3e-4
2.2e-4
5
2.0e-5
6.9c-
16
32
64
128
1024
32
64
250
$12
1024
16
32
128
256
1024
64
250
1024
16
1
0.09
0.027
0.008
0,002
7.3c-4
2.2e-4
6.60-5
2.0e-
6.9c-6
30
00 MONN SONT
102
INIO
Figure 13: Learning rate (y-axis) vs. model dimension d (x-axis) on the aggmin-fist-aarmmal case distinction task (with output across all tokens) architecture modification ablation study. the first two rows, RGB pixel values correspond to min- mean- and max max-acccracyy In the last two rows, RGB pixel values correspondto argmin-, frgmin-.fistt- and
1. row: Training accuracy (sequence length N = 128). 2. row: Validation accuracy when validating on sequences of half the length (N 64). 3. row: Training case accuracy (sequence length
128). 4. row: Validation case accuracy when validating sequences of half the length (N 64). Crosses indicate combination for best mean aceuracy, which are reported Table T
D.1 ARCHITECTURE MODIFICATION ABLATIONS
An empirical ablation of the modifications that lead from the BERT architecture to the MTE chitecture is given Figure 13
The plots are generated as described in Sections 5.I.I and 5.1.2
The first column in Figure 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 learn- ing rate up. Here that too high learning rates learn even less without learning rate
19

---

Under review as a conference paper at ICLR 2022
Table Ablation study accuracy values taken from the papee--rrameter combination that led the best mean overall accuracy, indicated by a cross in Figure 13
BERT
warm up
grad. clip
normalize
+ GELU
Overall
min 99.3%
99.4%
99.3%
99.2%
99.3%
Training
mean
99.4%
99.5%
99.4%
99.3%
99.3%
Accuracy
max
99.5%
99.6%
99.6%
99.5%
99.5%
Overall
min 96.5%
96.9%
96.9%
96.3%
96.7%
Validation
mean 97.2%
97.6%
97.2%
96.8%
97.3%
Accuracy
98.2%
98.2%
98.2%
98.4%
98.2%
Mean Case
argmin 99.3%
99.6%
99.5%
99.6%
99.5%
Accuracy
first 100%
100%
100%
100%
100%
Training
argmax
96.9%
98.1%
97.5%
96.5%
98.0%
Mcan Case
argmin 98.0%
98.0%
98.0%
98.1%
97.9%
Accuracy
first 100%
100%
100%
100%
100%
Validation
argmax 93.9%
93.8%
93.1%
91.2%
93.8%
Cases all tokens (9.42..1))
Cases token (99.220.19)
Mode (99.4.10.18)
0.3
0.09
0,027
0.008
0.002
7.3e-
2.2e-
6.6e-5
2.0e-5
6.90-6
300 sonreeo -MONii -eiing
Figure 14: Learning rate (y-axis) model dimension (x-axis) in the different task setups when just replacing the softmax in BERT with normalization. The plots from left to right are grouped according the tasks: The task with outputs across all tokens (ef. Section 5.1 S. the task with outputs from the first token (cf. Section 5.2) and the mode task (cf. Section 5.3) For the case tasks, the left sub-plot reports min, mean and max accuracies while the right sub-plot reports mean case accuracies as RGB pixel values. The mode plot reports min, mean and max accuracies. The plots show the accuracies for N 128.
warm up in the BERT architecture, hinting at necessity for learning rate warm 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 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 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 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 Section5511.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 remain similar throughout all modifications, cf. Table T
accuracies
For completeness, we also provide results for just replacing the softmax BERT with normalization in Figure 14] that is. 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 hyperparameter changes, albeit not performing well
20

---

Under review as a conference paper at ICLR 2022
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 focus local information.
D.2 REGULARIZATION EXPERIMENTS
To limit the number of variables which not accounted for in the experiments, focus 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, more devices create constant stream of data and training is more limited by the available training time than the available data. This regime allows In fact, our supplementary experiments below as well as related work (Lan et al. 2020) show that regularization does help in this regime. We leave a comparison of the architectures the limited data regime to future work.
omit regularization in all architectures over-fitting is not issue.
Here, 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-parametee (d = 128, N 128) and train random seeds per learning rate for models with regular-
2, M ization, varying the dropout rate in (0.0625, 0. 2,125, 0.25, 0.5} and the T.2 regularization weighting in {0.0001, 0.001, 0.01, 0.1}. Tables 2 and report the best mean accuracy achieved with the small number behind the accuracies indicating the regularization used, referring to the smallest, to the largest. We underline the results where regularization accuracy. Note however that these improvements should be taken with 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 be overestimated, the is taken 40 averages (4 regularization values times 10 learning rates) as compared to 10 averages (10 learning rates) in the unregulated case.
lead to an improvement in mean
Overall we note that none of the architectures consistently benefits from regularization in our setup and regularization often decreases mean performance. Further, point out that the best perfor- mance with regularization is most of the times achieved with the smallest regularization.
Table 2: Regularization results in the case distinction task with output taken across all tokens. The top three rows correspond to the best mean training accuracy, while the bottom three rows corre- spond to the best mean validation accuracy when validating sequences of half the length.
BERT
MTE
NAP
NON
sum
max
unregularized with dropout 98.1%!
99.3%
99.1%
99.3%
99.1%
98.9%
99.2%
97.3%
97.8%
97.5%
97.3%!
98.2%
with L2- L2-reuularization
99.3%?
99.2%
99.2%?
99.2%
98.9%
99.4%
unregularized
95.5%
95.5%
97.0%
95.3%
75.0%
97.1%
with dropout 94.4%1
94.6%
96.8%%
96.0%
83.1%
96.3%
with L2-regularization
97.2%2
93.6%
97.1%
96.1%?
67.7%
97.2%
Table 3: Regularization results in the case distinction task with output from the first token. The top three rows correspond to the best mean training accuracy, while the bottom three rows correspond the best mean validation accuracy when validating on sequences of half the length.
BERT
MTE
NAP
NON
max
unregularized
36.6%
66.5%
94.5%
23.2%
22.8%
97.8%
with dropout
44.9%1
44.3%1
85.0%1
23.2%1
22.6%1
92.6%1
with L2-regularization
36.0%²
55.3%!
93.8%?
22.8%
22.8%
95.4%!
unregularized
36.7%
50.6%
83.9%
29.6%
28.5%
88.5%
with dropout
41.4%2
40.7%1
74.6%
29.6%
28.9%
87.8%1
with 722rrgullarization
37.2%
45.7%1
82.5%
28.9%1
29.0%1
81.0%1
21

---

Under review as a conference paper at ICLR 2022
Table Regularization results in the mode finding task. The top three rows correspond 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
max
unregularized
99.6%
99.8%
99.6%
98.7%
99.8%
14.4%
with dropout
93.9%
93.3%1
94.3%1
91.8%
93.3%1
24.5%
with I22regularization
99.5%1
99.9%2
99.7%1
98.8%
99.9%1
14.4%
unregularized 95.3%
95.4%
94.9%
91.3%
95.8%
13.5%
with dropout
94.8%
95.4%%
93.8%
92.6%
95.7%
13.4%
with L2-regularization
94.7%
96.0%2
94.9%*
94.7%²
95.8%
13.7%
E CASE LEARNING CURVES
Figures I5 and 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, few additional insights can be noted: (1) Cases are mostly learned in the order of their occurrences (recall that 72.37% of the examples
from the argmin case, 20.09% 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 always given 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 curiosity that occurs when the model dimension 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 different learning rate. the opposite holds seen in the plot just below the highlighted plot. We highly encourage an interested reader to check out our code release which includes these results as well as visualization scripts to inspect them further.
'See supplementary material.
22

---

Under review as a conference paper at ICLR 2022
BERT
argmin
first
argmax
0.3
0.027
Initial Learning Rate
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0c-5
5.9c-6
16
64
256
512
1024
Model Dimension
MTE
argmin
first
argmax
0.3
0.09
0.027
Initial Learning Rate
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
16
32
64
128
256
512
1024
Model Dimension
Case accuracies over the course of training the frgmmin-fist-aargaax case distinction task with output across all tokens, cf. Section 5. Each small sub-plot shows the case accuracies (y-axis, bottom is set 0%, top to 100%) over the course of training (x-axis). Solid lines represent the mean accuracy over the random seeds while shaded areas fill the spread between min- and max-accuracy achieved. Models BERT and MTE are shown here, cf. Figures 16 and 17
Figure 15:
23

---

Under review as a conference paper at ICLR 2022
NAP
argmin
first
argmax
0.3
0.027
Initial Learning Rate
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0c-5
5.9c-6
16
64
256
512
1024
Model Dimension
NON
argmin
first
argmax
0.3
0.09
0.027
Initial Learning Rate
0.008
0.002
7.3c-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
16
32
64
128
256
512
1024
Model Dimension
Case accuracies over the course of training the frrmtin-fist-argmaax case distinction task with output across all tokens, cf. Section 5. Each small sub-plot shows the case accuracies (y-axis, bottom is set 0%, top to 100%) over the course of training (x-axis). Solid lines represent the mean accuracy over the random seeds while shaded areas fill the spread between min- and max-accuracy achieved. Models NAP and NON shown here, cf. Figures 15 and 17
Figure 16:
24

---

Under review as a conference paper at ICLR 2022
sum
argmin
first
argmax
0.3
0.027
Initial Learning Rate
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0c-5
5.9c-6
16
64
256
512
1024
Model Dimension
max
argmin
first
argmax
0.3
0.09
0.027
Initial Learning Rate
0.008
0.002
7.3c-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
16
32
64
128
256
512
1024
Model Dimension
Case accuracies over the course of training the frrmtin-fist-argmaax case distinction task with output across all tokens, cf. Section 5. Each small sub-plot shows the case accuracies (y-axis, bottom is set 0%, top to 100%) over the course of training (x-axis). Solid lines represent the mean accuracy over the random seeds while shaded areas fill the spread between min- and max-accuracy achieved. Models and max shown here, ef. Figures 15 and 16
Figure 17:
25

---

Under review as a conference paper at ICLR 2022
BERT (99.500.18)
MTE (9440118)
NAP 19440188)
NON (9440J8)
Sum (990400188)
max (99420019))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6c-5
2.0e-5
5.9c-6
-
-MoNinG
-MONI-
-moNIngg
-mONI
-morin
Figure 18: Learning rate (y-axis) vs. model dimension d (x-axis) on the aegmin-frst-argmme case distinction task (with output across all tokens). The pixels' R (red), G (green) and B (blue) values correspond to min-, mean- and max-accuracy, respectively, of the corresponding hyperparameter combination. The plot shows the accuracy when evaluating sequences of length N 128. Crosses indicate the combination for best mean accuracy, reported behind the model name.
BERT (98.10033)
MTE (9766.22))
NAP (9.01049))
NON (781044))
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e
2.0e-5
6.9e-6
8c CONTToN
DONO
-MMO
Ning
Figure 19: Rerun of the experiments of Figure 5 with L = instead of L = 2. Crosses indicate the combination for best validation accuracy, which we report behind the model name.
F CASE DISTINCTION TASK ADDITIONAL RESULTS
Figure 18 shows the accuracy in the case distinction task when evaluating with the same sequence length as used during training, i.e., = 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, specifi- cally and Transformer-layer deep models. The results do differ which is why 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 explicitly vary the depth up 64 Transformer-layers.
F.2 VARYING BATCI SIZE
In Figure we provide the case accuracy results of 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 size 128 batch size in which each batch contains in expectation several examples per case, but fewer updates are made the network param- eters. Besides replicating several insights made in the main text, this experiment shows: (1) smaller batches require smaller learning rate, supporting 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.
argument that parameters should not
26

---

Under review as a conference paper at ICLR 2022
BERT (98.240.25)
MTE 1911)000))
NAP (9664338)
NON (764058)
(84346228)
max (9784022))
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
6.9c-6
OO -MONIn Nine
King
PONT -iin
WONT 1N
SONT2OND
Figure 20: Rerun of the experiments of Figure with L instead = 2. Crosses indicate the combination for best validation accuracy, which report behind the model name.
BERT
MTE
NAP
NON
SWm
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
6.9e-6
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6c-5
2.0e-5
6.9c-6
case distinction task (with output across tokens). RGB pixel values correspond to argmin-, first- and argmax- case-accuracies, respectively. Top row: Training accuracy (sequence length N 128). Bottom row: Validation accuracy on sequences length N
Figure 21: Learning rate (y-axis) vs. batch size (x-axis) the
F.3 VARYING INITIALIZATION SCALE
Figure we provide the case accuracy results of an additional experiment, varying the initial- ization standard deviation of the truncated normal distribution with which initialize the weight matrices. The results show additionally to the observations made previously that the softmax based models BERT and MTE struggle learn if the initialization is large. This is likely due the fact that the softmax is already saturated the initialization too large. Also, larger initializations seem require 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 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 argmin- case is learned. We can see that BERT struggles to learn the argmin-case (which requires global information) soon it is not the majority case anymore.
27

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
Sum
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2c-
6.6e-5
2.0c-5
5.9e-6
3e
Ie
Se
0.02
0.00
0.28
5.12
3e-
1e:3
5e.
0.02
0.08
1.38
5.12
3c-4
ie
5e
0
0.28
5.12
3e
5e
0.0
0.00
5.12
3e-4
0.0
5.12
3e
1
5e
0.
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
6.9e-6
Figure 22: Learning rate (y-axis) initialization scale (x-axis) on the case distinction task (output across tokens). RGB pixel values correspond to the case-accuracies. Top row: Training accuracy (N 128). Bottom row: Validation (N 64).
BERT
MTE
NAP
NON
Sum
max
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6c-5
2.0c-5
5.9e-6
9.0%
18
26
72,4
86,2
03.1%
96.5
98.3 %
9. 3.0%
18.1
36.2%
12. 12.4%
06.2
33.1%
96 06.5%
98.3%
9.0%
18.
1624%
72.4%
86.2%
93 031
96.3
98.3%
9. 9.0%
18.1 16
36.2%
12.4 72.4%
86.2 26.256
15
93.1
06
08.3%
9.0%
18.1 96
16.2%
72.4%
86.2%
03.122
96.5
98.3%
9. 1.0%
18.1 16
36.2%
12.4 1139
85.2%
93.1
06.55
08.3%
98.5
Figure 23: Learning rate (y-axis) vs. percentage of identity-case (x-axis). R identity accuracy, G
and B argmin accuracy.
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/raainng accuracy (NV = 128) while the bottom 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.
corresponds to out-of-distribution generalization accrucay when validating
F.6 FIRST TOKEN OUTPUT VARYING DEPTH
In this section we investigate whether our results are tied the shallow architecture of L, Transformer layers. We therefore vary the number of Transformer-layers I. 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-layerr is increased to = 4. However, the performance degrades further increase the depth. (2) The NAP architecture achieves higher best mean accuracy and performs well on a wide range depths.
(3) The max architecture performs well on the biggest range of hyperparameters. This is due to the beneficial architectural prior discussed earlier.
28

---

Under review as a conference paper at ICLR 2022
BERT (72.700.%) MTE 19934188))
NAP (94111188)
NON (9.30-98))))
(3094-0608)
max (1177138))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6c-5
2.0e-5
5.9c-6
1024
1024
1N6
1000
-MON
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9e-6
Figure 24: Learning rate (y-axis) vs. model dimension d (x-axis) on the case distinction task with output from the first token. RGB pixel values correspond to min, mean and max accuracy. Top row: Training accuracy (sequence length N 128). Bottom row: Validation aceuracy when validating on sequences of half the length (N = 64). Crosses indicate the combination for best validation accuracy, which we report with standard deviation behind the model name.
BERT (77.34:1
MTE 182.721868))
NAP 11.12199))
NON (32.5156))
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2c-4
6.6e-5
2.0c-5
5.9e-6
C 2000
32
64
2
8
16
32
2
8
16
32
64
88
16
16
0.3
0.09
0.027
0.008
0.002
7.3c-
2.2c4
5.6e-5
2.0e-5
6
Figure 25: Learning rate (y-axis) vs. Transformer-layers L (x-axis) on the case distinction task (output from the first token). RGB pixel values correspond to min, mean and max accuracy. Top row: Training accuracy (sequence length = 128). Bottom row: Validation accuracy when validating on sequences of half the length (N = 64). Crosses indicate the combination for best validation accuracy, which we report with standard deviation behind the model name.
G MODE FINDING TASK ADDITIONAL RESULTS
G.1 VARYING VOCABULARY SIZE
Figure 26 shows the results of an additional experiment, varying the vocabulary size while keep- ing the sequence length N 128 constant during training. For this experiment, also vary the total number of training steps and set 400 S, to keep the number of examples seen per vocabu- lary token approximately constant. We also include zero-shot generalization results when testing
29

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
Swm
0.3
0.09
0.027
0.008
0.002
XXX
7.3e-4
2.2c-4
XXXX
6.6e-5
2.0c-5
5.9e-6
N NT
12
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
X
5.0e-5
6.9e-6
Figure 26: Learning rate (y-axis) vs. vocabulary size S (x-axis) on the mode finding task. RGB pixel values correspond to min, mean and max accuracy. Top row: Training accuracy (sequence length N 128). Bottom row: Validation accuracy when validating sequences of twice the length (N 256). Crosses indicate the learning rate for best mean accuracy, reported in Table
Table 5: Best mean accuracy per vocabulary size, taken from the combinations indicated in Fig- ure 26 First six rows correspond training accuracies, bottom six rows correspond validation accuracies. Bold numbers indicate min-accuracy higher than the best max accuracy of all other models.
S=2
S =
S 16
S 32
S 64
S 128
S 256
BERT
100%
99.9%
99.9 %
92.1%
72.5%
76.2%
77.4%
74.4%
MTE
100%
100%
99.9%
99.8%
99.3%
97.3%
73.3%
64.9%
NAP
100%
99.9%
99.8%
99.6%
99.7%
99.6%
97.4%
84.6%
NON
100%
99.9%
99.2%
97.3%
74.7%
71.5%
65.2%
61.5%
Sum
100%
100%
99.9%
99.8%
99.7%
99.2%
97.5%
60.6%
max
55.7%
30.1%
17.3%
10.4%
6.6%
4.6%
3.6%
3.1%
BERT
100%
98.2%
95.8 %
88.0%
65.7%
68.6%
68.0%
53.0%
MTE
99.2%
98.4%
96.1%
93.6%
90.5%
85.4%
61.6%
38.9%
NAP 99.6%
98.4%
95.8%
93.1%
90.6%
90.4%
84.3%
64.3%
NON
100%
97.7%
93.1%
85.7%
66.4%
58.3%
50.2%
46.4%
sum
99.0%
97.9%
96.7%
94.4%
91.6%
89.1%
85.9%
45.6%
max 53.8%
29.3%
16.1%
9.5%
6.0%
4.2%
3.0%
2.1%
sequences of twice the length (N 256). Compared to the case distinction task we can do such 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 tokens. This is expected, as the mean subtraction in the nor- malization effectively removes the task relevant information (the mode) in this case. Note however, that for high enough learning rate, the model learns to use the bias parameter b introduced in the normalization effectively reverting to 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, ef. Table
30

---

Under review as a conference paper at ICLR 2022
Cases all tokens (988850035)
Cases first token (88.2:0.35)
Mode (9.71.0.8)
0.3
0.09
0.027
0,008
0.002
X
7.3e4
2.2e4
6.6c-5
2.0-5
6.9c-6
1 D
30
MENTXON iing
102
~iie
Figure 27: Learning rate (y-axis) vs. model dimension d (x-axis) in the different task setups with learned aggregation weights. The plots from left right are grouped according the tasks: The case task with outputs across all tokens (cf. Section 5.1). the case task with outputs from the first token (ef. Section 5.2) and the mode task (cf. Section 5.3). For the case tasks, the left sub-plot reports min, mean and max accuracies while the right sub-plo reports mean case accuracies as RGB pixel values. The mode plot reports min, mean and max accuracies. The plots show the accuracies for N = 128.
BERT (0011)005)
MTE (911008)
NAP 19:12 188)
NON (99 1000))
(92. 42.31444)
(9799002))
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
6.9c-6
-MNO
NNN -MC -N0
NNN
Figure 28:
Full results of the GNN experiments on the intteiinnninnnn task. The plot shows variation of learning rate (y-axis) vs. model dimension (x-axis). RGB pixel values correspond to min, mean and max validation node classification F1-score. Crosses indicate the best combination (validation), from which we report the test performance behind the model name.
interaction
H LEARNED AGGREGATION WEIGHTS
Here we present the results of additional model for fixed aggregation size. In particular, ask how the results look if we learn the aggregation weights as parameters of the network. For this model, take the model and replace the sum-broadcast operation with fully connected linear layer across the sequence dimension. This is akin to the Random Synthesizer presented by TTy et al. (2021), but without the softmax. The results presented in Figure 27 confirm conjecture that the softmax limits the information flow, as this new model more robust to hyperparameters than BERT and MTE.
I PPI EXPERIMENT
For the PPI node classification task we train GNN based on our architectures with aggregation over neighbors instead of the whole sequence. Here use L transformer layers, as the number of layers specifies the k-hop neighborhood from which information is aggregated a GNN (k = L). The remaining settings are left the default (M 4 attention heads, regularization). We train with one graph per update for 32 epochs.
All runs repeated with 5 different random seeds. We track the validation performance over the course training. Figure I0 shows the highest validation performance each hyperparameter combination, with the results of all models given in Figure 28 For the test performance reported
31

---

Under review as a conference paper at ICLR 2022
BERT (99.3-00228) MTE (9440028)
NAP 19440118)
NON (9944038)
(9600-1228))
max (996-0228)
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
6.9c-6
XONTWONT
WONT ONt -MONIn Monkin
1624
SONTRONT NIN
Sin
Figure 29: Full results of the RL experiments. The plot shows a variation of learning rate (y-axis) vs. model dimension d (x-axis). RGB pixel values corresonnd min, mean and max success rate when evaluating on 10,000 newly generated games. Crosses indicate the combination for best success rate, which report with standard deviation behind the model name. Note here that training is stopped early a success rate of 99% is reached.
we take the hyperparameter combination with the best average validation performance and use the checkpoints with highest validation performance over the course of training.
J BABY AI EXPERIMENT
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 factored representation of the environment and let the agent learn, what it needs 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 simple linear layer each. Also, we replace the ReLU Iineariessss used with GeLUs. These modifications are done reflect similar setup investigated the rest of the paper. Similarly and set the "WMG Memo size" equal to the model dimension d. We chose level 3 of the Baby-AI environment (Chevalier-Boisse Boisvert al. 2019) based on the results presented by Loynd et al. (2020). That is, we chose 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
fix M 4, L
investigated are given in Figure 29 different random seeds.
Also here we ran every hyperparameter combination with
Table 6: Hyperparameters used in the reinforcement learning experiment. These are based on the hyperparameters which Loynd et al. (2020) report their Appendix in Table
Hyperparameter
Value
A3C tmaz
Adam eps
le-12
Discount factor
0.95
Entropy term strength 3 Gradient clipping threshold
0.1
128.0
Reward scale factor
8.0
WMG Memos
32

---

Under review as a conference paper at ICLR 2022
K BROADER IMPACT STATEMENT
Our work contrast different architectures on an abstract level. Due to the abstract nature of the study, there is direct risk associated with system failure. However, would like mention potential implications of the findings. Our main proposal is an architecture which is more robust hyperparameters. This has potential positive benefits on the environment as less computation needs to be invested to reach 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 underrepresented in the data. In general, we see it vital to get a better understanding how different architectures learn to represent different biases in the data and hope that our work provide stepping stone in this direction.
L FIGURE REPLICATION
Since way visualize the results uses the full color spectrum, 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 get in contact with us if they face any other difficulty reading our plots. We replicate Figures [0 in Figures [3 36) and the Figures of the Appendix thereafter. We also invite interested reader to check out our code release, which allows for further visualizations of all results.
BERT (177404%)
MTE (778)03))
NAP
NON
(8155638))
(97.8
0.09
0.027
0.008
0.002
4-3
2.2c-4
6.6e-5
2.0e-5
5.9e-6
16
32
102
66
Ning
1024
00
1966
512
1024
0.3
0.09
0.027
0.008
0.002
--3e-
2.2c-4
6.6e-5
2.0e-5
5.9e-6
286
32
16
178
1024
16
32
128
1024
10000 Nin
1024
8 XONT
102
8
256
512
1024
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2e4
6.6e-5
2.0e-
5.9e-6
1024
TOKE
Figure 30: Replica of Figure 5 -ignin-firt-aargmax case distinction task (with output across all tokens). The rows from top to bottom correspond to min- mean- and max-accuracy, respectively. The plots show the validation accuracy when validating on sequences of length N = 64. Crosses indicate the combination for best validation accuracy, which we report with standard deviation behind the model name.
Learning rate (y-axis) vs. model dimension d (x-axis) on the
33

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
SUM
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9e-6
+2
72
1
160
16
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9e-6
128
2
12
512
420
16
G
1 038
25
512
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2e4
6.6e-5
2.0e-5
6.9c-6
1 -N
-N -
-NI S
-NI -N
Figure 31: Replica of Figure (top row): Biased data results the case distinction task (with output across all tokens). The rows from top to bottom correspond to argmin, first and g accuracies, respectively, Shown is the learning rate (y-axis) vs. sequence length N (x-axis).
argmax-mean-case-
34

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9e-6
9.0%
18 11
36
72.4%
86 26
93.1
96 3%
98.3 13.3%
9 9.0%
18.1%
36 25
71.45 49.4%
86
93.12
26 96.5%
98. 98.3%
9.0%
18.1 8.1%
36 222
72.4%
86 26
93. 0311
96.35
98.3%
9.0%
18.1% 1%
36 26
21.49
86 24
93.1%
96
34.34
9.0%
18. 1.
36 24
72.4%
16.25%
03. 13.1
96.52
98.3%
9.0%
18.1%
36 25
21.4%
86 24
1.
96.5 3/
98. 36.34
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0e-5
6.9e-6
9.0%
18. 8.1%
36.2%
72.4%
86.2%
93.1%
96.5
98.3%
9.0%
18.1%
36.2%
72.4%
862
93.15
96 6..5%
98.3%
9.0%
18. 88155
3629
27.4%
86.25
0311%
96.5
98.3%
9.0%
18.1
3622%
71.49
86.2
93.1%
96.5%
98.3%
9.0%
18. 15
362 2525
72.4%
93.1%
96.5
98.3%
9.0%
18.1%
36.258
212455
36.2%
93.1%
96.3%
98.3
0.3
0.09
0.027
0.008
0.002
-.3e-4
2.2c
6.6e-5
2.0e-5
5.9e-6
9.0%
18
26
862
96.3
98.3%
9.0%
18.1
36.2
F1.4%6
862
93.1
96 5%
98.3%
9.0%
18.15 151
3622
72.4%
86.25
03. 5%
96.5%
98.3%
9.0%
18.1
36.2%
P1.45
86.2
9315%
06.3%
8.3%%
9.0%
18.15
3622
36625
96.55
98.3%
9.0%
18.1
36.2%
P1:45
35.2%
13.55
06.5 3665
98.3
Figure 32: Replica of Figure (bottom row): Biased data results on the case distinction task (with
output across all tokens). The rows from top to bottom correspond to argmin, first and argmar-mean-
aase-accura,ies, respectively. Shown the learning rate (y-axis) vs. percentage of argmin-case in the data (x-axis) with fixed N = 128.
35

---

Under review as a conference paper at ICLR 2022
BERT
MTE (75842533)
NAP
NON
0.09
0.027
0.008
0.002
7.3e-4
2.2e44
5.6e-5
2.0c-5
6.9e-6
1
32
64
28
1022
64
1024
16
64
512
1024
0.3
0.027
0.008
0.002
7.3c-4
2,2c-4
6.6e
2.0c-
6,9e-6
64
1024
60
3
1027
1024
1024
358
512
1024
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0=5
5,9e-6
1024
1027
1024
IMONN
1024
1024
-MON
1024
Learning (y-axis) model dimension d (x-axis) the case distinction task with output from the first token. The rows from top to bottom correspond to argmin, first and argmax- mrema--men-aeee mean- cacnn-ccceeaasaee accuracies. Crosses indicate the best accuracy, reported behind the model name.
Figure 33: Replica of Figure
7
36

---

Under review as a conference paper at ICLR 2022
BERT (9994019) MTE (000+000)
NAP
NON (994018)
(10004000))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
102
1024
1024
102
102
12
250
512
1024
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
1024
16
32
-12
1024
64
128
150
1024
236
$12
1024
16
64
1024
%
37
128
256
512
1024
0.09
0.027
0.008
0.002
7.3e-4
2.2e44
6.6e-5
2.0c 2.0e-
5.9e-6
ANinO 1
Figure 34: Replica of Figure Learning rate (y-axis) vs. model dimension d (x-axis) on the mode finding task. The rows from top bottom correspond to min, mean and max accuracy. Crosses indicate the combination for best accuracy, reported behind the model name.
37

---

Under review as a conference paper at ICLR 2022
BERT (9244004%
MTE (0091107))
NAP (08.2.044))
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0-55
5.9c-6
166
1
64
1024
16
32
64
128
256
312
1024
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
64
128
255
5
1024
32
64
128
256
512
1024
8
16
32
64
128
256
312
1024
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.20-4
6.6c-5
2.0e-
5.9c-6
1024
1024
$
1024
Figure 35: Replica Figure 9 Learning rate (y-axis) vs. attention heads M (x-axis) on the case
distinction task (first token output). The rows from top to bottom correspond to min, mean and max accuracy. Crosses indicate the combination for best accuracy, reported behind the model name.
38

---

Under review as a conference paper at ICLR 2022
BERT (9.0)))0.0%)
MTE (99.1)00%)
NAP (99.190.115)
BERT (99.3 0.22%)
MTE (99.4-10.2%
NAP (99.4
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.60-5
2.0e-5
5.9c-6
128
102
3
64
23
230
1024
39
23
102
102
125
250
312
1024
0.3
0.09
0.027
0.008
0.002
7.3e4
2.2e-4
6.6c-5
2.0e-5
6.9e-6
64
128
256
1024
64
78
1024
16
23
200
102
182
1024
1024
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2e-4
5.6e-5
2.0e-5
5.9e-6
8 9
TO0 MONN
102
ANinO
Figure 36: Replica of Figure Learning rate (y-axis) vs. model dimension (x-axis). The rows from top to bottom correspond to min, mean and max validation performance. Crosses indicate the best combination (validation), from which we report the test performance. Left: Protein-protein- interaction task. Shown is the node classification Fl Fl-score. Right: Altered working memory graph agent the Baby-Al level 3 reinforcement learning task. Shown the success rate.
39

---

Under review as a conference e paper at ICLR 2022
BERT
Warmun up
grad. clip
+
normalize
GELU
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
6.9e-6
1024
102
16
37
66
22
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
5.60-5
2.0e5
6.9c-
1022
0.3
0.09
0.027
0.008
0,002
7.3e-4
2.2e
6.60-
2.0e-5
5.9e-6
BERT
warm warmmmm up
grad. clip
normalize
GELU
0.3
0.09
0.027
0.008
0,002
7.3c-4
2.2e-4
6.60-5
2.0e-5
16
6
1024
102
1024
0.3
0.09
0.027
0.008
0,002
7.3c-4
2.2e-4
6.6c-5
2.0c-5
5.9c-6
16
1024
102
8
16
37
64
228
250
$12
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2e-4
6.6c-5
2.0-5
5.9c-
(top two rows): L Learning rate (y-axis) vs. model dimension d (x-
Figure 37: Replica of Figure 13
axis) on the case distinction task (with outpur across all lokens) ablation study. The rows from top to bottom corres min. mean and max  for N = 128 (lop 3 rows) and N 64 (bottom 3 rows). Crosses indicate the combinaama for best accuracy. reported in Table
Crossoon
40

---

Under review as a conference paper at ICLR 2022
BERT
Warmuu up
grad. clip
+
normalize
GELU
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
6.9c-6
16
128
1024
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
5
2.0e-
6.9c-
0.3
0.09
0.027
0.008
0,002
7.3e-4
2.2e
6.60-
2.0e-5
5.9e-6
BERT
warm warm up
grad. clip
normalize
GELU
0.3
0.09
0.027
0.008
0,002
7.3c-4
2.2e-4
6.60-5
2.0e-
5.9e-6
16
102
102
0.3
0.09
0.027
0.008
0,002
7.3c-4
2.2e-4
6.6c-5
2.0c-5
5.9c-
1
16
1024
102
6
250
(024
16
37
6
12
750
$12
(024
0.3
0.09
0.027
0.008
7.3c.
0.002
7.3c-4
2.2e-4
6.6c-5
2.ec-
5.9c-6
Figure 38: Replica of Figure 13 (bottom rows): Learning (y-axis) vs. model dimension (x-axis) on the case sase distinction task (with output across all tokens) ablation study. The rows rows from top to bottom correspond argmin, first and merma accuracy for N 128 (top 3 rows) and d
-oomt N = 64 (bottom tows). Crosses indicate the combingg for best aecurrayy, reported in Table
41

---

Under review as a conference paper at ICLR 2022
Cases all tokens (9.44019)
Cases first token (9224.18)
Mode (00.4220.1
10
102
16
125
1024
COMONIT -NNO
-NA
Learning rate (y-axis) vs. model dimension (x-axis) in the different task setups when just replacing the softmax in BERT with normalization. The plots from left to right are grouped according to the tasks: The case task with outputs across all tokens (cf. Section 5.1 the case task with outputs from the first token (cf. Section 5.2) and the mode task (cf. Section 5.3). For the case tasks, the left sub-plots report min, mean and max accuracies (from top to bottom) while the right sub-plots report mean accuracies (argmin, first and argmax from top to bottom). The mode plots report min, mean and max accuracies. The plots show the accuracies for N 128.
Figure 39: Replica Figure 14
42

---

Under review as a conference paper at ICLR 2022
BERT (99.520.19)
MTE (944018)
NAP 1944011)
NON (9940019)
sum (99000.19) 18)
max (9942011)
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.60-5
2.0c-5
5.9e-6
12
10
3
2ONTTSONN
1024
98
256
51
1024
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2e-4
6.6e-5
2.0e-5
5.9e-6
16
1024
10
1000000
16
1024
-00
IMO
100
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
Figure 40: Replica of Figure 18 Learning rate (y-axis) vs. model dimension d (x-axis) on the argmin-finttaargaa case distinction task (with output across all tokens). The from top to bottom show min, mean and max accuracy for N = 128. Crosses indicate the combination for best accuracy, reported behind the model name.
43

---

Under review as a conference paper at ICLR 2022
NON Mesaaaee sum (
BERT 33) MTE (920)
NAP (80+048)
max (0.5))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.60-5
2.0c-5
5.9e-6
6
1022
10
1022
10 10
10
256
1024
1022
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.60
2.0c-5
6.9e-
16
32
120000
1024
1024
1024
102
102
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5 5
2.0c-5
5.9e-6
BERT 082242255
MTE
AIEEER 1.00)
NAP 07.6433))
NON
sum (84 (8433-625))
max 078+022)
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
102
16
3
1024
KONTMONT
102
1627
1027
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
5.9c-6
64
1286
1024
16
32
64
1024
8
16
64
1024
16
64
128
256
512
1024
8
16
64
28
T024
125
250
1024
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
Learning rate (y-axis) model dimension d (X-axis) of fir----tttttaaa case distinction task (with output across all tokens) with L = 4 (top rows) and L 6 (bottom rows). The rows from
Figure 41: Replica of Figures 19 and 20 the argmin-first-un
N = 64. Crosses indicate the combinas from top bottom show min, mean and max accuracy for Conoinatioono for bes accuracy reported behind the model name.
44

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
Sum
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0e-5
5.98
6
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9c-6
1
64
128
128
128
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c
6.6e-5
5 9
2.
5
BERT
MTE. M
NAP
NON
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2c-4
6.6e-5
2.0e-5
5.9c-6
128
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2e-4
6.6e-5
2.0e-5
6.9c-6
64
128
N
0.3
0.09
0.027
0.008
0.002
7.3e4
2.2e-4
6.6e-5
2.0e-5
6
Learning rate (y-axis) vs. batch size (x-axis) on the case distinction task (with output across all tokens). The rows from top to bottom correspond to argmin, first and argmax
Figure 42: Replica of Figure 21
indicall the commaeee combination for best
for N = 128 (top 3 rows) and N 64 (bottom 3 rows). Crosses dest accuracy. reported behing the mode name.
45

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
SUm
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0e-5
6.9c-6
0.0
36
0.
5c
0.00
0.3
0.09
0.027
0.008
0.002
7.3e4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
0.00
0.00
0:200
5.12
1e.3
5e
0.02
0.00
0 0000
3e
0 0.00
0:
123
5.12
3e
1o
5e
0.02
0.08
0.23
5.12
3e
1c:
58
0
0. 0:35
3e
20.0
0.00
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.00 2.0e-5
5.9c-6
BERT
MTE
NAP
NON
Sum
max
0.09
0.027
0.008
0.002
7.3e-
2.2e-4
6.6e-5
2.0e-5
6.9e-6
le
1000
5.12
3e-4
le
5e.3
0.02
0.08
0.32
1.8
5.12
3e-
le
56
0.3
0.09
0.027
0.008
0.002
7.3e-
2.2e-4
6.6e-5
2,0e-5
6.9e-6
3e.
0.02
0.00
0.32
1288
5.12
3e.4
16-3
5e:
0.0
0.3 0,33
5.12
3c-4
le:
5.12
3e-4
1e
5e:
0.0 00080
0.5
Be-4 T
5.1
0.3
0.09
0.027
0.008
0.002
7.3e4
2.2e-4
6.6c-5
2.0e-5
5.96-
Figure 43: Replica of Figure 22 Learning rate (y-axis) vs. initialization scale (x-axis) on the case distinction task (output across tokens). The rows from top to bottom correspond to argmin, first and argmax- ceccucccss--cccurccies N = 128 (top 3 rows) and N 64 (bottom rows).
46

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
SUm
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
00
1
18.
36.26
72.4%
86.25
03.19
96.59
98.3%
9. 0.00%
18. 18.190
36.24 36.2%
72.49
86.25
93.15
96.5%
98.3% 38.3%
9.0%
18.1%
36.20
36.38
72.4%
86.2%
93.13
96.35
98.3%
9.0%
18. 18.12
36.2%
72.4%
66.2%
93.1%
96.
98.3%
9.0%
18.1% 19
36228
72.4%
$6.2%
93.1
96.36
98.3%
2.0%
18. 18.1%
36. 26.25
27.38
36.2%
93.1%
96.
98.30
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e
2.0c-5
6.9e-6
000
18
36.26
36.2
72.4 4.
86 2..
03.1%
96 5%
98.3%
9. ..0%%
18. 18.156 166
16.2%
72.4 71.4%
86.2 36.2%0
93.1
96 5.0
98. 88.3%
9.0%
18. 19
36 260
72.4%
86.2% 2..2%
93. 93.1%
06.5 5%
98.3%
9. 0.0%
18. 18.15e
36.2%
72.4 71.4%
256
93.1%
96.5
98.3%
9.0%
18. 18.1% %
72.4%
86.2%
03.1%
06 36.39
98.3 98.3%
9 1.0%%
18. 18.1% 166
36 756
22. 71.3%
93. 56
96
96. 96.3%
Figure 44: Replica of Figure 23
Learning rate (y-axis) vs. percentage of identity-cass (x-axis). Top
identily accuracy. Bottom
argmin accuracy.
47

---

Under review as a conference paper at ICLR 2022
NON 1aa1ee
(923) MTE 2533)
28)
sum 53))
max (98620.18) (1))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.60-5
2.0c-5
6.9c-6
64
128
250
$12
1024
16
32
64
128
256
1022
16
1027
2
1022
12
350
0.09
0.027
0.008
0.002
7 7-3c-4
2.2c-4
6.60
2.0c-5
6
1000
512
1024
16
28
250
1021
8
16
32
12
1025
102
102
60
166
2
1024
0.3
0.09
0.027
0.008
0.002
7. 7.3e-4
2.2c-4
6.6e-5 5
2.0c-5
5.9e-6
BERT (7274070 MTE 19334188))
NAP (811111))
NON (69 $2099)))))
sum (309+066)
max (91.7+33))
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
7024
8
16
37
64
23
250
$12
(024
32
64
102
$12
1024
1027
16
33
64
256
$12
0024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
6.6e-5
2.0c-5
5.9c-6
30
64
128
512
1024
16
32
64
128
256
(024
16
32
64
28
1024
16
32
64
128
256
512
1024
8
16
64
88
[024
33
128
250
1024
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0c-5
5.9e-6
Figure 45: Replica of Figure 24 distinction task (first token output). The rows from top to bottom show min, mean and max accuracy for N 128 and then N 64. Crosses indicate the combination for best accuracy, reported behind the model name.
Learning rate (y-axis) model dimension (x-axis) the
48

---

Under review as a confferenct paper at ICLR 2022 pape
BERT (95 1a-44a
MTE
NAP Soafeiiinn
(% Sassssee
NON
Aum (23.6
max (98.3
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c-4
5 5
2.0c-5
5.9c-6
0.3
0.09
0.027
0.008
0.002
4
2.28-
-
6.6c 5
5.9g.
2.0c.3
6
16
32
64
16
0.3
0.09
0.027
0.008
7.30- 1
0.002
4
6.6e-1 5
-
3
BERT T 17734400 Se D
MTTE TE
62274 1818
sum (30.6-000) (
max (93
0.3
0.09
0.027
0.008
0.002
2.2c-4 4
-
-
0.3
0.09
0.027
0.008
72
2,28 1
0380
0.002
4
6.60
$,98
2.0c-
32
163
32
0.3
0.09
0.027
0.008
222
0.002
2.2c4 -
6.6g-
2 -
5.9e-
Figure Figure 46: Replica of Figure 25 Learning rate (y-axis) vs. Transformer L (x-axis) on the case aattmnettt case
(eeespooie to min, rereedee and N = 64 (botton 3 rows) Crosse (oo
task
(
(output from the first token). The rows from top to bo rows)
to
the c Commattion for for mean mean and may
accurss f
for N = 128 (top 3 rows)
indicale
be be
mean mean aecuracy
sportted
behing the mode name
49

---

Under review as a conference paper at ICLR 2022
BERT
MTE
NAP
NON
Sum
max
0.09
0.027
0.008
0.002
*
XXXX
7.3e-4
2.2c-4
XXXX
6.6e-5
2.0e-5
5.9c-6
C CO 20
NTX0
500
-Ci
0.3
0.09
0.027
X
0.008
X XX
0.002
XXXX
7.3e-4
2.2c-4
6.6e-5
2.0e-5
5.9c-6
2
2
64
D3
256
0.3
0.09
0.027
0.008
0.002
7.3e-4
X
2.2e-4
6.6e-5
2.0e-5
5.9c-6
BERT
MTE
NAP
NON
Sum
max
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2e-4
6.6e-5
2.0e-5
5.9c-6
2 Cr
128
0.09
0.027
0.008
0.002
7.3e4
2.2e-4
6.6e-5
2.0e-5
5.9e-6
0.3
0.09
0.027
0.008
0.002 X
X
7.3e-
2.2e4
6.6e-5
XXXX
2.0e-5
6.9e-6
Figure 47: Replica of Figure 26 Learning rate (y-axis) vs. vocabulary size S (x-axis) on the mode
finding task. The rows correspond min, mean and max accuracy for N 128 (top rows) and N
256 (bottom 3 rows). Crosses indicate the learning rate for best mean accuracy, reported in Table
50

---

Under review as a conference paper at ICLR 2022
Cases all tokens 093%%
Cases first token (88.2+0.3%)
Mode (99 19.720.18)
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6c-5
2.0e-5
5.9c-6
1024
102
0.3
0.09
0,027
0.008
0.002
7.3c-4
2 4
6.60-5
2.0e-5
6.9c-6
30
64
102
16
32
64
128
312
1024
16
32
$17 1024
16
64
28
$1
1024
16
3
64
02
250
5
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2e-4
6.6c-5
2.0e-
5.9c-6
Learning rate (y-axis) model dimension (x-axis) in the different task setups with learned aggregation weights. The plots from left to right are grouped according the tasks: The task with outputs across all tokens (ef. Section 5.1). the case task
Figure 48: Replica of Figure
27
with outputs from the first token (ef. Section 5.2) and the mode task (cf. Section [.3) For the case tasks, the left sub-plots report min, mean and max accuracies (from top to bottom) while the right sub-plots report mean case accuracies (argmin, first and argmax from top to bottom). The mode plots report min, mean and max accuracies. The plots show the accuracies for N 128.
51

---

Under review as a conference paper at ICLR 2022
BERT (00.1)0.05)
MTE 19.1))0))
NAP (90.1:0 18)
NON (99 1000))
(92 42.1404))
(9794022)
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6e-5
2.0e-5
5.9c-6
2O0TOO 2000
102
300 1200000
102
102
12
256
51
1024
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c
6.6c-5
2.0e-5
6.9c-6
64
1024
16
1023
16
1024
16
64
8
1024
8
121
1024
32
64
128
256
$12
1024
0.3
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6c-5
2.0e-5
5.9c-6
Figure 49: Replica of Figure interaction task. The plot shows variation of learning rate (y-axis) model dimension d (x-axis). The rows from top to bottom correspond to min, mean and max validation node classification F1- score. Crosses indicate the best combination (validation), from which report the test performance behind the model name.
28
Full results of the GNN experiments on the protein-protein-
52

---

Under review as a conference paper at ICLR 2022
BERT 20)
MTE 29)
NAP 18)
NON (094103%)
(9601
max (9962022)
0.09
0.027
0.008
0.002
7.3e-4
2.2c-4
6.6c-5
2.0e-5
6.9c-6
300C 10 2000 SNN
KONT MONNTOON
0.3
0.09
0.027
0.008
0.002
7.3c-4
2.2c
6.6c-5
2.0e-5
6.9c-6
16
102
0.3
0.09
0.027
0.008
0.002
7.3e4
2.2c-4
6.6c-5
2.0e-5
6.9c-6
Figure 50: Replica of Figure 29 Full results the RL experiments. The plot shows a variation of learning rate (y-axis) vs. model dimension (x-axis). The rows from top bottom correspond to min, mean and max success rate when evaluating on 10,000 newly generated games. Crosses indicate the combination for best success rate, which report with standard deviation behind the model name. Note here that training is stopped early if success rate of 99% is reached.
53


Reviews:
[{"review_id": "gELQTQxAE2I", "reviewer": "Reviewer_rjM7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n+   The paper is clearly motivated by the problem that current self-attention is not robust to hyperparameter changing. The authors also give some intuitive examples on this, such as bias towards local information and gradient vanishing.\n+   The paper designs synthetic tasks and does a through analysis of NAP and other baseline models in this setting. Then the authors also experiment in some real-world settings such as GNN and RL.\n+   The visualization method of the results is clever. \n+   The experiment on finding the mode in the input data is inspiring, which implies that information is not effectively routing through all tokens in attention modules.\n+   The paper is well-written and easy to follow.\n\n\nWeaknesses:\n\n-    A few related works are not cited or compared to. For example, [1] replace the softmax in self-attention with a token-wise Guassian kernel. [2] compares self-attention with fast weight programmers and proposes to remove the softmax. Several other works also adopt a similar design [3, 4]. Although these works are focusing on different problems instead of robustness to hyperparameter as in your paper, could you give some discussion on the relation between these prior arts and your model?\n-    The authors did experiments in real-world settings (GNN and RL) in Sec 5.4, but the performance of NAP and BERT / MTE is somewhat similar. They all have the same validation results (99.1% in GNN, and 99.3%-99.4% in RL). The robustness to hyperparameters is also similar between MTE and NAP in RL task. Is it because the tasks are still too small in scale to show the difference between NAP and baseline models? \n-    The normalization on $l^{i, j}_m$ will force the attention weights $a^{i, j}_m$ to have a fixed variance. But for example, if the tokens are all very similar, the attention weights should be similar too (thus having a small variance). Is the normalized attention able to handle this situation?\n\n\nOther comments:\n\n-    Some other works have also focused on stabilize the training process of transformers [5]. Do they have any relation to your work? A short discussion on this line of work could be beneficial.\n\n\n\n\n[1] Lu J, Yao J, Zhang J, et al. SOFT: Softmax-free Transformer with Linear Complexity[J]. arXiv preprint arXiv:2110.11945, 2021.\n\n[2] Schlag I, Irie K, Schmidhuber J. Linear transformers are secretly fast weight programmers[C]//International Conference on Machine \nLearning. PMLR, 2021: 9355-9366.\n\n[3] Shen Z, Zhang M, Zhao H, et al. Efficient attention: Attention with linear complexities[C]//Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2021: 3531-3539.\n\n[4] Cao S. Choose a Transformer: Fourier or Galerkin[J]. arXiv preprint arXiv:2105.14995, 2021.\n\n[5] Xiao T, Singh M, Mintun E, et al. Early convolutions help transformers see better[J]. arXiv preprint arXiv:2106.14881, 2021.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well-motivated by the problem of sensitivity to hyperparameters and propose a sensible algorithm to solve the problem. The experiments on the synthetic task are sound. However, it seems lacking some convincing results on real-world dataset. And the relation to some prior works needs to be clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "bbOL3EDUTMn", "reviewer": "Reviewer_BH5E", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n1. The paper is clearly written and easy to follow.\n2. The implications of using softmax attention are thoroughly analyzed in theoretical and experimental settings.\n3. The set of synthetic experiments conducted highlights the differences between attention mechanisms reasonably well and supports the authors’ claims.\n\nCons:\n1. The authors do not apply their proposed mechanism to the established BERT benchmarks. Instead, they construct a set of synthetic tasks that aim to imitate real-world NLP tasks. Lack of real-world NLP benchmarks makes it harder to believe that the implications of the softmax use actually matter. As the authors mentioned in Appendix C, some softmax-attention constraints may actually be beneficial, and this can’t be confirmed or disproven without real world benchmarks. The synthetic tasks chosen are intuitively similar to the specific NLP problems, but in the end, it’s not clear how well the synthetic results will generalize.\n2. The experimental results demonstrate the advantage of NAP over sum and max pooling, but the case of NAP vs NON is less obvious. For many tasks, NON performs adequately well in terms of final quality and hyperparams sensitivity (in figure 26 it seemingly outperforms NAP for sensitivity). The paper would benefit from providing clear evidence (in addition to Figure 7) that despite all the additional layer norms in MTE, the dynamic normalization of attention logits is still necessary.\n\nSuggestions/questions:\n1. It’d be great to add the learned aggregation weights mechanism into the comparison. The authors explicitly try to avoid depending on the sequence length but having such a comparison may help prove the hypothesis that convex limitations are to blame and moving away from the convex hull is beneficial. For example, Figure 7 may benefit from such comparison.\n2. Due to the use of MTE architecture, it’s not clear how well NAP will perform as a drop-in replacement of the softmax attention. The paper would benefit from NAP experiments that are not MTE-based to confirm that NAP isn’t brittle and can be used in practice.\n3. (Figure 7) The explanation here can be more convincing. Why is MTE worse than traditional BERT on this task? Is it because local information bias is actually beneficial? Why is NON underperforming here and how can it be confirmed (NON doesn’t underperform like this in other tasks)? Figure 23 from the Appendix is more convincing, so it may make sense to move it here\n4. (Figure 9) It’d be great to move the NON, sum, and max results from the appendix here and add the explanation. \n5. (Style) Using RGB pixels for visualization purposes helps with the presentation but is not colorblind friendly. . It may be better to split the plots per-channel.  \n", "summary_of_the_paper": "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.", "summary_of_the_review": "While the practical viability of the proposed method is still uncertain, the insights on using softmax attention from this paper are valuable. The proposed synthetic evaluation framework is also a good foundation for further discussions on softmax-attention alternatives. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "RtPM_f7ETba", "reviewer": "Reviewer_fpkN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The finding of the paper is interesting - the arguments made about the restrictive nature of the softmax operator makes sense. The proposal to replace the softmax with normalization is very simple and could have a lot of implications. The experiments conducted on the synthetic data are thorough and the discussions were made well. However, the authors could have shown gone beyond synthetic experiments and shown more results on real world datasets.", "summary_of_the_paper": "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. ", "summary_of_the_review": "While I enjoyed reading the paper, my biggest concern is lack of real world experiments. I appreciate the authors on conducting experiments on multiple runs, repeating the experiments with different model sizes, visualizing the results effectively and conducting carefully crafted simple experiments that illustrates the specific points being made. The carefully experiments were all designed well. However, for practitioners, a pressing question would be if these architectures can be used as a replacement of the traditional transformer models. Can this be used for training language models or benchmark NLP tasks?\n\nTo this end, it would have been really nice if authors had shown one experiment on real world datasets. I understand that some institutions might not have enough resources to train large-scale models, but a proof of concept at a decent scale on a realistic dataset would have been nice. Such an experiment would tell whether these findings would translate to large scale datasets and tasks. I would also like to point out that the authors have one experiment on protein-protein interaction graph prediction, but I am not aware of this benchmark and I am unable to gain much insights from this experiment aside from the fact that the model works.\n\nWhile this issue is certainly concerning to me, I believe the findings in this paper are a reasonably good contribution that could help others in the community. It's just that the paper would have been really strong had the authors gone beyond synthetic dataset and done some benchmark NLP experiments.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2m0st8y_Ako", "reviewer": "Reviewer_xpRQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThe authors provide a theoretical analysis and proof on using softmax function in attention module for effect of gradient. I did not check the proof precisely. But analyzing neural network modules from a theoretical perspective is interesting and valuable to the community. \n\nThe whole paper is well organized and written, and easy to follow. \n\nThe authors give a comprehensive review and discussion to modern works related to transformer architectures and normalization and pooling methods.\n\nThe code is provided in the supplementary material, which would be helpful for reproducing the experiments in the paper.\n\nWeaknesses:\n\nThe proposed NAP operation is quite related to layer normalization (LN). Instead learning element-wise scale and bias parameters in LN, scalar parameters are learned in NAP, which is supposed to function analogous to LN.  Can it be used as replacement of all the LN layers in the architecture?  If not considering variable length, is using LN operations only (i.e., replacing NAP as LN in the Figure 3) an effective variant compared to MTE and NAP architectures?\n\nMost ablation studies and empirical analysis are conducted on synthetic tasks which are less convincing to demonstrate the effectiveness of the method. Although the authors present the experiments of exploring the proposed method into graph neural networks in PPI task and memory graph agent in reinforcement learning task, the performance improvements (in Fig. 9) are rather marginal compared to baseline methods.\n\n\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The authors theoretically analyze the limitation of using softmax function in attention module and propose an unconstrained surrogate in the paper. However, the empirical evidences shown in the paper are insufficient to validate the method. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies potential drawbacks in using softmax over attention in Transformers and evaluates other normalization approaches. Reviewers, while had been positive about the empirical analysis and the insights from the synthetic data experiments, agree that the paper lacks real world experiments/insights. I agree with that and believe the paper falls short in several areas.

1) Drawbacks of softmax: Paper states several generic drawbacks of softmax such as saturation issue leading to vanishing gradients. However the paper does not demonstrate if Transformer models used in practice suffer from this issue under standard training settings. Even the arguments that attention layer focuses on local information is quite vague and not well supported. Overall the analysis is quite weak without much concrete statements and demonstration in real settings. 

2) Experiments: The paper presents many synthetic experiments evaluating alternates to softmax varying from layer normalization to pooling. There are no experiments showing if the studied variations actually solve the issues discussed in earlier section. Finally due to the lack of any real world experiments (even small scale ones), it is not clear if the results apply in real world settings.

Overall I think the paper needs significant work in formalizing the drawbacks of using softmax in Transformers and demonstrating that the proposed solutions indeed solve this problem.
Decision: Reject