Backpropagation, the cornerstone of deep learning, is limited to computing gradients for continuous variables. This limitation poses challenges for problems involving discrete latent variables and sparse computations. To address the challenge posed by discrete latent variables, we first assess the Straight-Through (ST) heuristic, formally establishing it works as a first-order gradient approximation. Guided by our findings, we propose ReinMax, which achieves second-order accuracy by integrating Heun's method, a second-order numerical method for solving ODEs. ReinMax does not require Hessian or other second-order derivatives, thus having negligible computation overheads. Extensive experimental results on benchmark datasets demonstrate the superiority of ReinMax over the state of the art. 

In addition, the inherent sparse computation of the Mixture-of-Expert (MoE) models poses unique challenges on gradient estimations. To reconciles the dense backpropagation with sparse expert routing, we present SparseMixer, a gradient estimator specifically designed for MoE. 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. Also grounded in a numerical ODE framework, SparseMixer harnesses the mid-point method, a different second-order ODE solver, to deliver precise gradient approximations with negligible computational overhead. Applying SparseMixer to Switch Transformer on both pretraining and machine translation tasks, SparseMixer showcases considerable performance gain, accelerating training convergence by up to 2 times.