Approximating How Single Head Attention Learns

We use the dot-attention variant from Luong et al. (2015). The model maps from an input sequence $\{x_{l}\}$ with length $L$ to an output sequence $\{y_{t}\}$ with length $T$. We first use LSTM encoders to embed $\{x_{l}\}\subset\mathcal{I}$ and $\{y_{t}\}\subset\mathcal{O}$ respectively, where $\mathcal{I}$ and $\mathcal{O}$ are input and output vocab space, and obtain encoder and decoder hidden states $\{h_{l}\}$ and $\{s_{t}\}$. Then we calculate the attention logits $a_{t,l}$ by applying a learnable mapping from $h_{l}$ and $s_{t}$, and use softmax to obtain the attention weights $\alpha_{t,l}$:

Next we sum the encoder hidden states $\{h_{t}\}$ weighted by the attention to obtain the “context vector" $c_{t}$, concatenate it with the decoder $s_{t}$, and obtain the output vocab probabilities $p_{t}$ by applying a learnable neural network $N$ with one hidden layer and softmax activation at the output:

We train the model by minimizing the sum of negative log likelihood of all the output words $y_{t}$:

We define the lexical probe $\beta_{t,l}$ as:

which means “the probability assigned to the correct word $y_{t}$, if the network attends only to the input encoder state $h_{l}$". If we assume that $h_{l}$ only contains information about $x_{l}$, $\beta$ closely reflects KTIW, since $\beta$ can be interpreted as “the probability that $x_{l}$ is translated to the output $y_{t}$".

Heuristically, to minimize the loss, the attention weights $\alpha$ should be attracted to positions with larger $\beta_{t,l}$.¹¹1This statement is heuristical rather than rigorous. See Appendix A.1 for a counterexample. Hence, we expect the learning of the attention to be driven by KTIW (Figure 1 left). We then discuss how KTIW is learned.

To approximate how KTIW is learned early on in training, we build a proxy model by making a few simplifying assumptions. First, since attention weights are uniform early on in training, we 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 to its corresponding hidden state. Therefore, similar to Sun and Lu (2020), we replace $h_{l}$ with an input word embedding $e_{x_{l}}\in\mathbb{R}^{d}$, where $e$ represents the word embedding matrix and $d$ is the embedding dimension. Third, to simplify analysis, we assume $N$ only contains one linear layer $W\in\mathbb{R}^{|\mathcal{O}|\times d}$ before softmax activation and ignore the decoder state $s_{t}$. Putting these assumptions together, we now define a new proxy model that produces output vocab probability $p_{t}$:

On a high level, this proxy averages the embeddings of the input “bag of words", and produces a 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 statistics for this proxy.

The probe $\beta^{\mathrm{px}}$ can be similarly defined as:

We provide more intuitions on how this proxy learns in Section 4.

Binary classification can be reduced to “machine translation", where $T=1$ and $|\mathcal{O}|=2$. We drop the subscript $t=1$ when discussing classification.

We use the standard architecture from Wiegreffe and Pinter (2019a). After obtaining the encoder hidden states $\{h_{t}\}$, we calculate the attention logits $a_{l}$ by applying a feed-forward neural network with one hidden layer and take the softmax of $a$ to obtain the attention weights $\alpha$:

where $Q$ and $v$ are learnable.

We sum the hidden states $\{h_{l}\}$ weighted by the attention, feed it to a final linear layer and apply the sigmoid activation function ($\sigma$) to obtain the probability for the positive class

Similar to the machine translation model (Section 2.1), we define the “lexical probe":

where $y\in\{0,1\}$ is the label and $2y-1\in\{-1,1\}$ controls the sign.

On a high level, Sun and Lu (2020) focuses on binary classification and provides almost the exact same arguments as ours. Specifically, their polarity score “$s_{l}$" equals $\frac{\beta_{l}}{1-\beta_{l}}$ in our context, and they provide a more subtle analysis of how the attention mechanism is learned in binary classification.

We start by describing how to evaluate the agreement between quantities of interest, such as $\alpha$ and $\beta$. For any input-output sentence pair $(x^{m},y^{m})$, for each output index $t$, $\alpha^{m}_{t},\beta^{m}_{t},\beta^{\mathrm{px},m}_{t}\in\mathbb{R}^{L^{m}}$ all associate each input position $l$ 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,v\in\mathbb{R}^{L}$, we formally define the agreement of $v$ with $u$ as:

