Predictive Attention Transformer: Improving Transformer with Attention Map Prediction
Authors: Yujing Wang
Yaming Yang
Jiangang Bai
Mingliang Zhang
Jing Bai
Jing Yu
Ce Zhang
Yunhai Tong
Venue:  (2021)
Keywords: Transformer
Convolution
Attention Map

Abstract: Transformer is a ubiquitous model for natural language processing and has also attracted wide attentions in other domains such as computer vision. The self-attention maps, learned independently for each layer, are indispensable for a transformer model to encode the dependencies among input tokens, however, learning them effectively is still a challenging problem. In this paper, we address this problem and propose a novel approach to improve self-attention through supplementary prediction modules. The underlying assumption is that the attention structures in the current layer should not be completely independent from those in the previous layer. Instead, we model their dependencies via a chain of prediction models that take previous attention maps as input to predict the attention maps of a new layer through convolutional neural networks. Specifically, we propose Predictive Attention Transformer and obtain significant performance gains for various kinds of tasks on top of multiple state-of-the-art models. On GLUE benchmark, the average performances of BERT-Base and BERT-Large are lifted by 4.1 and 2.5 points respectively. For machine translation, it improves the BLUE score of a vanilla Transformer consistently on IWSLT'14 De-En dataset with different model sizes. For ImageNet classification, we achieve significant improvement over a strong backbone model with comparable capacity.

Under review as a conference paper at ICLR 2021
PREDICTIVE ATTENTION TRANSFORMER: IMPROVING TRANSFORMER WITH ATTENTION MAP PREDICTION
Anonymous authors Paper under double-blind review
ABSTRACT
Transformer is a ubiquitous model for natural language processing and has also attracted wide attentions in other domains such as computer vision. The self- attention maps, learned independently for each layer, are indispensable for a trans- former model to encode the dependencies among input tokens, however, learning them effectively still challenging problem. this paper, we address this prob- lem and propose novel approach to improve self-attention through supplemen- tary prediction modules. The underlying assumption is that the attention structures in the current layer should not be completely independent from those in the pre- vious layers and can be better modeled via a convolutional inductive bias. Specif- ically, propose Predictive Attention Transformer, which predicts the attention maps through a chain of convolutional neural networks and obtains significant performance gains for various kinds of tasks on top of multiple tate-offthe-art models. On GLUE benchmark, the average performances of BERT- Base, BERT- Large, RoBERTa-Large and T5-Base are lifted by 4.1, 2.5, 0.8 and 1.3 points re- spectively. For ImageNet classification, we achieve sigmificant improvement over multiple backbone models with different capacities.
INTRODUCTION
Transformer (Vaswani 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 al.] 2018) and time-series forecasting (Li al. 2019)) The performance of transformer model largely depends on its capability of inducing reasonable depen- dencies 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 apri ori knowledge. To cope with this problem, recent efforts have tried address the effectiveness of attention learning, such concatenating self-attention with CNN layers to obtain a better represen- tation (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 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 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 induce a better dependency structure. To this end, we propose Predictive Attention Transformer PPA-Transformer) which guides the learning of attention maps via a chain of convolution-bast 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 self- attention 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 [Tb), the attention map learned in the second PA-Transforma block correctly highlights the structure of horse with the help of inherited knowledge from previous layers. Specifically, the convolution-bai based attention prediction module captures key patterns from local perspective (probably owning the
prediction modules. In each block. PA- TA-Transformer takes

---

Under review as a conference paper at ICLR 2021
Block
Bloct
Feed-Forward Layen
Layers
Value Projectis
Attes tion Map G wrotinn
Atteetio Map Genuration
Jayer
Actention of layer
lookk1
Feod-Fonward Layess
Feof-Forward
Attencion Predictor
Projotttt
Attesticn Map Generation
Map
Ixpal uttention
Attention of 1oyer
layor
Vanilla Transformer
Predictive Attention Transformer
Input
Figure 1: Overview of model architectures with an exemplar case study
convolutional inductive bias), which guides the residual self-attention module to generated better attention maps. contrast, the vanilla transformer learns each self-attention layer separately and does not produce good attention maps (see Figure 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 trans- lation and image classification. We also apply the generic idea of attention map prediction to other state-of-te oftte-of-the-arr models such BERT (Devlin et al. 2019) and AA-ResNet (Bello al.) 2019). The experimental results demonstrate the superiority PA enhancedd 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), recent SOTA which has already encapsulated self-attention and convolutional layers jointly for image representa- tion. 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 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 chain con- volutional attention prediction models. Extensive experiments demonstrate consistent per- formance enhancement for various natural language and computer vision tasks.
Ablation study proves the effectiveness both residual connections among attention maps and convolutional inductive bias to induce better attention patterns. To the best of our knowledge, DA-Transformer is one of the first works that take attention maps as multi- channel input 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 transformer, including those pre-trained by large corpus and enhanced by CNN layers. Thus, we believe it will have further impacts to more advanced model architectures and broader range of applications.
2

---

Under review as a conference paper at ICLR 2021
2 RELATED WORK
Transformer 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 self- attention 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 edge text representation model, BERT (Devlin et al. 2019), is based on 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 auto- matically 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 atten- tion maps generated by the attention mechanism. |aganato et al. (2018) analyze the Transformer model for machine translation and show that some attention heads capture certain relations implic- itly: 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 Trans- former model is weaker than its recurrent neural network counterpart. There is debate on whether or not the intermediate representations offered by attention mechanisms may be useful to explain the reasons for model's prediction (Jain & Wallace 2019; Wiegrefte & Pinter 2019). Moreover, Synthesizer (Tay et al. 2020) develops 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 dedicated model for extracting attention patterns. Besides, there have been successful attempts to combine convolutional and self-attention layers to enrich image and text representa- tions (Bello et al., 2019; Wu al.) 2020). However, 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 prob- Kitaev al. (220)) propose Reformer which utilizes two techniques improve the efficiency of Transformers: (I) revising dot-product attention with locality
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 re- gion of pixels. In addition, there are other directions in the literature for the improvement of a Transformer model, such relative positional representations (Shaw al. 2018), adaptive masks for long-range information (Sukhbaatar et al. 2211)) tree-based transformer (Shiv & Quirk] [2019), and AutoML-based evolved transformer (So et al.) (219).. These works are orthogonal 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 E RNN where N denotes the se- quence 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 respec- tively. In order to apply a standard Transformer to the image representation, we flatten its shape individual token in the Transformer.
as X NCC where N HW and each pixel serves
A standard Transformer block is composed of a self-attention layer and two position-wise feed- forward layers, while each attention map is generated by 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-
3

---

