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

By Hanrui Wang, Zhekai Zhang, Song Han. 2020

563 citations recorded by the source corpus.

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, general-purpose platforms such as CPUs and GPUs are inefficient when performing attention inference due to complicated data movement and low arithmetic intensity. Moreover, existing NN accelerators mainly focus on optimizing convolutional or recurrent models, and cannot efficiently support attention. In this paper, we present SpAtten, an efficient algorithm-architecture co-design that leverages token sparsity, head sparsity, and quantization opportunities to reduce the attention computation and memory access. Inspired by the high redundancy of human languages, we propose the novel cascade token pruning to prune away unimportant tokens in the sentence. We also propose cascade head pruning to remove unessential heads. Cascade pruning is fundamentally different from weight pruning since there is no trainable weight in the attention mechanism, and the pruned tokens and heads are selected on the fly. To efficiently support them on hardware, we design a novel top-k engine to rank token and head importance scores with high throughput. Furthermore, we propose progressive quantization that first fetches MSBs only and performs the computation; if the confidence is low, it fetches LSBs and recomputes the attention outputs, trading computation for memory reduction.Extensive experiments on 30 benchmarks show that, on average, SpAtten reduces DRAM access by 10.0× with no accuracy loss, and achieves 1.6×, 3.0×, 162×, 347× speedup, and 1.4×, 3.2×, 1193×, 4059× energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (2022, background): “It proves to be effective for various convolutional neural works (CNNs) (Han et al., 2016; Jacob et al., 2018; Nagel et al., 2019; Wang et al., 2019; Lin et al., 2020) and transformers (Shen et al., 2020; Kim et al., 2021; Liu et al., 2021; Wang et al., 2020; Bondarenko et al., 2021).”
H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (2023, unknown): “And (Wang et al., 2021) follows the same heuristic to select the most important tokens via attention scores, yet focus on relatively small-scale tasks.”
Compression of Generative Pre-trained Language Models via Quantization (2022, unknown): “SpAtten (Wang et al., 2021) proposes a sparse model with algorithm and architecture co-design, which removes uninformative tokens and attention heads.”
Efficient Streaming Language Models with Attention Sinks (2023, background): “…et al., 2023) and improve training (Dao et al., 2022; Dao, 2023) and inference (Pope et al., 2022; Xiao et al., 2023; Anagnostidis et al., 2023; Wang et al., 2021; Zhang et al., 2023b) efficiency for lengthy inputs, the acceptable sequence length remains intrinsically finite , which doesn’t…”
Efficient Acceleration of Deep Learning Inference on Resource-Constrained Edge Devices: A Review (2023, background): “Lossless reduction of computation and memory access of 3.8 × and 1.1 × is reported in the NLP model SpAtten [227] by pruning unimportant tokens and heads.”
MLPruning: A Multilevel Structured Pruning Framework for Transformer-based Models (2021, background): “One of the promising approaches for addressing the inference time and power consumption issues of these large models is pruning [36, 27, 44].”
AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration (2023, unknown): “To find the salient weight channels, the insight is that we should refer to the activation distribution instead of the weight distribution, despite we are doing weight-only quantization: weight channels corresponding to larger activation magnitudes are more salient since they process more important…”