드롭-업사이클링: 부분 재초기화를 통한 희소 전문가 혼합 모델 학습

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) 아키텍처는 동일한 용량의 밀집 모델(dense model)에 비해 학습 및 추론 비용을 크게 절감합니다. 업사이클링(Upcycling)은 사전 학습된 밀집 모델을 사용하여 MoE 모델을 초기화하고 학습시키는 접근 방식입니다. 업사이클링은 초기 성능 향상을 가져오지만, 처음부터 학습을 시작할 때보다 학습 진행 속도가 느려 장기적으로는 최적의 성능을 달성하지 못하는 문제가 있습니다. 우리는 이 문제를 효과적으로 해결하는 Drop-Upcycling 방법을 제안합니다. Drop-Upcycling은 사전 학습된 밀집 모델의 지식을 활용하면서도 일부 가중치를 통계적으로 재초기화하는 두 가지 상반된 접근 방식을 결합합니다. 이 방법은 전문가(expert)의 특화를 전략적으로 촉진하여 MoE 모델의 지식 습득 효율을 크게 향상시킵니다. 대규모 실험을 통해 Drop-Upcycling이 장기적으로, 특히 수천억 개 이상의 토큰을 학습할 때 이전의 MoE 구축 방법을 크게 능가함을 입증했습니다. 결과적으로, 5.9B 활성 파라미터를 가진 우리의 MoE 모델은 동일 모델 계열의 13B 밀집 모델과 비슷한 성능을 달성하면서도 학습 FLOPs를 약 1/4로 줄였습니다. 재현성과 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.