Drop-Upcycling: Обучение разреженной смеси экспертов с частичной повторной инициализацией

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

February 26, 2025

Авторы: Taishi Nakamura, Takuya Akiba, Kazuki Fujii, Yusuke Oda, Rio Yokota, Jun Suzuki

cs.AI

Аннотация

Архитектура "Смесь экспертов" (Mixture of Experts, MoE) значительно снижает затраты на обучение и вывод по сравнению с плотной моделью эквивалентной мощности. Метод "апсайклинга" заключается в инициализации и обучении модели MoE с использованием предварительно обученной плотной модели. Хотя апсайклинг приводит к начальному улучшению производительности, процесс обучения происходит медленнее, чем при обучении с нуля, что в долгосрочной перспективе приводит к неоптимальным результатам. Мы предлагаем метод Drop-Upcycling, который эффективно решает эту проблему. Drop-Upcycling объединяет два, казалось бы, противоречивых подхода: использование знаний предварительно обученных плотных моделей при статистической повторной инициализации некоторых частей весов. Этот подход стратегически способствует специализации экспертов, значительно повышая эффективность модели MoE в усвоении знаний. Масштабные эксперименты показывают, что Drop-Upcycling значительно превосходит предыдущие методы построения MoE в долгосрочной перспективе, особенно при обучении на сотнях миллиардов токенов и более. В результате наша модель MoE с 5,9 миллиардами активных параметров достигает сопоставимой производительности с плотной моделью на 13 миллиардов параметров из того же семейства моделей, при этом требуя примерно 1/4 вычислительных затрат (FLOPs) на обучение. Все экспериментальные ресурсы, включая исходный код, обучающие данные, контрольные точки модели и логи, находятся в открытом доступе для обеспечения воспроизводимости и дальнейших исследований в области 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.