Original: M AI T: I NTEGRATING S PATIAL L OCALITY INTO I MAGE T RANSFORMERS WITH A TTENTION M ASKS
Revised: M AI T: I NTEGRATING S PATIAL L OCALITY INTO I MAGE T RANSFORMERS WITH A TTENTION M ASKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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 efﬁciency especially for embedded applications.
Revised: 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 efﬁciency especially for embedded applications.

Original: In this work, we address this issue by introducing attention masks to incorporate spatial locality into self-attention heads.
Revised: In this work, we address this issue by introducing attention masks to incorporate spatial locality into self-attention heads.

Original: Local dependencies are captured with masked attention heads along with global dependencies captured by original unmasked attention heads.
Revised: Local dependencies are captured with masked attention heads along with global dependencies captured by original unmasked attention heads.

Original: 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.
Revised: With Masked attention image Transformer – MaiT, top-1 accuracy increases by up to 2.5% compared to DeiT, without extra parameters/FLOPs, or external training data.

Original: Moreover, attention masks regulate the training of attention maps, which facilitates the convergence and improves the accuracy of deeper transformers.
Revised: Moreover, attention masks regulate the training of attention maps, which facilitates the convergence and improves the accuracy of deeper transformers.

Original: Masked attention heads guide the model to focus on local information in early layers and promote diverse attention maps in latter layers.
Revised: Masked attention heads guide the model to focus on local information in early layers and promote diverse attention maps in latter layers.

Original: Deep MaiT improves the top-1 accuracy by up to 1.5% compared to CaiT with fewer parameters andless FLOPs.
Revised: Deep MaiT improves the top-1 accuracy by up to 1.7% compared to CaiT with fewer parameters and FLOPs.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: Recently, vision transformers (Dosovitskiy et al., 2021; Touvron et al., 2021a) are gaining momentum translating the success from transformer-based models in natural language processing (NLP) tasks (Vaswani et al., 2017; Devlin et al., 2019).
Revised: Recently, vision transformers (Dosovitskiy et al., 2021; Touvron et al., 2021a) are gaining momentum translating the success from transformer-based models in natural language processing (NLP) tasks (Vaswani et al., 2017; Devlin et al., 2019).

Original: Self-attention heads excel at capturing the long-range dependencies in sequences but struggle at focusing on local information.
Revised: Self-attention heads excel at capturing the long-range dependencies in sequences but struggle at focusing on local information.

Original: Unlike CNNs which have gone through many iterations of optimization, vision transformers are not very efﬁcient and still require huge computing power and a large number of Flops.
Revised: Unlike CNNs which have gone through many iterations of optimization, vision transformers are not very efﬁcient and still require huge computing power and a large number of Flops.

Original: Naturally, combining the beneﬁts of both CNNs and vision transformers is promising to further boost performances of CV models.
Revised: Naturally, combining the beneﬁts of both CNNs and vision transformers is promising to further boost performances of CV models.

Original: The question remains how to effectively integrate inductive bias such as spatial locality into transformers.
Revised: The question remains how to effectively integrate inductive bias such as spatial locality into transformers.

Original: 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 CNN-transformer hybrid structures as in Li et al.
Revised: 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 CNN-transformer hybrid structures as in Li et al.

Original: (2021b);
Revised: (2021b);

Original: Srinivas et al.
Revised: Srinivas et al.

Original: ; Graham et al.
Revised: ; Graham et al.

Original: ; Wu et al.
Revised: ; Wu et al.

Original: (2021a)
Revised: (2021a)

Original: ; Yuan et al.
Revised: ; Yuan et al.

Original: (2021a)
Revised: (2021a)

Original: ; Guo et al.
Revised: ; Guo et al.

Original: However, forcefully inserting convolutional operations into transformers may potentially constrain the learning capacity of transformers.
Revised: However, forcefully inserting convolutional operations into transformers may potentially constrain the learning capacity of transformers.

Original: To capture the spatial information without signiﬁcantly changing the transformer model architecture,
Revised: To capture the spatial information without signiﬁcantly changing the transformer model architecture,

Original: Chu et al.
Revised: Chu et al.

Original: (2021b) introduce extra positional encoding.
Revised: (2021b) introduce extra positional encoding.

Original: Han et al.
Revised: Han et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2021) fuse local and global representations using multiple transformer blocks or branches to simultaneously process images at different scales such as pixel-level, small-patch or large patch.
Revised: (2021) fuse local and global representations using multiple transformer blocks or branches to simultaneously process images at different scales such as pixel-level, small-patch or large patch.

Original: Yuan et al.
Revised: Yuan et al.

Original: (2021b) apply a layerwise tokens-to-tokens transformation to capture local structure.
Revised: (2021b) apply a layerwise tokens-to-tokens transformation to capture local structure.

Original: These approaches usually come with the cost of extra parameters and model complexity, thus potentially lowering the inference speed.
Revised: These approaches usually come with the cost of extra parameters and model complexity, thus potentially lowering the inference speed.

Original: d’Ascoli et al.
Revised: d’Ascoli et al.

Original: ; Zhou et al.
Revised: ; Zhou et al.

Original: (2021) improve the self-attention for better representation with gated positional self-attention and learnable transformation matrix respectively.
Revised: (2021) improve the self-attention for better representation with gated positional self-attention and learnable transformation matrix respectively.

Original: Hu et al.
Revised: Hu et al.

Original: Ramachandran et al.
Revised: Ramachandran et al.

Original: (2019) adapt the self-attention module and improve the performance of CNN models.
Revised: (2019) adapt the self-attention module and improve the performance of CNN models.

Revised: Liu et al.

Revised: (2021) captures locality within shifted windows in a hierarchical structure.

Revised: Though it could save computation in some cases, the added model complexity may lower its potentials.

Original: Different from prior works, we try to incorporate spatial locality without changing its architecture or adding extra parameters/FLOPs.
Revised: Different from prior works, we try to incorporate spatial locality without changing its architecture or adding extra parameters/FLOPs.

Original: We propose attention masks to guide the attention heads to focus on local information.
Revised: We propose attention masks to guide the attention heads to focus on local information.

Original: Masked attention heads extract local dependencies more efﬁciently by allowing information aggregation only from the closest neighbors.
Revised: Masked attention heads extract local dependencies more efﬁciently by allowing information aggregation only from the closest neighbors.

Original: This liberates other unmasked heads to learn global information more effectively.
Revised: This liberates other unmasked heads to learn global information more effectively.

Original: We name the modiﬁed model, ”Masked attention image
Revised: We name the modiﬁed model, Masked attention image

