PaperlayerView exact revision on arXiv

Source-linked AI summary

Nearest Neighbor Machine Translation

Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, Mike Lewis

arXiv:2010.00710v2cs.CL

TL;DR

Neural MT models have limited direct access to large, adaptable example stores during decoding. kNN-MT adds nearest-neighbor retrieval over cached source-and-target contexts without further training, and improves translation across in-domain, out-of-domain, and multilingual settings. Its main practical limitation is substantially slower inference with billion-item datastores.

  • Problem

    Neural MT needs a way to use large amounts of cached data at test time while remaining adaptable and interpretable.

  • Method

    kNN-MT interpolates a pretrained neural MT distribution with a nearest-neighbor distribution retrieved from cached translation contexts using source and target context.

  • Results

    kNN-MT significantly improves performance across in-domain, out-of-domain, and multilingual evaluations, including 1.5 BLEU on a state-of-the-art German-English model and 9.2 BLEU average out-of-domain gains.

  • Takeaways & Limitations

    A single pretrained model can be adapted to diverse domains or specialized for language pairs by changing the datastore, without additional training.

  • Takeaways & Limitations

    Retrieving 64 keys from a datastore containing billions of items makes generation two orders of magnitude slower than the base MT system.

Abstract

from arXiv · show

We introduce $k$-nearest-neighbor machine translation ($k$NN-MT), which predicts tokens with a nearest neighbor classifier over a large datastore of cached examples, using representations from a neural translation model for similarity search. This approach requires no additional training and scales to give the decoder direct access to billions of examples at test time, resulting in a highly expressive model that consistently improves performance across many settings. Simply adding nearest neighbor search improves a state-of-the-art German-English translation model by 1.5 BLEU. $k$NN-MT allows a single model to be adapted to diverse domains by using a domain-specific datastore, improving results by an average of 9.2 BLEU over zero-shot transfer, and achieving new state-of-the-art results -- without training on these domains. A massively multilingual model can also be specialized for particular language pairs, with improvements of 3 BLEU for translating from English into German and Chinese. Qualitatively, $k$NN-MT is easily interpretable; it combines source and target context to retrieve highly relevant examples.

1 INTRODUCTION

kNN-MT adds nearest-neighbor retrieval to pretrained neural translation models without further training, using cached translation contexts to improve performance and enable domain and language-pair adaptation.

  • Method: kNN-MT interpolates a neural MT softmax distribution with a nearest-neighbor distribution over cached translation contexts.The cache contains complete source and target-prefix contexts indexed by base-model hidden states.
  • Properties: No further training is required, while the decoder can access billions of key-value pairs during inference.This makes the approach simpler than methods requiring training and more expressive than prior approaches with less inference-time data access.
  • Results: 1.5 BLEU improves a state-of-the-art German-English translation model after adding nearest-neighbor search.
  • Results: 9.2 BLEU average gains over the base model out-of-domain come from adding domain-specific datastores, which also outperform models trained on those domains.
  • Results: 3 BLEU improvements specialize a multilingual model for English-to-German and English-to-Chinese translation.

2 NEAREST NEIGHBOR MACHINE TRANSLATION

kNN-MT augments autoregressive decoding with nearest-neighbor retrieval over source-and-target translation contexts, converts retrieved targets into a distribution, and interpolates it with the base model.

  • Generation: At each generation step, kNN-MT retrieves similar contexts from a datastore and computes a distribution over their corresponding target tokens.The query representation conditions on the source input and previously generated target tokens.
  • Datastore: The offline datastore stores decoder representations of complete translation contexts as keys and corresponding ground-truth target tokens as values.Source-language tokens are encoded implicitly in the keys rather than stored as values.
  • Retrieval distribution: Nearest neighbors are selected by squared-L2 distance, then their target tokens are converted into a vocabulary distribution using a temperature-scaled softmax.A temperature greater than one flattens the distribution and reduces overfitting to the closest retrievals.
  • Interpolation: The retrieved distribution is interpolated with the pretrained MT distribution using a tuned parameter λ, improving robustness when cached examples are irrelevant.
  • Relation to kNN-LM: kNN-MT generalizes nearest-neighbor language modeling by conditioning representations on both source and target context.Its additional softmax temperature controls distribution flatness and retrieval diversity.

