Subset Retrieval Nearest Neighbor Machine Translation
Year: 2023

Abstract:
k-nearest-neighbor machine translation (kNNMT) (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 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 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 subset kNNMT 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.

heading: 1 Introduction
text: Neural machine translation (NMT) (Sutskever et al., 2014; Bahdanau et al., 2015; Luong et al., 2015; Wu et al., 2016; Vaswani et al., 2017) has achieved state-of-the-art performance and become the focus of many studies. Recently, kNNMT (Khandelwal et al., 2021) 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) (Jégou et al., 2011) 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.
2 kNN-MT
kNN-MT (Khandelwal et al., 2021) retrieves the k-nearest-neighbor target tokens in each timestep, computes the kNN probability from the distances of retrieved tokens, and interpolates the probability with the model prediction probability. The method consists of two steps: (1) datastore creation, which creates key–value translation memory, and (2) generation, which calculates an output probability according to the nearest neighbors
174
of the cached translation memory.
Datastore Construction A typical NMT model is composed of an encoder that encodes a source sentence x = (x1, x2, . . . , x|x|) ∈ V |x|X and a decoder that generates target tokens y = (y1, y2, . . . , y|y|) ∈ V |y|Y where |x| and |y| are the lengths of sentences x and y, respectively, and VX and VY are the vocabularies of the source language and target language, respectively. The t-th target token yt is generated according to its output probability P (yt|x,y<t) over the target vocabulary, calculated from the source sentence x and generated target tokens y<t. kNN-MT stores pairs of Ddimensional vectors and tokens in a datastore, represented as key–value memory M ⊆ RD × VY . The key (∈ RD) is an intermediate representation of the final decoder layer obtained by teacher forcing a parallel sentence pair (x,y) to the NMT model, and the value is a ground-truth target token yt. The datastore is formally defined as follows:
M = {(f(x,y<t), yt) | (x,y) ∈ D, 1 ≤ t ≤ |y|}, (1) where D is parallel data and f : V |x|X × V t−1Y → RD is a function that returns the D-dimensional intermediate representation of the final decoder layer from the source sentence and generated target tokens. In our model, as in (Khandelwal et al., 2021), the key is the intermediate representation before it is passed to the final feed-forward network.
Generation During decoding, kNN-MT generates output probabilities by computing the linear interpolation between the kNN and MT probabili-
ties, pkNN and pMT, as follows:
P (yt|x,y<t) = λpkNN(yt|x,y<t) + (1− λ)pMT(yt|x,y<t), (2)
where λ is a hyperparameter for weighting the kNN probability. Let f(x,y<t) be the query vector at timestep t. The top i-th key and value in the k-nearest-neighbor are ki ∈ RD and vi ∈ VY , respectively. Then pkNN is defined as follows:
pkNN(yt|x,y<t)
∝ k∑
i=1
1yt=vi exp (−∥ki − f(x,y<t)∥22 τ ) , (3)
where τ is the temperature for pkNN, and we set τ = 100. Note that this kNN search is seriously time-consuming1 (Khandelwal et al., 2021).
3 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).

heading: 3.1 Subset Retrieval
text: Sentence Datastore Construction In our method, we construct a sentence datastore that stores pairs comprising a source sentence vector
1In our experiments on the WMT’19 German-to-English, the datastore has 862M tokens, the vocabulary size is 42k, and the batch size was set to 12,000 tokens. While a normal Transformer translates 2,000 sentences in 7.5 seconds, kNNMT takes 2446.0 seconds. Note the kNN search is executed for each timestep in generating a target sentence.
and a target sentence. Concretely, a sentence datastore S is defined as follows:
S = {(h(x),y) | (x,y) ∈ D}, (4)
where h : V |x|X → RD ′
represents a sentence encoder, which is a function that returns a D′dimensional vector representation of a source sentence.
Decoding At the beginning of decoding, the model retrieves the n-nearest-neighbor sentences of the given input sentence from the sentence datastore S . Let Ŝ ⊂ S be the subset comprising nnearest-neighbor sentences. The nearest neighbor search space for target tokens in kNN-MT is then drastically reduced by constructing the datastore corresponding to Ŝ as follows:
M̂ = {(f(x,y<t), yt) | (h(x),y) ∈ Ŝ, 1 ≤ t ≤ |y|}, (5)
where M̂ ⊂ M is the reduced datastore for the translation examples coming from the n-nearestneighbor sentences. During decoding, the model uses the same algorithm as kNN-MT except that M̂ is used as the datastore instead of M. The proposed method reduces the size of the nearest neighbor search space for the target tokens from |D| to n (≪ |D|) sentences.

