GRIN: 그래디언트 정보를 활용한 MoE

GRIN: GRadient-INformed MoE

September 18, 2024

저자: Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, Vishrav Chaudhary, Zeqi Lin, Chenruidong Zhang, Jilong Xue, Hany Awadalla, Jianfeng Gao, Weizhu Chen

cs.AI

초록

전문가 모델의 혼합(Mixture-of-Experts, MoE)은 전문가 라우팅을 통해 희소 계산으로 인해 밀집 모델보다 효과적으로 확장됩니다. 이는 전문가 모듈의 작은 하위 집합만 활성화되도록 선택적으로 활성화하기 때문입니다. 그러나 희소 계산은 기존의 훈련 방법에 도전을 제기합니다. 왜냐하면 이산적인 전문가 라우팅이 표준 역전파를 방해하고 따라서 그라디언트 기반 최적화를 방해하기 때문입니다. MoE의 확장 능력을 더 잘 추구하기 위해, 우리는 GRIN(GRadient-INformed MoE training)을 소개합니다. 이는 전문가 라우팅을 위한 희소 그라디언트 추정을 통합하고 토큰 삭제를 피하기 위해 모델 병렬화를 구성합니다. 자기회귀 언어 모델링에 GRIN을 적용하여 상위 2개의 16배3.8B MoE 모델을 개발했습니다. 우리의 모델은 활성화된 매개변수가 6.6B개뿐이며, 7B 밀집 모델을 능가하며, 동일한 데이터로 훈련된 14B 밀집 모델의 성능과 일치합니다. 다양한 작업을 통해 수행된 포괄적인 평가는 GRIN이 MoE의 효과성을 크게 향상시킬 수 있는 잠재력을 보여주며, MMLU에서 79.4, HellaSwag에서 83.7, HumanEval에서 74.4, MATH에서 58.9의 성능을 달성했습니다.

English

Mixture-of-Experts (MoE) models scale more effectively than dense models due to sparse computation through expert routing, selectively activating only a small subset of expert modules. However, sparse computation challenges traditional training practices, as discrete expert routing hinders standard backpropagation and thus gradient-based optimization, which are the cornerstone of deep learning. To better pursue the scaling power of MoE, we introduce GRIN (GRadient-INformed MoE training), which incorporates sparse gradient estimation for expert routing and configures model parallelism to avoid token dropping. Applying GRIN to autoregressive language modeling, we develop a top-2 16times3.8B MoE model. Our model, with only 6.6B activated parameters, outperforms a 7B dense model and matches the performance of a 14B dense model trained on the same data. Extensive evaluations across diverse tasks demonstrate the potential of GRIN to significantly enhance MoE efficacy, achieving 79.4 on MMLU, 83.7 on HellaSwag, 74.4 on HumanEval, and 58.9 on MATH.