k NN - BOX : A Unified Framework for Nearest Neighbor Generation (2024)

Abstract: Augmenting the base neural model with a token-level symbolic datastore is a novel generation paradigm and has achieved promising results in machine translation (MT). In this paper, we introduce a unified framework kNN-BOX, which enables quick development and visualization for this novel paradigm. kNN-BOX decomposes the datastore-augmentation approach into three modules: datastore, retriever and combiner, thus putting diverse kNN generation methods into a unified way. Currently, kNN-BOX has provided implementation of seven popular kNN-MT variants, covering research from performance enhancement to efficiency optimization. It is easy for users to reproduce these existing work or customize their own models. Besides, users can interact with their kNN generation systems with kNN-BOX to better understand the underlying inference process in a visualized way. In experiment section, we apply kNN-BOX for machine translation and three other seq2seq generation tasks (text simplification, paraphrase generation and question generation). Experiment results show that augmenting the base neural model with kNN-BOX can bring large performance improvement in all these tasks. The code and document of kNN-BOX is available at https://github.com/NJUNLP/knn-box. The demo can be accessed at http://nlp.nju.edu.cn/demo/knn-box/. The introduction video is available at https://www.youtube.com/watch?v=m0eJldHVR3w.

## introduction equipping the base neural model with a symbolic datastore is a novel paradigm for enhancing generation quality. @xcite apply this paradigm in machine translation, known as knn-mt, and achieves promising results, especially in mt domain adaptation and multilingual mt. af-* equal contributions. terwards, the following work keep optimizing this approach, making it a more mature methodology, e.g., dynamically deciding the usage of retrieval results @xcite , building a light and explainable datastore @xcite , injecting knn knowledge into the neural model @xcite . however, we notice that these knn generation methods are implemented with diverse codebases, e.g., fairseq foot_0 , transformers foot_1 and joeynmt foot_2 , which hinders comparison between these methods and potential fusion of latest research advances. interpretability is another interesting point in knn research, as the community is curious why knn generation works and whether it is reliable. in this paper, we introduce a unified framework knn-box for nearest neighbor generation, which supports quick development and visualization anal-ysis. our framework decomposes the datastoreaugmentation approach into three modules: data-store, retriever and combiner, thus putting diverse knn generation methods into a unified way (figure 1 ). up till now, knn-box has released implementation of seven popular knn-mt models, covering research from performance enhancement @xcite @xcite to efficiency optimization @xcite @xcite , which can help users to quickly reproduce existing works. moreover, users can easily fuse advanced models with knn-box, for example, jointly using a better combiner and a lighter datastore, to achieve the best of both worlds. another useful feature of knn-box is supporting visualized interactive analysis. via our provided web service, users can interact with their knn model and observe its inference process, e.g. the content and distribution of its retrieval results (figure 3 ). we hope knn-box can help the community to better understand the interpretability of knn generation. experiment results on machine translation datasets show that knn-box is a reliable platform for model reproduction and development. in addition, we apply knn-box for three other seq2seq tasks, i.e., text simplification, paraphrase generation and question generation. experiment results show that augmenting the base neural model with knn-box is also beneficial in these tasks, showing the great potential of nearest neighbor generation and the wide usage of our knn-box toolkit. at the time of writing, we are happy to see that knn-box has been used as the backbone of this year's acl paper @xcite and emnlp papers @xcite , and we hope this toolkit to support more valuable research in the future. memorizing knowledge into datastore to extract translation knowledge, translation pair (x , y) is fed into m for teacher-forcing decoding. at time step t, the continuous representation of the translation context (x , @xmath0, i.e. the hidden state h t from the last decoder layer, is taken as key: and the target token y t is taken as value. each key-value pair explicitly memorizes the translation knowledge: generating the value token at the decoder hidden state key. with a single forward pass over the entire corpus, the full datastore d can be constructed: generating with memorized knowledge the constructed datastore is then combined with the base nmt model as an augmentation memory. during inference, the nmt model retrieves related knowledge from the datastore to adjust its own translation prediction. specifically, the nmt model uses the contextualized representation of the test translation context (x , @xmath1, y j )} k @xmath2. the retrieved entries are then converted to a distribution over the vocabulary: here, d denotes l 2 -square distance and t is the temperature. in the end, the output distribution of the nmt model and symbolic datastore are interpolated with the weight λ: recent advances in knn-mt to make knn-mt more effective, efficient and explainable, various methods have been devised. @xcite and @xcite propose to dynamically decide the usage of retrieval results to exclude potential noise in nearest neighbors. @xcite explore the setting of multi-domain adaptation and remedy the catastrophic forgetting problem. inspired by @xcite , @xcite introduce three ways to improve the efficiency of knn-mt, i.e. dimension reduction, datastore pruning and adaptive retrieval. later, @xcite propose to reduce dimension and prune datastore with a learnable network. recently, @xcite explore the interpretability issue in knn-mt and builds a light and more explainable datastore according to the capability of the nmt model. ## unified framework: knn-box this section describes how we design and implement knn-box, and introduce how users run knn-box for developing knn generation models and interacting with the deployed model visually. ## experiments to evaluate the effectiveness of knn-box, we conduct experiments on machine translation and three other seq2seq tasks.