Chunk-based Nearest Neighbor Machine Translation (2022)

Abstract: Semi-parametric models, which augment generation with retrieval, have led to impressive results in language modeling and machine translation, due to their ability to retrieve fine-grained information from a datastore of examples. One of the most prominent approaches, kNN-MT, exhibits strong domain adaptation capabilities by retrieving tokens from domain-specific datastores (Khandelwal et al., 2021). However, kNN-MT requires an expensive retrieval operation for every single generated token, leading to a very low decoding speed (around 8 times slower than a parametric model). In this paper, we introduce a chunk-based kNN-MT model which retrieves chunks of tokens from the datastore, instead of a single token. We propose several strategies for incorporating the retrieved chunks into the generation process, and for selecting the steps at which the model needs to search for neighbors in the datastore. Experiments on machine translation in two settings, static and “on-the-fly” domain adaptation, show that the chunk-based kNN-MT model leads to significant speed-ups (up to 4 times) with only a small drop in translation quality.

## introduction machine translation has seen remarkable advances due to increasingly powerful neural models @xcite @xcite . most deployed systems are fully-parametric (the training data is fully compressed into the parameters of a neural model), but they often struggle when translating rare words or out-of-domain sentences @xcite , commonly requiring several stages of finetuning to adapt to data drift or to new domains. recently, semi-parametric methods have shown great promise, by combining the strengths of parametric models with external databases of parallel sentences, such as translation memories @xcite @xcite @xcite @xcite . one of the most prominent semi-parametric models for machine translation is the k-nearest neighbor machine translation model (knn-mt) @xcite , which has led to impressive results, particularly in domain adaptation settings, without requiring fine-tuning. the knn-mt model constructs domain-specific datastores of parallel sentences and, at inference time, it retrieves similar examples from these datastores, which are used to improve the generation process, through the interpolation of probability distributions. however, knn-mt only retrieves single tokens-this is inefficient, since the model needs to consult the datastore at every generation step, an expensive operation. consequently, its decoding speed is around 8 times slower than that of a fully-parametric model. recent work has introduced several techniques to speed up knn-mt. @xcite proposed fast knn-mt, which constructs a different datastore for each example, by first searching for the nearest neighbors of the source tokens. @xcite introduced faster knn-mt, similar to fast knn-mt but with reduced memory requirements. @xcite proposed pruning the datastore, reducing the keys' representation size, and using a cache of retrieval distributions. however, despite leading to some decoding speedups, these methods are limited as they still retrieve a single token at each time step. in this paper, we propose a simple and efficient chunk-based knn-mt model. inspired by retro @xcite , the chunk-based knn-mt model retrieves chunks of tokens, instead of single tokens. but, similarly to knn-mt and unlike retro, it does not require any training or finetuning of the parametric component: it simply uses a combination of caching and interpolation of probability distributions to incorporate the retrieved tokens. by doing this, the model leads to a similar translation quality while having to search for neighbors in the datastore less often. this leads to decoding speeds up to 4 times faster than the ones achieved using the vanilla knn-mt model and only twice as slow as a fully-parametric model, but with considerably higher translation quality. in sum, our main contributions are: ## background in machine translation, a model is given a sentence or document in a source language, @xmath0, . . . , x l ], and the goal is to output a translation in a target language, @xmath1, . . . , y n ]. this is commonly done using a parametric sequence-tosequence model @xcite @xcite , in which the encoder receives the source sentence as input and outputs a set of hidden states. then, at each step t, the decoder attends to these hidden states and outputs a probability distribution over the vocabulary, p nmt (y t |@xmath2, x). finally, these probability distributions are used in a search procedure to generate the translation, typically using beam search @xcite . ## chunk-based knn-mt we now describe our chunk-based knn-mt model, illustrated in figure 1 . we first describe the datastore creation ( §3.1) and how we can retrieve chunks of tokens from it ( §3.2), describing how they are used by the model ( §3.2.2). finally, we describe how to select the steps in which the model performs retrieval ( §3.3). ## experiments to understand if the chunk-based knn-mt is able to maintain the translation quality while speedingup decoding, we performed experiments on domain adaptation ( §4.1) and on-the-fly adaptation ( §4.2).