Synergistic Intra- and Cross-Layer Regularization Losses for MoE Expert Specialization

Sparse Mixture-of-Experts (MoE) models scale Transformers efficiently but suffer from expert overlap, where different experts process similar tokens and learn redundant functions, resulting in ambiguous routing and underutilized capacity. While architectural solutions like DeepSeek-style shared experts promote specialization, they require substantial structural modifications and rely solely on intra-layer signals. We propose two plug-and-play auxiliary losses that enhance MoE specialization and routing efficiency without modifying routers or model architectures. First, an intra-layer specialization loss penalizes cosine similarity between experts' SwiGLU activations on identical tokens, encouraging experts to specialize in complementary functions. Second, a cross-layer dependency loss maximizes joint Top-$k$ routing probabilities across adjacent layers, establishing coherent expert pathways through network depth while reinforcing intra-layer specialization. Both losses are orthogonal to the standard load-balancing loss and compatible with shared-expert and vanilla Top-$k$ MoE architectures. We implement both losses as a drop-in Megatron-LM module. Extensive experiments across pre-training, fine-tuning, and zero-shot benchmarks demonstrate consistent task gains, higher expert specialization, and lower-entropy routing; together, these improvements translate into faster inference via more stable expert pathways.

Large AI language models are often made more efficient by dividing their work among many smaller “expert” components, so that only a few experts are used for each piece of text. However, these experts can gradually become too similar to each other during training. When this happens, the model wastes capacity, and it becomes less clear which expert should handle which input. This paper introduces two simple training signals that help the experts develop clearer roles. The first encourages experts working on the same input to produce different kinds of internal responses, so they do not all learn the same behavior. The second encourages the model to follow more stable expert pathways across different layers, making its expert choices more consistent. These methods can be added to existing mixture-of-experts language models without changing their architecture. Across several training and evaluation settings, they improve model performance, make expert behavior more specialized, and support more stable routing, which can also help inference run more efficiently.