Neural Machine Translation with Key-Value Memory-Augmented Attention

1 Introduction

The past several years have witnessed promising progress in Neural Machine Translation (NMT) [Cho et al., 2014; Sutskever et al., 2014], in which attention model plays an increasingly important role [Bahdanau et al., 2015; Luong et al., 2015; Vaswani et al., 2017]. Conventional attention-based NMT encodes the source sentence as a sequence of vectors after bi-directional recurrent neural networks (RNN) [Schuster and Paliwal, 1997], and then generates a variable-length target sentence with another RNN and an attention mechanism. The attention mechanism plays a crucial role in NMT, as it shows which source word(s) the decoder should focus on in order to predict the next target word. However, there is no mechanism to effectively keep track of attention history in conventional attention-based NMT. This may let the decoder tend to ignore past attention information, and lead to the issues of repeating or dropping translations [Tu et al., 2016]. For example, conventional attention-based NMT may repeatedly translate some source words while mistakenly ignore other words.

A number of recent efforts have explored ways to alleviate the inadequate translation problem. For example, Tu et al. [2016] employ coverage vector as a lexical-level indicator to indicate whether a source word is translated or not. Meng et al. [2016] and Zheng et al. [2018] take the idea one step further, and directly model translated and untranslated source contents by operating on the attention context (i.e., the partial source content being translated) instead of on the attention probability (i.e., the chance of the corresponding source word is translated). Specifically, Meng et al. [2016] capture translation status with an interactive attention augmented with a NTM [Graves et al., 2014] memory. Zheng et al. [2018] separate the modeling of translated (PAST) and untranslated (FUTURE) source content from decoder states by introducing two additional decoder adaptive layers.

Meng et al. [2016] propose a generic framework of memory-augmented attention, which is independent from the specific architectures of the NMT models. However, the original mechanism takes only a single memory to both represent the source sentence and track attention history. Such overloaded usage of memory representations makes training the model difficult [Rockt¨aschel et al., 2017]. In contrast, Zheng et al. [2018] try to ease the difficulty of representation learning by separating the PAST and FUTURE functions from the decoder states. However, it is designed specifically for the precise architecture of NMT models.

In this work, we combine the advantages of both models by leveraging the generic memory-augmented attention framework, while easing the memory training by maintaining separate representations for the expected two functions. Partially inspired by [Miller et al., 2016], we split the memory into two parts: a dynamical key-memory along with the updatechain of the decoder state to keep track of attention history, and a fixed value-memory to store the representation of source sentence throughout the whole translation process. In each decoding step, we conduct multi-rounds of memory operations repeatedly layer by layer, which may let the decoder have a chance of re-attention by considering the "intermediate" attention results achieved in early stages. This structure allows the model to leverage possibly complex transformations and interactions between 1) the key-value memory pair in the same layer, as well as 2) the key (and value) memory across different layers.

Experimental results on Chinese⇒English translation task show that attention model augmented with a single-layer keyvalue memory improves both translation and attention performances over not only a standard attention model, but also over the existing NTM-augmented attention model [Meng et al., 2016]. Its multi-layer counterpart further improves model performances consistently. We also validate our model on bidirectional German⇔English translation tasks, which demonstrates the effectiveness and generalizability of our approach.

2 Background

Attention-based NMT Given a source sentence x = {x1, x2, · · · , xn} and a target sentence y={y1, y2, · · · , ym}, NMT models the translation probability word by word:

where f(·) is a non-linear function, and st is the hidden state of decoder RNN at time step t:

ct is a distinct source representation for time t, calculated as a weighted sum of the source annotations:

where hj is the encoder annotation of the source word xj, and the weight at,j is computed by

where et,j = vT a tanh(Wa˜ st-1 +Uahj) scores how much ˜st-1 attends to hj, where ˜st-1 = g(st-1, yt-1) is an intermediate state tailored for computing the attention score with the information of yt-1.

The training objective is to maximize the log-likelihood of the training instances (xs , ys ):