Under review as a conference paper at ICLR 2021
lem, we propose convolution- prediction module that calculates attention maps for the cur- rent 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 previ- ous layers. Thus, the self-attention module in the current layer could dedicate itself to incorporate layer-specifie knowledge into residual attention maps.
We name the transformer network with proposed augmentation as Predictive Attention Transformer (PA-Transformer), the architecture of which illustrated in Figure [b).. Each PA-Transformer block consists of three modules as a standard transformer, namely Attention Map Generation, Value Projection and 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 standard self-attention layer. We omit layer norm and residual connections in the figure for brevity. the rest of this section, we will introduce the details of each module separately.
Layers. In addition, there are Convolution based Attention Prediction
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 = XW K = XW* where Q and K denote query and key matrices respectively, WQ and W are linear projection parameters. Then, the attention map is derived by scaled dot-product operation:
QK
A =Attention softmar
Vd
Here A denotes the attention map and is the hidden dimension size. To inject sequential informa- tion into the model, we incorporate positional encoding to the input representation. The positional encoding can be either absolute 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 al., 2018), the attention formulation can be re-written as:
QK
A =Attention (X) softmax
+ +R)
(2)
Vd
where R = (rir)) } is the matrix of relative positional encoding. For text data, we have Fly = q, c;-j+ where e;-j is 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).
H Fij q, en()) + q, (m(()-
(3)
where qi is the query representation for the 7th pixel, e# and w represent for trainable embedding vectors of height and width respectively, h(i)) and h(j) the height indices for the 7th and jth pixels, and w-) denotes the index in width.
3.3 CONVOLUTION-BASER ATTENTION PREDICTION
In vanilla Transformer, the attention maps are calculated by a single dot-product operation. How- ever, it may not be effective for constructing dependencies among tokens, since each attention map generated independently without sharing information with each other. To tackle this limitation, we propose convolution-based prediction module, which calculates attention maps for new block by applying convolutional neural network to existing ones. The prediction module expected to induce effective attention patterns in the guidance of generic knowledge. Thus, the self-attention module can dedicate itself to capture residual knowledge for a specific layer.
Assume there are K heads in each layer, then we have K attention map outputs from a Attention Map Generation module. They construct tensor A E N N'x (N is the sequence length), which adopt
can be viewed as N x N image with K input channels. On the basis of this input, several 2D-convolttional layers with x 3 kernels to predict the attention maps for the next block. The output channel is also set be K, the attention maps of all heads can be generated jointly.
4

---

Under review as a conference paper at ICLR 2021
We apply a ReLU activation after each 2D convolluutio layer provide non-linearity and sparsity. Finally, the predicted attention map is combined with that learned by dot- dot-product attention:
sof toftmalaa por.lgi))) + - ax) Attention X (ogitss
(4)
where Apre logsts is the matrix of attention logits (before softmax) from the previous layer; Attention ()) Koggits is the attention matrix of logits calculated by the current self-attention layer; [0, is hyperrparameete to balance the importance of two branches. In our experiments, a chosen empirically on the validation set. We will analyze the impact of different a values in Section We do not apply convolution- attention prediction to the first transformer block.
5
3.4
VALUE PROJECTION AND FEED-FORWARD LAYERS
Given the attention map A, the rest of a PA-Transformer block includes value projection and position-wid -wise
layers that are identical to standard transformer block. The value projection layer can be formulated as:
H. = A, H
H + ... Hx)WO
(5)
where A, is the attention map for the tt head, W. is the parameter of value projection, and H; the corresponding representation generated by value projection. Afterwards, the representation of each head is concatenated (denoted by ()) and fed into a linear projection layer with parameter Wo At last, the block is followed by two position- wise feed-forward layers:
PA- PA-Transformer(X) ReLU(HW, bi) b2
(6)
Conventionally, the dimension of W1 four times of Wo and W₂ to form a bottleneck structure.
4 EXPERIMENTS
4.1 NATURAL LANGUAGE PROCESSING
First, we start from scratch training scenarios. Then, we conduct experiments of fine-tuning based on pre-trained models. This is practical setting, as it will be exhaustive to test all model variants by pre-training them from scratch.
4.1.1 SCRATCH TRAINING SCENARIOS
We consider three NLP tasks which are commonly adopted for the evaluation of scratch training: sentiment classification, natural language inference and machine translation.
Datasets We adopt SST (Socher et al. 2013) and Yelp (Zhang et al.) 2013) datasets for sentiment classification, SNLI (Bowman al. (0155) for natural language inference, and IWSLT' 14 German- English (De-En) (Cettolo et al.] 2014) for machine translation. Max token length set to 64 for both SST and SNLL, 256 for Yelp, and 4096 for IWSLT" 14 De-En. Detailed descriptions available in Appendix[A]
Models. Existing works (Bello et al. 2019) have showed the effectiveness of convolutional neural network to be used in parallel with self-attention in multi-branch structure. Our work leverages CNN for reasoning on the attention maps, which is orthogonal existing ones. To examine the com- plementarity, we take a multi-branch network architecture that concatenates self-attention and con- volutional layers jointly as another baseline model. This architecture named as Cony-Transformet and its predictive counterpart is called PA-Conv- caan-Tranffomer (see Figure 4 in the appendix). To ex- amine the effectiveness of the proposed method, we compare the following model variants: (1) CNN contains M layers of ID-convolution of kernel size 3. (2) Transformer is stacked by M stan- dard Transformer blocks (Figure Ia)). (3) PA-Transformer is stacked by Predictive Attention Transformer blocks visualized in Figure [b).. (4) Conv-Transformer leverages a multi-branch architecture that concatenates self-attention and convolutional layers in each Transformer block. (5) PA- PA-CCom-Transformer is stacked by M Predictive Attention Convolutional Transformer blocks (Figure 4). We set M = by default except for the machine translation task, where we follow the implementation of Joulin et al. (207)) and the block number for both the encoder and
5

---