Original: Transformer” (MaiT), which is built on top of DeiT (Touvron et al., 2021a).
Revised: Transformer (MaiT), which is built on top of DeiT (Touvron et al., 2021a).

Original: MaiT gathers both local and global information at the same time from different heads.
Revised: MaiT gathers both local and global information at the same time from different heads.

Original: 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.
Revised: 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.

Original: We proved that less is more in this speciﬁc case with attention masks.
Revised: We proved that less is more in this speciﬁc case with attention masks.

Original: MaiT achieves up to 1.0% higher top-1 accuracy on ImageNet (Deng et al., 2009) with the same model architecture as DeiT (Touvron et al., 2021a).
Revised: MaiT achieves up to 2.5% higher top-1 accuracy on ImageNet (Deng et al., 2009) with the same model architecture as DeiT (Touvron et al., 2021a).

Original: Additionally,  MaiT outperforms CaiT (Touvron et al., 2021b) by up to 1.5% in top-1 accuracy with fewer parameters and simpler structure in deeper transformers.
Revised: Additionally, deep MaiT outperforms CaiT (Touvron et al., 2021b) by up to 1.7% in top-1 accuracy with fewer parameters and simpler structure.

Original: In summary, we make three major contributions in this work: 1).
Revised: In summary, we make three major contributions in this work: 1).

Original: we propose attention masks to encode  spatial locality into self-attention heads without structural change or extra parameter/computation, while improving model efﬁciency. 2).
Revised: we propose attention masks to encode the spatial locality into self-attention heads without structural change or extra parameter/computation, while improving model efﬁciency. 2).

Original: We present a quick and effective search strategy for the masking scheme exploration. 3).
Revised: We present a quick empirical search strategy for the masking scheme exploration, along with an automatic end-to-end search alternative. 3).

Original: We also reveal the importance of the locality across layers and the  impact of attention masks onfacilitating the training and convergence of deep transformer models.
Revised: We also reveal the importance of the locality across layers and the performance impact of attention masks on deep vision transformers.

Original: Note though MaiT is demonstrated with DeiT, the attention mask is applicable to other vision transformers as well.
Revised: Note though MaiT is demonstrated with DeiT/CaiT, the attention mask is applicable to other vision transformers as well.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Spatial locality is an integral part of the convolutional operation with weight ﬁlters attending to local regions of input feature maps.
Revised: Spatial locality is an integral part of the convolutional operation with weight ﬁlters attending to local regions of input feature maps.

Original: Vision transformer (ViT)
Revised: Vision transformer (ViT)

Original: (Dosovitskiy et al., 2021) is the ﬁrst pure transformer-based model on vision tasks, but it requires a large private labeled dataset JFT300M (Sun et al., 2017) to achieve competitive performances.
Revised: (Dosovitskiy et al., 2021) is the ﬁrst pure transformer-based model on vision tasks, but it requires a large private labeled dataset JFT300M (Sun et al., 2017) to achieve competitive performances.

Original: Data-efﬁcient image transformer (DeiT) (Touvron et al., 2021a) improves upon ViT models by introducing stronger data augmentation, regularization, and knowledge distillation.
Revised: Data-efﬁcient image transformer (DeiT) (Touvron et al., 2021a) improves upon ViT models by introducing stronger data augmentation, regularization, and knowledge distillation.

Original: Class-attention in image transformer (CaiT) (Touvron et al., 2021b) extends DeiT by increasing the number of transformer layers.
Revised: Class-attention in image transformer (CaiT) (Touvron et al., 2021b) extends DeiT by increasing the number of transformer layers.

Original: To overcome the difﬁculties of training deeper transformers, CaiT introduces LayerScale and class-attention layers, which increase the parameters and model complexity.
Revised: To overcome the difﬁculties of training deeper transformers, CaiT introduces LayerScale and class-attention layers, which increase the parameters and model complexity.

Original: Tokens-to-Token vision transformer (T2T) (Yuan et al., 2021b) proposes an image transformation by recursive token aggregation to capture local structure.
Revised: Tokens-to-Token vision transformer (T2T) (Yuan et al., 2021b) proposes an image transformation by recursive token aggregation to capture local structure.

Original: Stand-alone self-attention (Ramachandran et al., 2019) applies local self-attention layer to replace spatial convolution and outperform original
Revised: Stand-alone self-attention (Ramachandran et al., 2019) applies local self-attention layer to replace spatial convolution and outperform original

Original: ResNet models.
Revised: ResNet models.

Original: Even though sharing value and key spatially is parameter efﬁcient in this approach, content-based information is lost.
Revised: Even though sharing value and key spatially is parameter efﬁcient in this approach, content-based information is lost.

