Bahdanau attention: Neural Machine Translation by Jointly Learning to Align and Translate

A translator that rereads every source word, with learned weights, before it writes each word of the output.

In 2014 a neural translator read the source sentence into one vector of 1000 numbers and wrote the translation from that vector alone, and its quality fell on long sentences. This paper keeps one vector per source word, and for every output word the decoder takes a new weighted average of them. Those weights are what is now called attention.

Nobody tells the model which English word matches which French word. The weights are trained only by the translation loss, and they come out lining up the two sentences the way a bilingual reader would, including where French reverses the word order.

The model in this paper, which the authors call RNNsearch, translates English into French with two recurrent neural networks, like the sequence-to-sequence translator published the same month. It changes how the second network reads from the first, and that change is the additive attention now named after Bahdanau. The sections below build the model in the order its parts depend on each other:

1. the fixed-vector encoder–decoder it extends, and why its single vector was suspected to be a bottleneck;
2. a context vector computed separately for every output word, as a weighted average of per-word vectors;
3. the alignment model that scores each source word, and the softmax that turns scores into weights;
4. why those weights are soft, and what that does for training;
5. the bidirectional encoder that produces the per-word vectors;
6. one full decoder step at the paper's tensor sizes;
7. the results: BLEU scores, quality against sentence length, and the alignments the model learned.

Why one fixed-length vector limits translation

Neural machine translation in 2014 meant an encoder–decoder. An encoder RNN reads the English sentence one word at a time and leaves behind a single vector $c$. A decoder RNN then writes the French one word at a time, conditioned on $c$. In the baseline this paper compares against, the RNN Encoder–Decoder of Cho et al. (2014a), which the paper calls RNNencdec, the encoder has 1000 hidden units. A 5-word sentence and a 50-word sentence both have to be summarized in those 1000 numbers, and every French word, the first and the fiftieth, is written from the same $c$.

The authors conjecture that this fixed size is a bottleneck. Their evidence is a companion study (Cho et al., 2014b), which found that a basic encoder–decoder translates short sentences well and gets rapidly worse as sentences get longer. This paper's experiments test the conjecture; Figure 5, near the end of this page, plots BLEU against sentence length for both models.

Figure 1 draws the two designs. In the fixed-vector mode the box for $c$ keeps its width as you drag the sentence length $L$ up, so each word's share of it falls as $1/L$. In the soft-search mode the encoder keeps one vector per word, so the strip grows with the sentence and the decoder reads a different part of it for each output word.

The $1/L$ share is a sketch: a trained network does not split its 1000 numbers evenly between words. It does have to fit a 40-word sentence into the same space as a 4-word one, while in RNNsearch each of the 40 words keeps its own 2000-number vector and the memory grows with the sentence.

The encoder–decoder baseline

The paper writes the baseline in three equations. The encoder is an RNN, a function $f$ that takes the current word and the previous hidden state and returns the new state. After the last word, a function $q$ turns the hidden states into the context vector:

In (1), $x_t$ is the $t$-th source word, $h_t$ the encoder state after reading it, and $T_x$ the source length. Sutskever et al. used an LSTM for $f$ and let $q$ return the last state, $c = h_{T_x}$.

The decoder defines the probability of a French sentence one word at a time:

$s_t$ is the decoder's hidden state and $g$ is a small network that outputs a probability for every word in the 30,000-word French vocabulary. The $c$ in (3) has no subscript: the same vector conditions the first output word and the last. Training maximizes the log of (2) over the parallel corpus, so every French word's gradient reaches the encoder through that one vector.

How Bahdanau attention works: one context vector per output word

RNNsearch keeps the decoder's form and gives the context a subscript:

The next-word distribution for position $i$ and the state update both read $c_i$, a context computed for that position. To build it, the encoder no longer reduces the sentence to one vector. It keeps one vector per source word, $h_1, \dots, h_{T_x}$, which the paper calls annotations. Annotation $h_j$ describes word $j$ in the context of the whole sentence (the encoder that makes it comes two sections down; for now take $h_j$ as a 2000-number vector per word). The context for output word $i$ is a weighted sum of all of them:

