Source-linked AI summary
Sparse is Enough in Scaling Transformers
Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, Łukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, Jonni Kanerva
TL;DR
Large Transformers are costly to train and fine-tune and slow to decode. The paper proposes sparse Transformer layers and integrates them with sparse attention for long sequences; sparse models match dense-model performance while inference speed benefits grow with model size.
Problem
Large Transformers achieve strong results but are expensive to train or fine-tune and slow to decode.
Source: §1 · p. 1Method
Scaling Transformers sparsify QKV and feedforward layers, with Terraformer combining the approach with reversible layers and sparse attention for long sequences.
Source: §1 · pp. 1–2, §3.1 · p. 3, §3.2 · pp. 5–6Results
Takeaways & Limitations
Terraformer achieves results competitive with BigBird-Pegasus on arXiv summarization without using the Pegasus loss in pre-training.
Source: §1 · p. 2Reported practical speedups are for inference rather than training and for unbatched CPU inference, not batched GPU inference.
Source: §5 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Large Transformers deliver strong results but incur high costs, motivating sparse architectures that scale efficiently. Scaling Transformers sparsify QKV and feedforward layers, match dense-model perplexity, and speed decoding; Terraformer applies the approach to long-document summarization.
1 Introduction
- Large Transformer models incur substantial training, data, and serving costs, including reported BERT training costs of $4k-$12k and over 143B words served daily by Google Translate.The passage also cites GPT-3 emissions and BERT's 128B-token training run as examples of these costs.Source: §1 · p. 1
- Scaling Transformers combine sparse QKV layers with sparse feedforward blocks to create a fully sparse Transformer architecture.The paper describes separate sparsity mechanisms for the query, key, value, and output layers.Source: §1 · p. 1
- For the approximately 800M-parameter model, Table 1 reports up to 2x decoding speedup with FF+QKV sparsity over the dense baseline.The table caption also reports 20x speedup for a 17B-parameter model.Source: Table 1 · p. 2
- On large C4, fully sparse Scaling Transformers match the baseline Transformer; authors report over 2.6x speedup in wall-clock decoding time at 800M parameters and 20x at 17B parameters.Source: §1 · p. 2
- Terraformer combines Scaling Transformers with reversible layers and sparse attention, achieving results competitive with BigBird-Pegasus on arXiv summarization without Pegasus loss in pre-training.It is pre-trained on C4 and fine-tuned to summarize arXiv articles.Source: §1 · p. 2
2 Related Work
Prior efficiency work includes pruning, distillation, quantization, sparse attention, and feedforward sparsity, but the paper distinguishes its dynamic sparsity from methods that shrink or distill models. It presents Scaling Transformers as complementary to sparse attention and focused on reducing decoder scaling costs.
2 Related Work
- Model pruning, distillation, quantization, parameter sharing, efficient attention, and efficient feedforward are identified as classes of Transformer-efficiency techniques.Pruning reduces matrices by removing weights; the paper contrasts its approach with distillation, which speeds inference from pretrained large models.Source: §2 · pp. 2–3
- Scaling Transformers use dynamic sparsity by training full weight matrices and activating selected parts per token, unlike pruning approaches that remove weights or blocks.The paper says unstructured sparse matrices often fail to yield inference speedups on actual hardware, while structured pruning addresses that challenge.Source: §2 · p. 3
- Sparse attention methods improve efficiency, especially for long sequences, and the paper describes its approach as complementary to them.The paper also distinguishes its QKV approach from grouped convolutions by allowing embedding dimensions to be permuted across attention heads.Source: §2 · p. 3
- Tensor decomposition addresses quadratic scaling in dmodel, but the cited discussion says it is most effective for large input and output embeddings and can lower decoder-block performance.The paper frames the remaining scaling issue as quadratic complexity even after distillation, quantization, and pruning.Source: §2 · p. 3
- Mixture-of-experts methods reduce feedforward complexity by activating one of several experts per token, whereas this work trains a full matrix and activates selected parts during decoding.The related-work discussion says mixture-of-experts speedups are mostly measured in training and focus on feedforward blocks.Source: §2 · p. 3
3 Sparse is Enough
The section develops sparse equivalents for feedforward, attention, and loss layers to avoid dense components becoming decoding bottlenecks. Combining sparse feedforward and QKV layers speeds decoding while retaining comparable perplexity.
3.1 Sparse Feedforward Layer
- Dynamic sparsity trains a full feedforward weight matrix but selects only a token-specific fraction of its parameters for decoding.A controller selects one nonzero activation per block, allowing the corresponding W1 columns and W2 rows to be skipped.Source: §3.1 · pp. 3–4
- N=64 feedforward sparsity retains baseline-comparable neg-log-perplexity as decoding time falls; N=128 degrades perplexity slightly.MoE-style sparsity reaches 0.11s decoding time versus 0.09s for sparse FF, but its log-perplexity of 1.64 is worse than both sparse FF and the dense baseline.Source: §3.1 · pp. 4–5, Figure 3 · p. 5
3.2 Sparse QKV Layer
- Sparse QKV layers compose multiplicative and convolutional layers to replace Q, K, and V projections, with the output dense layer removed in the combined variant.The multiplicative layer lets each attention head access arbitrary embedding dimensions rather than a predetermined subset.Source: §3.2 · pp. 5–6
- Combining sparse QKV and feedforward layers yields a 3.05x decoding-block speedup with comparable perplexity; sparse QKV alone does not speed the model when dense feedforward dominates.The comparison uses a vanilla Transformer baseline; decoding speed is almost the same for a Reformer model as well.Source: §3.2 · pp. 6–7
- Across sparse QKV settings for S and F, log-perplexity stays similar to the dense baseline within variance, while the multi-layer variant can improve it.The ablation increases dff from 4096 to 6144 to preserve parameter count, and its decoding time remains similar to baseline because dense feedforward dominates.Source: §3.2 · p. 6, Figure 5 · p. 7
- On selected GLUE validation tasks, sparse QKV+FF models achieve accuracy similar to the baseline Transformer.The table also includes Terraformer, a model using sparse QKV+FF.Source: Table 3 · p. 7
3.3 Sparse loss layer.
- Replacing the final dense vocabulary layer with a multiplicative layer speeds decoding but may degrade perplexity.The section reports these results in the appendix.Source: §3.3 · p. 7
4 Sparsity for Long Sequences
The paper extends Scaling Transformers to long sequences with sparse attention, reversible layers, and recurrence. Terraformer combines these mechanisms and is evaluated for decoding speed, generalization, and scientific-paper summarization.
4 Sparsity for Long Sequences
4.1 Architecture for Long Sequences
- The decoder removes separate encoder-decoder attention, concatenates encoder representations with decoder tokens, and uses two attention mechanisms before the feedforward block.The redesigned architecture is reported as fast as the baseline Transformer while giving better results.Source: §4.1 · p. 7
4.2 Reversibility for Memory Efficiency
- Reversible layers are used in encoder and decoder blocks to support large-batch training and fine-tuning large models on single machines.For sparse reversible layers, storing discrete row-selection decisions for reversal avoids the worse results seen when recalculating them on the backward pass.Source: §4.2 · p. 8
4.3 Recurrence for Generalization
- Low-rank SRUs project dmodel to 32 dimensions and back, providing recurrence that the authors found sufficient for generalization.The SRU is used in the feedforward block; simple recurrent units were chosen because they parallelize well during training.Source: §4.3 · p. 8
- On tasks trained with at most 128 digits and evaluated at lengths of 256–300, Terraformer correctly handles a large portion, but not as perfectly as Neural GPU [16].The tasks are long addition and copying a decimal digit; the baseline Transformer does not generalize well.Source: §4.3 · p. 8, Table 4 · p. 9
4.4 Experiments
- Terraformer is competitive with strong baselines on arXiv summarization using neither Pegasus loss nor beam search.The experiments pretrain on C4 and fine-tune on arXiv papers, with each paper as input and its abstract as output.Source: §4.4 · p. 9
- 37x faster single-token decoding is reported for Terraformer with 17B parameters.The section also reports that Terraformer decodes at 0.086s versus 0.061s for a sparse Transformer with similar perplexity.Source: §4.4 · p. 9
5 Conclusion
Sparse models match dense counterparts in perplexity at the same parameter count and are many times faster at inference in the reported large-model C4 experiments. The authors describe the work as a first step, noting that practical speedups were demonstrated for unbatched CPU inference rather than training or batched GPU inference.
5 Conclusion
- Sparse models match dense counterparts in performance on C4 while being many times faster at inference, with benefits growing as models scale.The authors characterize this as evidence that sparse layers are enough, and say the results promise more sustainable and useful Transformers.Source: §5 · p. 9
- The demonstrated practical speedups apply to inference, not training, and to unbatched CPU inference rather than batched GPU inference.The authors view the results as a first step and anticipate further work to address these settings.Source: §5 · p. 10
- Training speed did not improve; the authors attribute this mainly to using Gumbel-Softmax when training the feedforward block.They identify Fedus et al. [8] as a promising alternative for future exploration.Source: §5 · p. 10
- The experiments used dense-Transformer layer sizes and hyperparameters, which the authors say are probably not optimal for Scaling Transformers.They suggest proper tuning and further improvements could help match GPT-3 accuracy while enabling laptop inference, framing this as a challenge rather than a demonstrated result.Source: §5 · p. 10
6 Appendix
The appendix establishes a permutation-representing multiplicative layer for Sparse QKV, examines sparse loss, and reports Terraformer’s architecture, pretraining comparison, and summarization examples.
6.1 Sparse QKV
- A multiplicative layer can represent any bijection from model dimensions to module and position pairs, enabling Sparse QKV to represent any permutation.The proof constructs binary weights from the bijection’s module and position functions.Source: §6.1 · p. 14
6.2 Sparse Loss
- Increasing sparsity in the loss layer slightly worsens perplexity while speeding decoding.Table 7 reports decoding times as the number of modules varies.Source: §6.2 · p. 14, Table 7 · p. 15
6.3 Sparsity Results Data
- The appendix points readers to a public TensorBoard experiment containing the Figure 1 results.The passage supplies the experiment URL.Source: §6.3 · p. 15
6.5 Pretrained Terraformer on C4 dataset
- Terraformer (800M) has log-perplexity similar to the dense baseline and Scaling Transformers, with single-token decoding 1.92x faster than baseline.It used a 4x larger batch size and caught up at around one third of the training steps, though it trained for more steps overall.Source: Table 8 · p. 15, §6.5 · p. 15
6.6 Finetuning Terraformer on summarization task
- Table 9 presents Terraformer-generated scientific-paper abstracts alongside examples whose Pegasus abstracts and scores appear in Tables I.25–27 of.The abstracts were decoded greedily with temperature T = 0.5.Source: §6.6 · p. 15, Table 9 · p. 22
6.4 Architecture for Terraformer
- Terraformer’s decoder blocks place two attention mechanisms before the feedforward block.Figure 8 depicts the architecture discussed in Section 4.1.Source: §6.4 · p. 15, Figure 8 · p. 16