heading: 3.2 Efficient Distance Computation Using Lookup Table
text: Subset kNN-MT retrieves the k-nearest-neighbor target tokens by an efficient distance computation method that uses a look-up table. In the original kNN-MT, inverted file index (IVF) is used
for retrieving kNN tokens. IVF divides the search space into Nlist clusters and retrieves tokens from the neighbor clusters. In contrast, in subset kNNMT, the search space varies dynamically depending on the input sentence. Therefore, clusteringbased search methods cannot be used; instead, it is necessary to calculate the distance for each key in the subset. For this purpose, we use asymmetric distance computation (ADC) (Jégou et al., 2011) instead of the usual distance computation between floating-point vectors. In ADC, the number of table lookup is linearly proportional to the number of keys N in the subset. Therefore, it is not suitable for searching in large datastore M, but in a small subset M̂, the search is faster than the direct calculation of the L2 distance.
Product Quantization (PQ) The kNN-MT datastore M may become too large because it stores high-dimensional intermediate representations of all target tokens of parallel data. For instance, the WMT’19 German-to-English parallel data, which is used in our experiments, contains 862M tokens on the target side. Therefore, if vectors were stored directly, the datastore would occupy 3.2 TiB when a 1024-dimensional vector as a key 2, and this would be hard to load into RAM. To solve this memory problem, product quantization (PQ) (Jégou et al., 2011) is used in both kNNMT and our subset kNN-MT, which includes both source sentence and target token search.
PQ splits a D-dimensional vector into M subvectors and quantizes for each DM -dimensional sub-vector. Codebooks are learned by k-means clustering of key vectors in each subspace. It is computed iteratively by: (1) assigning the code of a key to its nearest neighbor centroid (2) and updating the centroid of keys assigned to the code. The m-th sub-space’s codebook Cm is formulated as follows:
Cm = {cm1 , . . . , cmL }, cml ∈ R D M . (6)
In this work, each codebook size is set to L = 256. A vector q ∈ RD is quantized and its code vector q̄ is calculated as follows:
q̄ = [q̄1, . . . , q̄M ]⊤ ∈ {1, . . . , L}M , (7) q̄m = argmin
l ∥qm − cml ∥22, qm ∈ R
D M . (8)
23.2 TiB ≃ 862.6M tokens × 1024 dimension × 32 bits (float size)/8 bits (byte size)/10244
Asymmetric Distance Computation (ADC) Our method efficiently computes the distance between a query vector and quantized key vectors using ADC (Jégou et al., 2011) (Figure 2). ADC computes the distance between a query vector q ∈ RD and N key codes K̄ = {k̄i}Ni=1 ⊆ {1, . . . , L}M . First, the distance look-up table Am ∈ RL is computed by calculating the distance between a query qm and the codes cml ∈ Cm in each sub-space m, as follows:
Aml = ∥qm − cml ∥22. (9)
Second, the distance between a query and each key d(q, k̄i) is obtained by looking up the distance table as follows:
d(q, k̄i) =
M∑
m=1
dm(q m, k̄mi ) =
M∑
m=1
Amk̄mi . (10)
A look-up table in each subspace, Am ∈ RL, consists of the distance between a query and codes. The number of codes in each subspace is L and a distance is a scalar; therefore, Am has L distances. And the table look-up key is the code of a key itself, i.e., if the m-th subspace’s code of a key is 5, ADC looks-up Am5 . By using ADC, the distance is computed only once3 (Equation 9) and does not decode PQ codes into D-dimensional key vectors; therefore, it can compute the distance while keeping the key in the quantization code, and the k-nearest-neighbor tokens are efficiently retrieved from M̂.

