Source-linked AI summary
Adaptive Nearest Neighbor Machine Translation
Xin Zheng, Zhirui Zhang, Junliang Guo, Shujian Huang, Boxing Chen, Weihua Luo, Jiajun Chen
TL;DR
Fixed-k retrieval in kNN-MT can include noisy neighbors, while translation quality is sensitive to the choice of k. Adaptive kNN-MT uses a lightweight Meta-k Network to select and aggregate neighbor predictions for each target token. It improves on vanilla kNN-MT across domains, and a network trained on one domain transfers to others with strong performance.
Problem
Method
Results
Takeaways & Limitations
A Meta-k Network trained on one domain can be applied directly to other domains and obtain strong performance, supporting the method’s generality.
Source: §1 · p. 2, §4.2 · p. 4
Abstract
from arXiv · showhide
kNN-MT, recently proposed by Khandelwal et al. (2020a), successfully combines pre-trained neural machine translation (NMT) model with token-level k-nearest-neighbor (kNN) retrieval to improve the translation accuracy. However, the traditional kNN algorithm used in kNN-MT simply retrieves a same number of nearest neighbors for each target token, which may cause prediction errors when the retrieved neighbors include noises. In this paper, we propose Adaptive kNN-MT to dynamically determine the number of k for each target token. We achieve this by introducing a light-weight Meta-k Network, which can be efficiently trained with only a few training samples. On four benchmark machine translation datasets, we demonstrate that the proposed method is able to effectively filter out the noises in retrieval results and significantly outperforms the vanilla kNN-MT model. Even more noteworthy is that the Meta-k Network learned on one domain could be directly applied to other domains and obtain consistent improvements, illustrating the generality of our method. Our implementation is open-sourced at https://github.com/zhengxxn/adaptive-knn-mt.
1 Introduction
Retrieval-based methods augment NMT, while fixed-k kNN-MT can be sensitive to the choice of k and to noisy neighbors. Adaptive kNN-MT selects and aggregates neighbor information per target token, improving benchmark performance and transferring across domains.
1 Introduction
- Retrieval-based approaches complement NMT, with cited work addressing degradation on out-of-domain sentences and rare words.The passage also notes that access to a datastore during translation makes these methods scalable, adaptable, and interpretable.Source: §1 · p. 1
- Fixed k can include noisy neighbors and makes translation quality sensitive to k, weakening robustness and generalization.The paper motivates adapting k to each target token rather than using one fixed value.Source: §1 · p. 1
- A Meta-k Network estimates the importance of candidate k-nearest-neighbor results and aggregates them to determine each target token’s prediction.The approach dynamically evaluates and uses neighbor information conditioned on each target token.Source: §1 · p. 1
- Across four domains, Adaptive kNN-MT improves average BLEU by 1.44∼2.97 over vanilla kNN-MT when K ≥4.The lightweight Meta-k Network requires only thousands of parameters and can be trained with a few samples.Source: §1 · p. 1
- A Meta-k Network trained on one domain can be applied directly to other domains and obtain strong performance.The authors present this transfer as evidence of method generality.Source: §1 · p. 2
2 Background: kNN-MT
kNN-MT uses a datastore of target tokens paired with representations of their translation contexts. At each decoding step, it retrieves neighbors for the current context and combines their prediction distribution with the NMT model’s distribution.
2 Background: kNN-MT
- The datastore pairs each training target token with the hidden representation of its translation context.The datastore can be created with a single forward pass over the training set.Source: §2 · p. 2
- At each decoding step, kNN-MT queries the datastore for the k nearest neighbors to the current context representation.The query uses the already generated target tokens and the source context.Source: §2 · p. 2
- The final target-token probability interpolates the kNN and NMT prediction distributions using λ.The vanilla NMT prediction is denoted pNMT.Source: §2 · p. 2
3 Adaptive kNN-MT
Adaptive kNN-MT replaces a fixed neighbor count with token-specific choices, using retrieved-neighbor evidence to select and combine kNN predictions. The Meta-k Network is trained while the pre-trained NMT model remains fixed.
3 Adaptive kNN-MT
- Adaptive kNN-MT lets each untranslated token use a different number of neighbors to address noise in fixed-k retrieval.A small k can filter noise but may overfit other cases, while the best k varies across datastores.Source: §3 · p. 2
- Candidate k values are drawn from powers of 2 below an upper bound K, with k = 0 representing NMT-only prediction.The method evaluates the retrieved neighbors to estimate the importance of each candidate k.Source: §3 · p. 3
- The Meta-k Network uses neighbor distances and counts of distinct values among the top i neighbors as input features.The paper motivates distances as direct evidence of neighbor importance and value diversity as a cue for kNN prediction credibility.Source: §3 · p. 3
- The Meta-k output aggregates the NMT prediction with predictions from different kNN choices to produce the final prediction.Each candidate kNN prediction is weighted by its Meta-k probability.Source: §3 · p. 3
- Training fixes the pre-trained NMT model and optimizes only the Meta-k Network using cross-entropy loss and hundreds of training samples.Source: §3 · p. 3
4 Experiments
The experiments evaluate Adaptive kNN-MT in multi-domain translation, comparing it with fixed-k and uniform baselines. Results report stronger performance than vanilla kNN-MT, along with cross-domain transfer, robustness to domain mismatch, and effectiveness with limited training data.
4.1 Experimental Setup
- The domain-adaptation experiments use a pretrained general-domain NMT model with an in-domain datastore and compare Adaptive kNN-MT against vanilla and uniform kNN-MT across IT, Medical, Koran, and Law.BLEU is the evaluation metric; Table 1 presents dataset statistics for the domains.Source: §4.1 · p. 3, Table 1 · p. 3
4.2 Main Results
- Adaptive kNN-MT significantly outperforms vanilla kNN-MT across domains and is more robust to the choice of K; uniform kNN-MT performs close to vanilla kNN-MT with small k.The results are reported in Table 2; the authors attribute the advantage to evaluating and filtering noisy retrieved neighbors.Source: §4.2 · p. 4, Table 2 · p. 4
- An Meta-k Network trained on IT achieves comparable performance on other domains to one retrained in-domain, with K = 32 in both settings.The authors interpret this as evidence that the mapping from designed features to neighbor confidence is common across domains.Source: §4.2 · p. 4, Table 3 · p. 4
- In the IT-to-Medical mismatch setting, Adaptive kNN-MT filters noisy retrievals and prevents performance degradation as much as possible.The evaluation uses the medical datastore with IT test sentences and sets K = 32.Source: §4.2 · p. 4, Table 4 · p. 4
- A Medical-domain example shows that Meta-k selects k separately for each target token, helping Adaptive kNN-MT use the in-domain datastore for word choice and language style.Table 5 presents translation examples from the Medical domain with K = 32.Source: §4.2 · pp. 4–5, Table 5 · p. 5
- On IT with K = 8, both designed features contribute to performance, with the distance feature more important; the model also beats vanilla kNN-MT with 100 training sentences or a hidden size of 8.The hidden-size setting contains around 0.6k parameters; the analysis is summarized in Table 6 and Figure 2.Source: §4.2 · p. 5, Table 6 · p. 5, Figure 2 · p. 5
5 Conclusion and Future Works
Adaptive kNN-MT uses a lightweight Meta-k Network to choose how retrieved neighbors are used for each target token. The authors report noise filtering and gains over vanilla kNN-MT, and propose extending the approach to other tasks.
5 Conclusion and Future Works
- Adaptive kNN-MT dynamically determines how retrieved neighbors are used for each target token through a lightweight Meta-k Network.Source: §5 · p. 5
- Experiments on domain-adaptation machine translation tasks show that the method filters retrieval noise and significantly outperforms vanilla kNN-MT.Source: §5 · p. 5
- The authors plan to extend the method to language modeling and question answering, which can also benefit from kNN searching.Source: §5 · p. 5
A.1 Datastore Creation
The datastore stores training-set key-value pairs in NumPy arrays, while FAISS supports nearest-neighbor indexing and search.
A.1 Datastore Creation
- Training-set key-value pairs are saved as a datastore using NumPy arrays.Source: §A.1 · p. 7
- FAISS builds an index for each datastore to enable fast nearest-neighbor search.Source: §A.1 · p. 7
- For each domain, the authors learn 4k cluster centroids and search 32 clusters for each target token during decoding.Source: §A.1 · p. 7
A.2 Hyper-Parameter Tuning for kNN-MT
The vanilla kNN-MT results depend strongly on hyperparameter choices, while Adaptive kNN-MT is reported to avoid this instability.
A.2 Hyper-Parameter Tuning for kNN-MT
- Vanilla kNN-MT performance is unstable across hyperparameter settings, whereas Adaptive kNN-MT avoids this problem.Source: §A.2 · p. 7
- The authors fix T at 10 for IT, Medical, and Law, and at 100 for Koran across experiments.Source: §A.2 · p. 7
- For vanilla kNN-MT, k and λ are tuned separately for each domain, with the optimal choices reported in Table 9.Source: §A.2 · p. 7, Table 9 · p. 7
A.3 Decoding Time
The decoding-time comparison evaluates NMT, replicated vanilla kNN-MT, and Adaptive kNN-MT on the IT test set under different batch sizes.
A.3 Decoding Time
- The authors compare decoding time for NMT, replicated kNN-MT, and Adaptive kNN-MT on the IT test set.Source: §A.3 · p. 7
- The comparison covers different batch sizes during decoding.Source: §A.3 · p. 7
- Decoding uses beam size 4 and length penalty 0.6, with results summarized in Table 8.Source: §A.3 · p. 7