On the Dynamics of Training Attention Models
Authors: Haoye Lu
Yongyi Mao
Amiya Nayak
Venue: ICLR 2021 Poster (2021)

Abstract: The attention mechanism has been widely used in deep neural networks as a model component. By now, it has become a critical building block in many state-of-the-art natural language models. Despite its great success established empirically, the working mechanism of attention has not been investigated at a sufficient theoretical depth to date. In this paper, we set up a simple text classification task and study the dynamics of training a simple attention-based classification model using gradient descent. In this setting, we show that, for the discriminative words that the model should attend to, a persisting identity exists relating its embedding and the inner product of its key and the query. This allows us to prove that training must converge to attending to the discriminative words when the attention output is classified by a linear classifier. Experiments are performed, which validate our theoretical analysis and provide further insights.

Published as a conference paper at ICLR 2021
ON THE DYNAMICS OF TRAINING ATTENTION MODELS
Haoye Lu, Yongyi Mao & Amiya Nayak School Electrical Engineering and Computer Science University of Ottawa
Ottawa, KIN 6N5, Canada (hlu044, ymao, nayak) @uottawa.ca
ABSTRACT
The attention mechanism has been widely used in deep neural networks as model component. By now, has become critical building block in many state-of-the-art natural language models. Despite its great success established empirically, the working mechanism of attention has been investigated at sufficient theoretical depth to date. In this paper, we set up simple text classification task and study the dynamics of training simple attention based classification model using gradient descent. In this setting, should attend to, a persisting identity exists relating its embedding and the inner product of its key and the query. This allows to prove that training must converge to attending to the discriminative words when the attention output is classified by a linear classifier. Experiments performed, which validate our theoretical analysis and provide further insights.
show that, for the discriminative words that the model
INTRODUCTION
Attention-based neural networks have been broadly adopted in many natural language models for
machine translation (Bahdanau et al.] 2014; Luong et al. (015). sentiment classification (Wang et al.) 2016), image caption generation (Xu et al. 2015). and the unsupervised representation learning (Devlin et al.) 2019). etc. Particularly the powerful transformers (Vaswani al. 2017). attention its key ingredient.
Despite its great successes established empirically, the working mechanism attention has not been well understood (see Section 2). This paper sets up a simple text classification task and considers basic neural network model with the most straightforward attention mechanism. We study the model's training trajectory understand why attention can attend to the discriminative words (referred to as the topic words). More specifically, in this task, each sentence is treated as bag of words, and its class label, topic, is indicated by topic word. The model consider involves a basic attention mechanism, which creates weighting factors to combine the word embedding vectors into a "context vector"; the context vector is then passed a classifier.
In this setting, we prove closed-form relationship between the topic word embedding norm and the inner product of its key and the query, referred to as the "score", during gradient-descent training. It is particularly remarkable that this relationship holds irrespective of the classifier architecture or configuration. This relationship suggests the existence of a "synergy" in the amplification of the topic word score and its word embedding; that is, the growths of the two quantities promote each other. This, in turn, allows the topic word embedding to stand out rapidly in the context vector during training. Moreover, when the model takes fixed linear classifier, this relationship allows rigorous proof' of this "mutual promotion" phenomenon and the convergence of training to the topic words.
Our theoretical results and their implications are corroborated by experiments performed on synthetic dataset and real-world datasets. Additional insights are also obtained from these experiments. For example, low-capacity classifiers tend give stronger training signals to the attention module. The "mutual promotion" effect implied by the discovered relationship can also exhibit itself as "mutual suppression" in the early training phase. Furthermore, in the real-world datasets, where perfect

---

