Dense Backpropagation Improves Training for Sparse Mixture-of-Experts

Thank you for your detailed review and insightful questions. We address each point below:

How compact or diffuse are the clusters? What are the distances of each expert's outputs to each avg?

Since 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.

But what are the impacts of non-zero-centered error, other than just worse performance?

When 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.

In 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.

How does this interact, if at all, with other expert selection tricks like imposed balancing constraints or stochastic selection/gradient estimators?

Our 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.

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?

This 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.

As 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.

If the new expert is enabled, how close are the actual outputs to the mean? Could be oscillation at the edges?

Related 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.

We do not observe significant oscillation effects in routing; see our response to Reviewer 3 regarding router fluctuation.

We 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.

As mentioned above, to me it seems this has the potential to flatten the router distribution. Have you seen anything like this?

We 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.

"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?

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.

$$\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)$$

where $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.

Paragraph at l.240 about globally reduced load balancing loss: This paragraph didn't make any sense to me, but seems important.

Traditionally, 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:

Assume 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]$.

The 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.

In 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.

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

We 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.

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?

We 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.

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.

Thank you for providing these comments and suggestions. We are glad that you found our paper convincing and easy to follow.

Tuning \beta\ could be costly

As 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.

$$\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)$$

where $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.

Compare to XMoE, MoEUT, and AoE

Thank 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.

Plot eval scores throughout training

We 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)

Thank you for your review and constructive feedback. We appreciate you pointing out the scale of our experiments and depth of our ablation studies.

Default MoE should be compared to Dense MoE

We 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.

Experiment diversity

We 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.

Router fluctuation

To 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.

Thank you for your detailed review. We look forward to reading your response.

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.

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?

Our method outperforms standard TopK for the sparsest and largest model (the 1/64 sparsity 7B MoE).

You 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.

To 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.

I couldn't find any equation or detailed explanation of the actual "weighting" being used.

You’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.

$$\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)$$

where $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.

Why apply the EMA update before expert forward pass if the natural dependence of tokens may be broken? Why not update after?

You 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!

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

These are interesting points, and we would like to clarify a few details:

• 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.
• 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.
• We agree with your intuition on what would happen when router entropy is zero.

8c1 results on Lambada show % of almost 4000, is this correct?

We 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.

Figure 2d is a bit confusing without reading the caption.

We 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.

Did you ever consider using the same optimizer as the rest of the parameters by adding an auxiliary term to the NLL loss?

Yes, 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.

Typo in line 240? "dMoEs" -> "MoEs"?

Yes, thanks for catching this!

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?

Yes, $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.

Thank you for your detailed review; we genuinely enjoyed thinking about the questions you asked and look forward to reading your response.