3 EXPERIMENTAL SETUP

Experiments evaluate kNN-MT for single-pair translation, multilingual MT, and domain adaptation using German-English and multilingual Transformer models across several datasets.

  • Evaluation settings: Three settings are evaluated: single language-pair translation, multilingual MT, and domain adaptation.
  • Datasets: The domain-adaptation dataset contains German-English parallel data from five domains: Medical, Law, IT, Koran, and Subtitles.
  • Models: The single-pair and domain-adaptation experiments use the WMT’19 German-English translation winner, a Transformer encoder-decoder model.The model was additionally trained on over 10 billion backtranslation tokens and fine-tuned on earlier newstest sets.
  • Models and data: The multilingual model is a 418M-parameter Transformer trained for 100K updates on CCMatrix, which contains 79 languages and 1,546 language pairs.
  • Computational cost: Inference with billions-item datastores is two orders of magnitude slower than the base MT system, though speed can be improved with smaller searches or datastores.These changes involve relatively minor performance trade-offs according to the experiments.

4 EXPERIMENTS

Experiments show that kNN-MT improves translation across single-language, multilingual, cross-lingual, and domain-adaptation settings without additional training. Gains are generally larger with relevant datastore data, while datastore size introduces a speed–performance trade-off.

  • 4.1 SINGLE LANGUAGE-PAIR TRANSLATION: 1.5 BLEU improves a strong German-English model using a 770M-token training datastore, raising performance from 37.59 to 39.08 without additional training.The result demonstrates improvement even for a heavily tuned baseline.
  • 4.2 MULTILINGUAL MACHINE TRANSLATION: 1.4 BLEU is the average improvement across 17 language pairs, with up to 3 BLEU for English-German, English-Chinese, and Chinese-English.All gains use language-pair-specific training examples without additional training.
  • 4.2 MULTILINGUAL MACHINE TRANSLATION: 1 BLEU is the average improvement when English-source datastores support translation from other languages, making cross-lingual retrieval effective.The result is reported across five Ted Talks pairs and three newstest2019 pairs.
  • 4.3 DOMAIN ADAPTATION: 9.2 BLEU is the average domain-adaptation gain from a domain-specific datastore, with improvements up to 16 BLEU on Law and 14.5 BLEU on Medical without further training.kNN-MT also exceeds the best reported multi-domain average by 1.4 BLEU.
  • 4.3 DOMAIN ADAPTATION: Domain-specific datastores outperform large arbitrary out-of-domain datastores, while larger datastores improve performance but slow retrieval.Datastore size and BLEU gains are only weakly correlated across languages, and parallel-data quality may also affect improvements.

5 TUNING KNN-MT

The tuning analysis examines how temperature, neighbor count, and datastore size affect validation BLEU, revealing trade-offs between performance, robustness, and retrieval speed.

  • Softmax temperature: Temperatures above 1 flatten the neighbor distribution and help prevent overfitting to the most similar neighbor.A temperature of 1 produces significantly lower BLEU, while 10 or 100 is optimal in the experiments.
  • Number of neighbors per query: 64 neighbors per query is used, but retrieving more neighbors does not consistently improve performance and can introduce noise.In some cases, retrieving fewer neighbors improves performance, suggesting that k can benefit from further tuning.
  • Datastore size: Increasing datastore size improves translation performance, but larger datastores slow retrieval.Much of the benefit can be achieved with smaller, faster datastores.

6 QUALITATIVE ANALYSIS

Qualitative inspection shows that kNN-MT retrieves contextually related translation examples and can reconstruct unseen named entities from subword evidence.

  • Retrieval relevance: Retrieved examples often share source-side phrases and target-side local contexts even when their longer contexts differ substantially.For the example involving “tamed the military,” retrieved phrases express related concepts such as authority, control, or force over the military.
  • Retrieval relevance: All retrieved neighbors in one example predict the target word military despite describing different nations and historical periods.The retrievals include related local phrases but differ sharply in their broader context.
  • Subword reconstruction: kNN reconstructs a named entity absent from the datastore by combining neighbors for its separate subword tokens.For Haysom, neighbors containing Hayes or Haydn support Hay, while Grissom or Folsom support som.

7 RELATED WORK