Original: Transformer-iN-Transformer (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.
Revised: Transformer-iN-Transformer (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.

Original: ConViT (d’Ascoli et al., 2021) proposes the gated positional self-attention to incorporate soft convolutional biases.
Revised: ConViT (d’Ascoli et al., 2021) proposes the gated positional self-attention to incorporate soft convolutional biases.

Original: CrossViT (Chen et al., 2021) proposes a dual-branch transformer architecture for multi-scale feature extraction.
Revised: CrossViT (Chen et al., 2021) proposes a dual-branch transformer architecture for multi-scale feature extraction.

Original: For pixel-level prediction tasks such as semantic segmentation, object detection, Pyramid Vision
Revised: For pixel-level prediction tasks such as semantic segmentation, object detection, Pyramid Vision

Original: Transformer (PVT) (Wang et al., 2021a) introduces a progressive shrinking pyramid and spatialreduction attention with ﬁne-grained image patches.
Revised: Transformer (PVT) (Wang et al., 2021a) introduces a progressive shrinking pyramid and spatialreduction attention with ﬁne-grained image patches.

Original: DETR (Carion et al., 2020) adapts transformers for object detection tasks.
Revised: DETR (Carion et al., 2020) adapts transformers for object detection tasks.

Original: Swin Transformer (Liu et al., 2021) applies hierarchical transformer with shifted windows of varying sizes.
Revised: Swin Transformer (Liu et al., 2021) applies hierarchical transformer with shifted windows of varying sizes.

Original: Twins (Chu et al., 2021a)
Revised: Twins (Chu et al., 2021a)

Original: deploys interleaved locally-grouped self-attention and global sub-sample attention layers to improve performances.
Revised: deploys interleaved locally-grouped self-attention and global sub-sample attention layers to improve performances.

Original: To optimize transformer and save computation, Wu et al.
Revised: To optimize transformer and save computation, Wu et al.

Original: (2021b) uses centroid attention to extract and compress input information, Jaegle et al.
Revised: (2021b) uses centroid attention to extract and compress input information, Jaegle et al.

Original: (2021) iteratively distill inputs into latent space with attention bottlenecks, Wang et al.
Revised: (2021) iteratively distill inputs into latent space with attention bottlenecks, Wang et al.

Original: (2021b) dynamically adjusts the number of tokens with multiple cascading transformers, and Wu et al.
Revised: (2021b) dynamically adjusts the number of tokens with multiple cascading transformers, and Wu et al.

Original: (2020) introduced semantic token to replace pixel-based transformers to save computation.
Revised: (2020) introduced semantic token to replace pixel-based transformers to save computation.

Original: There are hybrid architectures fusing convolutional and transformer blocks, such as LocalViT (Li et al., 2021b), BoTNet (Srinivas et al., 2021), LeViT (Graham et al., 2021), BossNet (Li et al., 2021a), CvT (Wu et al., 2021a), CoaT (Xu et al., 2021), CMT (Guo et al., 2021) for higher accuracy and faster inference.
Revised: There are hybrid architectures fusing convolutional and transformer blocks, such as LocalViT (Li et al., 2021b), BoTNet (Srinivas et al., 2021), LeViT (Graham et al., 2021), BossNet (Li et al., 2021a), CvT (Wu et al., 2021a), CoaT (Xu et al., 2021), CMT (Guo et al., 2021) for higher accuracy and faster inference.

Original: Unlike prior literature, our work explores the intrinsic capability of pure transformer block on incorporating spatial locality and the impact of the locality along the depth direction.
Revised: Unlike prior literature, our work explores the intrinsic capability of pure transformer block on incorporating spatial locality and the impact of the locality along the depth direction.

Original: This work is also inspired by the emerging graph attention network (Veliˇckovi´c et al., 2018), borrowing the concept of message passing and information aggregation from nearest neighbors.
Revised: This work is also inspired by the emerging graph attention network (Veliˇckovi´c et al., 2018), borrowing the concept of message passing and information aggregation from nearest neighbors.

Original: Note attention masks have been used in NLP tasks as a sparsiﬁcation method to reduce the computation complexity, as well as to capture local information, as in Guo et al.
Revised: Note attention masks have been used in NLP tasks as a sparsiﬁcation method to reduce the computation complexity, as well as to capture local information, as in Guo et al.

Original: ; Child et al.
Revised: ; Child et al.

Original: Beltagy et al.
Revised: Beltagy et al.

Original: (2020); Ainslie et al.
Revised: (2020); Ainslie et al.

Original:  V ISION T RANSFORMER P RELIMINARIES
Revised:  V ISION T RANSFORMER P RELIMINARIES

Original: The transformer architecture introduced by Vaswani (Vaswani et al., 2017) inspired many model variants with remarkable success in NLP tasks.
Revised: The transformer architecture introduced by Vaswani (Vaswani et al., 2017) inspired many model variants with remarkable success in NLP tasks.

Original: ViT (Dosovitskiy et al., 2021) extends pure transformer-based architecture into CV applications.
Revised: ViT (Dosovitskiy et al., 2021) extends pure transformer-based architecture into CV applications.

Original: Instead of pixel-level processing, ViT splits the original images into a sequence of patches as inputs and transforms them into patch tokens , for better computation efﬁciency.
Revised: Instead of pixel-level processing, ViT splits the original images into a sequence of patches as inputs and transforms them into patch tokens, for better computation efﬁciency.

Original: In general, ViT consists of 3 fundamental modules: embedding layer, multi-head self-attention, and feed-forward network.
Revised: In general, ViT consists of 3 fundamental modules: embedding layer, multi-head self-attention, and feed-forward network.

Original: To process images in transformer, the original(224x224) RGB images  is ﬂattened into a sequence of
Revised: To process images in transformer, the original RGB images (224 × 224) is ﬂattened into a sequence of

Original: N (16x16) patches.
Revised: N (14 × 14) patches.

Original: Each patch has a ﬁxed size, typically 16x16 pixels.
Revised: Each patch has a ﬁxed size (typical 16x16 pixels).

Original: Patches are then transformed into patch embedding with hidden dimensions (D) of 192, 384, and 768 for tiny, small, and base models respectively in ViT/DeiT.
Revised: Patches are then transformed into patch embedding with hidden dimensions ( D ) of 192, 384, and 768 for tiny, small, and base models respectively in ViT/DeiT.

Original: In addition to patch tokens, the embedding layer also integrates positional information, classiﬁcation and knowledge distillation through the positional token, class token and distillation token, respectively.
Revised: In addition to patch tokens, the embedding layer also integrates positional information, classiﬁcation and knowledge distillation through the positional token, class token and distillation token, respectively.

Original: Positional token is added into the patch embedding with a trainable positional embedding.
Revised: Positional token is added into the patch embedding with a trainable positional embedding.

Original: However, this positional embedding is added only in the embedding layers.
Revised: However, this positional embedding is added only in the embedding layers.

Original: The spatial information is largely lost in the transformer layers since all-to-all attention is invariant to the order of the patches.
Revised: The spatial information is largely lost in the transformer layers since all-to-all attention is invariant to the order of the patches.

Original: The class token is another trainable vector (1xD), concatenated to the patch tokens (total N+1).
Revised: The class token is another trainable vector (1 × D ), concatenated to the patch tokens (total N +1).

Original: It is used to collect information from the patch tokens to make output predictions, while also spreading information among patches during training.
Revised: It is used to collect information from the patch tokens to make output predictions, while also spreading information among patch tokens during training.

Original: Distillation token is sometimes added for knowledge transfer from teacher models, such as a CNN model.
Revised: Distillation token is sometimes added for knowledge transfer from teacher models, such as a CNN model.

Original: When training the distilled version of the model, a distillation token is further concatenated to the patch token along with the class token (total N+2).
Revised: When training the distilled version of the model, a distillation token is further concatenated to the patch token along with the class token (total N +2).

Original: Multi-head self-attention (MHA) module has multiple parallel attention heads, where each of them comprises three main components: Key, Query, and Value. Key((N+1)xd) and Query((N+1)xd) are trained and multiplied to estimate how much weights on each corresponding token in Value((N+1)xd) for output ((N+1)xd):
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: Where softmax is applied to each row of the input product matrix ( QK T ) and √ d provides appropriate normalization.
Revised: Where softmax is applied to each row of the input product matrix ( QK T ) and √ d provides appropriate normalization.

Original: Multiple attention heads in MHA attend to different parts of the input simultaneously.
Revised: Multiple attention heads in MHA attend to different parts of the input simultaneously.

Original: Considering H heads in MHA layer, the hidden dimension D is split equally across all heads (D=Hxd).
Revised: Considering H heads in MHA layer, the hidden dimension D is split equally across all heads ( D = H × d ).

Original: Feed-forward network (FFN) follows after the MHA module, containing two linear transformation layers separated by Gelu activation.
Revised: Feed-forward network (FFN) follows after the MHA module, containing two linear transformation layers separated by Gelu activation.

Original: The hidden dimension expands by 4x after the ﬁrst linear layer from D to 4D, and is reduced back to D in the second linear layer.
Revised: The hidden dimension expands by 4x after the ﬁrst linear layer from D to 4 D , and is reduced back to D in the second linear layer.

Original: Both MHA and FFN use skipconnections with layer normalization as the residual operation.
Revised: Both MHA and FFN use skip-connections with layer normalization as the residual operation.

Original:  M ASKED A TTENTION H EAD
Revised:  M ASKED A TTENTION H EAD

Original: Spatial locality plays a crucial role in computer vision tasks.
Revised: Spatial locality plays a crucial role in computer vision tasks.

Original: CNN models capture it using the sliding ﬁlter of shared weights, typically with a receptive ﬁeld of 3x3, 5x5, or 7x7.
Revised: CNN models capture it using the sliding ﬁlter of shared weights, typically with a receptive ﬁeld of 3 × 3, 5 × 5, or 7 × 7.

Original: In contrast to CNN models, the locality is not intrinsically encoded in the transformer structure.
Revised: In contrast to CNN models, the locality is not explicitly introduced in the transformer structure.

Original: With attention masks, we can explicitly insert locality into self-attention modules without introducing any extra parameter or computation.
Revised: With attention masks, we can explicitly insert locality into self-attention modules without introducing any extra parameter or computation.

Original: The key idea is to apply a mask on the all-to-all attention products (i.e. QK T ) to reinforce the weights from the closest neighbors and allow informationaggregation only from tokens selected by the mask.
Revised: The key idea is to apply a mask on the all-to-all attention products (i.e. QK T ) to reinforce the weights from the closest neighbors by aggregating information only from tokens selected by the mask.

Original: Figure 1 illustrates an example of our proposed masking scheme.
Revised: Figure 1 illustrates an example of our proposed masking scheme.

Original: The orange box shows a 3 ×mask, where only the direct neighboring patches are selected.
Revised: The orange box shows a 3 ×mask, where only the direct neighboring patches are selected.

Original: Speciﬁcally, 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.
Revised: Speciﬁcally, 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.

Original: This is different from the typical all-to-all attention module, where Patch 16 attends to all 0-196 patches.
Revised: This is different from the typical all-to-all attention module, where Patch 16 attends to all 0-195 patches.

Original: We can easily expand the depth of the attention mask beyond the closest neighbors, to second-level neighbors (green box in Figure 1) or more.
Revised: We can easily expand the depth of the attention mask beyond the closest neighbors, to second-level neighbors (green box in Figure 1) or more.

Original: Note that the class token (and distillation token) still attends to all the patches to collect and spread information during forward and backward passes.
Revised: Note that the class token (and distillation token) still attends to all the patches to collect and spread information during forward and backward passes.

Original: 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.
Revised: 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.

Original: The attention mask is added before Softmax, regulating the distribution of attention maps to focus more on the closest neighbors:
Revised: The attention mask is added before Softmax, regulating the distribution of attention maps to focus more on the closest neighbors:

Original: [Equation]
Revised: [Equation]

Original: where M ∈ R ( N +1) × ( N +1) is a 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.
Revised: where M ∈ { 0 , 1 } ( N +1) × ( N +1) is a 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.

Original: Note it’s important to add the mask before Softmax because it allows theﬂexibility for the model to learn the importance of the locality.
Revised: Note it’s important to add the mask before Softmax because it allows the model to learn the importance of the locality ﬂexibly.

Original: More precisely, unselected patches appear as e 0 = 1 in the numerator of the softmax operation.
Revised: More precisely, unselected patches appear as e 0 = 1 in the numerator of the softmax operation.

Original: Thus, if the attention product result of the closest neighbors is meaningfully larger than zero (i.e. M (cid:12) QK T (cid:29) 0 ), it suggests that local information dominates.
Revised: Thus, if the attention product result of the closest neighbors is meaningfully larger than zero (i.e. M (cid:12) QK T (cid:29) 0 ), it suggests that local information dominates.

Original: However, if those results are negative or close to zero, it implies that local information is insigniﬁcant and global information is more important.
Revised: However, if those results are negative or close to zero, it implies that local information is insigniﬁcant and global information is more important.

Original: Therefore, inserting masks before the softmax operation allows models to enforce locality or bypass it.
Revised: Therefore, inserting masks before the softmax operation allows models to enforce locality or bypass it.

Original: 4. A TTENTION L OCALITY S CORE (ALS)
Revised: 4. A TTENTION L OCALITY S CORE (ALS)

Original: The softmax operation transfers the  QK Tproduct results into the probability space.
Revised: The softmax operation transfers the results of QK T into the probability space.

Original: As a result, each row of the attention map ( A ) sums up to 1.
Revised: As a result, each row of the attention map ( A ) sums up to 1.

Original: For each patch, the probability of focusing on local neighbors equals the sum of their neighbors’ attention map weights.
Revised: For each patch, the probability of focusing on local neighbors equals the sum of their neighbors’ attention map weights.

Original: If this number approaches one, the local information is crucial; whereas if it is close to zero, it means global information matters.
Revised: If this number approaches one, the local information is crucial; whereas if it is close to zero, it means global information matters more.

Original: For Patch n , we deﬁne ALS n as the Attention Locality Score (ALS) and the average from all  patches (N+1) as ALS for each attention head, as follow:
Revised: For Patch n , we deﬁne ALS n as the Attention Locality Score (ALS) . The mean of all N patches is the ALS for each attention head, as follow:

Original: [Equation]
Revised: [Equation]

Original: , where [Equation] n 
Revised: , where ALS n =

Original: [Equation]
Revised: [Equation]

Original: where M is the same attention mask deﬁned earlier, and A = Softmax (cid:16) Mask (cid:12) QK T / √ d (cid:17) is the attention map, or A = Softmax (cid:16) QK T / √ d (cid:17) for unmasked heads.
Revised: where M is the same attention mask deﬁned earlier, and A = Softmax (cid:16) M (cid:12) QK T / √ d (cid:17) is the attention map, or A = Softmax (cid:16) QK T / √ d (cid:17) for unmasked heads, and n, i denotes the row and column index of the attention map respectively.

Original: We later use the ALS metric to get some insights about the locality behavior of different attention heads  in our models.
Revised: We later use the ALS metric to get some insights about the locality behavior of different attention heads at various layers in our models.

Original: 4. M ASKING STRATEGY
Revised: 4. M ASKING STRATEGY

Original: With total H attention heads, h (cid:48) number of attention heads can be allocated to focus on local information.
Revised: With total H attention heads, h (cid:48) number of attention heads can be allocated to focus on local information.

Original: The rest of the unaltered attention heads ( H − h (cid:48) ) capture global dependencies.
Revised: The rest of the unaltered attention heads ( H − h (cid:48) ) capture global dependencies.

Original: Therefore, we can extract the local information through masked attention heads and global information through original attention heads at the same time, as shown in Figure 2.
Revised: Therefore, we can extract the local information through masked attention heads and global information through original attention heads at the same time, as shown in Figure 2.

Original: Besides the depth of the attention mask, the number of masked attention heads in MHA and the position (which layer to insert mask) are also hyper-parameters.
Revised: Besides the depth of the attention mask, the number of masked attention heads in MHA and the position (which layer to insert mask) are also hyper-parameters.

Original: Though masks encode locality into the attention heads, the regularization from masking (pruning attention map) can also limit the learning capacity.
Revised: Though masks encode locality into the attention heads, the regularization from masking (pruning attention map) can also limit the learning capacity.

Original: Consequently, where to add attention masks requires careful consideration.
Revised: Consequently, where to add attention masks requires careful consideration.

Original: The complexity to search for the best masking strategy is up to 2 36 for a 12-layer model with 3 heads, and thus we leave it for future work.
Revised: The complexity to search for the best masking strategy is up to 2 36 for a 12-layer model with 3 heads.

Original: Instead, we leverage the insights from ALS study to guide the mask placement, by using a masking strategy as follows:
Revised: Thus for a quick empirical search to obtain a 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:

Original: Add mask to only one attention head (Head 0) for all layers  2.
Revised: 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.

Original: Compute ALS 0 ,i for every layer i (a) If ALS 0 ,i is close to 1: add masks to all attention heads of layer i, proceed to 3.
Revised: Assignment : compute ALS 0 ,l for every layer l (a) If ALS 0 ,l is close to 1: add masks to all attention heads of layer l , proceed to(b)

Original: (b) If ALS 0 ,i is close to 0.5: keep the mask, add no more (c) If ALS 0 ,i is close to 0: remove the mask 3.
Revised: If ALS 0 ,l is close to 0.5: keep the mask, add no more (c) If ALS 0 ,l is close to 0: remove the mask 3.

Original: Following 2.a), compute ALS h,i for each head and remove the mask if ALS h,i is nearly
Revised: Calibration : following 2.a), compute ALS h,l for each head and remove the mask if ALS h,l is close to

