Harder Task Needs More Experts: Dynamic Routing in M o E Models (2024)

Abstract: In this paper, we introduce a novel dynamic expert selection framework for Mixture of Experts (MoE) models, aiming to enhance computational efficiency and model performance by adjusting the number of activated experts based on input difficulty. Unlike existing MoE approaches that rely on fixed TopK Routing, which activates a predetermined number of experts regardless of the input’s complexity, our method dynamically allocates experts based on the confidence level in expert selection for each input. This allows for more efficient utilization of computational resources, activating more experts for complex tasks requiring advanced reasoning and fewer for simpler tasks. Through extensive evaluations, our dynamic routing method demonstrates substantial improvements over Top2 Routing across various benchmarks, achieving an average improvement of 0.7% with less than 90% activated parameters. Further analysis shows our model dispatches more experts to tasks requiring complex reasoning skills, like BBH, confirming its ability to dynamically allocate computational resources in alignment with the input’s complexity.Our findings also highlight a variation in the number of experts needed across different layers of the transformer model, offering insights into the potential for designing heterogeneous MoE frameworks. The code and models are available at https://github.com/ZhenweiAn/Dynamic_MoE.

## introduction to effectively increase the model's parameter size, researchers have proposed the mixture of experts (moe) framework @xcite . by setting up multiple experts to enhance the model's overall capacity, moe models selectively activate a subset of parameters for use, thereby achieving more efficient parameter * equal contribution. songfang huang and yansong feng are the corresponding authors. most moe frameworks adopt a routing mechanism that dispatches a fixed number of experts for every input @xcite . the most famous method is topk routing @xcite , which initially calculates the probability of each expert being suited to the current input and then activates the topk suitable experts. empirically, previous works @xcite activate two experts per token, as activating more experts offers limited improvements in model performance but substantially increases training overhead. most of the subsequent studies @xcite can be seen as variants of topk routing, where different constraints are introduced to ensure that the number of tokens processed by different experts is as balanced as possible. almost all these efforts activate a fixed number of experts. the topk routing, though it achieves good performance on downstream tasks, overlooks the different difficulties of inputs. compared with simpler input, the more challenging input, e.g., tasks that require complex reasoning or logic inference, might need more parameters to solve. dispatching experts equally across inputs could lead to computational waste on simpler tasks and insufficient computational resources for more difficult ones. to fully leverage the potential of moe models, we propose a dynamic routing mechanism that adjusts the number of required experts based on the confidence level in the expert selection. when the model deems the currently selected experts as insufficient, it activates more experts. specifically, we first compute a probability distribution for selecting experts. if the highest probability for an expert exceeds a predefined threshold p, indicating high confidence, we activate only that one expert. otherwise, we progressively add additional experts until the cumulative probability of the selected experts exceeds the threshold p. this approach allows for a dynamic selection of experts, with the number of experts adjusted based on the input's complexity. evaluation across multiple common benchmarks has revealed that our method substantially outperforms moe models based on topk routing. compared with top2 routing, our dynamic routing achieves an average improvement of 0.7% with less than 90% activated parameters. further analysis has shown that our dynamic routing mechanism activates more experts in tasks requiring complex reasoning like bbh @xcite , while using fewer experts in relatively easier tasks such as hellaswag @xcite , confirming that our method indeed dynamically allocates experts based on the difficulty of the input. token-level analysis indicates that tokens with ambiguous semantics are more challenging for the model, typically activating more experts. another interesting finding is that the number of experts needed varies across different layers of the transformer. lower layers require more experts for combination, while the top layer needs only one. this may relate to the over-thinking phenomenon @xcite , which is widely observed in deep neural networks. our contributions can be summarized as follows: ## method in this section, we first briefly introduce the moe model with topk routing strategy, which activates a fixed number of experts for each token. as topk routing ignores the varying difficulty of different inputs and the different requirements for experts at different layers, we propose a dynamic routing mechanism that adjusts the number of activated experts according to the complexity of inputs. to avoid activating too many parameters through the dynamic routing mechanism, we also introduce a dynamic loss to encourage the model to activate only the necessary experts. ## efficiency of dynamic routing the greatest advantage of moe models is their ability to efficiently scale to larger models. the topk routing mechanism controls the number of parameters used by the entire model by activating a fixed number of experts. in contrast, our proposed dynamic routing mechanism removes the limitation of a fixed number of experts. naturally, there may be concerns that our method might assign too many experts to each token. another possible concern is that our dynamic routing mechanism might cost too many resources and make the whole model inefficient. to address these concerns, we demonstrate the efficiency of the dynamic routing mechanism from both training and inference perspectives. ## bottom layers need more experts an intriguing observation from our study is that our model achieves superior performance while activating fewer parameters. as shown in this unexpected phenomenon might be attributed to our method's more proper allocation of the experts across different layers, employing more experts at the bottom layers and fewer at the top. this layer-wise dynamic allocation, as opposed to the fixed number of experts per layer, somewhat mitigates the common issue of overthinking in deep neural networks, thereby enhancing performance. the overthinking refers to the situations where simpler representations of an input sample at an earlier layer, relative to the complex representations at the final layer, are adequate to make a correct prediction @xcite . previous works @xcite @xcite have demonstrated that shallower representations can achieve comparable, if not better, performance across various tasks than deeper representations. this could be due to deeper representations overfitting specific distributions, lacking generalizability, and being more vulnerable to attacks @xcite . it suggests that in some cases, acquiring a better shallow representation is more valuable than obtaining a more complex deep representation, which correlates to previous findings that removing top layers has a limited impact on the downstream tasks @xcite . compared with top2 routing, our dynamic routing activates more experts at the bottom layers to obtain better shallow representations and use the simpler network in the top layers to alleviate the overthinking issue. figure 3 displays the number of experts activated per token at different layers foot_1 . from the figure, we observe a gradual decrease in the average number of experts activated per token with the increase in layer depth. the lowest layer activates the most experts, up to 4 experts per token, enabling better shallow representations through a wider network, which is beneficial for various downstream tasks. at the topmost layer, the number of activated experts per token is reduced to even one. this phenomenon can prevent the model from being too complex and preserve generality in the final representation. it should be noted that we do not argue allocating most parameters to the bottom layers is the optimal strategy for a model. it is possible that one model could achieve better performance if more experts are activated in the intermediate layers. in our experiment, we observe the phenomenon that bottom layers need more experts and we think this is quite interesting and worthy of future exploration. beyond the series of overthinking works that we mention above @xcite , our finding also correlates with the conclusion that the upper layers of the model are tied to pre-training tasks @xcite . maybe, because the upper layer is specialized for the task of the next token prediction, the topmost layers only activate just one expert. ## related work the mixture of experts (moe) model is initially introduced by @xcite . recent studies have demonstrated sparsely gated moe models have substantial improvements in model capacity and efficiency, enabling superiors performance than dense models @xcite . particularly moe has shown great potential with the integration of transformer architectures @xcite . in previous moe architectures, a static number of experts are activated regardless of the varying complexity presented by input tokens. most of moe models activate top1 or top2 experts @xcite . there are many works @xcite make improvements based on the top2 routing @xcite . they make great efforts to distribute tokens among different experts more evenly during the training process and all of them activate static experts for each token. we choose top2 routing @xcite as the baseline because it's currently the most widely used method and the only one successfully applied in large language models (llm) thus far. for instance, @xcite , deepseek-moe @xcite and grok-moe (x.ai, 2024) all use this structure. there are some work that allocating different numbers of experts to different tokens. expert-choice moe model selects topk tokens for each expert @xcite and each token will be allocated with a different number of experts. however, on average every token activates exactly two experts, which is the same with top2 routing models and does not save computation. different from the previous works, our dynamic routing mechanism can activate fewer parameters on average by dynamically allocating experts to different tokens, which makes our method more efficient. recent works point out that in language modeling not all tokens require the same calculation cost @xcite and tasks across different domains require varying numbers of experts @xcite . our findings that harder task needs more experts are correlated with this line of work. ## conclusion our paper introduces a dynamic expert selection framework for mixture of experts (moe) models, surpassing existing static topk routing by adjust-ing expert activation based on input complexity. our approach not only improves computational efficiency but also model performance, evidenced by obvious gains over topk routing in our evaluations. our findings reveal the framework's effectiveness at dynamically dispatching different numbers of experts, particularly for complex reasoning tasks, and suggest the potential for developing more challenging heterogeneous moe models. in support of further research, we will open-source our models, contributing to advancements in the moe domain.