Nearest Neighbor Machine Translation

1 INTRODUCTION . Non-parametric methods have recently been successfully applied to tasks such as language modeling ( Khandelwal et al. , 2020 ) and question answering ( Guu et al. , 2020 ; Lewis et al. , 2020 ) . They allow models that are ( 1 ) expressive , because they can use an arbitrary amount of data at test time ; ( 2 ) adaptable , because predictions can be controlled by changing the datastore , and ( 3 ) interpretable , because the data used to make the prediction can be directly inspected . We introduce kNN-MT , a simple non-parametric method for machine translation ( MT ) using nearest neighbor retrieval . kNNMT can be added to any pre-trained neural translation model without further training , and significantly improves performance for in-domain , out-of-domain , and multi-lingual evaluations . More specifically , kNN-MT interpolates the target-token softmax distribution from a neural MT model with a multinomial generated using nearest neighbor search over examples cached in a data store . The cache is over translation contexts ( i.e . the complete source and prefix of the target ) , and is indexed by hidden states computed from the base MT model . We hypothesize that contexts which are close in representation space are more likely to be followed by the same target word . We show this is not only true for the original training data , thereby improving base model performance , but across a range of different bi-text corpora , allowing for simple and effective model adaptation . Our work builds upon recent results showing the effectiveness of nearest neighbor methods in unconditional language models ( Khandelwal et al. , 2020 ) . We generalize to conditional language models , by using both source and target context , and show nearest neighbour models can be effective for generation in addition to density estimation . Compared to prior work on non-parametric methods for MT , our approach is arguably simpler ( in that it requires no training , as compared to Gu et al . ( 2018 ) ) and more expressive ( in that it provides access to billions of key-value pairs during inference , as compared to Zhang et al . ( 2018 ) ; Gu et al . ( 2018 ) ) . Extensive experiments show that kNN-MT scales to datastores containing billions of tokens , improving results across a range of settings . For example , it improves a state-of-the-art GermanEnglish translation model by 1.5 BLEU . kNN-MT can also be used to adapt a single model to ∗Work done while the first author was interning at Facebook AI Research . diverse domains by simply adding a domain-specific datastore—improving results by an average of 9.2 BLEU over the base model out-of-domain , and even outperforming existing models that train on these domains . Finally , language-pair-specific datastores are used to adapt a multilingual model to particular language pairs , with improvements of 3 BLEU for translating English into German and Chinese . We find that retrievals from kNN-MT are typically highly contextually relevant . 2 NEAREST NEIGHBOR MACHINE TRANSLATION . kNN-MT involves augmenting the decoder of a pre-trained machine translation model with a nearest neighbor retrieval mechanism , allowing the model direct access to a datastore of cached examples . The translation is generated word-by-word ; at each time step , we find the most similar contexts in the datastore , and compute a distribution over the corresponding target tokens , as shown in Figure 1 . This distribution is then interpolated with the output distribution from the pre-trained MT model . More specifically , given an input sequence of tokens in a source language s = ( s1 , . . . , sM1 ) , a neural MT model outputs a sequence of tokens t = ( t1 , . . . , tM2 ) in the target language . When using autoregressive decoders , the output distribution for each token ti in the target sequence is conditioned on the entire source sequence as well as the previous target tokens , p ( ti|s , t1 : i−1 ) . Let ( s , t1 : i−1 ) be the translation context and ti be the target token . Datastore creation Our datastore is constructed offline and consists of a set of key-value pairs . The key is a high-dimensional representation of the entire translation context computed by the MT decoder , f ( s , t1 : i−1 ) , where f represents a mapping from input to an intermediate representation of the decoder . The value is the corresponding ground truth target token ti . For a parallel text collection ( S , T ) , the representations are generated by a single forward pass over each example and the complete datastore is defined as follows : ( K , V ) = { ( f ( s , t1 : i−1 ) , ti ) , ∀ti ∈ t | ( s , t ) ∈ ( S , T ) } ( 1 ) Tokens from the source language are not stored directly as values in the datastore . Conditioning on the source is implicit via the keys , and the values are only target language tokens . Generation At test time , given a source x , the model outputs a distribution over the vocabulary pMT ( yi|x , ŷ1 : i−1 ) for the target yi at every step of generation , where ŷ represents the generated tokens . The model also outputs the representation f ( x , ŷ1 : i−1 ) , which is used to query the datastore for the k nearest neighbors N according to squared-L2 distance , d. In practice , the search over billions of key-value pairs is carried out using FAISS ( Johnson et al. , 2017 ) , a library for fast nearest neighbor search in high-dimensional spaces . The retrieved set is converted into a distribution over the vocabulary by applying a softmax with temperature T to the negative distances and aggregating over multiple occurrences of the same vo- cabulary item . Using a temperature greater than one flattens the distribution , and prevents overfitting to the most similar retrievals . pkNN ( yi|x , ŷ1 : i−1 ) ∝ ∑ ( kj , vj ) ∈N 1yi=vj exp ( −d ( kj , f ( x , ŷ1 : i−1 ) ) T ) ( 2 ) While a pure kNN approach is effective , we improve results by interpolating with the base model distribution , which is more robust in cases without relevant cached examples . The model and kNN distributions are interpolated with a tuned parameter λ , resulting in the final kNN-MT distribution : p ( yi|x , ŷ1 : i−1 ) = λ pkNN ( yi|x , ŷ1 : i−1 ) + ( 1− λ ) pMT ( yi|x , ŷ1 : i−1 ) ( 3 ) The complete translation is generated using beam search . kNN-MT vs. kNN-LM kNN-MT is a generalization of kNN-LM applied to conditional sequence generation , with a few important differences . First , the keys are not only conditioned on prior context , but also on the source sequence ( here , in a different language ) . This means that the representations must encode both source and target context ; we show examples in Section 6 . Second , there is an additional tuned parameter , the softmax temperature . Higher temperatures flatten the distribution and allow for greater diversity without overfitting to the retrieved contexts , as shown in Section 6 . 3 EXPERIMENTAL SETUP . We experiment with kNN-MT in three settings : ( 1 ) single language-pair translation , ( 2 ) multilingual MT and ( 3 ) domain adaptation . Data We use the following datasets for training and evaluation . WMT ’ 19 : For the single language-pair experiments , we use WMT ’ 19 data for German-English . CCMATRIX : We train our multilingual model on CCMatrix ( Schwenk et al. , 2019 ) , containing parallel data for 79 languages and 1,546 language pairs . The parallel sentences are mined from cleaned monolingual commoncrawl data created using the ccNet pipeline ( Wenzek et al. , 2019 ) . Semantically similar sentences in different languages are aligned using a learned distance measure ; we use examples where the distance measure is at least 1.06 , resulting in 4 billion sentence-pairs . NEWSTEST : The newstest2018 and newstest2019 test sets from WMT ( Bojar et al. , 2018 ; Barrault et al. , 2019 ) are used as validation and test sets for the multilingual experiments . The same GermanEnglish validation and test sets are also used for evaluation in the single language-pair and domain adaptation experiments . TED TALKS : We use the Ted Talks data prepared by Qi et al . ( 2018 ) for evaluation in the multilingual setting , particularly to explore performance for language pairs that do no include English . MULTI-DOMAINS : We use the multi-domains dataset ( Koehn & Knowles , 2017 ) , re-split by Aharoni & Goldberg ( 2020 ) for the domain adaptation experiments . It includes German-English parallel data for train/validation/test sets in five domains : Medical , Law , IT , Koran and Subtitles . Models For the single language-pair and domain adaptation experiments , we use the WMT ’ 19 German-English news translation task winner ( Ng et al. , 2019 ) , available via the FAIRSEQ library ( Ott et al. , 2019 ) .1 It is a Transformer encoder-decoder model ( Vaswani et al. , 2017 ) with 6 layers , 1,024 dimensional representations , 8,192 dimensional feedforward layers and 8 attention heads . Apart from WMT ’ 19 training data , this model is trained on over 10 billion tokens of backtranslation data and fine-tuned on newstest test sets from years prior to 2018 . In this work , we do not use ensembles or n-best reranking . For multilingual MT , we trained a 418M parameter Transformer-based encoder-decoder model on the CCMatrix data for 100K updates . The model has embedding dimension 1,024 , hidden dimension 4,096 , 12 layers in both the encoder and decoder , with 16 attention heads . To balance the 1https : //github.com/pytorch/fairseq/tree/master/examples/translation training of different language pairs , which have various resource levels , we apply temperature upsampling with T = 5 ( Arivazhagan et al. , 2019 ) . The vocabulary is shared across all languages and consists of 128K subwords extracted using sentencepiece ( Kudo & Richardson , 2018 ) .2 All results use case-sensitive detokenized BLEU , measured using SACREBLEU ( Post , 2018 ) .We provide the SACREBLEU signatures , along with details on the statistical power of our experiments , in Appendix C. kNN-MT In this work , we use a FAISS index to represent the datastore and search for nearest neighbors . The keys are stored in clusters to speed up search and quantized to 64-bytes for space efficiency ( the full-precision keys are discarded ) . The index is constructed offline via a single forward pass over every example in the given parallel text collection . We use the 1024-dimensional representation input to the final layer feedforward network as the key . Building the index involves a training phase to learn the cluster centroids . We use 5M keys for learning 131K cluster centroids for the multilingual experiments , and 1M keys for 4K clusters for in-domain data in the domain adaptation experiments . During inference , we query the datastore for 64 neighbors while searching 32 clusters . The interpolation and softmax temperature parameters are tuned on the validation sets.3 Computational Cost While kNN-MT does not add trainable model parameters , it does add some computational overhead . The primary cost of building the datastore is a single forward pass over all examples in the datastore , which is a fraction of the cost for training on the same examples for one epoch . During inference , retrieving 64 keys from a datastore containing billions of items results in a generation speed that is two orders of magnitude slower than the base MT system . Generation speed can be improved by searching fewer clusters , using smaller beams , or querying smaller datastores , with relatively minor trade-offs in performance , as we will see in Section 5 . Developing faster nearest neighbor search tools remains an active area of research ( Guo et al. , 2020 ) .

Summary: This paper describes a nearest-neighbor enhancement to NMT, where internal token-level context representations are used to index into a large data store to find relevant (source, target prefix) pairs. Since the index representation is taken from a pre-softmax representation in the decoder network, no additional training of the NMT model is required. The authors show a diverse range of strong results, from improvements using a data store over the model’s own training data, to improvement from using a collection of domain-specific corpora not present during training used for domain adaptation, to language specific collections to improve capacity of multilingual models. They are also able to show by example how the model makes MT more interpretable.