Drop-Upcycling: Treinamento de uma Mistura Esparsa de Especialistas com Re-inicialização Parcial

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

February 26, 2025

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

cs.AI

Resumo

A arquitetura Mixture of Experts (MoE) reduz significativamente o custo de treinamento e inferência em comparação com um modelo denso de capacidade equivalente. O upcycling é uma abordagem que inicializa e treina um modelo MoE usando um modelo denso pré-treinado. Embora o upcycling resulte em ganhos de desempenho iniciais, o treinamento avança mais lentamente do que quando treinado do zero, levando a um desempenho subótimo a longo prazo. Propomos o Drop-Upcycling - um método que aborda efetivamente esse problema. O Drop-Upcycling combina duas abordagens aparentemente contraditórias: aproveitar o conhecimento de modelos densos pré-treinados enquanto reinicializa estatisticamente algumas partes dos pesos. Essa abordagem promove estrategicamente a especialização de especialistas, melhorando significativamente a eficiência do modelo MoE na aquisição de conhecimento. Experimentos extensivos em grande escala demonstram que o Drop-Upcycling supera significativamente os métodos anteriores de construção de MoE a longo prazo, especificamente ao treinar em centenas de bilhões de tokens ou mais. Como resultado, nosso modelo MoE com 5,9B de parâmetros ativos alcança desempenho comparável a um modelo denso de 13B na mesma família de modelos, enquanto requer aproximadamente 1/4 dos FLOPs de treinamento. Todos os recursos experimentais, incluindo código-fonte, dados de treinamento, checkpoints e logs do modelo, estão publicamente disponíveis para promover a reprodutibilidade e pesquisas futuras sobre 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.