Drop-Upcycling: Training von Sparse Mixture of Experts mit partieller Re-Initialisierung

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

February 26, 2025

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

cs.AI

Zusammenfassung

Die Mixture-of-Experts (MoE)-Architektur reduziert die Trainings- und Inferenzkosten im Vergleich zu einem dichten Modell gleicher Kapazität erheblich. Upcycling ist ein Ansatz, bei dem ein MoE-Modell mithilfe eines vortrainierten dichten Modells initialisiert und trainiert wird. Obwohl Upcycling anfängliche Leistungssteigerungen bewirkt, schreitet das Training langsamer voran als bei einem Training von Grund auf, was langfristig zu suboptimalen Ergebnissen führt. Wir schlagen Drop-Upcycling vor – eine Methode, die dieses Problem effektiv adressiert. Drop-Upcycling kombiniert zwei scheinbar widersprüchliche Ansätze: die Nutzung des Wissens vortrainierter dichter Modelle bei gleichzeitiger statistischer Reinitialisierung einiger Gewichtsteile. Dieser Ansatz fördert strategisch die Spezialisierung der Experten und verbessert dadurch die Effizienz des MoE-Modells bei der Wissensaneignung signifikant. Umfangreiche groß angelegte Experimente zeigen, dass Drop-Upcycling langfristig, insbesondere beim Training mit Hunderten von Milliarden Token oder mehr, frühere MoE-Konstruktionsmethoden deutlich übertrifft. Infolgedessen erreicht unser MoE-Modell mit 5,9 Milliarden aktiven Parametern eine vergleichbare Leistung wie ein 13 Milliarden Parameter umfassendes dichtes Modell derselben Modellfamilie, benötigt jedoch nur etwa 1/4 der Trainings-FLOPs. Alle experimentellen Ressourcen, einschließlich Quellcode, Trainingsdaten, Modell-Checkpoints und Protokolle, sind öffentlich verfügbar, um die Reproduzierbarkeit und zukünftige Forschung zu MoE zu fördern.

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.