Original: C ROSS - LAYER SIMILARITY
Revised: Note that optimal masking is a large space as opposed to a single unique placement. The model is capable of adjusting to different masking schemes ﬂexibly. Step 2 (a) relies on the observations that early layers (typically with ALS close to 1) focus heavily on local features and thus more masked heads are beneﬁcial. Step 3 will remove the redundant mask accordingly.

Original: To evaluate the impact of attention masks on the diversity of the attention maps across all layers, we apply a cross-layer similarity metric similarly deﬁned in DeepViT (Zhou et al., 2021):
Revised: To alleviate the complicated searching problem and automatically learn the mask placement, we introduce a learnable scale factor α ∈ (0 , 1) for each attention head. We name this approach Soft

Original: [Equation]
Revised: Masking, which replaces the zeros in the original masks with scale factors and keeps the ones. The binary masking approach is named Hard Masking.

Original: [Equation]
Revised: [Equation]

Original: where M i,jh,t is the attention map cosine similarity between layer i and layer j for attention head h and token
Revised: where M n,j , M (cid:48) n,j is the n-th row and j-th column element in the hard and soft mask respectively.

Original: t.
Revised: This scale factor penalizes the attention weights from non-neighboring patches. It allows non-local patch tokens to contribute differently. When α approaches 0, this attention head focuses more on local information. Otherwise when it’s close to 1, the attention head attends global information.

