Dense Backpropagation Improves Training for Sparse Mixture-of-Experts
This addresses training inefficiencies for researchers and practitioners using sparse MoE models in scalable pretraining, representing an incremental improvement over existing routing methods.
The paper tackles the problem of training instability and suboptimal performance in Mixture-of-Experts (MoE) models due to sparse backward updates by introducing a lightweight approximation method called Default MoE, which substitutes missing expert activations with default outputs to provide dense gradient updates, leading to significant improvements in training performance and outperforming standard TopK routing in various settings.
Mixture of Experts (MoE) pretraining is more scalable than dense Transformer pretraining, because MoEs learn to route inputs to a sparse set of their feedforward parameters. However, this means that MoEs only receive a sparse backward update, leading to training instability and suboptimal performance. We present a lightweight approximation method that gives the MoE router a dense gradient update while continuing to sparsely activate its parameters. Our method, which we refer to as Default MoE, substitutes missing expert activations with default outputs consisting of an exponential moving average of expert outputs previously seen over the course of training. This allows the router to receive signals from every expert for each token, leading to significant improvements in training performance. Our Default MoE outperforms standard TopK routing in a variety of settings without requiring significant computational overhead. Code: https://github.com/vatsal0/default-moe.