GaP-MoE : Explicit Expert Specialization in Mixture-of-Experts via Orthogonal Gradient Accumulation and Projection
Authors: Donghyun Jeong, Yoon-Yeong Kim
Organizations: Department of Statistics and Data Science University of Seoul Republic of Korea · Department of Statistics University of Seoul Republic of Korea
Sparse Mixture-of-Experts (SMoE) models improve scalability by routing input tokens to only a subset of expert networks, but they often suffer from representation collapse, in which multiple experts converge toward similar functions and lose their intended diversity. Prior work has addressed this limitation either by improving router architectures or by introducing auxiliary loss terms that implicitly penalize similarity among experts. However, router-based approaches remain vulnerable to routing collapses, and loss-based approaches can conflict with the task's original objective during optimization. In this work, we propose GaP-MoE (Gradient Accumulation and Projection in MoE), which induces expert specialization by directly manipulating expert output gradients rather than modifying the loss function. At each training step, GaP-MoE stochastically selects one of the TopK routed experts as a leader based on normalized routing scores, and projects the gradients of the remaining follower experts onto the orthogonal complement of the subspace spanned by the leader's recent gradients. This removes directional overlap between expert gradients without altering the task loss, allowing specialization to emerge as a direct consequence of the optimization dynamics rather than an externally imposed constraint. We evaluate GaP-MoE in both fine-tuning settings (BERT-based, on GLUE) and from-scratch pretraining settings (GPT-2-based, on WikiText-103, enwik8, and text8), and further confirm its compatibility as a drop-in module for existing SMoE architectures. Experimental results show that GaP-MoE achieves performance comparable to or better than standard SMoE.
Figures & tables
Figure 1: (a) GaP-MoE selects a leader expert from the TopK set using normalized routing probabilities pi . (b) The aggregated output gradients of the leader Ejl are stored in Qj as orthonormal basis vectors. When the similarity between their most recent basis vectors exceeds the threshold γ , the follower gradient is projected onto the orthogonal complement of the leader subspace and used to update Ekf .
Figure 2: Effect of gradient projection in the output representation space. The gray oval represents the past output gradient subspace of the leader expert ( gl(⋅)x ). The yellow and green lines represent the output vectors of the follower and leader experts, respectively, for the input token x , while the dotted lines represent their corresponding gradient update terms ( −ηg(t)x ). If these output vectors are aligned with the same subspace, the output gradient term −ηg(t)x is also aligned with that subspace. However, by projecting the follower’s gradient term, its output can escape from the existing subspace and encode different information, as illustrated by the red line.
Models
CoLA
SST-2
MRPC
STS-B
QQP
MNLI
QNLI
RTE
Avg.
SMoE
56.00
93.16
85.21
89.28
91.20
84.54
91.45
64.86
81.96
GaP-MoE( γ=−1 )
56.37
93.20
84.64
89.22
91.28
84.75
91.58
65.46
82.06
GaP-MoE( γ=0.7 )
56.44
93.00
85.29
89.27
91.28
84.52
91.45
64.62
81.98
Table 1: Evaluation results on the GLUE benchmark. For evaluation metrics, CoLA uses the Matthews correlation coefficient, STS-B uses Pearson’s correlation coefficient, and all others use accuracy. For all of them, the higher the score, the better. Results are averaged over three seeds.
Models
WikiText-103
enwik8
text8
SMoE
78.80
2.005
1.962
GaP-MoE( γ=0.7 )
76.50
1.953
1.945
Table 2: Evaluation results of from-scratch environments. For evaluation metrics, WikiText-103 uses perplexity, while enwik8 and text8 use BPC. For all of them, the lower the score, the better. Results are averaged over three seeds.
Table 5
Figure 3: Each bar represents, in order, the expert output similarity of SMoE and GaP-MoE with γ=−1 and γ=0.7 , respectively, trained using a single seed on COLA. The x-axis shows each token passed through the layers in the first sequence of COLA’s evaluation set.
Figure 5: In both cases, the top figure shows the first MoE layer, and the bottom figure shows the last MoE layer; the t∈[1,3,5,10] used to set the Gap indicates how many time steps back the expert was compared to.
Sparse Mixture-of-Experts (MoE) models scale capacity by routing each token to a small subset of experts. However, their routers exhibit a fundamental trade-off: strong load balancing can suppress expert specialization, while aggressive diversity often causes routing collapse. We propose Hi-MoE, a grouped MoE framework that decomposes routing control into two coupled levels: (i) inter-group balancing that enforces fair traffic across expert groups, and (ii) intra-group specialization that promotes complementary expert behaviors while preventing within-group collapse. Our analysis provides a principled explanation of how our hierarchical objectives reshape the router, thereby promoting stable specialization and mitigating collapse. We observe consistent improvements over recent sparse-routing and grouped-MoE baselines across NLP and vision benchmarks, and confirm robustness via scaling studies (model size, expert count) and targeted ablations. In large-scale pre-training on 58B tokens, Hi-MoE-7B achieves a 5.6% perplexity reduction and a 40% improvement in expert balance over OLMoE-7B across diverse evaluation domains.
Gleb Molodtsov, Alexander Miasnikov, Aleksandr Beznosikov
Sparse Mixture-of-Experts (MoE) architectures enable scaling LLM parameters under a fixed inference budget by activating only a small subset of experts via top-k routing. While this preserves causality and suits autoregressive language models, the discrete top-k operator is not differentiable, forcing a fixed number of active experts per input and resulting in inefficient use of computation. We propose SoftMoE, which replaces discrete routing with a truncated soft top-k LapSum relaxation, allowing gradient-based optimization of expert routing. We further parameterize the mean number of active experts per layer and impose a global budget constraint, enabling the model to learn how to allocate expert capacity across layers. SoftMoE remains fully compatible with autoregressive modeling and achieves performance comparable to or better than sparse MoE on language modeling and downstream tasks, while activating significantly fewer experts. Notably, the learned allocation is highly non-uniform, with later layers activating more experts. The source code is publicly available†.
Mikołaj Zasada, Łukasz Struski, Jacek Tabor +1
AGH University of Krakow, Poland · Faculty of Mathematics and Computer Science, Jagiellonian University, Poland · Jagiellonian University, Poland +1
Mixture-of-experts (MoE) models have become an important approach for scaling model parameters. Modern MoE architectures increasingly combine two design choices: fine-grained expert segmentation and shared experts. The former increases expressiveness by enabling more expert combinations at a fixed activation ratio, and the latter introduces experts that are active for every token. We identify an exploration dilemma between these designs: shared experts can make routing less exploratory by reducing the variability of token representations. This conflicts with the greater need for exploration induced by fine-grained expert segmentation due to the combinatorially enlarged routing space. Consequently, the additional expressiveness offered by fine-grained segmentation may not be fully realized in practice. MirrorMoE is proposed to reconcile this exploration dilemma by shrinking the exploration space through token-adaptive suppression of unnecessary experts. MirrorMoE consistently improves performance across model scales.
Bingcong Li, Hao Ma, Michael Muehlebach +2
ETH Zurich · The University of Hong Kong · Max Planck Institute for Intelligent Systems +1