[논문 리뷰] Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers

Rui Liu, Young Jin Kim외 2명|arXiv (Cornell University)|2022. 05. 28.

Topic Modeling인용 수 11

한 줄 요약

이 논문은 희소적으로 활성화되는 트랜스포머, 특히 믹스처 오브 익스퍼트(MoE) 모델을 위한 통신 효율적인 정규화 기법인 게이팅 드롭아웃(Gating Dropout)을 제안한다. 학습 중에 토큰의 라우팅을 임의로 생략함으로써 원격 전문가에게 보낼 수 있도록 하여 고비용의 크로스머신 전역 통신을 줄이고, 일반화 성능을 향상시키는 정규화 효과를 도입함으로써 수렴 속도를 빠르게 하고 다국어 번역 작업에서 최대 0.6 BLEU 점수 향상을 이룬다.

ABSTRACT

Sparsely activated transformers, such as Mixture of Experts (MoE), have received great interest due to their outrageous scaling capability which enables dramatical increases in model size without significant increases in computational cost. To achieve this, MoE models replace the feedforward sub-layer with Mixture-of-Experts sub-layer in transformers and use a gating network to route each token to its assigned experts. Since the common practice for efficient training of such models requires distributing experts and tokens across different machines, this routing strategy often incurs huge cross-machine communication cost because tokens and their assigned experts likely reside in different machines. In this paper, we propose \emph{Gating Dropout}, which allows tokens to ignore the gating network and stay at their local machines, thus reducing the cross-machine communication. Similar to traditional dropout, we also show that Gating Dropout has a regularization effect during training, resulting in improved generalization performance. We validate the effectiveness of Gating Dropout on multilingual machine translation tasks. Our results demonstrate that Gating Dropout improves a state-of-the-art MoE model with faster wall-clock time convergence rates and better BLEU scores for a variety of model sizes and datasets.

연구 동기 및 목표

• 모수의 라우팅을 원격 전문가에게 전달하는 데 자주 발생하는 전역 통신 작업으로 인해 발생하는 분산 MoE 모델 학습에서의 높은 크로스머신 통신 비용을 해결하기 위해.
• 특히 대규모 분산 학습 환경에서 모델 성능을 저하시키지 않으면서 통신 오버헤드를 줄이기 위해.
• 학습 중에 지역 전문가 사용을 강제함으로써 일반화 성능을 향상시키는 정규화 메커니즘을 도입하기 위해.
• MoE 기반의 다국어 기계 번역 모델에 대한 학습 효율성과 수렴 속도를 향상시키기 위해.

제안 방법

• 게이팅 드롭아웃은 학습 중에 라우팅 메커니즘을 무작위로 비활성화하여 토큰이 할당된 원격 전문가에게 보내지 않고 현지 머신에 그대로 유지되도록 한다.
• 이 방법은 게이팅 네트워크의 라우팅 결정에 드롭아웃 비율 p를 적용하며, 확률 p로 토큰이 라우팅되지 않고 현지에서 처리된다.
• 전역 통신 작업의 무작위 생략은 상호 머신 간 데이터 재정렬을 크게 줄여 통신 비용을 낮춘다.
• 정규화 효과는 전문가가 최적의 라우팅에 접근할 수 없기 때문에 더 일반화 가능한 특징을 학습하도록 강제당함으로써 발생한다. 이는 일반화 성능 향상에 기여한다.
• 이 방법은 k=1 및 k>1 라우팅 전략 모두와 호환되며, 어떤 MoE 기반 트랜스포머 아키텍처에도 적용 가능하다.
• 이 방법은 학습 중에만 적용되며 추론 시에는 비활성화되어 모델 정확도를 유지한다.

실험 결과

연구 질문

• RQ1 무작위로 라우팅 작업을 생략함으로써 분산 MoE 학습에서 통신 비용을 줄일 수 있을까? 이로 인해 모델 성능이 떨어지지 않을까?
• RQ2 게이팅 드롭아웃을 통한 라우팅의 무작위화가 MoE 모델의 일반화 성능 향상에 기여하는가?
• RQ3 드롭아웃 비율이 MoE 학습에서 통신 효율성, 수렴 속도, 모델 정확도 사이의 트레이드오프에 어떤 영향을 미치는가?
• RQ4 게이팅 드롭아웃은 다국어 번역 벤치마크에서 학습 처리량과 수렴 속도 향상에 기여하는가?
• RQ5 게이팅 드롭아웃은 자원이 적은 언어 번역 작업에 특별한 이점을 제공하는가?

주요 결과

• 기본 MoE 모델 대비 목표 BLEU 점수에 도달하는 데 소요되는 학습 시간을 최대 58% 줄였다.
• 일반화 성능 향상으로 다국어 번역 작업에서 최대 0.6 BLEU 점수 향상을 달성했다.
• 특히 높은 드롭아웃 비율에서 전역 통신이 감소함에 따라 처리량이 최대 16% 향상되었다.
• 게이팅 전문가 드롭(Gate-Expert-Drop)의 최적 성능는 드롭아웃 비율 0.2에서, 게이팅 드롭(Gate-Drop)의 최적 성능는 0.3에서 달성되었으며, 이보다 높은 비율은 성능 저하를 초래했다.
• 정규화 효과는 자원이 적은 언어 번역 작업에서 특히 유익하여 최소 0.3 BLEU 포인트의 향상을 이뤘다.
• 이 방법은 다양한 모델 크기와 데이터셋에서 효과적이며, 수렴성과 효율성 향상 측면에서 일관된 개선 효과를 보였다.