knn-seq: Efficient, Extensible kNN-MT Framework

TL;DR

knn-seq is an efficient, extensible framework for kNN-based machine translation that handles billion-scale datastores with reduced computational costs, maintaining translation quality.

Contribution

It introduces a scalable, plug-in compatible kNN-MT framework that significantly reduces construction time for large datastores while preserving translation performance.

Findings

Achieves comparable translation gains to original kNN-MT

Constructs billion-scale datastores in 2.21 hours

Runs efficiently with large-scale data

Abstract

k-nearest-neighbor machine translation (kNN-MT) boosts the translation quality of a pre-trained neural machine translation (NMT) model by utilizing translation examples during decoding. Translation examples are stored in a vector database, called a datastore, which contains one entry for each target token from the parallel data it is made from. Due to its size, it is computationally expensive both to construct and to retrieve examples from the datastore. In this paper, we present an efficient and extensible kNN-MT framework, knn-seq, for researchers and developers that is carefully designed to run efficiently, even with a billion-scale large datastore. knn-seq is developed as a plug-in on fairseq and easy to switch models and kNN indexes. Experimental results show that our implemented kNN-MT achieves a comparable gain to the original kNN-MT, and the billion-scale datastore construction took 2.21 hours in the WMT'19 German-to-English translation task. We publish our knn-seq as an MIT-licensed open-source project and the code is available on https://github.com/naist-nlp/knn-seq . The demo video is available on https://youtu.be/zTDzEOq80m0 .