Gradient-based Gradual Pruning for Language-Specific Multilingual Neural Machine Translation (2023)

Abstract: Multilingual neural machine translation (MNMT) offers the convenience of translating between multiple languages with a single model. However, MNMT often suffers from performance degradation in high-resource languages compared to bilingual counterparts. This degradation is commonly attributed to parameter interference, which occurs when parameters are fully shared across all language pairs. In this work, to tackle this issue we propose a gradient-based gradual pruning technique for MNMT. Our approach aims to identify an optimal sub-network for each language pair within the multilingual model by leveraging gradient-based information as pruning criterion and gradually increasing the pruning ratio as schedule. Our approach allows for partial parameter sharing across language pairs to alleviate interference, and each pair preserves its unique parameters to capture language-specific information. Comprehensive experiments on IWSLT and WMT datasets show that our approach yields a notable performance gain on both datasets.

## introduction in recent years, neural machine translation (nmt) based on the transformer architecture has achieved great success and become the dominant paradigm for machine translation @xcite . multilingual neural machine translation (mnmt), which learns a unified model to translate between multiple languages, has attracted growing attention in the nmt area @xcite . the reasons are: 1) from a practical perspective, it significantly reduces the training and inference cost and simplifies deployment in production; 2) utilizing data from multiple language pairs simultaneously can potentially help improve the translation quality of low-resource or even zero-resource language pairs by transferring knowledge across languages. despite these benefits, mnmt remains challenging as its performance degrades compared to bilingual counterparts in high-resource languages @xcite . previous studies @xcite attribute this degradation to parameter interference: each language is unique and therefore has distinct requirements from the model parameters to capture language-specific information, and fully sharing parameters causes negative interactions across language pairs as languages compete for model capacity. although naively increasing model size may alleviate parameter interference, large models often suffer from parameter inefficiency and overfitting @xcite . to address the parameter interference issue more effectively, researchers have explored various methods to allocate language-specific parameters to capture unique linguistic information for each language pair. one line of work focused on introducing extra parameters to the model. for instance, adapter-based approaches @xcite @xcite inject lightweight language-specific modules into the shared model. another well-known method @xcite utilizes a shared encoder but a different decoder for each target language. despite being effective, these methods can become parameter-inefficient as the number of languages in the multilingual model grows. another line of work focused on extracting a separate sub-network within the multilingual model for each language pair. in these approaches, each sub-network preserves exclusive parameters to capture the language-specific features and has some parameters shared with other languages @xcite @xcite . a recent work @xcite shows promising results by first training a multilingual model that covers all language pairs, then fine-tuning the trained model on each pair, and in the end, extracting the sub-networks via magnitude pruning in one step. although this paradigm is intuitive and straightforward, it can be suboptimal due to the following limitations: (1) magnitude pruning is demonstrated to be ineffective in a transfer learning context @xcite , potentially affecting nmt systems in the multilingual scenario. this work aims to mitigate the parameter interference issue in mnmt without adding extra parameters and overcome the aforementioned two limitations in @xcite . we propose gradient-based gradual pruning for languagespecific mnmt. more specifically, a multilingual base model is first trained, and subsequently, the model is simultaneously fine-tuned and pruned on each language pair. conducting finetuning and pruning concurrently allows the model to adapt and optimize the parameters while reducing the model size. instead of the widely used magnitude scores, we opt for the gradient-based scores as the pruning criterion. the percentage of pruned weights is gradually increased from zero to the target level throughout the pruning process. by optimizing the pruning criterion and schedule, we strive to identify optimal sub-networks and limit the parameter interference. lastly, the resulting sub-networks are integrated into the mnmt model through a final training phase in a language-aware manner. a large set of experiments is conducted on iwslt and wmt datasets, showing that our approach leads to a substantial performance gain of 2.06 bleu on iwslt and 1.41 bleu on wmt. our contributions can be summarized as follows: ## related work standard multilingual neural machine translation systems translate between multiple languages with a unified model. the model can be jointly trained on multiple language pairs by prepending a special token to the source sentence, informing the model about the desired target language @xcite . although fully sharing parameters across languages and joint training can enhance knowledge transfer, mnmt suffers from parameter interference and the lack of language-specific parameters for capturing language-specific information, resulting in performance degradation, especially in high-resource language pairs. various previous works have explored the idea of partially sharing parameters across languages while allowing for language-specific parameters. @xcite investigates different parameter sharing strategies. @xcite compares several methods of designing language-specific attention module. @xcite studies when and where language-specific capacity matters. additionally, another widely recognized technique, adapter-based, has attracted substantial interest in recent years. adapter-based approaches @xcite @xcite inject additional lightweight languagespecific modules for different language pairs to capture language-specific information. although effective, these methods increase the parameters for each language pair and thus result in an inference speed decrease. on the contrary, our work introduces no extra parameters to the model and thus has negligible to no impact on the inference speed. for more details, see appendix e. to avoid the inference speed decrease, several works @xcite @xcite explore designing language-aware modules inside the mnmt model. this way, the model can capture languagespecific information without sacrificing the inference speed. for instance, @xcite employs a single shared encoder across all languages but a unique decoder for each target language. however, the model can suffer from the parameter exploration issue as the number of languages in the multilingual model increases. in contrast, our approach maintains a consistent number of total parameters, irrespective of the number of languages. recent works allocate language-specific parameters by extracting a unique sub-network for each language pair within the multilingual model, which avoids introducing additional parameters to the model. different approaches are explored for subnetwork extraction, such as taylor expansion @xcite , parameter differentiation @xcite , and model pruning @xcite . in this paper, we focus on model pruning because it showed to be more effective and results in promising performance. in machine learning, model pruning is widely used to remove redundant weights from a neural network while preserving important ones to maintain accuracy @xcite @xcite @xcite . two key components of model pruning are the pruning criterion, which determines the relative importance of weights and which weights to prune, and the pruning schedule, which defines the strategy how the target pruning ratio is achieved throughout the pruning process. in terms of pruning criteria, magnitude pruning @xcite , which removes weights with low absolute values, is the most widely used method for weight pruning. a recent work on large language models @xcite proposed movement pruning, which scores weights by using the accumulated product of weight and gradient and removes those weights shrinking toward zero. we refer to this approach as gradient-based pruning in this work. regarding pruning schedules, two commonly considered options are one-shot pruning @xcite and gradual pruning @xcite . while one-shot pruning removes the desired percentage of weights in a single step after the completion of finetuning, gradual pruning incrementally removes weights, starting from an initial pruning ratio (often 0) and gradually progressing towards the target pruning ratio. in mnmt, the most similar technique to our approach for extracting sub-networks through model pruning has been proposed by @xcite . in their work, sub-networks are searched through magnitude one-shot pruning, which means the model is pruned based on the magnitude of weights in a single step after the completion of fine-tuning. different from their method, we opt for the gradientbased pruning criterion, which is proved to be more effective than magnitude pruning in a transfer learning context. besides, in contrast to pruning in one single step @xcite , we gradually increase the ratio from zero to the target value during fine- algorithm 1 gradient-based gradual pruning for mnmt 1: input: n bilingual corpora data d @xmath0, ds 2 →t 2 , ...ds @xmath1, language pairs @xmath2, @xmath3, ..., @xmath4: // phase 1: pretrain a multilingual model 3: training a multilingual model base: θ0 4: // phase 2: extract sub-network for each pair 5: for each pair @xmath5: if @xmath6: // finetune without pruning 8: fine-tune and keep @xmath7: else if @xmath8: calculate gradient-based importance scores as shown in eq. (2), and rank the scores. 11: calculate the current pruning ratio according to eq. (4) and prune the weights with the lowest scores accordingly. 12: else 13: finetune with target ratio until converge. 14: end if 15: extract pruning mask ms @xmath9: extract sub-network θs @xmath10: end for 18: // phase 3: structure-aware joint training 19: θ @xmath11, θs 2 →t 2 , ..., θs @xmath12: while θ all not converge do 21: for each pair @xmath13: further training θs @xmath14: ## methodology our approach includes three main phases. in the first one (phase 1 in algorithm 1), a multilingual model is initially trained using the parameter loglikelihood loss (see sec. 3.1). subsequently, the multilingual base model is simultaneously finetuned and pruned on each language pair (phase 2 in algorithm 1). in this phase, we adopt the gradientbased pruning criterion as the scoring function for each weight and the gradual pruning schedule to identify sub-network masks for language pairs (see sec. 3.2). the extracted masks are then jointly used during the final training (phase 3 in algorithm 1) as described in sec. 3.3. ## experiment results on iwslt this section shows the results of our approach on the iwslt dataset, along with a comparative analysis against the mnmt baseline and prior research works. we also provide a comprehensive analysis of our method from various perspectives. 5 ## experiment results on wmt on the highly imbalanced wmt dataset with the transformer-base model, we compare our method with the multilingual baseline and the prior work magnitudeoneshot, which shares the same philosophy as our method: both our approach and magnitudeoneshot introduce no additional parameters to the model and have no negative impact on the inference speed. the results in model low med high all baseline 14.01 17.54 24.16 18.76 magnitudeoneshot +0.70 +0.85 +0.99 +0.85 gradientgradual (ours) +0.91 +1.56 +1.67 +1.41 ## conclusion in a standard mnmt model, the parameters are shared across all language pairs, resulting in the parameter interference issue and compromised performance. in this paper, we propose gradient-based gradual pruning for multilingual translation to identify optimal sub-networks and mitigate the interference issue. extensive experiments on the iwslt and wmt datasets show that our method results in large gains over the normal mnmt system and yields better performance and stability than other approaches. additionally, we observe that the interference issue can be more severe in attention sublayers and it is possible to reconstruct a reliable phylogenetic tree of languages using the languagespecific sub-networks generated by our approach. all the experiments confirm the effectiveness of our approach and the need for better training strategies to improve mnmt performance.