S table M o E : Stable Routing Strategy for Mixture of Experts (2022)

Abstract: The Mixture-of-Experts (MoE) technique can scale up the model size of Transformers with an affordable computational overhead. We point out that existing learning-to-route MoE methods suffer from the routing fluctuation issue, i.e., the target expert of the same input may change along with training, but only one expert will be activated for the input during inference. The routing fluctuation tends to harm sample efficiency because the same input updates different experts but only one is finally used. In this paper, we propose StableMoE with two training stages to address the routing fluctuation problem. In the first training stage, we learn a balanced and cohesive routing strategy and distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the distilled router to determine the token-to-expert assignment and freeze it for a stable routing strategy. We validate our method on language modeling and multilingual machine translation. The results show that StableMoE outperforms existing MoE methods in terms of both convergence speed and performance.

## introduction in recent years, large-scale transformers @xcite @xcite @xcite have shown a striking ability to model languages. however, with the model scale growing, the training speed will go slower, and the extremely large memory requirement also introduces a heavy burden of engineering. mixture of experts (moe) @xcite @xcite , in a much easier way, enables transformers to scale up the number of parameters meanwhile introducing an affordable * contribution during internship at microsoft research. computational overhead. moe-based transformers have a set of expert modules, and only a few experts will be activated for each input token. in this way, we can expand the model scale by adding expert modules, which will keep the computational and memory overhead within a tolerable range. most existing moe methods @xcite @xcite decide the token-to-expert routing according to the dynamically changing token representations. however, we point out that they face the routing fluctuation problem. as shown in figure 1 , the same input may be assigned to different experts along with training. however, during inference, only one expert will be activated for the input. the routing fluctuation problem tends to harm sample efficiency because the same input updates different experts while only one is finally used. taking base layer @xcite as an example, during the whole training process, we examine the token-to-expert assignment for tokens in the validation set. for an input token, we define the last fluctuation step as the last step where its target expert is different from the final step. we plot the cumulative token percentage with regard to the last fluctuation step (annotated as its percentage accounting for all training steps) in figure 2 . we find that the last fluctuation step of 40.9% tokens exceeds 20%, which means 40.9% tokens do not have a stable target expert when 20% of all training steps have been done. furthermore, 29.1% tokens still change their target experts after half of the whole training process, and 15.4% tokens even change the target expert after 80% of all training steps, which is nearing the training ending. these statistics prove that the routing fluctuation problem indeed exists in previous moe methods. in this paper, we propose stablemoe with two training stages to address the routing fluctuation problem. in the first training stage, we follow the learning-to-route paradigm and aim to learn a training epochs expert 3 inference 𝑒 1 𝑒 2 𝑒 3 … … transformer blocks transformer blocks the car is expert 2 expert 1 expert 3 transformer blocks transformer blocks expert 2 expert 1 expert 3 transformer blocks transformer blocks expert 2 expert 1 transformer blocks transformer blocks expert 2 expert 1 expert 3 the car is the car is the car is 0% 20% 40% 60% 80% 100% last fluctuation step 0% 20% 40% 60% 80% 100% cumulative token percentage (20%, 40.9%) (50%, 29.1%) (80%, 15.4%) figure 2: cumulative token percentage with regard to the last fluctuation step of tokens for base layer @xcite . a substantial portion of tokens still change their target experts even if the training is nearing the end. our contributions are summarized as follows: ## background: mixture-of-experts for transformers we first introduce the moe mechanism designed for transformers @xcite . given a standard l-layer transformer model and an input sequence x containing t tokens, the transformer output h l is calculated by we implement moe for transformers by inserting moe layers, that are composed of a set of ffns, into two neighboring transformer blocks. at an moe layer, for each input token, only a few or one expert will be activated, controlled by a gating function g(•): transformer blocks transformer blocks this car is router expert 1 distilled router distilling car is expensive training stage 1 training stage 2 expert 2 expert 3 transformer blocks transformer blocks this car is distilled router (frozen) expert 1 car is expensive expert 2 expert 3 figure 3: illustration of two training stages in stablemoe. in training stage 1, we learn a routing strategy and distill it into a lightweight router. then, we freeze the distilled router for stable routing in training stage 2. during inference, we also use the frozen distilled router for consistent routing. ## conclusion in this paper, we point out the routing fluctuation problem that exists in previous learning-to-route moe methods. in order to address this problem, we propose stablemoe with two training stages. we first learn a balanced and cohesive routing strategy and synchronously distill it into a lightweight router decoupled from the backbone model. then, we freeze the distilled router for a stable routing strategy in the remaining training. we validate sta-blemoe on language modeling and multilingual machine translation. the results show that sta- blemoe outperforms existing moe methods in terms of both convergence speed and performance. 7 acknowledgement damai dai, zhifang sui, and baobao chang are supported by the national key research and development program of china 2020aaa0106701 and nsfc project u19a2065. experts most frequent tokens descriptions 5 my, his, her, year, years, day, life, week, family, days possessive case & time units 6 with, at, from, about, them, need, want, him, against, using prepositions & objective case 11 that, ?, !, which, )., .", that, "., .), !!, ?", !!!, :), â, !", ?, !, !), conjunctions & punctuations 12 one, what, some, any, two, many, $, use, 2, 1 numerals 13 information, support, experience, service, data, services, money, access, research nouns about technologies 17 world, government, state, country, community, city, 2018, united, us, law nouns about politics 22