HMoE: Heterogeneous Mixture of Experts for Language Modeling (2408.10681v1)

Abstract: Mixture of Experts (MoE) offers remarkable performance and computational efficiency by selectively activating subsets of model parameters. Traditionally, MoE models use homogeneous experts, each with identical capacity. However, varying complexity in input data necessitates experts with diverse capabilities, while homogeneous MoE hinders effective expert specialization and efficient parameter utilization. In this study, we propose a novel Heterogeneous Mixture of Experts (HMoE), where experts differ in size and thus possess diverse capacities. This heterogeneity allows for more specialized experts to handle varying token complexities more effectively. To address the imbalance in expert activation, we propose a novel training objective that encourages the frequent activation of smaller experts, enhancing computational efficiency and parameter utilization. Extensive experiments demonstrate that HMoE achieves lower loss with fewer activated parameters and outperforms conventional homogeneous MoE models on various pre-training evaluation benchmarks. Codes will be released upon acceptance.

HMoE: Heterogeneous Mixture of Experts for LLMing

The paper "HMoE: Heterogeneous Mixture of Experts for LLMing" presents an innovative approach to enhancing the effectiveness and efficiency of Mixture of Experts (MoE) models in the field of LLMs. Traditionally, MoE models employ homogeneous experts, each characterized by an identical structure and capacity. This homogeneity restricts the models’ ability to specialize and optimize parameter utilization effectively. The authors propose the Heterogeneous Mixture of Experts (HMoE) as a novel framework where experts differ in size, offering diverse capacities to handle varying input complexities. The proposed framework includes tailored training objectives aimed at balancing the expert activation, thus improving computational efficiency and performance.

Introduction

Mixture of Experts (MoE) models have garnered attention for their ability to enhance the performance of LLMs without linearly increasing computational costs. MoE models operate by selectively activating subsets of model parameters via multiple experts, each specializing in different tasks or data aspects. The key advantage of MoE lies in its scalability, allowing model parameters to increase without a corresponding rise in computational burden.

However, existing MoE models, featuring homogeneous experts, face significant challenges. These include convergence phenomena, which reduce the uniqueness and specialization potential of experts, and limitations in handling input data of varying complexity due to uniform expert capacity. This paper hypothesizes that heterogeneity among experts—characterized by differences in size—can address these limitations by promoting better specialization and optimal parameter utilization.

Methodology

To implement HMoE, the authors propose varying the sizes of experts to create heterogeneous capacities. However, initial explorations indicated that simply introducing heterogeneity was insufficient. Larger experts tended to dominate, reducing the representational capacity of the model. Consequently, a novel set of training objectives was developed to encourage the activation of smaller experts. This approach not only addresses imbalanced expert activation but also enhances computational efficiency.

The advanced HMoE framework was compared against conventional homogeneous MoE models in both Top-K and Top-P routing strategies. The key differentiators included:

1. Expert Size Distribution: Different strategies, such as geometric, arithmetic, and hybrid, were analyzed to identify optimal expert size distributions.
2. Parameter Penalty Loss: This was introduced to encourage the activation of smaller experts, balancing the load more effectively and ensuring efficient parameter utilization.
3. Routing Strategy Enhancements: Entropy loss was applied in Top-P routing to stabilize the number of activated parameters over time.

Experimental Results

Extensive experiments were conducted to evaluate the effectiveness of HMoE models. Key findings include:

• Performance and Efficiency: HMoE models consistently outperformed homogeneous MoE models across various pre-training evaluation benchmarks with fewer activated parameters. For instance, the HMoE-3B model with Top-P routing not only improved performance metrics but also showed a 2.6% average improvement compared to the Dense-1B model.
• Optimal Activation Parameters: The paper illustrated that, for larger training FLOPs, the optimal number of activated parameters in HMoE was lower than that in homogeneous MoE, indicating better efficiency in HMoE.
• Similarity and Synergy Analysis: Experts of similar sizes within HMoE displayed higher similarity in token distribution, indicating specialized functionalities. Smaller experts showed more generalized capabilities, further contributing to efficient model performance.
• Token Complexity Handling: Visualization of activated experts for tokens of varying complexities revealed that HMoE effectively allocated simpler tokens to smaller experts and complex tokens to larger experts, thereby balancing the computational load.

Conclusion and Implications

The introduction of heterogeneous experts in HMoE models significantly enhances both computational efficiency and model performance by allowing for specialized handling of input data of varying complexities. The novel training objectives proposed, including the parameter penalty loss, effectively address imbalances in expert activation, ensuring more efficient and effective use of model parameters. The findings suggest substantial potential for further optimization and widespread application of heterogeneous expert architectures in a broad array of natural language processing tasks.

The paper opens avenues for future research in optimizing expert size distributions and exploring additional routing strategies. This work establishes a foundation for the development of more nuanced and capable LLMs, paving the way for advancements in both theoretical understanding and practical applications of MoE frameworks.