heading: 3.3 Sentence Encoder
text: In our subset kNN-MT, a variety of sentence encoder models can be employed. The more similar sentences extracted from M, the more likely the subset M̂ comprises the target tokens that are useful for translation. Hence, we need sentence encoders that compute vector representations whose distances are close for similar sentences.
In this work, we employ two types of representations: neural and non-neural. We can employ pre-trained neural sentence encoders. While they require to support the source language, we expect that the retrieved sentences are more similar than other encoders because we can use models that have been trained to minimize the vector
3The direct distance computation requires N times calculations according to ∥q − k∥2. ADC computes the distance only L ≪ N times and just looks-up the table N times.
distance between similar sentences (Reimers and Gurevych, 2019). An NMT encoder can also be used as a sentence encoder by applying average pooling to its intermediate representations. This does not require any external resources, but it is not trained from the supervision of sentence representations. Alternatively, we can also use nonneural models like TF-IDF. However, it is not clear whether TF-IDF based similarity is suitable for our method. This is because even if sentences with close surface expressions are retrieved, they do not necessarily have similar meanings and may not yield the candidate tokens needed for translation.

heading: 4 Experiments

heading: 4.1 Setup
text: We compared the translation quality and speed of our subset kNN-MT with those of the conventional kNN-MT in open-domain settings that assume a domain of an input sentence is unknown. The translation quality was measured by sacreBLEU (Post, 2018) and COMET (Rei et al., 2020). The speed was evaluated on a single NVIDIA V100 GPU. We varied the batch size settings: either 12,000 tokens (B∞), to simulate the document translation scenario, or a single sentence (B1), to simulate the online translation scenario. The beam size was set to 5, and the length penalty was set to 1.0.
k-Nearest-Neighbor Search In kNN-MT, we set the number of nearest neighbor tokens to k = 16. We used FAISS (Johnson et al., 2019) to retrieve the kNN tokens in kNN-MT and for neighbor sentence search in subset kNNMT. The subset search and ADC were implemented in PYTORCH. We use approximate distance computed from quantized keys instead of full-precision keys in Equation 3, following the original kNN-MT (Khandelwal et al., 2021) implementation. The kNN-MT datastore and our sentence datastore used IVF and optimized PQ (OPQ) (Ge et al., 2014). OPQ rotates vectors to minimize the quantization error of PQ. The subset kNN-MT datastore is not applied clustering since we need to extract subset tokens. In this datastore, the 1024-dimensional vector representation, i.e., D = 1024, was reduced in dimensionality to 256-dimensions by principal component analysis (PCA), and these vectors were then
quantized by PQ. At search time, a query vector is pre-transformed to 256-dimensions by multiplying the PCA matrix, and then the kNN target tokens are searched by ADC. The subset of a datastore can be loaded into GPU memory since it is significantly smaller than the original kNNMT datastore, so we retrieved k-nearest-neighbor tokens from a subset on a GPU.
Sentence Encoder We compared 4 different sentence encoders: LaBSE, AvgEnc, TF-IDF, and BM25. LaBSE (Feng et al., 2022) is a pre-trained sentence encoder, fine-tuned from multilingual BERT. AvgEnc is an average pooled encoder hidden vector of the Transformer NMT model, which is also used for translation. TF-IDF (Jones, 1972) and BM25 (Jones et al., 2000) compute vectors weighted the important words in a sentence. We used the raw count of tokens as the term frequency and applied add-one smoothing to calculate the inverse document frequency, where a sentence was regarded as a document. We set k1 = 2.0, b = 0.75 in BM25 (Jones et al., 2000). Both TF-IDF and BM25 vectors were normalized by their L2norm and their dimensionality was reduced to 256- dimensions by singular value decomposition.

heading: 4.2 In-Domain Translation
text: We evaluated the translation quality and speed of subset kNN-MT in the WMT’19 De-En translation task (newstest2019; 2,000 sentences) and compared them with the kNN-MT baselines (Khandelwal et al., 2021; Meng et al., 2022). We used a trained Transformer big implemented in FAIRSEQ (Ott et al., 2019) as the base MT model. We constructed the datastore from the parallel data of the WMT’19 De-En news translation task with subword lengths of 250 or less and a sentence length ratio of 1.5 or less between the source and target sentences. The datastore contained 862.6M target tokens obtained from 29.5M sentence pairs. The subset size was set to n = 512.
Table 1 shows our experimental results. In the table, “tok/s” denotes the number of tokens generated per second. The table shows that, although kNN-MT improves 0.9 BLEU point from the base MT without additional training, the decoding speed is 326.1 times and 51.7 times slower with the B∞ and B1 settings, respectively. In contrast, our subset kNN-MT (h: LaBSE) is 111.8 times (with B∞) and 47.4 times (with B1) faster than kNN-MT with no degradation in the BLEU
score. Subset kNN-MT (h: AvgEnc) achieved speed-ups of 92.7 times (with B∞) and 38.9 times (with B1) with a slight quality degradation (−0.2 BLEU and −0.05 COMET), despite using no external models. We also evaluated our subset kNNMT when using non-neural sentence encoders (h: TF-IDF, BM25). The results show that both TFIDF and BM25 can generate translations with almost the same BLEU score and speed as neural sentence encoders. In summary, this experiment showed that our subset kNN-MT is two orders of magnitude faster than kNN-MT and has the same translation performance.