Memory-Augmented Attention Meng et al. [2016] propose to augment the attention model with a memory in the form of NTM, which aims at tracking the attention history during decoding process, as shown in Figure 1. At each decoding step, the NTM employs a network-based reader to read from the encoder annotations and output a distinct memory representation, which is subsequently used to update the decoder state. After predicting the target word, the updated decoder state is written back to the memory, which is controlled by a network-based writer. As seen, the interactive read-write operations can timely update the representation of source sentence along with the update-chain of the decoder state and let the decoder keep track of the attention history.

However, this mechanism takes only a single memory to both represent the source sentence and track attention history. Such overloaded usage of memory representations makes training the model difficult [Rockt¨aschel et al., 2017].

3 Model Description

Figure 2 shows the architecture of the proposed Key-Value Memory-augmented Attention model (KVMEMATT). It consists of three components: 1) an encoder (on the left), which encodes the entire source sentence and outputs its annotations as the initialization of the KEY-MEMORY and the VALUE-MEMORY; 2) the key-value memory-augmented attention model (in the middle), which generates the context representation of source sentence appropriate for predicting the next target word with iterative memory access operations conducted on the KEY-MEMORY and the VALUE-MEMORY; and 3) the decoder (on the right), which predicts the next target word step by step.

Specifically, the KEY-MEMORY and VALUE-MEMORY consists of n slots, which are initialized with the annotations of the source sentence {h1, h2, . . . , hn}. KVMEMATTbased NMT maintains these two memories throughout the whole decoding process, with the KEY-MEMORY keeping updated to track the attention history, and the VALUEMEMORY keeping fixed to store the representation of source sentence. For example, the j-th slot (vj) in VALUE-MEMORY stores the representation of the j-th source word (fixed after generated), and the j-th slot (kj) in KEY-MEMORY stores the attention (or translation) status (updated as translation goes) corresponding to the j-th source word. At step t, the decoder state st-1 first meets the previous prediction yt-1 to form a "query" state qt, which can be calculated as follows

where eyt-1 is the word-embedding of the previous word yt-1. The decoder uses the "query" state qt to address from the KEY-MEMORY looking for an accurate attention vector at, and reads from the VALUE-MEMORY with the guidance of at to generate the source context representation ct. After that the KEY-MEMORY is updated with qt and ct.

The memory access (i.e. address, read and update) in one decoding step can be conducted repeatedly, which may let the decoder have a chance of re-attention (with new information added) before making the final prediction. Suppose there are R rounds of memory access in each decoding step. The detailed operations from round r-1 to round r are as follows:

First, we use the "query" state qt to address fromK(r-1) (t)

to generate the "intermediate" attention vector ˜art

which is subsequently used as the guidance for reading from the value memory V to get the "intermediate" context representation ˜crt of source sentence

which works together with the "query" state qt are used to get the "intermediate" hidden state:

Finally, we use the "intermediate" hidden state to update K(r-1) (t) as recording the "intermediate" attention status, to finish a round of operations,

After the last round (R) of the operations, we use {˜sRt , ˜cRt , ˜aRt } as the resulting states to compute a final prediction via Eq. 1. Then the KEY-MEMORY K(R) (t) will be transited to the next decoding step t+1, being K(0) (t+1) . The details of Address, Read and Update operations will be described later in next section.

As seen, the KVMEMATT mechanism can update the KEY-MEMORY along with the update-chain of the decoder state to keep track of attention status and also maintains a fixed VALUE-MEMORY to store the representation of source sentence. At each decoding step, the KVMEMATT generates the context representation of source sentence via nontrivial transformations between the KEY-MEMORY and the VALUEMEMORY, and records the attention status via interactions between the two memories. This structure allows the model to leverage possibly complex transformations and interactions between two memories, and lets the decoder choose more appropriate source context for the word prediction at each step. Clearly, KVMEMATT can subsume the coverage models [Tu et al., 2016; Mi et al., 2016] and the interactive attention model [Meng et al., 2016] as special cases, while more generic and powerful, as empirically verified in the experiment section.

