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

TL;DR

This paper introduces two plug-and-play regularization losses for Mixture-of-Experts models that improve expert specialization and routing efficiency without architectural changes, leading to better performance and faster inference.

Contribution

It proposes novel intra- and cross-layer regularization losses that enhance MoE expert specialization and routing coherence without modifying existing architectures.

Findings

Improved expert specialization and routing efficiency.

Consistent performance gains across various benchmarks.

Faster inference due to more stable expert pathways.

Abstract

Sparse Mixture-of-Experts (MoE) models scale Transformers efficiently but suffer from expert overlap -- redundant representations across experts and routing ambiguity, resulting in severely underutilized model capacity. While architectural solutions like DeepSeekMoE promote specialization, they require substantial structural modifications and rely solely on intra-layer signals. In this paper, we propose two plug-and-play regularization losses that enhance MoE specialization and routing efficiency without modifying router 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 knowledge. Second, a cross-layer coupling loss maximizes joint Top-$k$ routing probabilities across adjacent layers, establishing coherent expert pathways through network depth while reinforcing intra-layer expert specialization. Both losses are orthogonal to the standard load-balancing loss and compatible with both the shared-expert architecture in DeepSeekMoE 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.