📂Machine Learning

Paper Review: Bahdanau Attention

Overview and Summary

• References, equation numbers, and notation follow the paper as closely as possible.

The paper “Neural Machine Translation by Jointly Learning to Align and Translate”¹ was presented at ICLR 2015, and its authors are Dzmitry Bahdanau, KyungHyun Cho, and Yoshua Bengio. It is regarded as the paper that first introduced the modern 🔒(26/08/07)attention idea used in the Transformer and elsewhere, and the form of attention proposed here is commonly called Bahdanau attention, after the first author, or additive attention.

The encoder-decoder architecture, i.e. the family of seq2seq models, which was the mainstream of neural machine translation at the time, compresses the entire source sentence into a single context vector of fixed dimension, from which the decoder then generates the translation. The authors conjecture that this fixed-length vector is a bottleneck for improving performance.

"In this paper, we conjecture that the use of a fixed-length vector is a bottleneck in improving the performance of this basic encoder–decoder architecture"

To solve this, the paper proposes a method that lets the decoder search on its own for the relevant parts of the source sentence each time it generates a word of the translation. Letting $h_{j}$ denote a vector summarizing the neighborhood of the $j$th word of the source sentence, the context vector $c_{i}$ for the $i$th target word is not fixed to a single vector but is freshly computed at every step as the following weighted sum.

$$ c_{i} = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j}, \qquad \alpha_{ij} = \frac{\exp (e_{ij})}{\sum_{k=1}^{T_{x}} \exp (e_{ik})}, \qquad e_{ij} = a (s_{i-1}, h_{j}) $$

The alignment model $a$ that computes the weights $\alpha_{ij}$ is a small neural network, trained jointly with the rest of the translation model. This is where the title comes from: the model jointly learns to align and to translate. They report that this extension alone significantly outperforms the existing encoder-decoder on English-to-French translation and achieves performance comparable to the then state-of-the-art phrase-based statistical translation system.

Modern Notation

Summarizing the paper in the notation of attention used today goes as follows. 🔒(26/08/07)Attention is the operation defined, for a query matrix $\mathbf{Q}$, a key matrix $\mathbf{K}$, a value matrix $\mathbf{V}$, and a score function $f$, as

$$ \operatorname{Attention} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) := \mathbf{V} \operatorname{Softmax} \left( f (\mathbf{Q}, \mathbf{K}) \right) $$

In Bahdanau attention, the query at the step where the decoder generates the $i$th word is just one vector, the decoder’s previous hidden state $s_{i-1}$. The keys and values are not distinguished from each other; both are the encoder’s hidden states (annotations).

$$ \mathbf{q}_{i} = s_{i-1}, \qquad \mathbf{K} = \mathbf{V} = \begin{bmatrix} h_{1} & h_{2} & \cdots & h_{T_{x}} \end{bmatrix} $$

The score function $f$ is a small neural network (= the alignment model $a$) that applies a linear transformation to the query and to each key, adds them, and takes $\operatorname{Tanh}$.

$$ \left[ f (\mathbf{q}_{i}, \mathbf{K}) \right]_{j} = a (s_{i-1}, h_{j}) = \mathbf{w}_{a}^{\mathsf{T}} \operatorname{Tanh} \left( W_{a} s_{i-1} + U_{a} h_{j} \right) $$

Here $W_{a}, U_{a}, \mathbf{w}_{a}$ are the learnable parameters of the alignment model, and the score vector $\mathbf{w}_{a}$ is a separate vector unrelated to the values $\mathbf{v}_{j}$. $\operatorname{Tanh}$ is the function that applies the hyperbolic tangent $\tanh$ to each component of its vector argument (the detailed definition of the alignment model $a$, following the paper’s notation, is given in §A.1.2).

The attention output at this step is exactly the context vector $c_{i}$.

$$ \mathbf{c}_{i} = \operatorname{Attention} (\mathbf{q}_{i}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \left( f (\mathbf{q}_{i}, \mathbf{K}) \right) = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j} $$

The correspondence between the paper’s notation and the modern notation is as follows.

Paper’s notation	Attention notation	Meaning
$s_{i-1}$	query $\mathbf{q}_{i}$	decoder’s previous hidden state
$h_{j}$	key $\mathbf{k}_{j}$ and value $\mathbf{v}_{j}$	encoder’s $j$th annotation
$e_{ij} = a (s_{i-1}, h_{j})$	score $\left[ f (\mathbf{q}_{i}, \mathbf{K}) \right]_{j}$	compatibility of query and key
$\alpha_{ij}$	attention weight	softmax of the scores
$c_{i}$	attention output	weighted average of the values, context vector

