Dense Backpropagation Improves Routing for Sparsely-Gated Mixture-of-Experts

Dear Reviewer XWKc,

Thank you for the constructive comments, and for appreciating the strengths of our evaluation. In this response, we want to outline how we have addressed your concerns in the revised paper.

W1: Motivation Clarification

Our method is motivated by previous works like SparseMixer that show approximating the dense gradient leads to improvements in training performance and stability. We develop a new estimator using the straight-through, and in Figure 1, Table 1 we demonstrate our improvements upon the baseline method as reported on a set of benchmarks. We include our improvement over the baseline through Section 4 (Evaluation).

W2: Unclear Methodology

We recognize that our methodology was not clearly explained with regard to using the approximation in the forward pass. To clarify, our method does not use the approximation in the forward pass. We experimented with this and found the performance was better when only using the approximation in the backward pass. To answer your third question, this does not affect the causality of the language model, as the approximation information is not passed between tokens in the forward pass.

W3: Modest Improvement

Thank you for your comments. You pointed out that while our lightweight method improves performance and load balancing, one weakness is the improvement being modest compared to the overhead increase. We have discussed how the improvement of our method over the baseline increases as we scale up training in Table 2. We also did a deep dive into the efficiency of our method in Section 4.4 “Efficiency”, with Figure 9 and Table 4. We believe that at scale, the improvement of our method is greater than the additional overhead, and also that the overhead can be further reduced by kernel specialists (which we are not).

W4: Dependence on large micro batch size

We only use this approximation to update gradients in the backward pass; moreover, we all-reduce the approximation gradients across data parallel workers. This gives us a large sample size of token gradients to estimate the dense gradient, regardless of the micro batch size on each worker. As a result, our method does not depend on large micro batches to produce a good approximation. We will add an ablation over the micro batch size to validate this by the end of the author response period.

Q1: How does Expert Group Approximation perform for K=1 or K>2?

To answer your question on values of K, we experimented with multiple of K and consistently observed an improvement in training. We report K=4 in Table 2 and K=1 in Table 3, and find that we still outperform the baseline with these other values of K.

Dear Reviewer XWKc,

We are following up with some initial results of the comparison between the STE-based dense-trained baseline and our method. We train an 8-expert MoE with 2B hidden parameters.

After 9B tokens, we have:

These results corroborate our existing results in Table 5, where we find that using our method yields similar improvements to activating another expert, without the added compute cost of actually activating another expert. The experiments to run these for 20B tokens are currently running and we expect we may be able to update the results tomorrow.

We will put these fleshed-out results into the camera ready. We believe that this STE-trained dense baseline serves as an upper bound for the performance of our method. This makes the intuition for our method clearer: we want to approximate what happens if we activate all the experts, because we know this will increase performance, and our goal is computational efficiency.

We also have initial results for an ablation of our method where we use Sparsemixer for the router, and only use our method to provide a dense gradient for the experts. This seems to improve over just using Sparsemixer, and while it lags behind using our method without Sparsemixer, this may be expected because Sparsemixer adds noise and therefore will not be better until about 0.5T tokens [1]. We train an MoE with 8 experts and 2B hidden parameters, activating K=2 experts, for 10B tokens.

We believe this shows that our method is actually compatible with Sparsemixer, further bolstering the case that our method can be a valuable addition to MoE pretraining.

[1] GRIN: Gradient Informed MoE. Liu et al. https://arxiv.org/abs/2409.12136

Dear Reviewer GiJM,

Thank you for your very useful comments, and for appreciating the importance of the problem we’ve chosen to study. In this response, we want to outline how we have addressed your concerns in the revised paper.

W1: Unclear how good our approximation is

Thank you for your valuable feedback. We want to clarify that our goal is not to approximate the individual expert activation, but the dense gradient during the backward pass. We mentioned using the approximation in the forward pass in the paper, but after later experiments found that this was not optimal. We have edited the paper to reflect this in lines 198-201. In Figure 6 we demonstrate that we are very accurately approximating this dense gradient using the gradients of other tokens. Thus, the motivation of our paper does hold in practice when we train an MoE with the dense gradient approximation.

W2: Poor presentation of experimental results

Thank you, we have rewritten the evaluation section. We now clearly state which method is being reported in each of our figures.

W3: No comparison to SparseMixer

Thank you for the references. We have added a comparison to SparseMixer in Figure 8. We reproduce the text here for convenience:

\citet{liu2023sparsebackpropagationmoetraining, liu2024gringradientinformedmoe} use Sparsemixer, which estimates the true router gradient without straight-through. Our method is distinct in that we also provide gradient updates for the expert weights(as seen in Figure 8). \citet{liu2024gringradientinformedmoe} note that Sparsemixer lags behind TopK~(which our method always outperforms) for the first 0.5T tokens, likely due to the noise that Sparsemixer adds. Unfortunately, we have not been able to invest 0.5T tokens into this baseline comparison, but we are currently running this and expect to have it done for the camera ready. In the 10B tokens that we were able to run Sparsemixer for, our method significantly outperforms it.

W4: Significance of the results

We have added Table 1 with precise numbers comparing the benchmark performance of our method and the baseline. Our baseline itself is a DeepSeek MoE with tuned coefficients for auxiliary loss and tuned LR, so we feel that we are beating a strong baseline. For the camera ready, we will add a row with the Sparsemixer results on these benchmarks as well.

W5: Runtime overhead

Thank you, we agree that we were not clear about the runtime overhead of our method. We have added Section 4.4 “Efficiency” doing a deep dive into this, with Figure 9 and Table 4. We have updated “does not increase runtime” to “does not significantly increase runtime”. Our attention method was proposed as another alternative for approximating the dense gradient, but after further experiments we found this to be unreasonably efficient. This section is now entirely removed from the main text and moved to the appendix.