3.1 Memory Access Operations

In this section, we will detail the memory access operations from round r-1 to round r at decoding time step t.

Key-Memory Addressing Formally, K(r) (t) ∈ Rn×d is the KEY-MEMORY in round r at decoding time step t before the decoder RNN state update, where n is the number of memory slots and d is the dimension of vector in each slot. The addressed attention vector is given by

where ˜art ∈ Rn specifies the normalized weights assigned to the slots in K(r-1) (t) , with the j-th slot being k(r-1) (t,j) . We can use content-based addressing to determine ˜art as described in [Graves et al., 2014] or (quite similarly) use the softalignment as in Eq. 4. In this paper, for convenience we adopt the latter one. And the j-th cell of ˜art is

Value-Memory Reading Formally, V ∈ Rn×d is the VALUE-MEMORY, where n is the number of memory slots and d is the dimension of vector in each slot. Before the decoder state update at time t, the output of reading at round r is ˜crt given by

where ˜art ∈ Rn specifies the normalized weights assigned to the slots in V.

Key-Memory Updating Inspired by the attentive writing operation of neural turing machines [Graves et al., 2014], we define two types of operation for updating the KEYMEMORY: FORGET and ADD.

FORGET determines the content to be removed frommemory slots. More specifically, the vector Frt ∈ Rd specifies the values to be forgotten or removed on each dimension in memory slots, which is then assigned to each slot through normalized weights wr t . Formally, the memory ("intermediate") after FORGET operation is given by

• Frt = σ(WrF , ˜srt ) is parameterized with WrF ∈ Rd×d , and σ stands for the Sigmoid activation function, and Frt ∈ Rd ;
• wr t ∈ Rn specifies the normalized weights assigned to the slots in K(r) (t) , and wr t,i specifies the weight associated with the i-th slot. wr t is determined by

ADD decides how much current information should be written to the memory as the added content

where Art = σ(WrA, ˜srt ) is parameterized with WrA ∈ Rd×d , and Art ∈ Rd . Clearly, with FORGET and ADD operations, KVMEMATT potentially can modify and add to the KEYMEMORY more than just history of attention.

3.2 Training

EOS-attention Objective: The translation process is finished when the decoder generates eostrg (a special token that stands for the end of target sentence). Therefore, accurately generating eostrg is crucial for producing correct translation. Intuitively, accurate attention for the last word (i.e. eossrc) of source sentence will help the decoder accurately predict eostrg. When to predict eostrg (i.e. ym), the decoder should highly focus on eossrc (i.e. xn). That is to say the attention probability of am,n should be closed to 1.0. And when to generate other target words, such as y1, y2, · · · , ym-1, the decoder should not focus on eossrc too much. Therefore we define

where

NMTObjective: We plag Eq. 17 to Eq. 5, we have

The EOS-attention objective can assist the learning of KVMEMATT and guiding the parameter training, which will be verified in the experiment section.

4 Experiments

4.1 Setup

We carry out experiments on Chinese⇒English (Zh⇒En) and German⇔English (De⇔En) translation tasks. For Zh⇒En, the training data consist of 1.25M sentence pairs extracted from LDC corpora. We choose NIST 2002 (MT02) dataset as our valid set, and NIST 2003-2006 (MT03-06) datasets as our test sets. For De⇔En, we perform our experiments on the corpus provided by WMT17, which contains 5.6M sentence pairs. We use newstest2016 as the development set, and newstest2017 as the testset. We measure the translation quality with BLEU scores [Papineni et al., 2002]. 1