In summary, Bahdanau attention is an attention in which (i) keys and values are not distinguished and the encoder’s annotations play both roles, (ii) queries are produced one at a time, sequentially at each decoding step, and (iii) the score function is additive.

1 Introduction

Most neural translation models proposed at the time the paper appeared belong to the encoder-decoder family, i.e. seq2seq models. The authors point out that having to compress all the information of the source sentence into a single fixed-length vector is a potential problem of this architecture. To address this problem, they introduce an extension of the encoder-decoder that jointly learns to align and translate. Alignment here refers to the function that selects where in the source sentence the word to be translated should look, which attention computes with the key vectors. Translation refers to the function that connects words of the source language to words of the target language, which attention computes with the value vectors. Each time the proposed model generates a word of the translation, it (soft-)searches for a set of positions in the source sentence where the most relevant information is concentrated, and predicts the next target word based on the context vectors associated with these positions and the target words generated so far.

The most important distinguishing feature from the basic encoder-decoder is that it does not attempt to encode the whole input sentence into a single fixed-length vector. Instead, it encodes the input sentence into a sequence of vectors and adaptively chooses a subset of these vectors while decoding the translation. The experiments report that the proposed approach of jointly learning alignment and translation significantly outperforms the basic encoder-decoder, and that the improvement is more apparent with longer sentences but is observed regardless of sentence length.

2 Background: Neural Machine Translation

From a probabilistic point of view, translation is equivalent to finding the translation $\mathbf{y}$ that maximizes the conditional probability $p(\mathbf{y} | \mathbf{x})$ given a source sentence $\mathbf{x}$. That is, it amounts to solving the following optimization problem.

$$ \hat{\mathbf{y}} = \argmax\limits_{\mathbf{y}} p(\mathbf{y} | \mathbf{x}) $$

Neural network-based machine translation approaches typically consist of two parts, an encoder and a decoder. The encoder encodes the input source sentence $\mathbf{x}$ into a suitable feature vector $\mathbf{c}$, and the decoder decodes it into the translation $\mathbf{y}$.

$$ \mathbf{x} \overset{\text{encoding}}{\mapsto} \mathbf{c} \overset{\text{decoding}}{\mapsto} \mathbf{y} $$

2.1 RNN Encoder–Decoder

We briefly summarize the RNN encoder-decoder (the seq2seq model), the framework proposed by Cho et al.(2014)² and Sutskever et al.(2014)³. The encoder encodes the input sentence, a sequence of vectors $\mathbf{x} = (x_{1}, \cdots, x_{T_{x}})$, into a vector $c$. The most common approach is to use a recurrent neural network (RNN).

$$ h_{t} = f (x_{t}, h_{t-1}) \tag{1} $$

$$ c = q \left( \left\{ h_{1}, \cdots, h_{T_{x}} \right\} \right) $$

Here $h_{t} \in \mathbb{R}^{n}$ is the encoder’s hidden state at time $t$, $c$ is a vector generated from the sequence of hidden states, and $f$ and $q$ are some nonlinear functions. For instance, Sutskever et al. (2014) used an LSTM as $f$ and set $q \left( \left\{ h_{1}, \cdots, h_{T} \right\} \right) = h_{T}$. The decoder is trained to predict the next word $y_{t^{\prime}}$ at time $t^{\prime}$ given the context vector $c$ and all previously predicted words $\left\{ y_{1}, \cdots, y_{t^{\prime}-1} \right\}$. In other words, it decomposes the joint probability of the translation $\mathbf{y} = \left( y_{1}, \cdots, y_{T_{y}} \right)$ into a product of ordered conditional probabilities.

$$ p(\mathbf{y}) = \prod_{t=1}^{T} p \left( y_{t} \mid \left\{ y_{1}, \cdots, y_{t-1} \right\}, c \right) \tag{2} $$

With an RNN, each conditional probability is modeled by some nonlinear function $g$ as follows.

$$ p \left( y_{t} \mid \left\{ y_{1}, \cdots, y_{t-1} \right\}, c \right) = g (y_{t-1}, s_{t}, c) \tag{3} $$

Here $g$ is a nonlinear function that outputs the probability of $y_{t}$, and $s_{t}$ is the decoder’s hidden state. In the existing methods, the source sentence $\mathbf{x}$ was encoded into a context vector $c$ of fixed dimension, and this paper argues that this is the bottleneck preventing performance improvement.

