Upcycling Large Language Models into Mixture of Experts (2410.07524v1)

Abstract: Upcycling pre-trained dense LLMs into sparse mixture-of-experts (MoE) models is an efficient approach to increase the model capacity of already trained models. However, optimal techniques for upcycling at scale remain unclear. In this work, we conduct an extensive study of upcycling methods and hyperparameters for billion-parameter scale LLMs. We propose a novel "virtual group" initialization scheme and weight scaling approach to enable upcycling into fine-grained MoE architectures. Through ablations, we find that upcycling outperforms continued dense model training. In addition, we show that softmax-then-topK expert routing improves over topK-then-softmax approach and higher granularity MoEs can help improve accuracy. Finally, we upcycled Nemotron-4 15B on 1T tokens and compared it to a continuously trained version of the same model on the same 1T tokens: the continuous trained model achieved 65.3% MMLU, whereas the upcycled model achieved 67.6%. Our results offer insights and best practices to effectively leverage upcycling for building MoE LLMs.

Overview of "Upcycling LLMs into Mixture of Experts"

The paper "Upcycling LLMs into Mixture of Experts" explores methodologies for transforming pre-trained dense LLMs into sparse Mixture of Experts (MoE) models. This conversion, referred to as upcycling, aims to leverage existing model parameters to increase model capacity efficiently, thereby enhancing performance without retraining from scratch.

Key Contributions

The authors introduce several novel techniques and provide a comprehensive analysis of upcycling methods. Key contributions are:

• Virtual Group Initialization and Weight Scaling: The paper proposes a "virtual group" initialization scheme to facilitate the conversion into fine-grained MoE architectures. Additionally, a weight scaling approach is introduced, improving MoE models’ performance by around 1.5% compared to those without it.
• Effective Routing Strategies: The paper examines routing techniques, finding that a softmax-then-topK strategy outperforms the topK-then-softmax approach. This finding contributes to enhancing accuracy and providing insights into routing efficiencies.
• Granularity in MoE Models: By increasing the granularity, where tokens are routed to more but smaller experts, the paper illustrates improvements in model accuracy. This approach, however, comes with the trade-off of increased computational complexity.
• Extensive Hyperparameter Study: The paper provides detailed experimental results regarding various hyperparameters like learning rates and batch sizes, underscoring their significance in MoE efficiency and accuracy.

Experimental Findings

The experimental results highlight several notable outcomes:

• Efficiency of Upcycling: Compared to continuously trained dense models, upcycled models exhibit superior performance. For instance, upcycling Nemotron-4 15B led to a notable increase in Multi-task Language Understanding (MMLU) score from 65.3% to 67.6%.
• Learning Rate and Batch Size Effects: A resetting learning rate strategy and larger batch sizes were shown to significantly benefit the upcycling process. Specifically, resetting the learning rate allowed the model to escape local minima, leading to better performance.
• Granularity and Scaling: The experiments showed diminishing returns beyond a certain granularity level, with the optimal setup being model-dependent. For instance, a configuration with 64 experts achieved better performance than one with 256 experts.

Implications and Future Directions

The research presents significant implications for the development of large-scale AI models. By supporting more efficient use of pre-trained models, the proposed upcycling methods offer a pathway to scaling LLMs with limited computational resources. The paper underscores the need for further exploration in expert diversity and utilization, potentially paving the way for integrating more complex routing mechanisms and architectural innovations.

Future developments could focus on optimizing the upcycling process for larger models and datasets. Enhancements in this area could significantly impact the efficiency and scalability of LLMs, facilitating their deployment in varied real-world applications.

In conclusion, this paper contributes valuable insights and methodologies that advance the understanding of upcycling techniques in AI. Its findings are crucial for researchers aiming to enhance model efficiency and performance while utilizing existing computational resources judiciously.