Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

  • 2021-01-11 16:11:52
  • William Fedus, Barret Zoph, Noam Shazeer
  • 221

Abstract

In deep learning, models typically reuse the same parameters for all inputs.Mixture of Experts (MoE) defies this and instead selects different parametersfor each incoming example. The result is a sparsely-activated model -- withoutrageous numbers of parameters -- but a constant computational cost. However,despite several notable successes of MoE, widespread adoption has been hinderedby complexity, communication costs and training instability -- we address thesewith the Switch Transformer. We simplify the MoE routing algorithm and designintuitive improved models with reduced communication and computational costs.Our proposed training techniques help wrangle the instabilities and we showlarge 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 obtainup to 7x increases in pre-training speed with the same computational resources.These improvements extend into multilingual settings where we measure gainsover the mT5-Base version across all 101 languages. Finally, we advance thecurrent scale of language models by pre-training up to trillion parametermodels on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over theT5-XXL model.

 

Quick Read (beta)

loading the full paper ...