heading: 4.3 Domain Adaptation
text: German-to-English We evaluated subset kNNMT on out-of-domain translation in the IT, Koran, Law, Medical, and Subtitles domains (Koehn and Knowles, 2017; Aharoni and Goldberg, 2020) with open-domain settings. The datastore was constructed from parallel data by merging all target domains and the general domain (WMT’19 De-En) assuming that the domain of the input sentences is unknown. The datastore contained 895.9M tokens obtained from 30.8M sentence pairs. The NMT model is the same as that used in Section 4.2 trained from WMT’19 De-En. The subset size was set to n = 256, and the batch size was set to 12,000 tokens.
Table 2 shows the results. Compared with base MT, kNN-MT improves the translation performance in all domains but the decoding speed is much slower. In contrast, our subset kNNMT generates translations faster than kNN-MT. However, in the domain adaptation task, there are differences in translation quality between those using neural sentence encoders and those using non-neural sentence encoders. The table shows
that the use of non-neural sentence encoders (TFIDF and BM25) causes drop in translation quality, whereas the use of neural sentence encoders (LaBSE and AvgEnc) do not. In addition, compared with kNN-MT, our subset kNN-MT with neural encoders achieves an improvement of up to 1.6 BLEU points on some datasets. In summary, these results show that neural sentence encoders are effective in retrieving domain-specific nearest neighbor sentences from a large datastore.
English-to-Japanese We also evaluated our model on English-to-Japanese translation. We used a pre-trained Transformer big model trained from JParaCrawl v3 (Morishita et al., 2022) and evaluated its performance on Asian Scientific Paper Excerpt Corpus (ASPEC) (Nakazawa et al., 2016) and Kyoto Free Translation Task (KFTT; created from Wikipedia’s Kyoto articles) (Neubig, 2011). The datastore was constructed from parallel data by merging ASPEC, KFTT, and the general domain (JParaCrawl v3). Note that ASPEC contains 3M sentence pairs, but we used only the first 2M pairs for the datastore to remove noisy data, following Neubig (2014). The datastore contained 735.9M tokens obtained from 24.4M sentence pairs. The subset size was set to n = 512, and the batch size was set to 12,000 tokens.
Table 3 shows the results. These show that kNN-MT improves out-of-domain translation performance compared with base MT on other language pairs other than German-to-English. On English-to-Japanese, subset kNN-MT improves the decoding speed, but subset kNN-MT with TFIDF and BM25 degrades the translation quality compared with kNN-MT. However, subset kNNMT still achieves higher BLEU scores than base MT without any additional training steps, and it is two orders of magnitude faster than kNN-MT.
In summary, subset kNN-MT can achieve better translation performance than base MT in exchange for a small slowdown in open-domain settings.

heading: 5 Discussion