Under review as a conference paper at ICLR 2021
Model
#Params
#FLOPs (SST)
SST
Yelp
SNLI
De-En
CNN
2.77M
317.38M
43.56 0.43
62.09 + 0.72
Transformer
2.79M
338.86M
47.64 0.70
63.23 + 0.39
83.22 + 0.39
31.64
PA- PA-Iransformer
2.79M
349.10M
48,60 = 0.41
63,90 = 0.18
84.63 = 0.29
32.51
Conv-Transfomuue
2.58M
321.34M
49.74 0.55
64.21 1 0.09
85,40 1 0.61
32.24
2.58M
331.39M
50.66 0.64
64.75 0.30
86.28 + 0.25
32.48
Table Performance for natural language processing tasks
Model
Avg
CoL.A
SST-2
MRPC
STS-B
QQP
MNL.I-m/-mm
QNLI
RTE
WNLI
BERT Base
77.4
51.7
93.5
87.2/82.1 1
86.7/85.4
91.1/89.0
84.3/83.7
90.4
67.2
65.1
PA-BERT-Base
81.5
59.8
93.7
88.9/90.8
89.3/89.2
91.4/88.3
84.8/85.2
92.0
68.6
65.1
BERT-Large
80.5
60.5
94.9
89,3/85
87.6/86.5
92.1/89.3
86.8/85.9
92.7
70.1
65.1
PA-BERT-Large
83.0
63.1
95.4
90.4/88.9
88.9/88.0
92.4/89.9
87.7/86.2
93.5
70.9
65.1
RoBERTs Large
83.1
63.8
96.3
91.0/89.4
72.9/90.2
92.7/90.1
89.5/89.7
94.2
84.2
65.1
PA Ra-RoBBKK--Large
83.9
65.4
96.5
91899066
73.6/90.3
93.0/90.1
90.3/89.7
95.0
85.2
65.1
TS-Base
83.2
51.1
95.2
90.7/87.5 5
89.4/88.6 6
72.6/89.4
87, 87.1/86.2
93.7
80.1
PA-TS-Base
84.5
53.3
92.4/89.5
89.6/89.1
89.4/91.5
86.0/85.
92.9
80.7
TS-Base (Dov)
83.5
53.1
92.0/88.7
89.1/88.9
88.2/91.2
64.7/85.0
91.7
76.9
Synthesizer-TS BS-Base (Dev)
84.1
53,3
92.2
91.2/88.77
89.3/88.9
88. 6/9
85.0/84.
92.3
81.2
PA-T5-Base (Dev)
84.4
53.6
92.4
92.2/89.1
89.7/89.1
89.7/91.7
85.4/84.9
92.5
80,7
Table 2: Comparison of different model backbones on GLUE benchmark.
decoder networks. We only apply TA-Trassfommer architecture to the encoder while its application decoder is considered as future work. We use one convolutional layer in each prediction module. More analysis can be found in Section 5
Settings. For sentiment classification and natural language inference tasks, the hidden dimension size is set to be 256 with 8 heads, and the first feed-forward layer in each block has hidden size of 1024. An initial Transformer. the hidden size is one half of the total dimension for both convolutional and self- attention branches, and the self-attention branch has 4 heads. We apply dense connections between different layers similar to DenseNet (Huang et al. 2017). The word embedding vectors are ini- tialized by GloVe (Pennington et al., 2014) and fixed during training. For machine translation, we follow recent work (Wu al. 2020) that removes the bottleneck structure and set the hidden size as 160 for all layers. Besides, we utilize 10K joint byte pair encoding (BPE) (Sennrich et al.] 2016) vocabulary in lower case. The hyrer-parameters are chosen empirically on the validation set of each dataset. Specifically, the value of a is selected by grid search from a search space of {0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.3, .... 0.9) on the validation set. Finally, we set it as 0.1. 0.1, 0.5 and 0.1 for both PA- PA-Transformer and PA
convolutional layer is applied to the input embedding. In Conv-
on SST, Yelp. SNLI and IWSLT' 14 De-En datasets respectively. All models are trained by the Adam optimizer (Kingma & Ba 2014). More settings are described in Appendix B.2.1
Results. The performances of different models on four natural language processing datasets are summarized in Table For the first three datasets, the numbers are obtained by running each model five times with random seeds, and their average accuracy and standard deviation are reported. For De-En machine translation. we leverage default seed and report the BLUE score. Note that the number of parameters and FLOPs changes for each dataset, and we just report the ones on SST for inference. As shown in Table I PA- Transsfoore outperforms vanilla Transformer significantly on all datasets. Consistent with previous findings (Wu et al., 2020), we observe steady improve- ments by concatenating convolutional and self-attention layers for text representation (refer Conv- Transformer). Furthermore, as indicated by Aa--Con--Transfore attention prediction modules into Conv-Transformer brings supplementary benefits. The compar- isons are fair since all models have similar numbers of parameters and FLOPs.
incorporating convolution-base
4.1.2 FINE-TUNING FROM PRE-TRAINED MODELS
Pre-trained language models like BERT (Devlin et al. 2019) become popular in recent years. These models are based on the bi-directional transformer architecture and pre-trained using Masked Lan-

---

Under review as a conference paper at ICLR 2021
Mary
Mayy
Mary
Marr
tried
tried
tried
tried
<atapo-
<xtapo-
ettago
<atepo
BERT-Base
(b) PA-BERT-Bas Base
(c) Predicted Attention
(d) Residual Self-Attenion
Figure 2: Visualization of attention maps for an exemplar case study
guage Model (MLM) loss on a large corpus. An interesting question is that if the proposed attention prediction module is still useful when fine-tuned from a pre-trained model checkpoint. To answer this question, we choose GLUE benchmark (Wang et al. 2018) for an empirical study.
Settings. The encoder network of BERT consists of multiple transformer blocks. In PA-BERT, We replace each transformer block with PA-Transformer proposed in this paper. We load the pre- trained checkpoints of BERT-Base, BERT-Large, RoBERTa-Large and T5-Base directly and fine- tune them for different downstream tasks. For PA-BERT, the additional parameters randomly and trained jointly during fine-tuning. The hyper- parameers are chosen on validation sets and reported in Appendix B.2.2 Finally, make predictions the test data and send the results to GLUE online evaluation service.
initialized
Results. The comparison between BERT-style models shown in Table PA-BERT's perform generally better than vanilla BERTs on mutiple downstream tasks. Specifically PA-BERT-BE Base, PA- achieve average scores of 81.5, 83.0 and 83.9 the GLUE
BERT-Large and PA-RoBERTa-Laar test set, increasing 4. 2.5 and 0.8 points based on BERT-Base, BERT-Large and RoBERTa-LLan respectively. We also compare vanilla Transformer, Synthesizer and PA Transformer architectures on T5-Base. Note that the corresponding average scores do not count for the WNLI dataset following the baseline setting (Tay et al. (2020). Impressively, PA-T5-Bard outperforms Synthesizer after being fine-tuned on each task individually from T5-Base checkpoint, whereas Synthesizer is pre-trained from scratch and fine-tuned in a multi-task framework. This is advantageous since PA-Transformer can be applied to improve any pre-trained transformer models without exhaustive pre-training again. In addition, worth mentioning that PA-BERT-Base boosts the score by 8.1 on CoLA dataset, indicating its superior generalization ability for small datasets.
Analysis. We compare the attention maps generated by BERT-Base and PA-BERT-BE for more insights. We find that the tionships for making the final decision. For instance, we have a case study in Figure visualizing the attention maps from the last layer. The sentence is "Mary tried John to go abroad.", and the task is to check the grammatical correctness. As shown in Figure aa.). BERT-Base only focuses on verbs and stop signs, leading mis-classification result. In contrast, PA-BERT-Base (Figure 2(b)) learns to attend to the relationships between "tried" and "John",, which correctly captures the error part and gives the right answer. In addition, we visualize the attention maps generated by the convolution- based prediction layer and the residual self-attention layer in Figure and (d) respec- tively. The design of generalizing inter-token patterns through convolutional inductive bias seems to be beneficial for emphasizing problematic parts, whereas self-attention captures complementary relationships, for example, between the words "go" and "abroad".
Base model can better highlight the key tokens and their rela-
4.2
IMAGE CLASSIFICATION
4.2.1
CIFAR
Settings, The model architecture used for CIFAR image classification similar to that introduced in Section except for 3 2D-convolution kernels applied image representation. All models are stacked by 12 blocks with 64 filters at the beginning. For every blocks, we reduce the image size via 2 x 2 average pooling layer and double the hidden size at the same time. Finally, apply

---

