Subset Retrieval Nearest Neighbor Machine Translation (2023)

Abstract: k-nearest-neighbor machine translation (kNN-MT) (Khandelwal et al., 2021) boosts the translation performance of trained neural machine translation (NMT) models by incorporating example-search into the decoding algorithm. However, decoding is seriously time-consuming, i.e., roughly 100 to 1,000 times slower than standard NMT, because neighbor tokens are retrieved from all target tokens of parallel data in each timestep. In this paper, we propose “Subset kNN-MT”, which improves the decoding speed of kNN-MT by two methods: (1) retrieving neighbor target tokens from a subset that is the set of neighbor sentences of the input sentence, not from all sentences, and (2) efficient distance computation technique that is suitable for subset neighbor search using a look-up table. Our proposed method achieved a speed-up of up to 132.2 times and an improvement in BLEU score of up to 1.6 compared with kNN-MT in the WMT’19 De-En translation task and the domain adaptation tasks in De-En and En-Ja.

## introduction neural machine translation (nmt) @xcite @xcite @xcite has achieved state-of-the-art performance and become the focus of many studies. recently, knn-mt @xcite has been proposed, which addresses the problem of performance degradation in out-of-domain data by incorporating example-search into the decoding algorithm. knn-mt stores translation examples as a set of key-value pairs called "datastore" and retrieves k-nearest-neighbor target tokens in decoding. the method improves the translation performance of nmt models without additional training. however, decoding is seriously timeconsuming, i.e., roughly 100 to 1,000 times slower than standard nmt, because neighbor tokens are retrieved from all target tokens of parallel data in each timestep. in particular, in a realistic opendomain setting, knn-mt may be significantly slower because it needs to retrieve neighbor tokens from a large datastore that covers various domains. we propose "subset knn-mt", which improves the decoding speed of knn-mt by two methods: (1) retrieving neighbor target tokens from a subset that is the set of neighbor sentences of the input sentence, not from all sentences, and (2) efficient distance computation technique that is suitable for subset neighbor search using a lookup table. when retrieving neighbor sentences for a given input, we can employ arbitrary sentence representations, e.g., pre-trained neural encoders or tf-idf vectors, to reduce the knn search space. when retrieving target tokens in each decoding step, the search space in subset knn-mt varies depending on the input sentence; therefore, the clustering-based search methods used in the original knn-mt cannot be used. for this purpose, we use asymmetric distance computation (adc) @xcite in subset neighbor search. our subset knn-mt achieved a speed-up of up to 132.2 times and an improvement in bleu score of up to 1.6 compared with knn-mt in the wmt'19 german-to-english general domain translation task and the domain adaptation tasks in german-to-english and english-to-japanese with open-domain settings. ## proposed model: subset knn-mt our subset knn-mt (figure 1 ) drastically accelerates vanilla knn-mt by reducing the knn search space by using sentence information (section 3.1) and efficiently computing the distance between a query and key by performing table lookup (section 3.2). ## related work the first type of example-based machine translation method was analogy-based machine translation @xcite . zhang et al. (2018); gu et al. (2018) incorporated example-based methods into nmt models, which retrieve examples according to edit distance. bulte and tezcan (2019) and xu et al. (2020) concatenated an input sentence and translations of sentences similar to it. both knn-mt and subset knn-mt retrieve knn tokens according to the distance of intermediate representations and interpolate the output probability. to improve the decoding speed of knn-mt, fast knn-mt @xcite constructs additional datastores for each source token, and reduces the knn search space using their datastores and word alignment. subset knn-mt requires a sentence datastore that is smaller than source token datastores and does not require word alignment. @xcite decreased the number of query times by retrieving chunked text; their model led to a speed-up of up to 4 times, compared with knn-mt. in contrast, subset knn-mt reduces the search space. @xcite reduced the knn search space by retrieving the neighbor sentences of the input sentence. they searched for neighboring sentences by bm25 scores with elas-ticsearch foot_3 , so our subset knn-mt with bm25 can be regarded as an approximation of their method. they also proposed "adaptive lambda", which dynamically computes the weights of the lambda of linear interpolation in equation 2 from the distance between the query and the nearest neighbor ## conclusion in this paper, we proposed "subset knn-mt", which improves the decoding speed of knn-mt by two methods: (1) retrieving neighbor tokens from only the neighbor sentences of the input sentence, not from all sentences, and (2) efficient distance computation technique that is suitable for subset neighbor search using a look-up table. our subset knn-mt achieved a speed-up of up to 132.2 times and an improvement in bleu of up to 1.6 compared with knn-mt in the wmt'19 de-en translation task and the domain adaptation tasks in de-en and en-ja. for future work, we would like to apply our method to other tasks.