Elamur Research⌕
☰
home / papers / 2020 / SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Paper · 2020 · Hanrui Wang, Zhekai Zhang and Song Han

Summary

Attention-based NLP models such as BERT and GPT-2 run slowly on CPUs and GPUs. The cost grows quadratically with input length, much of the time goes to moving data rather than arithmetic, and existing accelerators were built for CNNs and RNNs. SpAtten (HPCA 2021) is an algorithm-hardware co-design that cuts both computation and memory traffic in attention. It uses three input-dependent techniques:

  • Cascade token pruning drops tokens that receive little accumulated attention.
  • Cascade head pruning drops heads whose outputs have little accumulated magnitude. With both pruning methods, anything removed stays removed in all later layers.
  • Progressive quantization first loads only the high-order bits and fetches the low-order bits only when the attention distribution looks flat.

The custom hardware includes a high-throughput top-k engine, a specialized memory hierarchy and a fully pipelined datapath. Across 30 benchmarks, SpAtten reduces DRAM access by 10× on average without losing accuracy. It is also faster and more energy-efficient than the prior attention accelerators A3 and MNNFast, a TITAN Xp GPU and a Xeon CPU. Unlike those earlier accelerators, it also speeds up memory-bound generative models such as GPT-2, not only discriminative ones like BERT.

Key points

  • Average speedups: 1.6× over A3, 3.0× over MNNFast, 162× over TITAN Xp GPU and 347× over Xeon CPU.
  • Energy savings: 1.4×, 3.2×, 1193× and 4059× over the same four baselines.
  • DRAM reductions on eight GPT-2 models: token pruning gives 3.8× less DRAM access and head pruning gives 1.1×. Progressive quantization saves a further 5.1× in memory access.
  • Motivation:
    • Producing a 30-token sentence with GPT-2 takes 370 ms on a TITAN Xp.
    • Attention accounts for over 50% of end-to-end GPT-2 latency on the profiled platforms.
    • Data movement makes up about 73% of attention latency.
    • Generation is about 97% of GPT-2 latency when producing 32 tokens, and runs at roughly 0.5 ops/Byte, so it is memory-bound.
  • Example and code: with BERT-Base on SST-2, a sentence shrinks from 11 tokens and 12 heads in layer 1 to 2 tokens and 8 heads by layer 3, with the classification unchanged. Code is released at github.com/mit-han-lab/spatten.

Interesting ideas

  • I1. Use attention as a running importance score. Summing the attention probabilities each token receives across layers gives a cheap, training-free measure of importance. Pruning by this score avoids weight retraining and adapts to each input. The same signal could help with long-document compression or KV-cache trimming.
  • I2. Prune globally, not locally. Once a token is dropped, it disappears from all subsequent layers, including the feed-forward blocks. Savings therefore compound with depth instead of staying inside one attention layer.
  • I3. Scale pruning with input length. The paper treats longer sentences as more redundant and prunes a larger share of them. This is a simple heuristic worth trying in any per-input compute budget.
  • I4. Use model confidence to decide precision. The method first computes with coarse high-order bits. Full precision is loaded only if the largest attention probability falls below a threshold, which signals a flat, error-prone distribution. This trade of extra compute for less memory traffic suits any memory-bound inference setting.
  • I5. Prune only before fetching data. Pruning saves memory bandwidth only if the decision happens before the data is loaded. This is why SpAtten can help generative models where earlier accelerators could not, and it is a useful lens for judging other sparsity schemes.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Massachusetts Institute of TechnologyCambridge · United States243 works · 76,506 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

🧭 Topics 1Research subjects and languages the work is about.Efficient models & compression · 297 works

Cite

Cite as Hanrui Wang, Zhekai Zhang and Song Han. 2020. SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning. arXiv:2012.09852.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.1109/hpca51647.2021.00018 · arXiv 2012.09852

@misc{wang2020spatten,
  title = {{SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning}},
  author = {Hanrui Wang, Zhekai Zhang and Song Han},
  year = {2020},
  doi = {10.1109/hpca51647.2021.00018},
  eprint = {2012.09852},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.1109/hpca51647.2021.00018}
}

References

