Drop-Upcycling : Entraînement d'un mélange parcimonieux d'experts avec ré-initialisation partielle

Drop-Upcycling: Training Sparse Mixture of Experts with Partial Re-initialization

February 26, 2025

Auteurs: Taishi Nakamura, Takuya Akiba, Kazuki Fujii, Yusuke Oda, Rio Yokota, Jun Suzuki

cs.AI

Résumé

L'architecture Mixture of Experts (MoE) réduit considérablement les coûts d'entraînement et d'inférence par rapport à un modèle dense de capacité équivalente. Le recyclage (upcycling) est une approche qui initialise et entraîne un modèle MoE en utilisant un modèle dense pré-entraîné. Bien que le recyclage entraîne des gains de performance initiaux, la progression de l'entraînement est plus lente que lorsqu'il est effectué à partir de zéro, ce qui conduit à des performances sous-optimales à long terme. Nous proposons Drop-Upcycling - une méthode qui résout efficacement ce problème. Drop-Upcycling combine deux approches apparemment contradictoires : l'utilisation des connaissances des modèles denses pré-entraînés tout en réinitialisant statistiquement certaines parties des poids. Cette approche favorise stratégiquement la spécialisation des experts, améliorant ainsi significativement l'efficacité du modèle MoE dans l'acquisition de connaissances. Des expériences à grande échelle démontrent que Drop-Upcycling surpasse nettement les méthodes précédentes de construction de MoE à long terme, en particulier lors de l'entraînement sur des centaines de milliards de tokens ou plus. En conséquence, notre modèle MoE avec 5,9 milliards de paramètres actifs atteint des performances comparables à un modèle dense de 13 milliards de paramètres dans la même famille de modèles, tout en nécessitant environ 1/4 des FLOPs d'entraînement. Toutes les ressources expérimentales, y compris le code source, les données d'entraînement, les points de contrôle des modèles et les journaux, sont publiquement disponibles pour promouvoir la reproductibilité et les recherches futures sur MoE.

English

The Mixture of Experts (MoE) architecture reduces the training and inference cost significantly compared to a dense model of equivalent capacity. Upcycling is an approach that initializes and trains an MoE model using a pre-trained dense model. While upcycling leads to initial performance gains, the training progresses slower than when trained from scratch, leading to suboptimal performance in the long term. We propose Drop-Upcycling - a method that effectively addresses this problem. Drop-Upcycling combines two seemingly contradictory approaches: utilizing the knowledge of pre-trained dense models while statistically re-initializing some parts of the weights. This approach strategically promotes expert specialization, significantly enhancing the MoE model's efficiency in knowledge acquisition. Extensive large-scale experiments demonstrate that Drop-Upcycling significantly outperforms previous MoE construction methods in the long term, specifically when training on hundreds of billions of tokens or more. As a result, our MoE model with 5.9B active parameters achieves comparable performance to a 13B dense model in the same model family, while requiring approximately 1/4 of the training FLOPs. All experimental resources, including source code, training data, model checkpoints and logs, are publicly available to promote reproducibility and future research on MoE.