Sparse Backpropagation for MoE Training

Thank you for your comprehensive feedback. We appreciate the opportunity to clarify and expand on the strengths of our approach, with further elaborations to be included in the final paper.

Reply to presentation:

We appreciate your suggestion regarding the term. We have retained "MoE" in our manuscript, as it is a widely recognized term in the literature, predominantly referring to sparsely activated networks [1,2,3]. Also, we would like to acknowledge that "SMoE" could potentially lead to confusion, as it might be misconstrued as Soft MoE [4].

1. Adaptive mixtures of local experts. Neural Computation 1991
2. Hierarchical mixtures of experts and the em algorithm. Neural Computation 1994
3. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. 2021
4. From Sparse to Soft Mixtures of Experts. NeurIPS 2022

Reply to Efficiency Evaluation:

Table 3 in our initial submission summarized the wall-clock time per epoch (excluding checkpoint saving) and took into account all components of SparseMixer, including omega scaling. To provide further clarity, we have now added Figure 5, which compares the training FLOPs ratio between Switch and Switch+SparseMixer configurations.

Reply to Omega Scaling: To evaluate the impact of Omega scaling, we conducted a specific ablation study, which is presented as ablation-1 SparseMixer variant in Figure 1 of our submission. This variant excludes Omega scaling from SparseMixer, allowing for a direct comparison of its effects.

Reply to Research Scope: Early on, neural networks were trained using a variety of update rules, such as the classical Winnow algorithm and the Perceptron algorithm. Given that these algorithms were largely handcrafted, generalizing these update rules to accommodate deeper and more intricate neural networks posed significant challenges.

The introduction of the backpropagation algorithm and the stochastic gradient descent algorithm re-anchored the foundation of neural networks to the principles of gradient descent and the chain rule, mirroring the foundational ideas behind least squares computation. This paradigm shift has been instrumental, with the vast majority of neural network research post-2000 relying on this update rule.

However, backpropagation cannot be directly applied to sparse model training. Accordingly, many attempts have been made to derive ad-hoc proxy for gradient computation and gradient descent. However, there is no guarantee the resulting estimation can be viewed as a gradient approximation, and it is not clear doing gradient descent with these proxies can lead to a descent.

Our study reveals the underlying connection between the current MoE training practice and gradient approximation. It not only re-anchored MoE training as gradient descent at-scale paradigm, but also provides guidance on how to obtain better gradient estimation.

Since our algorithm provides better gradient estimation, the gradient-based optimizer would be more effective at more training, thus leading to better model training (accelerating model training by up to 2 times).

Reply to Experiment Design:

In our submission, we adhered to the experiment design of the existing study, focusing on normal-scale problems for controllability and resource efficiency. For example, choosing ELECTRA-base as the pre-training task allows us to conduct experiments with multiple settings. As a comparison, typical ELECTRA-base model training takes 125k steps and typical BERT-base model training takes 1M steps.

While our method has potential applications in different scales and MoE gating methods, the focus of this study is on foundational gradient estimation. Scaling and adapting our method for varied MoE paradigms presents both research and engineering challenges. We aim to explore these in future work, considering our current computational resource constraints.

Reply to SparseMixer Application on Top-k Routing:

SparseMixer is for estimating gradients involving discrete random variables with only sparse backpropagation. In principle, it is applicable to other MoE routing algorithms. The only bottleneck for applying SparseMixer is that it requires the expert index is explicitly sampled from a parameterized distribution, which requires us to apply some adaptation to Switch Transformer / top-1 routing.

I appreciate the authors' efforts addressing my concern. However, after going through the revision and the posted comments, I am not entirely convinced that this work is ready for publications. My concerns regarding the contributions and presentation (minor) still remains. Particularly:

• The newly added results in Figure 1 are poorly discussed, the Figure is placed 3 pages before it is mentioned.
• The contribution of $\omega-$scaling is poorly discussed. Given Figure 1, I could argue that in many cases, most of the contribution comes from the $\omega-$scaling while the proposed method only contribute marginally.
• The authors failed to report the total training time or explain why SparseMixer could achieve the same complexity despite introducing several computational steps.
• The replies to regarding the extension to $k\ge 2$ or application on different architectures are not satisfactory. I expected the authors to outline some steps or challenges to extend to cases $k\ge 2$ and provide a minimal experiment/insight with a different architecture. However, none of such results are presented.
• Naming convention (minor): the name Sparse Mixer has been used in an existing work [A] for a different method.

