Mixture of Diverse Size Experts (2024)

Abstract: The Sparsely-Activated Mixture-of-Experts (MoE) architecture has gained popularity for scaling large language models (LLMs) due to the sub-linearly increasing computational costs. Despite its success, most of the current structure designs face the challenge that the experts share the same size such that tokens have no chance to choose the experts with the most appropriate size to generate the next token. To migrate this defect, we propose Mixture of Diverse Size Experts (MoDSE), a new MoE architecture with designed layers where experts have different sizes. Analysis on difficult token generation tasks shows that experts with different sizes give better predictions, and the routing path of the experts tends to be stable after a period of training. The diversity of experts’ size will lead to load unbalancing. To tackle this limitation, we introduce an expert-pair allocation strategy to distribute the workload evenly across the GPUs. Comprehensive evaluations across multiple benchmarks demonstrate the effectiveness of MoDSE, surpassing existing MoEs by adaptively assigning the parameter budget to experts while maintaining the same total parameter size and number of experts.

## introduction large language models (llms) have demonstrated remarkable performance in a variety of nlp tasks and have become valuable assistants through a wide range of applications. the scaling law @xcite demonstrates that larger models exhibit superior performance. however, training larger models requires increased computational resources, posing a critical challenge. mixture-of-experts (moe) @xcite address this challenge by using sparse activation to scale up the trainable parameters while maintaining high training and inference efficiency. recent moe-based architectures, such as mixtral of experts @xcite , deepseekmoe @xcite , and openmoe @xcite have shown superior performance in various tasks. specifically, @xcite discuss two main issues in the design of the moe feed-forward networks (ffns) architecture: knowledge hybridity, where each expert covers diverse knowledge due to the limited number of experts, and knowledge redundancy, where multiple experts share common knowledge. to address these issues, they propose fine-grained expert segmentation by splitting the ffn intermediate hidden dimension and shared expert isolation by isolating certain experts to be always activated as shared experts. additionally, @xcite introduce hypernetworks and hyperexperts modules to capture the cross-expert and cross-layer knowledge. however, almost all existing moe architectures consist of experts with identical structures and sizes. this homogeneous architecture becomes a significant bottleneck when generating tokens with varying difficulty; some tokens are easier to predict, while others are more challenging. to deal with the varied difficulty, we propose the diverse size experts structure for each ffn layer, where each expert has a different parameter size to handle generating tasks of varying difficulty. note that we find a similar recent work called heterogeneous the mixture of experts @xcite , which shares a similar motivation and utilizes parameter penalty loss and router entropy loss to control the size and number of activated experts. our contributions are summarized as follows: ## preliminaries: mixture of experts moe models are usually constructed by replacing dense ffns layers in the transformer @xcite with moe layers. an moe layer typically consists of multiple experts n indicates the number of the experts. the gate model @xcite with trainable weight matrices w @xmath0, where h input is the dimension of input x and h is the dimension of the hidden layer. fedus et al. (2022) set k as one, while lepikhin et al. (2021); jiang et al. (2024) set as two. the outputs of experts are added with the noise to help with load balance. the noise generated from the input hidden vector x is multiplied by w n and processed by sof tplus and the root mean square layer normalization function (rmsnorm), where γ is a learnable coefficient. ## modse architecture predicting the next token is easier within frequently appearing token pairs in the corpus. tokens within the same word or phrase are easier to generate than those between two phrases or words. analogous to the human brain, the amount of thought required to generate the next word varies among different words. inspired by the fact that the difficulty of generating each next token varies, we propose modse as shown in figure 1 . in our work, the size of the expert parameters is used to quantify the amount of thinking involved. we assign experts a range of parameter sizes by setting the dimensions of the hidden layers to various lengths. however, the imbalance in expert size leads to an uneven workload. to address this issue, we propose a meticulously designed expert-pair allocation method to ensure each gpu node's workload is evenly distributed. ## conclusion in this paper, we propose modse, a novel structure for moe layers. inspired by the varying difficulties of next-token-generating tasks, we introduce the diverse size expert design, providing each expert with different prediction abilities. our analysis of token routing distribution shows that modse directs tokens to experts whose sizes are best suited for specific token generation tasks. this enhancement improves the moe model's performance in auto-regression tasks and demonstrates superior results compared to the conventional moe structure. additionally, we present the expert-pair allocation method to address the issue of load imbalances in the diverse size expert design, making the modse design more practical.