Efficient k -Nearest-Neighbor Machine Translation with Dynamic Retrieval (2024)

Abstract: To achieve non-parametric NMT domain adaptation, k -Nearest-Neighbor Machine Translation ( k NN-MT) constructs an external datastore to store domain-specific translation knowledge, which derives a k NN distribution to interpolate the prediction distribution of the NMT model via a linear interpolation coefficient \lambda . Despite its success, k NN retrieval at each timestep leads to substantial time overhead. To address this issue, dominant studies resort to k NN-MT with adaptive retrieval ( k NN-MT-AR), which dynamically estimates \lambda and skips k NN retrieval if \lambda is less than a fixed threshold. Unfortunately, k NN-MT-AR does not yield satisfactory results. In this paper, we first conduct a preliminary study to reveal two key limitations of k NN-MT-AR: 1) the optimization gap leads to inaccurate estimation of \lambda for determining k NN retrieval skipping, and 2) using a fixed threshold fails to accommodate the dynamic demands for k NN retrieval at different timesteps. To mitigate these limitations, we then propose k NN-MT with dynamic retrieval ( k NN-MT-DR) that significantly extends vanilla k NN-MT in two aspects. Firstly, we equip k NN-MT with a MLP-based classifier for determining whether to skip k NN retrieval at each timestep. Particularly, we explore several carefully-designed scalar features to fully exert the potential of the classifier. Secondly, we propose a timestep-aware threshold adjustment method to dynamically generate the threshold, which further improves the efficiency of our model. Experimental results on the widely-used datasets demonstrate the effectiveness and generality of our model.

## introduction as an effective paradigm for non-parametric domain adaptation, k-nearest-neighbor machine translation (knn-mt) @xcite derives from k-nearest-neighbor language model (knn-lm) @xcite and has garnered much attention recently @xcite @xcite . typically, knn-mt introduces translation knowledge stored in an external datastore to enhance the nmt model, which can conveniently achieve non-parametric domain adaptation by changing external datastores. in knn-mt, a datastore containing key-value pairs is first constructed with an off-the-shelf nmt model, where the key is the decoder representation and the value corresponds to its target token. during translation, the current decoder representation is used as a query to retrieve k nearest pairs from the datastore, where retrieved values are converted into a probability distribution. finally, via a linear interpolation coefficient λ, this distribution is used to adjust the prediction distribution of the nmt model. in spite of success, retrieving at each timestep incurs substantial time overhead, which becomes considerable as the datastore expands. to address this drawback, researchers have proposed two categories of approaches: 1) datastore compression that improves retrieval efficiency by reducing the size of datastores @xcite @xcite @xcite ; 2) retrieval reduction that skips some knn retrieval to speed up decoding. in this regard, the most representative work is knn-mt with adaptive retrieval (knn-mt-ar) @xcite ) that skips knn retrieval when the coefficient λ is less than a fixed threshold α. however, knn-mt-ar does not achieve desired results as reported in @xcite . in this work, we mainly focus on the studies of retrieval reduction, which is compatible with the other type of studies. to this end, we first reimplement knn-mt-ar @xcite and conduct a preliminary study to analyze its limi-tations. through in-depth analyses, we show that 1) the optimization gap leads to inaccurate estimation of λ for determining knn retrieval skipping; 2) with the increase in timesteps, the demand for knn retrieval diminishes, which proves challenging for the fixed threshold α to handle effectively. to overcome the above defects, we then significantly extend the vanilla knn-mt into knn-mt with dynamic retrieval (knn-mt-dr), which accelerates the model decoding in two aspects. concretely, instead of relying on the interpolation coefficient λ, we introduce a mlp-based classifier to explicitly determine whether to skip knn retrieval as a binary classification task. particularly, instead of using the decoder representation as the input of the classifier, we explore several carefullydesigned scalar features to fully exert the potential of the classifier. besides, we propose a timestepaware threshold adjustment method to dynamically generate the threshold, so as to further improve the efficiency of our model. to summarize, main contributions of our work include the following four aspects: retrieval reduction. in this regard, some knn retrieval is reduced to decrease time overhead for retrieval. for instance, @xcite adopt chunk-wise knn retrieval rather than timestepwise one, and @xcite explore two approaches to reduce the frequency of knn retrieval operations: 1) one introduces a caching mechanism to speed up decoding, where the cache mainly contains retrieved pairs from previous timesteps, and skip knn retrieval if the distance between the query and any cached key is less than a predefined threshold; 2) the other proposes to conduct knn retrieval when the interpolation coefficient λ is less than a predefined threshold α, which, however, does not achieve satisfactory results. our work mainly focuses on the second type of studies mentioned above. we first conduct a preliminary study to in-depth analyze two limitations of the λ-based knn retrieval skipping. to address these limitations, we introduce a classifier to explicitly determine whether to skip knn retrieval as a classification task. notably, almost concurrently with our work, @xcite also use a classifier to speed up model decoding, sharing a similar motivation with ours. however, our work not only achieves better results, but also significantly differs from theirs in the following three aspects: first, we explore several carefully-designed scalar features as the input for the classifier, which are crucial for achieving better performance. second, when training the classifier, we adopt more reasonable criteria to construct training samples. to be specific, in addition to skipping retrieval when the target token ranks the 1st position in the nmt prediction distribution, we believe that the model should also skip when the target token can not be obtained through knn retrieval. finally, based on the observation that the demand for knn retrieval diminishes as timesteps increase, we propose a timestep-aware threshold method to further improve the efficiency of our model. ## our model in this section, we significantly extend knn-mt into knn-mt-dr in the following two aspects.