[논문 리뷰] 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참고 문헌 51인용 수 361

한 줄 요약

이 논문은 Switch Transformer를 도입하여 토큰당 FLOPs를 일정하게 유지하면서 매개변수를 대대적으로 증가시키는 희소 활성 Mixture-of-Experts 모델을 제시한다. 이는 최대 7배의 사전 학습 속도 향상을 가능하게 하고, 트릴리언 파라미터 모델의 확장을 개선된 스케일링, 미세 조정, 다국어 결과, 밀집 소형 모델로의 증류를 가능하게 한다.

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.

연구 동기 및 목표

• 희소 활성 모델(MoE)을 통해 컴퓨팅과 무관하게 확장 가능한 매개변수 수를 동기화하고 학습 효율성 및 성능을 개선하려는 동기 부여.
• Mixture-of-Experts의 라우팅 및 통신을 단순화하여 TPU에서 안정적이고 확장 가능한 학습을 가능하게 함.
• 토큰당 FLOPs를 고정한 채 전문가(매개변수) 수를 증가시키면 사전 학습 속도가 빨라지고 샘플 효율성이 향상된다는 것을 보여줌.
• 미세 조정, 다국어 평가, 밀집 소형 모델로의 증류를 통해 다운스트림 이득을 입증함.
• 대규모 희소 모델의 안정화를 위한 실용적 학습 기법(정밀도 처리, 초기화, 정규화)을 제시함.

제안 방법

• 각 토큰을 단일 전문가로 라우팅하는 Switch 라우팅 메커니즘 도입(k=1)으로 라우팅 비용 및 통신을 감소시킴.
• auxiliary 로드 밸런싱 손실이 있는 differentiable gate p_i(x)를 사용하여 전문가 간 토큰 분포를 고르게 만듦.
• 고정된 전문 능력치와 용량 팩터를 사용한 희소 라우팅으로 토큰 디스패치와 오버플로우를 관리함.
• 탈선 없이 평가를 위한 벤치마크 perplexity를 위해 15% 토큰 드롭아웃과 함께 Colossal Clean Crawled Corpus(C4)에서 마스킹 언어 모델링으로 사전 학습.
• 저정밀도 학습의 안정화를 위해 선택적 정밀도 훈련(라우팅 계산은 float32, 나머지는 bfloat16 사용)
• 다양한 NLP 작업에서 미세 조정, 거대 희소 모델의 소규모 밀집 모델로의 증류, 101개 언어에 걸친 다국어 성능 평가를 수행.

실험 결과

연구 질문

• RQ1 제한된 FLOPs 내에서 Switch와 같은 단일 전문가 라우팅의 희소 활성 변환기가 동일한 컴퓨트 예산에서 밀집 모델 및 MoE 모델보다 품질이 우수하거나 동등한지를 확인할 수 있는가?
• RQ2 라우팅 단순화, 초기화 전략, 정밀도 기법이 대형 Switch Transformer 모델의 안정성 및 확장성에 어떤 영향을 미치는가?
• RQ3 Switch Transformer가 밀집 기반과 비교해 사전 학습 속도, 미세 조정 성능 및 다국어 설정에 일관된 향상을 제공하는가?
• RQ4 대형 희소 모델을 의미 있는 성능 향상을 유지하는 소형 밀집 모델로 증류할 수 있는가?
• RQ5 토큰당 고정된 FLOPs 하에서 전문가 수를 늘릴 때 Switch Transformer의 확장 특성은 무엇인가?

주요 결과

Model	Capacity	Quality after 100k steps	Time to Quality (hours)	Speed (examples/sec)
T5-Base	—	-1.731	Not achieved †	1600
T5-Large	—	-1.550	131.1	470
MoE-Base	2.0	-1.547	68.7	840
Switch-Base	2.0	-1.554	72.8	860
MoE-Base	1.25	-1.559	80.7	790
Switch-Base	1.25	-1.553	65.0	910
MoE-Base	1.0	-1.572	80.1	860
Switch-Base	1.0	-1.561	62.8	1000
Switch-Base+	1.0	-1.534	67.6	780

• Switch Transformer는 동일한 컴퓨트 예산 대비 밀집 및 MoE 기반의 기준보다 더 빠른 사전 학습을 달성했으며, 일부 설정에서 최대 7배의 속도 향상을 보였다.
• 전문가 수를 늘리되 토큰당 FLOPs를 일정하게 유지하면 페르플렉시티 및 샘플 효율성이 지속적으로 개선되어 더 크고 강력한 모델을 가능하게 한다.
• Switch-Base는 64개의 전문가에서 같은 컴퓨트 및 하드웨어 하에서 약 1/7의 시간으로 T5-Base와 유사한 품질에 도달하여 월드-클록 효율성을 입증했다.
• Switch 트랜스포머는 더 큰 밀집 모델과 비교해도 뛰어난 성능을 보이며, 예를 들어 FLOP-매칭 조건에서 Switch-Base가 T5-Large를 능가하고, Switch-Large가 여러 지표에서 T5-Large를 능가하는 경우가 많다.
• 선택적 정밀도(라우터 내부에서만 부동소수점 형식을 float32로 cast)가 Full bf16에 가까운 속도로 대규모 훈련의 안정성을 제공한다.
• 미세 조정 및 증류 실험은 GLUE, SuperGLUE, QA, 요약, 다국어 작업에서 Switch 변형의 다운스트림 이득을 크게 보여주며, 증류는 파라미터 수 1/20 학생에서도 이득의 약 30%를 보존한다.