[Paper Review] Better Datastore, Better Translation: Generating Datastores from Pre-Trained Models for Nearest Neural Machine Translation

Jiahuan Li, Shanbo Cheng+3 more|arXiv (Cornell University)|Dec 17, 2022

Natural Language Processing Techniques4 citations

TL;DR

This paper proposes PRED, a framework that leverages pre-trained models (PTMs) to generate higher-quality datastores for k-NN Neural Machine Translation (k-NNMT), improving retrieval accuracy and translation performance. By using PTMs for datastore construction and introducing a contrastive alignment objective to bridge representation gaps, PRED achieves superior translation results without increasing the NMT model size, outperforming baselines on bilingual and multilingual benchmarks.

ABSTRACT

Nearest Neighbor Machine Translation (kNNMT) is a simple and effective method of augmenting neural machine translation (NMT) with a token-level nearest neighbor retrieval mechanism. The effectiveness of kNNMT directly depends on the quality of retrieved neighbors. However, original kNNMT builds datastores based on representations from NMT models, which would result in poor retrieval accuracy when NMT models are not good enough, leading to sub-optimal translation performance. In this paper, we propose PRED, a framework that leverages Pre-trained models for Datastores in kNN-MT. Better representations from pre-trained models allow us to build datastores of better quality. We also design a novel contrastive alignment objective to mitigate the representation gap between the NMT model and pre-trained models, enabling the NMT model to retrieve from better datastores. We conduct extensive experiments on both bilingual and multilingual translation benchmarks, including WMT17 English $\leftrightarrow$ Chinese, WMT14 English $\leftrightarrow$ German, IWSLT14 German $\leftrightarrow$ English, and IWSLT14 multilingual datasets. Empirical results demonstrate the effectiveness of PRED.

Motivation & Objective

• To address the sub-optimal translation performance in k-NNMT caused by low-quality datastores built from underperforming NMT model representations.
• To improve Key-Value (K-V) consistency in datastores by leveraging the superior representations of pre-trained models (PTMs).
• To mitigate the representation discrepancy between the NMT model and PTM-generated datastores through a novel contrastive alignment objective.
• To maintain inference efficiency by avoiding any increase in the NMT model size while enhancing retrieval and translation quality.
• To demonstrate the effectiveness of PTM-based datastore generation across diverse bilingual and multilingual translation benchmarks.

Proposed method

• Construct datastores using representations from powerful pre-trained models (PTMs) instead of NMT models, enabling higher K-V consistency.
• Use the NMT model's decoding states as queries during inference, while retrieving from datastores built with PTM representations.
• Design a contrastive alignment objective that aligns the NMT model's query representations with corresponding keys in the PTM-generated datastore.
• Train the NMT model end-to-end with the contrastive objective to improve retrieval relevance despite representation distribution shifts.
• Apply temperature-scaled weighted fusion of NMT and k-NN probabilities during inference, using retrieved values from the PTM-built datastore.
• Ensure no increase in NMT model size by keeping the original NMT architecture and only modifying the datastore and retrieval process.

Experimental results

Research questions

• RQ1 Can leveraging pre-trained models for datastore generation significantly improve retrieval accuracy and translation quality in k-NNMT?
• RQ2 How does the representation gap between NMT models and pre-trained models affect retrieval performance, and can it be effectively mitigated?
• RQ3 Does using PTM-generated datastores lead to better K-V consistency compared to NMT-generated datastores?
• RQ4 Can a contrastive alignment objective improve retrieval performance when the query and key representations come from different models?
• RQ5 Does PRED maintain computational efficiency while outperforming existing k-NNMT methods on multilingual and bilingual benchmarks?

Key findings

• PRED achieves consistent improvements over k-NNMT and other baselines on WMT17 English↔Chinese, WMT14 English↔German, IWSLT14 German↔English, and IWSLT14 multilingual translation benchmarks.
• On WMT17 Chinese↔English, PRED achieves a BLEU score of 32.1, outperforming the strongest baseline by 1.2 BLEU points.
• On WMT14 English↔German, PRED improves BLEU by 1.0 over k-NNMT, demonstrating strong gains in low-resource and complex translation settings.
• The contrastive alignment objective significantly improves retrieval accuracy, especially for similar but non-interchangeable words like personal pronouns and linking verbs, as shown in qualitative analysis and Figure 6.
• PRED maintains the same NMT model size as k-NNMT, avoiding any increase in inference cost while achieving better translation quality.
• Ablation studies confirm that both PTM-based datastore generation and the contrastive objective are essential for performance gains, with ablation on either component leading to significant drops in BLEU.