Guiding the Experts: Semantic Priors for Efficient and Focused MoE Routing

TL;DR

This paper introduces a semantic-aware enhancement for Soft MoE models, aligning expert routing with semantic regions to improve efficiency, interpretability, and performance in vision tasks.

Contribution

It proposes a foreground-guided auxiliary loss and LayerScale mechanism to explicitly incorporate semantic priors into MoE routing, enhancing interpretability and accuracy.

Findings

Improved accuracy on ImageNet-1K and other benchmarks.

More interpretable expert routing patterns.

Seamless integration with existing Soft MoE frameworks.

Abstract

Mixture-of-Experts (MoE) models have emerged as a promising direction for scaling vision architectures efficiently. Among them, Soft MoE improves training stability by assigning each token to all experts via continuous dispatch weights. However, current designs overlook the semantic structure which is implicitly encoded in these weights, resulting in suboptimal expert routing. In this paper, we discover that dispatch weights in Soft MoE inherently exhibit segmentation-like patterns but are not explicitly aligned with semantic regions. Motivated by this observation, we propose a foreground-guided enhancement strategy. Specifically, we introduce a spatially aware auxiliary loss that encourages expert activation to align with semantic foreground regions. To further reinforce this supervision, we integrate a lightweight LayerScale mechanism that improves information flow and stabilizes optimization in skip connections. Our method necessitates only minor architectural adjustments and can be seamlessly integrated into prevailing Soft MoE frameworks. Comprehensive experiments on ImageNet-1K and multiple smaller-scale classification benchmarks not only showcase consistent performance enhancements but also reveal more interpretable expert routing mechanisms.