Under review as a conference paper at ICLR 2021
Model
#Params (CIFAR-10)
#FLOPs
CIFAR-10
CIFAR-100
CNN
5.36M
1.11G
95.90
78.17
Conv-Tranfformer
4.92M
1.05G
95.82
79.86
4.92M
1.14G
96.22
80.21
Table 3: Accuracy on CIFAR image classification datasets
Model
#Params
#GFL.OPs
Top-1
Top-5
ResNet-34
21.8M
7.4
73.79
91.43
AA-ResNet-34
20.7M
74.33
91.92
PA-AA-ResNet-34
20.7M
7.9
74.90
92.20
ResNet-50
25.6M
75.99
93.00
AA-ResNet-50
25.8M
8.3
77.15
93.52
PA-AA-ResNet-50
25 35.8M
8.7
77.55
93.81
ResNet-101
44.5M
15.6
77.40
93.65
AA-ResNet-
45,4M
16.1
78.31
94.16
PA-AA-ResNet-101
45.4M
17.2
78.49
94.23
Table 4: Accuracy on ImageNet dataset
global average pooling layer and feed the flattened representation the classifier. Note that the memory is exhaustive for the original image size (32 32). Thus, we only leverage self-attention after the image size reduced to x 8. We set a 0.01 empirically.
Results. As shown in Table Conv-Transfommer is par with CNN on CIFAR-10 and performs significantly better on the more difficult dataset, CIFAR- 100. Conv-Transformer a strong baseline, it matches the performance of 7M DenseNet (Huang et al. (017)) both CIFAR- 10 and CIFAR- 100 datasets with less than 5M parameters. After adding the attention prediction module, PA-Conv- Transformer boosts the performances on two datasets consistently with comparable parameters and computations, demonstrating the superiority our approach.
4.2.2 IMAGENET
AA-ResNet [Bello et al. 2019) proved that traditional CNN models could benefit from attention mechanisms. One may curious about if our revised attention model would bring additional advan- tages. Hence, take AA-ResNet as the backbone model for evaluation on ImageNet.
Settings. We follow the experimental protocol of AA-ResNet with standard ResNet architec- tures (He et al. 2016). We for each attention prediction module. All models are trained by 100 epochs on 8 TESLA V100 GPUs. Major hyper-parameterr are as follows: optimizer is SGD with momentum 0.9, batch size 32 worker, weight decay is le-4. For the first 5 epochs, the learning rate is scaled linearly from to 02128, and then it divided by at epoch 30, 60, 80 and 90 respectively.
0.5 for PA-AA-Reste
and adopt one convolutional layer
Results. As shown in Table 4 AA-ResNets outperform corresponding ResNet models by large margin. The proposed PA-AA-ResNer further boost the top-1 accuracy by 0.73%, 0.51% and 0.23% on top of Res -50 and -101 respectively. These numbers statistically significant under 95% confidence level. Arguably, the performance lifts are owing to better attention maps calculated by the proposed framework of convolution-basel predictive attention.
5 ANALYSIS
Effectiveness of convolutional layers. Our default configuration utilizes one convolutional layer for an attention prediction module. To examine the effectiveness of convolutional neural networks and the impact of layer number, we take layer number hyper-parameter and compare the results of different settings. When layer number is set to zero, the setting equals to a vanilla Transformer with residual connections (referred to "Transformer with RC"). The performances of vanilla
8

---

Under review as a conference paper at ICLR 2021
13.0
BLED
2.4
12.5
0.4
Transformer
PA-Tnnnfommet
12
Alpha
Siee()
(a) Analysis of
(h) Analysis of model capacities
Figure 3: Analysis on IWSLT' De-En dataset
Model
SST
SNLI
De-En
Model
ImageNet
Transformer
47.64
83.22
31.64
AA-ResNet-34
74.33
Transformer with RC
48.14
83.81
32.30
AA-ResNet-34 with RC
74,36
I-Layer PA-Transformer
48.62
84.63
32.51
1-Layer PA-AA-ResNet-34
74.90
2-Layer PA-Transformer
48.82
84.78
31.67
2-Layer PA-AA-ResNet-34
74.35
3-Layer PA- Transformer
48.22
84.65
31.55
Layer PA-AA-ResNer-s
Table 5: Analysis of PA-Transformer
Table 6: Analysis of AA-ResNet-34
Transformer, Transformer with RC, and PA- PA-Transformers consisting of different convolutional lay- ers are reported in Table We can see that the residual connection is beneficial, and we generally achieve strong performances all datasets through convolutional layer. Similar trends ob- served in the ImageNet experiments, where we analyze the effectiveness of residual connections and convolution-based prediction modules in Attention Augmented ResNet-34 architecture. As shown in Figure 6, the best performance is obtained with one predictive convolutional layer.
Sensitivity of a. In Equation (4). there is a ratio a E [0. to combine the attention maps gener- ated by self-attention and convolution- based prediction. Figure 3(a) and Figure 3 (in the appendix) visualize the test accuracy of different values of a on IWSLT" 14 De-En and SNLI datasets. Note that 0 corresponds to vanilla Transformer. Although the best varies in different datasets, the performance gains are stable for multiple choices of a. Thus, one can easily find a reasonable value by searching the validation set. In addition, want to point out that the value of is not necessarily corresponding to relative importance, but also has an effect of normalization. In the future work, we would like explore advanced strategies to get rid of this hyper-parameter.
Impact of model capacity. We further compare Transformer and PA- DA-Transfommer with different parameter sizes on De-En dataset in Figure 3(b) We follow the setting in Wu et al. (2020) to take three models with roughly 3M, 6M and 9M parameters. As illustrated by the figure, PA-Trans Transformer scales well with the growth of parameter size and outperforms vanilla Transformer consistently with different model capacities. Experiments on ImageNet (Table demonstrate a similar trend.
CONCLUSION
In this paper, propose novel transformer model, PA- PA-Transformer, which facilitates the learn- ing of attention maps by a chain of convolutional neural networks. demonstrates superior per- formances various tasks in NLP and CV domains. Future works are considered in three aspects. First, we aim to incorporate this idea into the decoder network and evaluate its effectiveness. Second, would like to explore its applications for more tasks and domains, such question answering, object detection and time-series forecasting. Last but not least, we plan to develop a more automatic strategy for choosing replacing the parameer

---

Under review as a conference paper at ICLR 2021
REFERENCES
Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shiens, and Quoe V Le. Attention augmented convolutional networks. Proceedings of the IEEE International Conference on Computer Vi- sion (ICCV), pp. 3286-3295, 2019.
Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. A large anno- tated corpus for learning natural language inference. In Empirical Methods in Natural Language Processing (EMNLP), pp. 632-642, 2015.
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. In Proceedings of the International Workshop on Spoken Language Translation, Hanoi, Vietnam, volume 57, 2014.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies (NAACL-HLT), 2019.
Terrance ries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint
2017.
Jonas Gehring, Michael Auli, David Grangier. Denis Yarats, and Yann Dauphin. Convolutional sequence to sequence learning. In International Conference on Machine Learning (ICML), pp- 1243-1252, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IFFE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinherger. Densely connected convolutional networks. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 4700-4708, 2017.
Sarthak Jain and Byron C Wallace. Attention is not explanation. In North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT) pp. 3543-3556, 2019.
Armand Joulin, Moustapha Cissé, David Grangier, Hervé Jégou, et al. Efficient softmax approxima- tion for gpus. In International Conference on Machine Learning (ICML), pp. 1302-1310. PMLR, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, arXiv preprint 1rXiv:1412.69800 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR), 2020.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. In Advances Neural Information Processing Systems (NeurIPS), pp. 5243-5253, 2019.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11169 1692,, 2019. URL /htttp://axxi....t www.aawwww.w..aaaaaoooooaaooaeoe hwwpgg/abbs/9999 11692
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. Interna- tional Conference on Learning Representations (ICLR), 2017.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning (ICML), 2018.
10

