MaiT: integrating spatial locality into image transformers with attention masks 
Authors: Ling Li
Ali Shafiee
Joseph H Hassoun
Venue: ICLR 2022 Submitted (2022)
Keywords: vision transformer
image classification
deep learning
computer vision

Abstract: Though image transformers have shown competitive results with convolutional neural networks in computer vision tasks, lacking inductive biases such as locality still poses problems in terms of model efficiency especially for embedded applications. In this work, we address this issue by introducing attention masks to incorporate spatial locality into self-attention heads of transformers. Local dependencies are captured with masked attention heads along with global dependencies captured by original unmasked attention heads. With Masked attention image Transformer – MaiT, top-1 accuracy increases by up to 1.0\% compared to DeiT, without extra parameters, computation, or external training data. Moreover, attention masks regulate the training of attention maps, which facilitates the convergence and improves the accuracy of deeper transformers. Masked attention heads guide the model to focus on local information in early layers and promote diverse attention maps in latter layers. Deep MaiT improves the top-1 accuracy by up to 1.5\% compared to CaiT with fewer parameters and less FLOPs. Encoding locality with attention masks requires no extra parameter or structural change, and thus it can be combined with other techniques for further improvement in vision transformers.

Under review as a conference paper at ICLR 2022
MAIT: INTEGRATING SPATIAL LOCALITY INTO IMAGE TRANSFORMERS WITH ATTENTION MASKS
Anonymous authors
Paper under double-blind review
ABSTRACT
Though image transformers have shown competitive results with convolutional neural networks in computer vision tasks, lacking inductive biases such as lo- cality still poses problems in terms of model efficiency especially for embedded applications. In this work, we address this issue by introducing attention masks incorporate spatial locality into self-attention heads. Local dependencies are captured with masked attention heads along with global dependencies captured by original unmasked attention heads. With Masked attention image Transformer Mai'T, top-1 accuracy increases by up to 2.5% compared to Dei'T, without extra parameters/FLOPS. or external training data. Moreover, attention masks regulate the training of attention maps, which facilitates the convergence and improves the accuracy of deeper transformers. Masked attention heads guide the model to fo- on local information in early layers and promote diverse attention maps in latter layers. Deep MaiT improves the top-1 accuracy by up to 1.7% compared to CaiT with fewer parameters and FLOPs. Encoding locality with attention masks requires no extra parameter or structural change, and thus it can be combined with other techniques for further improvement in vision transformers.
INTRODUCTION
Convolutional neural networks (CNNs) (Krizhevsky et al. 2012 He et al., 2016 Tan & Le 2019) have been the de facto model for computer vision (CV) tasks, which are inherently equipped with inductive biases such as translation equivariance and locality. Recently, vision transformers (Doso- vitskiy al. 202T, Touvron et al. (021a) are gaining momentum translating the success from Tranformer--aseee models in natural language processing (NLP) tasks (Vaswani et al.
cct al. 2019). Self-attention heads excel at capturing the long-range dependencies in sequences but struggle at focusing local information. Unlike CNNs which have gone through many iterations of optimization, vision transformers are not very efficient and still require huge computing power and large number of Flops.
Naturally, combining the benefits of both CNNs and vision transformers is promising to further boost performances of CV models. The question remains how effectively integrate inductive bias such as spatial locality into transformers. One direction is to utilize convolutional blocks to extract spatial information by adapting either the patch-token embedding layer, self-attention module or feed-forward layers, to form -NN-trranformmr hybrid structures as in in et al. (2021b); (rinivaas et al.)
(2211: Graham et al. (2021); D. Wu al.) (2021a): Yuan al. (2021a): Guo et al. (2021). However, forcefully inserting convolutional operations into transformers may potentially constrain the learning capacity of transformers.
To capture the spatial information without significantly changing the transformer model architecture,
Chu et al. (2021b) introduce extra positional encoding. [Han al. (202T); Chen et al. (202)) fuse local and global representations using multiple transformer blocks branches simultancously process images at different scales such as pixel-level, small-patch or large patch. Yuan et al. (002
transformation to capture local structure. These approaches ally come with the cost of extra parameters and model complexity, thus potentially lowering the inference speed.
apply layerwise
d' A'Ascoli et al. (2021): Zhou et al. 2021 improve the self-attention for better representation with gated positional self-attention and learnable transformation matrix respectively. Hu al. (2019)

---

Under review as a conference paper at ICLR 2022
Ramachandran et al. (2019) adapt the self-attention module and improve the performance of CNN models. Liu et al. (2021) captures locality within shifted windows in a hierarchical structure. Though it could save computation in some cases, the added model complexity may lower its potentials.
Different from prior works, we try to incorporate spatial locality without changing its architecture or adding extra parameters/FLLOPs. We propose attention masks guide the attention heads to focus local information. Masked attention heads extract local dependencies more efficiently by allowing information aggregation only from the closest neighbors. This liberates other unmasked heads to learn global information more effectively. We name the modified model, Masked attention image Transformer (MaiT), which is built on top of DeiT (Touvron et al., 2021a). MaiT gathers both local and global information at the same time from different heads. Moreover, the regularization effects from attention masks facilitate the training of deep transformers by guiding the attention map learning and promoting diversity across transformer layers.
We proved that less is more in this specific case with attention masks. MaiT achieves up to 2.5% higher top-1 accuracy ImageNet (Deng et al.) (009) with the same model architecture DeiT TTovrro et al. 202 Additionally, deep MaiT outperforms CaiT (Touvron et al. 2021 by up to 1.7% in top-1 accuracy with fewer parameters and simpler structure.
In summary, we make three major contributions this work: 1). we propose attention masks to encode the spatial locality into self-attention heads without structural change or extra parame- ter/computation, while improving model efficiency. 2). We present a quick empirical search strat- egy for the masking scheme exploration, along with automatic end-to-end search alternative. 3)- We also reveal the importance of the locality across layers and the performance impact of atten- tion masks on deep vision transformers. Note though MaiT is demonstrated with DeiT/CaiT, the attention mask is applicable to other vision transformers as well.
2 RELATED WORK
Spatial locality an integral part of the convolutional operation with weight filters attending to local regions of input feature maps. Vision transformer (ViT) (Dosovitskiy et al. (021) is the first pure transformer- based model on vision tasks, but it requires a large private labeled dataset JFT300M (Sun et al.] (017) to achieve competitive performances. Data-efficient image transformer (DeiT) (Touvron et al. 2021a) improves upon ViT models by introducing stronger data augmentation, reg- ularization, and knowledge distillation. Class-attention in image transformer (CaiT) (Touvron et al.) 202Tb) extends DeiT by increasing the number of transformer layers. To overcome the difficulties of training deeper transformers, CaiT introduces LayerScale and class-attention layers, which increase the parameters and model complexity.
Tokens-to-Token vision transformer (T2T) (uaan et al. 2021b) proposes an image transformation by recursive token aggregation to capture local structure. Stand-alona self-attention (Ramachandran et al. 2019) applies local self- layer to replace spatial convolution and outperform original ResNet models. Even though sharing value and key spatially is parameter efficient in this approach, content-based information is lost. Transformer-iN-Trrnstooo
(TNT) (Han et al., 2021) models both patch-level and pixel-level representations and applies outer and inner transformer blocks to extract global and local information respectively. Con (d' Ascoli al. (22I]) proposes the gated positional self-attention to incorporate soft convolutional biases. Cross ViT (Chen et al. (021 pro-
poses a
transformer architecture for multi-scale feature extraction.
For pixel-level prediction tasks such as semantic segmentation, object detection, Pyramid Vision Transformer (PVT) (Wang et al., 2021a) introduces progressive shrinking pyramid and spatial- reduction attention with fine-grained image patches. DETR (Carion et al.) 2020) adapts transformers for object detection tasks. Swin Transformer (Liu et al. 2021 applies hierarchical transformer with shifted windows of varying sizes. Twins (Chu et al. 202 deploys interleaved locally-grouph self-attention and global sub-sample attention layers to improve performances.
To optimize transformer and save computation, Wu et al. (202 uses centroid attention extract and compress input information, Jaegle et al. (2021) > iteratively distill inputs into latent space with attention bottlenecks, WWang al. (2021b) dynamically adjusts the number of tokens with multi- ple cascading transformers, and Wu et al. (2020) introduced semantic token to replace pixel-based transformers to save computation.
2

---

Under review as a conference paper at ICLR 2022
There are hybrid architectures fusing convolutional and transformer blocks, such Lical (Li et al. 2021b), BoTNet (Srinivas et al.] 2021). LeViT (Graham et al. (021). BossNet (Li et al.
202 CvT (Wuet al. 2021a), CoaT (Xu et al. 202I). CMT (Guo et al. 202 for higher accuracy
and faster inference.
Unlike prior literature, our work explores the intrinsic capability of pure transformer block incor- porating spatial locality and the impact of the locality along the depth direction. This work is also inspired by the emerging graph attention network (Veličković et al. 2018), borrowing the concept of message passing and information aggregation from nearest neighbors.
Note attention masks have been used in NLP tasks sparsification method reduce the computa- tion complexity, as well as to capture local information, as in Guo et al. (201)); Child et al. (2019); Beltagy al. (2020); AAinsli et al. (2020).
3 VISION TRANSFORMER PRELIMINARIES
The transformer architecture introduced by Vaswani (Vaswani et al. 2017) inspired many model variants with remarkable success in NLP tasks. ViT (Dosovitskiy et al., 2021) extends pure transformer- based architecture into CV applications. Instead of pixel-level processing, ViT splits the original images into sequence of patches inputs and transforms them into patch tokens, for better computation efficiency. In general, ViT consists of 3 fundamental modules: embedding layer, multi-head self self-attention, and feed-forward network.
To process images in transformer, the original RGB images (224x224) is flattened into a sequence of N (14x 14) patches. Each patch has fixed size (typical 16x pixels). Patches are then transformed into patch embedding with hidden dimensions (D) of 192, 384, and 768 for tiny, small, and base models respectively ViT/DeiT. In addition to patch tokens, the embedding layer also integrates positional information, classification and knowledge distillation through the positional token, class token and distillation token, respectively.
Positional token is added into the patch embedding with trainable positional embedding. However, this positional embedding added only in the embedding layers. The spatial information largely lost in the transformer layers since all-to-all attention is invariant to the order of the patches.
The class token another trainable vector (1 x D), concatenated the patch tokens (total N+1). is used to collect information from the patch tokens to make output predictions, while also spreading information among patch tokens during training.
Distillation token is sometimes added for knowledge transfer from teacher models, such as a CNN model. When training the distilled version the model, a distillation token is further concatenated to the patch token along with the class token (total N+2).
Multi head self-attention (MHA) module has multiple parallel attention heads, where each of them comprises three main components: Key (K), Query (Q), and Value (V), Key and Query are trained and multiplied to estimate how much weights on each corresponding token in Value for output:
QKT Vd
Attention(K: Q.V) = Softmax
V
(1)
Where softmax is applied to each row of the input product matrix (QKT) and vd provides ap- propriate normalization. Multiple attention heads in MHA attend different parts of the input simultancously. Considering heads in MHA layer, the hidden dimension D is split equally across all heads (D- H d).
Feed-forward network (FFN) follows after the MHA module, containing two linear transformation layers separated by Gelu activation. The hidden dimension expands by 4x after the first linear layer from D to 4D, and is reduced back to D in the second linear layer. Both MHA and FFN use skip-connections with layer normalization the residual operation.
3

---

Under review as a conference paper at ICLR 2022
4 MASKED ATTENTION HEAD
Spatial locality plays crucial role in computer vision tasks. CNN models capture it using the sliding filter of shared weights, typically with a receptive field of 3x3, 5x5, or 7x7. contrast CNN models, the locality is not explicitly introduced in the transformer structure. With attention masks, we can explicitly insert locality into self-attention modules without introducing any extra parameter computation. The key idea is to apply mask the all-to-all attention products (i.e. QK") to reinforce the weights from the closest neighbors by aggregating information only from tokens selected by the mask.
22 16
10
11
13
24
124
Atlention mask for Putchlis
3x5 (eepth =1]
47
48
51
52
53
3G
36
57
58
60
61
G2
64
63
66
67
57
5x5 (dupth
192
Figure 1: Attention mask with depth of 1 (orange) and 2 (green).
Figure illustrates an example of our proposed masking scheme. The orange box shows a 3 x mask, where only the direct neighboring patches selected. Specifically, Patch 16 only gathers information from the closest neighbors of Patch 1, 2, 3, 15, 17, 29, 30, and 31, and ignores the rest of patches. This is different from the typical all-to-all attention module, where Patch 16 attends to all 0-195 patches. We can easily expand the depth of the attention mask beyond the closest neighbors, second-level neighbors (green box in Figure IT more. Note that the class token (and distillation token) still attends to all the patches to collect and spread information during forward and backward passes. Since each attention product selected by the mask is calculated by Q and K, the masked attention head also retains the content-based locality information.
The attention mask added before Softmax, regulating the distribution of attention maps to focus more the closest neighbors:
M MOQKT QK
Masked Attention K, K,Q.V) Softmar
V
vd
where M {0, (N++) is binary attention mask, encoding the spatial locality into the attention head by passing through only the weights from close neighbors and setting the rest to zero.
Note it's important to add the mask before Softmax because it allows the model learn the impor- tance of the locality flexibly. More precisely, unselected patches appear as e = in the numerator of the softmax operation. Thus, if the attention product result of the closest neighbors is meaning- fully larger than zero (i.e. M C QKT > 0), it suggests that local information dominates. However, those results are negative or close zero, it implies that local information insignificant and global information is more important. Therefore, inserting masks before the softmax operation al- lows models enforce locality or bypass
4.1 ATTENTION LOCALITY SCORE (ALS)
The softmax operation transfers the results of QK" into the probability space. As a result, each row of the attention map (A) sums up 1. For each patch, the probability of focusing local neighbors equals the sum of their neighbors' attention map weights. this number approaches one, the local information is crucial: whereas if it is close zero, it means global information matters more. For
4

---

Under review as a conference paper at ICLR 2022
Patch n. define ALS,, the Attention Locality Score (ALS). The mean of all N patches is the ALS for each attention head, as follow:
ALS En ALS, N
N+1 where ALS, Σ (M A)n.i
(3)
where M is the same attention mask defined earlier, and A = Softmax (MOQKT/Va) is the attention map, or A = Softmax (QKT/Vá for unmasked heads, and n, denotes the row and column index the attention map respectively. We later use the ALS metric to get some insights about the locality behavior of different attention heads at various layers in our models.
4.2 MASKING STRATEGY
With total H attention heads, 1 number of attention heads can be allocated to focus on local infor- mation. The rest of the unaltered attention heads (H - h') capture global dependencies. Therefore, we can extract the local information through masked attention heads and global information through original attention heads at the same time, shown in Figure 2
Masked attention head: local dependencies
Original attention head: global dependencies
Parallel In' times
Parallel (H ht] times
Output
Output
Attention mask
Softmmax
MaskOak
OKT
Query
Key
Value
Query
Key
Value
Figure 2:
Masked attention heads and original attention heads for local and global dependencies.
Besides the depth of the attention mask, the number of masked attention heads MHA and the position (which layer to insert mask) are also hyper-parameters. Though masks encode locality into the attention heads, the regularization from masking (pruning attention map) can also limit the learning capacity. Consequently, where to add attention masks requires careful consideration. The complexity to search for the best masking strategy is up to 236 for 12-layer model with heads. Thus for quick empirical search to obtain decent mask placement within 2-3 training iterations, we leverage the insights from ALS study in Section 5.4 and follow a masking strategy below:
1. Initialization: add mask to only one attention head (Head 0) for all layers to explicitly introduce locality in every layer (forming a local path from the beginning to the end)
2. Assignment:
compute A A for every layer
(a) AL.So.l is close to I: masks to attention heads layer proceed
(b) If o.. is close 0.5: keep the mask, add no more
(c) AL is close O: remove the mask
3. Calibration: following 2.a), compute ALSh. l for each head and remove the mask if ALShI is close 0
Note that optimal masking is a large space as opposed to a single unique placement. The model is capable of adjusting different masking schemes flexibly. Step 2 (a) relies the observations that early layers (typically with ALS close to 1) focus heavily on local features and thus more masked heads are beneficial. Step 3 will remove the redundant mask accordingly.
4.3 SOFT MASKING
To alleviate the complicated searching problem and automatically lear the mask placement, we introduce learnable scale factor E (0, 1) for each attention head. We name this approach Soft
5

---

Under review as a conference paper at ICLR 2022
Masking, which replaces the zeros in the original masks with scale factors and keeps the ones. The binary masking approach is named Hard Masking.
Mn = {1
if is n's close neighbor jf is not n's close neighbor vs. Ma
if is n's close neighbor
(4)
if is not n's close neighbor where Maj- Mn is the n-th row and j-th column element in the hard and soft mask respectively.
This scale factor penalizes the attention weights from non-neighboring patches. It allows non-local patch tokens to contribute differently. When approaches O, this attention head focuses more on local information. Otherwise when it's close to 1, the attention head attends global information. As result, each attention head at every layer is able to flexibly determine the importance of locality. Note, this introduces negligible extra parameters during training. For example, it only introduces 36 extra parameters for 12-layer MaiT tiny model with 3 heads.
5 EXPERIMENTS
In this section, we report the results for MaiT, implemented by applying attention masks on top of DeiT. We evaluate the impact of different hyper parameters related the masking scheme including the mask depth, number of masked heads per layer, and number of transformer layers. We also analyze the importance of locality in the depth direction and the performance impact of masks for deep transformers. We merely focus on models with fewer than 50M parameters, which are more applicable embedded systems and mobile devices.
5.1 SETUP
We follow the same training procedure as DeiT, unless specified otherwise. The implementation based on the Timm library (Wightman/ 2019). Models are trained with ILSVRC-2012 ImageNet dataset (Deng et al. 2009) with the default batch size of 1024 on 4 GPUs for 300 epochs for 12-layer models and 400 epochs for models with more than 12 layers. The model parameters we adopted summarized in [T
Table I: Details of MaiT model parameters
Model
Extra tiny (XT)
Tiny (T)
Extra small (XS)
Small (S)
Hidden dim.
144 3
192
288
384
Heads
3
6
5.2 IMPACT OF MASK DEPTH
Intuitively, the field of view is larger with attention mask of larger depth, gather broader neighborhood information. However, in the extreme case with mask depth of N+1 (standard self- attention), all the patch token information is merged and positional information is lost. Nevertheless, considering the pruning effect from the mask, smaller depth for the attention mask potentially translates to fewer FLOPs, which might be useful for some hardware devices such as embedded CPUs. For example, with the 3x3 mask, we only compute 9 out of all (N+1) attention weights for each token. When N=196, this saves 95% of attention map computation.
To study the impact of the mask depth, we mask only one attention head for every MHA module, with the field of view as 3x3, 3x5, and 5x5 on tiny and small DeiT models. The results are summarized in Table 2 There is 0.1 0.3 % improvement compared to DeiT small. Various mask sizes only lead to marginal differences in accuracy. This is probably because each patch (original 16x16 pixels) is processed sub-image and already contains some locality information, and thus further away patch tokens barely provide much extra information. Therefore, 3x3 mask is the default setting for MaiT in the subsequent sections.
In contrast to DeiT, whose accuracy saturates beyond 300 epochs as reported in (Tourron et al. 202 Mmillsmall another 1.1% improvement when trained for 400 epochs. Accordingly,

---

Under review as a conference paper at ICLR 2022
Table 2: Top-1 accuracy ImageNet for DeiT and MaiT (one masked head for all 12 layers) with various attention mask depth. MaiT* adopts soft masking and is trained for 400 epochs.
Model
DeiT
MaiT((x3)
MaiT(3x5)
MaiT+ MaiT(5x5)
Tiny
72.2
72.6(+0.4)
72.2(-)
72.1(-0.1)
74.7(+2.5)
Small
79.8
79.9(+0.1)
80.1(+0.3)
80.0(+0.2)
80.9(+1.1)
suspect higher learning rate and lower drop path rate are beneficial for MaiT-tiny since attention masks already provide some regulation during training. Changing the batch size from 1024 to 3072 only results in small gain of 0.4% for MaiT-tiny with 3x3 mask. In comparison, adopting drop path rate and 2x learning rate, we obtain 2.5% improvement over DeïT-T.
5.3 IMPACT MASKING STRATEGY
We compare the performances of 24-layer MaiT models with the following three representative masking schemes (Sch.), to quantify the impact of the masking strategy. Sch. I is a naive way to combine local and global information every layer. Sch.2 results from the masking strategy in Section 4.2. Sch.3 is based on the soft masking.
Sch.1: :: one masked head for all 24 layers
Sch.2: Masked heads for Layer 0-7 is H-1, Layer 9-19 is one, and Layer 20-23 is none
Sch.3: Soft masking for Layer 0-20, no mask for Layer 21-23
Table 3:
Top-1 accuracy on ImageNet for 24-layer DeiT and MaiT with different masking schemes. 1ch.1 adopts only one masked head. Sch.2 applies a mixed masking and Sch.3 uses soft masking.
DeîT
MaiT(Sch.1) MaiT(Sch.2) MaiT(Sch.3)
Tiny Extra Small
78.3
79.1
79.1
79.3
81.0
81.4
81.3
As shown in Table 3 top-1 accuracy improves by 0.8% simply through adopting one masked head for MaïT-T. Another 0.4% improvement is achieved with mixed masking for MaiT-XS. Soft masking further increases the accuracy by 0.2% for MaïT-T. Note that the initialization of the scale factor for MaiT-XS subject to further tuning.
Empirically find Sch.2 yields relative good performances by masking heavily in the first third of layers to extract local features and only one mask for the middle layers. Note that there is a tradeoff between computation complexity and accuracy with hard masking and soft masking. Hard masking can translate to linear increase for attention map computation with respect to the number of tokens N, or O(N), if hardware leverage the structured sparsity. This is especially beneficial with high resolution images with large number of tokens. The computation complexity remains quadratic ( for soft masking.
5.4 IMPORTANCE OF SPATIAL LOCALITY
Attention Locality Score (AI.S) provides quantitative metric to study the importance of spatial locality for each head across all layers. Figure 3) illustrates the result of a 24-layer MaiT-TTny with one masked head. ALS the masked head is significantly higher than other heads Layer 0-22. This indicates the model learns to focus more local information with the masked attention head. ALS gradually declines from around 90% for the first 8 layers down to around 60% and below 5% for the last layer. Note ALS of the last layer is below 0.0051 (1/(N+1) 1/197), average score if the attention weights are distributed uniformly. This suggests that spatial locality largely ignored in the last layer where global information is more important.

---

Under review as a conference paper at ICLR 2022
Similarly, when learning the masking automatically through soft masking (Figure3 b)), the first layers show over 0.9 ALS for multiple heads and layers. ALS gradually decreases in the middle layer with at most one head focusing heavily local features. ALS is below 0.3 for the last 10 layers, where all heads tend to global information.
Overall, the spatial locality is more important in the early layers, and its importance gradually de- clines in deeper layer, which is consistent with d' Ascoli et al. (2021) Raghu et al. (2021)
MaiT-Tiny /witterd Masking) ALS
MaiT-Tiny (w/Soft Masking)
1.0
Masked HeadD
Masked Heado
Headl
Masked Headl
0.8
Head2
0.8
Masked Head2
0.5
0.6
0.4
0.4
0.2
0.2
0.0
0
12
36
20
24
36
20
Block Index
Block index
Figure 3: Attention Locality Scores (AI.)) for MaiT-Tiny: a) one masked head for all 24 layers; b)
soft masking for all 24 layers
5.5 DEEPER TRANSFORMERS
Another difficulty is train deeper models because naively stacking transformer layers fails to deliver the expected performance gain as shown in Touvron et al.) (2021b). One reason is attention collapse, i.e. the attention maps mixed masking is able to break the structural repetition in deep transformers and promotes diversity later layers.
alike among deeper layers (Zhou et al. (02T). We find
MaiT (tiny) depth scalability
mixed masking
80
accuracy
random masking
76
deeper MaiT deeper DeiT CalT
72
12
20
24
28
32
36
of transformer blocks
Figure 4: Top-1 accuracy on ImageNet for DeiT (red), MaiT (green), CaiT (yellow) and randomly masked DeiT (blue) with various numbers of transformer blocks. Green square MaiT applies one mask heads for all layers while green star MaiT masks all three heads for the first 12 layers and only one masked attention head for the rest 24 layers.
Performances of DeiT start to saturate after reaching certain depth. For example, increasing the number layers from 24 to 36, top-1 accuracy only increases by 0.25% for DeiT-tiny (Figure 4. In contrast, MaiT with mixed masking boosts the accuracy to 79.8%, extending the performance gain with deeper transformers. Interestingly, noticed that for 24-la model, CaiT is even worse than naively stacked DeiT-24, suggesting that the pareee-rrameters for 24-layer CaiTT-tiy are probably not optimal.
8

---

Under review as a conference paper at ICLR 2022
To rule out the pure pruning effects attention maps from masking, also apply random mask of the same drop-out rate to one head all 24 layers. The randomly masked model leads to 0.5% accuracy drop compared to 24-layer DeiT. This provides further proof that the performance gain with attention masks is indeed from better spatial locality aggregation.
To better understand the impact of attention masks in deep MaiT, compare the cross layer simi- larity heat map among 36 layers of DeïT, MaiT (one masked head), and MaïT (mixed masking) in Figure 7 Figure [8] and Figure 9 respectively. Effective masking lowers the cross-layer similarity and thus improves scalability of MaïT.
5.6 COMPARISON WITH OTHER STATE OF THE ART VISION TRANSFORMERS
Attention mask can be easily combined with other optimization techniques. For deep transform- ers, we adopt the LayerScale (Touvron et al. 2021b) to MaiT as MaiT*. As shown in Table 4 MaiT/MaiT obtain competitive performances across different model sizes. Specifically, MaiT-T outperforms CaiT-T by 1.7%. MaiT-S achieves 0.2% higher accuracy than CaiT-S with less param- eters/FLOPs and simpler structure, which translates to higher throughput. Moreover, simply adding the attention mask to CaiT leads to 0.5% and 0.2% higher accuracy for tiny and smaller models respectively.
Table 4: Top-1 accuracy for deep MaiT and some SOTA vision transformer models. AAdded Layer-
Scale on MaiT.
Model
Layers
Params(M)
GFLOPs
Top-1 Acc.
LocalViT-T (Li et
2021b)
5.9
1.3
74.8
MaiT-XT
24
6.3
1.5
75.5
PVT-T (Wang et al.
2021a)
13.2
1.9
75.1
CaiT-T Touvron
2021b
26
12
2.5
77.6
Masked Call
26
12
2.5
78.1
MaiT-T
24
11
2.5
79.3
PVT-S Waan et al.
2021a)
15
24.5
3.8
79.8
LocalViT-S Li et al.
(202Tb)
22.4
4.6
80.8
Swin-T (Liu et al. 2021)
12
29
4.5
81.3
T2T- 1:T-4 Yuan et al.
b021b
21.5
5.2
81.5
Twins-SVT-S (Chu et al. 2021a
18
24
2.8
81.7
CaiT--XS Touvron et al.
20216
26
26.6
5.4
81.8
MaiT-XS#
24
24.5
5.3
81.8
PVT-M (Wang et al. 2021a)
27
44.2
6.7
81.2
VIT-24 Yuan et 2021b
64.1
14.1
82.3
CaiT-S Touvron et al. 2021b)
26
46.9
9,4
82.7
Masked Carl-S
46.9
9.4
82.9
MaiT-S*
24
43.3
9.1
82,9
Swin-S Liu ct al.
2021
50
8.7
83.0
Twins-SVI-E Chu et al. 2021
24
56
8.3
83.2
CONCLUSION
In this work, we incorporate spatial locality into vision transformers by inserting masks to attention heads. Masked heads are able to focus on local information and liberate other unmasked heads to extract global information more effectively. We also introduce attention locality score to guide the mask search and rapidly evaluate various masking schemes. Attention masking is simple and effective technique, especially for small and deep transformer models. We observe that attention masks also serve regularizer to guide the training of attention maps for deeper transformers. Moreover, it is promising to further boost performances by combining this attention mask technique with other optimization approaches.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Joshua Ainslic, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Ete: Encoding long and structured inputs in transformers, 2020.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers, 2020.
Chun-Fu Chen, Quanfu Fan, and Rameswar Panda. Crossvit: Cross-attention multi-scale vision transformer for image classification, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers, 2019.
Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Twins: Revisiting the design of spatial attention in vision transformers, 2021a.
Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Conditional positional encodings for vision transformers, 202
Stéphane d' Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. NAACL, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x 16 words: Transformers for image recognition at scale. ICLR, 2021.
Ben Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, and Matthijs Douze. Levit: vision transformer in convnet's clothing for faster inference, 2021.
Jianyuan Guo, Kai Han, Han Wu, Chang Xu, Yehui Tang, Chunjing Xu, and Yunhe Wang. Cmt: Convolutional neural networks meet vision transformers, 2021.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer, 2019.
Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang. Transformer in transformer, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. CVPR, pp. 770-778, 2016.
Han Hu, Zheng Zhang, Zhenda Xie, and Stephen Lin. Local relation networks for image recognition, 2019.
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention, 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. NIPS, 1097-11055 2012.
Changlin Li, Tao Tang. Guangrun Wang, Jiefeng Peng, Bing Wang, Xiaodan Liang, and Xiaojun Chang. Bossnas: Exploring hybrid enn transformers with block-wisely self-supervised neural architecture search, 2021a.
10

---

Under review as a conference paper at ICLR 2022
Yawei Li, Kai Zhang, Jiezhang Cao, Radu Timofte, and Luc Van Gool. Localvit: Bringing locality to vision transformers, 2021b.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows, 2021.
Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks?, 2021.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens. Stand-alone self-attention in vision models. Advances in Neural Information Processing Systems, 32, 2019.
Aravind Srinivas, Tsung- Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck transformers for visual recognition, 2021.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable ef- fectiveness of data in deep learning era, 2017.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In ICML, pp. 6105-6114. PMLR, 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In ICML, volume 139, pp. 10347-10357, July 2021a.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. arXiv preprint arXiv:2 170.399,, 2021b.
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 in neural information processing systems, pp. 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks, 2018
Wenhai Wang. Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions, 2021a.
Yulin Wang, Rui Huang, Shiji Song, Zeyi Huang, and Gao Huang.
Not all images are worth 16x16 words: Dynamic vision transformers with adaptive sequence length, 2021b.
Ross Wightman. Pytorch image models. https://github..om http://www...m/mm/igtt.t.ttttt
pytorchhimage-modee
2019.
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision, 2020.
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions to vision transformers, 2021a.
Lemeng Wu, Xingchao Liu, and Qiang Liu. Centroid transformers: Learning abstract with atten- tion, 2021b.
Weijian Xu, Yifan Xu, Tyler Chang, and Zhuowen Tu. Co-scale conv-attentional image transform- ers, 2021.
Kun Yuan, Shaopeng Guo, Ziwei Liu, Aojun Zhou, Fengwei Yu, and Wei Wu. Incorporating con- volution designs into visual transformers, 2021a.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet, 2021b.
Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer, 2021.
11

---

Under review as a conference paper at ICLR 2022
A APPENDIX
A.1 HYPER PARAMETERS FOR MAIT
Table 5: Hyper parameters for 24-layer MaiT in Table 4 drop rate, the same in Touvron et al. (2021a). *denotes the constant drop out for all layers in Touvron et al. 2021b)
Drop rate refers to linear stochastic depth
Model
Drop rate
Batch size
Learning rate
MaïT-XT
0.0
1024
0,001 x xbatchsize/5 12
MaïT-T
0.05
1024
0.001 x batthsizze5112
MaïT-XS
0.1
1024
0,0005 batchsize/512
MaïT-S
0.1*
1024
0.0005 batchsize/512
A.2 ALS OF DEIT AND DEEP MAIT
DeiT-Tiny ALS
Head
0.8
Head
0.4
12
24
Figure 5: ALS for 24-layer DeiT tiny.
When trained without restriction from attention masks, DeiT emphasizes the locality in the first 8 layers with over 0.8 ALS for some heads, as shown in Figure [5] In comparison, Figure MaiT utilizes masked Head 0 to extract local information and ALS lower for unmasked Head 1 and Head 2.
MalT (diny) Head
1.0
12 layers
0.9
layors
0.0
36 layers
0.6
5
0.5
0.4
0.3
0.2
0.1
0.0
12
20
28
36
Block Index
Figure 6: ALS of deeper MaiT-tiny with 12, 24 36 layers.
12

---

Under review as a conference paper at ICLR 2022
When the number of transformer blocks increases, the importance of locality expands to deeper layers as shown in Figure 6 For the masked Head 0 in MaiT-tiny, 12-layer model shows more than 0.8 ALS for 6 layers out of the first 7 layers, the following 4 layers give ALS around 0.5, and the ALS of the last layer declines below 0.05. Likewise, for 24-layer and 36 layer models, 10 layers out of the first 12 layers show ALS higher than 0.8, and decreases 0.5 or 0.6 in the following layers before the last layer.
A.3 CROSS-LAYER SIMILARITY
To evaluate the impact of attention masks on the diversity of the attention maps across all layers, apply a cross-layer similarity metric similarly defined Deep ViT (Zhou et al. (001]:
ST
h.t.
B.
(5)
where M J is the attention map cosine similarity between layer and layer j for attention head h and token t. Ar.t.: measures the weight contribution for each token in input (Value) for output token T. Therefore, reflects the similarity of attention maps on information aggregation across all T input tokens. For example, if M reaches output token head attends all N+1 tokens in Value with the same probability for both layer i and layerj.
DeiT shows the highest cross-layer similarity at the last 8 stages, with an average of 0.75 (Figure This homogeneity
attention maps the late stage limits the model learning capability. Adding mask to one head for all 36 layers does not alleviate this problem since the same structure is repeated 36 times in the model. As result, the average cross-layer similarity for the last layers with unmasked heads is also 0.75 (Figure 8). Alternatively. if we mask all heads in the first 8 layers, they are naturally different from the rest of the partially masked layers, and the same structure only repeated 8 or 24 times. As shown in Figure 9 cross-layer similarity between the first 8 layers and the following 24 layer are mostly well below 0.3 for Head and 2. Additionally, a mixed masking scheme lower the cross-layer similarity among the last 8 layers by 0.1 to 0.65.
crozs layer larrtyy map Head 1
Head
Head 2
Boor nod wew. 2ac A
0.8
0.7
0.6
002
-0
Index
Figure 7: Attention map similarity across 36 layers for DeiT-tiny
In Figure 7 36-layer DeiT-tiny shows the highest similarity at the last 8 stages across all three heads, with an average similarity of 0.75, whereas the similarity is quite small among the first 16 layers and between the first 16 layers and the last 20 layers. This agreement with Zhou et al. 202I), indicating attention collapse at later stage of deep transformers.
In comparison, for MaiT with one masked head across all layers, the masked attention Head shows more uniform distribution with lower peak intensity compared to the the non-masked or partially masked (in early layers) Head I and Head (Figure For unmasked Head and Head 2, later stage especially the last 8 layers show highest similarity among each other, the average of them is also 0.75. This suggests one masked head for all transformer blocks is not as effective in 36-layer model in 24yle model. This indicates repeating the same structure for 36 times in the
13

---

Under review as a conference paper at ICLR 2022
ayer similarity map Heal
Head
Head
300 34 35 328 2032 2032 3333
0.0
0.8
Block Indos
0.4
003
62
61
12
2628
16,00 70 24.24
111-822077 2628 Index
Block
Figure 8: Attention map similarity across 36 layers for MaiT-tiny
MaiT tmiixed masking) cross layer similarity map
Heac
Head
0 6 8 301223 30 2022222 2028 20 3233
0.8
0.7
Block roaex
06
02
01
263820
Figure 9: Attention map similarity across 36 layers for MaiT-tiny with mixed masking scheme
depth direction is causing the attention collapse. The diversity among the heads within the same layer doesn't alleviate this issue.
14


Reviews:
[{"review_id": "uK2CQLY3HS", "reviewer": "Reviewer_71Sc", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 5, "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": "**Strength**: \n\nIncorporating local information efficiently into global self-attention is a promising research direction.\n\n**Weakness**:\n\n1. Writing is very poor, which makes it hard for me to understand what the authors want to express. \n\n2. Notations are frustratingly messy and I can only list a small part of mistakes here. For example, at the beginning of Section 3, this paper fixes the size of input images, patch size, hidden dimensions, etc. The authors should define symbols clearly and specify the configurations in the experiments. \nIn Eq. (1), what are the definitions of Q, K, V?  In Eq. (3), what’s the meaning of index i? I guess it does not represent the layer index here? But in Section 4.2, i indicates the layer index. \n\n3. The masking strategy, which is the core component of the paper, has serious technical problems.\n    - The paper proposes a heuristic strategy to add masks. However, in Transformers, different heads actually encode distinct patterns (e.g., [A1]). However, in this paper, all the heads assume to share the same mask which is not reasonable. For 2 (a), why add masks to all attention heads? For 2(b), if the score is around 0.5, why only add a mask to the 0-th head? More explanations are needed. \n    - The paper analyzes the search space but no search algorithm is proposed. To automatically search the masking strategy is important and it is not OK to leave it for future work.\n\n4. The advantage of the proposed masked attention over the existing efficient Transformers with local-global self-attention (e.g., PVT [Wang el al., 2021a], Swin [Liu et al., 2021a] and Twins [Chu et al., 2021a]) is unclear.  In addition, the paper should compare with these methods.\n\n5. The network configurations for MaiT-XT, Mai-T, MaiT-XS, MaiT-S are not clearly explained. \n\n6. The paper argues that the proposed masked attention heads can facilitate the convergence but I failed to find any training and validation loss curves to justify this statement. \n\n7. In Figure 1, why does the green mask correspond to the 72-th patch? In Figure 2, what does ‘-’ mean in the visualized attention mask?\n\nOverall, I think the submission is not complete at the current stage and needs to be significantly improved. \n\n**References**:\n\n[A1]: On the Relationship between Self-Attention and Convolutional Layers, in ICLR 2020", "summary_of_the_paper": "This paper proposes to introduce attention masks to incorporate spatial locality into part of self-attention heads, while keeping the remaining heads to capture global dependencies. Experiments show improved accuracy on ImageNet in comparison with some baselines.", "summary_of_the_review": "The advantage of the proposed masked attention over the existing literature is not properly discussed. In addition, paper writting is very bad, with substansive confusing descriptions and notations.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "qh3T2JINoqs", "reviewer": "Reviewer_wzwv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "strengths:\n- the paper proposes a masking strategy for attention mechanism, to incorporate the spatial locality into self-attention heads.\nweaknesses :\n- how is the proposed method differs from the multi-scale transformers or hierarchical vision transformers\n-  novelty for masking strategy is minimal.\n- Can the proposed method be extended to any kind of transformers.\n- it hard to understand how the masking strategy addresses the spatial locality-based attention mechanism\n", "summary_of_the_paper": "The paper proposes a method that increased DeiT performance by 1%, where the proposed method introduces attention to incorporate spatial locality into self-attention heads. By doing this paper claims that local dependencies are captured with masked attention heads along with global dependencies captured by original unmasked attention heads.", "summary_of_the_review": "In summary the proposed method has few strengths and weaknesses, it would be easy to understand the effectiveness of the proposed method, if the paper provides experiments showing the improvements when masking strategy is introduced to other transformer networks", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "54XSDA6mO6_", "reviewer": "Reviewer_PxTP", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "Pros:\n1. The paper is well-written.\n2. The experimental results shows that improvement of the proposed method compared with the baselines.\n\nCons:\n1. The idea of this paper is closely related to Swin Transformer. Swin transformers bring locality into network architecture by the window attention, which could save computation significantly. Yet, using attention masks does not have this additional benefit.\n2. Thus, whether the proposed method could be applied to other transformers with locality in the attention module is questionable.\n3. In the experiments, other transformers with locality mechanism (Swin transformer, LocalViT, T2T-ViT) should be compared.\n", "summary_of_the_paper": "This paper proposes to bring locality into the attention module of vision transformers. This locality mechanism is brought by the introduced attention masks. Basically, the attention masks are binary and is likely to restrict the attention to the local field of a token. The local attention mask results in a block matrix before the application of Softmax function. Tokens with zero mask values are given a constant attention value of 1. Thus, those tokens are still involved in the computation. It is claimed that the binary masks is able to keep the global connection when necessary.", "summary_of_the_review": "The main concern of this paper is the usefulness of the locality mechanism due to the existence of Swin transformers which are more computationally efficient. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "0ct5koujqpE", "reviewer": "Reviewer_kZat", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "Pros:\n\n- The idea of this paper is interesting. Though there have been many papers discussing the importance of introducing locality into vision transformers, this paper starts from using a mask to improve the self-attention.\n\n- Experiments on DeiT show that the proposed approach performs better than the baselines, especially when two masked heads are used as shown in Table 2.\n\n- Analysis is also given to demonstrate why involving attention mask could benefit training deeper vision transformers.\n\nCons:\n\n- The arguement in Page 4 seems not correct. The authors claim that \"the locality is not intrinsically encoded in the transformer structure.\"\nIn my view, it should be the locality is not explicitly introduced in transformers. As shown in the DeepViT paper, lower transformer blocks can indeed capture local information. I suggest the authors to reorganize the presentation here.\n\n- The mathematic symbols should be italic. I think the authors should take a correct attitude towards paper writing.\n\n- Experiments are insufficient. After the DeiT paper, there are a large number of papers working on vision transformers. In my view, taking the DeiT paper as baseline is already out-of-date. It would be better to show results on stronger baselines, such as CaiT, LV-ViT, which do not change the self-attention itself. If based on a stronger baseline, the proposed approach can still perform well, I think it should be a significant contribution.\n\n- The improvement over the baselines is also not significant. Most previous work, like CaiT, LV-ViT, T2T-ViT, improve the DeiT model by more than 1% in terms of Top-1 accuracy. However, the improvement shown in the paper is lower.\n\n- I do think a comparison with other methods for vision transformers should be added. I cannot find any tables showing such a comparison.", "summary_of_the_paper": "This paper attempts to improve the classic vision transformers, or specifically the DeiT model by introducing the Masked Attention Head. Instead of focusing on aggregating global information in the original self-attention heads, this paper introduce local information via the proposed Masked Attention Head into self-attention. Experiments on ImageNet show that the proposed approach does matter for lifting the model performance of the original DeiT model.", "summary_of_the_review": "Based on the concerns listed above, I think at this moment this paper does not reach the status for publication in ICLR.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a masking strategy to introduce the locality bias into the vision transformers. The experiments show the effectiveness of considering such inductive bias. The reviewers agreed on the importance of the research question and the simplicity of the algorithm. MaiT also has a straight-forward sparse attention extension that performs on the complexity of $O(n)$ rather than $O(n^2)$.

The reviewers also listed some common concerns of the paper:

(1) The novelty of such a masking approach is relatively low. I don't think the ALS or the soft masking adding too much contribution to that. Similar ideas have been explored in a number of papers.

(2) Reviewers also raise concerns about the experiments. Inductive biases often help more in small settings (fewer parameters and FGLOPs) and gain less in the large settings. When comparing with the STOA models, I think this is basically the trend shown in the paper as well. While I appreciate the authors’ efforts in including more comparisons, I have to say I really don’t think the performance gain is significant enough especially in the large settings. Needless to say that there are many other ways of encoding the same locality bias into the model.

Based on the reviewers' judgements and my own opinion, I therefore recommend rejection of this paper.
Decision: Reject