Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Faster k-Nearest-Neighbor Machine Translation

Xiangyu Shi, Yunlong Liang, Jinan Xu, Yufeng Chen

TL;DR

The paper addresses the high retrieval overhead of kNN-MT in cross-domain translation by introducing a lightweight MLP-based selector that predicts per-token whether a kNN lookup is necessary. The selector is trained with a dual objective—a retrieval-decision loss and a translation-guided loss—and integrated via a Gumbel-softmax mechanism to maintain differentiability. Empirical results across multiple benchmarks show substantial reductions in kNN overhead (up to 53%) with only modest declines in translation quality, and the method complements existing kNN-MT variants within the kNN-BOX framework. This approach offers a simple, generalizable path to faster, retrieval-augmented MT without requiring changes to the underlying NMT or datastore structure, enabling practical deployment in cross-domain scenarios.

Abstract

Recent works have proven the effectiveness of k-nearest-neighbor machine translation(a.k.a kNN-MT) approaches to produce remarkable improvement in cross-domain translations. However, these models suffer from heavy retrieve overhead on the entire datastore when decoding each token. We observe that during the decoding phase, about 67% to 84% of tokens are unvaried after searching over the corpus datastore, which means most of the tokens cause futile retrievals and introduce unnecessary computational costs by initiating k-nearest-neighbor searches. We consider this phenomenon is explainable in linguistics and propose a simple yet effective multi-layer perceptron (MLP) network to predict whether a token should be translated jointly by the neural machine translation model and probabilities produced by the kNN or just by the neural model. The results show that our method succeeds in reducing redundant retrieval operations and significantly reduces the overhead of kNN retrievals by up to 53% at the expense of a slight decline in translation quality. Moreover, our method could work together with all existing kNN-MT systems. This work has been accepted for publication in the jornal Advances in Artificial Intelligence and Machine Learning (ISSN: 2582-9793). The final published version can be found at DOI: https://dx.doi.org/10.54364/AAIML.2024.41111

Towards Faster k-Nearest-Neighbor Machine Translation

TL;DR

The paper addresses the high retrieval overhead of kNN-MT in cross-domain translation by introducing a lightweight MLP-based selector that predicts per-token whether a kNN lookup is necessary. The selector is trained with a dual objective—a retrieval-decision loss and a translation-guided loss—and integrated via a Gumbel-softmax mechanism to maintain differentiability. Empirical results across multiple benchmarks show substantial reductions in kNN overhead (up to 53%) with only modest declines in translation quality, and the method complements existing kNN-MT variants within the kNN-BOX framework. This approach offers a simple, generalizable path to faster, retrieval-augmented MT without requiring changes to the underlying NMT or datastore structure, enabling practical deployment in cross-domain scenarios.

Abstract

Recent works have proven the effectiveness of k-nearest-neighbor machine translation(a.k.a kNN-MT) approaches to produce remarkable improvement in cross-domain translations. However, these models suffer from heavy retrieve overhead on the entire datastore when decoding each token. We observe that during the decoding phase, about 67% to 84% of tokens are unvaried after searching over the corpus datastore, which means most of the tokens cause futile retrievals and introduce unnecessary computational costs by initiating k-nearest-neighbor searches. We consider this phenomenon is explainable in linguistics and propose a simple yet effective multi-layer perceptron (MLP) network to predict whether a token should be translated jointly by the neural machine translation model and probabilities produced by the kNN or just by the neural model. The results show that our method succeeds in reducing redundant retrieval operations and significantly reduces the overhead of kNN retrievals by up to 53% at the expense of a slight decline in translation quality. Moreover, our method could work together with all existing kNN-MT systems. This work has been accepted for publication in the jornal Advances in Artificial Intelligence and Machine Learning (ISSN: 2582-9793). The final published version can be found at DOI: https://dx.doi.org/10.54364/AAIML.2024.41111
Paper Structure (27 sections, 10 equations, 3 figures, 7 tables, 1 algorithm)

This paper contains 27 sections, 10 equations, 3 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Neural Machine Translation
  4. Domain Adaptation for NMT
  5. Retrieval-Argmented Text Generation
  6. $k$NN-MT
  7. Datastore
  8. Estimation of Probability Distribution
  9. Adaptive-kNN-MT
  10. PCK-kNN-MT
  11. Some Other $k$NN-MT Variants
  12. kNN-BOX
  13. Methodolgy
  14. Our Mehtod
  15. Training Procedures
  16. ...and 12 more sections

Figures (3)

  • Figure 1: Overview of our method. The selector predicts all the tokens to decide whether to revise the probability distribution. $k$NN retrievals only occur on parted tokens instead of all tokens, thus we speed up the $k$NN-MT systems.
  • Figure 2: Top-8 tokens which cause futile retrievals on multi-domain IT dataset. These tokens are always unchanged after $k$NN retrievals.
  • Figure 3: The training procedure of our method.