In training the neural networks, we limit the source and target vocabulary to the most frequent 30K words for both sides in Zh⇒En task, covering approximately 97.7% and 99.3% of two corpus respectively. For De⇔En, sentences are encoded using byte-pair encoding [Sennrich et al., 2016], which has a shared source-target vocabulary of about 36000 tokens. The parameters are updated by SGD and mini-batch (size 80) with learning rate controlled by AdaDelta [Zeiler, 2012] (ϵ = 1e-6 and ρ = 0.95). We limit the length of sentences in training to 80 words for Zh⇒En and 100 sub-words for De⇔En. The dimension of word embedding and hidden layer is 512, and the beam size in testing is 10. We apply dropout on the output layer to avoid over-fitting [Hinton et al., 2012], with dropout rate being 0.5. Hyper parameter λ in Eq. 19 is set to 1.0. The parameters of our KVMEMATT (i.e., encoder and decoder, except for those related to KVMEMATT) are initialized by the pre-trained baseline model.

4.2 Results on Chinese-English

We compare our KVMEMATT with two strong baselines: 1) RNNSEARCH, which is our in-house implementation of the attention-based NMT as described in Section 2; and 2) RNNSEARCH+MEMATT, which is our implementation of interactive attention [Meng et al., 2016] on top of our RNNSEARCH. Table 1 shows the translation performance.

Model Complexity KVMEMATT brings in little parameter increase. Compared with RNNSEARCH, KVMEMATTR1 and KVMEMATT-R2 only bring in 1.95% and 7.78% parameter increase. Additionally, introducing KVMEMATT slows down the training speed to a certain extent (18.39%∼39.56%). When running on a single GPU device Tesla P40, the speed of the RNNSEARCH model is 2773 target words per second, while the speed of the proposed models is 1676∼2263 target words per second.

Translation Quality KVMEMATT with one round memory access (KVMEMATT-R1) achieves significant improvements over RNNSEARCH by 1.92 BLEU and over RNNSEARCH+MEMATT by 0.65 BLEU averagely on four test sets. It indicates that our key-value memory-augmented attention mechanism can give more effective power for the Our end-to-end NMT systems attention via nontrivial transformations and interactions between the KEY-MEMORY and the VALUE-MEMORY. The two rounds counterpart (KVMEMATT-R2) can further improve the performance by 0.48 BLEU on average. It confirms our hypothesis that the decoder can benefit from the reattention process, which considers the "intermediate" attention result achieved in the early stage and then makes a more accurate decision. However, we find that adding more than two rounds of memory access operations into KVMEMATT does not lead to better translation performance (not shown in the table). One possible reason is that memory access with more rounds leads to more updating operations (i.e. attentive writing) on the KEY-MEMORY, which may be difficult to optimize within our current architecture design. We will leave it as future work.

1 For Zh⇒EEn task, we apply case-insensitive NIST BLEU mteval-v11b.pl. For De⇔En tasks, we tokenized the reference and evaluated the performance with case-sensitive multi-bleu.pl. The metrics are exactly the same as in previous work.

The contribution of adding EOS-attention Objective is to assist the learning of attention, and guiding the parameter training. It consistently improves translation performance over different variants of KVMEMATT. It gives a further improvement of 0.40 BLEU points over KVMEMATT-R2, which is 2.80 BLEU points better than RNNSEARCH.

Alignment Quality Intuitively, our KVMEMATT can enhance the attention and therefore improve the word alignment quality. To confirm our hypothesis, we carry out experiments of the alignment task on the evaluation dataset from [Liu and Sun, 2015], which contains 900 manually aligned ChineseEnglish sentence pairs. We use the alignment error rate (AER) [Och and Ney, 2003] as the evaluation metric for the alignment task. Table 2 lists the BLEU and AER scores. As expected, our KVMEMATT achieves better BLEU and AER scores (the lower the AER score, the better the alignment quality) than the strong baseline systems. Additionally, the results also indicate that the EOS-attention objective can assist the learning of attention-based NMT, since adding this objective yields better alignment performance. By visualizing the attention matrices, we found that the attention qualities are improved from the first round to the second round as expected, as shown in Figure 3.

