Title: Dense Backpropagation Improves Training for Sparse Mixture-of-Experts
Venue: NeurIPS 2025
Accepted: yes
Keywords: moe

mixture of experts

router
Abstract: Mixture of Experts (MoE) pretraining is more scalable than dense Transformer pretraining, because MoEs learn to route inputs to a sparse set of their feedforward parameters. However, this means that MoEs only receive a sparse backward update, leading to training instability and suboptimal performance. We present a lightweight approximation method that gives the MoE router a dense gradient update while continuing to sparsely activate its parameters. Our method, which we refer to as Default MoE, substitutes missing expert activations with default outputs consisting of an exponential moving average of expert outputs previously seen over the course of training. This allows the router to receive signals from every expert for each token, leading to significant improvements in training performance. Our Default MoE outperforms standard TopK routing in a variety of settings without requiring significant computational overhead.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an improvement for training MoE routers by including the moving average output of non-selected outputs in the router gradient calculation.  This better approximates the dense gradients, particularly at lower layers, and leads to improved convergence and/or test accuracy.  Validation is performed on MoE models around 2B params (500M active params) using FineWeb and a suite of language model benchmark tasks, finding good improvements in accuracy, perplexity by training length.\n\nSTRENGTHS AND WEAKNESSES:\nThis is a simple idea that explained well and demonstrated effectiveness.  It's only compared to Top-K in an apples-to-apples comparison, but quite effective in boosting its performance there (they mention unrelated globally reduced load balancing loss as why comparison systems underperform, but I appreciate their candidness of the explanation here).  Similarity plots in the ablations between expert averages and the approximated vs true dense gradient show some of the behavior the method nicely.\n\nThere are still lots of questions I have around the method's behavior and potential for interactions with other selection or balancing methods, when used in combination.  For example, what happens if additional load balancing constraints are imposed while this gradient approximation is used?  Do they work together, or does one lessen the need for the other?  And since gradient error is introduced as a problem, is it possible to know more about its impact, and the impact of non-zero-mean vs zero-mean error in particular, other than just better/worse performance?  I've listed these and a few more in the questions section below.  While some performance and behavior were demonstrated well, I think many others remain unexplored.\n\nQUESTIONS:\n* While distances between the means is shown in the paper, I'm curious about the distribution of distances of each actual output to each mean (including the distance to the expert that produced that output) --- how compact or diffuse are the clusters?  What are the distances of each expert's outputs to each avg?\n\n* This zero-centers the error of router gradients for non-selected experts.  But what are the impacts of non-zero-centered error, other than just worse performance?  In a case where all expert outputs are near their means, as a possibly illustrative case, what does the error from zeroing experts do to the update (not just the grad), and what behaviors can this contribute to in terms of future expert selection, selection collapse, etc?\n\n* How does this interact, if at all, with other expert selection tricks like imposed balancing constraints or stochastic selection/gradient estimators?\n\n* The average is computed only over x for which the expert is selected.  But then it's being used in the grad approximation for inputs in which it is *not* selected --- and its behavior could be very different for these x.  I'm not sure how much this matters, other than maybe the possibility to tend to flatten the router distribution, since the inputs that are most affected are those closest to being selected by the expert anyway, and thus closest to its average expectation's distribution.  It's a bit of a fine point, though, and I wonder what you think about this seeming inconsistency?\n\n* Relatedly, for x that are close to the boundary of enabling an expert from its mean comparison, if the new expert is enabled, how close are the actual outputs to the mean?  In particular, if we calculate the distance between the output x would have had for expert K+1 (sorted by softmax score), how does it compare to distance for experts 1..K, and in particular expert K, which K+1 will replace if selected?  Are these close, so that there could be oscillation at the edges?  And what are the effects of this --- it could smooth expert selection, but at some cost to specialization, or it could induce more clustering around the means to improve specialization, but right now the precise behavior is unclear.\n\n* As mentioned above, to me it seems this has the potential to flatten the router distribution.  Though this doesn't affect top-k, it might impact other selection methods (though temperature could be adjusted).  Have you seen anything like this?\n\n* l.187 \"we can actually automatically account for the impact of sparsity and 188 granularity on the default vector update by weighting the updates to the default vector by the router 189 logits.\" --- This seems important and interesting enough to go in the main text in method description.  I can imagine some sort of heuristic that uses a different parameter indicating the weight of the ema update per overall training token, then calculates beta depending on how often the expert is selected.  I actually haven't found a description of this yet looking over the appendix, either -- how does this work?\n\n\n\nSmaller questions/comments:\n\n\n* Paragraph at l.240 about globally reduced load balancing loss:  This paragraph didn't make any sense to me, but seems important.  I think it could be more clearly written, with more concrete descriptions of what distributions or behaviors change and in which directions, with a small summary of the prior work and its effects.\n\n* Fig 2d:  1/64 sparsity --- this jumped out as strange, but explained in the caption that it's for a different smaller model.  If it isn't a comparible point, I think it would make sense to either leave out or at least separate and better indicate in the figure it's for a different model\n\n* Table 2: in row for PubMedQA, 54.8_2.2 appears exactly in 3 of the 4 score cols --- is this correct or is there an error here building the table?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDefault-MoE addresses lack of gradient flow for MoE routers' unactivated paths by backpropagating a proxy output of such paths/experts. That default expert-output used for backpropagation, as the actual output is unavailable due to tokens not being routed to said expert, is a weighted average of historical expert outputs.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths:\n- Clearly written paper with an easy-to-follow flow.\n- Experiments are done in a more realistic scale with FineWeb.\n- In-depth ablation studies and method specifications.\n\n### Weaknesses:\n- It's perhaps clearer to pose the motivation to be the approximating truncation error that Sparse MoE incur with respect to Dense MoE (Eq. 7) in order to achieve Dense MoE's level of _performance_.\n  - Personally, I find the current argument to imply that SMoE is a simple practical truncation of MoE born from an engineering standpoint, while in reality all of the theory (e.g. equations in Sec. 3.1) can treat SMoE as a formal router option.\n  - If the goal is to estimate the gradient of the full Dense MoE (L266), a comparison with the Dense MoE variant on perplexity/computational overhead would be relevant.\n- While experiment size is relatively large, diversity is rather lacking.\n  - Default-MoE is only compared with Sparse MoE formally.\n  - For SparseMixer, only the training perplexity is reported, and only for the first 10B tokens. While it does seem like the test set perplexity would follow the trend of the training curve, actual numbers for a converged training would greatly improve persuasiveness.\n  - The same applies with Loss-Free Balancing, but with no plot this time and only a brief mention in L239. While I understand that the wasted resources in reproducing ReMoE would have prevented the replications of other baselines, the paper still needs to be a complete evaluation instead of a list of individual experiments.\n- As router entropy is considered, it would be great if there is also additional analysis on router fluctuation.\n- \"Default MoE\" naming is easy to be confused with Dense MoE.\n\nOverall I think the paper would be great if all of these concerns are addressed.\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nEfficiency concerns are listed in Sec 4.6. No societal impact was addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

comments: SUMMARY:
This work investigate the problem of ineffective learning of non-selected experts in Sparse Mixture of Expert (SMoE) training. The authors noticed that in the traditional SMoE, gradient of non-selected experts are zero, and proposed to replace it with the running average of the experts' activation, which is a non-bias estimator. This result in DefaultMoE, and has shown to achieve better performance than the traditional TopK routing on the pre-training task with two architectural choices.

STRENGTHS AND WEAKNESSES:
Strengths
- Overall this paper is easy to follow, the method is well-motivated and justified. The empirical study is also comprehensive and convincing.  

Weaknesses
- Cross-validating $\beta$ could be costly.
- Although it is nice to see the authors compared 2 architectural variants, I would prefer to see a few more routing strategy baselines such as XMoE [A], MoEUT [B] , or AoE [C]

[A] Chi, Zewen, et al. "On the representation collapse of sparse mixture of experts." Advances in Neural Information Processing Systems 35 (2022): 34600-34613.

[B] Csordás, Róbert, et al. "Moeut: Mixture-of-experts universal transformers." NeurIPS 2024

[C] Lv, Ang, et al. "Autonomy-of-Experts Models." ICML 2025

QUESTIONS:
Suggestion: Together with Figure 11, it would be nice to plot the average zero-shot evaluation scores throughout training.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nOne of the problems of standard Sparse Mixture-of-Experts (MoE) models is that they are discontinuous functions _and_ the gradients of the router are very sparse, which difficults its training. The paper suggests to learn a \"default output\" for every expert  (i.e. the average output among the tokens that selected it) in a Sparse MoE model, and use this as the output of an expert when it is not selected. This way, the gradient of the router parameters corresponding to the inactive experts still get non-zero gradients, due to the (non-zero) default output. The proposed method is simple to implement, and does not add any significant overhead (neither time or memory) to the standard MoE model.\n\nThe paper experiments with several models (with 8 and 32 experts, and a different number of experts/token), and presents results on a varietay of well-established benchmarks (e.g. MMLU, OpenBooksQA, Lambada, ...) showing that it typically (but not always) outperforms the scores of standard TopK routing.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The paper is very well motivated, structured, and written. It is (generally speaking) very easy to read and understand.\n- The proposed method seems very simple to implement (and the authors have released the source code), which allows other researchers to quickly adopt the idea, re-evaluate it, and build on top to improve it.\n- The method improves the standard TopK routing, which is still essentially the most used routing algorithm among state-of-the-art MoE models, over a wide range of datasets and models (with some caveats). \n- The proposed method seems very robust to different hyperameters (learning rate, choice of EMA's $\\beta$), which also increases the chances of successfully adopting it for other models & datasets.\n- The paper includes lots of ablation and analysis experiments, measuring the importance of different design choices and empirical obsergations. I want to highlight, in particular, the experiment in section 4.5, showing that the proposed method provides gradients that are more similar to those of a dense MoE, than the ones obtained when using standard TopK (see figure 4, left). This shows that the proposed approach is indeed a good method for approximating the dense gradients of a MoE.\n\n**Weaknesses**\n- My main concern relates to the observation detailed in lines 196-199 and figure 2a. It seems that the proposed method achieves better results, compared to standard TopK, when the sparsity is not too high. However, given that when $K \\rightarrow N$, both methods are equivalent (because no default vectors are used), one would expect the proposed method to be more beneficial with sparser models, given that TopK does a very bad job approximating the gradients of the full (dense) MoE, as shown in Figure 4a. This raises concerns about the ability of the proposed method to generalize to larger and sparser models, and begs the question: how do the authors explain this?\n- Some details of the proposed method are not clear. For instance, it was observed that after \"weighting the updates to the default vetor by the router logits\", the value of $\\beta$ is not as important (see lines 187-189, and Figure 8). However, I couldn't find any equation or detailed explanation of the actual \"weighting\" being used.\n- The paper uses EMA to compute a running average of the output of every expert based on the tokens that activated it. In Lines 130-132, the paper states: \"our method ensures that computed expert outputs at the current step are factored into the EMA update _before_ applying the EMA as a substitute for other tokens that did not activate the expert\". This is problematic with auto-regressive models, since the natural dependence of tokens may be broken. The proposed method still shows superior performance than TopK after pretraining, so it seems that it didn't learn to exploit this. However, one wonders if models with bigger capacity could eventually learn to exploit this, and why not simply preventing it by using the _previous_ EMA state as the experts' output of a given step, before updating it.\n- The statement in lines 262-264 is not obvious to this reviewer, one might think that the opposite is actually true. Let's suppose that the (conditional) entropy of the router is 0, then every expert will process a very distinct subset of tokens, so the different default vectors will be different. If the entropy of the router is high, very tiny differences in a given input token may cause the router to choose a different expert, which implies that all the experts may be processing relatively similar tokens, and then all default vectors would be similar too. However, it is also true that if the conditional entropy is close to 0, the importance/contribution of the default vectors to the output of the MoE is also 0 (because only the output of the top1 expert matters).\n- The \"improvement (%)\" columns in Table 2 are not clear at all. What is the baseline for this \"improvement\"? 8c1 results on Lambada show % of almost 4000, is this correct?\n- Figure 2d is a bit confusing without reading the caption. Perhaps labeling each point with the number of activated & total parameters of the corresponding model would be useful.\n\nQUESTIONS:\n- You experimented a lot with EMA to learn the default expert outputs. Did you ever consider using the same optimizer as the rest of the parameters by adding an auxiliary term to the NLL loss? For instance, something like this:  $\\mathcal{L} = \\mathcal{L}_{\\text{NLL}} + \\beta \\sum_i (\\hat{E}_i - \\text{stop-gradient}(\\mathbb{E}_x[E_i(x)]) )^2$\n- Typo in line 240? \"dMoEs\" -> \"MoEs\"?\n- Typo in Figure 4's caption: \"dense router gradient (K=8)\" should be \"(N=8)\", since it's the total number of experts, I assume. No?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-11"}