Although the approach is somewhat interesting, I believe the current draft still has many issues. Thus, I keep my original rating of this work and don't endorse the paper to be accepted at its current state.

[A] Lee-Thorp, James, and Joshua Ainslie. "Sparse Mixers: Combining MoE and Mixing to build a more efficient BERT." arXiv preprint arXiv:2205.12399 (2022).

Thank you for your comprehensive feedback. We appreciate the opportunity to clarify and expand on the strengths of our approach, with further elaborations to be included in the final paper.

Reply to weakness 1 (scope of study)

Our study primarily focuses on gradient estimation in MoE training.

Early on, neural networks were trained using a variety of update rules, such as the classical Winnow algorithm and the Perceptron algorithm. Given that these algorithms were largely handcrafted, generalizing these update rules to accommodate deeper and more intricate neural networks posed significant challenges.

The introduction of the backpropagation algorithm and the stochastic gradient descent algorithm re-anchored the foundation of neural networks to the principles of gradient descent and the chain rule, mirroring the foundational ideas behind least squares computation. This paradigm shift has been instrumental, with the vast majority of neural network research post-2000 relying on this update rule.

However, backpropagation cannot be directly applied to sparse model training. Accordingly, many attempts have been made to derive ad-hoc proxy for gradient computation and gradient descent. However, there is no guarantee the resulting estimation can be viewed as a gradient approximation, and it is not clear doing gradient descent with these proxies can lead to a descent.

Our work bridges this gap by revealing the intrinsic connection between current MoE training practices and gradient approximation. This not only repositions MoE training within the gradient descent paradigm at scale but also guides more accurate gradient estimation.

Reply to weakness 1 (significance of performance gain):

We believe the improvement brought by SparseMixer is significant. Regarding training convergence, SparseMixer accelerates model training by up to two times. Regarding downstream performance, the best performance achieved by SparseMixer is 28.72 on WMT’14 EN-DE (note our architecture configuration is based on Transformer-base) and 88.31 on GLUE (note our architecture configuration is based on Electra-base).

It's important to note that the improvements on GLUE are actually significant, as GLUE is a very competitive benchmark, Most leading methods on the GLUE benchmark (as in https://gluebenchmark.com/leaderboard) are based on the Electra-large setting and their improvements over the previous methods are usually within 0.2 absolute points. In our study, simply applying SparseMixer to Electra-base training without any pre-training hyper-parameter re-tuning, we achieve a consistent 0.5 absolute point improvements.

While our method has potential applications in different scales and MoE gating methods, the focus of this study is on foundational gradient estimation. Scaling and adapting our method for varied MoE paradigms presents both research and engineering challenges. We aim to explore these in future work, considering our current computational resource constraints.

Reply to weakness 2 (Section 3 writing) We appreciate your feedback on Section 3. The section has been revised for improved coherence, focusing on the rationale and intuition behind our techniques rather than merely describing experimental procedures.

Reply to question 1 In our submission, Table 3 summarizes the wall-clock time for the model training (measured by total training time per epoch without checkpoint saving, divided by number of updates). In our revised submission, Figure 5 summarizes the training FLOPs ratio between Switch and Switch+SparseMixer.

Reply to question 2 Following previous study, we used the nvidia/pytorch:22.04-py3 docker image for the pre-training task and the nvidia/pytorch:22.02-py3docker image for the NMT task. The training was primarily conducted on V100 GPUs, with some fine-tuning on P40/P100 GPUs.

Thank you for your comprehensive feedback. We appreciate the opportunity to clarify and expand on the strengths of our approach, with further elaborations to be included in the final paper.

Reply to weakness 1 (terminology) We choose ST as the abbreviation, since ST is a widely used terminology in the literature [1,2,3].

[1] Bridging Discrete and Backpropagation: Straight-Through and Beyond. NeurIPS 2023

[2] Training discrete deep generative models via gapped straight-through estimator. ICML 2022

[3] Rao-blackwellizing the straight-through gumbel softmax gradient estimator. ICLR 2021

Reply to weakness 2 (Definition of ODE in Abstract) Thanks for the suggestion! We have revised our abstract to include a clear definition of Ordinary Differential Equations (ODE) at its first mention.

