A Probabilistic Framework for Pruning Transformers Via a Finite Admixture of Keys

Pairwise dot product-based self-attention is key to the success of transformers which achieve state-of-the-art performance across a variety of applications in language and vision, but are costly to compute. It has been shown that most attention scores and keys in transformers are redundant and can b...

Full description

Saved in:
Bibliographic Details
Published inICASSP 2023 - 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) pp. 1 - 5
Main Authors Nguyen, Tan M., Nguyen, Tam, Bui, Long, Do, Hai, Nguyen, Duy Khuong, Le, Dung D., Tran-The, Hung, Ho, Nhat, Osher, Stan J., Baraniuk, Richard G.
Format Conference Proceeding
LanguageEnglish
Published IEEE 04.06.2023
Subjects
Acoustics
admixture models
Computational modeling
Data models
Memory management
Probabilistic logic
pruning
Signal processing
Transformers
Online AccessGet full text

Cover

More Information
Summary:Pairwise dot product-based self-attention is key to the success of transformers which achieve state-of-the-art performance across a variety of applications in language and vision, but are costly to compute. It has been shown that most attention scores and keys in transformers are redundant and can be removed without loss of accuracy. In this paper, we develop a novel probabilistic framework for pruning attention scores and keys in transformers. We first formulate an admixture model of attention keys whose input data to be clustered are attention queries. We show that attention scores in self-attention correspond to the posterior distribution of this model when attention keys admit a uniform prior distribution. We then relax this uniform prior constraint and let the model learn these priors from data, resulting in a new Finite Admixture of Keys (FiAK). The learned priors are used for pruning away redundant attention scores and keys in the baseline transformers, improving the diversity of attention patterns that the models capture. We corroborate the efficiency of transformers pruned with FiAK on the ImageNet object classification and WikiText-103 language modeling tasks. Our experiments demonstrate that transformers pruned with FiAK yield similar or better accuracy than the baseline dense transformers while being much more efficient in terms of memory and computational cost.
ISSN:2379-190X
DOI:10.1109/ICASSP49357.2023.10096107