Published as a conference paper at ICLR 2021
delimitation of topic and non-topic words does not exist, interesting training dynamics is observed. Due to length constraints, all proof' are presented in Appendix.
2 RELATED WORKS
Since 2019, a series of works have been published to understand the working and behaviour of attention. One focus of these works pertains to understanding whether an attention mechanism can
provide meaningful explanations (Michel et al. 2019, Voita et al.) 2019: Jain & Wallace 2019
Wiegreffe & Pinter 2019 [errano & Smith 2020: Vashishth et al.) 2020). Most of these works are empirical in nature, for example, by analyzing the behaviours of a well-trained attention- baasd
model (Clark et al.; 2019), or observing the impact of altering the output weights of the attention module or pruning a few heads (Michel et al. [2019] Voita et al. 2019), or a combination of them (Jain & Wallace 2019; Vashishth et al.) 2020). Apart from acquiring insights from experiments, Brunner et al. (2019) and Hahn (2020) show theoretically that the self-attention blocks lacks identifiability. where multiple weight configurations may give equally good end predictions. The non-uniqueness of the attention weights therefore makes the architecture lack interpretability.
As a fully connected neural network with infinite width can be seen as a Gaussian process (Lee et al., 2018). few works apply this perspective to understanding attention with infinite number of heads and infinite width of the network layers (Yang 20199) Hron et al. 2020). In this paper, we restrict our study to the more realist non-asymptotic regime.
PROBLEM SETUP
Learning Task To obtain insights into the training dynamics of attention models, we set up simple topic classification task. Each input sentence contains m non-topic words and one topic word indicating its topic. Note that topic may have multiple topic words, but a sentence is assumed to include only one of them. Assume that there are J topics that correspond to the mutually exclusive topic word sets T1, T2. ... Ty. Let T = U Ty be the set of all topic words. The non-topic words are drawn from a dictionary O, which are assumed not to contain any topic word.
The training set V consists of sentence-topic pairs, where each pair (x+y) is generated by (1) randomly pick a topic y E {1, 2, ... ./} (2) pick topic word from set Ty and combine it with m words drawn uniformly at random from θ to generate the sentence (or the bag of words) X. In this task, one aims to develop a classifier from the training set that predicts the topic y for a random sentence generated in this way.
We will consider the case that |0| >> [T), which implies that a topic word appears much more frequently in the sentences than a non-topic word.
Attention Model For this task, we consider a simple attention mechanism similar to the one proposed by Wang et al. (0116). Each word w is associated with two parameters: an embedding Vn E Rd and a key Kew E F.d Based on a global query g E Rdd, the context vector of sentence X is computed by p(x) Ewex Var exp(g" kw] where [xx) Σ w'Ex exp(g" "sp(g""wwr) Then p(x) is fed into classifier
7(x)
that predicts the sentence's 's topic in terms of distribution over all topics
Denote the loss function by ((x, y). Our upcoming analysis implies this attention model, although simple, may capture plenty of insight in understanding the training of more general attention models.
Problem Statement Our objective is to investigate the training dynamics, under gradient descent, of this attention model. In particular, we wish to understand if there is an intrinsic mechanism that allows the attention model to discover the topic word and accelerates training. Moreover. we wish to investigate, beyond this setup, how the model is optimized when there is no clear delimitation between topic and non-topic words, as in real-world data.
'The condition that the attention layer directly attends to the word embeddings merely serves to simplify the analysis in Section but this condition is not required for most results presented in Sections and 5 More discussions are given in Appendix Alin this regard.
2

---

Published as a conference paper at ICLR 2021
4 THEORETICAL ANALYSIS
is common to fix some parameters when we train a model with limited resources. Also
Lemma 1. Assume q * 0 when initialized. Fixing it does affect the attention block's 's capacity.
Thus, our upcoming discussion focuses on the case in which the query is fixed. Doing so also allows establish a closed-form expression connecting the word's ' embedding and the inner product of its key and the query. In Appendix B extra discussions and experimental results reveal that the trainability of the query does not affect the fundamental relationship we are about present.
For topic word E, let Vr denote the training samples involving it. Then, by gradient descent,
exp(q* Nx)
A
Σ
VF(()((6,))
7(x)
(G.)EV,
exp(g" kc) z(x)
AKr
Σ
g(u--(xx))
Vxxol(x.1)
(X-))E
where denote the learning rate. As it will turn out, important quantity in this setting is the inner product q" kw of query q and the key kw, which we denote by and refer to it as the score of the word w.
Denoting the = |1 rll4 la. 0(x)= = Ewex exp(sw] Ow and h(T(x):)) Voool(, v), for
a topic word I, the dynamics simplifies to
Av
7
Σ
h((():v)
exp($t)
(3)
T
Z(x)
(x.))E.
Ast
n
Σ
(c - ū(x)) h(((x::)
exp(s:)
(4)
V
Z(x)
(x.))e,
In the rest the paper, whenever we refer to the embedding of word f, we actually mean v not Vt-
Our analysis assumes the word embeddings sampled i..d. from distribution with mean zero and variance 4 where 2² is assumed close to zero. The word keys and the query are also sampled from zero mean distributions with a possibly different variance. We assume that this variance is so small that the initial word scores are approximately zero. This assumption of the initial configurations corresponds to the attention model starting how the model deviates from this initial setting with training. We also assume the derivative ((((x) of l is Lipschitz continuous in o(x) throughout training. Further the assumption in Section 3 that the number of non-topic words (0] is much larger than the number of topic words [T] implies that with a sufficient number of training samples, the occurrence rate of topic word is significantly higher than the non-topic ones. This then justifies the following assumption we will use throughout our analysis.
word-averaging model, and allows us investigate
The scores and the embeddings of the non-topic words are nearly unchanged compared to their counterparts for the topic words.
Assumption 1.
Hence, our upcoming analysis will treat the scores and embeddings of the non-topic words as constants. Assumption [
will be validated by experimental results presented in Section
By selecting a sufficiently small n. we can take the gradient-descent updates in Eq and Eq to its continuous time limit and get"
du: 17
exp(s:)
Σ h(((x):v)
(5)
dt
Z(x)
(xW)ev.
ds;
"
Σ
(u. - -(x)) h(T(x):0)
exp(st)
(6)
dt V
Z(x)
(x.iER,
Reverselyy Eq (3 is discretized approximation of Eq Me (t+ 1. dor(t') 1- de
Avx(t). The approximation becomes accurate if n, (t + 1 close to nt (t()), which can be achieved by choosing sufficiently small 17. Likewise, Eq (4) is discretized approximation of (
3

---

Published as a conference paper at ICLR 2021
We then characterize the update of the score and the embedding of topic word continuous- time dynamical system stated in Lemma [
The same technique has been used to analyze the training
of neural networks in other contexts (Saxe et al. 2014 Greydanus al. 2019).
Lemma 2. For sufficiently small and o², the score 8 and embedding of topic word t sattisfy
du:
n/
exp(st)
dt | h(((():0)
[(x)
V.
ds;
(vx- (0(x( (o(x) v.)
du dt
(xx((5)) + Z(x (t)
(8)
dt
Z(x (2)
where Z(x\t) = Ewexkty exp(sw). 6(x (1) = = =wexx(1 Uw explaw) Z(xa)
and (>),, denotes taking
sample mean over the set V..
Eq 7 implies the speed of moving vs along the direction of (h(*(x): v)),. is controlled by the attention weight exp(sr)
[(x) Eq (8) shows that u2 increases if and only if e has greater projection on (h(ū(x):v) than the weighted average of the non-topic word counterparts.
Consider simplified case where (h(6(x); 9) v, is fixed. Since the change of e, is much faster than the non-topic word counterparts, 22 will have larger projection on (h(ū (xx): y)) after a few epochs of training. Then 8, increases well its attention weight, which in turn speeds up the extension of the embedding Uz. This observation reveals a mutual enhancement effect between the score increment and the embedding elongation. In fact such effect exists in general, stated in the theorem below, irrespective of whether (f(o(x); v) is fixed.
Theorem 1. In the setting of Lemma 2 from epoch to to t1, the topic word score 8t and its embedding Ut satisfy
1 2 (le((t) (5(x \ ())
st(1) + (t)(s
7(x \ {t)
=
Ve
Following from Lemma 2 this theorem implies positive relationship between the topic word score and the distance between it and the non-topic word embedding average (ō(x ()),, Remarkably this result makes no reference to (h(#(x); v- hence independent of it. This implies the identity in Eq ( holds irrespective of the choice and setting of the classifier. Theorem [ further implies a score and embedding norm ("SEN" in short) relationship for the topic words:
Corollary 1. In the context of Theorem !
by setting to and t1 t. Eq is reduced to
exp st))s(t) 1 m
[|ex(t)ll2 V2 (t(())+
(10)
m
The corollary indicates that [|v. (t) is monotonically increasing with 8, (t). So, 8, increases if and only if the point v, departs from its initial location. That is, the norm of the topic word embedding increases, it will be attended to. This result is independent of the configuration of all other network layers. Thus, if (h(v(x): y)) v. has a gradient field that pushes ii away from its original location, the topic word is expected to be attended to. This statement can be made precise, in Theorem 2 when the model uses a linear classifier.
Theorem 2. Assume the model has a fixed classifier in the form c(ū(x)) = softmax(U" v(x)). where the columns of U are linearly independent, and the model is trained using gradient descent with the cross-entropy loss. As training proceeds, the model will attend to the topic word in every inpus sentence and have its training loss approach
It is notable that the theorem holds broadly for any arbitrary fixed linear classifier (subjective to the mild linear independence constraint of its parameter U). Additionally, we anticipate that this result holds for much wider family of classifiers including trainable and even nonlinear ones. But rigorous proof appears difficult to obtain in such settings, and we will corroborate this claim in an experimental study in Section 5
4

---

Published as a conference paper at ICLR 2021
To up, in this section, we have shown two main results: (a) there is closed-form positive relationship, the SEN relationship, between the topic word score and its embedding norm, which is independent of the configuration of the classifier. (b) the model, equipped with fixed linear classifier stated in Theorem 2 can be trained to have all topic words attended to.
5 EXPERIMENTAL STUDY
In this section, we first test our model on artificial dataset generated through the procedure introduced in Section 3 The test corroborates our theoretical results and validates their assumptions. Our test results suggest that the attention mechanism introduces a synergy between the embedding and the score of topic words.
Another experiment is performed on the real datasets SST2 and SST5 (Socher et al. 2013). The experiment results suggest that the SEN relationship of topic words holds at least in imitial training stages. As training proceeds, some words appear to deviate from the theoretical trajectories. Further analysis of this behaviour provides additional insights into the attention model's training dynamics on real-world datasets, often possessing much more complex structure as well as rich noise. We performed all experiments using PyTorch (Paszke et al. 2017).
We performed our experiments on three models, Attn-FC, Attn-TC and Attn-TL, having the same attention block but different classifiers. The first two have the classifier in form c(ō(x)) = softmax(U" ((x)) and the last in form c(ū(x)) softmax(U?] (()) +2)). Except that the U in Attn-FU FC is fixed, other parameters of the three models are trainable and optimized using the cross-entropy loss.
Since a real-world dataset does not have topic word as the sentence topic indicator, we introduce a word "topic purity" measurement to facilitate our discussion on the experiments performed on SST2.
Let t and 8 (w) respectively denote the portions of the positive and negative sentences among all training samples containing word w. Then the topic purity w is d(w) = 18+(w) - -8- 8-(w)]. 6(w) = 1, w is either pure positive or negative topic word. If 8(w) 0, o+ (ww) o d-(ww 0.5. which implies w has a completely random topic correspondence.
5.1 EXPERIMENTS ON SYNTHETIC DATASETS
The artificial dataset, consisting of 800 training and 200 test samples, is generated through the procedure introduced in Section 3 The dataset has four topies, and each contains two topic words. There are 20 non-topic words per sentence and the non-topic word dictionary size M = 5,000.
Our experiments use the same embedding dimension 15 for all three models. Regarding the classifiers, Attn-FC and Attn-TC adopt U R 15x4 while Attn-TL takes U1= E R 15x by E R 1o Un- E R 10x4 and by IR. For the validation Theorem 2 the columns of in Attn are set to orthonormal and thus linear independent. Unless otherwise stated, the scores are set to zero and the embeddings
are initialized by normal distribution with mean zero and variance e 10-6, We trained the models using gradient descent with learning rate 7 0.1 for 5K epochs before measuring their prediction accuracy on the test samples. When training completed, all three models achieve the training loss close to zero and the 100.0% test accuracy. which implies the trained models perfectly explain the training set's variations and have a good generalization on the test set.
Verification of Assumption 1 and validation of Corollary I and Theorem 2, We repeated the experiments for five runs and plotted the empirical score distributions of the non-topic and topic words of the three well-trained models with their 95% confidence intervals in the first two graphs of Fig Compared to the topic words, the scores of the non-topic words are nearly unchanged throughout the entire training process. Likewise, the next two plots show the embedding norms of the non-topic words are nearly constant, too. This implies Assumption [ indeed holds. Fig plots the empirical and the theoretical SEN curves of randomly picked topic word for the three models, where the theoretical curve has the expression stated in Eq (1O). The coincidence of the empirical and the theoretical curves in all three models validates the SEN relationship stated in Eq (10) and
5

---

Published as a conference paper at ICLR 2021
2.0
1.
Am,
10-
0/8
0.3
4ATTC
AAXUTL
0.-
6.8
644
0.4
6.2
0.2
02
6.0
0.0
022 2.11 0.8 316
5:10
0.
hagitapiceee sord cbiibiiiin
T:2 word dist buzon
Non-copic
Top woord wmooddiing narm datritatinn
Figure The first two histograms plot the distributions with the 95% confidence interval of the non-topic and topic word scores for the three well-trained models, Attn-FC, Attn-TC and Attn-TL by repeating the experiments five times. Likewise, the next two plot the ones of the embedding norms. The data of the models are overlaid in each histogram for easy comparisons.
Exp
Teer wAr enbrodtr rem
Ppp
Dp
Pp
Tep Topcant anboaaiis
Exp
Ex
3
Ex
Exp
word
Scate FQ
Scord
Scare
Figure 2: For Attn-FC, Attn-TC and Attn-TL, the graphs in order respectively plot the SEN relation- ship of a randomly picked topic word along with the theoretical predication derived in Corollary [ For each model, we repeated the experiments five times and selected the same topic word.
its independence of the later layers Moreover, Fig (left two) shows the scores of the topic words exceed the non-topic word counterparts by two orders of magnitude, which implies the topic words are attended to in well-trained model. As we have reported earlier, Attn-FC has the training loss roughly zero when the training is completed, Theorem i confirmed.
Lower-capacity classifiers result in stronger attention effects. The comparison, among the SEN distributions of the topic words in the three models, implies that lower-capacity classifier leads to greater topic word SEN, which means more drastic attention decision. This happens because the classifier of a larger capacity can explain more variations of the sample distribution and has more freedom to accommodate and absorb the correcting gradient signals. As result, the attention layer receives a weaker gradient on average, which makes the embeddings of the topic word extend less from the original point. Thus, implied by Eq (1O), the magnitudes of the scores of the topic words are dampened, and therefore a weaker attention effect will be expected. This observation hints that if we know attending the right words explain most of the variations in sample distributions, we should consider a low capacity classifier (or later layers in general). Alternatively, we may also freeze the classifier in the initial stage of a training process, forcing the attention layer explain more variations. Remarkably, all these modifications do not affect the relationship stated in Eq (10).
Synergy between score growth and embedding elongation in topic words. Eq 10 implies positive SEN relationship of the topic word. That is, larger topic word embedding norm results in larger score (and thus a larger attention weight), which turn makes the embedding extend faster. To corroborate this claim, we performed an ablation study by considering two variants of Attn-TC. The first has the scores frozen (referred as (Cn-TCC-FF) and the second has the embeddings fixed (referred as Attn-TC-EF). In this experiment, the embeddings of the models are initialized by normal distribution of mean zero and variance g² descent for 60K epochs with learning rate 17 = 0.1, which are sufficient for all three models to fully converge. All three trained models reached 100% test accuracy.
d 0.1. We trained all three models by gradient
II Appendix B trainability of queries do not affect the positive SEN relationship. Besides, the query fixed model has very similar training dynamics the with trainable query and large initial
the experiments with trainable queries are also implemented. The results indicate that the

---

Published as a conference paper at ICLR 2021
Mit 10-3
08 BRLTE
680
10.00
2003
2055
3950
1893 2000
300
Spc
3K
289
X1
Figure 3: The changes of the SEN and the training loss for Attn-TC, Attn-TC-KF and Attn-TC-ER. The first two graphs demonstrate the evolution of a topic word SEN when the mutual enhancement happens. The third shows the change of the training loss. The last two show how the topic word SEN change in a separate training when the mutual diminution occurs.
Trearitca
..X Thoorsical
02
Score stppl
Theaetti
Thortia
WeT semDen sm
Thasmiic Sppesveell
Some (N0
Figure 4: The SEN relationship in the training process of Attn-FC (left two columns) and Attn-TC (right two columns) on SST2 (upper row) and SST5 (lower row). The words, picked among the ones having the largest five scores, have the SEN curves that agree with their theoretical counterparts.
The first three graphs of Fig describe the evolution of the three models in the first 3K training epochs. For randomly picked topic word, the first plot shows its score in Attn-TC grows faster than the one in AttF-C-EFF. Note that the score in Attn-TC-EF finally surpasses the one Attn-TC because Attn-TC has converged at around 1K epochs. Likewise, the word embedding norm in Attn-TC increases more rapidly than the one in Attn-TC-KF before Attn-TC converges. The observations imply the attention introduces a mutual enhancement effect on training the topic word's 's score and its embedding, which makes Attn- ACC enjoy the fastest training loss drop as shown in the third plot.
The mutual enhancement could become the mutual diminution in the early training stage if the initial embedding of the topic word has a negative projection on the direction that will move along. This effect can be precisely characterized by Eq (8) Assume the embedding of topic word is initialized to have a smaller projection on the gradient, passed from the classifier, than the average of the non-topic words. The reversed order of the projections makes the score of the topic word decrease as its embedding has a "negative effect" on the training loss compared to the average of the non-topic word embeddings. This will, in turn, impede the elongation of the embedding vector or even make it shrink (see the last two plots of Fig The "negative effect" cannot last long because the topic word embedding moves along the gradient much faster than the non-topic words due its high occurrence rate in the training samples. By Eq (8) again, das will finally become positive. That is, the score of dt
the topic word starts to increase and its attention weight will surpass the one of the word-averaging model (see the second last plot of 3). Then, we start to observe the mutual enhancement effect, which is indicated in the last plot: the increase speed of the Attn-TC's embedding norm exceeds the Attn-TC MCn-TC-KF's since around the 370-th epoch.
5.2 EXPERIMENTS ON SST2 AND SST5
The second part of the experiment is performed on datasets SST2 and SST5, which contain movie comments and ratings (positive or negative in SST2 and one to five stars in SST5). For simplicity, we limit our discussion on Attn-FC and Attn-TC using the same configurations of our previous experiments except that the embedding dimension is set to 200. Remark that goal is not to find a state-of of-the-art algorithm but to verify our theoretical results and further investigate how an attention-base based network works.
7

---

Published as a conference paper at ICLR 2021
Scare MM porrd
J0P Bouth
0.2-
piverti
035
0.23
Eaperrerta
6.30
0.7
22: 50 15
15 205
415
+13 Byen
Figure 5: The juxtaposition of powerful and best during the training of Attn-FC (the first row) and Attn-TC (the second row). In each row, the first two graphs are the empirical and the theoretical SEN curves of words powerful and best. The next two plot their SEN as a function of the number of epochs. The last shows the changes of their weights in sentence that they co-occur.
For both SST2 and SST5, we trained the two models by gradient descent with learning rate 0.1 combined with the early stopping technique (Prechelt 2012) of patience 100. As Py Torch requires equal length sentences in batch, pad all the sentences to the same length and set the score of the padding symbol to the negative infinity. Under this configuration, the trained Attn-FC and Attn-TC reached 76.68% and 79.59% test accuracy on SST2 and 38.49% and 40.53% on SST5.
Validation of Corollary T. As the true topic words are unknown in a real dataset, we checked the words of the largest fifty scores after the training is completed. We observed that most of the words have their SEN curves close to our theoretical prediction. We picked two words for each model-dataset combination and plotted the curves with their theoretical counterparts in Fig 4
The competition of two topic word candidates of various occurrence rates and topic purity, To better understand how the attention block works, we investigated the case when the empirical and the theoretical SEN curves disagree. We limit our discussion on the model trained on SST2.
For both Attn-FC and Attn-TC, we noticed that there are mainly two types of deviations, as shown in the first two columns of 5 In the first one, the theoretical curves overestimate the score growth of powerful in terms of the embedding norm, while best shown in the second column experiences a score drop combined with failed theoretical prediction in the final training stage. These two types of the disagreement in fact occur in pairs, which caused by a pair of words that one of them frequently appears in the training samples but has a low topic purity, while the other has the opposite.
Regarding the (powerful, best) pair, best appears in 128 training samples while 103 of them are positive. In comparison, powerful appears in the 36 training samples, and all of them are positive. The large difference the number of samples makes the embedding norm of best extend much faster than powerful in the initial training process (Fig 5] third column). As the positive SEN relationship implics, a quicker embedding norm increase of best results in its faster score growth. Therefore, best will be more attended (Fgg5) 5 accelerates its SEN increase (this is the synergy relationship between the embedding and score that has been demonstrated in the 'ablation' study). This process does not stop until the gradient from the classifier diminishes because of its low topic purity (which is similar to the case when the label smoothing (Szegedy et al. 2016) is applied for alleviating the overfitting problem). In contrast, powerful ful is less trained initially due to its low occurrence rate. But its high topic purity makes the direction of the gradient stable, in which its embedding will steadily elongate. The elongation will, at last, let the embedding have a greater projection on the gradient vector than the average of the other words appearing the same sentence. Thus, the score of power ful starts to increase as shown by Eq (8) and plotted in the second last column of Fig 5 In contrast, as the gradient magnitude drops, the embedding of best will extend a decreasing speed; and its projection on the gradient, passed from the classifier, will finally be surpassed by the words co-occurring with it but having higher topic purity (like power fuwerful). Thus, its score starts to drop eventually.
last column), which thus further
The dynamics of topic purity of attended words The analysis of the inconsistency between the empirical and theoretical SEN curves hints that the disagreement are strongly related to the word's topic purity and its occurrence rate. To better characterize their dynamics, for word w, let (t)
8

---

Published as a conference paper at ICLR 2021
15
433
Apwss var Tpak 1000 Augult
steliw Annrd Asai?
vit Youa M0 Autme!
SEP AMaRO Awni
10--
10 Augad!
A
Aug
Figure 6: The average topic purity (red) and number of occurrence (blue) of the words attended to in the sentences containing the training of Attn-FC (first row) and Attn-TC (second row) proceeds. Word w picked from the set of words having the largest five scores in the well-trained models.
+10 Paurh
Figure 7: The average topic purity (red) and number of occurrence (blue) of the words attended to in all the training sentences the training of Attn-FC (left) and Attn-TC (right) proceeds.
be the list, at epoch t, that records the attended word (or of the largest score) in the sentences containing w. If multiple words in sentence have the same score, randomly pick of them. Note that Aw may contain repeating words as a word could he attended in multiple sentences. We selected w be the words having the largest five scores in the well-trained Attn-FC and Attn-TC, respectively. At various epoch t, Fig 6 plots how the average of topic purity 5(w') evolves for w/ E (.w(tt) as well as the average number of occurrence the training samples. For both models, they initially attend to the words that mostly have low topic purity with high occurrence rate. As the training proceeds, the average topic purity of the attended words increases while the average occurrence rate drops. At the end of the training, almost all the attended words have a close-to-one topic purity. Fig shows the evolution of the average topic purity and the average occurrence rate of the attended words over the entire training set. While similar changing pattern can be observed, the average topic purity is lower than the one presented in Fig 6 when the training is completed. We argue that this happens because topic purity words have a too low occurrence rate to be sufficiently trained.
sentences do not have any high topic purity words or their high
CONCLUSION
This paper investigated the dynamic training a series of attention-based bag-of-word classifiers on simple artificial topic classification task. We have shown a persisting closed-fone positive SEN relationship for the word to which the model should attend to. This result independent of the configurations of the later layers. Through the result, have proved that the model must converge in attending to the topic word with the training loss close to zero if the output the attention layer is fed into fixed linear classifier. A list of experiments has confirmed these results.
The experimental results indicate that the attention block intends to make a more drastic decision its later layers have a lower capacity. This hints the classifier's limited capacity may help if "selecting" the right word explains most of the variations in the training samples. An ablation study shows synergy between the topic word score's 's growth and its embedding elongation, leading training loss drop than the fixed score and fixed embedding variants. Besides, we have shown that this "mutual promotion" effect can also exhibit itself as "mutual suppression" in the initial training stage.
faster
We investigated the competition of two topic word candidates with large differences in the topic purity and the occurrence rate in the training samples. The words of higher occurrence rate but possibly low topic purity are more likely to be attended to initially. However, as the training proceeds, the attended words are gradually replaced by those of higher topic purity.
9

---

Published as a conference paper at ICLR 2021
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. In International Conference on Learning Representations (ICLR), 2014.
Sergei G. Bobkov and Christian Houdrè. Variance of Lipschitz functions and an isoperimetric problem for class product measures. Bernoulli, 2(3):249-255: 1996.
Gino Brunner, Yang Liu, Damián Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Watten- hofer. On Identifiability in Transformers. In International Conference on Learning Representations (ICLR), 2019.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What Does BERT Look at? An Analysis of BERT's Attention. arXiv preprint arXiv:1906.0453 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep
Bidirectional Transformers for Language Understanding. In Annual Conference of the North American Chapter of the Association for Computational Linguisties (NAACL-HLT), 2019.
Sam Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian Neural Networks. In Neural Information Processing Systems (NIPS), 2019.
Michael Hahn. Theoretical Limitations of Self-Attention Neural Sequence Models. Transactions of the Association for Computational Linguistics, 8:156-171 1711,,1,200 2020.
Jiri Hron, Yasaman Bahri, and Jascha Sohl-dickstein Roman. Infinite attention : NNGP and NTK for deep attention networks. In International Conference on Machine Learning (ICML), 2020.
Sarthak Jain and Byron C. Wallace. Attention not Explanation. In Annual Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HI.T), 2019.
Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks gaussian processes. In International Conference Learning Representations (ICLR), 2018.
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention- based neural machine translation. In Empirical Methods in Natural Language Processing (EMNLP), 2015.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Neural Information Processing Systems (NIPS), volume 32. Curran Associates, Inc., 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Detito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic diffcrentiation in pytorch. In Neural Information Processing Systems (NIPS), 2017.
Lutz Prechelt. Early Stopping But When?, pp. 53-67. Springer Berlin Heidelberg, Berlin, Heidelberg, 2012.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations (ICLR), 2014.
Sofia Serrano and Noah A. Smith. Is attention interpretable? Annual Meeting of the Association for Computational Linguistics (ACL), 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher Manning, Andrew Ng, and Christopher Potts. Parsing With Compositional Vector Grammars. In Empirical Methods in Natural Language Processing (EMNLP), 2013.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the Inception Architecture for Computer Vision. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
10

---

Published as a conference paper at ICLR 2021
Shikhar Vashishth, Shyam Upadhyay, Gaurav Singh Tomar, and Manaal Faruqui. Attention Inter- pretability Across NLP Tasks. In International Conference on Learning Representations (ICLR), 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems (NIPS), 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi 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, July 2019.
Yequan Wang, Minlie Huang, Xiaoyan Zhu, and Li Zhao. Attention- LSTM for aspect-level sentiment classification. In Empirical Methods in Natural Language Processing (EMNLP), pp. 606-615, Austin, Texas, November 2016. Association for Computational Linguistics.
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. 2019 Conference on Empirical Methods in Natural Language Processing and 9th International Joint Conference Natural Language Processing (EMNLP-IJCNILP), 2019.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International Conference on Machine Learning (ICML), 2015.
Greg Yang. Tensor Programs I: Wide Feedforward or Recurrent Neural Networks of Any Architecture are Gaussian Processes. In Neural Information Processing Systems (NIPS), 2019.
A THE RESULTS OF THE MODEL THAT THE ATTENTION BLOCK ATTENDS TO A CNN LAYER
The main text focuses on model that has the attention block directly attends to word embeddings. Such design simplifies our analysis but should not be considered as pre-condition to keep our results valid. In particular, the positive SEN relationship generally holds regardless of other components of the network, which can be justified as follows. There are two correlated sources of gradient signals, one back proopaate from the score to update key and query, the other back pack-propagates from the classifier loss to update word embedding. The correlation governs the positive relationship between score and embedding norm. Although the integration of other modules makes the analysis harder, the relationship should persist. Hence, all the results depending this relationship keep valid.
We empirically verify our claims by implementing an experiment similar to the one discussed in the main text. We modified Attn-TC (named Attn-TC-CNN) by adding two parallel CNN layers to respectively process the word embeddings and the keys before feeding them into the attention block. Then, construct an analogous data generation process introduced in Section Finally, empirically show that Assumption I and the positive SEN relationship still hold.
The
has the same configurations the Attn-TC (introduced in Section 5.1 except that we added two parallel CNN layers to preprocess the word embeddings and the keys. Regarding the CNN layer processing the word embeddings, it has the kernel of size d 2 and stride 1. We used cl kernels to keep the word embedding dimension unchanged. So, for two consecutive words in sentence, the CNN mixes their embeddings and produce a new one. Given that the sentence has m + 1 words, the input embedding matrix has shape d x (m + 1) and the output has shape d x m. Likewise, regarding the keys, the CNN layer processing them has the kernel size d' x and stride 1. And there are d' kernels in total.
Consider classification problem containing two topics and B. The sentences of the two topics are generated by two Markov chains, respectively, which are constructed as follows:
1. Let L (i= = A, B) be two mutually exclusive sets of ordered word pairs. The word pairs do not contain repeating words.
2. For i A, B:
11

---

Published as a conference paper at ICLR 2021
Enqqencc
Carlandd oraind sur dereatie
Genteet tapk word cinbeddinggg hore CDIFRAcION
Figure 8: The first two histograms plot the distributions of the non-topic and topic word pair scores along with their 95% confidence intervals for the dell-trained Attn-TC-NN,, Likewise, the next two plot the ones of the centered embedding norms.
Dxp0
Dxp
Exg4
Dpp
Exp
Thea
Scrre
Figure 9: The empirical and theoretical SEN curves of randomly picked topic word pair for model AC-C-TC-CNN.
(a) Initialize Markov Chain (MC;) of words in the dictionary such that from word, there is a uniform probability of moving to any words (including itself) in one step. (b) Group the word pairs in T. according to the leading words. For each group, let denote the shared leading word and e, moving from to e; in one step be J and those to the rest zero.
1, 2. n) the second. Set the probability of
We call the word pairs in L; (i = A, B) the topic word pairs. For each pair of words, a new embedding and new key are generated by feeding their original embeddings and keys into the CNNs. We refer to the new embedding as the topic word-pair embedding and the new key as the topic word-pair keys.
Likewise, for any other word pairs, the new generated embeddings and keys are referred as the non-topic word-pair embeddings and keys.
Assume the training dataset contains sentences of length m + We generate it by repeating the following procedure:
1. uniformly sample a topic i E {A, B}.
2. uniformly sample word pair from L; and let the first word be the starting word.
3. sample another m words by running the Markov process for m times.
In our experiments, the word dictionary has 200 words, and so there are 40, 000 word pairs. The sentence length is set to 10. For i = A, B, each L, contains two randomly picked word pairs as the topic word pairs. Note that we have ensured that LA and Lg mutually exclusive. We used 4, 000 samples for training and 1, 000 for testing. The model was trained for 3, 000 epochs and achieved 99.89% test accuracy.
We repeated the experiments for ten runs and plotted the distributions along with the 95% confidence interval of the scores and the embedding norms of the topic and the non-topic word pairs in Fig X Note that, the word pair embeddings and the keys are generated from the CNN layers. So the initial embeddings and the scores are not very closed to the origin. To facilitate the verification of Assumption we centered them by subtracting their initial values before we plot their distributions.
From Fig g we observe that the non-topic word pair scores and the embeddings are nearly unchanged in comparison their counterparts of the topic ones. Therefore, we have shown that Assumption I is largely held even we have added CNN layers to process the word embeddings and the keys before feeding them into the attention block.
12

---

Published as a conference paper at ICLR 2021
The
Tharddaul
Lan Ecarimenta
Depertennl
Las Expervans
1ith 701
Scar 10t6
Trearaue
The siical
Theaniica
Thardeall
TodL: EYOLEE
Figure 10: The empirical and the theoretical SEN curves of a randomly picked topic word for Attn- FC (first row), AttnTCC (second row) and Attn-TL (third row) with trainable query. The query is initialized by normal distribution N (0, o²). From left to right, o² = 0.1, 0.5, and 2, respectively.
Randomly picking topic word pair, we plotted its empirical and theoretical SEN curves in ten runs in Fig The figure shows that the positive SEN relationship holds even if the attention layer attends to other layers instead of the word embeddings directly. Therefore, all the results due to the relationship keep valid.
B DISCUSSION AND EXPERIMENTAL RESULTS OF THE MODELS WITH A TRAINABLE QUERY
In Section 4 we assumed a fixed and non-trainable query which allows the derivation of a clean closed-form "SEN" relation. But it worth emphasizing that the positive relationship between the score and the embedding norm in fact exists regardless of the trainability of the query. As we have mentioned in Appendix A there from the score to update key and query, the other back- prok-propagates from the classifier loss to update word embedding. This correlation governs the positive relationship between score and embedding norm. Whether the query is trainable does not alter the existence of this correlation, although trainable query makes the analysis more difficult. In particular, when the query norm is large, the update of query is relatively negligible; thence, the training behaviour is similar to having fixed query.
correlated sources of gradient signals,
pock-propagates
To verify our claims, we reimplemented the experiments introduced in Section 5.1 with the same configurations except that the query trainable. Fig[O we plot the empirical and the theoretical SEN curves of randomly picked topic word by training Attn-FC, Attn-TC and Att-FFC with trainable query. We initialize the entries of the query vector by a normal distribution N N o²). From left to right, x² increases well the initial query norm. We observe that regardless of how the query is initialized, the positive SEN relationship always preserves. Moreover, as the initial norm increases, the empirical curve approaches the theoretical curve having the expression in Eq (10). As we have discussed, this asymptotic approach happens since an increasing initial norm of the query makes its change negligible during the training process compared to its already big enough norm.
13

---

Published as a conference paper at ICLR 2021
C PROOFS OF THE RESULTS IN SECTION 4
Proof of Lemma Assume there are in total N words in the dictionary (including both the topic and the non-topic words). Sample the keys of the N words arbitrarily to generate E IR xN with the key vectors as its columns. Randomly pick q E IRd as query. To prove the lemma, it is sufficient to show for any non-zero q E IR", there exists K Rd'' such that q" K = q K. Since $ & O, without loss of generality, assume its first entry ğ1 is non-zero. Let [s1, $2, ..- SN] = 9 K.
For = .., let the i-th column of K be q" q" K.
#i 0, 0. "0'" Then we can easily check that
Proof of Lemma dynamics,
Picking sufficiently small 1). a continuous time limit can be taken to obtain the
du
1
Σ h(o(x):0)
exp(s:)
(11)
dt
M
2(x)
(xW)Ev,
ds:
"
Σ (x.lEe.
(vz (0--6()) h(T())=))
exp(s1)
(12)
dt
V
2(x)
which are equivalent to
de;
Vel
((((():))) exp(st)
(13)
dt
V
7(x) we
ds;
exp(s1)
(vz (uz-vth o(x)) k(T(x)s8)
(14)
dt
2(x) Ve
As h(((x);)) is assumed Lipschitz continuous, we have
(vr - 0(0)), (M(66)::)
exp(tr)
exp(v:)
Lvdo²
7(x)
7(x)
4
w1
(15)
where L is the Lipschitz constant. Choosing small enough の² so that Lv is close to zero, we have
(M(00):))
exp(ve) Z(x) v.
~ (1v= v(x)) (0x ((0x=x)))(((((0000) exp(ux)
(we- - 5(x)1
2(x) v.
Then, combining it with Eq [4) yields
ds;
17 Mpl
exp(s:)
(ve b(x)) h(7():v)
dt
2(x)
(v.
v.
n Vel
(((((x)i))
exp(s1)
(Vr -xx))
du
(16)
2(x)
dt
21
*The derivation is given in Appendix
14

---

Published as a conference paper at ICLR 2021
Remarkably, the relation stated in Eq (16) is independent of (8(()):)) So the relationship does not depend on the architecture of the classifier (or later layers in general). Expand Eq (16).
dsc dt
du
(vr - T(x)),
Th v. dt
T
exp($1)
exp($w)
de
21 + E
U1
Cw
2(x)
2(x)
dt
mExifet
exp( (w) (1s - Uwr)
du
Σ
z(x)
dt
(EXi(()
cxp(sw)
vwly. T de
E
(er
2(x)
dt
w
V1
(exp (s. /Z \ (t))v, (z(x)/Z(x) \ ()))
T de
(we Uwl v. dt
wEX\(()
The second last step due to the independence between the score and embedding initialization, while the approximation in the last step is made as we assume all the scores of the non-topic words maintain the same during the entire training process. Rearranging the equation yields
ds; dt
Σ WEX
(xp(sw)
(1r Vwr vw) I du eep( (t) + Z(x \t)
-1
Z(x
o, dt
Z(x ()
T du dt
(eep(()) + Z(x ((x10) Z(x \t)
= =(1 - (F(x \ v))
where J(x (t) = Ewexi())
exp(aw) Z(x))
Proof of Theorem
1 By Lemma 2 we have
exp(s;)= Z(x \1)
ds:
duz
(vx (ū(x \ \11))
((x |t)
dt
dt
W.
which is
1 + exp(st)
1
dsc = (vx (ō(x) \ ()),,))
)7 du
(17)
Z(x \t)
v, dt
dt
By the fundamental theorem of calculus, integrating on both sides from to yields,
f1
1 |00
3t exp(s:)
(0(x \ (t)), 11
(18)
Z ve to
Proof of Corollary constant, we have Z and (7(x \ t)),, 0. As is initialized with mean zero and very small variance.
7
Since the scores and the embeddings of the non-topic words considered
[00(00] = Then, Eq 9 can be written as
[lvx(t)]12 /2 sx(() exp (t) 1 m 772
15

---

Published as a conference paper at ICLR 2021
Proof of Theorem 2 (sketch). Without loss of generality, pick topic word t and assume corresponds to the 9-thh topic. We prove the theorem by showing that as the number of epochs increases, for any sentence X in 1r the ip-th entry equals one and the rest zeros. Let U (x).). Notice that the loss function (- l--fosssotmax (x))) v, convex in terms of x. As U is fixed, the loss function is also convex in terms of o(x). This implies, if the model is optimized by gradient descent, the gradient will lead v(x) to its optimal solution 8* (x). In our case, as the columns of U are linearly independent. there exists vector that are orthogonal to all the columns of & except Ug. Without loss of generality, assume Up n > 0 (otherwise, choose inverse). Then potential optimal solution 8" (x) An for À goes to the infinity as U* (x) U; = O for * &p and 8* (x) Up - cc, which implies softmax TO (x) ep- Combined with the fact that there cannot be an optimal solution e** (x) of a finite norm such that DultmaxU""5**
oo and softmax(U" (xx)(x))))
Cp. where Ex is the one-hot vector that
= ep. the gradient must lead f(x) to an optimal solution that arbitrarily far away from the origin, which also applies to e as it receives the same gradient up to a multiple according to Eq (7). As |/ell increases unboundedly, by Theorem 1
the score &t > OC as well. So we have softmax (UT 8 F(X) >> softmax ve) - Cg and thus the cross-entropy loss drops to zero.
D THE DERIVATION OF EQ (5
For a vector let (+) denote its --thh entry. As we ignore the changes of the scores and the embeddings of non-topic words, their distributions maintain the same as the initial ones. In particular, the scores of the non-topic words are always zero. So, for 2. d,
var (x)) = var
exp(s1) (p) 2(x)
Σ
exp (5 (11))
+
Z(x)
wExi(()
exp($1e)
2 0²
mo²
Σ
Z(x)
(Z(x))
wEX
Since M((x):)) is assumed Lipschitz continuous in o(x). there exists L IR such that for yp 1, ..
I((((x)) - h(0(x2);))()) < cl(((1) ((2))lI,,
where ū(1), 6(x2) E Rd and ||_|| denote the 1 ddditancce by taking the sum of the absolute values of the entry differences each dimension. So we also have
s [1 ()((x))) 1 (((())):)) -
1
6(x2)l1.
According to the work of Bobkov & Houdrè 1996), for &o = 1,,,,... .d.
- var (xl))
mo²
(Z
which is
var(b((((x::)) (v)) mo² C2
(x(x)
Then, the Cauchy-Schwarz inequality implies, for = 1,2, 1, .d,
(M(((x):)) exp(st)
.1(4) - 5(x)
(M(B(X):)) (1) exp(si)
cov
Cov
7(x)
2(x)
<< var
(M(F()):0)) (0) exp(s:)
)
1/2
var (x(6)(x))
1/2
mo
(19)
7(x)
(2
16

---

Published as a conference paper at ICLR 2021
By the triangle inequality,
(vz -(x)) (M(O()):0) exp(s1)
7(x) / -((n--bxx) -(x)) v)
exp(81)
7(x)
Σ
( ( -
v. ((((()))))(0)) (F) exp(ss)
- ((())= (p) exp(se)
2(x) v1
x(x) V.
d
(M(O(x):))(()) exp(sr)
-(10)) -
- (x))(x)))() exp(st)
(//)) (()
I(x) / v.
Z(x)
d
(M(OW):)) exp(st)
-000) -x())())
cav
2(x)
mo=(exx(ss)VV
by Eq 19
(2(x))
m exp(st)
LVdo²
x(x)=
((x)
Z(x)
-(1- exp($1)
exp(si). LV
x(x
7(x)
7(x)
Lvdo2
xx(x)
Lvdo²
4
The last line is due to =(x) > Ewex)(() exp(sw) m 2 1.
17


Reviews:
[{"review_id": "qKj2jEu4Ddh", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper provides theoretical insight into the mechanisms by which a simplified attention model trained with gradient descent learns to allocate more mass to relevant words in the input.\n\nIn a limited toy setting, the authors derive a closed form relationship between word-score and word embedding norm in a simplified, one layer attention model. The theoretical findings are verified empirically both on the toy task and on a more realistic sentiment classification benchmark. Due to the extreme simplicity of the setting considered, as well as the number of assumptions made, it is unclear to me what to make of these results. In particular, it seems that the setting considered (fixed query attention over bag of word embeddings) is very different from real use cases of attention.\n\n**Pros**\n- The closed-form relationship between attention score and embedding norm during SGD training is novel as far as I know\n- The theoretical results are well justified in experiments: in particular the predicted \"SEN\" relationship seems to match the prediction.\n\n\n**Cons**\n- Large number of assumptions, the validity of which is unclear in practice: in particular\n   1. The assumption that the query vector is (1) a parameter and not a function of the inputs (as in self attention or cross-sentence attention) and (2) is fixed. I don't know of many \"real world\" attention networks that work this way, after all one of the main appeals of attention is its \"content-based\" nature\n   2. Assumption 1 that the score and embeddings of non-topic words don't change during training. First, this seems like something that could be proven from the earlier assumption that the topic words are updated more frequently. And second it is unclear if it holds for a real task (and a different model where eg. the attention layer attends to higher layers rather than just the embeddings)\n- Confusing notation makes the paper hard to follow (see remarks for examples)\n- Unclear takeaway: what does this paper tell us about attention as it is used in practice? \n\n\n**Remarks**\n-  5.1: \"The “negative effect” cannot last long because the topic word embedding moves along the gradient much faster than the non-topic words due to its high occurrence rate\": This is true in the toy example in the paper, but is this the case in practice? For instance in sentiment classifications there are many words to describe sentiment that are infrequent (cue Zipf's law). Moreover, in realistic settings there will be non-topic words which appear very frequently (stop words such as \"the\", \"a\" in English). \n- Lemma 1: while it is true that fixing q doesn't change the capacity of the model, it will definitely change its training dynamics (which is very much the theme of the paper as per the title). How important is it to fix q from this perspective?\n- A lot of the math would be easier to read if the dependence of some variables (\\hat v, Z,...) on a specific sentence was marked explicitly (eg. Z_S instead of Z)\n- The notation in Lemma 2 was extremely confusing to me, due to the sudden introduction of the bracket notation and the awkward spacing with both equations on the same line. I would recommend at least putting both on separate line, and also reorganizing so that the LHS of the second equation is only ds_i/dt (move the mean to the other side)\n- In 3. I think using \"\\mathbf R\" for the dictionary is unfortunate (too similar to \\mathbb R). Overall I found the separation between topic and non-topic words dictionaries confusing. Why not have a global Vocabulary V, a set of topic words T and refer to the remaining words as V\\T?\n- In 2. \"[Hahn and Brunner] show theoretically that the self-attention blocks are severely limited in their computational capabilities when the input sequence is long, which implies that self-attention may not provide the interpretability that one expects.\": can you clarify this sentence? Limitation in computational capabilities does not seem to entail limited interpretability in general (see linear models for instance).\n- Typo in citations in 2.: \"Hahn (Hahn, 2020) and Brunner (Brunner et al., 2019)\" -> \"Hahn (2020) and Brunner et al. (2019)\"\n- Typo in 3. \"The training set [...] are\" -> \"The training set [...] is\"\n\n---\n\n**Post Rebuttal**\n\nIn my review, the main concerns were (1) validity of assumptions, (2) confusing writing/notation and (3) unclear takeaway. The rebuttal appropriately addressed (1), although I am looking forward to the revision to see how this is discussed in the paper itself. I cannot really say anything about any improvements on the writing (2) without seeing the revision, but I am confident that the authors can address most of the issues pointed out by myself and other reviewers. Regarding (3), unclear takeaway, after reading the authors' response as well as the other reviews, my concerns are somewhat assuaged (partly because the assumptions were addressed better), although I am still unsure how or if the results in this paper could be expanded to realistic attention models.\n\nThere are additional issues I raised during the discussion (general lack of citations in particular), however this can be fixed fairly easily for the camera ready so I am willing to give the benefit of doubt and raise my score to 6 (borderline accept)\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "_eJr-JvhLa", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the dynamics of attention in a task of simplified topic modeling, over the course of training for a specific model, where the context vector is the sum over words in a sentence of their embedding weighted by the exponential of the dot-product their key embedding with a global query vector, normalized. Due to the simplification of the topic modeling problem (two null-intersect sets of words: topic vs. non-topic), they consider the embeddings of the non-topic words to be fixed over the course of training for their theoretical analysis. The applicability of the theoretical result is close to zero, and a somewhat known property (e.g. in word2vec, Mikolov et al. 2013). The experimental results include two parts. One on a tiny synthetic dataset that matches the simplified topic modeling problem and serves as illustration. The other is on SST2 and SST5 (movie comments and ratings, sentiment analysis), where the results are poor (obviously, as the model is simple), e.g. yielding 79.59% on SST while the SOTA is 97.4, and BERT base is at 91.2. The analysis is interesting, but does not lead to new insights.\n\nFor a simple analysis, the paper is at times hard to follow, and could benefit from more structure (presenting \"what\" before \"how\") and better notation (e.g. \\nu, v, $v$ all attached to (forms of the) the context vector).\n\nOverall, the contribution does not seem sufficient enough for inclusion at ICLR. The paper could be a good fit for a workshop on topic modeling or attention-based models.", "extra_scores": {}, "extra_text": {}}, {"review_id": "MgaGm9f4Xs", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "(Summary)\n\nThe paper investigates the dynamics of attention mechanism by configurating a controlled experiment on a simple topic classification task and training via gradient descent. Each random sentence in the training data is synthesized to include only one topic word among many. Then the authors try to find an intrinsic mechanism that triggers the attention model to discover the topic word and accelerates training via mutual promotion. They further experiment the evolution of models during optimization when no clear distinction between topic and non-topic words exist like in real data.\n\n\n(Originality and Contribution)\n\nThe paper proposes an artificial topic classification task and shows a positive score-and-embedding-norm relationship for the topic words to which the model must attend to. The authors also show that attention mechanism is highly helpful when the classifier has only limited capacity. They also demonstrate mutual promotion effect that leads a faster dropping of training loss than the fixed score and fixed embeddings. This discovery sounds to be original and relevant contribution to the field.\n\n\n(Strength and Weakness)\n-\tStrength: Design and run novel controlled experiment. Extensive analysis.\n-\tWeakness: Too much notational overloading. Writing quality.\n\n\n(Concerns, Questions, and Suggestions)\n\n1) It is unclear why $M >> N$ implies that a topic word appears more frequently in the sentences than a non-topic word. Section 3 describes that each sentence consists of only one topic word, then combining with $m$ non-topic words drawn uniformly at random. When the total number of topics $N$ is much smaller than the size of non-topic word dictionary $M$, what increases frequency of topic word?\n\n2) Overall simplifying some notations and avoiding notational overloading would greatly increase the readability of the paper.\n\n3) To reduce confusion, it would be great to change the iterator of summation for the partition function Z into $\\sum_{w’ \\in S_k}$ rather than using the same $w$.\n\n4) In Lemma 1, assume $q \\neq 0$ -> $q \\neq \\vec{0}$.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8AywpQXCBYw", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper aims to prove and illustrate that attention components are defined during training by gradients that mutually amplify the embedding and score associated with crucial features. In particular, a word embedding with a high magnitude increases the gradient following the attention score for the same word, while a high attention score increases the gradient directed at the word's embedding. In addition to a proof that treats behavior during training as a dynamical system under a large suite of assumptions, they test the analytic predictions on a synthetic dataset following the same suite of assumptions. They then test on a natural language data set and discuss where it diverges from the analytic and synthetic findings,  concluding that the difference is a result of competition between different words associated with a label.\n\nPros:\n1. We currently lack any substantial theory about attention modules and why they work. Although their model is simplistic, it could provide essential groundwork for analytic understanding of these popular systems. I would even consider it fairly realistic relative to a lot of the assumptions required for theoretical results in training dynamics research. Currently theory of attention is grounded in infinite-width networks, an assumption this paper does not make. \n2. The synthetic results appear to substantiate this theoretical result.\n3. They find an interesting result that, in more realistic settings, the learning dynamics follow particular patterns on the words that are paired together with more versus less predictive words. The framing of these effects in terms of competition between possible topic words is clearly inspired by considering which assumptions behind their proof have failed, which is evidence that the thinking behind their proof is potentially valuable.\n\nCons:\n1. It's not clear how dynamics like these would generalize to multilayer attention networks like BERT.\n2. The assumptions behind the theoretical and synthetic empirical results are simplistic:  The existence of a large vocabulary of \"non-topic\" words required to keep the variance of embedding negligible in out of focus words, the presence of only one topic word. There is also the very common assumption of Lipschltz continuity.\n3. The natural language experiments make a specific claim about the different dynamics for competing words of different topic purity, but only presents an example of two words as evidence. I want to see quantitative evidence of the pattern.\n4. The synthetic results would be strengthened by including multiple runs with different initializations so they can include confidence intervals.\n\n Questions:\n1. Does this mutual amplification effect have any ramifications for the debate over whether attention weights can be used as a proxy for saliency?\n2. In Lemma 1, there is a reference to the attention block's capacity which is difficult to decipher. What do you mean here by capacity?\n3. The assumption that word embeddings are sampled from a distribution with small variance seems likely to apply early in training, but not later. Have you checked the actual variance that would be associated with word embeddings late in training?\n4. What is actually meant by a word being \"paired\" with another word in the natural language experiments?\n5. Did I misunderstand something in interpreting gradients amplifying the embedding and score as directed towards v and k respectively in this simplified model?\n\n Minor:\n1. Notation is difficult to follow at times because several unrelated concepts use almost the same symbols: $s_i$ indicates score, but $S_i$ indicates a sentence; $\\tau$ indicates learning rate, but $\\mathcal{T}$ (which looks identical as a subscript) indicates a set of sentences.\n2. In discussing early alignment of attention to syntax, Clark et al. 2019 was concurrent with https://www.aclweb.org/anthology/P19-1580/", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper investigates the training dynamics of simple neural attention
mechanisms, in a controlled setting with clear (but rather strict) assumptions. Some reviewers
expressed caution about the applicability of the assumptions in practice,
but nevertheless there is agreement that the results deepen our understanding and enrich
our toolkit for reasoning about attention.
In support of this, in the discussion period, it was emphasized that the work uses different techniques than
most current work in this direction. I am therefore confident that the paper will be useful, and recommend acceptance.

I strongly encourage the authors to improve the clarity of the work and thorough
citation, as suggested by the reviewers.
Decision: Accept (Poster)