Learning Decoupled Retrieval Representation for Nearest Neighbour Neural Machine Translation (2022)

Abstract: K-Nearest Neighbor Neural Machine Translation (kNNMT) successfully incorporates external corpus by retrieving word-level representations at test time. Generally, kNNMT borrows the off-the-shelf context representation in the translation task, e.g., the output of the last decoder layer, as the query vector of the retrieval task. In this work, we highlight that coupling the representations of these two tasks is sub-optimal for fine-grained retrieval. To alleviate it, we leverage supervised contrastive learning to learn the distinctive retrieval representation derived from the original context representation. We also propose a fast and effective approach to constructing hard negative samples. Experimental results on five domains show that our approach improves the retrieval accuracy and BLEU score compared to vanilla kNNMT.

## introduction conventional neural machine translation (nmt) cannot dynamically incorporate external corpus at inference once finishing training @xcite , resulting in bad performance when facing unseen domains, even if feeding millions or billions of sentence pairs for training @xcite . to address this problem, researchers developed retrieval-enhanced nmt (renmt) to flexibly incorporate external translation knowledge. early renmts leverage a search engine to find the similar bitext to improve the translation performance @xcite @xcite . however, the results of sentence-level retrieval with high similarity are generally sparse in practical applications, while noises in low similarity retrieval could lead to severe performance degradation @xcite . in this work, we attempt to decouple the context representation by learning an independent retrieval representation. to this end, we leverage supervised contrastive learning with multiple positive and negative samples to learn a good retrieval representation (called clknn). we also propose a fast and effective method to construct hard negative samples. experimental results on five domains show that our approach outperforms the vanilla knn-mt in terms of bleu and retrieval accuracy. ## background vanilla nmt given a source sentence @xmath0, x 2 , . . . , x |x| } and a target prefix @xmath1, y 2 , . . . , y t-1 }, the vanilla nmt predicts the next target word y t by: ## approach motivation according to eq. 1-4, we can see that the context representation h simultaneously plays two roles in knn-mt: (1) the semantic vector for p c ; (2) the retrieval vector for p r . we note that coupling the same h in the two scenes is sub-optimal. recall that h in the translation model is generally learned through cross-entropy loss, which only pays attention to the gold target token and ignores others. foot_0 however, a good retrieval vector should be able to distinguish between different tokens, especially those owning similar representations. therefore, we attempt to derive a new retrieval vector z from h for better retrieval performance. ## experiments setup to fairly compared with previous work @xcite , we use wmt'19 german-english news translation task winner @xcite commentary v14 and europarl v9, and randomly sample 5m samples out of 65.7m samples from its datastore. see ## analysis effect of the number of contrastive samples one of the main differences between @xcite and us is that we use multiple positive and negative samples in our training objective. we vary the number of m and n and report the bleu scores in retrieval accuracy intuitively, our approach can learn more accurate retrieval representation than vanilla knn-mt. to validate this hypothesis, we use it validation as the datastore and plot the retrieval accuracy on top-k in figure 3 . we can see that clknn has more robust retrieval accuracy than knn-mt no matter how k changes. it indicates that the performance improvement comes from our better retrieval representation. visualization we visually present the differences between baseline and clknn on embedding space. specifically, we split three categories according to the word frequency in it training set: high(the first 1%), middle(40%-60%) and low(the last 1%) foot_5 . we uniformly sample 10 unique words in each category and randomly sample 10 unique vector representations from the training datastore. we use t-sne to plot these representations, as illustrated in figure 2 . we can see that: (1) highfrequency words' representations are prone to distinguish for both baseline and clknn; (2) clknn has more close distances in the same vocabulary than baseline; (3) clknn has more robust accuracy for low-frequency words. ## conclusion in this work, we proposed to use supervised contrastive learning to decouple the context representation from vanilla knn-mt. experimental results on several tasks show that our approach outperforms knn-mt and learns a more accurate retrieval representation.