🤖 AI Summary
This work addresses the challenge of gradient starvation and knowledge loss in Mixture-of-Experts (MoE) models during supervised fine-tuning, which arises from sparse activation leading to underutilization of long-tail experts. The authors propose an end-to-end fine-tuning framework that eliminates the need for auxiliary load-balancing losses by integrating bias-driven sparse routing with a gating mechanism that keeps a subset of experts consistently active. This approach effectively preserves the activity of task-relevant experts and mitigates gradient starvation without relying on conventional balancing objectives. Evaluated on large-scale MoE models, the method substantially outperforms existing approaches such as DenseMixer and ESFT, achieving average performance gains exceeding 2.5% on mathematical reasoning and CommonsenseQA benchmarks.
📝 Abstract
Despite MoE models leading many benchmarks, supervised fine-tuning (SFT) for the MoE architectures remains difficult because its router layers are fragile. Methods such as DenseMixer and ESFT mitigate router collapse with dense mixing or auxiliary load-balancing losses, but these introduce noisy gradients that often degrade performance. In preliminary experiments, we systematically pruned experts and observed that while certain super experts are activated far more frequently, discarding less used experts still leads to notable performance degradation. This suggests that even rarely activated experts encode non-trivial knowledge useful for downstream tasks. Motivated by this, we propose an auxiliary-loss-free MoE SFT framework that combines bias-driven sparsification with always-active gated condenser experts. Rather than enforcing balanced activation across all experts, our method encourages task-relevant experts to remain active while pushing long-tailed experts toward inactivity. The condenser experts provide a persistent, learnable pathway that alleviates gradient starvation and facilitates consolidation of information that would otherwise remain fragmented across sparsely activated experts. Analysis further suggest that this design better preserves long-tailed expert information under sparse routing. Experiments on large-scale MoE models demonstrate that our approach outperforms state-of-the-art SFT baselines such as DenseMixer and ESFT, achieving average gain of 2.5%+ on both mathematical reasoning and commonsenseQA benchmarks.