heading: 5.1 Case Study: Effects of Subset Search
text: Translation examples in the medical domain are shown in Table 4 and the search results of the top3 nearest neighbor sentences are shown in Table 5. In the table, the subset kNN-MT results are obtained using a LaBSE encoder. Table 4 shows that subset kNN-MT correctly generates the medical term “Co-administration”. The results of the nearest neighbor sentence search (Table 5) show that “Co-administration” is included in the subset. In detail, there are 30 cases of “Co-administration” and no case of “A joint use” in the whole subset consisting of k = 256 neighbor sentences. Base MT and kNN-MT have the subwords of “Coadministration” in the candidates; however, the subwords of “A joint use” have higher scores. Table 6 shows the negative log-likelihood (NLL) of the first three tokens and their average for each model. The second token of subset kNN-MT, “- ” (hyphen), has a significantly lower NLL than the other tokens. The number of “joint” and “- ” in the subset were 0 and 101, respectively, and the k-nearest-neighbor tokens were all “-” in subset kNN-MT. Therefore, the NLL was low because pkNN(“-”) = 1.0, so the joint probability of a beam that generates the sequence “Coadministration” is higher than “A joint use”.
In summary, the proposed method can retrieve more appropriate words by searching a subset that consists only of neighboring cases.
Input Eine gemeinsame Anwendung von Nifedipin und Rifampicin ist daher kontraindiziert. Reference Co-administration of nifedipine with rifampicin is therefore contra-indicated. Base MT A joint use of nifedipine and rifampicin is therefore contraindicated. kNN-MT A joint use of nifedipine and rifampicin is therefore contraindicated. Subset kNN-MT Co-administration of nifedipine and rifampicin is therefore contraindicated.
Table 4: Translation examples in the medical domain.
S-1 Die gemeinsame Anwendung von Ciprofloxacin und Tizanidin ist kontraindiziert. S-2 Rifampicin und Nilotinib sollten nicht gleichzeitig angewendet werden. S-3 Die gleichzeitige Anwendung von Ribavirin und Didanosin wird nicht empfohlen.
T-1 Co-administration of ciprofloxacin and tizanidine is contra-indicated. T-2 Rifampicin and nilotinib should not be used concomitantly. T-3 Co-administration of ribavirin and didanosine is not recommended.
Table 5: Top-3 neighbor sentences of our subset kNNMT in Table 4. “S-” and “T-” denote the top-n neighbor source sentences and their translations, respectively.

heading: 5.2 Diversity of Subset Sentences
text: We hypothesize that the noise introduced by sentence encoders causes the difference in accuracy. In this section, we investigate whether a better sentence encoder would reduce the noise injected into the subset. In particular, we investigated the relationship between vocabulary diversity in the subset and translation quality in the medical domain. Because an output sentence is affected by the subset, we measured the unique token ratio of both source and target languages in the subset as the diversity as follows:
number of unique tokens number of subset tokens . (11)
Table 7 shows the BLEU score and unique token ratio for the various sentence encoders, in which “source” and “target” indicate the diversity of the neighbor sentences on the source-side and target-side, respectively. The results show that the more diverse the source-side is, the more diverse the target-side is. It also shows that the less diversity in the vocabulary of both the source and target languages in the subset, the higher BLEU score.
We also investigated the relationship between sentence encoder representation and BLEU scores. We found that using a model more accurately represents sentence similarity improves the BLEU score. In particular, we evaluated translation quality when noise was injected into the subset by retrieving n sentences from outside the nearest neighbor. Table 8 shows the results of various n-selection methods when LaBSE was used as the sentence encoder. In the table, “Top” indicates the n-nearest-neighbor sentences, “Bottom
of 2n” the n furthest sentences of 2n neighbor sentences, and “Random of 2n” n sentences randomly selected from 2n neighbor sentences. The “Bottom of 2n” and “Random of 2n” have higher diversity than the “Top” on both the source- and target-sides, and the BLEU scores are correspondingly lower. These experiments showed that a sentence encoder that calculates similarity appropriately can reduce noise and prevent the degradation of translation performance because the subset consists only of similar sentences.

heading: 5.3 Analysis of Decoding Speed
text: Efficiency of ADC Subset kNN-MT computes the distance between a query vector and key vectors using ADC as described in Section 3.2. The efficiency of ADC in WMT’19 De-En is demonstrated in Table 9. The results show that “w/ ADC” is roughly 4 to 5 times faster than “w/o ADC”.
Effect of Parallelization The method and implementation of our subset kNN-MT are designed for parallel computing. We measured the translation speed for different batch sizes in WMT’19 De-En. Figure 3(a) shows that subset kNN-MT (h: LaBSE) is two orders of magnitude faster than kNN-MT even when the batch size is increased.
Subset Size We measured the translation speed for different subset sizes, i.e., the number of nnearest-neighbor sentences in WMT’19 De-En. Figure 3 (b) shows the translation speed of subset kNN-MT (h: LaBSE). Subset kNN-MT is two orders of magnitude faster than kNN-MT even when
the subset size is increased. The results also show that the speed becomes slower from n = 256 compared with base MT. We also found that 71.7% of the time was spent searching for the kNN tokens from the subset when n = 2048. Although ADC lookup search is slow for a large datastore, it is fast for kNN search when the subset size n is not large (Matsui et al., 2018), e.g., n = 512.
Figure 3(c) shows the results for translation quality on the development set (newstest2018). The results show that a larger n improves BLEU up to n = 512, but decreases for greater values of n. In terms of both the translation quality and translation speed, we set n = 512 for WMT’19 De-En.

