Title
Accelerating attention through gradient-based learned runtime pruning
Abstract
Self-attention is a key enabler of state-of-art accuracy for various transformer-based Natural Language Processing models. This attention mechanism calculates a correlation score for each word with respect to the other words in a sentence. Commonly, only a small subset of words highly correlates with the word under attention, which is only determined at runtime. As such, a significant amount of computation is inconsequential due to low attention scores and can potentially be pruned. The main challenge is finding the threshold for the scores below which subsequent computation will be inconsequential. Although such a threshold is discrete, this paper formulates its search through a soft differentiable regularizer integrated into the loss function of the training. This formulation piggy backs on the back-propagation training to analytically co-optimize the threshold and the weights simultaneously, striking a formally optimal balance between accuracy and computation pruning. To best utilize this mathematical innovation, we devise a bit-serial architecture, dubbed LeOPArd, for transformer language models with bit-level early termination microarchitectural mechanism. We evaluate our design across 43 back-end tasks for MemN2N, BERT, ALBERT, GPT-2, and Vision transformer models. Post-layout results show that, on average, LeOPArd yields 1.9×and 3.9×speedup and energy reduction, respectively, while keeping the average accuracy virtually intact (< 0.2% degradation).
Year
DOI
Venue
2022
10.1145/3470496.3527423
ISCA: International Symposium on Computer Architecture
Keywords
DocType
ISSN
Transformer, Learned Pruning, Gradient-Based Optimization, Attention Mechanism, Self-attention, Neural Processing Units, Accelerators, Deep Learning
Conference
1063-6897
Citations 
PageRank 
References 
1
0.34
3
Authors
5
Name
Order
Citations
PageRank
Zhenge Li110.34
Soroush Ghodrati210.68
Amir Yazdanbakhsh324115.28
H. Esmaeilzadeh4144369.71
Mingu Kang510.34