Memory-efficient NLLB -200: Language-specific Expert Pruning of a Massively Multilingual Machine Translation Model (2023)

Abstract: The recently released NLLB-200 is a set of multilingual Neural Machine Translation models that cover 202 languages. The largest model is based on a Mixture of Experts architecture and achieves SoTA results across many language pairs. It contains 54.5B parameters and requires at least four 32GB GPUs just for inference. In this work, we propose a pruning method that enables the removal of up to 80% of experts without further finetuning and with a negligible loss in translation quality, which makes it feasible to run the model on a single 32GB GPU. Further analysis suggests that our pruning metrics can identify language-specific experts.

## introduction the transformer @xcite has become the dominant modeling paradigm in natural language processing tasks. many subsequent advances in the field came from increasing the computational budget, training data, and model size. neural machine translation was not an exception, where massively multilingual nmt @xcite @xcite demonstrated promising results, while attempting to overcome the curse of multilinguality @xcite ) by scaling up model size. however, increasing the parameter size exacerbates the cost of training @xcite @xcite and hurts the memory footprint and inference latency @xcite @xcite . sparselygated mixture-of-experts (moe) models are an efficient alternative to dense models @xcite @xcite . for example, @xcite demonstrates that an moe language model results in a 7x larger model compared to gpt-3, but requires only 30% of its * work done during an internship at naver labs europe energy for training and half of its flops at inference. mixture-of-experts models are neural networks whose set of parameters is partitioned into experts. contrary to dense models, where all network parameters are used for every input, an moe model activates different parts of the network, the experts, depending on the input, which is typically done by a gating mechanism at the token level. moe models are computationally efficient due to expert parallelism @xcite across a large number of gpus, by having each gpu hold a subset of all experts and communicate with the other gpus when it needs expert outputs for its local batch. in nllb-200 foot_0 @xcite , a load balancing regularizer in the objective function @xcite promotes equal distribution of the tokens across experts. this encourages the model to use all the experts and ensures that all gpus are used equally for the sake of computational efficiency. however, considering a large number of experts, it does not guarantee that all experts will be equally activated for a particular pair of languages at inference. it raises a research question: are there language-specific experts in multilingual moe models? if this is the case, we may be able to prune such models without loss of translation quality for the language pairs of our interest. reducing memory usage would be useful for a model like nllb-200, which normally requires at least four 32gb gpus at inference. in this work, we define metrics to assess the importance of each expert and prune the least important experts at inference. we aim to avoid finetuning because of its computational cost. in an ideal scenario, we would like to be able to identify the important experts in an moe model so that practitioners can deploy large models, such as nllb-200, on a single gpu. we summarize our main contributions as follows: the concept of mixture-of-experts models in machine learning dates back to the works of @xcite ; jordan and @xcite . most recent versions were inspired by @xcite , who achieved state-of-the-art language modeling and translation results with the largest model at that time. combined with the transformer model, moe models grew in popularity @xcite . beyond natural language processing, moe models showed a large success in computer vision @xcite , speech recognition @xcite , multi-modal learning @xcite , and diffusion models @xcite to name a few. for a more detailed survey of moe models, we refer readers to @xcite and @xcite . despite the recent successes, large moe models require a lot of memory and the contribution (or roles) of experts is under-explored. @xcite showed that the contributions of experts of a pre-trained moe model in different tasks such as mnli, cola, and squad are quite different. moreover, they converted a large sparse moe model pre-trained on a general task to a singleexpert dense model by fine-tuning the most 'professional' expert and dropping the other experts. it demonstrates that experts do not contribute equally to the performance and some are more important than others. @xcite also studied dif-ferent expert specializations such as sentinel tokens, punctuation, conjunctions and articles, and even languages. they concluded that experts in the encoder exhibit specialization, in contrast to the decoder, but not by language. according to the authors, their mechanism of token routing and load balancing prevents language specialization. @xcite train study routing mechanisms at different levels of granularity and show that task-level experts (i.e., per language) can achieve similar performance as token-level experts. however, this work assumes that the model is trained this way, while our own work attempts to prune an existing token-level moe model at inference without re-training it. there have been a number of attempts to compress existing massively multilingual nmt models @xcite . however, to the best of our knowledge, none of them explicitly studied expert pruning and the emergence of language-specific experts in a large moe model like we do. there has been a related line of works on pruning attention heads in transformer models @xcite , demonstrating linguistically-interpretable roles of attention heads @xcite and the emergence of languagespecific attention heads @xcite . understanding the role of attention heads helps carefully remove the least important ones without damage to translation quality. closest to our work, @xcite tried to prune a machine translation moe model by keeping the most activated experts, 3 but did not manage to preserve performance without further fine-tuning. even though it has been shown that multilingual nmt models benefit from a larger number of experts @xcite , to the best of our knowledge, our work is the first to study whether any language-specific experts emerge in a massively multilingual mixture-of-expert model for nmt, and how can redundant (or non-relevant) experts be pruned. ## our approach we experiment with different experts' pruning metrics and strategies that allow us to select the most relevant experts per language or language pair, and thus significantly reduce the memory usage at inference time of nllb-200. ## conclusion in this paper, we study expert pruning in the nllb-200 mixture-of-experts mt model. we propose expert pruning metrics based on gate statistics collected while decoding. we study several pruning strategies and demonstrate that it is possible to prune up to 80% of experts with a negligible loss in performance, which makes it possible to decode on a single 32gb gpu. we compare pruning at three levels of granularity: per language direction, per language, or global. language-specific and language-pair pruning perform the same but the former is the most convenient. global pruning (i.e., pruning always the same experts regardless of the source and target languages) performs surprisingly well but worse than language-specific pruning, which suggests that there are indeed some language-specific experts. this latter hypothesis is confirmed by our analysis of the selected experts.