Fast Nearest Neighbor Machine Translation (2022)

Abstract: Though nearest neighbor Machine Translation ( k NN-MT) (CITATION) has proved to introduce significant performance boosts over standard neural MT systems, it is prohibitively slow since it uses the entire reference corpus as the datastore for the nearest neighbor search. This means each step for each beam in the beam search has to search over the entire reference corpus. k NN-MT is thus two-orders slower than vanilla MT models, making it hard to be applied to real-world applications, especially online services. In this work, we propose Fast k NN-MT to address this issue. Fast k NN-MT constructs a significantly smaller datastore for the nearest neighbor search: for each word in a source sentence, Fast k NN-MT first selects its nearest token-level neighbors, which is limited to tokens that are the same as the query token. Then at each decoding step, in contrast to using the entire corpus as the datastore, the search space is limited to target tokens corresponding to the previously selected reference source tokens. This strategy avoids search through the whole datastore for nearest neighbors and drastically improves decoding efficiency. Without loss of performance, Fast k NN-MT is two-orders faster than k NN-MT, and is only two times slower than the standard NMT model. Fast k NN-MT enables the practical use of k NN-MT systems in real-world MT applications. The code is available at https://github.com/ShannonAI/fast-knn-nmt .

## introduction machine translation (mt) is a fundamental task in natural language processing @xcite , and the prevalence of deep neural networks has spurred a diverse array of neural machine translation (nmt) models to improve translation quality @xcite @xcite . the recently proposed k nearest neighbor (knn) mt model @xcite has proved to introduce significant performance boosts over standard neural mt systems. the basic idea behind knn-mt is that at each decoding step, the model is allowed to refer to reference target tokens with similar translation contexts in a large datastore of cached examples. the corresponding reference target tokens provide important insights on the translation token likely to appear next. one notable limitation of knn-mt is that it is prohibitively slow: it uses the entire reference corpus as the datastore for the nearest neighbor search. this means each step for each beam in the beam search has to search over the entire reference corpus. knn-mt is thus two-orders slower than vanilla mt models. the original paper of knn-mt @xcite suggests using fewer searching clusters, smaller beams and smaller datastores for generation speedup, but to achieve satisfactory results, carefully tuning on these factors under different tasks and datasets is still required according to analyses in @xcite . the computational overhead introduced by knn-mt makes it hard to be deployed on real-world online services, which usually require both model performance and runtime efficiency. in this work, we propose a fast version of knn-mt -fast knn-mt, to tackle the aforementioned issues. fast knn-mt constructs a significantly smaller datastore for the nearest neighbor search: for each word in a source sentence, fast knn-mt first selects its nearest token-level neighbors, which is limited to tokens of the same token type. then at each decoding step, in contrast to consulting the entire corpus for nearest neighbor search, the datastore for the currently decoding token is limited within the tokens of reference targets corresponding to the previously selected reference source tokens, as shown in figure 1 . the chain of mappings from the target token to the source token, then to its nearest source reference tokens, and last to cor-responding target reference tokens, can be obtained using fastalign @xcite . fast knn-mt provides several important advantages against vanilla knn-mt in terms of speedup: ## background: knn-mt given an input sequence of tokens @xmath0, • • • , x n } of length n, an mt model translates it into a target sentence in another language @xmath1, • • • , y m } of length m. a common practice to produce each token y i on the target side is to obtain a probability distribution over the vocabulary p mt (y i |x, y 1:i-1 ) from the decoder and use beam search for generation. the combination of the complete source sentence and prefix of the target sentence (x, y 1:i-1 ) is called trans-lation context. knn-mt interpolates this probability distribution with a multinomial distribution p knn (y i |x, y 1:i-1 ) derived from the k nearest neighbors of the current translation context (x, y 1:i-1 ) from a large scale datastore s: more specifically, knn-mt first constructs the datastore s using key-value pairs, where the key is the high-dimensional vector of the translation context produced by a trained mt model f (x, y 1:i-1 ), and the value is the corresponding gold target token y i , forming @xmath2, v)} = {(f (x, y 1:i-1 ), y i )}. the context-target pairs may come from any parallel corpus. then, using the dense representation of the current translation context as query @xmath3, y 1:i-1 ) and l 2 distance as measure, knn-mt searches through the entire datastore s to retrieve k nearest translation contexts along with the corresponding target tokens @xmath4, v j } k @xmath5. last, the retrieved set is transformed to a probability distribution by normalizing and aggregating the negative l 2 distances, -d, using the softmax operator with temperature t , which can be expressed as follows: integrating eq.( 2 ) into eq.( 1 ) gives the final probability of generating token y i for time step i. note that the above knn search-interpolating process is applied to each decoding step of each beam, and each iteration needs to run on the full datastore s. this gives a total time complexity of o(|s|bm), where b is the beam size and m is the target length. in order for faster nearest neighbor search, knn-mt leverages faiss @xcite , an toolkit for efficient similarity search and clustering of dense vectors. ## method: fast knn-mt the time complexity of knn-mt before search optimization is o(|s|bm) foot_1 , which is prohibitively slow when the size of the datastore s or the beam size b is large. we propose strategies to address this issue. the same as vanilla knn-mt, fast knn-mt system is built upon a separately trained mt encoder-decoder model. to get a better illustration of how fast knn-mt works, we give a toy illustration in figure 1 . we use the capitalized characters to denote source tokens and lower-cased letters to denote target tokens. given the training set, which is: ## conclusion in this work, we propose a fast version of knn-mt -fast knn-mt -to address the runtime complexity issue of the vanilla knn-mt. during decoding, fast knn-mt constructs a significantly smaller datastore for the nearest neighbor search: for each word in a source sentence, fast knn-mt selects its nearest tokens from a large-scale cache. the selected tokens are the same as the query token. then at each decoding step, in contrast to using the entire datastore, the search space is limited to target tokens corresponding to the previously selected reference source tokens. experiments demonstrate that this strategy drastically improves decoding efficiency while maintaining model performances compared to vanilla knn-mt under different settings including bilingual machine translation and domain adaptation. comprehensive ablation studies are performed to understand the behavior of each component in fast knn-mt. in future work, we plan to further improve the efficiency of fast knn-mt by applying clustering techniques to build the datastore.