Original: A h,t, : measures the weight contribution for each token in input (Value) for output token T. Therefore, it reﬂects the similarity of attention maps on information aggregation across all T input tokens.
Revised: As a result, each attention head at every layer is able to ﬂexibly determine the importance of locality. Note, this introduces negligible extra parameters during training.

Original: For example, if M i,jh,t reaches 1, output token t at head h attends to all N+1 tokens in Value with the same probability for both layer i and layer j.
Revised: For example, it only introducesextra parameters for a 12-layer MaiT tiny model with 3 heads.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we report the results for MaiT, implemented by applying attention masks on top ofDeiT.
Revised: In this section, we report the results for MaiT, implemented by applying attention masks on top of

Original:  We consider different hyper parameters related to the masking scheme including  mask depth, number of masked heads per layer, and number of transformer layers.
Revised: DeiT. We evaluate the impact of different hyper parameters related to the masking scheme including the mask depth, number of masked heads per layer, and number of transformer layers.

Original: We also evaluate the importance of locality in the depth direction and the  impact of masks for deep transformers.
Revised: We also analyze the importance of locality in the depth direction and the performance impact of masks for deep transformers.

Original: In our evaluation, we merely focus on models with fewer than 50M parameters, which are more applicable to embedded systems and mobile devices.
Revised: We merely focus on models with fewer than 50M parameters, which are more applicable to embedded systems and mobile devices.

Original: 5. S ETUP
Revised: 5. S ETUP

Original: We follow the same training procedure as DeiT, unless speciﬁed otherwise.
Revised: We follow the same training procedure as DeiT, unless speciﬁed otherwise.

Original: The implementation is based on the Timm library (Wightman, 2019).
Revised: The implementation is based on the Timm library (Wightman, 2019).

Original: 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.
Revised: 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.

Original: The model parameters we adopted is summarized in Table 1. 
Revised: The model parameters we adopted is summarized in Table 1. [Table] dim.

Original: [Table]
Revised: [Table]

Original: 5. I MPACT OF MASK DEPTH
Revised: 5. I MPACT OF MASK DEPTH

Original: Intuitively, the ﬁeld of view is larger with an attention mask of larger depth, to gather broader neighborhood information.
Revised: Intuitively, the ﬁeld of view is larger with an attention mask of larger depth, to gather broader neighborhood information.

Original: However, in the extreme case with mask depth of N+1 (standard selfattention), all the patch token information is merged and positional information is lost.
Revised: However, in the extreme case with mask depth of N +1 (standard selfattention), all the patch token information is merged and positional information is lost.

