Efficient Attention using a Fixed-Size Memory Representation

The standard content-based attention mechanism typically used in sequence-to-sequence models is computationally expensive as it requires the comparison of large encoder and decoder states at each time step. In this work, we propose an alternative attention mechanism based on a fixed size memory representation that is more efficient. Our technique predicts a compact set of K attention contexts during encoding and lets the decoder compute an efficient lookup that does not need to consult the memory. We show that our approach performs on-par with the standard attention mechanism while yielding inference speedups of 20% for real-world translation tasks and more for tasks with longer sequences. By visualizing attention scores we demonstrate that our models learn distinct, meaningful alignments.

The most popular approaches are based on an encoder-decoder architecture consisting of two recurrent neural networks (RNNs) and an attention mechanism that aligns target to source tokens (Bahdanau et al., 2014;Luong et al., 2015). The typical attention mechanism used in these architectures computes a new attention context at each decoding * Equal Contribution. Author order alphabetical.
step based on the current state of the decoder. Intuitively, this corresponds to looking at the source sequence after the output of every single target token.
Inspired by how humans process sentences, we believe it may be unnecessary to look back at the entire original source sequence at each step. 1 We thus propose an alternative attention mechanism (section 3) that leads to smaller computational time complexity. Our method predicts K attention context vectors while reading the source, and learns to use a weighted average of these vectors at each step of decoding. Thus, we avoid looking back at the source sequence once it has been encoded. We show (section 4) that this speeds up inference while performing on-par with the standard mechanism on both toy and real-world WMT translation datasets. We also show that our mechanism leads to larger speedups as sequences get longer. Finally, by visualizing the attention scores (section 5), we verify that the proposed technique learns meaningful alignments, and that different attention context vectors specialize on different parts of the source.

Sequence-to-Sequence Model with Attention
Our models are based on an encoder-decoder architecture with attention mechanism (Bahdanau et al., 2014;Luong et al., 2015). An encoder function takes as input a sequence of source tokens x=(x 1 ,...,x m ) and produces a sequence of states s=(s 1 ,...,s m ) .The decoder is an RNN that predicts the probability of a target sequence y =(y 1 ,...,y T |s). The probability of each target token y i ∈ {1,...,|V |} is predicted based on the recurrent state in the decoder RNN, h i , the previous words, y <i , and a context vector c i . The context vector c i , also referred to as the attention vector, is calculated as a weighted average of the source states.
Here, f att (h i , s) is an attention function that calculates an unnormalized alignment score between the encoder state s j and the decoder state h i . Variants of f att used in Bahdanau et al. (2014) and Luong et al. (2015) are: where W a and v a are model parameters learned to predict alignment. Let |S| and |T | denote the lengths of the source and target sequences respectively and D denoate the state size of the encoder and decoder RNN. Such content-based attention mechanisms result in inference times of O(D 2 |S||T |) 2 , as each context vector depends on the current decoder state h i and all encoder states, and requires an O(D 2 ) matrix multiplication. The decoder outputs a distribution over a vocabulary of fixed-size |V |: The model is trained end-to-end by minimizing the negative log likelihood of the target words using stochastic gradient descent.

