anthology-of-the-sota

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

2 min read

Properties5
StatusActive
Tagsmodel-architecture
ExtendsLIT-187 — GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
Compared againstLIT-709 — Primer: Searching for Efficient Transformers for Language Modeling, LIT-711 — Do Transformer Modifications Transfer Across Implementations and Applications?
Cited as “source” bySOTA-150 — Make the feed-forward layers a sparse mixture of experts once the model is large enough to be compute-bound

Graph View

Fedus et al. (2021) — ARXIV-2101.03961

Key takeaways

  • Top-1 routing. LIT-187 argued top-2 was the minimum for a useful gradient to the router; this shows one expert per token works, and is cheaper in routing computation, in the expert’s batch size, and in communication.
  • Stability was the blocker, and the fix is small. Sparse models diverged in low precision; selectively casting only the router to float32 while the rest stays bfloat16 makes them trainable, without the communication cost of full float32.
  • 7× pretraining speedup over T5-Base at the same compute budget, and a 1.6T-parameter model — the first trillion-scale sparse model, five years before the record’s frontier entries reach that scale.
  • Distillation back to dense preserves about 30% of the sparse model’s gain, which is the honest version of “you can have this without serving an MoE”.

Standing in the anthology

The third and last of the originating MoE papers this record was missing, and the one that establishes the scale claim: a trillion parameters in 2021.

Its top-1 routing is the other pole of the granularity argument LIT-170 makes. GShard says two experts, Switch says one, DeepSeekMoE says many small ones and take more of them — three answers to how finely the router should choose, which is more interesting read together than any of them alone.

Two later papers measured Switch as a baseline. In Narang et al.’s audit of transformer modifications (LIT-711), a 1.1B-parameter Switch Transformer at roughly the vanilla model’s FLOPs reached the lowest final loss in the table, 1.758 against 1.838, with early loss 2.135 ± 0.007. That came at about 9% fewer steps per second, in the codebase Switch was built in. Primer (LIT-709) added squared ReLU and depthwise-convolution attention to Switch at 110M on C4. That gave a 1.5× speedup in compute to reach plain Switch’s final quality, so Primer’s changes combine with top-1 sparsity rather than replacing it.

The low-precision router note is worth remembering beside LIT-186: sparse models were the first place where precision stopped being an implementation detail and became an architecture constraint.

Backlinks