Original: On theother hand, attention masks can be considered as a structured pruning on attention maps as well. A smaller depth for the attention mask potentially translates to fewer FLOPs, which might be useful for some hardware devices such as embedded CPUs.
Revised: Nevertheless, considering the pruning effect from the mask, a smaller depth for the attention mask potentially translates to fewer FLOPs, which might be useful for some hardware devices such as embedded CPUs.

Original: For example, with 3x3 mask,instead of calculating all 197 attention weights, we only compute 9  of them, a 95% reduction for  attention map computation.
Revised: For example, with the 3 × 3 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.

Original: To study the impact of the mask depth, we mask only one attention head for every MHA module, with the ﬁeld of view as 3x3, 3x5, and 5x5 on DeiT-small model.
Revised: To study the impact of the mask depth, we mask only one attention head for every MHA module, with the ﬁeld of view as 3 × 3, 3 × 5, and 5 × 5 on tiny and small DeiT models.

Original: The results are summarized in Table 2.
Revised: The results are summarized in Table 2.

Original: There is 0.1 - 0.3 % improvement compared to DeiT small.
Revised: There is 0.1 - 0.3 % improvement compared to DeiT small.

Original: Various mask sizes only lead to marginal differences in accuracy.
Revised: Various mask sizes only lead to marginal differences in accuracy.

Original: This is probably because each patch (original 16x16 pixels) is a processed sub-image and already contains some locality information, and thus further away patch tokens barely provide much extra information.
Revised: This is probably because each patch (original 16x16 pixels) is a processed sub-image and already contains some locality information, and thus further away patch tokens barely provide much extra information.

Original: Therefore, 3x3 mask is the default setting for MaiT in the subsequent sections.
Revised: Therefore, 3x3 mask is the default setting for MaiT in the subsequent sections.

Revised: In contrast to DeiT, whose accuracy saturates beyond 300 epochs as reported in (Touvron et al., 2021a), MaiT-small sees another 1.1% improvement when trained for 400 epochs.

Revised: Accordingly, we

Original: [Table]
Revised: [Table]

Original: In contrast to DeiT, whose accuracy saturates beyond 300 epochs as reported in (Touvron et al., 2021a), MaiT-small sees another 0.9% improvement when trained for 400 epochs. Accordingly, we suspect a higher learning rate is beneﬁcial for MaiT since attention masks  provide some guidance for attention map training. 
Revised: suspect a higher learning rate and lower drop path rate are beneﬁcial for MaiT-tiny since attention masks already provide some regulation during training. Changing the batch size from 1024 toonly results in small gain of 0.4% for MaiT-tiny with 3 × 3 mask. In comparison, adopting 0 drop path rate and 2 × learning rate, we obtain 2.5% improvement over DeiT-T.

Original: If we change the batch size from 1024 to 3072 for MaiT-tiny, essentially increasing the [-learning rate ( lr = 0 . 
-]
Revised: We compare the performances of 24-layer MaiT models with the following three representative masking schemes (Sch.), to quantify the [-learning rate ( lr = 0 . 
-]impact of the masking strategy. Sch.1 is a naive way to combine local and global information at every layer. Sch.2 results from the masking strategy in

Original: × batchsize / 512 ) by 3x, the top-1 accuracy for MaiT-tiny increases by 0.4% to 72.6 after 300 epochs.
Revised: 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-

Revised: 1 adopts only one masked head.

Revised: Sch

Revised: 2 applies a mixed masking

Revised: Sch masking

Revised: [Table]

Revised: As shown in Table 3, top-1 accuracy improves by 0.8% simply through adopting one masked head for MaiT-T. Another 0.4% improvement is achieved with mixed masking for MaiT-XS.

Revised: Soft masking further increases the accuracy by 0.2% for MaiT-T. Note that the initialization of the scale factor for

Revised: MaiT-XS is subject to further tuning.

Revised: Empirically we ﬁnd Sch.2 yields relative good performances by masking heavily in the ﬁrst onethird of layers to extract local features and only one mask for the middle layers.

Original: Therefore, the number of training epochs and learning rate for MaiT is subject to further tuning for performance improvement.
Revised: 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 can leverage the structured sparsity. This is especially beneﬁcial with high resolution images with a large number of tokens. The computation complexity remains quadratic as O ( N 2 ) for soft masking.

Original: 5. I MPORTANCE OF SPATIAL LOCALITY
Revised: 5. I MPORTANCE OF SPATIAL LOCALITY

Original: To explore the importance of  localityacross transformer layers, we start by masking one head in each layer and compute the ALS for all the heads in all layers.
Revised: Attention Locality Score (ALS) provides a quantitative metric to study the importance of spatial locality for each head across all layers.

Original: Figure 3 a) illustrates the result for 12-layer MaiT-small.
Revised: Figure 3 a) illustrates the result of a 24-layer MaiT-Tiny with one masked head.

Original: ALS of the masked head is signiﬁcantly higherfor the ﬁrst 10 out oflayers, than the unmasked heads.
Revised: ALS of the masked head is signiﬁcantly higher than other heads in Layer 0-22.

Original: This indicates the model learns to focus  on local information with the masked attention head and global information with unmasked heads.
Revised: This indicates the model learns to focus more on local information with the masked attention head.

Original: We observe three stages in the ALS across all layers for themasked head: 1. high ALS, where more than 70% of attention weights are allocated to closest neighbors in the ﬁrst 4 layers; 2. mid ALS, where the probability to focus on locality decreases to around 50% forLayer 4 to Layer 9; 3. low ALS, where the chance of attending to local patches are less than 0.5% for the last 2 layers.
Revised: ALS gradually declines from around 90% for the ﬁrst 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), an average score if the attention weights are distributed uniformly. This suggests that spatial locality is largely ignored in the last layer where global information is more important.

Original: Notice that in the third layer, even for unmasked heads, two heads focus heavily on closest neighbors with ALS larger than 0.7. Similarly, among the unmasked heads, the ﬁrst 4 layers show higher ALS, which decreases in the following 6 layers and down to an average of 0.0051 (1/(N+1) = 1/197) if the attention weights are distributed uniformly.
Revised: Similarly, when learning the masking automatically through soft masking (Figure3 b)), the ﬁrstlayers show over 0.9 ALS for multiple heads and layers. ALS gradually decreases in the middle layer with at most one head focusing heavily on local features. ALS is below 0.3 for the lastlayers, where all heads tend to global information.

Original: This indicates the  locality is more important in the early layers (ﬁrst 4 layers) and becomes trivial in the last two layers, which is consistent with d’Ascoli et al.
Revised: Overall, the spatial locality is more important in the early layers, and its importance gradually declines in deeper layer, which is consistent with d’Ascoli et al.

Original: (2021), Raghu et al.
Revised: (2021), Raghu et al.