3 Learning to Align and Translate

The proposed new architecture uses a bidirectional RNN as the encoder (§3.2) and, as the decoder, a structure that emulates searching through the source sentence while decoding the translation (§3.1).

3.1 Decoder: General Description

Figure 1

In the new architecture, each conditional probability in $(2)$ is defined as follows.

$$ p(y_{i} \mid y_{1}, \dots, y_{i-1}, \mathbf{x}) = g (y_{i-1}, s_{i}, c_{i}) \tag{4} $$

$s_{i}$ is the decoder’s hidden state at time $i$, computed as $s_{i} = f (s_{i-1}, y_{i-1}, c_{i})$. Note that, unlike the existing encoder-decoder which uses the fixed context vector $c$ when predicting every $y_{i}$, the probability uses a distinct context vector $c_{i}$ for each target word $y_{i}$.

The context vector $c_{i}$ depends on the sequence of annotations $(h_{1}, \cdots, h_{T_{x}})$ into which the encoder encodes the input sentence. ($h_{j}$ is the concatenation of the two hidden states of the BRNN encoder.) Each annotation $h_{j}$ is a vector that contains information about the whole input sequence with strong attention on the parts surrounding the $j$th word; how it is computed is covered in the next section. The context vector is computed as a linear combination of the annotations.

$$ c_{i} = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j} \tag{5} $$

The weight $\alpha_{ij}$ of each annotation $h_{j}$ is computed in softmax form.

$$ \alpha_{ij} = \frac{\exp (e_{ij})}{\sum_{k=1}^{T_{x}} \exp (e_{ik})} \tag{6} $$

Here $e_{ij} = a (s_{i-1}, h_{j})$ is an alignment model that scores how well the inputs around position $j$ and the output at position $i$ match. The score is based on the decoder’s hidden state $s_{i-1}$ just before emitting $y_{i}$ and the $j$th annotation $h_{j}$ of the input sentence. The alignment model $a$ is parametrized as a feedforward neural network and trained jointly with all the other components of the system. Unlike traditional machine translation, the paper explains, making the alignment a neural network allows the alignment model and the whole translation model to be trained simultaneously.

The weighted sum of all the annotations can be understood, they say, as an expected annotation over possible alignments. If we view $\alpha_{ij}$ as the probability that the target word $y_{i}$ is aligned to (translated from) the source word $x_{j}$, then $c_{i}$ is the expectation of the annotations under that probability. And in the very next passage of the paper, the word attention is mentioned.

"Intuitively, this implements a mechanism of attention in the decoder. The decoder decides parts of the source sentence to pay attention to. By letting the decoder have an attention mechanism, we relieve the encoder from the burden of having to encode all information in the source sentence into a fixed-length vector."

3.2 Encoder: Bidirectional RNN for Annotating Sequences

The ordinary RNN of $(1)$ reads the input in order, from the first word $x_{1}$ to the last $x_{T_{x}}$, but in the scheme of §3.1 the output $y_{i}$ at time $i$ refers to the entire sequence of input vectors. This is possible because a bidirectional RNN, BRNN is used as the encoder architecture. (The paper writes it as BiRNN.) In a BRNN, the output $y_{i}$ refers to all the input vectors $\left\{ x_{1}, \dots, x_{T_{x}} \right\}$ regardless of the time step.

A BRNN consists of a forward RNN $\overrightarrow{f}$ and a backward RNN $\overleftarrow{f}$. The forward RNN receives the data in order (from $x_{1}$ to $x_{T_{x}}$) and computes the forward hidden states $\left\{ \overrightarrow{h}_{1}, \cdots, \overrightarrow{h}_{T_{x}} \right\}$, while the backward RNN receives the data in the reverse order and computes the backward hidden states $\left\{ \overleftarrow{h}_{1}, \cdots, \overleftarrow{h}_{T_{x}} \right\}$. The annotation of each word $x_{j}$ is obtained by concatenating the two.

$$ h_{j} = \begin{bmatrix} \overrightarrow{h}_{j} \\ \overleftarrow{h}_{j} \end{bmatrix} $$

In this way, the annotation $h_{j}$ contains the summaries of both the preceding words and the following words. Since the hidden state of an RNN tends to represent nearby data better, $h_{j}$ is said to naturally focus on the words around $x_{j}$.

A Model Architecture