Subjective Evaluation We did a subjective evaluation to investigate the benefit of incorporating KVMEMATT to NMT, especially on alleviating the issues of over- and undertranslations. Table 3 lists the translation adequacy of the RNNSEARCH baseline and our KVMEMATT-R2 on the randomly selected 100 sentences from test sets. From Table 3

we can see that, compared with the baseline system, our approach can decline the percentages of source words which are under-translated from 13.1% to 9.7% and which are overtranslated from 2.7% to 1.3%. The main reason is that our KVMEMATT can keep track of attention status and generate more appropriate source context for predicting the next target word at each decoding step.

4.3 Results on German-English

We also evaluate our model on the WMT17 benchmarks on bidirectional German⇔English translation tasks, as listed in Table 4. Our baseline achieves even higher BLEU scores to the state-of-the-art NMT systems of WMT17, which do not use additional synthetic data. 2 Our proposed model consistently outperforms two strong baselines (i.e., standard and memory-augmented attention models) on both De⇒En and En⇒De translation tasks. These results demonstrate that our model works well across different language pairs.

5 RelatedWork

Our work is inspired by the key-value memory networks [Miller et al., 2016] originally proposed for the question answering, and has been successfully applied to machine translation [Gu et al., 2017; Gehring et al., 2017; Vaswani et al., 2017; Tu et al., 2018]. In these works, both the key-memory and value-memory are fixed during translation. Different from these works, we update the KEY-MEMORY along with the update-chain of the decoder state via attentive writing operations (e.g. FORGET and ADD).

Our work is related to recent studies that focus on designing better attention models [Luong et al., 2015; Cohn et al., 2016; Feng et al., 2016; Tu et al., 2016; Mi et al., 2016; Zhang et al., 2017]. [Luong et al., 2015] proposed to use a global attention to attend to all source words and a local attention model to look at a subset of source words. [Cohn et al., 2016] extended the attention-based NMT to include structural biases from word-based alignment models. [Feng et al., 2016] added implicit distortion and fertility models to attention-based NMT. [Zhang et al., 2017] incorporated distortion knowledge into the attention-based NMT. [Tu et al., 2016; Mi et al., 2016] proposed coverage mechanisms to encourage the decoder to consider more untranslated source words during translation. These works are different from our KVMEMATT, since we use a rather generic key-value memory-augmented framework with memory access (i.e. address, read and update).

2 [Sennrich et al., 2017] obtains better BLEU scores than our model, since they use large scaled synthetic data (about 10M). It maybe unfair to compare our model to theirs directly.

Our work is also related to recent efforts on attaching a memory to neural networks [Graves et al., 2014] and exploiting memory [Tang et al., 2016; Wang et al., 2016; Feng et al., 2017; Meng et al., 2016; Wang et al., 2017] during translation. [Tang et al., 2016] exploited a phrase memory stored in symbolic form for NMT. [Wang et al., 2016] extended the NMT decoder by maintaining an external memory, which is operated by reading and writing operations. [Feng et al., 2017] proposed a neural-symbolic architecture, which exploits a memory to provide knowledge for infrequently used words. Our work differ at that we augment attention with a specially designed interactive key-value memory, which allows the model to leverage possibly complex transformations and interactions between the two memories via single- or multi-rounds of memory access in each decoding step.

6 Conclusion and FutureWork

We propose an effective KVMEMATT model for NMT, which maintains a timely updated key-memory to track attention history and a fixed value-memory to store the representation of source sentence during translation. Via nontrivial transformations and iterative interactions between the two memories, our KVMEMATT can focus on more appropriate source context for predicting the next target word at each decoding step. Additionally, to further enhance the attention, we propose a simple yet effective attention-oriented objective in a weakly supervised manner. Our empirical study on Chinese⇒English, German⇒English and English⇒German translation tasks shows that KVMEMATT can significantly improve the performance of NMT.

For future work, we will consider to explore more rounds of memory access with more powerful operations on keyvalue memories to further enhance the attention. Another interesting direction is to apply the proposed approach to Transformer [Vaswani et al., 2017], in which the attention model plays a more important role.