Drop-Upcycling: Entrenamiento de Mezclas Escasas de Expertos con Reinicialización 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

Resumen

La arquitectura Mixture of Experts (MoE) reduce significativamente los costos de entrenamiento e inferencia en comparación con un modelo denso de capacidad equivalente. El reciclaje (upcycling) es un enfoque que inicializa y entrena un modelo MoE utilizando un modelo denso preentrenado. Si bien el reciclaje conduce a mejoras iniciales en el rendimiento, el entrenamiento avanza más lento que cuando se entrena desde cero, lo que resulta en un rendimiento subóptimo a largo plazo. Proponemos Drop-Upcycling, un método que aborda eficazmente este problema. Drop-Upcycling combina dos enfoques aparentemente contradictorios: utilizar el conocimiento de modelos densos preentrenados mientras se reinicializan estadísticamente algunas partes de los pesos. Este enfoque promueve estratégicamente la especialización de los expertos, mejorando significativamente la eficiencia del modelo MoE en la adquisición de conocimiento. Experimentos extensos a gran escala demuestran que Drop-Upcycling supera significativamente a los métodos anteriores de construcción de MoE a largo plazo, específicamente cuando se entrena con cientos de miles de millones de tokens o más. Como resultado, nuestro modelo MoE con 5.9 mil millones de parámetros activos logra un rendimiento comparable al de un modelo denso de 13 mil millones en la misma familia de modelos, mientras requiere aproximadamente 1/4 de los FLOPs de entrenamiento. Todos los recursos experimentales, incluidos el código fuente, los datos de entrenamiento, los puntos de control del modelo y los registros, están disponibles públicamente para promover la reproducibilidad y futuras investigaciones 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.