[논문 리뷰] EvoMoE: An Evolutional Mixture-of-Experts Training Framework via Dense-To-Sparse Gate

Xiaonan Nie, Xupeng Miao외 8명|arXiv (Cornell University)|2021. 12. 29.

Expert finding and Q&A systems인용 수 7

한 줄 요약

EvoMoE는 전문가 다각화 및 게이트 희박화 단계를 거쳐 단일 전문가를 점진적으로 희박한 전문가의 혼합(MoE) 모델로 진화시키는 이중 단계 훈련 프레임워크를 제안한다. 이는 처음에는 밀집되어 있던 후에 점차 희박해지는 DTS-Gate(밀집-희박 게이트)를 도입하여 훈련의 안정성과 성능을 향상시키며, GLUE 벤치마크에서 기준 모델 대비 최대 0.562 및 0.403의 성능 향상을 달성하여 최신 기술 수준을 확립한다.

ABSTRACT

Mixture-of-experts (MoE) is becoming popular due to its success in improving the model quality, especially in Transformers. By routing tokens with a sparse gate to a few experts (i.e., a small pieces of the full model), MoE can easily increase the model parameters to a very large scale while keeping the computation cost in a constant level. Most existing works just initialize some random experts, set a fixed gating strategy (e.g., Top-k), and train the model from scratch in an ad-hoc way. We identify that these MoE models are suffering from the immature experts and unstable sparse gate, which are harmful to the convergence performance. In this paper, we propose an efficient end-to-end MoE training framework called EvoMoE. EvoMoE starts from training one single expert and gradually evolves into a large and sparse MoE structure. EvoMoE mainly contains two phases: the expert-diversify phase to train the base expert for a while and spawn multiple diverse experts from it, and the gate-sparsify phase to learn an adaptive sparse gate and activate a dynamic number of experts. EvoMoE naturally decouples the joint learning of both the experts and the sparse gate and focuses on learning the basic knowledge with a single expert at the early training stage. Then it diversifies the experts and continues to train the MoE with a novel Dense-to-Sparse gate (DTS-Gate). Specifically, instead of using a permanent sparse gate, DTS-Gate begins as a dense gate that routes tokens to all experts, then gradually and adaptively becomes sparser while routes to fewer experts. Evaluations are conducted on three popular models and tasks, including RoBERTa for masked language modeling task, GPT for language modeling task and Transformer for machine translation task. The results show that EvoMoE outperforms existing baselines, including Switch, BASE Layer, Hash Layer and StableMoE.

연구 동기 및 목표

• 무작위로 초기화된 전문가와 희박한 게이트로 인해 기존 MoE 훈련에서 발생하는 불안정성과 수렴 불량 문제를 해결하기 위해.
• 전문가와 게이트 학습을 분리하여 전문가 및 게이트 학습을 별도로 수행함으로써 훈련 효율성과 모델 품질을 향상시키기 위해.
• 내재적 희박화를 통해 밀집 라우팅에서 희박 라우팅으로의 부드러운 전환을 가능하게 하기 위해.
• 기존 MoE 방법들과 비교해 더 높은 성능과 FLOPs 효율성을 달성하기 위해.
• 대규모 MoE 모델을 위한 더 안정적이고 확장 가능한 훈련 패러다임을 탐색하기 위해.

제안 방법

• 전문가 다각화 단계에서는 먼저 단일 기본 전문가를 훈련한 후, 무작위 마스킹을 사용하여 구조적 다양성을 도입함으로써 다수의 전문가로 다각화한다.
• 게이트 희박화 단계에서는 밀집-희박 게이트(DTS-Gate)를 도입하며, 처음에는 모든 전문가를 대상으로 소프트하고 밀집된 라우팅을 수행하고, Gumbel-Softmax의 온도 감쇠를 통해 점차적으로 희박해지도록 한다.
• DTS-Gate는 콘텐츠 기반 라우팅 메커니즘을 사용하여 게이트 가중치가 학습된 임계값을 초과하는 전문가만 활성화함으로써 적응적이고 동적인 전문가 선택을 가능하게 한다.
• 이 프레임워크는 전문가와 게이트 학습을 분리한다: 희박한 라우팅 도입 이전에 전문가들이 공유 방식으로 사전 훈련된다.
• Gumbel-Softmax의 온도 감쇠 스케줄을 적용하여 게이트 출력의 점진적 희박화를 유도한다.
• 이 방법은 RoBERTa, GPT, Transformer를 포함한 다양한 모델에 걸쳐 다수의 자연어 처리 작업에서 종단 간(end-to-end)으로 적용된다.

실험 결과

연구 질문

• RQ1 단일 전문가에서 희박한 MoE 아키텍처로의 점진적 진화가 훈련 안정성과 모델 성능 향상에 기여하는가?
• RQ2 밀집-희박 게이트 메커니즘이 고정된 희박한 게이트(예: Top-k)보다 더 나은 수렴성과 라우팅 정확도를 제공하는가?
• RQ3 공유된 기본 전문가에서 유래한 전문가 다각화가 지식 일반화를 향상시키고 훈련 불안정성을 감소시키는가?
• RQ4 Switch, BASELayer, StableMoE와 같은 기존 MoE 방법들과 비교해 EvoMoE는 성능과 FLOPs 효율성에서 어떤가?
• RQ5 전문가와 게이트 학습을 분리함으로써 MoE 모델의 학습 역학은 얼마나 향상되는가?

주요 결과

• EvoMoE는 GLUE 벤치마크에서 기준 모델 대비 평균 점수 기준 최대 0.562 및 0.403 향상되며, 뛰어난 모델 품질을 입증한다.
• DTS-Gate 메커니즘은 더 부드러운 라우팅 적응을 가능하게 하며, 고정된 Top-k 게이트에 비해 무작위 초기화에 대한 민감도를 감소시킨다.
• 전문가 다각화 단계는 더 안정적이고 효과적인 전문가 특화를 이끌어내어 전체 모델의 강건성을 향상시킨다.
• EvoMoE는 밀집 MoE보다 더 뛰어난 FLOPs 효율성을 달성했으며, 희박 MoE 기준 모델들과 경쟁 가능한 성능을 보였다.
• 마스크 언어 모델링, 언어 모델링, 기계 번역을 포함한 다양한 작업에서 일관된 성능 향상이 관찰되었다.
• 제거 실험을 통해 전문가 다각화 및 게이트 희박화 단계가 최적 성능을 내기 위해 필수적임을 확인하였다.