heading: 6 Related Work
text: The first type of example-based machine translation method was analogy-based machine translation (Nagao, 1984). 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 kNNMT 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 (Meng et al., 2022) 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. Martins et al. (2022) 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. Dai et al. (2023) reduced the kNN search space by retrieving the neighbor sentences of the input sentence. They searched for neighboring sentences by BM25 scores with ElasticSearch4, 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
4https://github.com/elastic/ elasticsearch
key vectors. However, adaptive lambda requires an exact distance and cannot employ datastore quantization and the ADC lookup. To improve the translation performance of kNN-MT, Zheng et al. (2021) computed the weighted average of kNN probabilities pkNN over multiple values of k. Each weight is predicted by “meta-k network”, trained to minimize cross-entropy in the training data. For the other tasks, kNN-LM (Khandelwal et al., 2020), Efficient kNN-LM (He et al., 2021), and RETRO (Borgeaud et al., 2022) used kNN search for language modeling (LM). Our subset search method cannot be applied to LM because the entire input cannot be obtained.
In the field of kNN search, Matsui et al. (2018) allowed search in dynamically created subsets, whereas conventional search methods assume only full search. Subset kNN-MT retrieves kNN tokens from a subset depending on a given input. In our subset kNN-MT, the decoding speed is slow when the subset size n is large. The bottleneck is the lookup in the distance table, and this can be improved by efficient look-up methods that uses SIMD (André et al., 2015; Matsui et al., 2022).

heading: 7 Conclusion
text: 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.
Limitations
This study focuses only on improving the speed of kNN-MT during decoding; other problems with kNN-MT remain. For example, it still demands large amounts of memory and disk space for the target token datastore. In addition, our subset kNN-MT requires to construct a sentence datastore; therefore, the memory and disk requirements are increased. For example, the quantized target token datastore has 52GB (|M| = 862,648,422) and our sentence datastore has 2GB (|S| = 29,540,337) in the experiment of WMT’19 De-En (Section 4.2). Although subset kNN-MT is faster than the original kNN-MT in inference, datastore construction is still time-consuming. The decoding latency of our subset kNN-MT is still several times slower than base MT for large batch sizes. The experiments reported in this paper evaluated the inference speed of the proposed method on a single computer and single run only; the amount of speed improvement may differ when different computer architectures are used.
Ethical Consideration
We construct both kNN-MT and subset kNN-MT datastores from open datasets; therefore, if their datasets have toxic text, kNN-MT and our subset kNN-MT may have the risk of generating toxic contents.

heading: Acknowledgements
text: This work was partially supported by JSPS KAKENHI Grant Number JP22J1127 and JP22KJ2286.

heading: A Datasets, Tools, Models
text: Datasets Parallel data of the WMT’19 De-En translation task can be used for research purposes as described in https://www.statmt.org/ wmt19/translation-task.html. The five domain adaptation datasets in De-En can be used for research purposes as described in the paper (Aharoni and Goldberg, 2020). ASPEC can be used for research purposes as described in https://jipsti.jst.go.jp/ aspec/. KFTT is licensed by Creative Commons Attribution-Share-Alike License 3.0.
Tools FAIRSEQ and FAISS are MIT-licensed.
Models We used the following pre-trained NMT models implemented in FAIRSEQ.
• De-En: https://dl. fbaipublicfiles.com/fairseq/ models/wmt19.de-en.ffn8192. tar.gz
• En-Ja: http://www.kecl.ntt. co.jp/icl/lirg/jparacrawl/ release/3.0/pretrained_models/ en-ja/big.tar.gz
The De-En model is included in FAIRSEQ and it is MIT-licensed. The Ja-En model is licensed by Nippon Telegraph and Telephone Corporation (NTT) for research use only as described in http://www.kecl.ntt.co.jp/ icl/lirg/jparacrawl/.
We used the pre-trained LaBSE model licensed by Apache-2.0.

