Predictive Attention Transformer: Improving Transformer with Attention Map Prediction

1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) is the state-of-the-art for sequential modeling which achieves superior performances in multiple domains , including natural language understanding ( Devlin et al. , 2019 ) , image generation ( Parmar et al. , 2018 ) and time-series forecasting ( Li et al. , 2019 ) . The performance of a transformer model largely depends on its capability of inducing reasonable dependencies among input tokens . However , as demonstrated by previous work ( Jain & Wallace , 2019 ) , it is difficult for a vanilla attention layer to capture the dependencies effectively without any apriori knowledge . To cope with this problem , recent efforts have tried to address the effectiveness of attention learning , such as concatenating self-attention with CNN layers to obtain a better representation ( Bello et al. , 2019 ; Wu et al. , 2020 ) , or synthesizing the attention maps directly ( Tay et al. , 2020 ) . In this paper , we consider another question , can we improve the learning of attention maps via a dedicated prediction model ? As we will see , it is possible through augmenting the transformer architecture by a chain of convolutional modules for attention map prediction . For a multi-layer transformer , the self-attention maps in each layer are learned independently , which introduces a huge amount of parameters and hurts the generalization ability . Our motivation is that we can bridge the attention maps from different layers , while a succeeding layer can take the knowledge from previous layers directly to induce a better dependency structure . To this end , we propose Predictive Attention Transformer ( PA-Transformer ) , which guides the learning of attention maps via a chain of convolution-based prediction modules . In each block , PA-Transformer takes all attention maps generated by the previous block as a multi-channel image . Then , by performing 2D-convolution over that image , the attention maps for the current block can be predicted effectively and efficiently . In this way , the general patterns of inter-token dependencies are shared across all blocks , benefiting the generalization ability of a multi-layer Transformer . Meanwhile , the selfattention layer in each block is guided by the predicted attention patterns and can be learned to capture complementary relationships . As shown by a real case of image classification in Figure 1 ( b ) , the attention map learned in the second PA-Transformer block correctly highlights the structure of a horse with the help of inherited knowledge from previous layers . Specifically , the convolution-based attention prediction module captures key patterns from a local perspective ( probably owning to the convolutional inductive bias ) , which guides the residual self-attention module to generated better attention maps . In contrast , the vanilla transformer learns each self-attention layer separately and does not produce good attention maps ( see Figure 1 ( a ) ) . We evaluate the performance of PA-Transformer on plenty of tasks in both natural language and computer vision domains , including text classification , natural language inference , machine translation and image classification . We also apply the generic idea of attention map prediction to other state-of-the-art models such as BERT ( Devlin et al. , 2019 ) and AA-ResNet ( Bello et al. , 2019 ) . The experimental results demonstrate the superiority of PA-enhanced transformers in terms of accuracy , memory and computational costs . In particular , for ImageNet classification task , PA-AA-ResNets achieve strong accuracy improvement on top of AA-ResNet ( Bello et al. , 2019 ) , a recent SOTA which has already encapsulated self-attention and convolutional layers jointly for image representation . Moreover , we examine the generality of PA-Transformer by incorporating it into BERT-style pre-trained models . Impressively , the average GLUE scores are lifted by 4.1 , 2.5 , 0.8 and 1.3 points on top of BERT-Base , BERT-Large , RoBERTa-Large and T5-Base respectively , while little extra parameters and computational costs are introduced to the original models . The contributions of this paper are highlighted as follows . • In this paper , we propose a novel Transformer architecture augmented by a chain of convolutional attention prediction models . Extensive experiments demonstrate consistent performance enhancement for various natural language and computer vision tasks . • Ablation study proves the effectiveness of both residual connections among attention maps and convolutional inductive bias to induce better attention patterns . To the best of our knowledge , PA-Transformer is one of the first works that take attention maps as multichannel input to a deep neural network for explicit modeling . This sheds new lights on the attention mechanism . • We empirically show that the proposed PA-enhanced method is complementary to existing works on transformer , including those pre-trained by a large corpus and enhanced by CNN layers . Thus , we believe it will have further impacts to more advanced model architectures and a broader range of applications . 2 RELATED WORK . Transformer is first introduced by Vaswani et al . ( 2017 ) for machine translation and then widely adopted in numerous tasks in natural language ( Devlin et al. , 2019 ) , computer vision ( Parmar et al. , 2018 ; 2019 ) and time-series ( Li et al. , 2019 ) domains . Transformer is solely composed of selfattention and feed-forward layers . It is much more parallelizable than Recurrent Neural Networks ( RNNs ) and demonstrates extreme superiority in large-scale training scenarios . Most notably , the cutting-edge text representation model , BERT ( Devlin et al. , 2019 ) , is based on an architecture of deep bidirectional Transformer . After pre-trained on a large-scale language corpus through the “ masked language model ” ( MLM ) loss , BERT can be fine-tuned with just one additional output layer to create state-of-the-art performances for a wide range of text-based applications . The assumption behind Transformer is that the intra-sequence relationships can be captured automatically through self-attention . But in practice , it is questionable if a self-attention layer learns reasonable dependencies among input tokens . There are many endeavors trying to analyze attention maps generated by the attention mechanism . Raganato et al . ( 2018 ) analyze the Transformer model for machine translation and show that some attention heads capture certain relations implicitly : lower layers tend to learn more about syntax while higher layers tend to encode more about semantics . As suggested by Tang et al . ( 2018 ) , the ability of inducing syntactic relations for a Transformer model is weaker than its recurrent neural network counterpart . There is a debate on whether or not the intermediate representations offered by attention mechanisms may be useful to explain the reasons for a model ’ s prediction ( Jain & Wallace , 2019 ; Wiegreffe & Pinter , 2019 ) . Moreover , Synthesizer ( Tay et al. , 2020 ) develops a strategy to replace the dot-product attention by synthesized attention maps . It is argued that explicit token-token interaction is not that important . In short , the attention maps learned by existing attention models are far from perfect , which motivates us to propose a dedicated model for extracting attention patterns . Besides , there have been successful attempts to combine convolutional and self-attention layers to enrich image and text representations ( Bello et al. , 2019 ; Wu et al. , 2020 ) . However , to the best of our knowledge , our work is one of the first that takes attention maps as multi-channel inputs and utilizes a dedicated deep neural network model for explicit modeling . We believe this is a promising direction that deserves more investigations in the future . Another limitation of Transformer lies in its prohibition for modeling long sequences , as both the memory and computation complexities are quadratic to the sequence length . To address this problem , Kitaev et al . ( 2020 ) propose Reformer which utilizes two techniques to improve the efficiency of Transformers : ( 1 ) revising dot-product attention with locality-sensitive hashing ; ( 2 ) replacing residual layers with reversible ones . Moreover , Gehring et al . ( 2017 ) leverage an architecture based entirely on convolutional neural networks for sequence to sequence learning , where the number of non-linearities is fixed and independent of the input length . Parmar et al . ( 2019 ) apply stand-alone self-attention layers to image classification by restricting the attention operations within a local region of pixels . In addition , there are other directions in the literature for the improvement of a Transformer model , such as relative positional representations ( Shaw et al. , 2018 ) , adaptive masks for long-range information ( Sukhbaatar et al. , 2019 ) , tree-based transformer ( Shiv & Quirk , 2019 ) , and AutoML-based evolved transformer ( So et al. , 2019 ) . These works are orthogonal to ours , and most of them can be encapsulated in our framework freely . 3 PREDICTIVE ATTENTION TRANSFORMER . 3.1 OVERVIEW . The representation of a text sequence can be written as X ∈ RN×C , where N denotes the sequence length and C is the dimension size . For an image representation , the conventional shape is ( H , W , C ) , where H , W and C denote height , width and channel size of the image respectively . In order to apply a standard Transformer to the image representation , we flatten its shape as X ∈ RN×C , where N = HW and each pixel serves as an individual token in the Transformer . A standard Transformer block is composed of a self-attention layer and two position-wise feedforward layers , while each attention map is generated by a self-attention layer separately without sharing information with each other . We argue that a single and independent self-attention layer is not generalizable to capture the underlying dependencies among tokens . To alleviate this prob- lem , we propose a convolution-based prediction module that calculates attention maps for the current layer based on the attention map from the previous layer . Our conjecture is that this module would predict effective attention maps in the guidance of generalized attention patterns from previous layers . Thus , the self-attention module in the current layer could dedicate itself to incorporate layer-specific knowledge into residual attention maps . We name the transformer network with proposed augmentation as Predictive Attention Transformer ( PA-Transformer ) , the architecture of which is illustrated in Figure 1 ( b ) . Each PA-Transformer block consists of three modules as a standard transformer , namely Attention Map Generation , Value Projection and Feed-Forward Layers . In addition , there are Convolution-based Attention Prediction modules , which take the attention maps from the previous block as input to predict the ones in the next block through convolutional neural networks . Note that Attention Map Generation and Value Projection are two sub-modules of a standard self-attention layer . We omit layer norm and residual connections in the figure for brevity . In the rest of this section , we will introduce the details of each module separately . 3.2 ATTENTION MAP GENERATION . Given the input representation X , the attention map can be calculated as follows . First , we compute the query and key matrices for each attention head through linear projections , i.e. , Q = XWQ , K = XWK , where Q and K denote query and key matrices respectively , WQ and WK are linear projection parameters . Then , the attention map is derived by a scaled dot-product operation : A =Attention ( X ) = softmax ( QK > √ d ) ( 1 ) Here A denotes the attention map and d is the hidden dimension size . To inject sequential information into the model , we incorporate positional encoding to the input representation . The positional encoding can be either absolute or relative , and we follow the original implementation for different backbone models . For absolute positional encoding ( Vaswani et al. , 2017 ) , it is added to the input representation X directly . For relative positional representation ( Shaw et al. , 2018 ) , the attention formulation can be re-written as : A =Attention ( X ) = softmax ( QK > √ d +R ) ( 2 ) where R = { rij } is the matrix of relative positional encoding . For text data , we have rij = qTi ei−j , where ei−j is a trainable embedding vector in terms of the relative index of two tokens . For image data , we adopt two separate embedding vectors for height and width ( Bello et al. , 2019 ) . rij = q > i e H h ( j ) −h ( i ) + q > i e W w ( j ) −w ( i ) ( 3 ) where qi is the query representation for the ith pixel , eH and eW represent for trainable embedding vectors of height and width respectively , h ( i ) and h ( j ) are the height indices for the ith and jth pixels , and w ( · ) denotes the index in width .

Summary: This paper proposed a modification to the classical transformer architecture and demonstrated significant performance gain on multiple benchmark tasks in both natural language processing and computer vision. Specifically, the authors propose to introduce a convolution-based attention map prediction module, so the dependencies of attention maps across different layers can be captured. With the extensive experiments, the proposed modification is quite effective on improving the model's performance.