Memory-Based Attention Model
Our proposed model is shown in Figure 1. During encoding, we compute an attention matrix C ∈R K×D , where K is the number of attention vectors and a hyperparameter of our method, and D is the dimensionality of the top-most encoder state. This matrix is computed by predicting a score vector α t ∈ R K at each encoding time step t. C is then a linear combination of the encoder states, weighted by α t : where W α is a parameter matrix in R K×D . The computational time complexity for this operation is O(KD|S|). One can think of C as compact fixed-length memory that the decoder will perform attention over. In contrast, standard approaches use a variable-length set of encoder states for attention. At each decoding step, we similarly predict K scores β ∈R K . The final attention context c is a linear combination of the rows in C weighted by the scores. Intuitively, each decoder step predicts how important each of the K attention vectors is.
Here, h is the current state of the decoder, and W β is a learned parameter matrix. Note that we do not access the encoder states at each decoder step. We simply take a linear combination of the attention matrix C pre-computed during encoding -a much cheaper operation that is independent of the length of the source sequence. The time complexity of this computation is O(KD|T |) as multiplication with the K attention matrices needs to happen at each decoding step.
Summing O(KD|S|) from encoding and O(KD|T |) from decoding, we have a total linear computational complexity of O(KD(|S| + |T |). As D is typically very large, 512 or 1024 units in most applications, we expect our model to be faster than the standard attention mechanism running in O(D 2 |S||T |). For long sequences (as in summarization, where -S-is large), we also expect our model to be faster than the cheaper dot-based attention mechanism, which needs O(D|S||T |) computation time and requires encoder and decoder states sizes to match.
We also experimented with using a sigmoid function instead of the softmax to score the encoder and decoder attention scores, resulting in 4 possible combinations. We call this choice the scoring function. A softmax scoring function calculates normalized scores, while the sigmoid scoring function results in unnormalized scores that can be understood as gates.

Model Interpretations
Our memory-based attention model can be understood intuitively in two ways. We can interpret it as "predicting" the set of attention contexts produced by a standard attention mechanism during encoding. To see this, assume we set K ≈ |T |. In this case, we predict all |T | attention contexts during the encoding stage and learn to choose the right one during decoding. This is cheaper than computing contexts one-by-one based on the decoder and encoder content. In fact, we could enforce this objective by first training Figure 1: Memory Attention model architecture. K attention vectors are predicted during encoding, and a linear combination is chosen during decoding. In our example, K =3. a regular attention model and adding a regularization term to force the memory matrix C to be close to the T ×D vectors computed by the standard attention. We leave it to future work to explore such an objective.
Alternatively, we can interpret our mechanism as first predicting a compact K × D memory matrix, a representation of the source sequence, and then performing location-based attention on the memory by picking which row of the matrix to attend to. Standard location-based attention mechanism, by contrast, predicts a location in the source sequence to focus on (Luong et al., 2015;Xu et al., 2015).

Position Encodings (PE)
In the above formulation, the predictions of attention contexts are symmetric. That is, C i is not forced to be different from C j =i . While we would hope for the model to learn to generate distinct attention contexts, we now present an extension that pushes the model into this direction. We add position encodings to the score matrix that forces the first few context vector C 1 ,C 2 ,... to focus on the beginning of the sequence and the last few vectors ...,C K−1 ,C K to focus on the end (thereby encouraging in-between vectors to focus on the middle).
Explicitly, we multiply the score vector α with position encodings l s ∈R K : To obtain l s we first calculate a constant matrix L where we define each element as adapting a formula from (Sukhbaatar et al., 2015).
Here, k ∈ {1, 2, ..., K} is the context vector index and S is the maximum sequence length across all source sequences. The manifold is shown graphically in Figure 2. We can see that earlier encoder states are upweighted in the first context vectors, and later states are upweighted in later vectors. The symmetry of the manifold and its stationary point having value 0.5 both follow from Eq. 10. The elements of the matrix that fall beyond the sequence lengths are then masked out and the remaining elements are renormalized across the timestep dimension. This results in the jagged array of position encodings {l ks }.   Table 1 shows the BLEU scores of our model on different sequence lengths while varying K. This is a study of the trade-off between computational time and representational power. A large K allows us to compute complex source representations, while a K of 1 limits the source representation to a single vector. We can see that performance consistently increases with K up to a point that depends on the data length, with longer sequences requiring more complex representations.

Results
The results with and without position encodings are almost identical on the toy data. Our technique learns to fit the data as well as the standard attention mechanism despite having less representational power. Both beat the non-attention baseline by a significant margin.  That we are able to represent the source sequence with a fixed size matrix with fewer than |S| rows suggests that traditional attention mechanisms may be representing the source with redundancies and wasting computational resources. This makes intuitive sense for the toy task, which should require a relatively simple representation.
The last column shows that our technique significantly speeds up the inference process. The gap in inference speed increases as sequences become longer. We measured inference time on the full validation set of 1,000 examples, not including data loading or model construction times. Figure 3a shows the learning curves for sequence length 200. We see that K =1 is unable to fit the data distribution, while K ∈{32,64} fits the data almost as quickly as the attention-based model. Figure 3b shows the effect of varying the encoder and decoder scoring functions between softmax and sigmoid. All combinations manage to fit the data, but some converge faster than others. In section 5 we show that distinct alignments are learned by different function combinations.

Machine Translation
Next, we explore if the memory-based attention mechanism is able to fit complex real-world datasets. For this purpose we use 4 large machine translation datasets of WMT'17 5 on the following language pairs: English-Czech (en-cs, 52M examples), English-German (en-de, 5.9M examples), English-Finish (en-fi, 2.6M examples), and English-Turkish (en-tr, 207,373 examples). We used the newly available pre-5 statmt.org/wmt17/translation-task.html processed datasets for the WMT'17 task. 6 Note that our scores may not be directly comparable to other work that performs their own data pre-processing. We learn shared vocabularies of 16,000 subword units using the BPE algorithm (Sennrich et al., 2016). We use newstest2015 as a validation set, and report BLEU on newstest2016.

Training Setup
We use a similar setup to the Toy Copy task, but use 512 RNN and embedding units, train using 8 distributed workers with 1 GPU each, and train for at most 1M steps. We save checkpoints every 30 minutes during training, and choose the best based on the validation BLEU score.

Results
Table 2 compares our approach with and without position encodings, and with varying values for hyperparameter K, to baseline models with regular attention mechanism. Learning curves are shown in Figure 4. We see that our memory attention model with sufficiently high K performs on-par with, or slightly better, than the attention-based baseline model despite its simpler nature. Across the board, models with K = 64 performed better than corresponding models with K = 32, suggesting that using a larger number of attention vectors can capture a richer understanding of source sequences. Position encodings also seem to consistently improve model performance. Table 3 shows that our model results in faster decoding time even on a complex dataset with a large    vocabulary of 16k. We measured decoding time over the full validation set, not including time used for model setup and data loading, averaged across 10 runs. The average sequence length for examples in this data was 35, and we expect more significant speedups for tasks with longer sequences, as suggested by our experiments on toy data. Note that in our NMT ex-amples/experiments, K ≈T , but we obtain computational savings from the fact that K D. We may be able to set K T , as in toy copying, and still get very good performance in other tasks. For instance, in summarization the source is complex but the representation of the source required to perform the task is "simple" (i.e. all that is needed to generate the abstract). Figure 5 shows the effect of using sigmoid and softmax function in the encoders and decoders. We found that softmax/softmax consistently performs badly, while all other combinations perform about equally well. We report results for the best combination only (as chosen on the validation set), but we found this choice to only make a minor difference.

Visualizing Attention
A useful property of the standard attention mechanism is that it produces meaningful alignment between source and target sequences. Often, the attention mechanism learns to progressively focus on the next source token as it decodes the target. These visualizations can be an important tool in debugging and evaluating seq2seq models and are often used for unknown token replacement.
This raises the question of whether or not our proposed memory attention mechanism also learns to generate meaningful alignments. Due to limiting the number of attention contexts to a number that is generally less than the sequence length, it is not immediately obvious what each context would learn to focus on. Our hope was that the model would learn to focus on multiple alignments at the same time, within the same attention vector. For example, if the source sequence is of length 40 and we have K =10 attention contexts, we would hope that C 1 roughly focuses on tokens 1 to 4, C 2 on tokens 5 to 8, and so on. Figures 6 and 7 show that this is indeed the case. To generate this visualization we multiply the attention scores α and β from the encoder and decoder. Figure  8 shows a sample translation task visualization. Figure 6 suggests that our model learns distinct ways to use its memory depending on the encoder and decoder functions. Interestingly, using softmax normalization results in attention maps typical of those derived from using standard attention, i.e. a relatively linear mapping between source and target tokens. Meanwhile, using sigmoid gating results in what seems to be a distributed representation of the source sequences across encoder time steps, with multiple contiguous attention contexts being accessed at each decoding step.

Related Work
Our contributions build on previous work in making seq2seq models more computationally efficient. Luong et al. (2015) introduce various attention mechanisms that are computationally simpler and perform as well or better than the original one presented in Bahdanau et al. (2014). However, these typically still require O(D 2 ) computation complexity, or lack the flexibility to look at the full source sequence. Efficient location-based attention (Xu et al., 2015) has also been explored in the image recognition domain. Wu et al. (2016) presents several enhancements to the standard seq2seq architecture that allow more efficient computation on GPUs, such as only attending on the bottom layer. Kalchbrenner et al. (2016) propose a linear time architecture based on stacked convolutional neural networks. Gehring et al. (2016) also propose the use of convolutional encoders to speed up NMT. de Brébisson and Vincent (2016) propose a linear attention mechanism based on covariance matrices applied to information retrieval. Raffel et al. (2017) enable online linear time attention calculation by enforcing that the alignment between input and output sequence elements be monotonic. Previously, monotonic attention was proposed for morphological inflection generation by Aharoni and Goldberg (2016).

Conclusion
In this work, we propose a novel memory-based attention mechanism that results in a linear computational time of O(KD(|S|+|T |)) during decoding in seq2seq models. Through a series of experiments, we demonstrate that our technique leads to consistent inference speedups as sequences get longer, and can fit complex data distributions such as those found in Neural Machine Translation. We show that our attention mechanism learns meaningful alignments despite being constrained to a fixed representation after encoding. We encourage future work that explores the optimal values of K for various language tasks and examines whether or not it is possible to predict K based on the task at hand. We also encourage evaluating our models on other tasks that must deal with long sequences but have compact representations, such as summarization and question-answering, and further exploration of their effect on memory and training speed.