Sparse Backpropagation for MoE Training
This addresses a bottleneck in MoE training for scalable deep learning, offering a method to improve efficiency without sacrificing performance.
The paper tackles the challenge of dense backpropagation in Mixture-of-Expert (MoE) models by introducing SparseMixer, a scalable gradient estimator that bridges sparse expert routing with backpropagation, resulting in up to 2 times faster training convergence on Switch Transformer tasks.
One defining characteristic of Mixture-of-Expert (MoE) models is their capacity for conducting sparse computation via expert routing, leading to remarkable scalability. However, backpropagation, the cornerstone of deep learning, requires dense computation, thereby posting challenges in MoE gradient computations. Here, we introduce SparseMixer, a scalable gradient estimator that bridges the gap between backpropagation and sparse expert routing. Unlike typical MoE training which strategically neglects certain gradient terms for the sake of sparse computation and scalability, SparseMixer provides scalable gradient approximations for these terms, enabling reliable gradient estimation in MoE training. Grounded in a numerical ODE framework, SparseMixer harnesses the mid-point method, a second-order ODE solver, to deliver precise gradient approximations with negligible computational overhead. Applying SparseMixer to Switch Transformer on both pre-training and machine translation tasks, SparseMixer showcases considerable performance gain, accelerating training convergence up to 2 times.