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

Abstract

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, attention becomes the compution bottleneck because of its quadratic computational complexity to input length, 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.0x with no accuracy loss, and achieves 1.6x, 3.0x, 162x, 347x speedup, and 1,4x, 3.2x, 1193x, 4059x energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

Published as a conference paper in HPCA 2021; 15 pages, 23 figures

References in corpus (13)

• Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding
• AMC: AutoML for Model Compression and Acceleration on Mobile Devices
• End-To-End Memory Networks
• Non-Autoregressive Neural Machine Translation
• Pointer Sentinel Mixture Models
• GCN-RL Circuit Designer: Transferable Transistor Sizing with Graph Neural Networks and Reinforcement Learning
• SpArch: Efficient Architecture for Sparse Matrix Multiplication
• HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
• Faster CNNs with Direct Sparse Convolutions and Guided Pruning
• TensorDash: Exploiting Sparsity to Accelerate Deep Neural Network Training and Inference
• FPSA: A Full System Stack Solution for Reconfigurable ReRAM-based NN Accelerator Architecture
• Learning to Design Circuits
• TensorDIMM: A Practical Near-Memory Processing Architecture for Embeddings and Tensor Operations in Deep Learning

Cited by in corpus (31)

• QuantumNAS: Noise-Adaptive Search for Robust Quantum Circuits
• Keyword Transformer: A Self-Attention Model for Keyword Spotting
• Enable Deep Learning on Mobile Devices: Methods, Systems, and Applications
• Hardware Acceleration of Sparse and Irregular Tensor Computations of ML Models: A Survey and Insights
• SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification
• An Algorithm-Hardware Co-Optimized Framework for Accelerating N:M Sparse Transformers
• IANUS: Integrated Accelerator based on NPU-PIM Unified Memory System
• QuantumNAT: Quantum Noise-Aware Training with Noise Injection, Quantization and Normalization
• A Length Adaptive Algorithm-Hardware Co-design of Transformer on FPGA Through Sparse Attention and Dynamic Pipelining
• A Survey on Hardware Accelerators for Large Language Models
• QOC: Quantum On-Chip Training with Parameter Shift and Gradient Pruning
• Duplex: A Device for Large Language Models with Mixture of Experts, Grouped Query Attention, and Continuous Batching
• Mobile Foundation Model as Firmware
• Mokey: Enabling Narrow Fixed-Point Inference for Out-of-the-Box Floating-Point Transformer Models
• ITA: An Energy-Efficient Attention and Softmax Accelerator for Quantized Transformers
• QuEst: Graph Transformer for Quantum Circuit Reliability Estimation
• Oaken: Fast and Efficient LLM Serving with Online-Offline Hybrid KV Cache Quantization
• Constraint-aware and Ranking-distilled Token Pruning for Efficient Transformer Inference
• Sparse-DySta: Sparsity-Aware Dynamic and Static Scheduling for Sparse Multi-DNN Workloads
• Towards Pareto Optimal Throughput in Small Language Model Serving
• FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design
• An Analog and Digital Hybrid Attention Accelerator for Transformers with Charge-based In-memory Computing
• EdgeBERT: Sentence-Level Energy Optimizations for Latency-Aware Multi-Task NLP Inference
• Magic Pyramid: Accelerating Inference with Early Exiting and Token Pruning
• Delta Keyword Transformer: Bringing Transformers to the Edge through Dynamically Pruned Multi-Head Self-Attention
• Long Exposure: Accelerating Parameter-Efficient Fine-Tuning for LLMs under Shadowy Sparsity
• PIVOT- Input-aware Path Selection for Energy-efficient ViT Inference
• Approximations in Deep Learning
• Armour: Generalizable Compact Self-Attention for Vision Transformers
• Focus on the Core: Efficient Attention via Pruned Token Compression for Document Classification
• LEGO: Spatial Accelerator Generation and Optimization for Tensor Applications