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.

通过 SparseMixer 建立了稀疏计算与反向传播之间的桥梁，提供可靠的梯度估计，并加速了 Switch Transformer 的训练收敛速度。