Related work spans retrieval for translation, broader text generation, and domain adaptation, while kNN-MT uses retrieved distributions directly without additional training.

  • Retrieval in Translation: Earlier translation systems retrieved words, phrases, translation histories, or similar examples to guide generation.The cited approaches include n-gram retrieval, cache-based histories, edit-distance matching, local trigram contexts, phrase memories, and syntactic features.
  • Illustrative retrieval: The qualitative example retrieves source and target contexts related to control over the military, producing the target word military.The three nearest retrievals are semantically related to the source context.
  • Domain adaptation: Prior domain-adaptation methods use lexicons, domain-adaptive training, or source-side neighbors, and require additional training.kNN-MT keeps the model fixed while allowing different datastores during decoding in different domains.
  • Retrieval in Text Generation: Other generation systems retrieve dialogue examples, question-answering contexts, or structured-output examples, often concatenating retrievals or jointly training retrievers and editors.kNN-MT instead converts retrieval results directly into a vocabulary distribution for generation.

8 CONCLUSION

The paper presents kNN-MT as a simple, training-free augmentation for neural MT that improves translation and supports domain and language-pair specialization.

  • Conclusion: kNN-MT applies nearest-neighbor retrieval to any neural MT model without further training.It interpolates the base model with a nearest-neighbor classifier because similar embedding-space contexts tend to have similar next words.
  • Conclusion: The approach improves a state-of-the-art model in-domain, produces large gains out-of-domain, and specializes multilingual models for specific language pairs.These conclusions summarize the reported uses of datastore-based nearest-neighbor augmentation.
  • Conclusion: Future work should improve retrieval efficiency, including by down-sampling frequent target words in the datastore.This recommendation addresses the efficiency of the datastore-based approach.

A HYPERPARAMETER TUNING

The experiments tune two hyperparameters for multilingual and domain-adaptation settings while fixing several search and decoding choices. Validation results indicate that kNN-MT with λ = 0.5 performs similarly to or better than the base model across 17 language pairs.

  • A HYPERPARAMETER TUNING: Only the interpolation parameter λ and softmax temperature T were tuned on validation sets.The number of neighbors k was fixed to 64, searched clusters to 32, and beam size to 5.
  • A HYPERPARAMETER TUNING: λ = 0.5 makes kNN-MT perform similarly to or better than the base MT model across 17 evaluated language pairs.The validation comparison reports no performance decrease at this fixed interpolation value.
  • A HYPERPARAMETER TUNING: Values of 10 or 100 were optimal for the softmax temperature T in all experiments.The temperature was tuned alongside the interpolation parameter on validation sets.
  • A HYPERPARAMETER TUNING: Figure 5 shows neighbors sharing the changed target token while matching local source and target contexts but differing in global context about Papua.The example illustrates the role of local context in retrieval.

B ADDITIONAL EXAMPLES

Additional examples illustrate that kNN-MT retrieves neighbors using local source and target context. When target-side context is very short, source context can resolve ambiguity and support the correct target token.

  • B ADDITIONAL EXAMPLES: kNN-MT uses local contexts on both the source and target sides to retrieve nearest neighbors.The example is described as illustrating retrieval behavior based on both contextual sides.
  • B ADDITIONAL EXAMPLES: With very little target-side prior context, kNN-MT mainly relies on source context to retrieve the best neighbors.This behavior is shown in an example where the available target context is ambiguous.

C BLEU SCORES

The paper reports BLEU using case-sensitive detokenized SACREBLEU measurements and discusses statistical-power checks for translation experiments. The examples include multilingual validation outputs and short generated-token sequences.

  • C BLEU SCORES: Results use case-sensitive detokenized BLEU measured with SACREBLEU.The paper specifies separate tokenization signatures for general, Chinese, and Japanese evaluation.
  • C BLEU SCORES: Statistical-power assumptions were checked for Chinese-English and English-Chinese experiments.The reported values are P0 = 0.13 and b0 = 12 for Chinese-English, and P0 = 0.07 and b0 = 16 for English-Chinese.
  • C BLEU SCORES: One test example begins with the generated token sequence “We.”The input is “Wir werden das Beste tun, mit dem, was wir haben.”
  • C BLEU SCORES: Figure 6 presents an ambiguous target-side context where source context helps generate the correct token, “will.”The example demonstrates source-context use when target-side context is short.
Loading 2010.00710v2…