Predicting Attention Sparsity in Transformers (2022)

Abstract: Transformers’ quadratic complexity with respect to the input sequence length has motivated a body of work on efficient sparse approximations to softmax. An alternative path, used by entmax transformers, consists of having built-in exact sparse attention; however this approach still requires quadratic computation. In this paper, we propose Sparsefinder, a simple model trained to identify the sparsity pattern of entmax attention before computing it. We experiment with three variants of our method, based on distances, quantization, and clustering, on two tasks: machine translation (attention in the decoder) and masked language modeling (encoder-only). Our work provides a new angle to study model efficiency by doing extensive analysis of the tradeoff between the sparsity and recall of the predicted attention graph. This allows for detailed comparison between different models along their Pareto curves, important to guide future benchmarks for sparse attention models.

## introduction transformer-based architectures have achieved remarkable results in many nlp tasks @xcite @xcite . however, they also bring important computational and environmental concerns, caused by their quadratic time and memory computation requirements with respect to the sequence length. this comes in addition to the difficulty of interpreting their inner workings, caused by their overparametrization and large number of attention heads. there is a large body of work developing ways to "sparsify" the computation in transformers, either by imposing local or fixed attention patterns @xcite @xcite , by applying low-rank kernel approximations to softmax @xcite , * work done at instituto de telecomunicações. correspondence to marcos.treviso@tecnico.ulisboa.pt or by learning which queries and keys should be grouped together @xcite @xcite . most of the existing work seeks to approximate softmaxbased attention by ignoring the (predicted) tails of the distribution, which can lead to performance degradation. an exception is transformers with entmax-based sparse attention @xcite , a content-based approach which is natively sparse -this approach has the ability to let each attention head learn from data how sparse it should be, eliminating the need for heuristics or approximations. the disadvantage of this approach is that it still requires a quadratic computation to determine the sparsity pattern, failing to take computational advantage of attention sparsity. in this paper, we propose sparsefinder, which fills the gap above by making entmax attention more efficient ( §4). namely, we investigate three methods to predict the sparsity pattern of entmax without having to compute it: one based on metric learning, which is still quadratic but with a better constant ( §4.3), one based on quantization ( §4.4), and another based on clustering ( §4.5). in all cases, the predictors are trained offline on ground-truth sparse attention graphs from an entmax transformer, seeking high recall in their predicted edges without compromising the total amount of sparsity. figure 1 illustrates our method. more precisely, to evaluate the effectiveness of our method across different scenarios, we perform experiments on two nlp tasks, encompassing encoder-only and decoder-only configurations: machine translation (mt, §5) and masked language modeling (mlm, §6), doing an extensive analysis of the tradeoff between sparsity and recall (i.e., performance on the attention graph approximation), and sparsity and accuracy (performance on downstream tasks). we compare our method with four alternative solutions based on efficient transformers: longformer @xcite , bigbird (zaheer et al., 2020), reformer @xcite , and routing transformer @xcite , along their entire pareto curves. we complement these experiments by analyzing qualitatively what is selected by the different attention heads at the several layers and represented in different clusters/buckets. overall, our contributions are: foot_0 ## sparsefinder we now propose our method to extract sparse attention graphs and learn where to attend by exploiting a special property of α-entmax: sparse-consistency ( §4.1). we design three variants of sparsefinder to that end, based on metric learning ( §4.3), quantization ( §4.4), and clustering ( §4.5). ## conclusions we proposed sparsefinder, a method to identify the sparsity pattern of entmax-based transformers while avoiding full computation of the score matrix. our method learns a low-dimensional projection of queries and keys with a contrastive objective, and comes with three variants: distance, quantization, and clustering-based. we compared these variants against competing approaches on two tasks: machine translation and masked language modeling. we obtained favorable sparsity-recall and sparsityaccuracy tradeoff curves. our theoretical sparsity provides a lower bound for how much computational sparsity can be achieved, and may guide future research on efficient transformers. training and model. we replicated the sentence-level model of @xcite with the exception that we used α-entmax with α = 1.5 instead of softmax in all attention heads and layers.