43 references, 13 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. A. Vaswani et al., “Attention is all you need,” in NeurIPS, 2017.
  2. J. Devlin et al., “Bert: Pre-training of deep bidirectional transformers for language understanding,” arXiv, 2018.
  3. A. Radford et al., “Language models are unsupervised multitask learners,” OpenAI Blog, 2019. find
  4. A. Wang et al., “GLUE: A multi-task benchmark and analysis platform for natural language understanding,” in EMNLP Workshop BlackboxNLP, 2018.
  5. E. Voita et al., “Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned,” ACL, 2019.
  6. T. J. Ham et al., “A3: Accelerating attention mechanisms in neural networks with approximation,” HPCA, 2020. find
  7. H. Jang et al., “Mnnfast: A fast and scalable system architecture for memory-augmented neural networks,” in ISCA, 2019. find
  8. P. Rajpurkar et al., “Squad: 100,000+ questions for machine comprehension of text,” EMNLP, 2016. C. Efficient Natural Language Processing The large computation of attention-based NLP models raised much interest in improving their efficiency [20], [96]– [98]. GOBO [95] proposed to compress BERT model…
  9. S. Merity et al., “Pointer sentinel mixture models,” arXiv, 2016.
  10. M. P. Marcus et al., “Building a large annotated corpus of English: The Penn Treebank,” Computational Linguistics, 1993.
  11. C. Chelba et al., “One billion word benchmark for measuring progress in statistical language modeling,” arXiv, 2013.
  12. E. F. Tjong Kim Sang et al., “Introduction to the CoNLL-2003 shared task: Language-independent named entity recognition,” in HLT-NAACL, 2003.
  13. R. Socher et al., “Recursive deep models for semantic compositionality over a sentiment treebank,” in EMNLP, 2013.
  14. D. E. Knuth, The Art of Computer Programming, Volume 3: (2nd Ed.) Sorting and Searching. Addison Wesley Longman Publishing Co., Inc., 1998. find
  15. Y. Kim et al., “Ramulator: A fast and extensible dram simulator,” IEEE Computer architecture letters, 2015. find
  16. P. Nilsson et al., “Hardware implementation of the exponential function using taylor series,” in NORCHIP, 2014. find
  17. S. Salehi et al., “Energy and area analysis of a floating-point unit in 15nm cmos process technology,” in SoutheastCon, 2015. find
  18. N. Muralimanohar et al., “Cacti 6.0: A tool to model large caches,” HPL, 2015. find
  19. M. O’Connor et al., “Fine-grained dram: energy-efficient dram for extreme bandwidth systems,” in MICRO, 2017. find
  20. H. Wang et al., “Hat: Hardware-aware transformers for efficient natural language processing,” in ACL, 2020.
  21. S. Han et al., “Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding,” in ICLR, 2016. find
  22. T. Wang et al., “Apq: Joint search for network architecture, pruning and quantization policy,” in CVPR, 2020. find
  23. S. Pal et al., “Outerspace: An outer product based sparse matrix multiplication accelerator,” in HPCA, 2018. find
  24. Z. Zhang, H. Wang et al., “Sparch: Efficient architecture for sparse matrix multiplication,” in HPCA, 2020. find
  25. H. Wang et al., “Gcn-rl circuit designer: Transferable transistor sizing with graph neural networks and reinforcement learning,” DAC, 2020. find
  26. H. Mao et al., “Park: An open platform for learning-augmented computer systems,” NeurIPS, 2019. find
  27. H. Wang, J. Yang, H.-S. Lee, and S. Han, “Learning to design circuits,” NeurIPS MLSys Workshop, 2018. find
  28. Y. He, X. Zhang, and J. Sun, “Channel pruning for accelerating very deep neural networks,” ICCV, 2017. find
  29. Y. He et al., “Amc: Automl for model compression and acceleration on mobile devices,” in ECCV, 2018.
  30. J. Park et al., “Faster cnns with direct sparse convolutions and guided pruning,” arXiv, 2016. find
  31. M. H. Mubarik et al., “Printed machine learning classifiers,” in MICRO, 2020. find
  32. A. K. Ramanathan et al., “Look-up table based energy efficient processing in cache support for neural network acceleration,” in MICRO, 2020. find
  33. Z. Song et al., “Vr-dann: Real-time video recognition via decoder-assisted neural network acceleration,” in MICRO, 2020. find
  34. H. Kim et al., “Duplo: Lifting redundant memory accesses of deep neural networks for gpu tensor cores,” in MICRO, 2020. find
  35. L. Liu et al., “Duet: Boosting deep neural network efficiency on dual-module architecture,” in MICRO, 2020. find
  36. H. Mo et al., “Tfe: Energy-efficient transferred filter-based engine to compress and accelerate convolutional neural networks,” in MICRO, 2020. find
  37. J. Weng, S. Liu, V. Dadu, Z. Wang, P. Shah, and T. Nowatzki, “Dsagen: Synthesizing programmable spatial accelerators,” in ISCA, 2020. find
  38. M. Vilim et al., “Gorgon: accelerating machine learning from relational data,” in ISCA, 2020. find
  39. L. Ke et al., “Recnmp: Accelerating personalized recommendation with near-memory processing,” in ISCA, 2020. find
  40. R. D. Evans et al., “Jpeg-act: accelerating deep learning via transform-based lossy compression,” in ISCA, 2020. find
  41. U. Gupta et al., “Deeprecsys: A system for optimizing end-to-end at-scale neural recommendation inference,” ISCA, 2020. find
  42. Y. Zhao et al., “Smartexchange: Trading higher-cost memory storage/access for lower-cost computation,” ISCA, 2020. find
  43. E. Baek et al., “A multi-neural network acceleration architecture,” in ISCA, 2020. [77] M. Zhu et al., “Sparse tensor core: Algorithm and hardware co-design for vector-wise sparse neural networks on modern gpus,” in MICRO, 2019. [78] Y. Kwon et al., “Tensordimm: A practical near-memory processing ar… — In-Datacenter Performance Analysis of a Tensor Processing Unit

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2020-spatten-efficient-sparse-attention-architecture/