[Paper Review] Learning Decoupled Retrieval Representation for Nearest Neighbour Neural Machine Translation

Qiang Wang, Rongxiang Weng+1 more|arXiv (Cornell University)|Sep 19, 2022

Natural Language Processing Techniques4 citations

TL;DR

This paper proposes Clknn, a method that decouples retrieval representation from standard NMT context representations using supervised contrastive learning to improve k-Nearest Neighbor Neural Machine Translation. By learning a dedicated retrieval vector via a lightweight adapter and hard negative sampling, Clknn achieves higher retrieval accuracy and BLEU scores across five domains, with minimal inference speed cost compared to vanilla kNN-MT.

ABSTRACT

K-Nearest Neighbor Neural Machine Translation (kNN-MT) successfully incorporates external corpus by retrieving word-level representations at test time. Generally, kNN-MT 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 kNN-MT.

Motivation & Objective

• To address the sub-optimal performance of kNN-MT when using shared context representations for both translation and retrieval tasks.
• To learn a dedicated, more discriminative retrieval representation that better distinguishes between similar target words.
• To improve retrieval accuracy and downstream translation performance (BLEU) in low-resource and out-of-domain settings.
• To develop an efficient method for constructing hard negative samples to enhance contrastive learning.
• To ensure the proposed method maintains inference speed comparable to vanilla kNN-MT.

Proposed method

• A lightweight feedforward network adapter is used to transform the standard NMT context representation h into a dedicated retrieval representation z.
• Supervised contrastive learning is applied with multiple positive and negative samples to train the adapter, optimizing for retrieval performance.
• Hard negative samples are constructed by clustering retrieval vectors and sampling from the nearest K=128 clusters to increase learning signal.
• The retrieval representation z is used in inference to find k-nearest neighbors in the datastore, replacing the original h as the query vector.
• The final translation probability combines the standard NMT output and the retrieval-based probability via a learnable interpolation coefficient λ.
• A PCA projection is applied to reduce dimensionality to 128 for efficient vector retrieval using FAISS.

Experimental results

Research questions

• RQ1 Can decoupling the retrieval representation from the standard NMT context representation improve retrieval accuracy and translation performance in kNN-MT?
• RQ2 Does using multiple positive and negative samples in contrastive learning lead to better retrieval representations than single-sample methods?
• RQ3 Can a fast and effective hard negative sampling strategy improve contrastive learning without significant computational cost?
• RQ4 Does the proposed method maintain inference speed comparable to vanilla kNN-MT despite the added adapter?
• RQ5 How robust is the method across in-domain and out-of-domain test sets?

Key findings

• Clknn improves BLEU scores by over 1 point on average across five domains compared to vanilla kNN-MT, with gains of up to 1.64 BLEU points on the Medical domain.
• The retrieval accuracy of Clknn consistently outperforms kNN-MT across all top-k values, demonstrating more robust and accurate neighbor retrieval.
• Using M=2 positive samples and N=32 negative samples yields the best performance, with increasing N providing stronger learning signals than increasing M.
• The method achieves comparable inference speed to kNN-MT (97% ± 2% of baseline), confirming minimal overhead from the adapter.
• The performance gap between in-domain and out-of-domain settings is reduced to only 0.3 BLEU points, indicating strong generalization and reduced reliance on in-domain data.
• t-SNE visualization confirms that Clknn learns more compact and separable representations, especially for low-frequency words, compared to the baseline.