Learning Adaptive Axis Attentions in Fine-tuning: Beyond Fixed Sparse Attention Patterns (2022)

Abstract: We present a comprehensive study of sparse attention patterns in Transformer models. We first question the need for pre-training with sparse attention and present experiments showing that an efficient fine-tuning only approach yields a slightly worse but still competitive model. Then we compare the widely used local attention pattern and the less-well-studied global attention pattern, demonstrating that global patterns have several unique advantages. We also demonstrate that a flexible approach to attention, with different patterns across different layers of the model, is beneficial for some tasks. Drawing on this insight, we propose a novel Adaptive Axis Attention method, which learns—during fine-tuning—different attention patterns for each Transformer layer depending on the downstream task. Rather than choosing a fixed attention pattern, the adaptive axis attention method identifies important tokens—for each task and model layer—and focuses attention on those. It does not require pre-training to accommodate the sparse patterns and demonstrates competitive and sometimes better performance against fixed sparse attention patterns that require resource-intensive pre-training.

## introduction the wide adoption of the transformer architecture @xcite in contextual language representations such as bert @xcite has spurred interest in making transformers more efficient via sparse attention patterns @xcite @xcite @xcite . the typical process for learning a transformer model (e.g., bert) with a sparse attention pattern is to replace the full attention calculation with that pattern, then pre-train the model with the usual pre-training task and fine-tune the model to downstream tasks. the use of sparse attention pattern does not necessarily significantly improve the run time of the models foot_0 but it does reduce the model memory requirement during inference time. this reduction is helpful when deploying models on mobile devices or other memory-limited devices. in this paper we offer an extensive analysis of attention patterns, organized around the following questions: (1) is pre-training essential or is it possible to employ sparse patterns during fine-tuning only? (2) which types of attention patterns are important? (3) should the same attention pattern be applied to different downstream tasks and to all layers of the model? the answer to the first question carries critical implications for the practical adoption of sparse attention approaches. most current transformerbased approaches learn fixed patterns during pretraining and then apply these to fine-tuning as well. however, it is costly and impractical to pre-train a new model from scratch when a different attention pattern is expected to be more appropriate for a task. learning the sparse attention pattern model during fine-tuning is more reasonable. with this motivation in mind, we perform a controlled experiment on the eight tasks in the glue @xcite benchmark. we find that pre-training with sparse patterns is not a crucial ingredient for good performance-learning the model solely during fine-tuning sacrifices only one or two performance points on most tasks. grounded in this finding, we perform all other experiments efficiently, starting with the same pretrained model and varying sparse attention patterns during fine-tuning alone. we start to answer the second question by analyzing the two most popular patterns: local and global @xcite . local patterns allow each token to attend only to other tokens within a given window. global patterns allow some specially designated tokens to attend to all other tokens while the remaining tokens are allowed to attend only to the specially designated tokens. we show that global pattern exhibits unique and complementary strengths that local patterns cannot capture. this finding is aligned with the design choices for recent models that benefit from the combination of both patterns @xcite . for the third question, we extend sparse-bert @xcite to an adaptive diagonal attention model. with this model, we are able flexibly learn task-wise and/or layer-wise diagonal patterns. adapting attention patterns to tasks and layers improves performance over fixed attention pattern baselines and yields equivalent memory gains/sparsity levels. motivated by these findings, we design an adaptive sparse pattern that is learned during fine-tuning and that adapts to the task, layer as well as to the input sample. our pattern is an instance of axis patterns (figure 1(d )), which are a more general form of global patterns; we name it adaptive axis attention (aaa). aaa samples the important tokens by applying a fully connected layer that is followed by gumbel softmax @xcite applied to the token representations on each transformer layer. the tokens identified as important are then designated as the global tokens and are used to form an axis-aligned attention pattern. through extensive experiments we verify that learning such an adaptive axis attention can outperform the fixed patterns adopted in longformer @xcite , bigbird @xcite and sparsebert @xcite . aaa rivals or outperforms the fixed patterns even when compared with their pre-trained variants, which require extensive time and resources for pre-training. we also show that aaa can be integrated into lightweight models, e.g., mobilebert @xcite . the benefits for mobilebert indicate that our work is complementary to other methods for reducing hidden dimensions or attention heads. our comprehensive study of different sparse attention patterns in transformers advances the field with several key insights. here we highlight some of the core definitions related to self-attention and describe prior work on sparse self-attention. ## fixed sparse attention: a comprehensive analysis in this section, we address the first two questions related to fixed attention patterns: (i) is pre-training with these really necessary or does fine-tuning alone suffice, and (ii) what are the strengths and complementary aspects of local and global patterns. ## beyond fixed sparse attention in this part, we discuss the importance of adaptiveness and propose an adaptive axis attention pattern. ## conclusion in this paper, we present a comprehensive analysis of sparse attention patterns. we demonstrate that while pre-training with sparse attention does improve performance on many tasks, using sparse attention only in fine-tuning sacrifices a bit of per-formance for a big gain in time and computational resource savings. we compare the popular local and global patterns and conclude that either type provide an advantage depending on the task. we also show that allowing sparse patterns to be adaptive to the task or layers improves performance. finally we present aaa which incorporated all these insights and learns important tokens during fine-tuning. our model is consistently and considerably better than other sparse attention pattern models and rivals models that require extensive pre-training. for future work, we anticipate to integrate the adaptive diagonal pattern with our adaptive axis pattern to construct a fully learnable pattern.