W6: Expert group approximation

Since submitting the paper, we identified and fixed an error in our “accurate” and “viable” approximations. We have added new results for this in Table 3, and discuss it in Lines 450-462.

Q1: Does the main motivation hold in practice?

We believe our gradient approximation similarity figure (Figure 6) confirms that our motivation does hold in practice, because our approximated gradient is more similar to the true dense gradient than the baseline TopK gradient is to the true dense gradient.

Q2: More rigorous investigation of "accurate" and "viable" approximations?

We have now presented a more rigorous investigation in Table 3, and Lines 450-462.

Q3: Do we beat the SOTA on zero-shot benchmarks?

We outperform a fine-grained TopK MoE trained with auxiliary loss, which is the setup used by SOTA MoEs including DeepSeek, on a range of benchmarks in Table 1.

Q4: Please make the claims regarding the complexity more accurate.

We have carefully revised Section 4 per your guidance, and spun off the efficiency discussion into Section 4.4, backed by Figure 9 and Table 4.

We hope that we have addressed your concerns and that the revisions made to the paper answer your questions.

Dear Reviewer 3TBL, Thank you for your insightful comments, and for highlighting the novelty of our approximation method. In this response, we want to outline how we have addressed your concerns in the revised paper.

W1: More Evaluation on Tasks

We agree that evaluation beyond perplexity is important. We now present the results of our model on various benchmarks including MathQA, LogiQA, LogiQA 2, MMLU, OpenBookQA, ARC easy, ARC challenge, HellaSwag, COPA, and PiQA, in Figure 1 and Table 1.

W2: No analysis of Dense Backprop on inference workloads

We have modified our method to only apply the dense approximation in the backward pass. We observed this led to improved results in training and evaluation. This means that at inference, no additional forward computation is necessary. So for inference workloads, our method is identical to the baseline. We have updated the paper to reflect this in lines 198-201.

W3: Should not mention Attention approximation methods

We agree with your concern on presenting the two additional attention methods. Although they demonstrated strong preliminary results, ultimately we came to the conclusion that they cannot be made efficient enough to justify the improvement. We moved the entire section on our attention to the appendix as optional supplementary material.

W4: Compare to more recent work in MoE Routing

Thank you for the references. We have updated the paper to compare to more recent work in MoE routing. Our baseline is DeepSeek (we use their fine-grained MoEs) and we use auxiliary loss. GRIN uses SparseMixer, so we have added a comparison to SparseMixer in Figure 8. We reproduce the text here for convenience:

\citet{liu2023sparsebackpropagationmoetraining, liu2024gringradientinformedmoe} use Sparsemixer, which estimates the true router gradient without straight-through. Our method is distinct in that we also provide gradient updates for the expert weights(as seen in Figure 8). \citet{liu2024gringradientinformedmoe} note that Sparsemixer lags behind TopK~(which our method always outperforms) for the first 0.5T tokens, likely due to the noise that Sparsemixer adds. In the 10B tokens that we were able to run Sparsemixer for, our method significantly outperforms it.

Unfortunately, we have not been able to invest 0.5T tokens into this baseline comparison, but we are currently running this and expect to have it done for the camera ready. By the end of the author response period, we will update this number with a longer run.

Q1: Ablate K

To answer your question on values of K, we experimented with multiple of K and consistently observed an improvement in training. We report K=4 in Table 2 and K=1 in Table 3, and find that we still outperform the baseline with these other values of K.

Q2: Finetuning

We are currently working on setting up fine-tuning experiments and will have the results of these by the camera ready. Our intuition is that by allowing the router and experts to receive a dense gradient, we can potentially overcome suboptimal routing decisions that were learned during pretraining.

Dear Reviewer X235,

Thank you for your generous review. We appreciate you highlighting the quality and efficiency of our method, along with the overall writing of our paper. In this response, we aim to clarify a couple of questions.

W1: Averaging may not work well when the experts are non-linear

You are correct in mentioning that the experts are non-linear. To answer your question, we are not assuming that the experts will be linear – our assumption for the approximation is that nearby tokens will have a similar expert output and in a large training batch they will span the missing expert output. Additionally, our primary motivation is not to approximate any individual expert output, but instead to approximate the entire batch gradient. We demonstrate the efficacy of this batch approximation in Figure 6. A higher-order approximation is possible, but it would be more computationally expensive and our goal is to minimize overhead

W2: Perplexity drop in Table 4 is small

The perplexity drop is small, and the average benchmark improvement (in Table 1) is also small. However, we believe the important factor for our method is that it is a free lunch so long as the benefit outweighs the overhead. We have carefully documented how the performance improves as we scale up in Table 2, and we find that we improve over the baseline more as we scale up training. And this is scaling in our constrained academic budget, so real deployments will be much larger and we anticipate the improvement of our method over the baseline will be far greater. We also document the overhead scaling in Section 4.4, through Figure 8 and Table 4, and find that for larger models, the overhead decreases. We also note that our main perplexity and benchmark improvement results are without the method variants that we have introduced, and as we see in Table 3, these variants improve our method even more. In summary: our improvement is marginal at small scale, but grows as we scale up training.

Q1: Do we assume the experts are linear?

Our expert group approximation does not assume that the experts are linear. Rather, it assumes that an expert output for a given token is in the span of the expert outputs of other tokens.

Q2: Could something be done to approximate the expert gradients too?

Our method actually also sends an additional gradient update to experts using this method. This is explained in the new draft of the paper in lines 266-301.