HMoE: Heterogeneous Mixture of Experts for Language Modeling

TL;DR

Mixture of Experts models enable parameter-efficient scaling, but traditional MoE uses homogeneous experts with equal capacity, limiting specialization. This paper introduces HMoE, a heterogeneous MoE where experts differ in size to better match token complexity, and proposes training objectives, including a parameter penalty loss and router-entropy considerations, to bias activation toward smaller experts while preserving ability for hard tasks. Empirical results on RedPajama-based pretraining show HMoE achieves lower loss with fewer activated parameters and outperforms homogeneous MoE across multiple pre-training evaluation benchmarks, under both Top-K and Top-P routing. The findings suggest heterogeneous expert sizes promote specialization and efficiency, enabling more scalable, bias-aware MoE-based language models; future work will explore broader heterogeneity strategies and hardware-friendly training.

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.

Figures (13)

• Figure 1: Comparisons of our heterogeneous MoE-3B with conventional homogeneous MoE-3B. Our proposed HMoE is superior on both performance and efficiency.
• Figure 2: Two distinct model structures for Mixtures of Experts (MoE) are compared: (a) conventional homogeneous MoE model structure with all experts having identical parameter sizes; (b) our proposed heterogeneous MoE model structure characterized by substantial variations in parameter sizes of each expert, incorporating a parameter penalty loss during training to promote utilization of Experts with smaller parameter volumes. In our heterogeneous MoE, harder tokens are assigned to larger experts, while easier tokens are assigned to smaller experts. In conventional homogeneous MoE, all tokens are assigned to the same size experts regardless of their difficulty.
• Figure 3: Experimental results of intuitive exploration on HMoE. The left figure compares the performance of intuitive HMoE and conventional Homogeneous MoE. The Homogeneous MoE adapts load balancing loss while the intuitive Hetergeneous MoE does not utilize any auxiliary loss. The right figure shows the activated ratio of experts in the intuitive HMoE. The relative expert sizes in HMoE are $\{9,11,13,15,17,19,21,23\}$, matching experts a to h.
• Figure 4: Analysis of isoFLOP for conventional MoE (Top-P) and our poposed HMoE (Top-P). The left figure depicts the optimal activated model parameters for various FLOPs. The right figure illustrates the variations in loss as FLOPs increase, given the optimal settings.
• Figure 5: Average activated parameters across training FLOPs (left) or different layers (right).