Sparse is Enough in Scaling Transformers

By Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, Jonni Kanerva. 2021

125 citations recorded by the source corpus.

Large Transformer models yield impressive results on many tasks, but are expensive to train, or even fine-tune, and so slow at decoding that their use and study becomes out of reach. We address this problem by leveraging sparsity. We study sparse variants for all layers in the Transformer and propose Scaling Transformers, a family of next generation Transformer models that use sparse layers to scale efficiently and perform unbatched decoding much faster than the standard Transformer as we scale up the model size. Surprisingly, the sparse layers are enough to obtain the same perplexity as the standard Transformer with the same number of parameters. We also integrate with prior sparsity approaches to attention and enable fast inference on long sequences even with limited memory. This results in performance competitive to the state-of-the-art on long text summarization.

Subfields: Artificial Intelligence, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Yi Tay
2. Mostafa Dehghani
3. Donald Metzler
4. Quoc V. Le
5. Mohit Bansal
6. Luke S. Zettlemoyer
7. Jianfeng Gao
8. Dara Bahri

In Model Architectures

1. Yi Tay
2. Mostafa Dehghani
3. Donald Metzler
4. Dara Bahri
5. William Fedus
6. Joseph Dean
7. Barret Zoph
8. Quoc V. Le

In Efficiency & Compression

1. Yi Tay
2. Donald Metzler
3. Mostafa Dehghani
4. Dara Bahri
5. Barret Zoph
6. Joseph Dean
7. William Fedus
8. Georgios Evangelidis

Citations in the corpus

• Efficiently Scaling Transformer Inference (2022, positive): “, 2017), and adaptive computation techniques that allocate different amounts of compute per input and generation timestep (Jaszczur et al., 2021; Schuster et al., 2022), promise to reduce FLOPs per token of Transformer models.”
• A Review of Sparse Expert Models in Deep Learning (2022, methodology): “Jaszczur et al. (2021) sparsify all the layers (e.g. dense and self-attention) of a Trans-former model to achieve 37 × inference speedups for a special-case of single-example inference (unbatched).”
• Fast Inference from Transformers via Speculative Decoding (2022, background): “Numerous techniques try to make inference more efficient for all tokens, e.g. distillation (Hinton et al., 2015), sparcification (Jaszczur et al., 2021), quantization (Hubara et al., 2016), and architecture modification (So et al., 2021; Shazeer, 2019).”
• The case for 4-bit precision: k-bit Inference Scaling Laws (2022, neutral): “For example, without optimizations for the attention operations, multi-head attention can be a large chunk of the inference latency footprint (Jaszczur et al., 2021; Pope et al., 2022).”
• Efficient Transformers: A Survey (2020, background): “, 2020) O(N(2)) X Sparse Scaling Transformers (Jaszczur et al., 2021) O(N(2)) X Sparse GLaM (Du et al.”
• RULER: What's the Real Context Size of Your Long-Context Language Models? (2024, methodology): “Various sparse attention mechanisms (Child et al., 2019; Jaszczur et al., 2021) such as shifted sparse attention (Chen et al., 2024), dilated attention (Ding et al., 2023), and attention sinks (Han et al., 2023; Xiao et al., 2024b) were employed to enable efficient context scaling.”
• EfficientFormer: Vision Transformers at MobileNet Speed (2022, background): “Following this success, several efforts explore the potential of ViT by improving training strategies [3, 4, 5], introducing architecture changes [6, 7], redesigning attention mechanisms [8, 9], and elevating the performance of various vision tasks such as classification [10, 11, 12], segmentation [”