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: The paper is an extension of [1]. The task in [1] is Language Modeling, while this paper is doing machine translation with the similar idea. The authors propose a non-parametric method for machine translation via a k-nearest-neighbor (KNN) classifier. Specifically, it predicts tokens with a KNN classifier over examples cached in a so-called datastore and this method can be applied to any pre-trained neural machine translation model without further training. The experiments show that it improves results across a range of settings (in-domain, out-of-domain, and multi-lingual evaluations).