Merging Experts into One: Improving Computational Efficiency of Mixture of Experts (2023)

Abstract: Scaling the size of language models usually leads to remarkable advancements in NLP tasks. But it often comes with a price of growing computational cost. Although a sparse Mixture of Experts (MoE) can reduce the cost by activating a small subset of parameters (e.g., one expert) for each input, its computation escalates significantly if increasing the number of activated experts, limiting its practical utility. Can we retain the advantages of adding more experts without substantially increasing the computational costs? In this paper, we first demonstrate the superiority of selecting multiple experts and then propose a computation-efficient approach called Merging Experts into One (MEO), which reduces the computation cost to that of a single expert. Extensive experiments show that MEO significantly improves computational efficiency, e.g., FLOPS drops from 72.0G of vanilla MoE to 28.6G (MEO). Moreover, we propose a token-level attention block that further enhances the efficiency and performance of token-level MEO, e.g., 83.3% (MEO) vs. 82.6% (vanilla MoE) average score on the GLUE benchmark. Our code will be released upon acceptance. Code will be released at: https://github.com/Shwai-He/MEO .

## introduction scaling language models has achieved promising progress in the field of nlp @xcite . to further increase the model size under a computational budget, sparsely activated networks @xcite only employ a few parameters for each input. a widely studied approach is the mixture-of-experts (moe, @xcite , which trains multiple expert networks but only selects a subset of them for a specific input @xcite . compared to dense networks of the same model size, moe effectively reduces computational costs. although increasing the experts selected for each input can improve the representation diversity @xcite and downstream task performance @xcite , it usually comes with a price of significantly growing computational cost. our empirical study (figure 1 and our work aims to improve the computational efficiency of moe inference with multiple experts selected, for greatly rejuvenating the compromised performance. the computation involved in moe primarily consists of the inference on each selected expert and the summation of their outputs, with the former dominating the cost. hence, the cost linearly grows with the number of selected experts. to overcome the computational bottleneck, we instead propose merging experts into one (meo), 14685 which alters the calculation order of the two operations, i.e., first merging the parameters of the selected experts into one expert followed by inference on the merged expert. since the parameter merging only requires summation, an meo layer (approximately) only consumes the computation of single-expert inference, no matter how many experts are selected. this leads to a nearly constant inference cost when scaling up the model capacity (i.e., the number of selected experts) to improve the performance. meo can be applied as a drop-in replacement for moe, which has been deployed at various levels, e.g., selecting experts for each token @xcite , each sequence @xcite , each task @xcite , etc. on the sequence/task level, our empirical studies demonstrate that replacing moe with meo significantly improves computational efficiency, e.g., reducing flops from 72.0g to 28.9g, without hurting the performance. in addition, we propose a token-level attention mechanism that further enhances the efficiency and performance, e.g., from 82.6% (moe) to 83.3% (meo) on bert-base (figure 1 ). ## methodology review of mixture of experts. given a token moe performs at various levels, e.g., token, sequence, and task, where moe selects experts based on a single token, input sequence, or task embed-ding (or task ids): in however, selecting more experts leads to a substantial increase in flops (e.g., 74.1g v.s. 7.5g when increasing m from 1 to 32). this phenomenon urges us to reflect whether there exists an efficient approach to achieve both high performance and computational efficiency. our goal is to ensure consistent computational cost, regardless of the number of selected experts. merging experts into one. the computation cost of moe primarily involves the computation of individual experts (i.e., @xmath0.e., o(g) and o( @xmath1. notably, the computation of individual experts plays a dominant role, with even the cost of a single expert being significantly outweighing that of the mixture: on the other hand, as the number of selected experts m increases, the term @xmath2, whereas the increase in o( @xmath3. therefore, it is essential to address the growing trend of @xmath4. as illustrated in figure 2 , we propose the method called merging experts into one (meo), where the key idea is to leverage the gating scores to aggregate the parameters of the selected experts (which is akin to the simple weighted model fusion mechanism @xcite ): meo at different levels. in the case of sequence and task level meo, all tokens within a sequence share the same gating scores, as well as the aggregated parameters ŵ and b 1 . this property allows for easy adoption of meo at these levels. however, when directly applying meo at the token level, the situation is different. since the gating scores of each token within a sequence are unique, the straightforward usage of meo would require the aggregation of multiple sets of weights and biases, resulting in increased deployment cost. therefore, we refine and enhance the framework of token-level meo specifically. token-level meo. our proposed token-level meo aims to incorporate token-level information with minimal extra computational cost. specifically, the expert selection is performed at the sequence level, thereby preserving context information and eliminating the necessity of aggregating multiple weights and biases for individual tokens. to capture the identification of each token, we leverage the token attention mechanism inspired by @xcite ; @xcite . specifically, given the input sequence @xmath5, we employ a specialized bottleneck block, inspired by adapter-like structures @xcite . the bottleneck layer incorporates down-projection weights w @xmath6, an activation function f and up-projection weights w @xmath7, with reduce factor @xmath8. by operating on each token individually, the bottleneck applies token-level attention to the input sequence x: with the inclusion of token identification in the updated input, meo performs aggregation of ŵ and b through sequence-level expert selection. subsequently, these aggregated ŵ and b are used to compute the output in conjunction with the input. ## empirical evaluation experimental setup. experiments were conducted on four widely-used benchmarks, spanning understanding and generation tasks: (1) glue @xcite , containing understanding tasks like natural language inference, sentiment analysis, and sentence similarity evaluation; (2) xsum @xcite , a summarization dataset where the models are required to generate a short summary for a given article; (3) wikitext-2 @xcite , a collection of over 100 million tokens extracted from the set of verified good and featured articles on wikipedia where the models are utilized to generate the next tokens; (4) ## limitations despite the progress we have made, there are still limitations in our work. while our architecture for the mixture of experts demonstrates improved efficiency, there is a need for further exploration in terms of its deployment. specifically, determining the optimal number of experts in specific layers and selecting different levels of moes require additional investigation. we believe that with the implementation of efficient deployment strategies, our method has the potential to become even more competitive.