StableMoE: Stable Routing Strategy for Mixture of Experts

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 7085 - 7095 May 22-27, 2022 c⃝2022 Association for Computational Linguistics STABLEMOE: Stable Routing Strategy for Mixture of Experts Damai Dai†‡∗, Li Dong‡, Shuming Ma‡, Bo Zheng‡, Zhifang Sui†, Baobao Chang†, Furu Wei‡ †MOE Key Lab of Computational Linguistics, Peking University ‡Microsoft Research {daidamai,szf,chbb}@pku.edu.cn {lidong1,shumma,v-zhebo,fuwei}@microsoft.com 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 is- sue, 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 in- ference. The routing fluctuation tends to harm sample efficiency because the same input up- dates different experts but only one is finally used. In this paper, we propose STABLEMOE with two training stages to address the rout- ing fluctuation problem. In the first training stage, we learn a balanced and cohesive rout- ing strategy and distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the dis- tilled 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 transla- tion. The results show that STABLEMOE outperforms existing MoE methods in terms of both convergence speed and performance. The code is available at https://github. com/Hunter-DDM/stablemoe. 1 Introduction In recent years, large-scale Transformers (Devlin et al., 2019; Dong et al., 2019; Raffel et al., 2020; Clark et al., 2020; Bao et al., 2020; Brown et al., 2020) have shown a striking ability to model lan- guages. However, with the model scale grow- ing, the training speed will go slower, and the extremely large memory requirement also intro- duces a heavy burden of engineering. Mixture of Experts (MoE) (Jacobs et al., 1991; Jordan and Ja- cobs, 1994; Shazeer et al., 2017), 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 Transform- ers 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 (Lepikhin et al., 2021; Fedus et al., 2021; Lewis et al., 2021) decide the token-to-expert routing according to the dynam- ically 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 (Lewis et al., 2021) as …

Content truncated to 500 words