Fast Vocabulary Projection Method via Clustering for Multilingual Machine Translation on GPU (2022)

Abstract: Multilingual Neural Machine Translation has been showing great success using transformer models. Deploying these models is challenging because they usually require large vocabulary (vocab) sizes for various languages. This limits the speed of predicting the output tokens in the last vocab projection layer. To alleviate these challenges, this paper proposes a fast vocabulary projection method via clustering which can be used for multilingual transformers on GPUs. First, we offline split the vocab search space into disjoint clusters given the hidden context vector of the decoder output, which results in much smaller vocab columns for vocab projection. Second, at inference time, the proposed method predicts the clusters and candidate active tokens for hidden context vectors at the vocab projection. This paper also includes analysis of different ways of building these clusters in multilingual settings. Our results show end-to-end speed gains in float16 GPU inference up to 25% while maintaining the BLEU score and slightly increasing memory cost. The proposed method speeds up the vocab projection step itself by up to 2.6x. We also conduct an extensive human evaluation to verify the proposed method preserves the quality of the translations from the original model.

## introduction neural machine translation (nmt) has witnessed significant advances by the introduction of the transformer model @xcite where excellent performance has been shown for bilingual translation initially, mainly to and from english. later, the model has been extended to multiple language pairs e.g. @xcite and is referred to as multilingual neural machine translation (mnmt). the multilingual model usually comes with a significant increase in the number of parameters and the vocabulary (vocab) size to accommodate various languages. this work focuses on implementing efficient vocab projection in transformer models. the problem becomes very important for mnmt. while bilingual nmt typically use vocab of around 32k sub-words, mnmt has significantly larger vocab size to adequately cover all the languages. for example, the recent mnmt submission from microsoft to the wmt21 shared task that covers 100+ languages has a vocab size of 250k sub-words @xcite . the increased vocab makes the projection very compute-intensive and it can take up to 25% of the total computation in our internal benchmark. hence, the need to speedup this operation becomes even more important. fast vocab projection is well-studied for nmt and natural language processing (nlp). in @xcite , the authors propose two methods to reduce the effective vocab for nmt. the first formulates the problem as a nearest neighbor search and uses locality sensitive hashing (lsh) to speedup the computation while the second uses alignment information to select a subset of the vocab based on the input. it is found that the alignment-based method leads to around 2x speedup. the work @xcite proposes an efficient screening model to exploit the clustering structure of context features right before vocab projection. the authors formulate a joint optimization problem to learn the clusters and their corresponding vocab subspace. in the same paper, a thorough comparison to the existing literature is done including graph-based nearest-neighbor search @xcite , svd approximation @xcite , a modified version of hierarchical softmax @xcite and locality sensitive hashing (lsh) for maximum inner product search (mips) @xcite . it is shown that the clustering solution can outperform the latter techniques in terms of accuracy-speed trade-off in bilingual lstms on cpu. this paper extends the clustering based approach from @xcite to mnmt for both dense and sparse transformer models on gpu. to simplify integration and accessibility in the mnmt large-scale model settings, we propose to use kmeans clustering to split the vocab search space into disjoint clusters given the hidden context vector of the decoder output. the multilingual extension comprises experimenting with different ways to build the cluster maps as well as testing different configurations such as number of centroids and corresponding vocab subspace to make the method efficiently work at large scale. in @xcite , it is shown how to extend lsh to run efficiently within beam search on gpu but here we use and develop optimized kernels. these kernels are developed for fastertransformer, a highly optimized transformer-based encoder and decoder implementation offered by nvidia foot_1 . serving gpu inference is essential because gpus give several orders of magnitude speedups relative to cpus for large transformer models foot_2 . experimental results show end-to-end speed gains in float16 gpu inference up to 25% while maintaining the bleu score and slightly increasing memory cost. the proposed method speeds up the vocab projection step itself by up to 2.6x. we also conduct an extensive human evaluation to verify the proposed method preserves the quality of the translations from the original model. the paper is organized as follows. section 2 describes the proposed method focusing on the multilingual aspect and gpu implementation. this is followed by experimental results for both dense and sparse transformer models and for 6 translation directions in section 3. finally, section 4 concludes the paper. ## experimental results in this section, we present the evaluation results for the proposed vocab projection method in a multi-lingual transformer setting on gpu. we first introduce the experimental setup, then assess the impact of the proposed method on vocab projection itself as well as show the speed and accuracy performance results. we conducted an extensive human evaluation to verify the translation quality is preserved. finally, we study how different configurations such as the number of centroids and k likely tokens affect the performance of our method. ## conclusion this paper proposes a fast vocab projection method via clustering for multilingual neural machine translation with large vocab sizes and practically applied models on gpu. the method splits the large vocab search space into smaller subspaces to run efficient gpu inference. results reveal end-to-end speed improvements up to 25% while maintaining the bleu, and up to 2.6x speed improvement on the vocab projection step. we conducted human evaluations to verify the translation quality. in the future, we wish to explore our method in an x-y machine translation scenario.