---

Under review as a conference paper at ICLR 2021
Niki Parmar, Prajit Ramachandran, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens. Stand-alone self-attention in vision models. In Advances in Neural Information Processing Sys- tems (NeurIPS), pp. 68-80 28-80,200 2019.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP), pp. 1532-1543, 2014.
Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International Conference on Machine Learning (ICML), pp- 4095-4104. 2018.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint urXiv:/9/00 10683, 2019.
Alessandro Raganato, Jörg Tiedemann, et al. An analysis of encoder representations in transformer- based machine translation. In EMNLP Workshop BlackboxNI.P: Analyzing and Interpreting Neu- ral Nenworks for NLP, 2018.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 11(())211-222, 2015. doi: 10.1007/s 1263-015-08166yy.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 1715-1725, 2016.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position represen- tations. In North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), pp. 464-468, 2018.
Vighnesh Shiv and Chris Quirk. Novel positional encodings to enable based transformers. In Advances in Neural Information Processing Systems (NeurIPS), pp. 12081-12091 2019.
David So, Quoc Le, and Chen Liang. The evolved transformer. In International Conference on
Machine Learming (ICML), pp. 5877-5886, 2019.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality sentiment treebank. Empirical Methods in Natural Language Processing (EMNLP), pp. 1631-1642, 2013.
Sainbayar Sukhbaatar, Édouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings of the 57th Annual Meeting of the Association for Compu- tational Linguistics (ACL), pp. 331-335, 2019.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pallern recognition (CVPR), pp. 2818-2826, 2016.
Gongbo Tang, Mathias Müller, Annette Rios Gonzales, and Rico Sennrich. Why self-attention? a targeted evaluation of neural machine translation architectures. In Empirical Methods in Natural Language Processing (EMNLP), 2018.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. arXiv preprint arXiv:2005.00 00743, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems (NeurIPS), pp. 5998-6008, 2017.

---