heading: B Pseudo Code for ADC lookup
text: Algorithm 1 shows the pseudo code for the ADC lookup described in Section 3.2. The function COMPUTE_DISTANCES calculates the squared Euclidean distances between a query vector and each quantized key vector by looking up the distance table.

heading: C Tuning of the Subset Size in Domain Adaptation
text: Section 5.3 showed that n = 256 and 512 are in balance between speed and quality. To tune the
Algorithm 1 ADC lookup Require:
query; q ∈ RD quantized keys; K̄ = {k̄i}Ni=1 ⊆ {1, . . . , L}M codebook; C = {C1, . . . , CM},
where Cm = {cml }Ll=1 ⊆ R D M
Ensure: distances; d ∈ RN
subset size n in the domain adaptation task, we evaluated for n = 256 and 512 on the development set of each domain, and the choice of n was judged by the averaged BLEU. Table 10 and 11 show the results of the domain adaptation translation on each development set. We tuned the subset size by using LaBSE for the sentence encoder. Finally, we chose n = 256 for the German-toEnglish and n = 512 for the English-to-Japanese domain adaptation tasks.

heading: D Details of Translation Quality
text: We evaluated all experiments by BLEU, COMET, and chrF scores.
Table 12, 13, and 14 show the results of the WMT’19 De-En translation task, the domain adaptation task in De-En, and En-Ja, respectively. Note that Table 13 only shows COMET and chrF scores and the BLEU scores are shown in Table 2 due to space limitations.
E Details of kNN Indexes.
The details of the kNN indexes are shown in Table 15.

heading: F Domain Adaptation with Closed Domain Settings
text: We carried out the German-to-English domain adaptation experiments faithful to the original kNN-MT settings. In this experiment, the datastore for each domain was created only from the parallel data of the target domain, assuming a scenario where the domain of the input sentences is known. Note that the general domain data, i.e., the training data of the WMT’19 De-En translation task, is not included in the datastores.
Table 16 shows the German-to-English domain adaptation translation results in closed-domain settings. The original kNN-MT is faster than that of open-domain settings because the datastore is smaller; however, our subset kNN-MT is still 10 times faster than the original kNN-MT.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
After Conclusion ("Limitations" section)
3 A2. Did you discuss any potential risks of your work? After Limitations ("Ethical Consideration" section)
3 A3. Do the abstract and introduction summarize the paper’s main claims? Section 1
3 A4. Have you used AI writing assistants when working on this paper? We use tools that only assist with language: deepl, grammarly.
B 3 Did you use or create scientific artifacts? Section 4
3 B1. Did you cite the creators of artifacts you used? Section 4
3 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? Appendix (Section A: Dataset, Tools, Models)
3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Appendix (Section A: Datasets, Tools, Models)
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? We noted in the Ethical Consideration section that our used data may contain toxic contents.
3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Section 4
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Section 4
C 3 Did you run computational experiments? Section 4 and 5
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Section 4
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? Section 4 and 5
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? We report the experimental results of just a single run and that is noted in Limitations section.
3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Section 4
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: For future work, we would like to apply our method to other tasks. Limitations
This study focuses only on improving the speed of kNN-MT during decoding; other problems with kNN-MT remain. For example, it still demands large amounts of memory and disk space for the target token datastore. In addition, our subset kNN-MT requires to construct a sentence datastore; therefore, the memory and disk requirements are increased. For example, the quantized target token datastore has 52GB (|M| = 862,648,422) and our sentence datastore has 2GB (|S| = 29,540,337) in the experiment of WMT’19 De-En (Section 4. 2). Although subset kNN-MT is faster than the original kNN-MT in inference, datastore construction is still time-consuming. The decoding latency of our subset kNN-MT is still several times slower than base MT for large batch sizes. The experiments reported in this paper evaluated the inference speed of the proposed method on a single computer and single run only; the amount of speed improvement may differ when different computer architectures are used. Ethical Consideration
We construct both kNN-MT and subset kNN-MT datastores from open datasets; therefore, if their datasets have toxic text, kNN-MT and our subset kNN-MT may have the risk of generating toxic contents