The weights for one output word are nonnegative and add up to 1 over the source positions, so $c_i$ is a weighted average of the annotations. All the weight on one annotation makes $c_i$ equal to it; equal weights make it the plain average of the sentence.

A concrete case: the source is the deal was signed in August and the output is l'accord a été signé en août. When the decoder is about to write accord, a trained model should put most of the weight on deal. With the weights from Figure 2 below,

a 2000-number vector that is mostly the annotation of deal. One word later, writing a, the largest weight (0.717) sits on was, and $c_3$ is mostly $h_3$.

The paper reads (5) as an expected annotation: treat $\alpha_{ij}$ as the probability that French word $i$ is aligned to, or translated from, English word $j$, and $c_i$ is the expected value of the annotation under that distribution. The network computes this expectation exactly, as a deterministic weighted average under its own learned weights. Nothing is sampled, and there is no posterior over some true hidden alignment.

Alignment scores and the softmax

The weights come from one score per source word, which the paper calls the energy $e_{ij}$, passed through a softmax over the source positions:

The exponential makes every weight positive, the division makes the weights sum to 1 across the $T_x$ source words, and the result changes smoothly with every energy, so the gradient of the loss reaches whatever computed the energies. Adding the same constant to all energies leaves the weights unchanged; only the differences between source words matter.

The scoring function $a$ is the alignment model. It reads the decoder state $s_{i-1}$, the state before word $i$ is written, and the annotation $h_j$. Scoring against $s_{i-1}$ means the model chooses where to read before it commits to word $i$; the context it builds then feeds both the choice of $y_i$ and the next state $s_i$, as in (4). The paper makes $a$ a network with one hidden layer (App. A.1.2):

