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 proposes a novel approach to improve self-attention through by bridging the attention maps from different layers via a chain of convolution-based prediction modules.  In particular, it proposes to augment the existing works on Transformer through supplementary prediction modules by CNN-based attention prediction layers.  The main contribution of this paper is the introduction of CNN-based attention prediction to enhance model predictions. Empirical studies are performed to show the superiority of the proposed model PA-Transformer over several SOTA approaches on NLP and image classification tasks.