Reply to weakness 3 (Definition of $\nabla_0$ and $\nabla_1$) $\nabla_0$ and $\nabla_1$ are defined in Section 2 (Equation 2). In response to your feedback, we have added references in Section 3 of the revised manuscript for easier navigation and understanding.

Reply to weakness 4 (Clarification on the Mechanism of SparseMixer) The improved training speed is a result of our enhanced gradient estimation.

Early on, neural networks were trained using a variety of update rules, such as the classical Winnow algorithm and the Perceptron algorithm. Given that these algorithms were largely handcrafted, generalizing these update rules to accommodate deeper and more intricate neural networks posed significant challenges.

The introduction of the backpropagation algorithm and the stochastic gradient descent algorithm re-anchored the foundation of neural networks to the principles of gradient descent and the chain rule, mirroring the foundational ideas behind least squares computation. This paradigm shift has been instrumental, with the vast majority of neural network research post-2000 relying on this update rule.

Since backpropagation cannot be directly applied to sparse model training, many attempts have been made to derive ad-hoc proxy for gradient computation and gradient descent. However, there is no guarantee the resulting estimation can be viewed as a gradient approximation, and it is not clear doing gradient descent with these proxies can lead to a descent.

Our study reveals the underlying connection between the current MoE training practice and gradient approximation. It not only re-anchored MoE training as gradient descent at-scale paradigm, but also provides guidance on how to obtain better gradient estimation.

Since our algorithm provides better gradient estimation, the gradient-based optimizer would be more effective at more training, thus leading to better model training.

Reply to weakness 5 (Exploration of Other MoE Gating Methods) While our method has potential applications in different scales and MoE gating methods, the focus of this study is on foundational gradient estimation. Scaling and adapting our method for varied MoE paradigms presents both research and engineering challenges. We aim to explore these in future work, considering our current computational resource constraints.

Thank you for your comprehensive feedback. We appreciate the opportunity to clarify and expand on the strengths of our approach, with further elaborations to be included in the final paper.

Reply to weakness 1 (experiment design and performance on GLUE):

In our submission, we adhered to the experiment design of the existing study, focusing on normal-scale problems for controllability and resource efficiency. For example, choosing ELECTRA-base as the pre-training task allows us to conduct experiments with multiple settings. As a comparison, typical ELECTRA-base model training takes 125k steps and typical BERT-base model training takes 1M steps.

As to performance, we believe the improvement brought by SparseMixer is significant. Regarding training convergence, SparseMixer accelerates model training by up to two times. Regarding downstream performance, the best performance achieved by SparseMixer is 28.72 on WMT’14 EN-DE (note our architecture configuration is based on Transformer-base) and 88.31 on GLUE (note our architecture configuration is based on Electra-base).

It's important to note that the improvements on GLUE are actually significant, as GLUE is a very competitive benchmark. Most leading methods on the GLUE benchmark (as in https://gluebenchmark.com/leaderboard) are based on the Electra-large setting and their improvements over the previous methods are usually within 0.2 absolute points. In our study, simply applying SparseMixer to Electra-base training without any pre-training hyper-parameter re-tuning, we achieve a consistent 0.5 absolute point improvements.

Expanding our experiments to larger-scale problems or different MoE training paradigms is an exciting direction for future research. We acknowledge the combined research and engineering challenges this entails and look forward to exploring these possibilities as resources permit.

Reply to weakness 2 (Gradient Computation and Connection to Current Practice) Thanks for the insightful suggestion! In our revised manuscript, we delve deeper into the relationship between our proposed first-order approximation and existing MoE training methodologies. Our analysis reveals that the common practice of ignoring certain gradient components is akin to reducing gradient scaling, a factor that adaptive optimizers like Adam can inherently adjust to. This insight not only validates our approach but also elucidates the underlying mechanics of current MoE training practices.

To demonstrate this, please see the comparison between Switch and SparseMixer-ablation-2 (it fixes such additional bias by changing the sampling and avoiding the routing gradient downscale) in Figure 1. For more detailed discussion, please see Section 3.3 in the revised submission.

Reply to weakness 3 (Computational Overheads of SparseMixer) SparseMixer has the same compute complexity and memory complexity with the current practice (dropping $\nabla_0$ completely). To further demonstrate this, we visualized the FLOPS ratio of Switch and SparseMixer in Figure 5 for $\\{2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384\\}$. The computation overheads brought by SparseMixer only compose for less than $0.1\%$ of the training FLOPS.