INK : Injecting k NN Knowledge in Nearest Neighbor Machine Translation (2023)

Abstract: Neural machine translation has achieved promising results on many translation tasks. However, previous studies have shown that neural models induce a non-smooth representation space, which harms its generalization results. Recently, kNN-MT has provided an effective paradigm to smooth the prediction based on neighbor representations during inference. Despite promising results, kNN-MT usually requires large inference overhead. We propose an effective training framework INK to directly smooth the representation space via adjusting representations of kNN neighbors with a small number of new parameters. The new parameters are then used to refresh the whole representation datastore to get new kNN knowledge asynchronously. This loop keeps running until convergence. Experiments on four benchmark datasets show that INK achieves average gains of 1.99 COMET and 1.0 BLEU, outperforming the state-of-the-art kNN-MT system with 0.02x memory space and 1.9x inference speedup.

## introduction neural machine translation (nmt) have achieved promising results in recent years @xcite @xcite . the target of nmt is to learn a generalized representation space to adapt to diverse scenarios. however, recent studies have shown that neural networks, such as bert and gpt, induce non-smooth representation space, limiting the generalization abilities @xcite @xcite . in nmt, we also observe a similar phenomenon in the learned representation space where low-frequency tokens disperse sparsely, even for a strong nmt model (more details are described in section experiments). due to the sparsity, many "holes" could be formed. when it is used to translate examples from an unseen domain, the performance drops sharply @xcite recently, k-nearest-neighbor machine translation (knn-mt) @xcite provides an effective solution to smooth predictions by equipping an nmt model with a key-value datastore. for each entry, the value is the target token and key is the contextualized representation at the target position. it requires a training set to record tokens and representations. by aggregating nearest neighbors during inference, the nmt model can achieve decent translation results @xcite @xcite . despite the success, knn-mt also brings new issues with the increasing scale of training data. retrieving neighbors from a large datastore @xcite ) at each decoding step is time-consuming @xcite . furthermore, once the datastore is constructed, representations can not be easily updated, limiting the performance ceiling of knn-mt. given above strengths and weaknesses of knn-mt, we propose to directly smooth the representation space with a small number of parameters. in this paper, we propose a training framework ink, to iteratively refine the representation space with the help of extracted knn knowledge (fig. 1 ). specifically, we adjust the representation distribution by aligning three kinds of representations with kullback-leibler (kl) divergence to train a small number of adaptation parameters. first, we align the contextualized representation and its target embedding to keep semantic meanings. second, we align the contextualized representations of a target token and align the extracted knn contextualized representations to address the sparsely dispersing problem. after a training epoch, we refresh the datastore asynchronously with refined models to update knn representations. during inference, we only load the off-the-shelf nmt model and tune adaptation parameters. we conduct experiments on four benchmark datasets. experiment results show that our framework brings average gains of 1.99 comet and 1.0 bleu. compared with the state-of-the-art knn-mt method (i.e. robust knn-mt; jiang et al. 2022), ink achieves better translation performance with 0.02× memory space and 1.9× inference speed. our contributions can be summarized below: ## background this section briefly introduces the working process of knn-mt and the architecture of adapter @xcite . for the latter, we will use it to improve the representation space in our framework. ## approach: ink this section introduces our training framework ink. the key idea of the proposed approach is to use knn knowledge to smooth the representation space. the training process is built on a cycled loop: extracting knn knowledge to adjust representations via a small adapter. the updated parameters are then used to refresh and refine the datastore to get new knn knowledge. we define three kinds of alignment loss to adjust representations, which are described in section 3.1, section 3.2, and section 3.3. an illustration of the proposed framework is shown in figure 2 . ## analysis and discussion we conduce more analysis in this section to better understand our ink system. ink greatly refines the representation space of the nmt model inspired by @xcite , we evaluate the quality of the representation space by computing mean knn accuracy, which measures the ratio of k-nearest representations sharing the same target token with the query representation. ideally, all of the representations in a neighborhood should share the same target token. here, we use the contextualized representations from the unseen development set as the query. for each query, the nearest representations from the training set will be checked. ablation study to show the necessity of different proposed techniques in our ink framework, we conduct ablation study in this section. in ink enjoys faster inference speed after refining the representation space, our adapted system no longer need to querying datastore during inference. we compare the inference speed foot_7 of ink and r-knn. considering that decoding with large batch size is a more practical setting @xcite , we evaluate their inference speed with increasing batch sizes. to make our evaluation results more reliable, we repeat each experiment three times and report averaged inference speed. systems @xmath0.0 26.1 29.4 ink 19.9 46.4 55.1 speedup 1.4× 1.8× 1.9× @xcite propose to prune datastore entries and conduct dimension reduction to compress the datastore. @xcite propose to in-advance narrow down the search space with word-alignment to accelerate retrieval speed. @xcite propose to retrieve a chunk of tokens at a time and conduct retrieval only at a few decoding steps with a heuristic rule. however, according to their em-pirical results, the translation performance always declines after efficiency optimization. to exclude noise in the retrieval results, @xcite propose to dynamically decide the usage of retrieved nearest neighbors with a meta-k network. @xcite propose to dynamically calibrate the knn distribution and control more hyperparameters in knn-mt. @xcite propose to build datastore with more powerful pre-trained models, e.g. xlm-r @xcite . however, all of this methods rely on a full datastore during inference. when the training data becomes larger, the inference efficiency of these approaches will becomes worse. overall, it remains an open challenge to deploy a high-quality and efficient knn-mt system. using knn knowledge to build better nmt models as datastore stores a pile of helpful translation knowledge, recent research starts exploring to use knn knowledge in the datastore to build a better nmt model. as an initial attempt, @xcite try to from scratch train a better nmt model by distilling knn knowledge into it. different from their work, we focus on smoothing the representation space of an off-the-shelf nmt model and enhancing its generalization ability via a small adapter. besides, in our devised injectand-refine training loop we keep datastore being asynchronously updated, while they use a fixed datastore. ## conclusion in this paper, we propose a novel training framework ink, to iteratively refine the representation space of the nmt model according to knn knowledge. in our framework, we devise a inject-andrefine training loop, where we adjust the representation distribution by aligning three kinds of representation and refresh the datastore asynchronously with the refined representations to update knn knowledge. experiment results on four benchmark dataset shows that ink system achieves an average gain of 1.99 comet and 1.0 bleu. compared with the state-of-the-art knn system (robust knn-mt), our ink also achieves better translation performance with 0.02× memory space and 1.9× inference speed up.