Bridging the Domain Gaps in Context Representations for k -Nearest Neighbor Neural Machine Translation (2023)

Abstract: k -Nearest neighbor machine translation ( k NN-MT) has attracted increasing attention due to its ability to non-parametrically adapt to new translation domains. By using an upstream NMT model to traverse the downstream training corpus, it is equipped with a datastore containing vectorized key-value pairs, which are retrieved during inference to benefit translation.However, there often exists a significant gap between upstream and downstream domains, which hurts the datastore retrieval and the final translation quality.To deal with this issue, we propose a novel approach to boost the datastore retrieval of k NN-MT by reconstructing the original datastore.Concretely, we design a reviser to revise the key representations, making them better fit for the downstream domain. The reviser is trained using the collected semantically-related key-queries pairs, and optimized by two proposed losses: one is the key-queries semantic distance ensuring each revised key representation is semantically related to its corresponding queries, and the other is an L2-norm loss encouraging revised key representations to effectively retain the knowledge learned by the upstream NMT model. Extensive experiments on domain adaptation tasks demonstrate that our method can effectively boost the datastore retrieval and translation quality of k NN-MT.Our code is available at https://github.com/DeepLearnXMU/Revised-knn-mt .

## introduction the recently proposed k-nearest neighbors machine translation (knn-mt) @xcite is increasingly receiving attention from the community of machine translation due to its advantage on non-parametric domain adaptation @xcite @xcite . given an upstream nmt model, knn-mt first uses figure 1 : an example of datastore retrieval, where news and it are the upstream and downstream domains, respectively. we first build a downstream nmt model by fine-tuning the upstream nmt model on the downstream training corpus. then, we use the downstream nmt model to re-traverse the downstream training corpus, constructing a downstream datastore. finally, we reuse the upstream and downstream nmt model to conduct retrieval on the original and downstream datastores, respectively. the result shows that the nearest neighbors retrieved by the same query are quite different, and only the retrieved nearest neighbors from the downstream datastore contain the ground-truth token "spreadsheet". the downstream training corpus to establish a datastore containing key-value pairs, where each key is the representation of the nmt decoder and its value is the corresponding target token. during inference, it uses the current decoder representation as a query to retrieve n k nearest key-value pairs from the datastore. afterwards, the retrieved values are transformed into a probability distribution based on the query-key distances, denoted as knn distribution. finally, this distribution is interpolated with the prediction distribution of the nmt model to adjust the prediction translation. by doing so, the upstream nmt model can be easily adapted to diverse domains by equipping domain-specific datastores without additional parameters. to avoid confusion in subsequent descriptions, we name the datastore in conventional knn-mt as the original datastore. however, there often exists a significant domain gap between the upstream nmt model and the downstream training corpus @xcite . the learned key representations of the original datastore deviate from the ideal distribution of downstream-domain key representation. as shown in figure 1 , in the original datastore built by the news domain nmt model, the nearest neighbors of the query contain the out-domain token "table" rather than the target token "spreadsheet" from the it domain. this hurts the datastore retrieval of knn-mt. to alleviate the negative impact of the retrieval error, previous studies resort to dynamically estimating the weight of knn distribution for the final prediction @xcite @xcite . however, these studies ignore the key representation learning, which is the basis of constructing datastore, and low-quality key representations tend to result in retrieval errors. to bridge the domain gap, a natural choice is to fine-tune the nmt model on the downstream training corpus to obtain the downstream nmt model and then use it to build a downstream datastore. however, this method has two serious defects: 1) it is required to deploy multiple domain-specific nmt models when dealing with multi-domain translations, involving huge system deployment overhead. for example, in the commonly-used knn-mt datasets (aharoni and goldberg, 2020) involving four downstream domains, this method has to construct four nmt models with datastores, consuming 37.2g gpu memory with 1,028m parameters. by contrast, knn-mt involves only one nmt model and four datastores, consuming 11.3g gpu memory with 257m parameters; 2) it tends to be affected by the notorious catastrophic forgetting problem, weakening the adaptability of knn-mt. this may result from the fine-tuned nmt model tending to forget previous upstream-domain knowledge and are therefore challenging to adapt to other domains. thus, how to make more effective domain adaptation using knn-mt remains a problem worth exploring. in this paper, we propose a novel approach to boost the datastore retrieval of knn-mt by reconstructing the original datastore. concretely, we design a key representation reviser that revises the key representations in an offline manner, so that they can better adapt to the retrieval from the downstream domain. this reviser is a two-layer feed-forward (ffn) with a relu function, which is fed with the information about a key representa-tion k, and outputs an inductive bias ∆k to revise k as @xmath0. to train the reviser, we first use the downstream nmt model to extract semanticallyrelated key-queries pairs from the downstream datastore, and then use their counterparts in the upstream nmt model and original datastore as supervision signals of the reviser. for each key-queries pair, we introduce two training losses to jointly optimize the reviser: 1) the semantic distance loss, which encourages each revised key representation to be adjacent to its semantically-related queries; 2) the semantic consistency loss, which avoids the revised key representation to be far from the original one, and thus preserving the knowledge learned by the upstream nmt model. to summarize, our contributions are as follows: ## preliminary study in this section, we first briefly introduce knn-mt @xcite , and then conduct a group of experiments to study the domain gap in knn-mt. ## our method to alleviate the influence of the domain gap on the datastore retrieval of knn-mt, we propose a simple yet effective approach to directly revise the original datastore, of which revised key representations are required to satisfy two properties: 1) they are more in line with the ideal distribution of downstream-domain key representation; 2) they can effectively retain the translation knowledge learned by the upstream nmt model. to this end, we design a key representation reviser to revise the key representations of the original datastore. to train this reviser, we first identify some key-queries pairs from the original datastore and upstream nmt model as the training data, where each key is expected to be semantically related to its corresponding queries. then, we propose two training losses to jointly train the reviser. using the reviser to reconstruct the original datastore, the original datastore can also effectively capture the semantically related key-queries pairs contained in the downstream datastore and nmt model, and thus is more suitable for the downstream translation task. ## experiments to investigate the effectiveness of our method, we conduct experiments in the task of nmt domain adaptation. ## related work our related work mainly includes two aspects: domain adaptation for nmt, and non-parametric retrieval-augmented approaches for nmt. @xcite , dominant methods in this aspect can be roughly divided into two categories: 1) model-centric approaches that focus on carefully designing nmt model architecture to learn targetdomain translation knowledge @xcite @xcite , or refining the training procedures to better exploit context @xcite @xcite ; 2) data-centric methods resorting to leveraging the target-domain monolingual corpus @xcite , synthetic corpus @xcite @xcite or parallel corpus @xcite to improve the nmt model via fine-tuning. ## conclusion in this paper, we first conduct a preliminary study to investigate the impact of the domain gap on the datastore retrieval of knn-mt. furthermore, we propose a reviser to refine the key representations of the original knn-mt datastore in an offline manner, making them more suitable for the downstream domain. this reviser is trained on the collection of key-queries pairs, where the key of each pair is expected to be semantically related to its corresponding queries. particularly, we introduce two losses to train the reviser, ensuring that the revised key representations conform to the downstream domain while effectively retaining their original knowledge. through extensive experiments, we demonstrate the effectiveness of our method. besides, in-depth analyses reveal that: 1) the performance improvement achieved by our method is positively correlated with the degree of the domain gap; 2) this improvement is primarily attributed to the enhancement of the datastore quality; 3) our method is able to compatible with existing adaptive knn-mt. to further verify the generalization of our method, we will extend our method to knn-lm or other text generation tasks, such as controllable generation.