Reproduce “Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts” on OpenResearch

About this paper

Chaitanya Dwivedi, Binxuan Huang, Himanshu Gupta, Pratik Jayarao, Neeraj Varshney, Bing Yin

Researchers from Amazon Stores Foundation AI and Carnegie Mellon University introduce "expert upcycling," a two-phase training method that progressively increases the number of experts in a Mixture-of-Experts (MoE) model during continued pre-training while maintaining per-token inference cost. This approach significantly reduces the GPU-hours required to achieve the quality of large MoEs trained from scratch, demonstrating compute savings of 24-32% and achieving comparable or superior performance on validation loss and downstream tasks.

Key insights

• Warm initialization, achieved by copying parameters of existing experts and perturbing router biases, is crucial for efficient MoE capacity expansion, leading to significantly lower initial loss than random initialization.
• Expert specialization and diversification in upcycled models primarily arise organically through stochastic gradients and load balancing during continued pre-training, rather than requiring complex diversity-inducing initialization heuristics.
• The compute-efficient frontier for MoEs can be shifted by amortizing training costs, allowing models to reach the quality of large MoEs with substantial savings, particularly when leveraging existing checkpoints.

Summary by alphaXiv.