Omitting the content on the experimental results, we describe in detail the equations of the architecture of RNNsearch, the model used in the experiments.

A.2.1 Encoder

The model takes as input the source sentence, a sequence of one-hot vectors, and outputs the translation, likewise a sequence of one-hot vectors.

$$ \mathbf{x} = (x_{1}, \dots, x_{T_{x}}), \quad x_{i} \in \mathbb{R}^{K_{x}}, \qquad \mathbf{y} = (y_{1}, \dots, y_{T_{y}}), \quad y_{i} \in \mathbb{R}^{K_{y}} $$

$K_{x}$ and $K_{y}$ are the vocabulary sizes of the source language and the target language, respectively, and $T_{x}$ and $T_{y}$ are the lengths of the source sentence and the translation. The forward hidden states of the bidirectional recurrent neural network encoder are computed as follows.

$$ \overrightarrow{h}_{i} = \begin{cases} (1 - \overrightarrow{z}_{i}) \odot \overrightarrow{h}_{i-1} + \overrightarrow{z}_{i} \odot \overrightarrow{\tilde{h}}_{i} & \text{if } i > 0 \\ 0 & \text{if } i = 0 \end{cases} $$

Here $\odot$ is the Hadamard product, meaning componentwise multiplication of vectors. Beware that the paper writes it as $\circ$, which usually means composition of functions. Each term is as follows.

$$ \begin{align*} \overrightarrow{\tilde{h}}_{i} &= \tanh \left( \overrightarrow{W} \overline{E} x_{i} + \overrightarrow{U} \left[ \overrightarrow{r}_{i} \odot \overrightarrow{h}_{i-1} \right] \right) \\ \overrightarrow{z}_{i} &= \sigma \left( \overrightarrow{W}_{z} \overline{E} x_{i} + \overrightarrow{U}_{z} \overrightarrow{h}_{i-1} \right) \\ \overrightarrow{r}_{i} &= \sigma \left( \overrightarrow{W}_{r} \overline{E} x_{i} + \overrightarrow{U}_{r} \overrightarrow{h}_{i-1} \right) \end{align*} $$

$\overline{E} \in \mathbb{R}^{m \times K_{x}}$ is the word embedding matrix, $\overrightarrow{W}, \overrightarrow{W}_{z}, \overrightarrow{W}_{r} \in \mathbb{R}^{n \times m}$ and $\overrightarrow{U}, \overrightarrow{U}_{z}, \overrightarrow{U}_{r} \in \mathbb{R}^{n \times n}$ are weight matrices, $m$ is the embedding dimension, and $n$ is the number of hidden units. Since $x_{i}$ is a 1-of-$K$ vector, $\overline{E} x_{i}$ is the operation of picking out the column of the embedding matrix corresponding to that word.

The backward hidden states $\overleftarrow{h}_{1}, \cdots, \overleftarrow{h}_{T_{x}}$ are computed in the same way, reading the input in the reverse order. The embedding matrix $\overline{E}$ is shared between the forward and backward RNNs, but the weight matrices are not shared. Concatenating the two hidden states yields the annotation described in §3.2.

$$ h_{i} = \begin{bmatrix} \overrightarrow{h}_{i} \\ \overleftarrow{h}_{i} \end{bmatrix} \in \mathbb{R}^{2n} $$

A.2.2 Decoder

The decoder’s hidden state $s_{i}$ is computed from the annotations produced by the encoder as follows.

$$ s_{i} = (1 - z_{i}) \odot s_{i-1} + z_{i} \odot \tilde{s}_{i} $$

$$ \begin{align*} \tilde{s}_{i} &= \tanh \left( W E y_{i-1} + U \left[ r_{i} \odot s_{i-1} \right] + C c_{i} \right) \\ z_{i} &= \sigma \left( W_{z} E y_{i-1} + U_{z} s_{i-1} + C_{z} c_{i} \right) \\ r_{i} &= \sigma \left( W_{r} E y_{i-1} + U_{r} s_{i-1} + C_{r} c_{i} \right) \end{align*} $$

$E$ is the word embedding matrix of the target language, and $W, W_{z}, W_{r} \in \mathbb{R}^{n \times m}$, $U, U_{z}, U_{r} \in \mathbb{R}^{n \times n}$, $C, C_{z}, C_{r} \in \mathbb{R}^{n \times 2n}$ are weight matrices.

