In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) models defy this and instead select different parameters for each incoming example. The result is a sparsely-activated model – with an outrageous number 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 mitigate the instabilities, and we show large sparse models may be trained, for the first time, with lower precision formats. We design models based off T5-Base and T5-Large (Raffel et al., 2019) 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.

Barret Zoph is a senior research scientist on the Google Brain team. He has worked on a variety of deep learning research topics ranging from neural architecture search (NAS), data augmentation, semi-supervised learning for computer vision and model sparsity. Prior to Google Brain he worked at the Information Sciences Institute working on machine translation.

Liam Fedus is a senior research scientist at Google Brain and a PhD candidate at the University of Montreal. His prior research spanned generative models, reinforcement learning, and his current focus is on understanding and designing effective sparse expert models.