Dense Backpropagation Improves Training for Sparse Mixture-of-Experts

TL;DR

This work tackles the training instability of sparse Mixture-of-Experts (MoE) models by enabling dense gradient signals to flow through the router without sacrificing sparse forward computation. It introduces Default MoE, which uses per-expert exponential moving average (EMA) default outputs to substitute non-activated expert activations in the backward pass, yielding a dense router gradient with minimal overhead. Empirically, Default MoE outperforms standard Top-K routing across multiple configurations and datasets, achieving faster convergence (e.g., ~9% fewer tokens to reach a target perplexity) while maintaining throughput and memory efficiency. This approach enhances MoE scalability and training stability, providing a practical improvement for large-scale pretraining on sparse architectures.

Abstract

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.

Figures (44)

• Figure 1: Overview of Routing with Dense Approximations. The original MoE router only receives gradients corresponding to experts the token is routed to, because there is no output from other experts. Our approach provides the router with a complete (dense) gradient by letting non-activated experts contribute a default vector that approximates its output without the cost of a forward pass. As indicated by the dashed green arrows, the approximated gradients are not actually connected to the token in the computation graph; instead, they are artificially applied in the backward pass.
• Figure 2: Comprehensive Ablations. We report the perplexity after training for 10B tokens for Default MoE and Top-K while varying: (a) MoE configuration across five MoE setups. (b) Learning rate while keeping the MoE configuration at 8c1. (c) Total model size while keeping the MoE configuration at 8c1. (d) Sparsity, where the final 1/64 sparsity is for a 7B model, therefore the perplexity is lower than the other sparsity factors that are for 2B models.
• Figure 3: Analysis. Left: Router Entropy after training for 10B tokens. Middle, Right: Heatmap of the cosine similarities between default vectors for the 8c1 DefaultMoE at layers 0 and 16.
• Figure 4: Dense Router Gradient Similarity. We plot similarity to the dense router gradient (K=8) for our 8ck Top-K and DefaultMoEs pretrained on 160B tokens. The DefaultMoE's router gradient is much more similar to the dense router gradient.
• Figure 5: Comparison of Default MoE and SparseMixer. We compare Default MoE with SparseMixer, both configured with 8 experts and Top-K=2 active experts. The results report training perplexity throughout training, demonstrating that Default MoE consistently outperforms SparseMixer.