which means “whether the highest-valued position (dimension) in $u$ is in the top 5% highest-valued positions in $v$". We average the $\mathcal{A}$ values across all output words on the validation set to measure the agreement between two model properties. We also report Kendall’s $\tau$ rank correlation coefficient in Appendix 3 for completeness.

We denote its random baseline as $\hat{\mathcal{A}}$. $\hat{\mathcal{A}}$ is close to but not exactly $5\%$ because of integer rounding.

Even when the attention mechanism has not been learned, KTIW can still be learned. We train the same model architecture with the attention weights frozen to be uniform, and denote its lexical probe as $\beta^{\mathrm{uf}}$. Across all tasks, $\mathcal{A}(\alpha,\beta^{\mathrm{uf}})$ and $\mathcal{A}(\beta^{\mathrm{uf}},\beta^{\mathrm{px}})$ ³³3Empirically, $\beta^{\mathrm{px}}$ converges to the unigram weight of a bag-of-words logistic regression model, and hence $\beta^{\mathrm{px}}$ does capture an interpretable notion of “keywords”. (Appendix A.10.) significantly outperform the random baseline $\hat{\mathcal{A}}$, and the contextualized agreement $\xi(\alpha,\beta^{\mathrm{uf}})$ 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.

We consider a simple task of copying from the input to the output, and each input is 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.⁴⁴4We provide more intuitions on this in Section 4 As a 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 3 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 as training samples, the model successfully learns (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 of two disjoint sets of words (Figure 3 right), and Appendix A.3 illustrates the intuition.

For binary classification, it follows from the model definition that attention mechanism cannot be learned if KTIW cannot be learned, since

and the model needs to attend to positions with higher $\beta$, in order to predict correctly and minimize the loss. For completeness, we include results where we freeze $\beta$ and find that the learning of the attention fails in Appendix A.6.

We continue from the end of Section 2.3. For each input word $i$ and output word $o$, we are interested in understanding the probability that $i$ assigns to $o$, defined as:

This quantity is directly tied to $\beta^{\mathrm{px}}$, since $\beta^{\mathrm{px}}_{t,l}=\theta^{\mathrm{px}}_{x_{l},y_{t}}$. Using super-script ^m to index sentence pairs in the dataset, the total loss $\mathcal{L}$ is:

Suppose each $e_{i}$ or $W_{o}$ is independently initialized from a normal distribution $\mathcal{N}(0,I_{d}/d)$ and we minimize $\mathcal{L}$ over $W$ and $e$ using gradient flow, then the value of $e$ and $W$ are uniquely defined for each continuous time step $\tau$. By some straightforward but tedious calculations (details in Appendix A.2), the derivative of $\theta_{i,o}$ when the training starts is:

where $\overset{p}{\to}$ means convergence in probability and $C^{\mathrm{px}}_{i,o}$ is defined as

Equation 15 tells us that $\beta^{\mathrm{px}}_{t,l}=\theta^{\mathrm{px}}_{x_{l},y_{t}}$ is likely to be larger if $C_{x_{l},y_{t}}$ is large. The definition of $C$ seems hard to interpret from Equation 16, but in the next subsection we will find that this quantity naturally corresponds to the “count table" used in the classical IBM 1 alignment learning algorithm.

The classical alignment algorithm aims to learn which input word is responsible for each output word (e.g. knowing that $y_{2}$ “movie" aligns to $x_{2}$ “Film" in Figure 1 upper left), from a set of input-output sentence pairs. IBM Model 1 Brown et al. (1993) starts with a 2-dimensional count table $C^{\mathrm{IBM}}$ indexed by $i\in\mathcal{I}$ and $o\in\mathcal{O}$, denoting input and output vocabs. Whenever vocab $i$ and $o$ co-occurs in an input-output pair, we add $\frac{1}{L}$ to the $C^{\mathrm{IBM}}_{i,o}$ entry (step 1 and 2 in Figure 1 right). After updating $C^{\mathrm{IBM}}$ for the entire dataset, $C^{\mathrm{IBM}}$ is exactly the same as $C^{\mathrm{px}}$ defined in Equation 16. We drop the super-script of $C$ to keep the notation uncluttered.

Given $C$, the classical model estimates a probability distribution of “what output word $o$ does the input word $i$ translate to" (Figure 1 right step 3) as

In a pair of sequences ($\{x_{l}\},\{y_{t}\}$), the probability $\beta^{\mathrm{IBM}}$ that $x_{l}$ is translated to the output $y_{t}$ is:

and the alignment probability $\alpha^{\mathrm{IBM}}$ that “$x_{l}$ is responsible for outputting $y_{t}$ versus other $x_{l^{\prime}}$" is

which monotonically increases with respect to $\beta^{\mathrm{IBM}}_{t,l}$. See Figure 1 right step 5.

Figure 1 (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 a distribution of vocab permutations (Section 3.3) in Figure 4, and how word polarity is learned in binary classification (Section 2.4) in Figure 5.

We use gradient based method Ebrahimi et al. (2018) to approximate the influence $\Delta_{l}$ for each input word $x_{l}$. The column $\mathcal{A}(\Delta,\beta^{\mathrm{uf}})$ reports the agreement between $\Delta$ and $\beta^{\mathrm{uf}}$, 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.

We saw in Section 3.3 that learning to copy sequences under a 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 $c_{t}$ obtained by each head. Suppose there are $K$ heads each indexed by $k$, similar to Section 2.1:

and the context vector and final probability $p_{t}$ defined as:

where $W^{(k)}$ are different learn-able parameters.

We call $W^{(k)}_{init}$ a 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 3 scenarios: (1) all head initializations are bad, (2) only one initialization is good, and (3) initializations are sampled independently at random.

Figure 6 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). In 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.

We use a toy classification task to show that early on in training, expectantly, $\beta^{\mathrm{uf}}$ is larger near positions that contain the keyword. However, unintuitively, $\beta^{\mathrm{uf}}_{L}$ ($\beta$ at the last position in 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,\dots,40\}$ uniformly at random; a sequence is positive if and only if the keyword “1” appears in the sequence. We restrict “1" to appear only once in each positive sequence, and use rejection sampling to balance positive and negative examples. Let $l^{*}$ be the position where $x_{l^{*}}=1$.

For the positive sequences, we examine the log-odd ratio $\gamma_{l}$ before the sigmoid activation in Equation 8, since $\beta$ will be all close to 1 and comparing $\gamma$ would be more informative: $\gamma_{l}:=\log\frac{\beta^{\mathrm{uf}}_{l}}{1-\beta^{\mathrm{uf}}_{l}}.$

We measure four quantities: 1) $\gamma_{l^{*}}$, the log-odd ratio if the model only attends to the key word position, 2) $\gamma_{l^{*}+1}$, one position after the key word position, 3) $\bar{\gamma}:=\frac{\sum_{l=1}^{L}\gamma_{l}}{L}$, if attention weights are uniform, and 4) $\gamma_{L}$ if the model attends to the last hidden state. If the $\gamma_{l}$ only contains information about word $x_{l}$, we should expect:

However, if we accept the conventional wisdom that hidden states contain information about nearby words Khandelwal et al. (2018), we should expect:

To verify these hypotheses, we plot how $\gamma_{l^{*}},\gamma_{l^{*}+1},\bar{\gamma}$, and $\gamma_{L}$ evolve as training proceeds in Figure 7. Hypothesis 2 is indeed true when training starts; however, we find the following to be true asymptotically:

which is wildly different from Hypothesis 2. If we train under uniform attention weights for too long, the information about keywords can freely flow to other non-local hidden states.

In Section 2.1 we assumed that attention weights $\alpha$ behave like free variables that can assign arbitrarily high probabilities to positions with larger $\beta$. However, $\alpha$ is produced by a model, and sometimes learning the correct $\alpha$ can be challenging.

Let $\pi$ be a random permutation of integers from 1 to 40, and we want to learn the function $f$ that permutes the input with $\pi$:

Input $x$ are randomly sampled from a vocab of size 60 as in Section 3.3. Even though $\beta^{\mathrm{uf}}$ 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 a built-in inductive bias to learn monotonic attention.