Original: When the number of transformer blocks increases, the importance of locality expands to deeper layers as shown in Figure 7.

Original: The trends remain for deeper MaiT of 24 or 36 layers.

Original: Accordingly in Figure3 b), we mask all heads for the ﬁrst one-third of total layers, only one masked head for 8 layers in the middle (8-15), and no mask for the rest of the layers.

Original: For fully masked layers, ALS of most heads are larger than 0.5. Though one head in every other layer of the ﬁrst 8 blocks shows lower than 0.1 ALS. This suggests at least one head is trying to gather global information even in the ﬁrst 8 layers.
Revised: 5. DEEPER TRANSFORMERS

Original: Thus it’s beneﬁcial to keep one global attention head even for the initial few layers.

Original: 5. I MPACT OF ATTENTION MASK ON DEEPER TRANSFORMERS

Original: Besides lacking inductive bias, another issue with transformers is the difﬁculty  to train deeper models. Naively stacking transformer layers fails to deliver the expected performance gain as shown in Touvron et al.
Revised: Another difﬁculty is to train deeper models because naively stacking transformer layers fails to deliver the expected performance gain as shown in Touvron et al.

Original: (2021b).
Revised: (2021b).

Original: One reasonfor the performance degradation is attention collapse, i.e. the attention maps are more alike among deeper layers (Zhou et al., 2021).
Revised: One reason is attention collapse, i.e. the attention maps are more alike among deeper layers (Zhou et al., 2021).

Original: We ﬁnd the attention mask is able to break the attention map homogeneity and facilitates the convergence of the deep transformers.

Original: The mixed masking scheme breaks the structural repetition in deep transformers and promotes diversity in later layers.
Revised: We ﬁnd mixed masking is able to break the structural repetition in deep transformers and promotes diversity in later layers.

Original: When increasing transformer blocks from 24 to 36,even though the model capacity is largely increased, top-1 accuracy only increases by 0.25% for DeiT-tiny (Figure 4).
Revised: 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).

