Title: Advancing Expert Specialization for Better MoE
Venue: NeurIPS 2025
Accepted: yes
Keywords: Mixture of Experts

Load Balancing

Optimization

Supervised Fine-Tuning
Abstract: Mixture-of-Experts (MoE) models enable efficient scaling of large language models (LLMs) by activating only a subset of experts per input. 
However, we observe that the commonly used auxiliary load balancing loss often leads to expert overlap and overly uniform routing, which hinders expert specialization and degrades overall performance during post-training.
To address this, we propose a simple yet effective solution that introduces two complementary objectives: (1) an orthogonality loss to encourage experts to process distinct types of tokens, and (2) a variance loss to encourage more discriminative routing decisions.
Gradient-level analysis demonstrates that these objectives are compatible with the existing auxiliary loss and contribute to optimizing the training process.
Experimental results over various model architectures and across multiple benchmarks show that our method significantly enhances expert specialization. 
Notably, our method improves classic MoE baselines with auxiliary loss by up to 23.79\%, while also maintaining load balancing in downstream tasks, without any architectural modifications or additional components. We will release our code to contribute to the community.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors address the challenge of poor expert specialization in standard MoE models trained with load-balancing loss. The authors argue that the conventional load-balancing auxiliary loss leads to expert overlap or uniform routing. The paper proposes to mitigate that using two additional loss functions: 1) an orthogonality loss to encourage experts to learn unique functions and 2) a variance loss to promote more discriminative routing. The newly added loss functions are used for finetuning three pre-existing MoE models without any architecture modifications. Results are presented for 11 benchmarks including language, reasoning, science, math, and code evaluation.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n- The paper is well-written and addresses a significant, well-known conflict in MoE training.\n- The results show the model outperforms alternative regularization methods for the finetuned models.\n- The ablation showing the impact of different loss functions on the behavioural states of the models is very informative.\n\n**Weaknesses:**\n- **Missing pre-finetuning baseline results**: The authors do not report the performance of the original, pre-finetuning MoE models. Based on public figures for models like Moonlight and DeepSeek, the finetuning process described in the paper appears to significantly degrade performance on key benchmarks (e.g., MMLU, BBH). For example, the original accuracy of Moonlight MoE model on MMLU, MMLU pro, MBPP, and BBH are 70, 42.4, 63.8, 65.2, which are above the best reported numbers in Table 1. Similar behaviour is observed for DeepSeek models: MMLU 58.3 compared to the 35.59 of the proposed method, or BBH 44.1 compared to 38.84. \n\n- **LoRA mitigating performance loss**: The authors focus exclusively on LoRA finetuning of high-quality, pre-converged models, but this crucial detail is confined to the appendix. Given the high performance of the original base models, it should be clarified whether their method is simply mitigating performance loss during finetuning rather than providing a net gain over the original models.\n\n- **Complexity of multi-objective training with 4 loss functions:** While the authors only focus on finetuning already converged high-quality MoE models, it is unclear how the complex optimization of four separate loss functions behaves when training MoEs from scratch or when upcycling dense models, which are more challenging and common scenarios. The paper also lacks a sensitivity analysis for the loss coefficients (`α, β, γ`), making reproducibility and application difficult.\n\n- **Computational Overhead:** The paper does not discuss the computational overhead of the proposed loss functions. In particular, the orthogonality loss, which requires pairwise comparisons of expert outputs for the top-k selections, appears prohibitively expensive, especially with the top-6 expert configuration used in the experiments. A discussion of the added training cost and memory is essential.\n\n- **Limited MoE architectures:** All three models evaluated (DeepSeek, DeepSeek-Lite, Moonlight) share an identical MoE structure (64 experts, top-6 selection, 2 shared experts). To strengthen their claims, the authors should test their method on more diverse architectures like Mixtral and Phi-MoE, which use different routing strategies and do not include shared experts.\n\nQUESTIONS:\nSee the section above.\n\nMinor point:\nLine 66: Two periods at the end of the sentence.\n\nLIMITATIONS:\nThe authors do not sufficiently address the limitations of their work, including discussion of the computational overhead for the added loss functions and a sensitivity analysis for MoE training with 4 loss functions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the issues of uniform expert routing and insufficient expert specialization in MoE models during the post-training stage, primarily caused by the auxiliary load-balancing loss. To address this, the authors propose two additional regularization objectives that respectively promote expert diversity and enhance router confidence. Experimental results demonstrate that the proposed approach improves model performance on downstream tasks.\n\nSTRENGTHS AND WEAKNESSES:\n### Strength:\n1. The idea of jointly optimizing expert specialization and routing distribution via additional regularization terms is interesting.\n\n2. The paper is well written, with clear motivations and a comprehensive discussion of the identified issues.\n\n3. The effectiveness of the proposed method is supported by empirical results.\n\n### Weakness:\n\n1. While Lemma 1 and Lemma 2 provide theoretical guarantees for the existence of a desirable routing configuration, their practical realization strongly relies on sufficient diversity in the input tokens. The paper lacks discussion on whether this assumption holds in real-world scenarios or how it may affect performance.\n\n2. The orthogonality loss \\( \\mathcal{L}_o \\), while effective in promoting expert specialization, introduces significant computational overhead due to its quadratic complexity with respect to the number of experts. Specifically, it requires computing pairwise projections between all top-\\(k\\) expert outputs for each input token, resulting in an overall cost of \\( O(N \\cdot k^2 \\cdot d) \\), where \\( N \\) is the token count, \\( k \\) the number of selected experts, and \\( d \\) the hidden size. This can become a bottleneck in large-scale MoE settings with many experts or long sequences. \n\n3. Currently, the proposed regularization losses are applied only during post-training. However, prior work such as OLMoE suggests that auxiliary losses may not be necessary during post-training at all. It would be valuable to discuss whether similar regularization should be incorporated during pretraining to promote expert specialization and more confident routing from the outset.\n\nQUESTIONS:\nSee Weakness.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper makes an observation that experts often lack specialization (e.g., the input gets routed to experts uniformly) due to load balancing, and introduces two auxiliary loss terms that encourage specialization without sacrificing load balancing – (1) an orthogonality loss, which reduces the similarity between experts that receive the same token, and (2) a diversity loss, which discourages expert weights from being uniformly distributed. Experiments on 11 benchmarks, four baseline training methods, and four base models show that the proposed model consistently and substantially outperforms all baselines, while maintaining effective load balancing.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- Overall, I found this paper to be very well written and strongly motivated. The subtle tension between expert specialization and load balancing is an important and underexplored issue, and this paper does an excellent job articulating it with clear and convincing explanations. While MoEs are widely studied in general, this problem of encouraging expert specialization is underexplored, and remaining having balanced load is very critical for practical reasons, e.g., training stability. Addressing both together is very impressive.\n- The paper also clearly explains the key observations and the technical proposals that follow from them. Admittedly, I did not carefully check every math detail line by line, but I was able to understand the underlying motivation and that the math equation aligns with that motivation.\n- Experiments are very thorough and comprehensive, quite beyond typical standards for a conference paper. Evaluating on eleven benchmarks, four baselines, and three base models, with consistent and substantial improvements across the board, is very impressive. The detailed analysis is also very valuable.\n\nWeaknesses\n- My most critical concern — despite all the strengths noted above — is that the actual training setup and data used are very unclear. First of all, it only appears in the *appendix* that the paper doesn’t actually do full training but it is doing LORA fine-tuning. This came as a surprise, as the rest of the paper is written as full training was performed. This is very different from taking the existing base models and applying LoRA tuning. This should be clearly stated throughout the paper — including the abstract, introduction, methodology, and experiments. And, this leads to a follow up question: how much (a small number of) additional training steps with LoRA fundamentally changes how experts are trained, as expert weights themselves won’t be updated?\n- In fact, even the fact that the paper is mainly targeting post-training was unclear, e.g., the abstract and the intro mentions the word “post-training” once each without much contextualization, so if readers miss this word, the entire paper is written as it is pre-training. Actually, is there a reason there’s a fundamental difference between post-training and pre-training, so that expert specialization is more important in post-training than pre-training, given that pre-training and post-training are actually the same training process except the data and where the loss is applied to are different?\n- Aside from these, other training setups are unclear either. Is the model trained in a multi-task setup using a single training dataset across all elevent benchmarks, or is there a separate training dataset and a separate trained model for each benchmark? How many training steps are used? What is the performance of the base model before LoRA fine-tuning (e.g., in Table 1)? These are important information that impact the implication of the methods and experimental results drastically, that are currently missing.\n\nMinor weaknesses/questions\n- Orthogonality loss: just to confirm my understanding – this loss encourages the experts that receive the same token to be less similar to each other, so that a single token benefits from greater diversity? If that’s correct, it feels more about “expert diversification” rather than “expert specialization”. To me, specialization implies that an expert becomes responsible for (one or more) domains or token types.\n- I think Figure 1 can significantly be improved - it is quite complex and is not intuitive enough.\n\nQUESTIONS:\n(Questions are in the Strengths And Weaknesses section)\n\nLIMITATIONS:\n(Limitations are in the Strengths And Weaknesses section)", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses a critical challenge in Mixture-of-Experts (MoE) models: the trade-off between load balancing and expert specialization. The authors argue that the standard auxiliary load balancing loss often leads to overly uniform routing and expert overlap, which hinders model performance, especially during post-training adaptation. To tackle this issue, the paper proposes a novel method that introduces two complementary loss functions: an orthogonality loss to encourage experts to learn distinct representations for different tokens, and a variance loss to promote more discriminative routing decisions. The authors provide a theoretical analysis demonstrating the compatibility of these new losses with the existing auxiliary loss. Through extensive experiments on three different MoE models and eleven diverse benchmarks, the paper shows that the proposed method significantly enhances expert specialization and improves downstream task performance by up to 23.79% compared to baselines, without requiring any architectural modifications or compromising load balancing.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- Novel and Impactful Contribution: The paper presents a simple yet effective solution to a well-recognized problem in MoE models. The proposed orthogonality and variance losses offer a principled way to directly encourage expert specialization, a crucial aspect for unlocking the full potential of MoE architectures. The potential for significant performance gains without architectural changes makes this work highly valuable for the community.\n\n- Comprehensive Experimental Evaluation: The empirical evaluation is a major strength of this paper. The authors have conducted experiments on a diverse set of MoE models (DeepSeek-Moe-16B, Moonlight-16B-A3B, and DeepSeek-V2-Lite) and a wide range of 11 benchmarks covering mathematics, multi-domain tasks, and code generation. The consistent improvements across different models and tasks provide strong evidence for the effectiveness and robustness of the proposed method. The inclusion of four relevant baselines further strengthens the experimental setup.\n\n- In-depth Analysis and Ablation Studies: The paper provides a good level of analysis, including a theoretical discussion on the compatibility of the proposed loss functions.\n\nWeaknesses:\n- Lack of Hyperparameter Sensitivity Analysis: The paper states that the hyperparameters for the new losses (β and γ) are set to the same value as the auxiliary loss weight (α=0.001). A sensitivity analysis or at least a discussion on the tuning process for these hyperparameters would significantly enhance the practical value of the proposed method.\n- Generalizability: The benchmarks are primarily language- and code-focused. Broader validation in multimodal or low-resource settings (hinted at in limitations) would strengthen the claims.\n\nQUESTIONS:\n1. On hyperparameter sensitivity: The choice of hyperparameters β and γ for the new loss functions seems crucial for the method's performance. Could you provide analysis on the sensitivity of your method to these hyperparameters?\n2. On compareing baselines: The paper compares against classical MoE methods (e.g., GShard, ST-MoE) and auxiliary loss-based balancing. Could you add more recent baseline results?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-05"}