Adaptive Gating in Mixture-of-Experts based Language Models (2023)

Abstract: Large language models have demonstrated exceptional language understanding capabilities in many NLP tasks. Sparsely activated mixture-of-experts (MoE) has emerged as a promising solution for scaling models while maintaining a constant number of computational operations. Existing MoE models adopt a fixed gating network where each token is computed by the same number of experts. This contradicts our intuition that the tokens in each sequence vary in terms of their linguistic complexity and, consequently, require different computational costs. Little is discussed in prior research on the trade-off between computation per token and model performance. This paper introduces adaptive gating in MoE, a flexible training strategy that allows tokens to be processed by a variable number of experts based on expert probability distribution. Adaptive gating preserves sparsity while improving training efficiency. We further draw upon curriculum learning to better align the order of training samples and maximize the training time savings. Extensive experiments on diverse NLP tasks show that adaptive gating reduces at most 22.5% training time while maintaining inference quality. Moreover, we conduct a comprehensive analysis of the gating decisions and present our insights on which tokens are inherently difficult to process, depending on the specific language task.

## introduction the field of natural language processing (nlp) has undergone a remarkable revolution driven by the rapid advancements in language models (cha; @xcite @xcite . they exhibit socalled "emergent" capabilities for a wide variety of applications @xcite . however, as demands for these applications continue to grow, scalability of these models poses an increasingly challenging hurdle due to constraints in computational resources, memory capacity, interconnect bandwidth, etc. @xcite . sparsely-activated mixture-of-experts (moe) is a promising paradigm to address the scalability issue while maintaining a constant number of computation flops @xcite . moe utilizes an ensemble of experts to collectively tackle the learning task. each input activates a subset of experts, resulting in a dynamically-changing and sparse computation graph. this method effectively distributes the computation among experts, increases model capacity and improves training efficiency @xcite . very recently, there has been quite some prior work on improving the performance of transformers using moe @xcite @xcite . despite moe's benefit in scalability, it suffers from suboptimal training efficiency. in particular, we focus on the gating mechanism that selects the experts for each token in this work. existing moe models adopt a fixed top-2 gating in training while employing top-1 gating during inference for shorter response times. top-2 gating entails twice the computational cost per token and doubles the data transfer size of all-to-all operations compared to top-1. yet, it remains unclear whether top-2 gating actually leads to performance gains that could justify the additional overheads. therefore, a comprehensive analysis of the trade-off between training efficiency and model performance is increasingly crucial. more practically, how to construct an moe language model that effectively balances training efficiency and performance, is of great interest and imminent value. towards this end, we present our first attempt to empirically characterize and improve the efficiency of the gating mechanism in moe. we observe that across various models and tasks, a large number of tokens display simple linguistic characteristics or a single dominant feature, which allows them to be effectively processed using just the top-1 expert. this observation suggests that the current top-2 gating strategy incurs unnecessary computation costs for a significant number of tokens. motivated by this insight, we further introduce adaptive gating in moe that enables tokens to be processed by a flexible number of experts depending on the gating decision. our approach, in contrast to conventional moe models, preserves the sparsity of moe models while enhancing flexibility in token handling. we incorporate a threshold within the gating network to conduct adaptive token routing based on the distribution of expert probabilities. with adaptive gating, the majority of tokens use simple top-1 gating; top-2 gating is selectively applied only when necessary and beneficial, thus significantly reducing the computation cost. however, the training efficiency cannot achieve the same improvement as the computation cost due to the fact that tokens with top-2 gating always incur a longer training step, thus becoming the bottleneck. therefore, to enhance training efficiency even further, we leverage the idea of curriculum learning by strategically adjusting the order of training data samples. we conduct extensive experiments on six nlp tasks with different encoder and decoder models. the results show that our approach can effectively reduce the end-to-end training time by at most 22.5%, while achieving comparable inference quality with top-2 gating moe models. moreover, we show that the tokens routed to two experts are coupled with the nature of each nlp task. for sentiment analysis, those are the tokens expressing neutral opinions; translation task pays attention to sentences with complex structure; question and answer connects key words in question and context and assign both with top-2 gating; summarization puts more effort in understanding the pronouns and finding tokens expressing central idea; top-2 routing decision changes along with the token to generated in text completion task and conversational tokens in dialogue response task use top-2 experts frequently. empirically, we find that a small threshold value (i.e. 0.1, 0.2) in adaptive gating can lead to a similar performance as top-2 gating. our contributions are as follows: ## design we now discuss the design of adaptive gating in moe for training. ## evaluation we evaluate adaptive gating in moe on six nlp tasks using various encoder and decoder models. we then analyze the gating decision to better understand the effectiveness of adaptive gating. ## limitation choice of k. adaptive gating in moe currently is limited to top-k gating, where k can be either 1 or 2. this is built on the common practice in extensive prior work that top-2 gating shows a promissing resut in moe. further evaluation is necessary to validate the performance of a wider range of k values. our experiments were conducted on a diverse set of nlp tasks and datasets, but it is essential to note that the effectiveness and efficiency of adaptive moe may vary depending on the specific task characteristics. different tasks may exhibit distinct patterns and complexities, which can impact the performance and generalizability of the proposed approach. further investigation and evaluation on a wider range of tasks would provide a more comprehensive understanding of the limitations and applicability of adaptive moe. ## conclusion this paper demonstrates the effectiveness and flexibility of adaptive gating in moe models for a wide range of natural language processing tasks. by dynamically adjusting the number of experts based on token characteristics, we achieve improved training efficiency without compromising inference performance. additionally, the integration of curriculum learning allows us to tackle the challenge of varying execution times, thereby reducing training costs. our research sheds light on the trade-off between training efficiency and model performance in sparse and dynamic moe networks, offering valuable insights for the development of more scalable and adaptable language models.