[논문 리뷰] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

William Fedus, Barret Zoph외 1명|arXiv (Cornell University)|2021. 01. 11.

Topic Modeling인용 수 700

한 줄 요약

이 논문은 스위치 트랜스포머를 소개합니다. 이는 1-전문가 라우팅을 갖춘 희소 활성화 Mixture-of-Experts 모델로, 수조 파라미터에 이르는 대규모 파라미터 수를 달성하고, 학습 안정성을 개선하며, 고정 FLOPs 하에서 더 빠른 프리트레이닝이 가능하고, 품질을 크게 유지하는 상태로 컴팩트한 밀집 모델로 증류될 수 있습니다.

ABSTRACT

In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers of parameters -- but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability -- we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over the T5-XXL model.

연구 동기 및 목표

• 트랜스포머 모델의 확장을 파라미터 수를 늘리되 토큰당 계산량은 일정하게 유지하여 동기를 부여한다.
• TPU/GPU 하드웨어에서 확장 가능한 희소 모델을 가능하게 하기 위해 Mixture-of-Experts 라우팅을 단순화하고 안정화한다.
• 혼합 정밀도와 새로운 초기화 방식으로 학습 안정성을 입증한다.
• 프리트레이닝, 파인튜닝, 다국어 설정 전반에서 실질적 이점을 보인다.

제안 방법

• Mixture-of-Experts 계층에서 각 토큰을 단일 전문가로 라우팅하는 Switch 라우팅 scheme 제안(k=1).
• 토큰을 전문가들 간에 분배하기 위해 부하 균형 보조 손실이 있는 미분 가능한 라우터를 사용.
• 희소 FFN을 토큰을 독립적으로 처리하는 Switch FFN으로 운용하며 용량 계수와 오버플로우 처리 포함.
• 학습 안정화를 위해 선택적 정밀도 학습(라우팅 계산은 float32, 다른 곳은 bfloat16)을 적용.
• 더 큰 전문가 수와 안정적인 파인튜닝 가능성을 위해 초기화 스케일링과 전문가 정규화를 도입.
• 밀집 및 MoE 베이스라인 대비 FLOP 매칭 비교를 제공하고 프리트레이닝, 파인튜닝, 다국어 과제에서의 결과를 보고한다.

실험 결과

연구 질문

• RQ1 전통적인 MoE에 비해 라우팅 비용을 줄이면서 단일 전문 가이드(Switch)로도 모델 품질을 유지하거나 향상시킬 수 있는가?
• RQ2 토큰당 FLOPs를 고정한 채 전문가 수를 늘리면 학습 속도와 샘플 효율성에 어떤 영향을 주는가?
• RQ3 대형 희소 모델의 안정화를 위해 필요한 학습 기법(정밀도, 초기화, 규제)은 무엇인가?
• RQ4 Switch Transformers가 프리트레이닝, 파인튜닝, 다국어 설정 전반에서 실질적 이점을 제공하는가?
• RQ5 대형 희소 모델을 품질 손실 없이 더 작은 밀집 모델로 증류할 수 있는가?

주요 결과

• Switch Transformers는 동일한 컴퓨트 예산에서 튜닝된 T5 베이스라인 대비 7배 이상 프리트레이닝 속도향상을 달성한다.
• 64개의 전문가를 가진 Switch-Base는 T5-Base보다 더 빨리 학습되며 비슷하거나 더 나은 품질을 달성, 속도-품질 이점이 강력함을 보여준다.
• wall-clock 기준으로 Switch Transformers는 동일 FLOPs의 밀집 베이스라인을 능가하며 상당한 시간 절약을 보인다(예: 64-전문가 Switch-Base가 T5-Base의 약 1/7 시간만에 유사 품질 달성).
• Switch-Large는 T5-Large와 FLOP를 맞추면 더 큰 밀집 베이스라인보다 우수한 확장 및 파인튜닝 성능을 보인다.
• 101개 언어에서 보편적인 다국어 이점이 나타나며, 언어의 91%가 mT5 대비 4배 이상 속도 향상을 얻었다.
• 대형 희소 모델은 컴팩트한 밀집 모델로 증류될 수 있으며, 희소 모델의 개선 중 약 30%를 보존하고 파라미터는 약 1/20 수준을 사용한다.
• 선택적 정밀도 학습(로컬 라우터 계산을 float32로)으로 학습을 안정화하면서 거의 bf16 속도를 유지한다.
• 초기화 및 정규화 전략은 수조 파라미터 규모의 모델의 안정적인 학습을 가능하게 한다.