with $W_a \in \mathbb{R}^{n' \times n}$, $U_a \in \mathbb{R}^{n' \times 2n}$ and $v_a \in \mathbb{R}^{n'}$, and $n = n' = 1000$ in every experiment. That is 1,000,000 parameters in $W_a$, 2,000,000 in $U_a$ and 1,000 in $v_a$, about 3.0 million for the scorer. The decoder state is projected by $W_a$, the annotation by $U_a$, the two 1000-number projections are added and passed through tanh, and $v_a$ reduces the result to one number. Because the two inputs are added inside the tanh, this form is called additive attention. The same annotation $h_j$ is the vector that gets scored here and the vector that gets averaged in (5); there are no separate key and value projections.

$U_a h_j$ does not depend on the output position $i$, so the paper computes it once per source word and reuses it for every French word. Per output word, the only new matrix product is $W_a s_{i-1}$, and each source word then costs one tanh over 1000 numbers and one 1000-term dot product.

Figure 2 runs one sentence through this. The energies are illustrative numbers, and every weight printed below them is their exact softmax. Writing accord, the energies are (+0.2, +3.1, −0.5, −0.3, −1.6, −1.1). deal's energy is 2.9 above the runner-up, the (+0.2), so deal gets about $e^{2.9} \approx 18$ times the's weight; after the softmax divides by the sum over all six words, its weight is 0.877. Press Play and follow the weights along the source as the French is written; at été they split between was (0.53) and signed (0.35).

The weight row is recomputed from scratch for every French word, from a decoder state that already reflects the words written so far. No rule makes the weights move left to right; they follow the source in order here because English and French put these words in the same order.

Why the alignment is soft

Statistical translation before this paper treated word alignment as a discrete hidden variable: French word $i$ comes from exactly one English word $j$, and the system searches or sums over the choices. RNNsearch could have done the same and used only the annotation of the highest-scoring source word. The paper states that its alignment is not a latent variable: the alignment model computes a soft alignment, the gradient of the cost passes through it, and the alignment model trains jointly with the rest of the translator.

A hard pick is a step function of the energies. Change the scores slightly and the same word stays picked, so the derivative of the loss with respect to the energies is zero almost everywhere and backpropagation gives the alignment model nothing to learn from. Later hard-attention models, such as Xu et al.'s image captioner (2015), sample a position and train with a REINFORCE-style estimator, which gets a usable but noisy gradient from the samples. With the softmax, $\alpha$ is a smooth function of the energies and $c_i$ a smooth function of $\alpha$, so ordinary backprop delivers the gradient on every energy.

The paper also gives a translation reason. In its Figure 3(d), the man becomes l'homme. A hard alignment maps the to l' and man to homme, but the French article depends on the noun after it: the can become le, la, les or l'. With soft weights the decoder reads both words while writing l'; the published weights for that row are 0.32 on the and 0.50 on man. Soft weights also handle phrases of different lengths in the two languages without aligning leftover words to a dummy NULL word.

Figure 3 takes that l' row, renormalized over four nearby source words, and multiplies its energies by a scale $\beta$. At $\beta = 0$ every energy is 0 and each weight is 0.25. The real model starts there: the paper initializes $v_a$ to zero and draws every element of $W_a$ and $U_a$ from a Gaussian with mean 0 and variance $0.001^2$ (App. B.1), so every energy $v_a^\top \tanh(\cdot)$ is 0 and every context is the plain average of the sentence. At $\beta = 1$ you get the published weights. As $\beta$ grows toward 10 the softmax approaches a hard pick of man, and the curve shows the size of $\partial \alpha / \partial e$ falling from 0.45 at $\beta = 1$ to 0.02 at $\beta = 10$.

At the uniform starting point the softmax passes gradient well (0.43, close to the curve's peak), and $v_a$ gets a nonzero gradient through $\tanh(W_a s_{i-1} + U_a h_j)$, so the energies start to separate after the first updates. From there, gradient descent sharpens a row where a sharper weight lowers the translation loss and can leave it spread where that lowers the loss, as in the l' row.

Annotations from a bidirectional RNN

An annotation should describe its word together with the words on both sides of it, and a left-to-right RNN state at position $j$ has read only words 1 to $j$. So the encoder runs two RNNs: a forward one from $x_1$ to $x_{T_x}$, producing $\overrightarrow{h}_1, \dots, \overrightarrow{h}_{T_x}$, and a backward one from $x_{T_x}$ back to $x_1$, producing $\overleftarrow{h}_1, \dots, \overleftarrow{h}_{T_x}$. The annotation stacks the two states for the same position:

Each RNN has 1000 units, so $h_j$ has 2000 numbers, which is why $U_a$ above is 1000 by 2000. The two RNNs share the source word-embedding matrix and have separate recurrent weights (App. A.2.1). The paper argues that since RNNs represent recent inputs best, both halves are dominated by words near $j$, so $h_j$ summarizes the whole sentence with its focus on the neighborhood of word $j$. Because of that focus, a weight on $h_j$ can be read as a weight on word $j$ in the alignment plots later.

The decoder needs a state before it has read any context. The paper sets $s_0 = \tanh(W_s \overleftarrow{h}_1)$. The backward RNN reaches position 1 last, after reading every word from the end of the sentence back to the start, so $\overleftarrow{h}_1$ is the encoder state that has seen the whole source. Fix every $c_i$ to the last forward state $\overrightarrow{h}_{T_x}$ instead of computing it from (5), and RNNsearch turns back into RNNencdec (App. A.2.2).

In Figure 4, pick a word to see which inputs fed its forward state (itself and the words to its left) and its backward state (itself and the words to its right). The fade with distance is an illustration of the paper's recency argument, not a measurement. The decoder-start button shows $s_0$.

One decoder step, with the paper's sizes

Every RNN in the paper uses the gated hidden unit of Cho et al. (2014a), now called the GRU. The paper notes that LSTM units would also work, as in Sutskever et al. For the decoder (App. A.2.2), with $e(y_{i-1})$ the 620-number embedding of the previous French word:

# decoder GRU update for output word i (App. A.2.2, biases omitted)
# s_prev = s_{i-1} (1000), emb = e(y_{i-1}) (620), c = c_i (2000)
z  = sigmoid(Wz @ emb + Uz @ s_prev + Cz @ c)   # update gate, 1000
r  = sigmoid(Wr @ emb + Ur @ s_prev + Cr @ c)   # reset gate, 1000
s_ = tanh(W @ emb + U @ (r * s_prev) + C @ c)   # proposed state, 1000
s  = (1 - z) * s_prev + z * s_                  # new state s_i

The update gate z sets, per unit, how much of the old state to replace; where z is near 0 the unit copies $s_{i-1}$ forward unchanged, which gives gradients a path across many steps. The reset gate r sets how much of the old state enters the proposal. The context enters all three lines through 1000-by-2000 matrices C, Cz, Cr.

The output layer is a "deep output" with one maxout layer: a 1000-number vector is computed from the previous state, the previous word and the context, each adjacent pair of entries is reduced to its maximum, leaving $l = 500$ numbers, and a softmax over the 30,000-word vocabulary gives the next-word probabilities. One full step for output word $i$, in order:

# one decoder step: write output word i (biases omitted)
# H: Tx x 2000 annotations; UaH = H @ Ua.T, Tx x 1000, computed once per sentence
# s_prev = s_{i-1} (1000); y_prev = id of the previous French word
q     = Wa @ s_prev                       # 1000: the decoder's part of every score
e     = tanh(UaH + q) @ va                # Tx energies e_ij
alpha = softmax(e)                        # Tx weights, sum to 1          eq. (6)
c     = alpha @ H                         # 2000: context c_i             eq. (5)
emb   = E[:, y_prev]                      # 620: previous word embedding
t     = Uo @ s_prev + Vo @ emb + Co @ c   # 1000 = 2l
t     = t.reshape(500, 2).max(axis=1)     # maxout over adjacent pairs, l = 500
p     = softmax(Wo @ t)                   # 30000 next-word probabilities
s     = gru(s_prev, emb, c)               # new state s_i (listing above)
# loss for this word: -log p[y_i]; gradient reaches Wa, Ua, va, both encoders

Eq. (4) in the main text writes the output as $g(y_{i-1}, s_i, c_i)$, with the new state $s_i$. Appendix A.2.2 computes it from $s_{i-1}$, and so does the authors' GroundHog code; the listing follows the appendix and the code.

Training minimizes $-\log p(y_i)$ summed over the French words of each sentence pair. Every matrix above, including $W_a$, $U_a$, $v_a$ and both encoder RNNs, gets its gradient from that one loss; there is no separate alignment objective. The optimizer is minibatch SGD with Adadelta ($\epsilon = 10^{-6}$, $\rho = 0.95$), 80 sentence pairs per minibatch, and the gradient's L2 norm clipped to at most 1. Each step costs time proportional to the longest sentence in the minibatch, so before every 20th update the authors pulled 1600 sentence pairs, sorted them by length and split them into 20 minibatches. Each regular model trained for about five days (RNNsearch-50: 111 hours, 2.2 epochs, on a Quadro K-6000 GPU); RNNsearch-50★, the same model kept training until the development score stopped improving, ran for 252 hours (5.0 epochs). Translations are decoded with beam search, which keeps several of the most probable partial translations at each step instead of committing to the single most probable next word, and a translation ends when the model emits the end-of-sentence token.

The alignment model runs $T_x \times T_y$ times per sentence pair, once for each source word and output word. At 30 words each way that is 900 evaluations of a tanh over 1000 numbers plus a 1000-term dot product, about 0.9 million multiply-adds, next to 60 million for computing $U_a h_j$ for the 30 source words. The paper calls the quadratic term acceptable because most translation sentences are 15 to 40 words long, and flags it as a possible limit for other tasks. At 1000 tokens each way it would be a million evaluations per sentence pair. Transformer attention inherited the same quadratic growth; FlashAttention makes that computation faster and its memory linear without removing the quadratic compute.

Results: BLEU on English-to-French and on long sentences

The task is WMT'14 English-to-French. The parallel corpora total 850M words, reduced to 348M with the data selection method of Axelrod et al. (2011), and no monolingual data is used. Each language keeps its 30,000 most frequent words and maps every other word to [UNK]. The development set is news-test-2012 and news-test-2013 combined; the test set is news-test-2014, 3003 sentences. Both architectures are size-matched at 1000 hidden units (RNNsearch has a forward and a backward encoder RNN of 1000 units each), and each is trained twice, on sentences of up to 30 words (-30) and up to 50 words (-50).

Quality is measured in BLEU, which counts the 1- to 4-word sequences of the output that also appear in a reference translation, combines the four precisions, applies a penalty for output shorter than the reference, and reports the result on a 0 to 100 scale. Table 1 of the paper:

The right column scores only sentences with no unknown word in the source or the reference, and there the models were not allowed to output [UNK]. Every RNNsearch beats the RNNencdec trained on the same lengths: RNNsearch-50 scores 26.75 against 17.82, nearly nine BLEU. RNNsearch-30, trained only on sentences up to 30 words, also beats RNNencdec-50, 21.50 to 17.82, which the paper cites as further confirmation. Both numbers in the starred row, 28.45 and 36.15, come from the longer-trained run.

On sentences without unknown words, RNNsearch-50★ scores 36.15 against the phrase-based system Moses's 35.63. On the full test set Moses wins, 33.30 to 28.45, and Moses also used a separate 418M-word monolingual corpus. The paper claims only that on known-word sentences a single neural network performs as well as the phrase-based system. The difference between the two columns points at the 30,000-word vocabulary, and the paper's conclusion names unknown and rare words as one of the challenges left.

Table 1 averages over all lengths. The paper's Figure 2 splits BLEU by source sentence length and is the direct test of the bottleneck conjecture. Figure 5 re-plots its four curves, traced by eye from the published plot, so read the values to about half a BLEU point.

At 10 words RNNsearch-50 scores about 24.6 and RNNencdec-50 about 21.0. At 30 words the fixed-vector model is down to about 19.2 while RNNsearch-50 is at 27.5. At 50 words RNNencdec-50 is near 9 and RNNsearch-50 near 26, a gap of about 17 BLEU. RNNsearch-50 shows no deterioration even past 50 words, which is the paper's main evidence that the fixed vector was the bottleneck. At the very shortest lengths on the plot, RNNencdec-50 is slightly ahead.

RNNsearch-30 tracks RNNsearch-50 up to about 20 words, then declines, to about 23 at 30 words and about 10 at 50, level with RNNencdec-50. It never saw a training sentence longer than 30 words, and soft-search did not make it translate long sentences well. It stays above RNNencdec-30, which falls to about 3, at every length.

Section 5.2.2 of the paper shows the difference on one test sentence: An admitting privilege is the right of a doctor to admit a patient to a hospital or a medical centre to carry out a diagnosis or a procedure, based on his status as a health care worker at a hospital. RNNencdec-50 translates it correctly up to a medical centre and then drifts, ending with en fonction de son état de santé, "based on his state of health". RNNsearch-50 keeps the meaning to the end: … selon son statut de travailleur des soins de santé à l'hôpital.

What the learned alignments look like

Each row of weights is a distribution over source words, so one sentence's weights form a matrix that can be drawn as an image: a row per French word, a column per English word, brightness for weight. The paper's Figure 3 shows four such matrices from RNNsearch-50. Figure 6 reproduces two of them, with each cell's weight read from the published grayscale images; the rows of that reading sum to between 0.96 and 1.00.

Panel (a) is mostly a diagonal, because English and French put most words in the same order. The exception is European Economic Area, which French writes zone économique européenne, noun first. The model writes zone reading Area (0.71), then économique reading Economic (0.76), then européenne reading European (0.97): it jumps over two words and walks back one word at a time. The row for la puts 0.61 on the and 0.27 on Area; one plausible reading, which the paper does not offer, is that the article has to agree in gender with zone, the word that translates Area.

Panel (d) holds the l' row from Figure 3 above (0.32 on the, 0.50 on man), and the next row, homme, puts 0.98 on man.

From RNNsearch to Luong attention and the Transformer

Luong, Pham and Manning (2015) moved the score to the current decoder state $s_i$ and tried multiplicative scores such as $s_i^\top W h_j$, which replace the tanh layer with one bilinear product. The Transformer (Vaswani et al., 2017) removed the RNNs. Each position computes a query, a key and a value with three separate projections, the score is a dot product of query and key divided by the square root of their width, and attention alone moves information between positions. The Transformer keeps (5) and (6) from this paper: a softmax over scores, then a weighted average, recomputed for every output position and trained by the same loss as the rest of the network.

The paper itself uses the word attention three times, in three consecutive sentences of Section 3.1, where it describes the decoder deciding which parts of the source to attend to. Its own names for the mechanism are the alignment model and (soft-)search, and the model is RNNsearch.