Original: With MaiT changing one head to focus on local information, top-1 accuracy is 0.8% higher over DeiT of the same layers up to 32 layers. However, at 36 layers, the performance of MaiT degrades to the same level as DeiT. Beyond 32 layers, a mixed masking strategy can be adopted to extend the performance gain. If the 36-layer MaiT applies fully 3 masked heads in the ﬁrst 12 layers and only one masked head for therest 24 layers, the top-1 accuracyincreases by 1.2% (from 78.6% to 
Revised: In contrast, MaiT with mixed masking boosts the accuracy to 79.8%, extending the performance gain with deeper transformers.

Original: Interestingly, we noticed that for 24-layer transformers, CaiT shows even worseperformances than naively stacked DeiT-24, suggesting that the hyper-parameters for 24-layer CaiT-tiny are probably not optimal.
Revised: Interestingly, we noticed that for 24-layer model, CaiT is even worse than naively stacked DeiT-24, suggesting that the hyper-parameters for 24-layer CaiT-tiny are probably not optimal.

Original: To rule out theimpact of pure pruning effects on attention maps from masking, we also apply a random mask of the same drop-out rate to one head of all 24 layers.
Revised: To rule out the pure pruning effects on attention maps from masking, we also apply a random mask of the same drop-out rate to one head of all 24 layers.

Original: The randomly masked model leads to 0.5% accuracy dropas compared to 24-layer DeiT.
Revised: The randomly masked model leads to 0.5% accuracy drop compared to 24-layer DeiT.

Original: This provides further proof that the performance gain with attention masks is indeed from better  locality information aggregation instead of pruning on attention maps.
Revised: This provides further proof that the performance gain with attention masks is indeed from better spatial locality aggregation.

Original: To better understand the impact of attention masks in deep MaiT, we compare the cross layer similarity heat map among 36 layers of DeiT, MaiT (one masked head), and MaiT (mixed masking) inFigure 8, Figure 9, and Figure 5 respectively.
Revised: To better understand the impact of attention masks in deep MaiT, we compare the cross layer similarity heat map among 36 layers of DeiT, MaiT (one masked head), and MaiT (mixed masking) in

Original: Each point in the heat map denotes the averaged similarity between the y-axis indexed layer and the layer from x-axis. DeiT shows the highest cross-layer similarity at the last 8 stages, with an average of 0.75 (Figure 8). This homogeneity in attention maps at 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.
Revised: Figure 7, Figure 8, and Figure 9 respectively.

Original: As a result, the average cross-layer similarity for the last 8 layers with unmasked heads is also 0.75 (Figure 9). Alternatively, if we mask all heads in the ﬁrst 8 layers, they are naturally different from the rest of the partially masked layers, and the same structure is only repeated 8 or 24 times. As shown in Figure 5, cross-layer similarity between the ﬁrst 8 layers and the following 24 layer are mostly well below 0.3 for Head 1 and 2. Additionally, a mixed masking scheme lower the cross-layer similarity among the last 8 layers by 0.1 to 0.65.
Revised: Effective masking lowers the cross-layer similarity and thus improves scalability of MaiT .

Original: [Figure] C OMPARISON WITH D EI T AND C AI T In Table
Revised: 5. C OMPARISON WITH OTHER STATE OF THE ART VISION TRANSFORMERS

Original: [Table] we ﬁnd the attention mask is more effective with small hidden dimensions on deeper transformer models. To improve DeiT-T
Revised: Attention mask can be easily combined with other optimization techniques. For deep transformers, we adopt the LayerScale (Touvron et al., 2021b) to MaiT as MaiT*.

Original: [Table]
Revised: As shown in Table 4,

Original: we introduce MaiT-XT (24 layers) with similar parameters and FLOPs, achieving 1.5% higher top-1 accuracy. 24-layer MaiT-T outperforms CaiT-T by 1.5% even with fewer parameters and less FLOPs. However, with increased hidden dimension, MaiT-XS and MaiT-S are 0.8% and 0.7% lower than CaiT-XS and CaiT-S, respectively.
Revised: MaiT/MaiT

Original: Note that CaiT applies multiple optimization techniques to improve accuracy.
Revised: * obtain competitive performances across different model sizes.

Original: In comparison, by simply adding attention masks, MaiT-S improves the accuracy by 1.1% and thus bridges the gap between CaiT-S  and naively stacked DeiT-S. Moreover, if we add LayerScale to MaiT-S, its accuracy surpasses CaiT-S by 0.1% with fewer parameters and FLOPs.
Revised: Speciﬁcally, MaiT-T outperforms CaiT-T by 1.7%. MaiT-S* achieves 0.2% higher accuracy than CaiT-S with less parameters/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.

Original: It’s promising for more gain when combining other optimization techniques with MaiT.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we incorporate spatial locality into vision transformers by inserting masks to attention heads.
Revised: In this work [Table] we incorporate spatial locality into vision transformers by inserting masks to attention heads.

Original: Masked heads are able to focus on local information and liberate other unmasked heads to extract global information more effectively.
Revised: Masked heads are able to focus on local information and liberate other unmasked heads to extract global information more effectively.

Original: We also introduce attention locality score to guide the masking strategy search and rapidly evaluate various masking schemes.
Revised: We also introduce attention locality score to guide the mask search and rapidly evaluate various masking schemes.

Original: Attention masking is a simple and effective technique, especially for smaller models (i.e., < 6 M parameters).
Revised: Attention masking is a simple and effective technique, especially for small and deep transformer models.

Original: We observe that attention masking also serves as a regularizer to guide the training of attention maps for deeper transformers.
Revised: We observe that attention masks also serve as a regularizer to guide the training of attention maps for deeper transformers.

Original: Moreover, it is promising to further boost performances by combining this attention mask technique with other optimization approaches.
Revised: Moreover, it is promising to further boost performances by combining this attention mask technique with other optimization approaches.

Original: A A PPENDIX A. ALS OF D EI T AND DEEP M AI T
Revised: A A PPENDIX A. H YPER PARAMETERS FOR M AI T

Original:  When trained without restriction from attention mask, DeiT-small emphasizes the locality by allocating more than 50% of attention weights on closest neighbors for 3 heads at second and third layer, even close to 90% at the second layer as shown in Figure 6.
Revised: [Table] A. ALS OF D EI T AND DEEP M AI T When trained without restriction from attention masks, DeiT emphasizes the locality in the ﬁrstlayers with over 0.8 ALS for some heads, as shown in Figure 5.

Original: In comparison, in Figure 3, MaiT utilizes masked Head 0 to extract local information from the ﬁrst layer to the ninth layer.
Revised: In comparison, in Figure 3, MaiT utilizes masked Head 0 to extract local information and ALS is lower for unmasked Head 1 and

Original: Another notable change between MaiT and DeiT is the ALS from Layer 4 to 9, masked head 0 elevates ALS above other heads to focus on local information. Head 1 to Head 5 in MaiT show lower ALS compared to DeiT, and focus more on global information.
Revised: Head 2.

Original: When the number of transformer blocks increases, the importance of locality expands to deeper layers as shown in Figure 7.
Revised: When the number of transformer blocks increases, the importance of locality expands to deeper layers as shown in Figure 6.

Original: For the masked Head 0 in MaiT-tiny, 12-layer model shows more than 0.8 ALS for 6 layers out of the ﬁrst 7 layers, the following 4 layers give a ALS around 0.5, and
Revised: For the masked Head 0 in MaiT-tiny, 12-layer model shows more than 0.8 ALS for 6 layers out of the ﬁrst 7 layers, the following 4 layers give a ALS around 0.5, and

Original: the
Revised: the

Original: ALS of the last layer declines below 0.05.
Revised: ALS of the last layer declines below 0.05.

Original: Likewise, for 24-layer and 36 layer models, 10 layers out of the ﬁrst 12 layers show ALS higher than 0.8, and decreases to 0.5 or 0.6 in the following layers before the last layer.
Revised: Likewise, for 24-layer and 36 layer models, 10 layers out of the ﬁrst 12 layers show ALS higher than 0.8, and decreases to 0.5 or 0.6 in the following layers before the last layer.

Original: A. H YPER PARAMETERS FOR M AI T
Revised: A. C ROSS - LAYER SIMILARITY

Revised: To evaluate the impact of attention masks on the diversity of the attention maps across all layers, we apply a cross-layer similarity metric similarly deﬁned in DeepViT (Zhou et al., 2021):

Original: [Table]
Revised: [Equation]

Original: A. C ROSS LAYER SIMILARITY HEAT MAP
Revised: where M i,jh,t is the attention map cosine similarity between layer i and layer j for attention head h and token t. A h,t, : measures the weight contribution for each token in input (Value) for output token T.

Original: In Figure 8, 36-layer DeiT-tiny shows the highest  similarity at the last 8 stages across all three heads, with an averagesimilarity of 0.75, whereas the similarity is quite small among the ﬁrstlayers and between the ﬁrst 16 layers and the last 20 layers.
Revised: Therefore, it reﬂects the similarity of attention maps on information aggregation across all T input tokens. For example, if M i,jh,t reaches 1, output token t at head h attends to all N+1 tokens in Value with the same probability for both layer i and layer j. DeiT shows the highest cross-layer similarity at the last 8 stages, with an average of 0.75 (Figure 7). This homogeneity in attention maps at the late stage limits the model learning capability.

Revised: 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.

Revised: As a result, the average cross-layer similarity for the last 8 layers with unmasked heads is also 0.75 (Figure 8).

Revised: Alternatively, if we mask all heads in the ﬁrst 8 layers, they are naturally different from the rest of the partially masked layers, and the same structure is only repeated 8 or 24 times.

Revised: As shown in Figure 9, cross-layer similarity between the ﬁrst 8 layers and the following 24 layer are mostly well below 0.3 for Head 1 and 2.

Revised: Additionally, a mixed masking scheme lower the cross-layer similarity among the last 8 layers by 0.1 to 0.65.

Revised: 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 ﬁrstlayers and between the ﬁrst 16 layers and the last 20 layers.

Original: This is agreement with Zhou et al.
Revised: This is agreement with Zhou et al.

Original: (2021), indicating an attention collapse at later stage of deep transformers.
Revised: (2021), indicating an attention collapse at later stage of deep transformers.

Original: In comparison, for MaiT with one masked head across all layers, the masked attention Headshows a more uniform distribution with a lower peak intensity as compared to the the non-masked or partially masked (in early layers)
Revised: In comparison, for MaiT with one masked head across all layers, the masked attention Headshows a more uniform distribution with a lower peak intensity as compared to the the non-masked or partially masked (in early layers)

Original: Head 1 and Head 2 (Figure 9.
Revised: Head 1 and Head 2 (Figure 8.

Original: For unmasked Head 1 and Head 2, later stage especially the last 8 layers show highest similarity among each other, the average of them is also 0.75.
Revised: For unmasked Head 1 and Head 2, later stage especially the last 8 layers show highest similarity among each other, the average of them is also 0.75.

Original: This suggests one masked head for all transformer blocks is not as effective in 
Revised: This suggests one masked head for all transformer blocks is not as effective in 36-layer model as in 24-layer model. This indicates repeating the same structure for 36 times in the depth direction is causing the attention collapse. The diversity among the heads within the same layer doesn’t alleviate this issue.