Drop-Upcycling: Addestramento di Sparse Mixture of Experts con Re-inizializzazione Parziale

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

February 26, 2025

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

cs.AI

Abstract

L'architettura Mixture of Experts (MoE) riduce significativamente i costi di addestramento e inferenza rispetto a un modello denso di capacità equivalente. L'upcycling è un approccio che inizializza e addestra un modello MoE utilizzando un modello denso pre-addestrato. Sebbene l'upcycling porti a guadagni iniziali di prestazioni, l'addestramento progredisce più lentamente rispetto a un addestramento da zero, portando a prestazioni subottimali a lungo termine. Proponiamo il Drop-Upcycling, un metodo che affronta efficacemente questo problema. Il Drop-Upcycling combina due approcci apparentemente contraddittori: sfruttare la conoscenza dei modelli densi pre-addestrati mentre reinizializza statisticamente alcune parti dei pesi. Questo approccio promuove strategicamente la specializzazione degli esperti, migliorando significativamente l'efficienza del modello MoE nell'acquisizione di conoscenza. Esperimenti su larga scala dimostrano che il Drop-Upcycling supera significativamente i precedenti metodi di costruzione di MoE a lungo termine, in particolare quando si addestra su centinaia di miliardi di token o più. Di conseguenza, il nostro modello MoE con 5,9 miliardi di parametri attivi raggiunge prestazioni comparabili a un modello denso da 13 miliardi nella stessa famiglia di modelli, richiedendo circa 1/4 dei FLOP di addestramento. Tutte le risorse sperimentali, inclusi codice sorgente, dati di addestramento, checkpoint del modello e log, sono pubblicamente disponibili per promuovere la riproducibilità e future ricerche su 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.