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

A model that does not use all of itself

Most language models are dense: every token of every input passes through all of the model's parameters, so making a model larger also makes each forward pass more expensive. A Mixture of Experts (MoE) model takes a different route. It selects different parameters for each incoming example, which produces a sparsely activated model: the parameter count can be very large, but the computational cost of processing an example stays constant^[1]. Only a small slice of the network is engaged for any single input, and which slice is engaged is decided per example rather than fixed in advance.

The idea is attractive, because it separates two things that dense models tie together: how much the model knows, measured roughly by how many parameters it holds, and how much work it does per example, measured in floating-point operations. But MoE models were not widely adopted. Three problems stood in the way: the routing algorithm was complicated, communication costs were high, and training was unstable^[1].

What the paper does

William Fedus, Barret Zoph and Noam Shazeer submitted Switch Transformers to arXiv on 11 January 2021^[1]. The paper keeps the sparse-activation idea and attacks the obstacles directly, simplifying the MoE routing algorithm so that the complexity, the communication costs and the training instability that had held the approach back become manageable^[1].

A second contribution is on the numerical side. The training techniques proposed in the paper allow large sparse models to be trained, for the first time, with lower precision formats — specifically bfloat16^[1]. bfloat16 is a 16-bit floating-point format that keeps the exponent range of the 32-bit format usually used for training while storing far fewer bits of the fraction. Because it halves the memory footprint of activations and the volume of data moved between devices, it matters especially for sparse models, whose cost is dominated by moving parameters and activations around rather than by arithmetic alone.

Results

The gains are reported against the T5 family, a line of text-to-text models used here as the dense baseline. Models based on T5-Base and T5-Large obtain up to 7x increases in pre-training speed with the same computational resources^[1]. The advantage is not confined to English: the gains extend into multilingual settings, across all 101 languages measured against mT5-Base^[1].

The authors also push the scale of language models up, pre-training models of up to a trillion parameters on the Colossal Clean Crawled Corpus (C4), a large public collection of web text, and achieving a 4x speedup over the T5-XXL model^[1]. The largest model in the paper, Switch-C, has 1,571 billion parameters while using 890 GFLOPs per sequence^[1]. A GFLOP is one billion floating-point operations, so that figure is the arithmetic bill for pushing a single sequence through the model — a bill that is fixed by how much of the network is switched on, not by how many parameters the network stores.

Why it matters

Put together, the numbers describe a trade: parameter count grows by orders of magnitude, while the cost per sequence stays in the same range, because only a fraction of the parameters is used for any one example^[1]. That is what makes a trillion-parameter model trainable at all, and it is why the paper's simplification of routing is the central move rather than an implementation detail. A simpler routing decision is cheaper to compute and easier to distribute across many devices, which in turn reduces the communication that had made earlier MoE systems awkward to train^[1].

Switch-C illustrates the endpoint of the exercise: a model with 1,571 billion parameters that still spends 890 GFLOPs per sequence^[1]. Whether such models are better in every respect is not something the paper claims; what it shows is a workable recipe for growing parameters without growing per-example compute, and a set of training techniques that make large sparse models trainable in bfloat16^[1]. Those two pieces — simplified routing and low-precision training — are the parts of the work that later sparse and mixture-of-experts language models build on.