Now let us see, starting from the first step, how the translation is emitted from the decoder. The initial hidden state is computed as follows from the last hidden state of the backward RNN, i.e. the summary obtained by reading the source sentence from the end, with $W_{s} \in \mathbb{R}^{n \times n}$.

$$ s_{0} = \tanh \left( W_{s} \overleftarrow{h}_{1} \right) $$

The context vector $c_{1}$ is computed as below.

$$ c_{1} = \sum_{j=1}^{T_{x}} \alpha_{1j} h_{j}, \qquad \alpha_{1j} = \frac{\exp (e_{1j})}{\sum_{k=1}^{T_{x}} \exp (e_{1k})}, \qquad e_{1j} = v_{a}^{\mathsf{T}} \tanh \left( W_{a} s_{0} + U_{a} h_{j} \right) $$

Suppose the initial value $y_{0}$ is given as some token signaling the start of the sentence. The decoder’s hidden state $s_{1}$ is computed as below.

$$ \begin{align*} z_{1} &= \sigma \left( W_{z} E y_{0} + U_{z} s_{0} + C_{z} c_{1} \right) \\ r_{1} &= \sigma \left( W_{r} E y_{0} + U_{r} s_{0} + C_{r} c_{1} \right) \\ \tilde{s}_{1} &= \tanh \left( W E y_{0} + U \left[ r_{1} \odot s_{0} \right] + C c_{1} \right) \\ s_{1} &= (1 - z_{1}) \odot s_{0} + z_{1} \odot \tilde{s}_{1} \end{align*} $$

Given the decoder state $s_{1}$, the context vector $c_{1}$, and the previously generated word $y_{0}$, the probability vector (probability distribution) of the next word is defined as follows.

$$ p (\cdot \mid s_{1}, y_{0}, c_{1}) = \operatorname{softmax} \left( W_{o} t_{1} \right) \in \mathbb{R}^{K_{y}} $$

$$ t_{1} = \begin{bmatrix} \max \left\{ \tilde{t}_{1, 1}, \tilde{t}_{1, 2} \right\} & \cdots & \max \left\{ \tilde{t}_{1, 2l-1}, \tilde{t}_{1, 2l} \right\} \end{bmatrix}^{\mathsf{T}} \in \mathbb{R}^{l} $$ $$ \tilde{t}_{1} = U_{o} s_{0} + V_{o} E y_{0} + C_{o} c_{1} $$

Here $\tilde{t}_{i,k}$ is the $k$th component of the vector $\tilde{t}_{i}$. $W_{o} \in \mathbb{R}^{K_{y} \times l}$, $U_{o} \in \mathbb{R}^{2l \times n}$, $V_{o} \in \mathbb{R}^{2l \times m}$, $C_{o} \in \mathbb{R}^{2l \times 2n}$ are weight matrices. Pairing up the components of the vector $\tilde{t}_{i} \in \mathbb{R}^{2l}$ and keeping only the larger of each pair is a maxout layer with a single hidden layer, and multiplying its output $t_{i} \in \mathbb{R}^{l}$ by $W_{o}$ produces scores of dimension $K_{y}$, the vocabulary size, which become probabilities after normalizing with the softmax. Selecting the word at the index with the largest value here gives exactly $y_{1}$.

$$ y_{1} = \operatorname*{argmax}_{1 \le w \le K_{y}} \left[ p (\cdot \mid s_{1}, y_{0}, c_{1}) \right]_{w} $$

Repeating this process afterwards yields the translation $\left\{ y_{1}, \dots, y_{T_{y}} \right\}$.

A.1.2 Alignment Model

The alignment model must be designed with the consideration that it needs to be computed $T_{x} \times T_{y}$ times for each sentence pair of lengths $T_{x}$ and $T_{y}$. To reduce the computation, the paper uses a single-layer perceptron with one hidden layer.

$$ a (s_{i-1}, h_{j}) = v_{a}^{\mathsf{T}} \tanh \left( W_{a} s_{i-1} + U_{a} h_{j} \right) $$

Here $W_{a} \in \mathbb{R}^{n \times n}$, $U_{a} \in \mathbb{R}^{n \times 2n}$, $v_{a} \in \mathbb{R}^{n}$ are weight matrices. Since $U_{a} h_{j}$ does not depend on $i$, precomputing it once per sentence minimizes the computational cost. This form, in which a linear transformation of the query $s_{i-1}$ and a linear transformation of the key $h_{j}$ are added inside $\tanh$, is the reason, mentioned earlier, why Bahdanau attention is also called additive attention.