Approximating How Single Head Attention Learns
Authors: Charlie Victor Snell
Ruiqi Zhong
Dan Klein
Jacob Steinhardt
Venue: Submitted to ICLR 2023 (2023)
Keywords: NLP
training dynamics
attention

Abstract: Why do models often attend to salient words, and how does this evolve throughout training? We approximate model training as a two stage process: early on in training when the attention weights are uniform, the model learns to translate individual input word `i` to `o` if they co-occur frequently. Later, the model learns to attend to `i` while the correct output is o because it knows `i` translates to `o`. To formalize, we define a model property, Knowledge to Translate Individual Words (KTIW) (e.g. knowing that `i` translates to `o`), and claim that it drives the learning of the attention. This claim is supported by the fact that before the attention mechanism is learned, KTIW can be learned from word co-occurrence statistics, but not the other way around. Particularly, we can construct a training distribution that makes KTIW hard to learn, the learning of the attention fails, and the model cannot even learn the simple task of copying the input words to the output. Our approximation explains why models sometimes attend to salient words, and inspires a toy example where a multi-head attention model can overcome the above hard training distribution by improving learning dynamics rather than expressiveness. We end by discussing the limitation of our approximation framework and suggest future directions.

Under review as a conference paper at ICLR 2023
APPROXIMATING How SINGLE HEAD ATTENTION LEARNS
Anonymous authors Paper under double-blind review
ABSTRACT
Why do models often attend to salient words, and how does this evolve throughout training? We approximate model training as a two stage process: early on in training when the attention weights input word to o if they co-oceur frequently. Later, the model learns to attend to while the correct output is o because knows translates to o. To formalize, we define model property, Knowledge to Translate Individual Words (KTIW) (e.g. knowing that i translates to o), and claim that drives the learning of the attention. This claim is supported by the fact that before the attention mechanism is learned, KTIW can be learned from word co-occurrence statistics, but not the other way around. Particularly, we can construct training distribution that makes KTIW hard to learn, the learning of the attention fails, and the model cannot even learn the simple task of copying the input words to the output. Our approximation explains why models sometimes attend to salient words, and inspires a toy example where a multi-head attention model overcome the above hard training distribution by improving learning dynamics rather than expressiveness. We end by discussing the limitation of our approximation framework and suggest future directions.
uniform, the model learns translate individual
INTRODUCTION
The attention mechanism underlies many recent advances in natural language processing, such as machine translation BBahdanau al. (2015) and pretraining Devlin al. (2019). While many works focus analyzing attention in already- trained models Jain & Wallace (201);; Vashishth et al. (2019); Brunner et al. (2019); 9h: Elhage et al. [202T] T); Olsson et al. (2022), little understood about how the attention mechanism is learned via gradient descent training time.
These learning dynamics are important, as standard, gradient-trainee models can have very unique inductive biases, distinguishing them from more esoteric but equally accurate models. For example, in text classification, while standard models typically attend to salient (high gradient influence) words Serrano & Smith (2019), recent work constructs accurate models that attend to irrelevant words instead Wiegreffe & Pinter (219):: Pruthi et al. (2020). In machine translation, while the standard gradient descent cannot train high-accuracy transformer with relatively few attention heads, can construct one by first training with more heads and then pruning the redundant heads Voita et al. (2019); Michel et al. (2019). To explain these differences, we need to understand how attention is learned at training time.
Our work opens the black box of attention training, focusing on attention in LSTM Seq2Seq models Luong et al. (2015) (Section 2. Intuitively, if the model knows that the input individual word i translates to the correct output word o, it should attend minimize the loss. This motivates us to investigate the model's knowledge to translate individual words (abbreviated as KTIW), and we define a lexical probe 3 to measure this property.
We claim that KTTW drives the attention mechanism to be learned. This is supported by the fact that KTIW can learned when the attention mechanism has not been learned (Section 3.2). but not the other way around (Section 3.3 Specifically, even when the attention weights are frozen to be uniform, probe 3 still strongly agrees with the attention weights of standardly trained model. On the other hand, when KTIW cannot be learned, the attention mechanism cannot be learned. Particularly, we can construct a distribution where KTIW is hard to learn; as a result, the model fails to learn simple task of copying the input to the output.

---

Under review as a conference paper at ICLR 2023
Now the problem of understanding how attention mechanism is learned reduces to understanding how KTIW is learned. Section 2.3 builds a simpler proxy model that approximates how KTIW is learned, and Section 3.2 verifies empirically that the approximation is reasonable. This proxy model is simple enough to analyze and we interpret its training dynamics with the classical IBM Translation Model (Section 4.2 which translates individual word to o if they co-occur more frequently.
To collapse this chain of reasoning, we approximate model training in two stages. Early on in training when the attention mechanism has not been learned, the model learns KTIW through word co-occurrence statistics; KTIW later drives the learning of the attention.
Using these insights, explain why attention weights sometimes correlate with word saliency in binary text classification (Section 5.1]: the model first learns to "translate" salient words into labels, and then attend to them. We also present toy experiment (Section[5.2 where multi-head attention improves learning dynamics by combining differently initialized attention heads, even though single head model can express the target function.
Nevertheless, "all models are wrong". Even though our framework successfully explains and predicts the above empirical phenomena, it cannot fully explain the behavior of attention-based models, since approximations are after all less accurate. Section 6 identifies and discusses two key assumptions: (1) information of a word tends to stay in the local hidden state (Section 6. and (2) attention weights are free variables (Section 6.2 We discuss future directions in Section
2
MODEL
Section 2.1 defines the LSTM with attention Seq2Seq architecture. Section 2.2 defines the lexical probe 3, which measures the model's knowledge to translate individual words (KTIW). Section 2.3 approximates how KTIW is learned early on in training by building a "bag of words" proxy model. Section 2.4 shows that our framework generalizes to binary classification.
Classical Alignment Learning Procedure
Translate
Dieser Film ist großartig
This movie great
1-1 1-2 1-3 1-4
Dilser Filmo groflontig
Film
6:25
4,28
0,2
(I) We count all co- the input and cutput words.
andoolyy
8.25
9:20
8:39
6.35
9.55
0,35
0,25
Attention-basee based Model Learning Dynamics
This monie ie great t-2 tn
whmnt
Herer
6:34
BM
prglannig
File
6.34
635
6.36
(2)
1
impulccuppt sentence pains and adl them.
model first
This mumie had
trunslation under
uniform attention when
aaining stats.
h1
h4
Mutrix
(3) Trunslationn
9:11
0:34
100
Diour
57
1
200
Transo Plon
("movie" linpu)
0:54
extmmated by toch "Filn" likely
0:33
671
821
140
36
Thene
trandltic to "morie"
3.00
025
000
Ne
Nnl The
[otber]
(2) "Fiilm" is more likely Ios
Pe .04
Pas -
32 A 005
Mir -
.04
(4) "Filin" is more likely canslate "movie".
trassiite "msovie".
Dieser
Bus
0.13 Diewer
Seseeeee www.w.e.........................
gralinnig
This
momi
groar
(5) Alignment or how much each word
(3) Attention a then atrracted the ward "Fitm""
contiiattss
Alignmonty
2, 2)
2nd output "ssovie" is attracted to "Filn".
32
04 + 32 4......4. 64466.1
Figure 1: Attention mechanism in recurrent models (left, Section 2.1 and word alignments in the classical model (right, Section 4.2) are learned similarly. Both first learn how to translate individual words (KTIW) under uniform attention weights/alignment at the start of training (upper, blue background), which then drives the attention mechanism/alignmenn to be learned (lower, red background).
2

---

Under review as a conference paper at ICLR 2023
2.1 MACHINE TRANSLATION MODEL
We use the dot-attention variant from Luong et al. 2015). The model maps from an input sequence {ari) with length L an output sequence (3), J with length T. We first LSTM encoders to embed {mr) C and {vr} C O respectively, where and O are input and output vocab space, and obtain encoder and decoder hidden states {f4] and {st}- Then we calculate the attention logits at by applying learnable mapping from h4 and and use softmax to obtain the attention weights d
g
1111 = Wh: O1
Σrz e
Next we sum the encoder hidden states {he] weighted by the attention obtain the "context vector" Ct, concatenate it with the decoder 8t, and obtain the output vocab probabilities pt by applying learnable neural network N with one hidden layer and softmax activation at the output, and train the model by minimizing the sum of negative log- ligeliilddood of all the output words y/-
Pt [N[(( se]);
L
log Pt,ee-
2.2
LEXICAL PROBE 3
We define the lexical probe Br,1 as:
Be,1 = N(h,, sel)y.
which means "the probability assigned the correct word yt. if the network attends only to the input encoder state hi". If we assume that bu only contains information about 21, 3 closely reflects KTIW, since 3 can interpreted as "the probability that Fı translated to the output y."
Heuristically, to minimize the loss, the attention weights ar should be attracted to positions with larger Sr.1. Hence, we expect the learning of the attention to be driven by KTIW (Figure Ieft). We then discuss how KTIW is learned.
2.3 EARLY DYNAMICS OF LEXICAL KNOWLEDGE
To approximate how KTIW is learned early on in training, we build proxy model by making a few simplifying assumptions. First, since attention weights uniform early in training, replace the attention distribution with a uniform one. Second, since we are defining individual word translation, we assume that information about each word is localized its corresponding hidden state. Therefore, similar to Sun & Lu (2020), we replace h with an input word embedding ea; E Rd where represents the word embedding matrix and d the embedding dimension, Third, to simplify analysis, assume N only contains ignore the decoder state 8;. Putting these assumptions together, we now define a new proxy model that produces output vocab probability pt == =() Et=Wez,)) Wez,).
linear layer W E [ before softmax activation and
On a high level, this proxy averages the embeddings of the input "bag of words", and produces distribution over output vocabs to predict the output "bag of words". This implies that the sets of input and output words for each sentence pair are sufficient statisties for this proxy. The probe 3ox can be similarly defined as 3Px == oWWe,, ezy Iy.
We provide more intuitions on how this proxy learns in Section
2.4 BINARY CLASSIFICATION MODEL.
Binary classification can be reduced "machine translation", where T 1 and |O| 2. We drop the subscript f = 1 when discussing classification.
We use the standard architecture from Wiegreffe Pinter 2019). After obtaining the encoder hidden states {h,), we calculate the attention logits aj by applying a feed-forward neural network with one hidden layer and take the softmax of a to obtain the attention weights a::
= (ReLU(QQ))) an
(4)
E1=1 esp
3

---

Under review as a conference paper at ICLR 2023
where Q and are learnable.
We sum the hidden states {hr} weighted by the attention, feed it to a final linear layer and apply the sigmoid activation function (o) to obtain the probability for the positive class
p O( ayh) o
on WoWTThy
(5)
Similar to the machine translation model (Section 2.1 we define the "lexical probe": Br == a ((2 hu),,
(6)
where y e {0, 1} is the label and 2y [-1, controls the sign.
On a high level, Sun & Lu (2020) focuses binary classification and provides almost the exact same arguments as ours. Specifically, their polarity score "sy" equals - in our context, and they provide more subtle analysis of how the attention mechanism is learned in binary classification.
3 EMPIRICAL EVIDENCE
We provide evidence that KTIW drives the learning of the attention early in training: KTIW can be learned when the attention mechanism has not been learned (Section 3.2), but not the other way around (Section
3.1 MEASURING AGREEMENT
We start by describing how to evaluate the agreement between quantities of interest, such and 3. For any input-output sentence pair (x" (x"","") yy"" for each output index I, az ,18h" 13p E RR4 all associate each input position with a real number. Since attention weights and word alignment tend to be sparse, we focus on the agreement of the highest-valued position. Suppose u, e E R we formally define the agreement of with as:
A(u, v) = [[(i >> Varg max ) < 5%L].
which means "whether the highest- valued position (dimension) in u is in the top 5% highest-valued positions in v". We average the values across all output words on the validation set to measure the agreement between two model properties. We also report Kendall's rank correlation coefficient in Appendix 2 for completeness.
We denote its random baseline as Å. Á is close to but not exactly 5% because of integer rounding.
Contextualized Agreement Metric. However, since different datasets have different sentence length distributions and variance of attention weights caused by random seeds, it might be hard directly interpret this agreement metric. Therefore, we contextualize this metric with model performance. We use the standard method to train model till convergence using T steps and denote its attention weights as a next we train the same model from scratch again using another random seed. We denote its attention weights at training step as â(7) and its performance as P(T). Roughly speaking, when 7 < T, both A(a, a(7)) and p(T) increase as T increases. We define the contextualized agreement & as:
E(u, r) == ((inf ((T)) A(u,v)))
(8)
In other words, we find the training step where its attention weights á(m) and the standard attention weights a agrees more than u and v agrees, and report the performance this iteration. We refer to the model performance when training finishes (T T) E*.
Datasets. We evaluate the agreement metrics A and & on multiple machine translation and text classification datasets. For machine translation,
Multi-30k (En-De), IWSLT' 14 (De-En), and News Commentary v14 (En-NI, En-Pt, and It-Pt). For text classification, we use IMDB Sentiment Analysis, AG News Corpus, 20 Newsgroups (20 NG), Stanford Sentiment Treebank, Amazon review,
4

---

Under review as a conference paper at ICLR 2023
Task
A(a, S
A(gur, _g0x) )p
A(A, guf)
A(a, B)
A
E(a, 3"r)
E(a, 3)
8*
IMDB
82
62
60
5
87
87
90
AG News
39
55
43
48
6
94
95
96
20 NG
65
41
65
63
5
91
85
94
SST
20
34
22
25
8
78
82
84
Multi30k
31
34
27
49
7
43
49
66
IWSLT14
36
39
28
55
7
36
44
67
News It-Pt
29
39
25
52
6
22
25
55
Table 1: The tasks above the horizontal line are classification and below are translation. The (contextualized) agreement metric A(E) is described in Section 3.1 Across all tasks, A(a, 8). A(a, frf). and A(3rtf ((x)) significantly outperform the random baseline A and the corresponding contextualized interpretations are also non-trivial. This implies that the proxy model from Section 2.3 approximates well how KTIW learned, 2) attention weights a and the probe 3 of KTIW strongly agrees, and 3) KTIW can still be learned when the attention weights uniform.
and Yelp Open Data Set. All of them are in English. The details and citations of these datasets can be seen in the Appendix A.5 We use token accuracy to evaluate the performance of translation models and accuracy to evaluate the classification models.
Due to space limit we round to integers and include subset of datasets in Table for the main paper. Appendix Table Tuble4finee includes the full results.
3.2 KTIW LEARNS UNDER UNIFORM ATTENTION
Even when the attention mechanism has not been learned, KTTW can still be learned. We train the same model architecture with the attention weights frozen to be uniform, and denote its lexical probe as guf Across all tasks, A(a, 3 and A(3ut (3x) significantly outperform the random baseline Á, and the contextualized agreement E(ar, .gff) is also non-trivial. This indicates that 1) the proxy we built in Section 2.3 approximates KTIW and 2) even when the attention weights are uniform, KTIW is still learned.
3.3 ATTENTION FAILS WHEN KTIW FAILS
We consider simple task copying from the input to the output, and each input a permutation of the same set of 40 vocab types. Under this training distribution, the proxy model provably cannot learn: every input-output pair contains the exact same set of input-output words As result, our framework predicts that KTIW is unlikely to be learned, and hence the learning of attention is likely to fail.
The training curves of learning to copy the permutations are in Figure left, colored in red: the model sometimes fails to learn. For the control experiment, if we randomly sample and permute 40 vocabs from 60 vocab types training samples, the model successfully Iearns (blue curve) from this distribution every time. Therefore, even if the model is able to express this task, it might fail to learn it when KTIW is not learned. The same qualitative conclusion holds for the training distribution that mixes permutations two disjoint sets words (Figure middle), and Appendix A.3 fillustrates the intuition.
For binary classification, it follows from the model definition that attention mechanism cannot be leamedd if KTIW cannot be leareed, since
p comect
ano (31)); (xx) I+e
Apppendix Tables
and include results for BLEU.
'Empirically, 8rx converges to the unigram weight bag-of-words logistic regression model, and hence 8tx does capture an interpretable notion of "keywords". (Appendix A.10)
We provide intuitions this Section
5

---

Under review as a conference paper at ICLR 2023
and the model needs to attend to positions with higher B, in order to predict correctly and minimize the loss. For completeness. we include results where we freeze / and find that the learning of the attention fails in Appendix A.6
Copina Curos 1 18lc
Seqaece Capying Trakiig Cares, 2 Biccs.
Inforacell Iadeaked Balliiddiinn
48
13
vead pemziiis
Toikane Stpp
Figure 2: Each curve represents accuracy on the test distribution vs. number of training steps for different random seeds (20 each). Left and Middle are accuracy curves for single head attention models. When trained on a distribution of permutation of 40 vocabs (red) (Left) or a mixture of permutations (Middle), the model sometimes fails to learn and converges slower. The right figure is for multi-head attention experiments. If all head initializations (head-init) are bad (red), the model is likely to fail; if one of the head-init is good (blue), is likely learn; with high chance, at least one out of eight random head-init is good (green).
4 CONNECTION TO IBM MODEL 1
Section 2.3 built a simple proxy model to approximate how KTIW is learned when the attention weights uniform early in training, and Section 3.2 verified that such approximation is empirically sound. However, it still hard to intuitively reason about how this proxy model learns. This section provides more intuitions by connecting its initial gradient (Section 4.1 to the classical IBM Model alignment algorithm Brown et al. (1993) (Section 4.2
4.1 DERIVATIVE AT INITIALIZATION
We continue from the end Section 2.3 For each input word and output word o, we are interested in understanding the probability that assigns to defined as:
0 :- o o(We es)o.
(10)
This quantity is directly tied to grx, since 3 = 6px in the dataset, the total loss C is:
Using super-script to index sentence pairs
st.ye"
m
L=
log(a(
1
Σ
We Werr)or)
(11)
L
Suppose each or Wa is independently initialized from a normal distribution (0, Ia/d) and we minimize L over W and e using gradient flow, then the value of and W are uniquely defined for each continuous time step т. By some straightforward but tedious calculations (details in Appendix A.2h. the derivative of O, when the training starts is:
00px
lim d
(T 0) 3 2(
Σ O'EO
)
(12)
Ør
O
where 4 means convergence in probability and CPX is defined as
L T
CPp ΣΣΣ
1| i] o).
(13)
Equation 12 tells us that 3Px
0px is likely to be larger if Ca is large. The definition of C 21.M
seems hard to interpret from Equation 13 but in the next subsection we will find that this quantity naturally corresponds to the "count table" used in the classical IBM I alignment learning algorithm.

---

Under review as a conference paper at ICLR 2023
4.2 IBM MODEL 1 ALIGNMENT LEARNING
The classical alignment algorithm aims to learn which input word is responsible for each output word (e.g. knowing that y/2 "movie" aligns D2 "Film" in Figure[ left), from of input-output sentence pairs. IBM Model 1 Brown et al. 1993) starts with 2-dimensional count table CIBM indexed by 2 and O, denoting input and output vocabs. Whenever vocab and o co-occurs in an input-output pair, we add 1, to the ",o entry (step 1 and 2 in Figure | right). After updating CIBM for the entire dataset, C1BM is exactly the same C defined in Equation 13 super-script of C to keep the notation uncluttered.
We drop the
Given C, the classical model estimates probability distribution of "what output word does the input word i translate to" (Figure right step 3) as
Ci
Tra)s(oli) = Ed Cid
(14)
In pair of sequences ((xt), (vr)), the probability BIBM that on is translated to the output y/ is:
Pt Trans(y:
(15)
and the alignment probability a IBM that "z1 is responsible for outputting versus other xy is
3IBM
(1,
(16)
3IBM
which monotonically increases with respect to BLBM See Figure I right step 5.
4.3 VISUALIZING AFOREMENTIONEL TASKS
Figure (right) visualizes the count table C' for the machine translation task, and illustrates how
KTIW is learned and drives the learning of attention. We provide similar visualization for why
KTIW is hard to learn under distribution of vocab permutations (Section 3.3 in Figure 3 and how
word polarity learned in binary classification (Section 2.4p in Figure
0.25 0.25 0.25
0.25
ARCD
0.25
0.25
0.25
0.25
A' C'
0.25
0.25
0.25
0.25
A²
8'
C"
13"
M/4
M/4
M/4
M/4
Trans(C| .25
Trans(C ()=
TransiC (0) .25
M/4
M/4
M/4
M/4
Trans(C"|)) .25
A'
D'
C'
D'
DCBA
0.25
0.25
0.25
0.25
[In total M sequences)
0.25
0.25
0.25
0.25
B'A'
0.25
0.25
0.25
0.25
Alignment no preference over any of
these words, the prohabilities are unifform
Figure 3: Co-occurrence table C is non-informative under distribution of permutations. Therefore, this distribution is hard for the attention-based model to learn.
APPLICATION
5.1 INTERPRETABILITY IN CLASSIFICATION
We use gradient based method Ebrahimi et al. (2018) to approximate the influence A; for each input word zj. The column A(A, 3TH reports the agreement between A and 3"f, and it significantly outperforms the random baseline. Since KTIW initially drives the attention mechanism to be learned, this explains why attention weights are correlated with word saliency on many classification tasks, even though the training objective does not explicitly reward this.
7

---

Under review as a conference paper at ICLR 2023
Ned
0
Input: great movie
great
Label: Positive
Pontine Negulivv
Trans(pontlll =
0.5
Trankjoostie loresg
0.5
Transjuostiie [novid) u5
0.5
Plarttie
Megusliv
Inpur: bad movie
Label: Negative
Alignment is attracted to the word "gread". since the lurgest.
Postitive Negeniive
Figure 4: The classical model first learns word polarity, which later attracts attention.
5.2 MULTI-HEAD IMPROVES TRAINING DYNAMICS
We saw in Section 3.3 that learning copy sequences under distribution of permutations is hard and the model can fail to learn; however, sometimes it is still able to learn. Can we improve learning and overcome this hard distribution by ensembling several attention parameters together?
We introduce a multi-head attention architecture by summing the context vector obtained by each head. Suppose there are heads each indexed by k, similar to Section 2.1
(4)
(k) 1.
of (1())) bu
(k)
e
(17)
E e E.
and the context vector and final probability pt defined as:
K hus p = PN=N(
(b)
del),
(18)
where (kk) are different bearn-able parameters.
We call wx(k) init good initialization if training with this single head converges, and bad otherwise. We use rejection sampling to find good/bad head initializations and combine them to form 8-head (K 8) attention models. We experiment with scenarios: (1) all head initializations bad, (2) only one initialization is good, and (3) initializations are sampled independently at random.
Figure 2 right presents the training curves. If all head initializations are bad, the model fails to converge (red). However, as long as one of the eight initializations is good, the model can converge (blue). As the number of heads increases, the probability that all initializations are bad is exponentially small if all initializations are sampled independently: hence the model converges with very high probability (green). this experiment, multi-head attention improves not by increasing expressiveness, since one head is sufficient to accomplish the task, but by improving the learning dynamics.
ASSUMPTIONS
We revisit the approximation assumptions used in our framework. Section 6.1 discusses whether the lexical probe St,1 necessarily reflects local information about input word T₁, and Section 6.2 discusses whether attention weights can be freely optimized to attend to large 3. These assumptions are accurate enough to predict phenomenon in Section and but they are not always true and hence warrant more future researches. We provide simple examples where these assumptions might fail.
6.1 3 REMAINS LOCAL
We toy classification task show that early on in training, expectantly, guf is larger near positions that contain the keyword. However, unintuitively, 3y (3 at the last position the sequence) will become the largest if we train the model for too long under uniform attention weights.
In this toy task, each input is a length-40 sequence of words sampled from {1...., 40) uniformly at random; sequence is positive if and only if the keyword "1" appears in the sequence. We restrict "1"
8

---

Under review as a conference paper at ICLR 2023
to appear only once in each positive sequence, and use rejection sampling balance positive and negative examples. Let 1* be the position where 2/- = 1.
For the positive sequences, examine the log-odd ratio "/ before the sigmoid activation in Equation 5 since & will be all close to 1 and comparing would be more informative: := log 1-p
We measure four quantities: 1) " the log-odd ratio if the model only attends the key word position, 2) +1++1, one position after the key word position, 3) 5 := Σi= if attention weights uniform, and 4) 7/r. if the model attends to the last hidden state. If the only contains information about word 27, we should expect:
Hypothesis 1 :
7 22 +1+
(19)
However, we accept the conventional wisdom that hidden states contain information about nearby words Khandelwal et al. (2018), we should expect:
Hypothesis 2 -2- > +1 > 7 7L-
(20)
To verify these hypotheses, plot how "t- +1+ 7, and Ty evolve training proceeds in Figure S Hypothesis is indeed true when training starts; however, we find the following to be true asymptotically:
Observation 3 TL >
(21)
which is wildly different from Hypothesis 2. If we train under uniform attention weights for too long, the information about keywords can freely flow other non-local hidden states.
6.2 ATTENTION WEIGHTS ARE FREE VARIABLES
In Section 2.1 we assumed that attention weights behave like free variables that can assign arbi- trarily high probabilities to positions with larger 3. However, is produced by model, and sometimes learning the correct can chal- lenging.
Evolution of under Uniform Attention
V
Let ı be a random permutation of integers from I to 40, and we want to learn the function f that permutes the input with x:
log--n codds
V2
f(x11, T2. ool)) (1), Tx(2) 2 (10) J.
(22)
Input are randomly sampled from vocab of size 60 as in Section 3.3 Even though gur behaves exactly the same for these two tasks, sequence copying is much easier to learn than permutation function: while the model always reaches perfect accuracy in the former setting within 300 iterations, it always fails in the latter. LSTM has built-in inductive bias to learn monotonic attention.
Training step
Figure 5: When training begins, Hypothesis (Equation 19) is true; however, asymptotically, Oberservation 3 (Equation 21) true.
7 CONCLUSIONS
Our work tries to understand the black box of attention training. Early on training, the LSTM attention models first learn how to translation individual words from bag of words co-occurrence statistics, which then drives the learning of the attention. Our framework explains why attention weights obtained by standard training often correlate with saliency, and how multi-head attention can increase performance by improving the training dynamics rather than expressiveness. These phenomena cannot be explained if we treated the training process black box.

---

Under review as a conference paper at ICLR 2023
8 ETHICAL CONSIDERATIONS
We present a new framework for understanding and predicting behaviors of an existing technology: the attention mechanism in recurrent neural networks. We do not propose any new technologies or any new datasets that could directly raise ethical questions. However, it is useful to keep in mind that our framework far from solving the question of neural network interpretability. and should be interpreted as ground truth in high stake domains like medicine or recidivism. We are aware and very explicit about the limitations our framework, which we made clear Section 6
REPRODUCABILITY STATEMENT
To promote reproducibility, we provide extensive results the appendix and describe all experiments in detail. We also attach source code for reproducing all experiments to the supplemental of this submission.
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL httpp::/pp.i.././////a
409 1409.0473
Peter Brown, Stephen A Della Pietra, Vincent Della Pietra, and Robert Mercer. The mathematics of statistical machine translation: Parameter estimation. Computational linguistics, 19(2):263-311 1993.
Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Watten- hofer. On identifiability transformers. In International Conference on Learning Representations, 2019.
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. 2015.
Jacob Devlin, Ming- Riig-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/N19 1423. URL https
Minneapolis, Minnesota, June
www..aslwee... www.aaawwww.....nooooooooooooeoe /ttp:///////.aaoooooooooooaoaaa/
Javid Ebrahimi, Anyi Rao, Daniel Lowd, and Dejing Dou. HotFlip: White-box adversarial ex- amples for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics ( (Volume 2: Short Papers), pp. 31-36, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10. 1/18653/11/18--2006 URL
/tttss://ww..r ac www.wwlwwwwb.... http://www.....ooroaaaoooooooogo
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transformer.ciccitts
html..
Desmond Elliott, Stella Frank, Khalil Sima' an, and Lucia Specia. Multi3KK: Multilingual English- German image descriptions. In Proceedings of the 5th Workshop on Vision and Language, pp. 70-74, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: a /RI.httpss:///// web. wtt..orgg//athhoogg///l.sseeeeee
N0.18653/v1/W163222
Sarthak Jain and Byron C. Wallace. Attention is not Explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics:
10

---

Under review as a conference paper at ICLR 2023
Human Language Technologies, Volume (Long and Short Papers), pp. 3543- 3556, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19- 1357.
URL ht https: www. acwweb. wttp http:////////aaaaoooooo////----0
Urvashi Khandelwal, He He, Peng Qi, and Dan Jurafsky. Sharp nearby, fuzzy far away: How neural language models context. In Proceedingy of the 56th Annual Meeting of the Association for
Computational Linguistics (Volume 1: Long Papers), pp. 284-294, Melbourne, Australia, July
2018. Association for Computational Linguisties. doi: 10.18653/v1/P18- 1/22 URL https
7/ /www.oo www.wb.wwwbb...... wattaaoooooooo--1o1o
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in
Natural Language Processing, pp. 1412-1422 1421, Lisbon, Portugal, September 2015. Association for Computational Linguisties. doi: 10. 10.18653/1//D155 166.
https:/ /ww. aclweb. org/
anthology/615-1166
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL Intpps
//www.acl crg. hwt.rrggattaloooggg 2atthallgyy/111101111
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Advances in Neural Information Processing Systemx, pp. 14014-14024, 2019.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Saufield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. http:/aannoooeeeeciicccusssuupaa http://////ionoooetoooeoe context-learnnng -aniiinn------iie indduttio--eeadd//dde... heads/inde..here htmml.
Danish Pruthi, Mansi Gupta, Bhuwan Dhingra, Graham Neubig, and Zachary C. Lipton. Learning to deceive with attention-based explanations. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguisties, pp. 4782-4793, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/11/2220.alllm aciiin.43c-- URL https: //www aelweb
org/ haa/aanncccgy//22/2 aci.mmin.mmmmam..mmmmmmm........
Sofia Serrano and Noah A. Smith. Is attention interpretable? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2931 12951, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/P1-11 1282. URL https
/ wwwww www. www.t/t..gaaaaaooaeeoooeo.coo.a loythology/PPt oggy/P 1282
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp. 1631-1642, 2013.
Xiaobing Sun and Wei Lu. Understanding attention for text classification. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 1--8882, Online, July 2020. Association for Computational Lingussics. doi: 1a1116663///02222lcllmimmml https://www http://wwwwwwebe..o/raaaaaoooooa act-://:/22222....mmmm 312
URL
Shikhar Vashishth, Shyam Upadhyay, Gaurav Singh Tomar, and Manaal Faruqui. Attention inter- pretability across nlp tasks. arXiv preprint arXiv: / 2019.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18663/v 1/- 1580. URL
https: /t/////ww www..wawmmmmmmm web. wrg//ntho ogy 1580

---

Under review as a conference paper at ICLR 2023
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 11 11-20, Hong Kong, China, November 2019. Association for Computational Linguisties. doi: 10. 1053//1////111 URL
https: hwwa:://www.co web www.oorggg/ooooooo www.or/ggannooogoooooooeeee ogy/019-113000 002]
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Advances in neural information processing systems, pp. 649-657, 2015.
12

---

Under review as a conference paper at ICLR 2023
A APPENDICES
A.1 HEURISTIC THAT a ATTENDS TO LARGER 8
a heuristic rather than rigorous theorem that attention attracted larger 3. There are two reasons. First, there is non-linear layer after the averaging the hidden states, which can interact in an arbitrarily complex way to break this heuristic. Second, even there are no non-linear operations after hidden state aggregation, the optimal attention that minimizes the loss does not necessarily assign any probability to the position with the largest 3 value when there are more than two output vocabs.
Specifically, we consider the following model:
Pt = oWe
Σ anay + Wost)
Σ an + 7a),
(23)
where We and Wa are learnable weights, and 7 defined as:
Wohg;
7s S. Ba, ( + 7s)m.
(24)
Consider the following scenario that outputs probability distribution p over 3 output vocabs and % set to
p σ (a171 0272 + a373),
(25)
where 1=222 E R are the logits, a a valid attention probability distribution, the softmax, and is the probability distribution produced by this model. Suppose
71 = [0, 0j. = [0, - 10, 5], 73 -005.-10]
(26)
and the correct output is the first output vocab (i.e. the first dimension). Therefore, we take the softmax of 7/ and consider the first dimension:
B1- 1-==5- - -
(27)
We calculate "optimal a" appt; the optimal attention weights that can maximize the correct output word probability po and minimize the loss. We find that 0x2 ope opt 0.5, while aj = 0. In this 03 example, the optimal attention assigns weight to the position with the highest 3r.
Fortunately, such pathological examples rarely occur in real datasets, and the optimal are usually attracted to positions with higher 3. We empirically verify this for the below variant of machine translation model on Multi30K.
As before, we obtain the context vector Gt Instead of concatenating e₁ and de and pass it into a non-linear neural network N. add them and apply linear layer with softmax after it to obtain the output word probability distribution
Pt = (a + dc)).
(28)
This model is desirable because delay the proof to the end of this subsection). Additionally, this model has comparable performance with the variant from our main paper (Section 2.1). achieving 38.2 BLEU score, vs. 37.9 for the model in our main paper. We use ptt to denote the attention that can minimize the loss, and we find that opt opt 3) 0.53. do strongly agree with opt.
now provably find the optimal using gradient descent (we
Now we are left to show that we can use gradient descent to find the optimal attention weights to minimize the loss. We can rewrite pt as
Pr = σ
an hy + Wc)..
(29)
We define
31 Why := Wd.
(30)
13

---

Under review as a conference paper at ICLR 2023
Without loss of generality, suppose the first dimension of L. Y all 0, and the correct token want to maximize probability for is the first dimension, then the loss for the output word is
f = log(1 + g(a)),
(31)
where
g(a) == Σ e 7o
(32)
oE0,of0
where
% h No yk,ol E R.
(33)
Since is defined within the convex probability simplex and g(a) is convex with respect to a, the global optima a opt can found by gradient descent.
A.2 CALCULATING 89..c
BF
We drop the super-script of I to keep the notation uncluttered. We copy the loss function here to remind the readers:
L Weep )or:)
C
log(o(
(34)
Lm 1=1
and since we optimize W and with gradient flow,
ow
de
C
(35)
θr aw
De'
We first define the un-normalized logits & and then take the softmax.
= We,
(36)
then
00 O(We) Or Or - W de
aw ÔT E.
(37)
dT
We first analyze W de Since [[ we analyze each entry 6,,-- Since differentiation operation and left multiplication by matrix W is linear, we analyze each individual loss term in Equation a4d then sum them up.
We define
m
nd 0 T,
W eam)
(38)
and
m = W. 5Lm
C = log(n );
(39)
De:
Hence,
T
C 6,0
m
(40)
tá,00*
Therefore,
m
10
1
1/2
( pm
p- W.).
(41)
De;
L m
1=1
Hence,
WT m.m
1 1x i
(42)
8pth
de,
I.
10
W pm
112
Po W T Wo y1
0=1
14

---

Under review as a conference paper at ICLR 2023
while for o' + vi(
Gidd = -W, dCT
L 1 10 =
(43)
de;
WW
Po Wy
If W, and e; each sampled iid. from Ia/d). then by central limit theorem:
#o vdw. W. 10,,1)
(44)
Vo, VdW2'. 3 N0,,)..
(45)
and
Vo, valll - 1) 4 (0,
(46)
Therefore, when т
D"
lim m 2 d+000
1
i|(1
of
o )
(47)
L
Summing over all the t. terms, we have that
Go = Cs
1
Σ
Ci.or
(48)
10
where C' is defined as
1. Tm
Ci.: = ΣΣΣ
1| = 1| = o).
(49)
L
m
We find that
8W converges exactly the value. Hence
00.. OWe ØT dT 2(C,,0
1 Eg,o) C.,r)
(50)
Since limg-+0 0(T 0) &
1 p
| by chain rule,
lim 0 (T= B, 2(Ci, - 101 O'00
Σ Cio.
(51)
d Or
A.3 MIXTURE OF PERMUTATIONS
For this experiment, each input is either a random permutation of the set (1...40) 40]. or a random permutation of the set (41 (41...80) 80). The proxy model easily learn whether the input words less than 40 and decide whether the output words are all less than 40. However, 30x is still the same for every position; as a result, the attention and hence the model fail lear. The count table C can be see in Figure
A.4 ADDITIONAL TABLES FOR COMPLETENESS
We report several variants of Table I
We chose to use token accuracy to contextualize the agreement metric in the main paper, because the enrors would accumulate much more if we use not-fully trained model to regr-ssgrysyvvely generate output words.
Table 2 contains the same results as Table | except that its agreement score A(u, v) is now Kendall Tau rank correlation coefficient, which more popular metric.
15

---

Under review as a conference paper at ICLR 2023
0
0
0
0.25
0.29
0.25
0.25
0.25
0.25
0.25
0.25
ABCD C'
0
0.25
0.25
0.25
0.25
FGEH
0.50.250.25 0.25
0.25.250.25
0.25
0
C
0
0
0
FGER
0.260.50.25
0.25
RP
87
C
11
+
P
0
0
0
0
P
D
0
0.25
0.25
0.25
0.25
DCBA
HGFE
D'A'
HGFE
0.250.250.25
0.25
0
0
0200.200..25 25
0.25
A'
R
[In total sequences]
[In total K sequences]
K/4
K/4
K/4
K74
0
0
K/4
K/44
K/4
K/4
85
K/4
M/4
M/4
M/4
M/4
0
Trans<C"]]]
M/4
M/4
M/4
M14
0
19)) 25
M/4
M/4
M/4
M/4
0
Trans(C"]
Alignment has no preference over any of these words, since the probabilities uniform over the input words "A", "B", "C",""" "D".
Figure The training distributions mixes random permutation of disjoint set of words (left and right, respectively). From the count table, gvx could learn that the set of input words {A, B, C, D} corresponds to the set of output words (A', B', "',C", D'}. but its B value for each input position is still uniformly 0.25.
16

---

Under review as a conference paper at ICLR 2023
Task
A(a. grr)
A(gut, .30x)
A(A, ,
A
IMDB
12.77
33.31
12.56
0.00
Yelp
20.38
36.75
20.98
0.00
AG News
26.31
36.65
20.55
0.00
20 NG
16.06
22.03
6.50
0.00
SST
11.68
31.43
15.01
0.00
Amzn
15.21
35.84
9.33
0.00
Muti30k
7.89
27.54
3.93
0.00
IWSLT14
8.64
22.56
2.72
0.00
News It-Pt
4.82
17.16
1.63
0.00
News En-NI
4.53
20.35
2.08
0.00
News En-Pt
4.65
18.20
2.17
0.00
Task
E(a, ,3u))
E(gut, (Px)
<(A, g
IMDB
70.60
80.50
70.60
89.55
Yelp
87.44
93.20
87.44
96.20
AG News
89.31
93.54
85.85
96.05
20 NG
60.75
60.75
60.75
94.22
SST
76.69
83.53*
76.69
83.53
Amzn
69.91
88.07
57.36
90.38
Muti30k
22.94
36.61
22.94
66.29
IWSLT14
29.07
32.98
29.07
67.36
News It-
14.01
18.25
8.10
55.41
News En-NI
9.60
18.59
9.60
62.90
News En-Pt
14.10
14.10
7.71
67.75
Table 2: Table
except with agreement defined by Kendall Tau. Section A.4
Task
Ala, gut)
A(gut _Fx
A(A, Buf)
A
Muti30k
8.68
27.54
4.24
0,00
IWSLT14
8.64
22.56
2.72
0.00
News It-Pt
4.82
17.16
1.63
0.00
News En-NI
4.53
20.35
2.08
0.00
News En Sn
4.41
18.20
2.05
0.00
Task
E(a, (8u))
{(gut _BP)
É(A, grr)
Muti30k
1.99
6,91
1.99
37.89
IWSLT14
5.38
5,31
5.38
32.95
News It-Pt
0.09
0.55
0.04
24.71
News En-NI
0.01
0.94
0.01
29.42
News En-Pt
0.01
0.22
0.01
37.04
Table 3: translation results from Table 2 except with performance measured by bleu rather than token
accuracy. Section A.4
Table 4 contains the same results as Table except that results are now rounded to two decimal places.
of the same results as Table I except that the statistics is calculated over the training rather than the validation set.
and Table 7 contain the translation results from the above 3 mentioned tables respectively, except that á defined as BLEU score rather than token accuracy, and hence the contextualized metric interpretation changes correspondingly.
Table 3 Table 5]
A.5 DATASET DESCRIPTION
We summarize the datasets that we use for classification and machine translation. See Table for details train/test splits and median sequence lengths for each dataset.
17

---

Under review as a conference paper at ICLR 2023
Task
A(a, 3rF)
A(put _Pxx
A(A, gut)
A(a, 9)
IMDB
52.55
81.60
61.55
60.35
5.30
Yelp
17.55
75.38
58.90
35.00
5.80
AG News
39.24
55.13
43.08
48.13
6.20
20 NG
65.08
41.33
64.82
63.07
5.11
SST
19.85
33.57
22.45
25.33
8.39
Amzz
52.02
76.78
49.68
62.13
5.80
Muti30k
31.02
34.43
27.06
48.78
7.11
IWSLT14
35.75
39.09
27.69
55.25
6.52
News It-Pt
29.13
38.62
25.45
52.48
6.17
News En-NI
35,53
41.72
29.15
60.15
6.36
News En-Pt Pt
35.90
37.37
30.23
65.49
6.34
Task
E(a, (3t)
E(gut (Px)
<(A, 8
E(a, (8)
IMDB
86.81
88.88
86.81
86.81
89.55
Yelp
90.39
95.22
95.31
93.59
96.20
AG News
93.54
96.05*
94.32
94.50
96.05
20 NG
91.16
60.75
84.57
84.57
94.22
SST
78.16
83.53*
78.16
82.38
83.53
Amzn
82.48
90.38
82.48
88.07
90.38
Muti30k
43.45
43.45
43.45
48.58
66.29
IWSLT14
35.82
35.82
32.98
44.09
67.36
News It-Pt
21.82
25.06
21.82
25.06
55,41
News En-NI
18.59
23.21
18.59
26.79
62.90
News En-Pt
19.12
19.12
19.12
27.85
67.75
Table 4: Table ! with 2 decimal results. Section A.4
Task
A(a, g
A(guf pDx)
A(A, ,zf)
A(a, 3)
Á
Muti30k
30.77
34.43
27.24
48.70
7.19
IWSLT14
35.75
39.09
27.69
55.25
6.52
News It-Pt
29.13
38.62
25.45
52.48
6.17
News En-NI
35.53
41.72
29.15
60.15
6.35
News En-Pt
35.77
37.37
30.37
64.94
6.34
Task
((ar/84
{l (PXx
E(A, g
E(a; 9
Muti30k
11.43
11.43
11.43
16.41
37.89
IWSLT114
6.71
6.71
5.31
9.89
32.95
News It-Pt
1.29
2.16
1.29
2.16
24.71
News En-NI
0.94
2.39
0.94
4.12
29.42
News En-Pt
0.74
0.74
0.74
4.28
37.04
Table 5: translation results from Table 4 excepp with performance measured by bleu rather than token accuracy. Section A.4
18

---

Under review as a conference paper at ICLR 2023
Task
A(a, gut)
A _30x)
A(A, gut)
A(ar, (3)
IMDB
51.52
80.10
42.85
64.88
5.29
Yelp
11.15
76.12
55.50
37.63
5.85
AG News
36.97
53,95
43.11
46.89
6.17
20 NG
72.36
38.69
71.73
69.47
5.32
SST
21.82
29.35
20.48
28.50
8.48
Amzn
51.95
77.18
40.15
61.78
5.91
Muti30k
32.89
34.67
28.36
56.39
7.21
IWSLT14
36.61
38.95
28.37
57.71
6.52
News It-Pt
31.03
38.70
27.11
64.81
6.15
News En-NI
37.86
41.91
31.11
67.68
6.39
News En-Pt
37.43
37.23
31.76
71.96
6.35
Task
8"),
E(B"t, (8NN)
E(A, 3
E(a. 3)
g
IMDB
90.40
99.95
90.40
95.01
99.95
Yelp
75.61
96.54
96.19
94.44
98.22
AG News
93.57
98.42*
94.63
95.54
98.42
20 NG
100.00
65.40
100.00
100.0
100.00
SST
97.72
100.00"
84.11
100.0"
100.00
Amzn
87.96
99.58*
80.98
91.09
99.58
Muti30k
43.27
43.27
43.27
51.97
80.76
IWSLT14
35.94
35.94
35.94
44.18
71.18
News It-Pt
22.69
25.96
22.69
39.98
77.10
News En-NI
18.85
23.56
18.85
40.09
74.49
News En-Pt
19.33
19.33
19.33
42.41
77.97
Table 6: Table 1 excepp with correlations and performance metrics taken over the training set instead of the validation set. Section A.4
Task
A(a,,3u)
A(But BPX)
A(A, gut)
Á
Muti30k
32,89
34.67
28.36
7.16
IWSLT14
36.61
38.95
28.37
6.54
News It-Pt
31.03
38.70
27.11
6.17
News En-NI
37.86
41.91
31.11
6.38
News En-Pt
37.43
37.23
31.76
6.37
Task
E(a, g3u))
{(gut [PP
<(A, g
E*
Muti30k
11.87
11.87
11.87
52.28
IWSLT14
6.82
6.82
6.82
36.23
News It-Pt
1.30
2.30
1.30
42.40
News En-NI
1.11
2.29
1.11
39.40
News En-Pt
0.83
0.83
0.83
46.57
Table 7: translation results from Table Table 6 except with performance measured by bleu rather than token accuracy. Section A.4
19

---

Under review as a conference paper at ICLR 2023
IMDB Sentiment Analysis Maas et al. (2011) A sentiment analysis data set with 50,000 (25,000 train and 25,000 test) IMDB movie reviews and their corresponding positive or negative sentiment.
AG News Corpus Zhang al. 2015) 120,000 news articles and their corresponding topic (world, sports, business, or science/tech) We classify between the world and business articles.
20 Newsgroups A news data set containing around 18,000 newsgroups articles split between 20 different labeled categories. We classify between baseball and hocky articles.
Stanford Sentiment Treebank Socher et al. (2013) A data set for classifying the sentiment of movie reviews, labeled on scale from I (negative) to (positive). We remove all movies labeled as 3, and classify between 4 or 5 and 1 or 2.
Multi Domain Sentiment Data set Approximately 40,000 Amazon reviews from various product categories labeled with corresponding positive negative label. Since some of the sequences particularly long, we only use sequences of length less than 400 words.
Yelp Open Data Set 20,000 Yelp reviews and their corresponding rating from to 5, We classify between reviews with rating x and 2 4.
(016) English German translation. The data is from translation image Multi-30k Elliott al. captions.
[estto et al. 2015) German English translation. The data is from translated TED talk transcriptions.
News Commentary v14 al. (2015) A collection of translation news commentary datasets in different languages from WMT We use the following translation splits: English- Duth (En-NI), English-Portugune (En-Pt), and Ialian-Portuguess (It-Pt), In pre-processing for this dataset, removed all purely numerical examples.
A.6 a FAILS WHEN FROZEN
For each classification task we initialize random model and freeze all parameters except for the attention layer (frozen 3 model). We then compute the correlation between this trained attention (defined as at and the normal attention a. Table this correlation the iteration where a is most correlated with a on the validation set. As shown in Table D the left column is consistently lower than the right column. This indicates that the model can learn output relevance without attention, but not vice versa.
A.7 TRAINING guf
We find that A(a, (8 first increases and then decreases as training proceeds (i.e. 7 increases), so chose the maximum agreement report in Table [ over the course of training. Since this trend is consistent across all datasets, our choice minimally inflates the agreement measure, and is comparable to the practice of reporting dev results. As discussed in Section 6. training under uniform attention for too long might bring unintuitive results,
A.8 MODEL AND TRAINING DETAILS
Classification Our model uses dimension 300 VloVe pre-trained embeddings to initialize the token embeddings where they aligned with our vocabulary. The sequences are encoded with a layer bidirectional L.STM of dimension 256. The rest of the model, including the attention mechanism, is exactly as described in 2.4 Our model has 1,274,882 parameters excluding embeddings. Since each classification has a different vocab size each model has a slightly different parameter count when considering embeddings: 19,376,282 for IMDB, 10,594,382 for AG News, 5,021,282 for 20
'http://qwonc.co http:////wwccco///////////////// wawmmm/aaaa2oossssssaaeo jason/20Newssrroups/
httpp::/wwwwes.../ wttp://wwww......uu.u madredoo/datassstssettnene
"htpps:/www.ycch wwt..ycccppc.ctaeee.......
9/trasslationttaskhhee html
20

---

Under review as a conference paper at ICLR 2023
Data
median train seq len
train
Data
median val seq len
val
IMDB
181
25000
IMDB
178
4000
AG News
40
60000
AG News
40
3800
NewsG
183
1197
NewsG
207
796
SST
5130
SST
1421
Amzn
71
32514
Amzn
72
4000
Yelp
74
88821
Yelp
74
4000
IWSLT14
(23 src, 24 trg)
160240
IWSLT14
(22 src, 23 trg)
7284
Multi-30k
(14 src, 14 trg)
29000
Multi-30k
(15 src, 14 trg)
1014
News-en-nl
(30 src, 34 trg)
52070
News-en-nl
(30 src, 34 trg)
5786
News-en-pt
(31 src, 35 trg)
48538
News-en-pt
(31 src, 35 trg)
5394
News-it-pt
(36 stc, 35 trg)
21572
News-it-pt
(36 src, 36 trg)
2397
Data
vocab size
IMDB
60338
AG News
31065
NewsG
31065
SST
11022
Amzn
37110
Yelp
41368
IWSLT14 8000
Multi-30k
8000
News-en-nl 8000
News-en-pt
8000
News-it-pt 8000
Table 8: statistics for each dataset. Median sequence length in the training set and train set size. Note: refers to the input "source" sequence, and trg refers to the output "target" sequence. Section A.5
Dataset IMDB
Ala, at)
A(a, gut)
53
AG News
17
39
20 NG
19
65
SST
14
20
Amzn
15
52
Yelp
8
18
Table 9: We report the correlation between a and a on classification datasets, and compare it against
A(a, B"f), the same column defined in Table Section A.6
Newsgroups, 4,581,482 for SST, 13,685,282 for Yelp, 12,407,882 for Amazon, and 2,682,182 for SMS.
Translation We use a a bidirectional two layer bi-LSTM of dimension 256 to encode the source and the use last hidden state h2 as the first hidden state of the decoder. The attention and outputs are then calculated as described in 2 The learn-able neural network before the outputs that is mentioned in Section 2 a hidden layer model with ReLU non-linearity. The hidden layer dimension 256. Our model contains 6. 6.132.544 parameters excluding embeddings and 8, 180,544 including embeddings on all datasets.
Permutation Copying We use single directional single layer LSTM with hidden dimension 256 for both the encoder and the decoder.
Classification Procedure For all classification datasets we used a batch size of 32. We trained
for 4000 iterations each dataset. For each dataset train the pre-defined training set if the dataset has one. Additionally, a dataset had a predefined test set, we randomly sample at most 4000 examples from this test set for validation. Specific dataset split sizes are given in Table
21

---

Under review as a conference paper at ICLR 2023
Classification Evaluation We evaluated each model at steps 0, 10, 50, 100, 150, 200, 250, and then every 250 iterations after that.
Classification Tokenization We tokenized the data at the word level. We mapped all words occurring less than times in the training set to <unk>, For 20 Newsgroups and AG News we mapped all non-single digit integer "words" to <unk>> For 20 Newsgroups we also split words with the" " character.
Classification Training We trained all classification models on a single GPU. Some datasets took slightly longer train than others (largely depending on average sequence length), but each train took at most 45 minutes.
Translation Hyper Parameters For translation all hidden states in the model are dimension 256. We use the sequence to sequence architecture described above. The LSTMs used dropout 0.5.
Translation Procedure For all translation tasks we used batch size 16 when training. For IWSLT' 14 and Multi-30k we used the provided dataset splits. For the News Commentary v14 datasets did 90-10 split of the data for training and validation respectively.
Translation Evaluation We evaluated each model steps 0, 50, 100, 500, 1000, 1500, and then every 2000 iterations after that.
Translation Training We trained all translation models single GPU. IWSLT" 14, and the News Commentary datasets took approximately 5-6 hours to train, and multi-30k took closer to hour to
train.
Translation Tokenization We tokenized both translation datasets using the Sentence- Piece tok- enizer trained the corresponding train set to vocab size of 8,000. We used single tokenization for source and target tokens. And accordingly also used the same matrix of embeddings for target and source sequences.
A.9 A NOTE ON SMS DATASET
In addition to the classification datasets reported in the tables, we also ran experiments on the SMS Spam Collection V.1 dataset [ The attention learned from this dataset was very high variance, and so two different random seeds would consistently produce attentions that did not correlate much. The dataset itself also bit of outlier; it had shorter sequence lengths than any of the other datasets (median sequence length 13 on train and validation set), it also had the smallest training set out of all our datasets (3500 examples), and had by far the smallest vocab (4691 unique tokens). We decided not to include this dataset in the main paper due to these unusual results and leave further exploration future works.
A. LOGISTIC REGRESSION PROXY MODEL
Our proxy model can be shown to be equivalent to bag-of- logistic regression model in the classification case. Specifically, we define a bag-of-word -forddds logistic regression model to be:
vt, P: x))
(52)
where x E I Z glog E [ and a is the softmax function. The entries in x: are the number of times each word occurs in the input sequence, normalized by the sequence length, and glog is learned. This is equivalent to:
htpp://www.de http:///wwww..fc...cccm...mppp.p br/ tiggo/smsspameellecttonn
22

---

Under review as a conference paper at ICLR 2023
Task
A(But ,Px)
A(But BBo8)
IMDB
0.81
0.84
Yelp
0.74
0.76
AG News
0.57
0.58
20 NG
0.40
0,45
SST Amzn
0.39
0.46
0.53
0.60
Table 10:
report A(3 glok) demonstrate its effective equivalence to A(8"i g0x). (PX)). These values are not exactly the same due to differences in regularization strategies.
Vt. P: = o
3l0R
(53)
Here 8/08 indicates the column of glo8; these are the entries glog corresponding to predictions for the ith word in the vocab. Now it is easy to arrive at the equivalence between logistic regression and our proxy model. we restrict the rank of glug to be at most min(d, [0),14) by factoring it as glog = E where E Rlolxd and Rdx(2), then the logistic regression looks like:
Pt σ
W Ex.)
(54)
1=1
which is equivalent to proxy model:
VI, o T, Wen))
(55)
Since d = 256 for the proxy model, which is larger than |0| 2 in the classification case, the proxy model is not rank limited and hence fully equivalent to the logistic regression model. Therefore the 38% can be interpreted as "keywords" in the same way that the logistic regression weights can.
To empirically verify this equivalence, we trained logistic regression model with E2 regularization on each of our classification datasets. To pick the optimal regularization level, did a sweep of regularization coefficients across ten orders of magnitude and picked the one with the best validation accuracy. We report results for A(g"f, ,lo0) in comparison A(guf. ,Pxx) in Table 10
Note that these numbers are similar but not exactly equivalent. The reason is that the proxy model did not use /2 regularization, while logistic regression did.
Theeee numbers were obtained from retrain of all the models in the main table, so for instance, the LSTM model used to produce gu might not he exactly the same the one used for the results in all the other tables due random seed difference.
23


Reviews:
[{"review_id": "gVw-WH_pcfD", "reviewer": "Reviewer_AtcN", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "This paper presents a framework explaining how attention might be learned. First, the model would get the knowledge to translate individual words (KTIW) based on word co-occurences, which can be learned if the attention weights are uniform. KTIW then drives the learning of the attention mechanism.", "summary_of_the_review": "The paper proposes an interesting and plausible framework of how attention is learnt. However, it does not clearly show that the 2 stages of learning happen in practice (is it really sequential?). Some of the assumptions may be too restrictive.", "strength_and_weaknesses": "Strengths\n\nThe experiment showing that copying sequences is a difficult task under some constraints is insightful.\n\nThe paper shows how multi-head attention can improve learning dynamics.\n\nWeaknesses\n\nThe paper presents a plausible 2-stage learning approach for attention, but it doesn't really show that this happens in practice. For example, we could monitor the entropies of the output distributions and of the attention weights during training.\n\nThe simplifying assumptions are quite strong. Under the proposed framework, we may not be able to explain phenomena such as Figure 9 in [1], where attention does not necessarily match work alignment.\n\n[1] Koehn and Knowles. Six Challenges for Neural Machine Translation. First Workshop on Neural Machine Translation. 2017", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "M83_iemXCc", "reviewer": "Reviewer_8y4q", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "The paper aims to understand the training dynamics of the Attention mechanism through lexical prob $\\beta$ and its learning proxy model.\n", "summary_of_the_review": "All the observations made are interesting, and the intuitive connections are also very fascinating. However, I still fear that these analyses might not have many practical applications and even won't hold in many different learning dynamics.", "strength_and_weaknesses": "Strength:\n1. The paper did some rigorous analysis of the attention training mechanism and drew a connection to the word alignment of the classical models.\n2. There is a connection to interpretable classification for the mechanism described in the paper.\n\nWeakness:\n1. The applicability of such an understanding is now clear.\n2. The analysis has a lot of underlying assumption (almost all of them are pointed out by the authors) which does not usually hold.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6nDKruhWf7", "reviewer": "Reviewer_H7cW", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "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": "", "summary_of_the_paper": "This paper investigates how seq2seq attention weights are learned in a machine translation setting. The paper defines a proxy to measure a model’s ability to translate individual words, which they shorten to KTIW. The paper claims that this measurable quantity is a driver for the model to learn to attend well. The evidence presented in favour of this hypothesis is that KTIW can be learned (it is similar to learned attention weights) when attention weights are frozen and  uniform, but that when KTIW is not learned, neither are the attention weights. This suggests a sort of necessary condition on KTIW for attention learning, which the paper claims is a causal driver for attention learning.\n\nThe paper suggests that it is therefore possible to reduce the problem of understanding attention learning to the problem of understanding KTIW (which does not seem to be entirely true, as KTIW may be necessary but not sufficient condition for learning of attention weights).\nThe paper thus proposes a proxy model for approximating KTIW learning (that is, a proxy for the proxy for attention weight learning), and verifies this. The paper claims that attention weights are learned in two stages: first KTIW is learned through word co-occurence statistics, and second, a learned KTIW drives the learning of attention.\n", "summary_of_the_review": "This paper investigates how seq2seq attention weights are learned in a machine translation setting. The paper defines a proxy to measure a model’s ability to translate individual words, which they shorten to KTIW. I found this paper very hard to follow, as the writing, justifications and logical flow of the word hid the scientific contributions. The paper would benefit from significant clarifying and removal of experiments (e.g. section 5.3) that detract from the points the authors wish the readership to takeaway - this might enable more readers to see the scientific contributions of the paper to the wider literature more clearly.", "strength_and_weaknesses": "*Strengths:*\n- Aiming for a theoretical understanding of learning dynamics and attention is an admirable goal. The high level topic would be of interest to many in the community.\n\n*Weaknesses:*\n-  The major weakness of this work is that I found the text and presentation of results really difficult to parse. It is also difficult to establish what exactly each experiment contributes to the logic of the paper, why they were designed in that way, or the takeaways from each expreriment, which made it quite a difficult read.\n\n- What are the implications of the results in this paper? I think I am missing the major takeaway from this work. It is not clear to me how defining this proxy for machine translation is useful for interpretability and understanding machine translation learning dynamics.\n\n- The KTIW proxy seems to assume a 1:1 ratio of input to output words for translation, which surely cannot hold true in the general MT case? Perhaps I am missing something here. \n\n- If I understand correctly, the logic is section 5.3 seems to be circular. The paper defines the initialisation of an attention head as ‘good’ if training with that single head converges (on a task that requires only a single attention head). The paper then constructs multi-head attention scenarios with all bad initialisations, all random, or only-one-good head initialisation. They present the result that if all initialisations are bad, they model does not learn, but if at least on is ‘good’ then the model will converge. As the task is learnable with a single attention head, this logic seems circular and trivially true. It is not clear to me what it adds to the paper and seems off-topic.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject