Dense Backpropagation Improves Training for Sparse Mixture-of-Experts
Authors: Ashwinee Panda
Vatsal Baherwani
Zain Sarwar
Benjamin Thérien
Sambit Sahu
Tom Goldstein
Supriyo Chakraborty
Venue: NeurIPS 2025 poster ()
Keywords: moe
mixture of experts
router
Primary area: deep_learning

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.

Dense Backpropagation Improves Training for Sparse Mixture-of-Experts
Ashwinee Panda
Vatsal Baherwani!
Zain Sarwar Benjamin Therien 2,4
Sambit Sahu²
Tom Goldstein
Supriyo Chakraborty²
University of Maryland
Capital One
University of Chicago
Mila Quebec Al Institute
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 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 exponential moving average of expert outputs previously seen over the course 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.
Introduction
The sparsely activated Mixture-of- (MoE) Transformer architecture [Shazeer et al., 2017] has been used by many industry deployments [DeepSeck-Al Team, 2024b, xAI, 2024, Databricks, 2024, Mistral Team, 2024, Snowflake, 2024, DeepSeek-At AI Team, 2024a] because MoEs have been shown to scale even better than dense Transformers [Clark et al., 2022, Du et al., 2022, Lepikhin et al., 2020, Fedus et al., 2022]. MoEs learn a routing function that selectively activates the Top-K subset their modules, experts, most relevant et al., 1991, Jordan and Jacobs, 1994] allows us to multiplicatively increase the model parameter count without significantly increasing the cost of training inference. At train time, the sparse router enables very large MoEs to be trained with relatively little computation per token, but also presents challenge. The router does receive gradient update from experts that it does activate. This may slow down learning, as the gradient update is unable to adjust the router to promote the optimal expert for each token. This can be corrected by activating all experts during training, enabling all experts and routing directions to be optimized at once, although this requires much more compute.
given input. This conditionally sparse activation [Jacobs
In this work propose DefaultMoE: new method that strikes middle ground between dense and sparse routing; our proposed router can receive and differentiate contributions from all experts, while the computational cost remains virtually identical to standard Top- router. In our proposed method, the router receives contributions from all experts for every token. However, a forward pass is only computed the Top-K experts for each token, while other experts contribute "default vector" that represents the expected output of a typical token. Our method adds minimal computational overhead compared standard Top-K router while improving performance range of standard benchmarks for 2 billion total parameter MoEs trained on 160 billion tokens.
** denotes equal contribution. Correspondence to: ashwinee uudd.cdu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2 Background & Related Work
MoEs. The MoE layer replaces the feedforward networks (FFN) of Transformers and consists of two components: 1) N FENs (experts), Ex(x), (r),.... (x(x) and 2) a router that assigns tokens to experts. Each input to the MoE layer is processed by K experts where K < N, and this is the source of sparsity in MoEs. The K experts are chosen by the router, which a learnable component that maps each token to set of weights over the experts. The router performs linear transformation Rd + RN' which produces logits; these are normalized using softmax, resulting in a probability distribution over the experts. With the router's linear transformation parameterized by matrix W. we can represent the expert weights т in the following way:
π E RRA. Softmax xW
Once we have these expert weights, we apply a routing function to decide which of K experts to route and process this token through. We consider Top-K routing because it is the most popular.
Top-K routing. standard method to select K out of N experts given the expert weights is to select the experts corresponding to the K highest weights. Top-K K routing [Fedus et al., 2022] passes the token to the K selected experts and averages the expert outputs using these weights to produce the final output. Experts not selected by the Top-K routing function do not process the token, and this introduces sparsity MoEs. By representing the K chosen experts as the set A, we can express the output of the MoE layer an average of expert outputs weighted by the router scores:
y = Ex())
(2)
The expert weights serve two roles. They used by the routing function to decide which of the K experts to process a token through, and also provide the weights for combining the expert outputs. Top-K routing makes the MoE layer desirable for training large, compute-efficient neural networks. allows models to be scaled up, by way of increasing the total number of experts, while keeping the compute per token constant (as it is function of K and not N).
The Router Gradient. Consider the gradient of the MoE layer's output y with respect to the router parameters W. We express y a function of W by combining Equation (1) and Equation (2). With the chain rule, we can backpropagate through this function by considering the gradient at each respective step:
Oy θy θm OW
(3)
3x OW
The second term in Equation (3), JW is straightforward to compute because the steps in Equation (1) are easily differentiable, as they consist of linear operations and activations. But the first term, är is not differentiable because in Equation (2) Top-K expert selection transforms the continuous router weights E RV into a discrete set of selected experts A with (x) possible values. One way to address backpropagation of nondifferentiable operations is to use the straight-through estimator [Bengio et al., 2013], which treats the operation as the identity function. With straight-through we bypass the Top-K K routing function and Equation (2) becomes the dot product between and the vector of all E (z) with the following gradient:
dy θm
= [F) (2(x) ... [x(
(4)
This dense gradient requires the output of all of the experts for given token. Passing token through all the experts will destroy the sparsity of the MoE layer, thus impeding the scalability of this architecture. In this work, we develop a method for applying the ttraighttthrouge estimator while maintaining the sparsity of the MoE layer by substituting the non-activated expert outputs in the dense gradient term with a default vector. This default vector is a running average of previously computed expert values, and it approximates the missing expert output without incurring the cost of forward pass. Notably, this enables non-activated experts to contribute to the router's 's gradient update.

---

Not Rosted
Not Routed
Added Gradient
X 0.3
x 0.4
X
X 0.4
Ej
3
E4
E
B
E3
0.3
0.1
0.4
0.2
0.3
0.1
0.4
Token
Router
Token
Router
(a) Original Router
(b) Dense Approximation Router (ours)
Figure 1: Overview of Routing with Dense Approximations. The original MoE router only receives gradients corresponding to experts the token is routed to, because there is no output from other experts. Our approach provides the router with complete (dense) gradient by letting non-activated experts contribute a defauls vector that approximates its output without the cost forward pass. As indicated by the dashed green arrows, the approximated gradients are actually connected the token in the computation graph; instead, they are artificially applied in the backward pass.
Designing Dense Backpropagation
In this section we design new MoE router that can receive a dense gradient without requiring additional forward passes through experts. Equation (4) describes the dense gradient of the router for an individual token. Each expert output (x) in the router gradient corresponds updating the embedding for expert i in the routing layer (i.e. the ith row of the router parameters W). In standard sparse MoE, the gradient only includes expert outputs E, for the experts E AA selected by the Top-K K routing function. The other terms are essentially O, which means that non-activated experts can not influence the router's update for a token.
Our goal is to approximate the dense gradient in Equation (4) without directly computing the missing terms Eg (x) for i' & AA. This enables the router receive signal from all experts its gradient update, without actually activating all of these experts. The router can then consider information from all experts when learning to route tokens instead of being limited to only receiving feedback from experts that were activated. As result, we expect an improved router which learns to allocate tokens to different experts more effectively, leading improved MoE training performance.
Figure 1 presents an overview of our method for updating the router. Since we wish to avoid additional computational overhead, we fill in the missing terms in Equation (4) with estimates É, for non-activated experts E,, This simple addition ensures that the router updates its weights for all experts instead of only those selected by the Top-K K routing.
3.1 Approximating Missing Expert Outputs
Backpropagating through the sparse MoE activation in Equation (2) does not align with the router's s true dense gradient specified in Equation (4). Specifically, the gradient in TopK routing leads to an error based on the non-activated experts. Let us consider the gradient of the router parameters W with respect the loss C. For clarity, examine the gradient term for single input x; our approach also translates to practically training the router with a batch of inputs.
Using the chain rule, we decompose the router gradient to isolate the nondifferentiable By Equation (3):
similar to 8
ac ac dy ÔT
(5)
aW Dy Om aW
3

---

In standard Top-K routing, the gradient term effectively uses zero for non non-activated experts:
dy
(x
i A if
(6)
θT:
By The error is introduced by the missing terms in
UL
TTopK dy
(x)
aw
(7)
UEA
To address this error, we need a meaningful estimate for (x) without activating the non-activated experts Eg. Our approach to this involves maintaining default vector for each expert É; = [x)] that represents the expected value of the expert output. In practice, we will compute this estimate by taking the sample average existing expert outputs Ex(x') from other tokens x' for which E, was actually activated. Considering this estimator, our gradient term for non-activated experts now nonzero:
dy θT;
(x(x) = E.
if A i
(8)
Note that É, is not function of æ; it is universally applied to complete the router gradient for all such tokens in the batch. The gradient error term with this new default vector is as follows:
Oc
Ôπ
Edefault
Oy 1
Σ (x)
-E, [E)(())
Ow
(9)
By The error term corresponding to 3i (F,((((()) - (x))) is 0 in expectation for all i. As a result, default vector method corrects the gradient by filling in missing expert activations with the mean expert output.
3.2 Default Expert Outputs with Exponential Moving Averages
To enable dense backpropagation without sacrificing the computational efficiency of sparse MoE forward passes, we propose Default MoE. Default MoE provides a default value for the outputs of non-selected experts using exponential moving averages (EMAs). The EMA accounts for the fact that experts are being updated, and thus their outputs for inputs x will change over the course of training. For each expert E,, we maintain an EMA of its average output:
ES0) = -1- (x)
(10)
where 9 e [0, 1] is the decay rate and (x) the sample average of expert outputs E, for all tokens x for which expert was activated. We compute Ex(x) during the forward pass by collecting all such expert outputs Exr) and taking the average. This operation is computationally trivial, since the outputs themselves are already computed as part of the standard MoE forward pass. Thus, the EMA serves lightweight proxy for the expert's expected output.
During the forward pass, for given input x, and an expert (where there are N total experts) compute the output of the MoE as:
El(x)
if TopK(T)
T
(11)
otherwise
Specifically, we first compute the forward pass for activated experts F.(x:) for each token. Then, we perform the EMA update step using these outputs. After applying the EMA update, we use f to approximate missing outputs for non-activated experts. Our method ensures that computed expert outputs the current step are factored into the EMA update before applying the EMA as substitute for other tokens that did not activate the expert.

---

This formulation allows the router to receive meaningful gradients for all experts while maintaining the computational benefits of sparse activation. The EMA provides reasonable approximation of what non-activated experts would have computed, based on their historical outputs for other tokens. Importantly, this approximation requires only (11) additional memory per expert and requires minimal additional forward pass computation for the EMA update.
This approach enables dense backpropagation through the router while preserving the sparse com- putational pattern that makes MoE architectures efficient. The router receives gradient information about all possible routing decisions, not just the selected experts. We now evaluate DefaultMo..
Evaluation
We describe the experimental setup in Section 4.1, benchmark our Default MoE Section 4.2, show that our improvements hold across multiple configurations in Section 4.3, discuss prior work in Section 4.4, analyze our Default MoE in Section 4.5, and find in Section 4.6 that our method does not significantly impact throughput or the memory footprint of training, so the improvements we observe are -as far as we can tell- free lunch. We conclude with a cohesive explanation of why Default MoE beats TopKMoE in Section 5.
4.1 Experimental Setup
Model Architectures. We use NcK to refer to a model with total experts and K active experts. We train both standard and finegrained [DeepSeek-AII Team, 2024a] MoEs. All of our models have 1.96 billion total parameters, 366 million of which are non-MoE parameters. This leaves 1.6B MoE parameters, and the number of active parameters depends sparsity. All parameter counts specified in Table 1. We ablate the model architecture, hidden dim, number of total experts, and number of active experts in Section 4.3.
MoE Config.
8c1
8c2
32c1
32c2
32c4
Active Params.
565M
764M
416M
466M
565M
Table 1: Active parameter counts for each MoE configuration; the total parameter count is 1.96B.
Dataset. We train on FineWeb-Edu [Lozhkov et al., 2024] (for Table and Figure 11) and WibeWeb [Penedo et al., 2024] (for all other plots and results) with the Llama3 tokenizer [Llama Team, 2024]. For our main results, we train for 160B tokens, which is tokens-per seeens-per-paaeeee per-parameter ratio of ~ 283; at this level of overtraining, many spurious improvements should wash out, so we can be confident that any improvements we observe are real gains.
We have open-sourced our training code. All hyperparameters for the models trained in Table can be found in our config file, and all hyperparameters are the same between the bascline and our DefaultMoE. All further experimental details can be found in Appendix A.2.
4.2 Main Results
Pretraining Benchmarks. Table we compare Default MoEs to TopK MoEs. All models are trained for 160B tokens on FineWeb-Edl and have 1.96B params. For both finegrained MoEs (32c4) and standard MoEs (8c1), our Default MoEs outperform TopK MoEs on standard benchmarks. We conduct all evaluations with the Im-eval harness [Gao et al., 2024].
Due to space constraints we defer pretraining curves to the Appendix, but discuss them here. In Fig- ure 11 we compare each MoE's 's loss curves. Without introducing significant overhead, our method reduces the tokens required reach a target perplexity of 12.18 by 9%. As the model has just 500M active parameters, we are training for significantly more than the compute-optimal number of tokens Hoffmann al. [2022]. Therefore, we can confident that we are not just seeing our method converge faster to worse minima; this is a real improvement in terms of both speed of convergence and the quality of the converged model.
We use the 'Random word from passage" baseline for LAMBADA [Papemo et al., 2016].
5

---

32el
Random Baseline
Benchmark
Topk
Score Default
Difl (%)
Topk
Improvement (5)
Default
Dil (8)
Topk
Default
Score
Tupk
Impervement (9)
Delault
Diil (3)
(Seone)
LogiQA
26.0..
26.9..
+3.6%
3.8
7.5
+96.0%
27.2
27.8..s
-2.3%
8.8
11.2
+28. 1%
25.0
MathQA
26.2as
8o.s
24.4u.
36,
+8.0%
-2.4
MMLU
31.8
32. 32.3o.4
+1.6%
27.3
29-4
+7.5%
33.5o.4
32.46.4
-3.15
33.9
29.8
-12.3%
25.0
QpenBookQA
37.2g2
38.22.
+2.7%
48.8
52.8
+8.2%
34.2g.1
39.42.2
+15.2%
36.8
57.6
+56.5%
25.00
Lambada
41.2p.
1761
4016.0
43.8o.7
-1.1%
4374.6
4276.1
-1.1%
1.0"
SacialIQA
41.0...
33.2%
24.2
119.1%
41.4
4.
-1.6%
25.3
HellaSwag
40.46.
41.2%
+2.0%
61.6
64.9
+5.4%
43.1a5
43.40.5
+0.7%
72.3
73.5
+1.7%
25.0
ARC
45.70m
47. 47.40.
+3.7%
82.8
89.6
+8.35%
49.6p.
49.40.9
-0.4%
98.4
97.7
-0.75
25.0
Winegrande
53.61.4
54.8
+2.2%
133.0%
57.8,4
55.61.4
-3.7%
15.5
-27 4%
PubMedQA
55.2aa
52.4g.a
-5.1%
58.8
-1.66
58.4.
57.6
-1.4%
77.0
BoalQ
58.50n
62.06.
+6.18
17.0
24.1
141.7%
61.30.
61.40.
-0.2%
22.5
22.8
+1.4%
50.0
PIQA
70.9
71 4:.1
+0.6%
41.9
42.8
12.1%
71.82.
73.11.0
+1.8%
43.6
46.2
16.0%
50.0
S:0
86.2
87.8
+..9%
244.8
251.2
+2.6%
89.0,
90.
-1.3%
256.0
260.8
+1.9%
25.0
Average
46.96c
47.9%
+2.0%
52.3
54.8
+4.8%
48.968
49.
-0.9%
57.3
59.5
+3.8%
Table 2: Comparison of Pretraining Benchmark Scores for TopK MoE and Default MoE. After
training for 160 billion tokens, Default MoE outperforms TopK MoE across a range of benchmarks for both a standard MoE (right) and a finegrained MoE. (left).
Pretraining on Web. In Table 2 and Figure 11 we pretrain on FineWed-Edu [Lozhkov et al., 2024]. We also pretrain models on Web [Penedo et al., 2024], which is less curated. The rest of the plots in this paper (our ablations) are on FineWeb, in order to enforce that on some level we are not simply overfitting to the pretraining dataset that we use to collect benchmark scores. In Figure 12 we show that when using the learning rate of 7 x 10-4 that we sweep Figure 2, our DefaultMoE outperforms the baseline TopK. Therefore, our improvements are not due to any unfair hyperparameter tuning when compared to the baseline.
4.3 Ablations
We now ablate every setting in our experimental setup, from the MoE config. to the learning rate.
Tuning 3. We initially tuned 3 for every configuration independently, and the results of this can be found in Appendix A.4. We found that more sparse configurations. such as 32c1, required a lower 3 = 0.65 whereas a much higher beta of B = 0.999 was necessary for 32c4 to achieve its best performance. However, we can actually automatically account for the impact of sparsity and granularity on the default vector update by weighting the updates to the default vector by the router logits. When we do this, different values of B, that would otherwise be suboptimal, all converge to the same good performance. Therefore, in our final method, the value of / does not seem to matter as shown in Figure 8.
MoE Architecture Ablations. In Figure 2(a), we train both a standard TopK model and Default MoE with the configurations 8c1, 8c2, 32c1, 32c2, and 32c4, which correspond sparsity factors of 1/8, 1/4, 1/32, 1/16 and 1/8 respectively. Default MoE outperforms a standard TopK MoE at all sparsity configurations. Our improvement more pronounced for the lowest sparsity (8c2) but still significant for higher sparsity MoEs such as 32c1. Notably, Default MoE takes longer to "warm up" for sparser models; for example, we notice in Figure 13 that the 32c1 Default MoE rapidly catches up and surpasses the TopK MoE around the 10 billion token mark, while the 32c2 MoE is clearly ahead of Topk even at billion tokens.
Tuning the Learning Rate. We want to ensure that we are comparing against a tuned baseline, so we tune the learning rates in Figure 2(b). We use the learning rate that achieves the best performance for the TopK baseline, 7 x 10-1, for our results in Figure 12. We observe that while it performs better than Default MoE earlier in training (across our learning rate sweep experiments). the results in Figure 12 show that even in this case, Default MoE eventually surpasses TopK. Moreover, we note that the best learning rate for the Default MoE is the larger learning rate of 10-4 By contrast, 9 x 10-4 is much too large for the baseline. This indicates that our Default MoE is more stable to train, likely because we are updating the entire router. we train the baseline with such large learning rate, we see that single iteration with a very imbalanced load will lead to a large loss spike; this iteration is very noticeable because, since we train dropless MoEs, it is also slower. We never observe this for the Default MoF. It is unsurprising that our Default MoE's best learning rates may be different from the baseline, but importantly, our method outperforms the baseline at all learning rates consider.

---

Default MoE Remains Superior As We Increase the Model Size. Figure 2c compares the two methods as we increase the size of the model from 557M total parameters hidden dimension of 512, to 7.33 total parameters when the hidden dimension is 2048. Our method outperforms the baseline across all model sizes. All other results in the paper use a hidden dimension of 1024.
23.0
TopK Default MoE
Default Vector
Final Perplestty
TopK
22.67
Perplexity
22.5
22.32
22.00
22.0
22
23.54
21.5
21
20
3e-04
5e-04
7e-04
9e-04
3264
32c2
32c1
8c2
8c1
Sparsity Configuration
Learning Rate
(a) MoË Configurations
(b) Learning Rate Sweep
32.5
TopK
Default MoE
30.0
Perplexity
Perplex
20
25.0
22.5
18
Defoult McE
20.0
0.57
1.96
7.3
1/4
1/8
1/16
1/32
1/64
Total parameters (Billion)
MoE Sparsity Factor
(c) Ablation over model parameters
(d) Sparsity Factor Ablation
Figure 2: Comprehensive Ablations. We report the perplexity after training for 10B tokens for Default MoE and Top-K while varying: (a) MoE configuration across five MoE setups. (b) Learning rate while keeping the MoE configuration at 8c1. (c) Total model size while keeping the MoE configuration at 8c1. (d) Sparsity, where the final 1/64 sparsity is for a 7B model, therefore the perplexity is lower than the other sparsity factors that are for 2B models.
We present additional ablations regarding the EMA, including the hyperparameter 3, the EMA initialization, and passing the EMA value forward, in Appendix A.4.
4.4 Comparison to Other Routing Methods
We consider three alternatives to TopK routing for comparison: Sparesemixer [Liu et al., 2023, 2024], ReMoE [Wang et al., 2025], and Loss-Free Balancing [Wang et al., 2024].
Default MoE Beats Sparsemixer. Liu al. [2023] proposes Sparsemixer, which estimates the true router gradient without straight-through. Liu et al. [2024] note that Sparsemixer lags behind TopK (which our method always outperforms) for the first 0.57 tokens, likely due to the noise that Sparsemixer adds. After 24B tokens, Sparsemixer obtains validation perplexity of 19.77, worse than both TopK and our method. We present the training curve in Figure
ReMoE Fails to Converge. We cannot compare our zero-shot benchmark scores to ReMoE because they train on The Pile. which unavailable. We train ReMoEs FineWe--Edu. ReMoE reports that the first stage of their method is expected to run slowly while their method starts to work, as it is essentially acting as a dense model for the first few hundred steps. Following that, they report that ReMoE exhibits a significant loss spike around billion tokens, after which recovers and starts to outperform the baseline. In our setting, we indeed were able to replicate the first two results, where ReMoE runs very slowly while acting as a dense model, but the sparsity coefficients start to converge, and the loss spikes at 2 billion tokens. However, the loss never recovered. We invested over 1000 GPU hours into trying to get the ReMoE model to train, but did not have any success.
7

---

Loss-Free Balancing Underperforms. We first sweep the bias update coefficient, which is the hyperparameter in the loss-free balancing method used Wang et al. [2024]. and find that 0.0001 provides the best results. However, after training the 8c1 MoE for 10B tokens, Loss-Free Balancing obtains validation perplexity of 24.22, whereas our Default MoE. obtains 23.48.
Why do Prior Methods Underperform? Prior work claims to improve upon standard dMoEs, but there is one difference in our experimental setup: we use globally reduced load balancing loss [Qiu et al., 2025] in all experiments. Corroborating their conclusions, we find that globally reducing the auxiliary loss significantly improves the load balancing, and subsequently the performance, of the baseline, therefore making the improvement of prior methods less noticeable. This is because the "free lunch" of slightly improving routing is more less eaten by the global auxiliary loss reduction, which is trivial to implement and uses little compute, and prior work frequently targets this same free lunch as their source of improvement. Qiu al. [2025] directly evaluate Loss-Free Balancing and report that it underperforms proper globally reduced auxiliary loss, supporting our results.
4.5 Empirical Analysis
In Appendix B validate that the load balance, expert coactivation, and domain specialization of the trained MoEs do not deteriorate when we implement our DefaultMoE. Our results corroborate those of Qiu et al. [2025], that even bascline model trained with globally reduced auxiliary loss will have good load balance and domain specialization.
1.00
Layer
Layer
13
15
9
7
18
19
5
20
4
21
22
23
In(2) In(4) In(8)
-1.00
Router Entropy
Figure 3: Analysis. Left: Router Entropy after training for 10B tokens. Middle, Right: Heatmap of the cosine similarities between default vectors for the 8c1 DefaultMoE layers and 16.
In Figure 3(a) we observe that the router entropy highest for the early layers and consistently goes down, with the final layer's router entropy being very low. Intuitively, the router shouldn's 't quite know where to send a token at the very start of the model, but the routing decisions should become more and more fixed as proceed through the model.
Similarity of Default Vectors. Figure 3(a) we plot heatmaps of the cosine similarities between default vectors for experts in different layers of the 8c1 DefaultMoE. The default vectors at layer O are not very similar, indicating that the default vectors are learning something nontrivial. However, at layer 16 the default vectors are actually very similar. This follows the trend of the decline in router entropy of Section 4.5; the router very confident where token should go, the default vectors may learn similarly trivial things, but if the router has high entropy, the default vectors should learn faithful nontrivial approximations of the expert activations.
Similarity to Dense Gradient. Our entire design of DefaultMoE is based on the premise that we can better approximate the dense gradient corresponding to simultancous activation of all the experts. In Figure 4, we validate that our DefaultMoE's router gradient is more similar to the dense gradient than the TopK MoE's router gradient is. We take the 8e1 MoE pretrained on 160B tokens and compute the dense router gradient, by activating all the 8 experts. We then vary K between I and 7 (the full range can be found Appendix B) and report the similarity of the sparse router gradient to the dense gradient, for both DefaultMoE, and TopK MoE. At K=1, DefaultMoF: is clearly much more similar the dense gradient for the first few layers, while TopK is not close at all. This is because the routing distribution has high entropy at the early layers, so in DefaultMoE, the early
8

---

Similarily to Dense Gradient at K=1
Similarity to Dense Gradient at K=7
14
12
11
10
14
12
10
15
7
18
18
19
19
5
20
20
21
3
21
22
23
Topk Moll
Topk Moll
Defoult MoE (Ours)
Default MoE (Ours)
Figure 4: Dense Router Gradient Similarity. We plot similarity to the dense router gradient (K=8) for our 8ck Top-K and DefaultMoEs pretrained on 160B tokens. The DefaultMoE's router gradient is much more similar to the dense router gradient.
layers will make significant use of the default vectors. This increases the gradient flow to the router which is reflected in the closer approximation of the dense gradient. Whereas even if the router is allocating near-equal mass all experts and TopK MoE only selects one, its router gradient will be quite dissimilar from the dense gradient. Even at K=7, TopK MoE is still not quite approximating the dense gradient, indicating importance of the missing information from the least activated expert.
4.6 Efficiency
Our Method Does Not Significantly Reduce Throughput. When we train the 7.33B 8-expert MoE, the per-node GPU throughput is 1.393 tokens per second for TopKMoE and 1,391 tokens per second MefaultMoE, with a variation of about 1-2 tokens per second on each of these numbers; this is throughput decrease of just 0.18% which is just variance. For larger models, more of the time will be spent in matmuls and the overhead of our method will be insignificant.
Our Method Does Not Significantly Increase the Memory Footprint. The only additional memory required by our method the EMA buffers themselves. For each expert in each layer, we store buffer of the same size as the model's hidden dimension. The total number of parameters in an expert hidden _sie intermediate size; for our MoE this 1024 x 2816. We increase this by 1024, which is negligible 1/2816 0.03% increase in the number of MoE parameters.
Table 3: Throughput Comparison. We compare the throughput between DefaultMoF, and TopK for different model sizes, measured in tokens per second (sequence length 2048) on GPU,
Hidden Dim
Model
Tokens per second
Overhead vs TopK
Size
TopK
Ours
1024
1.968
26,393
25,913 1.391
-1.85%
2048
7,33B
1,393
-0.18%
Discussion
We have validated that DefaultMoEs beat tuned TopKMoEs across multiple settings, but why? Our explanation is as follows. Our default vectors are serving as nontrivial (Figure 3) approximations of their respective expert activations. Our gradient error should be corrected, which validated by De- faultMoE's 's router gradient being closer to the dense router gradient than that of TopKMoE (Figure 4). Less gradient error means that DefaultMoB can tolerate a higher learning rate than TopKMoE (Fig- ure 2), or simply converge faster at TopKMoE's best learning rate (Figure 11) and score higher on benchmarks (Table 2). This improvement is efficient (Table 3), and does not require additional hyperparameter tuning (Figure 8). We hope that our work will help the community train better MoEs.

---

References
Alex Andonian, Quentin Anthony, Stella Biderman, Sid Black, Preetham Gali, Leo Gao, Eric Hallahan, Josh Sevy-Krammer, Connor Leahy, Lucas Nestler, Kip Parker, Michael Pieler, Jason Phang. Shivanshu Purohit, Hailey Schoelkopf, Dashiell Stander, Tri Songz, Curt Tigges, Benjamin Thérien, Phil Wang, and Samuel Weinbach. GPT-NeoX: Large Scale Autoregressive Language Modeling in PyTorch, 2023. URL petts://wwrt www.gggthtn github. com/ www//eleuttteeoooooooooooooooeee www.eetatgtttteeeeeeeeeeeeeeeeee
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL bttps://arzee org/ abs/ 10/.1160700000 06450.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville, Estimating propagating gradients through stochastic neurons for conditional computation, 2013. URL https: / /aaxxiiivv..../ wwwrg/bbs/1300 3432.
Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack Rae, Erich Elsen, Koray Kavukeuoglu, and Karen Simonyan. Unified scaling laws for routed language models,
2022. URL aattps://arte /arxxv.xivv. org/ wabb./222..0111 01169.
Databricks. Dbrx, 2024.
URL
introducing- werreiiiinnndbxx---eee wwwp---ee-tttttrt----oooo-----
DeepSeek-AN Team. Deepseek-v2: strong, economical, and efficient mixture-of-experts language
model, 2024a. httpps////axxiv http://aaxxvv.../abbbaaaaa barrg/bbs/// haabs//405.0444
DeepSeek-AI Team. Deepseek-v3 technical report, 2024b. URL htt.httss//axxi.ooggaaatt 19437.
Nan Du, Yanping Huang, Andrew Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-cxpers In International Conference on Machine Learning, pages 5547-5569, PMLR, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. URL
org/ abs/2/0s//111// 03961.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-of-exppeeee 2022. URL https: // arxiv org/ abs 2211.12221111112 15841.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL httpt:ttt/teeeeeeeeeeee //tp:///eedooooo//e org/ http////oords//eeeee/ee sotooddss/220805//
Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024. URL https: //wpvv/rxvvvv//////////////... wwgg/os/////ooooooooooo..ota.stt /bbs/2403. /t/////43300500.00.
Jordan Hoffmaann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL https: http///axxvvvv.e/ee org/ abs/2203. 15556.
Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, Siyu Zhu, Steven Shimizu, Shivam Sahni, Haowen Ning, and Yanning Chen. Liger kernel: Efficient triton kernels for llm training. arXiv preprint arXiv:2410.1 10989, 2024. URL https: //t/x//xvvv///////////ee///aaav/ org/abs/ a/aabb/2411 10989,
Hunyuan Team. Hanguan-large: An open-source moe model with 52 billion activated parameters by tencent, 2024. URL httpt:///aartii org/at/ hts//bb//21111/ 02265.
10

---

Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive Mixtures of Local Experts. Neural Computation, 3(1):79- 87, 03 1991. ISSN 0899-7667. doi: 16./eo2/ 1991.3.1.79, URL hattps://dol org/ .000..11221ma 1162/neco. 1991.3.1.74 3...79.
M. I. Jordan and R. A. Jacobs. Hierarchical mixtures experts and the em algorithm. In Maria Marinaro and Pietro G. Morasso, editors, ICANN '94, pages 479-486, London, 1994. Springer London. ISBN 978- 1-411------
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 2020.
Liyuan Liu, Jianfeng Gao, and Weizhu Chen. Sparse backpropagation for moe training, 2023, URL https://arxix: org/ aat/abbs/22111 22300.2300..2 00811.
Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, Vishrav Chaudhary, Zeqi Lin, Chenruidong Zhang, Jilong Xue, Hany Awadalla, Jianfeng Gao, and Weizhu Chen, Grin: Gradient-informed moe, 2024. URL https: aarxiviivvv org/ ab2/2499. 12136.
Llama 3 Team. The Ilama herd of models, 2024. URL
21783.
Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https: 05101.
//arxiv.
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: the finest collection of educational content, 2024. URL https:/ co datasetss
HugginngFaceF///Eeeeeeeddod
Mistral Team. Mixtral of experts, 2024. URL htttt://///xii http://://aaa.i..oorrb////////// 24088.00088..... 04088.
Toan Q. Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention. 2019, doi: 10.528 10.5281/ZENOD0.555545 URL http::://eenooooorgg/eeeeee 3525484.
OpenAl Team. Team, Gpt-4 technical report, 2024. URL
08774.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context, 2016. URL attps://arxxs org/ abs/ 1606.06031.
Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. https://a/
org/ /b///////// /aa//200... 17557.
Phi Team. Phi-3 technical report: A highly capable language model locally your phone, 2024. wttp://...rr/aaa//44..... URL hwtphhttp//tt..... http///axxvvv..e www.iiowg//...............ooorae abs///44/.2222////////////////// 14219.
Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, and Junyang Lin. Demons in the detail: On implementing load balancing loss for training specialized mixture-of-exper models. arXiv preprint arXiv:2501.11877 2025.
Noam Shazeer. Glu variants improve transformer, 2020. URL https://arxiv.org/ 20tpvvvv.rr//bs//002.... 05202.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gan mixture-of-expprtt layer, 2017. URL
.rgg 1bb/ 06538.
Snowflake. Arctic, 2024.
URI. https://ww/ http://wwww...fffl///// snowflake www.....e.e..........a.. com/en/blog/ -tt.roooppeeefiiii-----eeee-eeee fowndation waw.:attiioo-o-aaaa-------dooooo mode-s- aag.uugge-mdllll--oollleeeeee snowflake/.
11

---

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023. URL https: http://axxxv...gaabb.// org/abs/ 0904.09864.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix. Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL https: waxxiv.orgee /tt.2bbb22222237...
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers 1,000 layers, 2022. URL hattps://aate //t:vv//a.ivvvvvaa /rgg/bb///03.....// 00555.
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliary-loss-fre load balancing strategy for mixture-of-experts, 2024.
http:///wwxivv.ooabbb////saa 15664.
Ziteng Wang. Jun Zhu, and Jianfei Chen. Remoe: Fully differentiable mixture-of-experts with relu routing, 2025. URL httph:ht/////tttt httpps///aaxvvvvve//ee httpiv..rr//bs//41222 14711.
xAI. Grok-1, 2024. URL https: ttttps://github.... wtwpth/u..oo/.iaoo//roro.oe grok- wwwb--readdeeoo---eeeeeeee
12

---

A
Appendix
A.1 Limitations
Throughout our work we only train dropless MoEs. These are ideal for the academic research setting, because they train efficiently in data-parallel and do require cumbersome implementations of model or pipeline parallelism. However, dropless MoEs may not be well-suited for industry, because reshaping the tensors to provide consistent shapes the MoE kernels incurs many host-device synes. This is limitation that we share with prior work, but an important future direction will be extending our results to token-dropping MoEs.
We report results on a range of standard language modeling benchmarks, but we have not invested enough compute to performance on LLM benchmarks such as MMLU-Pro, and our data is not sufficiently domain-spe -specifi- to report scores on coding benchmarks (vulanEval, MBPP) or math (GSM8K). In the same vein, our models have not been trained to reason. We focus heavily pretraining in this paper, but an important future direction would be seeing whether we can finetune MoEs with method.
A.2 Experimental Setup Details
Architecture. We use SwiGLU [Shazeer, 2020] MLPs following Llama [TTuuroon al., 2023], 16 attention heads with dimension of 64, LayerNorm [Ba et al., 2016] and RoPE [Su et al., 2023].
Hyperparameters. We use the initialization from Wang et al. [2022] for residual branch merge layers and the initialization from Nguyen and Salazar [2019] for all other layers. We use the Adam W optimizer [Loshchilov and Hutter, 2019]. We sequence length of 2048 and global batch size of 1024, resulting in a global token batch size of 2 We use standard cosine decay schedule.
MOE-Specific Training Details. We set the auxiliary loss [[eduu et al., 2022] to 0.01. We do not use z-loss or jitter, because we find that at this scale they do not improve training for the baseline model. Following DeepSeek-Al Team [2024b], we set the first layer to be dense. Following Liu et al. [2024], we compute the aux loss across nodes. We train dropless MoEs.
Implementation. We train with the gpt-neox library [Andonian et al., 2023] integrated with Megablocks [Gale et al., 2022] and augmented with Triton kernels from [Hsu et al., 2024].
A.3 More details on other Routing Mechanism Comparisons
We see in Figure that our method significantly outperforms SparseMixer for at least the first 10B tokens.
32
30
A.4 Ablating EMA design choices
28
Tuning S. The lone hyperparameter introduced by our method is the parameter of the EMA. While we use B = 0.9 for the 8c1 and 8c2 MoEs in Section 4.5, , requires more careful tuning for the 32 expert MoEs. For N 32, we use 3 0.65, 0.95, and 3 = 0.999 for
26
24
Default MoE 8c2
SparseMixer 8c2
22
2B
4B
8B
10B
1,
Tokens
and K = respectively. In other words, sparser MoEs require lower 3. We believe this due to each expert receiving less tokens at each step. which leads to sparser "history" for each default vector. As result, the default vector's s estimate for the average expert output improves when assigning higher weight to the current batch.
Figure 5: Comparison of Default MoE and Mixrse We compare Default MoE with SparseMixer, both configured with experts and Top-K=2 active experts. The results report train- ing perplexity throughout training, demonstrat- ing that Default MoE consistently outperforms SparseMixer.
We vary the 3 used the Default MoE in Figure for an 8 expert MoE and find that both 3 = 0.9 and / = 0.999 perform equally well. Intuitively, /3 parametrizes how fast our EMA adjusts its model of the sample mean as we train the model. In principle the optimal value of 3 might depend on the
13

---

learning rate, which defines how fast the model is changing. the batch size, which controls how many tokens are actually incorporated into the EMA, and even the data distribution. For example, we need to tune 3 more carefully for 32 expert MoE. In Figure 7 we sweep across multiple values of /3 for 32c1, 32c2, and 32c4 Default MoF.
32
Beta 0.999
30
Beta 0.9
Beta 0.99
28
26
24
2B
4B
6B Tokens
8B
10B
Figure 6: Selecting the Best EMA Beta Parameter. We compare different values of the exponential moving average (EMA) smoothing factor 3 for an 8c1 Default MoE, evaluating 3 = 0.999, 3 = 0.99, and A 0.9. The results show that / 0.9 provides the most stable and effective training behavior, leading to lower perplexity. Based on these findings, we use 3 = 0.9 to train our 8c1 model for 320B tokens.
Weighted Updates Remove Sensitivity EMA Decay Rate (3). We explore different values the EMA decay coefficient 3 used to update the default vector. Naively, performance is sensitive to B. However, we find that we accumulate activations in the the default vector weighted by their router score performance becomes robust to the choice of B. In fact, multiple 3 values converge to the similar final performance as shown Figure 8.
Initializing the EMA. We experiment with initializing default vector EMA in two ways: zero initialization and random (Gaussian) initialization. Zero initialization leads to the default vector starting off with zero signal in the earliest training steps. Random initialization, on the other hand, immediately provides some signal to fill in missing expert outputs but this signal starts as pure noise. Given these tradeoffs, we try both approaches and outline our results in Figure 9, Initializing our EMA with zeroes demonstrates slight improvement compared to random initialization. We believe this is due to the adverse effects from providing a random signal to the router early in training. However, the impact is minor, even without the use of any explicit bias correction term in our EMA.
Passing the Default Vector Forward. Our goal is to provide a default expert activation for unactivated experts so that the router can receive a gradient for all experts. We do this by passing the default vector forward through the network, and the router gradient is automatically computed. However, we can in principle do this without passing the default vector forward, and just manually writing the gradient update for the backward pass. One reason why we pass the default vector forward because our error in estimating the true dense gradient, as written in Eq. 9, is scaled by the loss. Therefore if the default vector activations actually improve the model's output, our error is smaller. We validate this Figure 10, where we find that passing the default vector forward does improve over only updating the router gradients in the backward pass with the default vector.
14

---

32
30
28
28
$=0.5
8=0.95
8=0.5
8=0.95
8=0.65
8=0.99
26
8=0.65
8=0.99
26
8=0.8
0=0.995
8=0.8
8=0.995
8=0.9
0.999
24
8=0.9
B=0.999
24
2B
4B
6B
8B
10B
2B
4B
6B
8B
10B
Tokens
Tokens
(a)
(b)
30
28
8=0.5
8=0.95
26
8=0.65
8=0.99
8=0.8
8=0.995
24
0.9.9
B 0.999
2B
4B
6B
8B
10B
Tokens
(c) 4
Figure 7: Sweeps for the | hyperrparaeterr for Default MoE with 32 experts. For K the optimal 3 is 0.65; for K = 2 it is 0.95, for = 4 is 0.999. Notably, as K decreases, 3 must also decrease. With roughly 10 billion tokens of training, can clearly distinguish the best for any given configuration.
B Detailed Empirical Analysis
B.1 Domain Specialization
We plot the domain specialization for 8c1 TopKMoE (Figure 14,Figure 15,Figure 16), 8c1 Default- MoE (Figure 20, 21, Figure 22), 32c4 TopKMoE (Figure 17.Figure 18.Fiurre 19) and 32c4 DefaultMoE (Figure 23,Figure 24,Figure 25). All models are pretrained on FineWeb-Edu for 160B tokens.
B.2
Load Balancing
We plot the load balancing for 8c1 TopE (Figure 26). 32c4 TopKMoE (Figure 27), 8c1 Default- MoE (Figure 28) and 32c4 DefaultMoE (Figure 29). All models are pretrained FineWeb-Edu for 160B tokens.
B.3 Expert Coactivation
We plot the expert coactivation for 32c4 DefaultMoE at Layer 0 (Figure 30) and Layer 16 (Figure 31).
All models are pretrained FineWeb-Eod for 160B tokens.
15

---

Beta- 0.9
Beta 0.999
Scored Beta 0.999
Scored Beta 0.9
22
Perplexity
18
16
12
0
20
40
60
80
Takens (B)
Figure 8: We compare exponential moving average (EMA) updates to the default vector with and without scoring. In the scored version, each expert's contribution is weighted by its router probability before being added to the EMA update. Without scoring, the choice of 3 (e.g.. 0.9 vs. 0.999) significantly affects performance. With scoring, this sensitivity disappears both scored variants converge to the same strong final performance. The unscored 3 = 0.999 curve (orange) is nearly hidden behind the scored curves.
B.4 Dense Gradient Similarity
We plot the similarity the dense gradient for 8c1 TopKMoE and DefaultMoE at K=1 (Figure 32), K=2 (Figure 33), K=3 (Figure 34), K=4 (Figure 35), K=5 (Figure 36), K=6 (Figure 37) and K=7 (Fig- ure 38). All models are pretrained on FineWeb-Edd for 160B tokens.
B.5 Default Vector Similarities
We plot the cosine similarity of default vectors in 8c1 DefaultMoF. (Figure 39, Figure 40, Figure 41) and 32c4 DefaultMoE (Figure 42, Figure 43, Figure 44), All models are pretrained FineWeb-Edu for 160B tokens.
B.6 Note on References
To avoid extremely long bibliography, abbreviated the bibliography entries for very large teams of authors (eeepSeek-Al Team, 2024b, Mistral Team, 2024, Phi Team, 2024, DeepSeek-At Team, 2024a, Gemini Team, 2024, OpenAl Team, 2024, Llama Team, 2024, Hunyuan Team, 2024]. Full author lists are available at the links provided in the bibliography.
16

---

32
Zero init
30
Random init
28
26
24
2B
4B
6B
8B
10B
Tokens
Figure 9: EMA Buffer Initialization Strategies for Default MoE. We compare two approaches for initializing the exponential moving average (EMA) buffer for Default MoE: zero initialization and random initialization from a Gaussian distribution. The results indicate that zero initialization is more effective, leading lower perplexity. Based on this finding, adopt initialization for experiments.
32
30
28
26
Default MoE Bwd
24
Default MoE Fwd + Bwd
2B
4B
6B
8B
10B
Tokens
The importance of applying the EMA during the forward pass. We perform ablation study to examine the effect of applying the exponential moving average (EMA) during both the forward and backward passes versus only the backward pass. The results show that using the EMA both passes leads to consistently lower training perplexity, whereas applying the EMA only in the backward pass-which which affects only the gradient update-results in inferior performance. This highlights the importance of incorporating EMA throughout training to maximize its stabilizing effect.
Figure 10:
17

---

20
TopK (321.1B)
TopK (107.6B) (B))))
18
Default MoE (292.6g (B)
30.0
Default MoE (93.0B)
Target Perplexity 18)
27.5
Target Perplexity (18)
25.0
14
22.5
12
20.0
0B
50B
100B
150B 200B 250B 300B 350B
Tokens
08
20B
40B
60B
RDB
100B
120B
Tokens
Figure 11: Default MoE Beats TopK. Our De- fault MoE reaches perplexity of x 12 about 9% faster than the baseline Topk MoE, with- out introducing any additional overhead. Both MoEs are configured with experts and Top- K=1 active experts.
Figure 12: Default MoE beats a tuned TopK baseline on FineWeb, Our method reaches the target perplexity of 18 about 15% faster than the baseline TopK method, without introducing any additional overhead.
32
Default MoE
32c2
30
TopK
8c1
32c4
8c2
32c1
28
26
24
22
4B
6B
8B
10B
Tokens
Figure 13: Comparison of MoE Configurations. We compare Default MoE and Top-K across five configurations: 8c1, 8c2, 32c1, 32c2, and 32c4. The results show that Default MoE outperforms Top-K in all MoE configurations. Both 8c1 and 8c2 use B 0.9; however, 32c1, 32c2, and 32c4 require more careful tuning of B. We use 3 0.65, 3 = 0.95, and f = 0.999 for 32c1, 32c2, and 32c4, respectively. We detail choices of in Appendix A.4.
Lascc
PubMic Cuntral
Soucibuchaag
Figure 14: Domain Specialization for 8c1 TopKMoE at Layer 0.
18

---

Layec
ILMMIC Cenrul
Figure 15: Domain Specialization for 8c1 TopKMoE at Layer 8.
Pilo
Prloec Contral
Figure 16: Domain Specialization for 8c1 TopKMoF Layer 16.
Laper
8.35
Edbba
Expar
Figure 17: Domain Specialization for 32c4 TopKMoE at Layer 0.
Laywr
Figure 18: Domain Specialization for 32c4 TopKMoE at Layer 8.
AcEa
Pi
Pullied Contral
Sraddiihegge
26-20
20.000
Figure 19: Domain Specialization for 32c4 TopKMoE at Layer 16.
Gablak
Iepertt
Figure 20: Domain Specialization for 8c1 DefaultMoB. at Layer
19

---

Pla
Prlseed Cental
Glbba
Papprt
Figure 21: Domain Specialization for 8c1 DefaultMoF, Layer 8.
Tayyr Ini
PuMMec Central
Ioert.D
Figure 22: Domain Specialization for 8c1 DefaultMoE at Layer 16.
ArEx
Pall Contral
Wraidholwage
18-11-111 20 25-2 213035
Figure 23: Domain Specialization for 32c4 DefaultMoE at Layer 0.
Laper
8.25
Cabak
Expart
Figure 24: Domain Specialization for 32c4 DefaultMoE at Layer
Layer 16
8.23
18:33 20
213033
Figure 25: Domain Specialization for 32c4 DefaultMoE at Layer 16.
20

---

1.25
0
1.20
1.15
5
1.10
Layer
10
1.05
1.00
15
0.95
0.90
20
0.85
0.80
0
5
Expert ID
Figure 26: Routing Frequency for 8c1 TopKMoF.
1.25
1.20
1.15
Layer
10
1.05
1.00
20
0.90
20
30
Expert
0.85
0.80
Figure 27: Routing Frequency for 32c4 MopK
21

---

1.25
0
1.20
1.15
5
1.10
Layer
10
1.05
1.00
15
0.95
0.90
20
0.85
0.80
0
5
Expert ID
Figure 28: Routing Frequency for 8c1 DefaultMoF.
1.25
1.20
1.15
Layer
10
1.05
1.00
20
0.90
20
30
Expert
0.85
0.80
Figure 29; Routing Frequency for 32c4 DefaultMoE.
22

---

1.0
Layer 0
0.8
10
0.6
Expert
15
0.4
25
0.2
30
10
20
30
Expert
0.0
Figure 30: Expert Coactivation for 32c4 DefaultMoB at Layer 0.
1.0
Layer 16
0.8
10
0.6
Expert
0.4
25
0.2
30
10
20
30
Expert ID
0.0
Figure 31: Expert Coactivation for 32c4 DefaultMoB, at Layer 16.
23

---

Similarity to Dense Gradient at K=1
14
10
13
9
16
17
18
19
5
20
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 32: Similarity Dense Gradient for 8c1 TopKMMo and DefaultMoR. K=1.
Similarity to Dense Cradient at K=2
10
13
9
16
17
18
19
5
20
21
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 33: Similarity to Dense Gradient for 8c1 TopKMoo and DefaultMoB K-2.
24

---

Similarity to Dense Gradient at K=3
10
13
9
16
17
18
19
5
20
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 34: Similarity Dense Gradient for 8c1 TopKMMo and DefaultMoF. K=3.
Similarity to Dense Cradient at K=4
10
13
9
16
17
18
19
5
20
21
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 35: Similarity to Dense Gradient for 8c1 TopKMoo and DefaultMoB K=4.
25

---

Similarity to Dense Gradient at K=5
10
13
9
16
17
18
19
5
20
21
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 36: Similarity Dense Gradient for 8c1 TopKMMo and DefaultMoF. K=5.
Similarity to Dense Cradient at K=6
10
13
9
16
17
18
19
5
20
21
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 37: Similarity to Dense Gradient for 8c1 TopKMoo and DefaultMoB K=6.
26

---

Similarity to Dense Gradient at K=7
10
13
9
16
17
18
19
5
20
3
22
2
23
1
0.0
0.5
1.0
TopK MoE
Default MoE (Ours)
Figure 38: Similarity Dense Gradient for 8c1 TopKMoE and DefaultMoR. K=7.
1.00
Layer 0
0.25
Expert
0.00
00.25
-0.50
-0.75
0
2
4
6
Expert ID
-1.00
Figure 39: Similarities Between Default Vectors for 8c1 DefaultMoB. Layer 0.
27

---

1.00
Layer 8
0.50
0,25
Expert
0.00
-0.55
-0.50
0
2
4
6
Expert ID
-1.00
Figure 40: Similarities Between Default Vectors for 8c1 DefaultMoF. Layer 8.
1.00
Layer 16
0.25
Expert
0.00
-0.25
-0.50
00.75
0
2
4
Expert ID
-1.00
Figure 41: Similarities Between Default Vectors for 8c1 DefaultMoE at Layer 16.
Layer
1.00
0.7
0.55
Expeet Ib
0.00
20
0.50
25
-1.00
10
15
20
25
Expert ID
Figure 42: Similarities Between Default Vectors for 32c4 DefaultMoE. Layer 0.
28

---

Layer
0.75
0.25
Expeet 1D
0.00
20
0.50
25
-0.75
-1.00
10
15
20
25
Expert ID
Figure 43: Similarities Between Default Vectors for 32c4 DefaultMoF, Layer 8.
Layer
0.7
0.55
Expert ID
0.00
20
0.50
23
15
20
Expert
Figure 44: Similarities Between Default Vectors for 32c4 DefaultMoB Layer 16.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the motivation, summarize method- ology, and state the key contributions made in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: The first section of the Appendix states limitations. Other limitations, regarding the assumptions made are also discussed throughout the paper.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
30

---

Answer:
[NA]
Justification:
We do not provide theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: In the main paper, each experiment details the setup and datasets used. In addition, we have open sourced our code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path reproducing verifying the results. 5. Open access to data and code
31

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide link the code and also explain the experimental setup in the main body. The data used in our experiments is publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurlPS code and data submission guidelines (https: / wwccc/pulccceecccccccccccecclell aellic/guiies gww.guudde//eee wwwpddeSumisssinnnnee.
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe experimental details in the main body of the paper and also in our open sourced code.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: w[Yes]
Justification: The benchmark results report the standard error given by Im-eval-harness. The experiments focus large-scale language model pretraining, where each run is computa- tionally expensive and typically conducted once per configuration. As is common this setting, we report complete training and validation results over training rather than multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
32

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The LLM experiments were on 64 GPUs the AWS cluster. Our code
contains configuration files which can be used to estimate compute resources.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't itdn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttpas:////ups../eeee wwtpeerww......u.pppeeeeeepeeeee www.blic//thissuiieeeeooo
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is methodological in nature, proposing improvements to training efficiency and stability in Mixture-of of-Exeerts models. As such, it does not directly engage with societal use cases their potential impacts.
Guidelines:
33

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
This work does not involve the release pretrained models, datasets, or tools with direct misuse risk. It focuses on architectural and training improvements for LLMS trained on publicly available data.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We do cite creators and original owners of assets in paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
34

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com/ comddataees has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets are introduced as part of this work.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to NeurlPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research involving human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
35

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The research presented in this manuscript does not involve the use of LLMs as important, original, or non-standard component.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: //nerrurps.... caaapps.cc/oooofeeeeeeee ccc//onferecces/225//1111 for what should or should not be described.
36


Reviews:
[{"review_id": "qLAzgbIGmO", "reviewer": "Reviewer_zKLq", "summary": "One 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.", "questions": "- 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?", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors have successfully addressed my concerns and questions. I don't think the method or the results achieved can be qualified as a \"groundbreaking impact\", but it is a very solid paper. Thus, I definitely recommend the AC accepting the paper.", "rebuttal": "Thank you for your detailed review and valuable feedback. We are glad you found our paper easy to read, and appreciate you recognizing the simplicity and effectiveness of our method. We address each of your questions and concerns below.\n\n> It seems that the proposed method achieves better results, compared to standard TopK, when the sparsity is not too high. How do the authors explain this?\n\n**Our method outperforms standard TopK for the sparsest and largest model (the 1/64 sparsity 7B MoE).**\n\nYou are correct in suggesting that higher sparsity presents more potential for improvement with Default MoE. In section 4.1 (line 151) we point out that in our 2B parameter MoE, roughly 1.6 billion of these parameters are allocated to the expert MLPs. In an MoE with $N=32$ experts and $K=1$, only 50 million of these MLP parameters are active (in contrast to 200M active expert parameters for $N=8$). We believe **it is difficult for our method to offer much of an improvement at such a limited scale**. \n\nTo further demonstrate this point, we note that our 7 billion parameter Default MoE with a sparsity of $1/64$ significantly improves over TopK, with a much wider performance gap than the 2 billion parameter 32c1 MoE (Figure 2d). In other words, **unless the model size is prohbitively small, our method does outperform TopK at higher sparsity levels**.\n\n> I couldn't find any equation or detailed explanation of the actual \"weighting\" being used.\n\nYou’re right, and we will update Section 3.2 to explain this weighting method more clearly. The following equation expresses how we weight the default vector update using the router logits. In doing so, we reduce the method's sensitivity to the $\\beta$ hyperparameter.\n\n$$\n\\hat E_i^{(t)} = \\beta \\hat E_i^{(t-1)} + (1-\\beta)\\frac{1}{K}\\sum_{j\\in R_i^{(t)}} S_{ij} E_i(x_j)\n$$\nwhere $R^{(t)}\\in\\mathbb{N}^{B\\times K}$ denotes the indices of the $K$ routed experts for each of the $B$ batch tokens at training step $t$, and $S\\in\\mathbb{R}^{B\\times N}$ denotes the router weights corresponding to all $N$ experts for each token. \n\n> Why apply the EMA update before expert forward pass if the natural dependence of tokens may be broken? Why not update after?\n\nYou are right in bringing up this potential issue. We update the EMA before the forward pass in the submitted paper, but we also experimented with applying it after the forward pass after the submission and found no difference in performance across the model scales we evaluate. We have updated the method to apply the EMA after the forward pass, and will edit lines 130-132 to reflect this in the camera ready. Thank you for this great suggestion!\n\n> 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\n\nThese are interesting points, and we would like to clarify a few details:\n* When the router entropy conditioned on an input is high, it is true that slight changes in the input can lead to a different routing decision. But if the inputs are similar, this does not necessarily mean the expert outputs will also be similar, as the expert MLP is nonlinear.\n* In the high-entropy setting, the router may make choose an expert based on a marginal difference in the router weights. This may not be the optimal expert choice, yet the TopK router will not receive feedback from other experts. We claim that the default vectors are *more useful* when entropy is high because they provide information from all experts when the router's uncertainty is high. \n* We agree with your intuition on what would happen when router entropy is zero.\n\n> 8c1 results on Lambada show % of almost 4000, is this correct?\n\nWe measure the normalized improvement using the expected random score as a baseline. For example our MMLU score is 32.5 compared to 31.8 for TopK MoE. This seems like a 2% improvement. But, MMLU is multiple choice so just a random guess would get 25%. So we are $\\frac{32.5-25}{25}=30\\%$ better than the random baseline on this metric, compared to the 27% normalized improvement of TopK. The normalized LAMBADA scores are unexpectedly high because the random baseline is 1%, much lower than that of the other benchmarks. Compared to 1%, our score of 41.0 is a 40x improvement. We realize this seems out-of-place so we have a footnote in the paper addressing it; if you can point us to a better random baseline for LAMBADA we are happy to use it. \n\n> Figure 2d is a bit confusing without reading the caption.\n\nWe agree, and thank you for pointing this out. All but the rightmost data point in Fig. 2d correspond to models with 2 billion active parameters. We will separate the rightmost data point (from a 7B model) into its own result in the camera ready revision.\n\n> Did you ever consider using the same optimizer as the rest of the parameters by adding an auxiliary term to the NLL loss? \n\nYes, specifically we experimented with training the default vector $\\hat E_i$ to be a learnable parameter. In an ablation experiment we found that this underperforms our proposed method. The core idea of our method is to approximate the *expected expert output*, but directly optimizing the default vector does not account for this heuristic. We'll add more detail on what we tried here in the camera ready.\n\n> Typo in line 240? \"dMoEs\" -> \"MoEs\"?\n\nYes, thanks for catching this!\n\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\nYes, $N=8$ is the total number of experts. Note that we are specifically comparing to a dense model as an upper bound.  We vary $K$ between $1$ and $8$ to evaluate all possible sparsity levels. In a dense model all of the experts are active so $K=N=8$ and there is no sparsity. We will edit the figure caption to clarify this further.\n\n*Thank you for your detailed review; we genuinely enjoyed thinking about the questions you asked and look forward to reading your response.*"}, {"review_id": "Cv8hP5RIav", "reviewer": "Reviewer_V11s", "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.", "questions": "Suggestion: Together with Figure 11, it would be nice to plot the average zero-shot evaluation scores throughout training.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n- Overall this paper is easy to follow, the method is well-motivated and justified. The empirical study is also comprehensive and convincing.  \n\nWeaknesses\n- Cross-validating $\\beta$ could be costly.\n- 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]\n\n[A] Chi, Zewen, et al. \"On the representation collapse of sparse mixture of experts.\" Advances in Neural Information Processing Systems 35 (2022): 34600-34613.\n\n[B] Csordás, Róbert, et al. \"Moeut: Mixture-of-experts universal transformers.\" NeurIPS 2024\n\n[C] Lv, Ang, et al. \"Autonomy-of-Experts Models.\" ICML 2025", "was_revised": true, "final_justification": "Overall a solid work, and the author addressed my concerns in the rebuttal. Therefore, I recommend acceptance.", "rebuttal": "Thank you for providing these comments and suggestions. We are glad that you found our paper convincing and easy to follow.\n\n> Tuning $\\beta\\$ could be costly\n\nAs we note in Line 184, we don't actually need to tune $\\beta$ because we find that when the EMA update includes the router weights. The following equation expresses how we weight the default vector update using the router logits. In doing so, we reduce the method's sensitivity to the $\\beta$ hyperparameter, as shown in Figure 8 in the paper. We recognize that this is an important detail regarding the practicality of our method, and we will add an additional explanation in Section 3 to clarify this.\n\n$$\n\\hat E_i^{(t)} = \\beta \\hat E_i^{(t-1)} + (1-\\beta)\\frac{1}{K}\\sum_{j\\in R_i^{(t)}} S_{ij} E_i(x_j)\n$$\nwhere $R^{(t)}\\in\\mathbb{N}^{B\\times K}$ denotes the indices of the $K$ routed experts for each of the $B$ batch tokens at training step $t$, and $S\\in\\mathbb{R}^{B\\times N}$ denotes the router weights corresponding to all $N$ experts for each token. \n\n> Compare to XMoE, MoEUT, and AoE\n\nThank you for the useful related works; we have started the process of comparing against these methods, and will include the results in the camera ready, along with updating our open-source codebase to include the code of our reimplementations.\n\n> Plot eval scores throughout training\n\nWe agree with your suggestion to plot the evaluation scores throughout training. We have all the training checkpoints available, and will add the eval benchmarks plotted over time in the camera ready revision. (Please note that we can’t add figures in the rebuttal)"}, {"review_id": "je6JzAIiiJ", "reviewer": "Reviewer_hP5u", "summary": "Default-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.", "questions": "See Weaknesses.", "limitations": "Efficiency concerns are listed in Sec 4.6. No societal impact was addressed.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "While the paper shows promise and the rebuttal does address _some_ of the concerns, they do not provide any substantial improvement to the original version of the paper. The arguments are very handwavy, with only a few numbers included. For that I believe my evaluation of the paper remains unchanged.", "rebuttal": "Thank you for your review and constructive feedback. We appreciate you pointing out the scale of our experiments and depth of our ablation studies. \n\n> Default MoE should be compared to Dense MoE\n\nWe agree with your suggestion to compare Default MoE against the Dense MoE as a baseline. After 10 billion tokens of training, a TopK 8c1 MoE achieves perplexity of $23.786$. A Dense 8c8 MoE achieves perplexity of $21.499$. Our Default MoE (8c1) achieves perplexity $23.192$. So, after 10 billion tokens, we recover $1 - \\frac{23.192 - 21.499}{23.786 - 21.499}\\approx 26\\%$ of the performance gap between standard TopK and Dense MoE. We again emphasize that the computational cost imposed by this improvement is negligible. We will include this result and a loss curve comparison in the camera ready revision. We note that while 26% of the performance gap may not seem like a lot, approximating Dense MoE's gradient 100% faithfully will not reproduce the same performance as the Dense 8c8 MoE unless all 8 experts are activated, and this would incur significant computational overhead.\n\n> Experiment diversity\n\nWe appreciate your suggestion to include more extensive results with SparseMixer and Loss-Free Balancing. The SparseMixer authors note that their method requires at least 500 billion tokens of training to surpass the TopK baseline (arXiv:2310.00811). While this scale is outside the scope of our compute budget, we can conclude that prior to 500 billion tokens Default MoE will outperform SparseMixer since it already outperforms TopK. We will provide extended training results with Loss-Free Balancing to compare its performance over many more tokens. These results will be included in the camera ready version of our paper. Additionally, we will include a comparison to the related works suggested by Reviewer 2 to further expand the breadth of our evaluation.\n\n> Router fluctuation\n\nTo address your idea on router fluctuation, we empirically measured the similar metric of router saturation method proposed in OLMoE (arXiv:2409.02060). Router saturation involves the router decisions for identical token sequences at various training checkpoints, and measures the rate at which these decisions match those of the final trained model. For example, a router saturation of $0.8$ at a checkpoint signals that the router's decisions will remain the same for roughly $80\\%$ of the tokens throughout training, and will change/fluctuate for the other $20\\%$. When observing router saturation over time, we observed no significant difference between TopK and Default MoE after the earliest steps of training. In other words, **the router in TopK and Default MoE \"fluctuate\" at similar rates.**\n\n*Thank you for your detailed review. We look forward to reading your response.*"}, {"review_id": "rI22lCuBfp", "reviewer": "Reviewer_HK1B", "summary": "This 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.", "questions": "* 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?", "limitations": "-", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This 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.", "was_revised": true, "final_justification": "This is a simple idea with demonstrated effectiveness.  Most of the questions in my initial review were answered well.", "rebuttal": "Thank you for your detailed review and insightful questions. We address each point below:\n\n> How compact or diffuse are the clusters? What are the distances of each expert's outputs to each avg?\n\nSince the router's conditional entropy decreases over time, we know that expert clusters become more compact (since each token's distance to other experts decreases). We will measure the exact distances over the course of training and include a plot of this trend in the camera ready revision.\n\n> But what are the impacts of non-zero-centered error, other than just worse performance?\n\nWhen the gradient error has nonzero expected error, the router only receives feedback from the expert(s) it selected, regardless of whether it was the optimal decision or other expert(s) were better choices. As an example, assume the router selects expert $2$ and receives a positive gradient signal. This will lead the router to increase expert $2$'s weight as to reinforce that selection. Assume also that had the router selected expert $1$, it would receive a positive signal gradient signal that is twice as strong. In fact, expert $1$ would be the optimal expert for the token and if the router received this additional information it would adjust its decision accordingly. \n\nIn this example, the absence of gradient signal from expert $1$ is a component of the error. The TopK router would promote only expert $2$ because it received a positive gradient signal. But this error prevents the router from seeing that expert $1$ is in fact a better option. When we adjust the gradient to be unbiased, we remove this pitfall of the standard TopK router.\n\n> How does this interact, if at all, with other expert selection tricks like imposed balancing constraints or stochastic selection/gradient estimators?\n\nOur implementation is designed to be flexible and add on to any existing MoE routing method. For example, we experimented with adding Default MoE interacting with SparseMixer. We were able to run this experiment successfully, although we did not evaluate this or any other combination of methods on a large-scale training run.\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. I wonder what you think about this seeming inconsistency?\n\nThis is an interesting point and we agree with your intuition regarding the closeness of inputs to an expert. The gradient signal supplied by the default vector has the strongest effect for inputs that have a high probability for that given expert. In other words, the effect of the default vector is proportional to how \"close\" that input is from the expert's domain. While there is a slight inconsistency in also applying this default vector to \"farther\" inputs this effect is much less pronounced. \n\nAs for flattening the router distribution, we don't believe this is a strong cause for concern because we observe empirically that conditional entropy gradually decreases (i.e. this distribution sharpens) over the course of training.\n\n> If the new expert is enabled, how close are the actual outputs to the mean? Could be oscillation at the edges? \n\nRelated to your first comment, the closeness of outputs to the mean depends on the time step. Expert clusters become compact over time as experts specialize, so the hypothetical output if a token were routed to an alternate expert grows father from the true mean expert output. But this is not a issue in practice because the default vector will have less of an effect on the input's routing decision as it grows farther. In the camera ready, we will include a figure depicting the measured distances between the expert output of an input on the boundary and the true mean expert output.\n\nWe do not observe significant oscillation effects in routing; see our response to Reviewer 3 regarding router fluctuation.\n\nWe also evaluated expert specialization (in the Appendix, Figs 14-25) between TopK and Default MoE and found no significant difference. So, **Default MoE does not harm expert specialization**.\n\n> As mentioned above, to me it seems this has the potential to flatten the router distribution. Have you seen anything like this?\n\nWe tracked conditional router entropy, which is positively associated with a flat router distribution. In all of our experiments we observed a monotonically decreasing router entropy, and did not see any signs of a flattened router distribution.\n\n> \"we can actually automatically account for the impact of sparsity and granularity on the default vector update by weighting the updates to the default vector by the router logits.\" I actually haven't found a description of this yet looking over the appendix, either -- how does this work?\n\nThe following equation expresses how we weight the default vector update using the router logits. In doing so, we reduce the method's sensitivity to the $\\beta$ hyperparameter, as shown in Figure 8 in the paper. We recognize that this is an important detail regarding the practicality of our method, and we will add an additional explanation in Section 3 to clarify this.\n\n$$\n\\hat E_i^{(t)} = \\beta \\hat E_i^{(t-1)} + (1-\\beta)\\frac{1}{K}\\sum_{j\\in R_i^{(t)}} S_{ij} E_i(x_j)\n$$\nwhere $R^{(t)}\\in\\mathbb{N}^{B\\times K}$ denotes the indices of the $K$ routed experts for each of the $B$ batch tokens at training step $t$, and $S\\in\\mathbb{R}^{B\\times N}$ denotes the router weights corresponding to all $N$ experts for each token. \n\n> Paragraph at l.240 about globally reduced load balancing loss: This paragraph didn't make any sense to me, but seems important.\n\nTraditionally, the load balancing loss is computed on each individual GPU's microbatch, using the microbatch's routing counts. The gradients from the load balancing loss are then all-reduced at the optimizer step. Alternatively, the globally reduced load balancing loss first all-reduces the routing counts, then applies the same \"global\" gradient to each microbatch. To see why this is important we present a motivating example:\n\nAssume we are training with 32 GPUs and 32 experts. Also assume on each GPU $i$, expert $i$ receives $131$ tokens and all other experts receive $99$ tokens. For example, the microbatch on GPU 0 has local routing counts $[131, 99, 99, \\dots, 99]$.\n\nThe global routing counts are $[3200, 3200, \\dots 3200]$ which means the distribution is perfectly uniform. So, the global load balancing loss should be zero. But, if we calculate the load balancing loss on each GPU individually the total loss will be nonzero because each individual microbatch is imbalanced. \n\nIn practice, this is problematic because individual microbatches's routing counts are noisy due to small sample size. At a large enough scale, we end up overestimating the penalty to apply for load balancing. Globally reduced load balancing loss avoids this by first reducing the routing counts, and then applying the loss. See Section 3.3 of GRIN MoE (arXiv:2409.12136) for more details. We will include a similar explanation in the camera ready revision.\n\n> Fig 2d: 1/64 sparsity --- 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\nWe agree, and thank you for bringing up this point of confusion. To clarify, this data point is for a larger (7 billion parameter) model compared to the other points corresponding to 2 billion parameter models. We will separate this individual result from the rest of the figure for a clearer comparison.\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\nWe checked the OpenReview submission; actually, in the first column, the PubMedQA score is 54.8_2.2, but in the third and fourth columns, it is 58.4_2.2. In this case, the DefaultMoE and TopK MoE got the same score for the 32c4 MoE configuration. (We went back and checked the LM-Eval-Harness outputs, and the stderrs are actually different, 2.2227 for the first column and 2.2064 for the third and fourth column) Please let us know if this is what you see when you look at the PDF on OpenReview. \n\n*Thank you very much for your thoughtful read of our paper. Your questions will make great additions to the paper, and we look forward to reading your response.*"}]

Author rebuttal:

Meta-review: This paper presents a method that improves sparse MoE training by substituting missing expert activations with exponential moving averages of previously seen expert outputs, enabling dense gradient updates to routers while maintaining sparse computation. The approach consistently outperforms standard TopK routing across multiple model scales and benchmarks with negligible computational overhead. 

While reviewers appreciated the method's simplicity, clear motivation, and comprehensive ablation studies, they noted limitations including incomplete comparisons with other routing methods and evaluation primarily at smaller scales (up to 7B parameters). On the other side, during the rebuttal, authors successfully addressed most technical concerns by clarifying implementation details, providing additional experimental results, and promising expanded baseline comparisons for the camera-ready version. Despite these limitations, the paper received unanimous support for acceptance due to its solid technical contribution and practical applicability.
Decision: Accept (poster)