Under review as a conference paper at ICLR 2021
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations (ICLR), 2018.
Sarah Wiegreffe and Yuval Pinter. Attention not not explanation. In Empirical Methods Natural Language Processing (EMNLP), pp. 11-20, 2019.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer. Patrick von Platen, Clara Ma, Yacine Jermite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's 's transformers: State- of-the-art natural language processing. ArXiv, abs/1910.03771, 2019.
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. In International Conference on Learning Reprexentations (ICLR), 2020.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character- convolutional networks for text clas- sification. In Advances in neural information processing systems (NeurlPS), pp. 649-657, 2015.
A DATASETS
SST. The Stanford Sentiment Treehank (SST) dataset (Socher al., (013) contains more than 10,000 movie reviews collected from vottentomatoes.com website, while each piece of review labeled by one of five fine-grained categories of sentiment polarity. We utilize the standard train/dev/test split of this dataset.
Yelp. The Yelp dataset (Zhang et al. 2015) contains business reviews from yelp.com which also classified into five polarities. We select 10% of the training data randomly for validation.
SNLL The Stanford Natural Language Inference (SNLI) datase (Bowman et al. 2015) contains 570k human-written English sentence pairs manually labeled for balanced classification with the following labels: entailment, contradiction, and neutral. We follow the original split of the released data.
IWSLT'14 De-En The IWSLT' German German-English (De-En) machine translation dataset (Cettolo et al. 2014) comes from translated TED talks. The dataset contains roughly 153K training parallel sentences, 7K development sentences, and 7K test sentences. We apply 3PE/[Sennricl et al. (2016) to the source language corpus and the target language corpus separately and the final vocabulary sizes of German and English are 8848 and 6632, respectively.
GLUE The GLUE (General Language Understanding Evaluation) benchmark contains several dif- ferent NLP tasks in different fields. In STS-B, each pair of sentences is annotated by a similarity score from to 5, and the task is to predict the target score by regression. All other tasks pervised classification ones, where MRPC and QQP are similarity and paraphrase tasks; CoLA and SST-2 are single-sentence classification tasks; MNLI, RTE, and QNLI are pairwise inference tasks. Such a wide range of NL. tasks can be used to assess the generalization ability of different methods.
CIFAR. CIFAR contains colored natural images with 32 32 pixels. CIFAR-10 10 consists of images drawn from 10 classes and CIFAR-100 contains 100 classes in total. The training and test sets contain 50,000 and 10,000 images respectively, and 5,000 training images are hold out for validation.
ImageNet. The ImageNet 2012 classification datase comes from ImageNet Large Scale Visual Recognition Challenge (Russakovsky al. 2015), which consists of 1000 classes with 1.28 million training images and 50k validation images. One can submit the inference results to online evaluation service get the test accuracy.
httpss//gluebeercharkk.cck
https://www.c... http://www...orrrroooeoo/////fff
http::/www/ wwwp://ww.w.nee....eoo..... net.org/
12

---

Under review as a conference paper at ICLR 2021
Block (1-1)
Feed-Forward
Concat
Projection
10/20 Convolution
Attention Map Generation
Block
Convolution based Attantion Prediction
Food Forward Layers
Concat
Projection
10/20 Convolution
Attension Map Gemeration
Figure 4: Architecture PA-Con--Transforeer
B DETAILED SETTINGS FOR REPRODUCTION
B. PA-CON--RRANSFORMER
The architecture of Predictive Attention Convolutional Transformer -e--oon--Tranffmmnnn is il- lustrated in Figure Compared to non-convolutional counterpart, we use one half of dimension for the convolutional layer and the other half for the self-attention layer. The head number of self- attention is set as 4.
B.2 NATURAL LANGUAGE PROCESSING
B.2.1 SCRATCH TRAINING
The Natural Language Inference (NLI) task takes pair of sentences as input. Similar to BERT- concatenate sentences into single by putting special token in the middle and adding field embedding vector to the input representation. Based on this design, can apply the same architecture as sentiment classification to this task.
style implementation,
We train SST and Yelp sentiment classification datasets for 10 epochs with cosine learning rate decay (Loshchilov & Hutter 2017), where the initial learning rate set as 4e-4 and the end learning rate is set to be 1e-6. For SNLI dataset, the initial learning rate is set as 4e-4 and is divided by 10 every 10 epochs until 30 epochs are finished. For all datasets, the batch size is 64 and l2 norm 2e- 6. We adopt 0.4 and 0.2 dropout ratios for the embedding and fully-connected layers respectively. These typer-parameters are chosen empirically on the validation set using a vanilla Transformer. To avoid exhaustive search, do not tune the hyper-parameters variants, except for the linear combination ratio in PA-Transformer and Con---ooonnooommnn We adopt relative positional encoding in Shaw et al. (2018).
different datasets and model
We train IWSLT' 14 De-En dataset using Adam optimizer (Kingma Ba 2014) with 31 = 0.9, 82 0.98 and inverse square root learning rate scheduling with linear warmup. The learning rate is 5e-4 and the warmup step is set to 4000 as in /Vaswani et al. (2017). Also, we use dropout apply early stopping the training procedure
rate p 0.2 and label smoothing 0.1.
13

---

Under review as a conference paper at ICLR 2021
Task
Training Epochs
Batch Size
Learning Rate
Adam Epsilon
Dropout Rate
CoLA
3
2c-5
le-8
0.1
0.2
SST-2
le-5
le-8
0.1
0.1
MRPC
2c-5
0.1
0.2
STS-B B
2e-5
le-8
0.1
0,4
QQP
16
2e-5
le-8
0.1
0.2
MNLI
2c-5
0.1
0.4
QNLI
3
16
2c-5
le-8
0.1
0.4
RTE
2e-5
le-8
0.1
0.2
WNLI
2e-5
0.1
0.2
Table 7: Detailed hyper-parameter settings for GLUE fine-tuning.
with a patience 5 epochs. In the evaluation phase, we average the final 10 checkpoints and conduct beam search with size 5. We adopt absolute positional encoding according to the original implementation (Vaswani et al. 2017).
B.2.2 PRE-TRAINING
As introduced in Devlin al. 2019), BERT-Base and PA-BERT-Base have 12 layers and 12 atten- tion heads with hidden dimension 768. BERT-large and Conv-BERT-large have 24 layers and 16 heads, while hidden dimension for each intermediate layer of the final fully-connected layer before Softmax is set to be 2000. We download the officially re- leased checkpoints of BERT-Base and BERT-Large and initialize the additional parameters for PA-BERT-Base and PA DA-BERT-Large randomly.
1024. The hidden dimension
Meanwhile, we also conducted TS-Base (Raffel et al. 2019) as the backbone models. We add the idea of predictive attention map to and PA-T5-
of comparative tests with RoBERTa-Large (Liu et al. 2019) and
the attention structure these models, named as PA-RoBERTa-- For RoBERTa and PA-RoBERTH they have 24 layers with 16 attention heads. The total hidden size of all heads 1024, and the hidden dimension of the final fully-connected layer 4096. Follow- ing (Raffel et al., 2019), we use NLP library Transformers Wolf et al., (019) implemented by the huggingface team to implement the base version of T5, which has 220 million parameter. In order to fine-tune the tasks in the glue benchmark, RoBERTa-Large and T5 T5-Basc as our start checkpoints.
respectively.
downloaded the official pre-trained parameters for
R
We use the Adam optimizer (Kingma & Ba 2014) with epsilon 1c-8. The dropout rate is set 0.1 empirically. We used grid search to optimize the values of hyper-paramett
on validation data. We search the learning rate in (1e-4, 1e-5, 2e-5), batch size in {8, 16), training epochs in (2, 3, and of equation 4 in 10 0.2, 0.4). We find that the following setting is the best choice for most tasks: learning rate 2e-5, batch size 8, training epoch number and 0.2. The specific hyper-parameter for each task is shown in T
B.3 IMAGE CLASSIFICATION
B.3.1 CIFAR
We have two bottleneck feed-forward layers on top of each LUNNCoon-Transfommerrlll -on- Transformer block. There dense connections between different layers like DenseNet (Huang et al. 2017). The training batch size is set as 50 for all models. We adopt cosine learning rate
schedule according to Pham al. (011).
hmin + 0.5 - +
11-12
H- 766 768_A-12.21p
http:://ooooaaaa.o......oooooeer 1.-24
024_A-16.2ip
models/base
14

---

Under review as a conference paper at ICLR 2021
355
6469
Accuracy
3542
Accuracy
0355
319
0859
Transformer
Transformer
04
0d
Alpha Rate
Model Size(M)
Figure 5: Different values of a on SNLI
Figure 6: Different model capacities on SNLI
Where is the learning rate and Teur is the epoch ID in the current cycle. We set lmne 0.05 and hmin = 0.0001. For CIFAR-10, 10, the value of T is set as 10 in the first cycle and is multiplied by a factor of 2 at the end of each cycle. The dropout rates of convolutional, self-attention and MLP layers are set to be 0.2, 0.2 and 0.2 respectively; l2 norm is set as Ic-4. Each model is trained by 630 epochs which consist of six complete cycles. For CIFAR-100, the value of T is in the first cycle and is multiplied by factor of 2 at the end of each cycle. The dropout rates of convolutional, self-attention and MLP layers are set as 0.2, 0.0 and 0.2 respectively; la norm is set as 1e-4. Each model is trained by 620 epochs which consist of five complete cycles. Cutout size (De & Taylor 2017) is set as for both CIFAR-10 and CIFAR-100 datasets.
B.3.2 IMAGENET
We follow a common strategy (Szegedy et al. 2016) for data augmentation. For ResNet (He et al. 2016), we adopt the implementation in tensorflow CNN benchmark") For AA-ResNet, modify the ResNet by augmenting 3x3 convolution with self-attention. specifically, we apply attention augmentation to each residual block in the last stages when the shapes of activation maps become 28x28, 14x14 and 7x7, We refer to Bello et al. (2019) for more details. The implementation is from the official repository and we simply add attention map prediction module the code base. not specified, we adopt the same setting with parameter search for the value of a and the number convolutional layers. The final performance of PA-AA-ResNet is obtained when the number of convolution layer is 2, and a is set to be 0.5.
e.g. k = 2 and 0.2. We conduct hyper-
C ANALYSIS
C.I ANALYSIS ON SNLI
The analysis of different a values and model capacities on SNLI dataset shown in Figure and respectively. All the result numbers are average of 5 different runs with random seeds. We search for multiple layer numbers in (3. 6, 9, 12} and hidden dimension sizes in (256, 384, 512, 768). and find that the following settings are the best among competitors of similar model sizes: (3, 256), (6, 256, 6), (3, 512), (6. 384), (3, 768). Therefore, we take these five configurations to compare Transformer and TA-Trrasformer for different model capacities. From the figures, can draw conclusion similar to the machine translation dataset: the performance is not very sensitive to a small change of and the proposed PA- Transformer model generalizes well for various model capacities.
15

---

Under review as a conference paper at ICLR 2021
0.43
0.5
0.1
0.90
Accuracy
Accuracy
0.07
Transformer
Transformer
044
PA- Transffoorr
PA- PA-Transformer
0.5
Conv- Transformer
0.99
Conv- Transformer
PA- CaaCCoooo--annnnnn
PA-Comv-Transfore
Epachs
Epochs
(a) SST
Yelp
0.85
0.2
Accurscy
BL.B0
0.70
Transformer
Transformer
PA- Transforner
PA- Transformer
0.2
Conv- Transformer
Conv- Transformer
PA Can--ooooTTTnonnee
PA- PP-Coov-Traanforme Tramsformer
Epochs
Epochs
(c) SNLI
(d) IWSLT Do-En
Figure 7: Learning curve comparison
C.2 LEARNING CURVE COMPARISON
Figure compares the learning curves of different models multiple datasets, where the x-axis denotes epoch number and the y-axis denotes the test performance (BLUE score for IWSLT' 14 De-En and accuracy for others). One observe that process (especially for the smallest dataset, SST) and consistently outperforms other baselines at convergence.
stabilizes the training
C.3 ATTENTION MAP VISUALIZATION
In order get insight into the attention prediction mechanism, visualize exemplar attention maps for both text and image inputs and find some interesting evidences.
C.3.1 TEXT ATTENTION
We choose BERT-Base and PA-BERT-Bass models for comparison on the CoLA dataset, a task for judging the grammatical correctness of a sentence. We select the sentence "Mary tried John to go abroad." for visualization. Obviously, this sentence is grammatically wrong, and an effective model should capture the error part "tried John to" in order to give the answer.
htt.ss//tttha.oeoeeeoootttteeeee 11__o_ot...._.a_a__.________ ble
16

---

Under review as a conference paper at ICLR 2021
Mary
tried
tried
John
abroad
abroad
abread
abr bor
(a) BERT #2
(b) PA-BERT #2
(c) Prediction #2
(d) self-Attention #2
tried
tried
tried
tried
to
to
go
go
abroad
abroadl
abroad
abread
<slop>
<alop
<alup>
(c) BERT #11
(f) PA-BERT #11
gg) Prediction #11
(h) Self-Attention #11
Mary
Mary
Mary
Mary
tried
tried
tried
tried
John
John
John
John
00
go
go
go
abroad
abroad
<stop>
kstepo
<stop>>
<stop>
(i) BERT #12
ii PA-BERT #12
(k) Prediction #12
(I) Self-Attention #12
Figure 8: Attention maps of layer #2, #11 and #12 for "Mary tried John to abroad."
In Figure 8 visualize related attention maps for three layers (#2, #11 and #12) in BERT-Bass and PA-BERT-Base models. The second layer is the first layer that utilizes the input of the convo- lutional prediction model, and #11 and #12 the last layers. For each layer, the attention maps from vanilla BERT and PA-BERT in the first and second columns respectively, then the predicted attention maps and supplementary self-attention maps are visualized in the third and fourth column. It should be noted that the second column is the linear fusion result of the third column and the fourth column according to Equation 4
first show
Consider layer #2, both BERT (Figure 8(a)) and PA- BERT (Figure 8(b)) pay major attentions on the verb phrase "go abroad". As shown in Figure 8(b) PA-BERT puts additional stress on the relation between word "tried" and the stop sign. This is reasonable because the stop sign is responsible of capturing sentence-level semantics and "tried" is a key word leading to the grammatical emror. As observe in Figure ((c) the attention on this part actually comes from the convolution-based prediction module, which is somewhat complementary to the self-attention map.
In order ensure that the information obtained by the convolution is beneficial to the task, visualize the last attention layer (#12) which is close to the classification output (see Figure ( i-)). In Figure 8()) layer of transformer. The attention to the wrong phrase "tried John to" is still weak, which directly leads to misclassification result for this case. In contrast, the attention scores between "tried" and
observe that BERT-Base still focuses verbs and stop signs in the very last
17

---

Under review as a conference paper at ICLR 2021
"John" become very high in PA-BERT (Figure 8(). largely owning to the predicted attention map illustrated Figure 8(k)
We also visualize the attention maps of the #11 layer, which serves the imput of the #12 layer. To analysis the evolution attention maps, we compare the difference between Figure 8() and Figure 8(k) the latter is the output of convolutional prediction module by taking the former input. We find that the convolutional prediction module helps to reason about the importance of word "John" based the previous attention input. Specifically, it weaken the attention scores of the correct part and raises higher importance to the wrong part. As illustrated in Figure [8(k) the attention can be clearly seen the upper left corner of the attention map where the error occurs. this way, the error is fully captured in the final representation layer, assisting the model to generate a correct answer.
C.3.2 IMAGE ATTENTION
In Figure we compare the attention maps Con-TTrasffommr and Conv-Trrnsformer for CIFAR- 100 image classification. Compared to Conv-Transformee our proposed method could cap- ture better global information and at the same time emphasize the important local information. Specifically, the self-attention layer prefers extract features from the global perspective, while the predicted attention highlights local features, which assists the self-attention mechanism to depict more accurate outline. As shown by the visualized examples, Conv-Transformer fails compute a explainable attention map. In contrast, with the help of attention prediction modules, PA-Conv- Transformer successfully identifies the objects in images.
dog
horse
Conv- Transfomerr
PA-Cony- Transformffrrrr Trannee-rrn
Predicted Attention
Self- Self-Atten Attention
Conv- Transfomer
PA-Conv-TI PA-Conv-Transen Transforan
Predicted Attention
Self-Att Belf-Attention
Figure 9: Attention maps of Conv- Ton-TTrasformer and Aa-Conv-Transfoon
18


Reviews:
[{"review_id": "qcAYZpe-h7", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nReasons for score: \n \n I like the idea of a chain of attention prediction to learn attention dependencies from the previous block. My major concern is about the clarity of the paper. Hopefully, the authors can address my concern in the rebuttal period. \n\nPros: \n \n1. The paper addresses one of the most important issue of transformer: attention dependencies cross blocks or layers. For me, the problem itself is real and practical. \n2. The proposed predictive attention transformer (PA-Transformer) is novel for capturing the attention dependencies transformer layers and address the problem of the self-attention maps learned independently for each layer. The design for using the PA-Transformer to tasks of NLP and image classification is reasonable and interesting. \n3. This paper provides comprehensive experiments, including both NLP and image classification results, to show the effectiveness of the proposed framework.  \n\nCons: \n \n1. The paper claims that Multi-channel is one of the first works to take attention maps as multi-channel images for explicit modelling in the section of Introduction. It is better to clarity this point and give the difference between Multi-channel and multi-branch on method section.\n2. Why does the paper call the proposed model as predictive attention? From my understanding, it is a type of residual connection for attention. Are the attention results used to predict some kind of tasks in intermediate layers?\n3. Is the source codes available to reproduce the work? \n4. It would be more convincing if the authors can provide a set of experiments about a baseline just using residual connection to bridge the layers, instead of CNN-based attention map prediction module, in the rebuttal period. \n\nQuestions during rebuttal period: \n \nPlease address and clarify the cons above \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "XLTc-sefSjl", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nThis paper's main topic is the enhancement of Transformer models for improving performance in a task-independent way.\n\n\nThis paper first points out that the attention maps are trained independently among layers in the Transformer models.\nThe authors then hypothesize that the performance could be improved by integrating an additional module for estimating attention maps.\nTherefore, they propose a method that yields the attention maps based on the lower layer's attention maps.\n\nThe experimental results on several different datasets from different domains show that the proposed method consistently improves the performance.\nThey are somewhat surprising results.\n\n\nThe following are the questions and concerns of this paper.\n\n\n1,\nThe proposed method seems rather strange; why can the proposed method yield better attention map predictions?\n\nI understand such a phenomenon if we provide the correct attention maps for model training.\nHowever, it seems that the proposed method does not require any additional information on correct attention maps. \nThis is the largest mystery for me about this method.\nPlease elaborate on what architecture or mechanism enables the proposed model to provide better attention maps theoretically or empirically to support the authors' claim?\nIf I did not miss something, there are no clear explanations about it.\n\n\n\n2, \nthe source of the effectiveness:\nThis additional module seems to also work as a sort of skip or short cut connections between layers.\nIn my feeling, the performance gain could be just the direct linking between attention mechanisms in each layer and not be caused by a better attention map prediction.\nAs a recent common knowledge in the community, the correct attention map can be, but not necessarily, a strong correlation to performance.\nPlease reveal the actual source of the performance gain to prove the correctness of the authors' claim.\nOtherwise, there may be a risk of providing the wrong knowledge to the community.\n\n\n\n3, \nRelated to the above two questions, this paper's main concern is that this paper does not provide more in-depth analyses of the proposed method that tell model behaviors or characteristics.\nThere are no intuitive and motivational examples of what kind of situation the proposed method successfully works.\n\n\n\n\n4,\nThe current version does not have the \"Conclusion\" section, which most scientific papers have.\nOf course, there is no rule that the paper always needs to have a Conclusion section.\nHowever, I would like to know why the authors decided not to provide the Conclusion section.\n\n\nI am willing to change my score if I got reasonable answers for all the questions and concerns written in the above reviews.\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "TS8M5EHlDcx", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n**Summary:**\n \nThis 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.\n\n\n**Reasons for score:**\n\nThe idea of bridging attention maps across layers in Transformers is intuitive, as later layers can benefit from dependency structures learned from earlier layers. The model is validated on several benchmarks in both NLP and CV, and show some consistent improvements over baselines. However there are many unclear expressions and claims in the paper, and some ablations are missing which might be critical to better understand the module. Besides the paper definitely needs more proof-reading.\n\n\n**Pros:**\n\n1. The idea of treating multi-headed attention maps as multi-channel images is interesting, and the proposed convolution-based attention prediction module is a natural choice under such settings.\n2. The proposed methods are validated with extensive experiments, and the performance gain is consistent and quite significant. This shows the effectiveness of the proposed approach.\n \n**Cons:**\n\n1. When applying transformer architecture to images, the image of shape $H \\times W \\times C$ is flattened as $X \\in R^{N\\times C}$ where $N=H \\times W$. However for regular images the resolution is quite large, e.g., ImageNet is usually used in 224x224, then the N would be ~50k. In Section 4.2.1 it seems even for CIFAR it will be OOM (out of memory), but this is not mentioned for experiments with ImageNet in Section 4.2.2.\n2. On Page 2 the authors claim that \"... experimental results demonstrate the superiority of PA-Transformer in terms of accuracy, memory cost and computational efficiency ...\" but I was only able to see the accuracy improvements; did I miss the experiments for computational efficiency? If it refers to the #Params and #FLOPs in e.g. Table 2, I'm actually curious: according to Table 1, Transformer and PA-Transformer has <0.01K (if not identical) FLOPs. Could you explain how this is calculated? Because if we count multiplications and additions as FLOPs, I think the PA module will definitely introduce more than 10 (=0.01K) FLOPs. \n3. The authors claim several times (e.g. last sentence of Page 3) that self-attention module could \"dedicate itself to incorporate layer-specific knowledge into *residual* attention maps\". It seems arguable since in most cases the self-attention is dominating the generated attention map ($\\alpha$ is usually small). Also, I'm wondering if the 0-layer PA in Table 5 corresponds to a direct skip-connection, i.e. simply copying the $A_{\\text{pre-logits}}$ over to next layer. In fact I think it's a quite important ablation experiment, e.g. replacing \"predicted attention\" with \"attention from previous layer\", and may worth showing the results on other tasks too (e.g. on CIFAR and ImageNet).\n4. Just curious: the $\\alpha$ seems to have large variations across datasets, e.g. on CIFAR it is set to 0.01 but on SNLI and ImageNet it is set to 0.5. In addition to empirical validation results, is there any explanations for this?\n5. In Table 1 and Table 2 since each experiments are replicated for five times, it's better to show the standard deviation (confidence interval) together with the mean value.\n\n**Questions during rebuttal period:**\n\nI've listed my questions in the cons section, and hopefully can be addressed during the rebuttal.\n  \n\n**Some typos and minor issues:**\n\n-- The phrase \"except for\" is misspelled as \"expect for\" in several places (e.g., Page 5, 7th line of Section 4.1.1 \"Models\", Page 7, 2nd line of Section 4.2.1 Settings). \n-- Multiple typos: \"ResetNet-50\" -> \"ResNet-50\", \"Noe\" -> \"Note\", etc. This paper needs more polishing.", "extra_scores": {}, "extra_text": {}}, {"review_id": "-HeKlmM1KVi", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The idea proposed in the paper is simple - \"predict\" future attention weights using past attention weights. A 2D CNN is used to mix previous N layer's attention maps. Strictly, this is also not \"predicting\" but instead generating\". There is no supervised loss here. \n\nThe authors introduce a PA-Transformer model. The idea is to use previous attention maps to augment future attention weights. A stack of N previous attention weights is modeled with 2D CNN to generate future attention maps. The idea of predicting attention weights (or generating them) is not new (see https://arxiv.org/abs/2005.00743). The difference here is that there is a 2D CNN to model relationships between N previous layers. This is somewhat a pretty incremental extension of the Synthesizer-Transformer model. \n\nThere is also insufficient convincing evidence that using previous layer's attention to generate future attention weights is beneficial. \n\nI think the experiments are lacking. The experiments on GLUE are only comparing against BERT (the least the authors could do is to compare side-by-side with at least a few other models). Machine translation datasets are tiny and ablation studies are unconvincingly run on SST and SNLI. The authors only run experiments using a preloaded checkpoint of BERT and do not apply their architecture to actually pretrain BERT which is also one weakness of this work. Hence, the paragraph beginning with \"pretraining\" is misleading\". The results on GLUE are also weak and could be a result of variance over the existing BERT model. \n\nThe authors should also discuss how this can be implemented in a decoder setting since the current setup will disable causal attention.\n\nOverall, I recommend a clear rejection. I think the key selling point and hypothesis behind this paper (using prev N layer attention) is not well supported. Experimental settings are also weak and there are insufficient convincing experiments to feel that this architecture is doing something useful. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Multiple reviewers point out the interesting improvement to mix attention maps at different layers via convolution based prediction modules. This module is sufficient to show improvements only on encoder side while comparing to concurrent work Synthesizer.
However, the novelty of the work is limited as compared to other papers and the results though improved did not convince the reviewers fully to gain a strong accept.

Decision: Reject