INK: Injecting kNN Knowledge in Nearest Neighbor Machine Translation

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 15948–15959 July 9-14, 2023 ©2023 Association for Computational Linguistics INK: Injecting kNN Knowledge in Nearest Neighbor Machine Translation Wenhao Zhu1 , Jingjing Xu2 , Shujian Huang1 , Lingpeng Kong3, Jiajun Chen1 1 National Key Laboratory for Novel Software Technology, Nanjing University 2 Shanghai AI Laboratory 3 The University of Hong Kong zhuwh@smail.nju.edu.cn, jingjingxu@pku.edu.cn huangsj@nju.edu.cn, lpk@cs.hku.hk, chenjj@nju.edu.cn Abstract Neural machine translation has achieved promising results on many translation tasks. However, previous studies have shown that neu- ral models induce a non-smooth representation space, which harms its generalization results. Recently, kNN-MT has provided an effective paradigm to smooth the prediction based on neighbor representations during inference. De- spite promising results, kNN-MT usually re- quires large inference overhead. We propose an effective training framework INK to directly smooth the representation space via adjusting representations of kNN neighbors with a small number of new parameters. The new parame- ters are then used to refresh the whole repre- sentation datastore to get new kNN knowledge asynchronously. This loop keeps running un- til convergence. Experiments on four bench- mark datasets show that INK achieves average gains of 1.99 COMET and 1.0 BLEU, outper- forming the state-of-the-art kNN-MT system with 0.02× memory space and 1.9× inference speedup1. 1 Introduction Neural machine translation (NMT) have achieved promising results in recent years (Vaswani et al., 2017; Ng et al., 2019; Qian et al., 2021b). The target of NMT is to learn a generalized representa- tion space to adapt to diverse scenarios. However, recent studies have shown that neural networks, such as BERT and GPT, induce non-smooth rep- resentation space, limiting the generalization abil- ities (Gao et al., 2018; Ethayarajh, 2019; Li et al., 2020). In NMT, we also observe a similar phe- nomenon in the learned representation space where low-frequency tokens disperse sparsely, even for a strong NMT model (More details are described in Section Experiments). Due to the sparsity, many 1Code will be released at https://github.com/ OwenNJU/INK Figure 1: The overview of our training loop. We refine the representation space of an NMT model according to the extracted kNN knowledge. The new parameters are then used to refresh the datastore to update kNN knowledge asynchronously. “holes” could be formed. When it is used to trans- late examples from an unseen domain, the perfor- mance drops sharply (Wang et al., 2022a,b) Recently, k-Nearest-Neighbor Machine Trans- lation (kNN-MT) (Khandelwal et al., 2021) pro- vides an effective solution to smooth predictions by equipping an NMT model with a key-value data- store. For each entry, the value is the target token and key is the contextualized representation at the target position. It requires a training set to record tokens and representations. By aggregating near- est neighbors during inference, the NMT model can achieve decent translation results (Khandel- wal et al., 2021; Zheng et al., 2021; Jiang et al., 2022). Despite the success, kNN-MT also brings new issues with the increasing scale of training data. Retrieving …

Content truncated to 500 words