What to Hide from Your Students: Attention-Guided Masked Image Modeling

Introduction

Self-supervised learning (SSL) has attracted significant attention over the last years. Recently, several studies are shifting towards adapting SSL to transformer architectures. Originating in natural language processing, where self-supervised transformers have revolutionized the field, these architectures were introduced to computer vision with the vision transformer (ViT) as an alternative to convolutional neural networks . ViT formulates an image as a sequence of tokens obtained directly from raw patches and then follows a pure transformer architecture. Despite the absence of image-specific inductive bias, ViT shows strong image representation learning capacity.

Considering that transformers are data-hungry, many studies advocate pre-training them on unsupervised pretext tasks, determined only by raw data. A prominent paradigm is to mask a portion of the input tokens—words in text or patches in images—and train the transformer to predict these missing tokens . This paradigm, called masked language modeling (MLM) in the language domain , is remarkably successful and extends to the vision domain as masked image modeling (MIM) .

MIM-based self-supervised methods have already shown impressive results on images. However, an important aspect that has not been well explored so far is how to choose which image tokens to mask. Typically, the selection is random, as has been the norm for text data. In this work, we argue that random token masking for image data is not as effective.

In text, random word masking is likely to hide high-level concepts that describe entire semantic entities such as objects (nouns) and actions (verbs). By contrast, an image has much more tokens than a sentence, which are highly redundant, and random masking is less likely to hide “interesting” parts; or when it does, the remaining parts still easily reveal the identity of the visual concepts. As shown in Figure 1(b-d), unless masking is very aggressive, this is thus less likely to form challenging token reconstruction examples that would allow the transformer to develop strong comprehension skills.

The question we ask is this: Can we develop a masking strategy that addresses this limitation and makes informed decisions on which tokens to mask?

To this end, we propose to exploit the intrinsic properties of ViT and in particular its self-attention mechanism. Given an input sequence of image patches, we forward it through the transformer encoder, thereby obtaining an attention map in its output. We then mask the most attended tokens. As shown in Figure 1(f-g), the motivation is that highly-attended tokens form more coherent image regions that correspond to more discriminative cues comparing with random tokens, thus leading to a more challenging MIM task.

This strategy, which we call attention-guided masking (AttMask), is an excellent fit to popular distillation-based self-supervised objectives, because it is the teacher encoder that sees the entire image and extracts the attention map, and the student encoder that sees the masked image and solves the reconstruction task. AttMask thus incurs zero additional cost.

We introduce a novel masking strategy for self-supervised learning, called AttMask, that exploits the intrinsic properties of ViT by leveraging its self-attention maps to guide token masking (subsection 3.2).

We show how to efficiently incorporate this above masking strategy into teacher-student frameworks that use a MIM reconstruction objective and demonstrate significant performance improvements over random masking.

Through extensive experimental evaluation, we confirm that AttMask offers several benefits: it accelerates the learning process; it improves performance on a data-limited regime (subsection 4.2) and on a variety of downstream tasks (subsection 4.3); it increases the robustness against background changes, thus revealing that it reduces background dependency.

Related Work

Transformers are based on self-attention and require pretraining on large unlabelled corpora . Their adaptation to vision tasks is not straightforward. Representing pixels by tokens is impractical due to the quadratic complexity of self-attention, giving rise to approximations . The idea of representing image patches by tokens is proposed in , where patches are of size $2\times 2$, and is further studied in ViT , where patches are $16\times 16$. Despite the absence of image-specific inductive bias, ViT is competitive to convolutional neural networks for ImageNet and other smaller benchmark datasets . Since it is pretrained on a large and private dataset , authors of DeiT question its efficiency and propose an improved data-efficient version, which however is based on a strong teacher instead .

Self-supervised Learning.

Early self-supervised learning methods follow the paradigm of training on an annotation-free pretext task, determined only by raw data . This task can be e.g. the prediction of patch orderings or rotation angles . Starting from instance discrimination and contrastive predictive coding , contrastive learning has become very popular . These methods pull positives together and push negatives apart, where positives are typically determined by different views of the same example. Alternatively, contrastive learning often relies on clustering . The requirement of negatives is eliminated in BYOL , OBoW , SimSiam and DINO , where the challenge is to avoid representation collapse, most notably by a form of self-distillation .

Using transformers, MIM as a pretext task is proposed in BEiT , which maps the images to discrete patch tokens and recovers tokens for masked patches, according to a block-wise random strategy. Other than that, MIM methods use continuous representations: SimMIM randomly masks large patches and predicts the corresponding pixels by direct regression; MAE randomly masks a large portion of patches and predicts the corresponding pixels using an autoencoder; MST masks low-attended patches and reconstructs the entire input with a decoder; iBOT extends the self-distillation loss of DINO to dense features corresponding to block-wise masked patches. Here, we advocate masking of highly-attended patches, in a sense the opposite of MST, and we exhibit this idea in the context of DINO and iBOT.

Regularization and Augmentation.

As the complexity of a task increases, networks with more and more parameters are introduced. But with increased representational power comes increased need for more data or risk of overfitting. Several regularization and data augmentation methods have been proposed in this direction , combined with standard supervised tasks.

In this context, feature masking is introduced by Dropout , which randomly drops hidden neuron activations. To address the strong spatial correlation in convolutional feature maps, SpatialDropout randomly drops entire channels. DropBlock generalizes Dropout—or constrains SpatialDropout—by dropping features in a block, i.e., a square region of a feature map. Attention Dropout makes use of self-attention to mask the most discriminative part of an image. Feature-space masking, guided by attention from another network or branch, has been extensively studied as a mechanism to explore beyond the most discriminative object parts for weakly-supervised object detection . Our work is a natural evolution of these ideas, where attention is an intrinsic mechanism of transformers; and the task becomes that of densely reconstructing the masked features. This is a pretext task, without need for supervision.

Method

A simplified overview of the method is shown in Figure 2. We first discuss in subsection 3.1 preliminaries and background on vision transformers and self-supervision with distillation-based masked image modeling. In subsection 3.2, we then detail our attention-guided token masking strategy, called AttMask, and how we incorporate it into masked image modeling.

We are given an input image $X\in{}^{h\times w\times c}$, where $h\times w$ is the spatial resolution and $c$ is the number of channels. The first step is to tokenize it, i.e., convert it to a sequence of token embeddings. The image is divided into $n=hw/p^{2}$ non-overlapping patches $P_{i}\in{}^{p\times p\times c}$ for $i=1,\dots,n$, where $p\times p$ is the patch resolution. Each patch is flattened into a vector in ${}^{p^{2}c}$ and projected to an embedding vector $\mathbf{z}_{i}\in{}^{d}$ using a linear layer, where $d$ is the embedding dimension. A learnable embedding $\mathbf{z}^{\textsc{[cls]}}\in{}^{d}$ of a “classification” token [cls] is then prepended to form the tokenized image

where “;” denotes row-wise stacking. The role of this special token is to represent the image at the output. A sequence of position embeddings is added to $Z$ to retain positional information. The resulting sequence is the input to the transformer encoder. Each layer of the encoder consists of a multi-head self-attention (MSA) block followed by a multi-layer perceptron (MLP) block. Through all of its layers, the encoder uses a sequence of fixed length $n+1$ of token embeddings of fixed dimension $d$, represented by a $(n+1)\times d$ matrix. The embedding of the [cls] token at the output layer serves as the image representation.

An MSA block consists of a number $H$ of heads, each computing a scaled dot-product self-attention , i.e., the relevance of each image patch to others, encoded as an $(n+1)\times(n+1)$ attention matrix. As discussed in subsection 3.2, we average attention matrices over all the heads of the last encoder layer and we use the row corresponding to the [cls] token to generate token masks.

Distillation-based Masked Image Modeling.

Self-distillation, using a moving average of the student as teacher , is studied for self-supervision in BYOL and extended to vision transformers in DINO , which applies the distillation loss globally on the [cls] token. iBOT turns this task into masked image modeling (MIM) by applying the loss densely on masked tokens.

for $i=1,\dots,n$, where $\mathbf{z}^{\textsc{[mask]}}\in{}^{d}$ is a learnable embedding of a “mask” token [mask]. Following the strategy of BEiT , the mask vector is generated with random block-wise token sampling, that is, defined in terms of random rectangles in the 2D layout of the $n$ tokens as a $(h/p)\times(w/p)$ matrix.

Following DINO , the transformer encoder is followed by a head that includes an MLP and scaled softmax, such that output token embeddings can be interpreted as probabilities. We denote by $f_{\theta}$ the mapping that includes the addition of the position embeddings, the encoder and the head, while $\theta$ is the set of learnable parameters. Given a tokenized image $Z$, masked or not, we denote by $f_{\theta}(Z)\in{}^{(n+1)\times d}$ the output token sequence and by $f_{\theta}(Z)_{i},f_{\theta}(Z)^{\textsc{[cls]}}\in{}^{d}$ the embedding of the $i$-th and [cls] token respectively. The teacher parameters $\theta^{\prime}$ are obtained from the student parameters $\theta$ by exponential moving average (EMA) according to $\theta^{\prime}\leftarrow\alpha\theta^{\prime}+(1-\alpha)\theta$.

For each input image, two standard resolution augmented global views are generated, with tokenized images $Z^{a},Z^{b}$ and mask vectors $\mathbf{m}^{a},\mathbf{m}^{b}$. For each view $v$ in $V=\{a,b\}$ and for each masked token, the MIM objective is to minimize the reconstruction loss between the student $f_{\theta}$ output for the masked input $\widetilde{Z}^{v}$ and the teacher $f_{\theta^{\prime}}$ output for the non-masked input $Z^{v}$:

Following DINO , a similar loss is applied globally on the [cls] tokens between the student output for one masked view $\widetilde{Z}^{v}$ and the teacher output for the other non-masked view $Z^{u}$:

Finally, as detailed in the Appendix Appendix 0.B, a multi-crop strategy applies, giving rise to a loss $L_{\textsc{lc}}$ (A13) between local crops and global views. The overall loss of iBOT is a weighted sum of $L_{\textsc{mim}}$ (3) and $L_{\textsc{g}}$ (4) $+$ $L_{\textsc{lc}}$ (A13). DINO itself uses the sum $L_{\textsc{g}}$ (4) $+$ $L_{\textsc{lc}}$ (A13) without masking.

2 AttMask: Attention-guided Token Masking

Prior MIM-based self-supervised methods use random or block-wise random token masking. In this section we describe our attention-guided token masking strategy, which hides tokens that correspond to the salient regions of an image and thus define a more challenging MIM objective.

Given an input sequence $Y\in{}^{(n+1)\times d}$, a multi-head self-attention (MSA) layer uses three linear layers to map $Y$ to the query $Q_{j}$, key $K_{j}$ and value $V_{j}$ sequences for $j=1,\dots,H$, where $H$ is the number of heads, $Q_{j},K_{j},V_{j}\in{}^{(n+1)\times d^{\prime}}$ and $d^{\prime}=d/H$. Then, it forms the $(n+1)\times(n+1)$ attention matrix, where softmax is row-wise:

To generate token masks from any layer of the transformer encoder, we average the attention matrices over all heads:

Now, each row of an attention matrix is a vector in n+1, that corresponds to one token and, excluding the diagonal elements, determines an attention vector in n over all other tokens. We focus on the attention vector of the [cls] token, which comprises all but the first elements of the first row of $\overline{A}$:

where $\overline{a}_{i,j}$ is the element $i,j$ of $\overline{A}$. This vector can be reshaped to $(h/p)\times(w/p)$ attention map, to be visualized as a 2D image, indicating the regions of the input image that the [cls] token is attending.

Mask Generation: Highly-attended Tokens.

There is a permutation $\sigma_{\downarrow}:\{1,\dots,n\}\rightarrow\{1,\dots,n\}$ that brings the elements of $\overline{\mathbf{a}}^{\textsc{[cls]}}$ in descending order, such that $\overline{a}^{\textsc{[cls]}}_{\sigma_{\downarrow}(i)}\geq\overline{a}^{\textsc{[cls]}}_{\sigma_{\downarrow}(j)}$ for $i<j$, where $\overline{a}^{\textsc{[cls]}}_{i}$ is the $i$-th element of $\overline{\mathbf{a}}^{\textsc{[cls]}}$. Choosing a number $k=\left\lfloor{rn}\right\rfloor$ that is proportional to the total number $n$ of tokens with mask ratio $r\in$, we define

as the set of indices of the top-$k$ most attended tokens. We thus define the high-attention mask vector $\mathbf{m}^{H}$ with elements

for $i=1,\dots,n$. This masking strategy, which we call AttMask-High, essentially hides the patches that correspond to the most discriminative or salient regions of an image. By AttMask we shall refer to this strategy as default.

Low-attended Tokens.

Highly-attended with Hints.

Finally, because AttMask-High may be overly aggressive in hiding the foreground object of an image, especially when the mask ratio $r$ is high, we also examine an alternative strategy that we call AttMask-Hint: While still masking highly attended tokens, we allow a small number of the most highly attended ones to be revealed, so as to leave hints about the identity of the masked object. In particular, we remove from the initial set $M^{H}$ a small number $m=\left\lfloor{sn}\right\rfloor$ of tokens with show ratio $s<r$. These $m$ tokens are randomly selected from the $\left\lfloor{s_{\max}n}\right\rfloor$ most attended tokens in $M^{H}$, where $s_{\max}>s$. An example comparing AttMask-Hint with AttMask-High is illustrated in Figure 3.

Incorporating AttMask into Self-supervised Methods.

We argue that masking highly attended regions using $\mathbf{m}^{H}$ helps in learning powerful representations. In section 4, we also experiment with low-attended regions using $\mathbf{m}^{L}$, supporting further our argument.

AttMask can be incorporated into different methods to either replace the block-wise strategy of BEiT or introduce masking. For iBOT , we use $\widetilde{Z}^{v}$ in $L_{\textsc{mim}}$ (3) and $L_{\textsc{g}}$ (4). For DINO , we introduce masking by using $\widetilde{Z}^{v}$ for global views in $L_{\textsc{g}}$ (4), but not for local crops in the $L_{\textsc{lc}}$ (A13) loss (see Appendix Appendix 0.B).

Experiments

We pretrain iBOT and DINO on 20% and 100% of the ImageNet-1k training set. For 20%, we select the first 20% of training samples per class. We evaluate on ImageNet-1k validation set by $k$-NN or linear probing. For linear probing, we train a linear classifier on top of features using the same training protocol as in DINO . With linear probing, we also validate robustness against background changes on ImageNet-9 (IN-9) . For $k$-NN , we freeze the pretrained model and extract features of training images, then use a $k$-nearest neighbor classifier with $k=20$. We also perform the same $k$-NN experiment, now extracting features only from $\nu\in\{1,5,10,20\}$ examples per class. This task is more challenging and is similar to few-shot classification, only the test classes are the same as in representation learning.

We downstream to other tasks either with or without finetuning. We finetune on CIFAR10 , CIFAR100 and Oxford Flowers for image classification measuring accuracy; on COCO for object detection and instance segmentation measuring mean average precision (mAP); and on ADE20K for semantic segmentation measuring mean Intersection over Union (mIoU). Without finetuning, we extract features as with $k$-NN and we evaluate using dataset-specific evaluation protocol and metrics. We test on revisited $\mathcal{R}$Oxford and $\mathcal{R}$Paris for image retrieval measuring mAP ; on Caltech-UCSD Birds (CUB200) , Stanford Cars (CARS196) , Stanford Online Products (SOP) and In-Shop Clothing Retrieval (In-Shop) for fine-grained classification measuring Recall@$k$ ; and on DAVIS 2017 for video object segmentation measuring mean region similarity $\mathcal{J}_{m}$ and contour-based accuracy $\mathcal{F}_{m}$ .

In Appendix Appendix 0.A, we provide more benchmarks, visualizations and ablations.

Implementation Details.

As transformer encoder, we use ViT-S/16 . The attention map (7) is generated from the last layer of the teacher encoder by default, i.e., layer 12. We mask the input with probability $p=0.5$, while the mask ratio $r$ is sampled uniformly as $r\sim U(a,b)$ with $[a,b]=[0.1,0.5]$ by default. For AttMask-Hint, we set $s_{\max}=0.1$ and the show ratio $s$ is sampled uniformly from $[s_{\max}a,s_{\max}b]=[0.01,0.05]$. Following , we apply multi-crop scheme, as detailed in Appendix Appendix 0.B. The overall loss of iBOT is a weighted sum of $L_{\textsc{mim}}$ (3), with weight $\lambda$, and $L_{\textsc{g}}$ (4) $+$ $L_{\textsc{lc}}$ (A13) (DINO ), with weight 1, where $L_{\textsc{lc}}$ (A13) is the multi-crop loss. By default, $\lambda=1$. Hyperparameters are ablated in subsection 4.4. Training details are given in the Appendix Appendix 0.B.

2 Experimental Analysis

We provide an analysis on 20% of ImageNet-1k training samples, incorporating AttMask into distillation-based MIM or self-distillation only . We also provide results on robustness against background changes.

We explore a number of masking strategies using distillation-based MIM, by incorporating AttMask into iBOT . We compare AttMask with random block-wise masking , which is the default in iBOT, random patch masking with the same ratio, as well as with a more aggressive ratio, following MAE . AttMask masks the most attended tokens (AttMask-High) by default, but we also consider the least attended (AttMask-Low) and the most attended with hints (AttMask-Hint).

We evaluate performance using $k$-NN and linear probing evaluation protocol on the validation set, along with a fine-tuning evaluation on CIFAR10 and CIFAR100. As shown in Table 1, the AttMask-High outperforms all other masking strategies on all the evaluation metrics. In particular, AttMask-High achieves an improvement of +3.0% on $k$-NN and +1.5% on linear probing compared with the default iBOT strategy (random block-wise).

Interestingly, random patch masking outperforms the default iBOT strategy, while the more aggressive MAE-like strategy is inferior and AttMask-Low performs the lowest. Intuitively, this means that masking and reconstruction of non-salient regions does not provide a strong supervisory signal under a MIM objective. By contrast, our AttMask creates the more aggressive task of reconstructing the most salient regions and guides the model to explore the other regions. In this setup, AttMask-Hint is slightly lower than AttMask-High.

Data and Training Efficiency.

Self-supervised methods on vision transformers typically require millions of images, which is very demanding in computational resources. We advocate that being effective on less data and fast training are good properties for a self-supervised method. In this direction, we assess efficiency on less data and training time, still with iBOT training. In Table 2 we observe that our AttMask-High consistently outperforms the default random block-wise masking strategy of iBOT at lower percentage of ImageNet-1k training set. In addition, in Table 2, AttMask-High achieves the same performance as random block-wise with 42% fewer training epochs.

Masking Strategies: Self-distillation Only.

Here, we compare masking strategies using distillation only, without MIM reconstruction loss, by incorporating AttMask into DINO . That is, we apply only the cross-view cross-entropy loss on the [cls] token (4). In Table 3, AttMask-High improves $k$-NN by +0.5 compared with the default DINO (no masking), while AttMask-Low is inferior. This reveals that AttMask is effective even without a MIM loss. Moreover, AttMask-Hint is slightly better than AttMask-High in this setting.

Robustness Against Background Changes.

Deep learning models tend to depend on image background. However, to generalize well, they should be robust against background changes and rather focus on foreground. To analyze this property, we use ImageNet-9 (IN-9) dataset , which includes nine coarse-grained classes with seven background/foreground variations. In four datasets, the background is altered: Only-FG (OF), Mixed-Same (MS), Mixed-Rand (MR), and Mixed-Next (MN). In another three, the foreground is masked: No-FG (NF), Only-BG-B (OBB), and Only-BG-T (OBT).

In Table 4, we evaluate the impact of background changes on IN-9 and its variations, training iBOT under different masking strategies. We observe that, except for O.BB. and O.BT, AttMask-High is the most robust. On OBB and OBT where the foreground object is completely missing, AttMask-Hint exploits slightly better the background correlations with the missing object.

In the Appendix subsection 0.A.3, we provide visualizations of attention maps in Figure A5 and masking examples in Figure A7.

3 Benchmark

We pre-train iBOT with AttMask-High and AttMask-Hint on 100% of ImageNet-1k and compare it with baseline iBOT and other distillation-based methods.

As shown in Table 5(a), AttMask-High brings an improvement of 1% $k$-NN and 1.3% linear probing over baseline iBOT and is better than prior methods. AttMask-High is thus effective for larger datasets too. Table 5(b) shows results of the more challenging task where only $\nu\in\{1,5,10,20\}$ training examples per class are used for the $k$-NN classifier. In this case, AttMask-High is very effective, improving the baseline iBOT masking strategy by 3-4%, demonstrating the quality of the learned representation. In this setup, AttMask-Hint offers a further small improvement over AttMask-High. For simplicity though, we use AttMask-High by default as AttMask.

More results are given in the Appendix. In particular, in Table A13, we provide results similar to Table 5 but with pre-training for 300 epochs. Also, in Table A14 we report further supervised finetuning on ImageNet-1k. In Table A10, we investigate the quality of the patch features by using global average pooling (GAP) rather than the [CLS] token embeddings. In Table A11, we study the effect of masking salient image parts at inference.

Downstream Tasks with Fine-tuning.

We fine-tune the pre-trained models with iBOT and iBOT with AttMask for image classification on CIFAR10 , CIFAR100 and Oxford Flowers , object detection and instance segmentation on COCO , and semantic segmentation on ADE20K . In Table 6, we observe that AttMask brings small improvement on the baseline iBOT masking strategy on image classification fine-tuning in all cases. Furthermore, we observe that AttMask improves clearly the scores by 0.6% APb on object detection and 0.4% mIoU on semantic segmentation.

Downstream Tasks without Fine-tuning.

Without finetuning, we use the pretrained models with iBOT and iBOT with AttMask to extract features as with $k$-NN and we evaluate using dataset-specific evaluation protocol and metrics. As shown in Table 7(a,b), AttMask is very effective on image retrieval, improving by 1-3% mAP the baseline iBOT masking strategy on $\mathcal{R}$Oxford and $\mathcal{R}$Paris , on both medium and hard protocols. More impressive the performance on fine-grained classification, improving by 2-6% R@1 on all datasets, as shown in Table 8. Finally, AttMask improves on video object segmentation on DAVIS 2017 on all metrics, as shown in Table 7(c). These experiments are very important because they evaluate the quality of the pretrained features as they are, without fine-tuning and without even an additional layer, on datasets of different distribution than the pretraining set. AttMask improves performance by a larger margin in this type of tasks, compared with ImageNet.

In Table A12 in the Appendix, we additionally provide results for scene classification with linear probing on Places205 .

4 Ablation Study

We provide an ablation for the main choices and hyperparameters of our masking strategy and loss function, incorporating AttMask into iBOT and pre-training on 20% of ImageNet-1k training samples. We provide additional ablations in the Appendix. In Table A15, we examine the MIM loss weight. In Table A16, we ablate both the masking strategy and the mask ratio $r$.

The attention map (7) is generated from the last layer of the teacher encoder by default, that is, layer 12 of ViT-S. In Table 9(a), we aim to understand the impact of other layer choices on AttMask. We observe that the deeper layers achieve the highest $k$-NN performance. Although layer 11 works slightly better, we keep the choice of layer 12 for simplicity, since layer 12 embeddings are used anyway in the loss function.

Masking Probability and Mask Ratio.

We mask the global views with probability $p=0.5$ by default. Table 9(b) reports on other choices and confirms that this choice is indeed best. Therefore, it is useful that student network sees both masked and non-masked images.

The mask ratio $r$ is sampled uniformly as $r\sim U(a,b)$ with $[a,b]=[0.1,0.5]$ by default. Table 9(c) shows the sensitivity of AttMask with respect to the upper bound $b$, along with a fixed ratio $r=0.3$. AttMask is relatively stable, with the default interval $[0.1,0.5]$ working best and the more aggressive choice $[0.1,0.7]$ worst. This is possibly due to the foreground objects being completely masked and confirms that masking the most attended patches is an effective strategy. The added variation around the fixed ratio $r=0.3$ is beneficial.

Conclusion

By leveraging the self-attention maps of ViT for guiding token masking, our AttMask is able to hide from the student network discriminative image cues and thus lead to more challenging self-supervised objectives. We empirically demonstrate that AttMask offers several benefits over random masking when used in self-supervised pre-training with masked image modeling. Notably, it accelerates the learning process, achieves superior performance on a variety of downstream tasks, and it increases the robustness against background changes, thus revealing that it reduces background dependency. The improvement is most pronounced in more challenging downstream settings, like using the pretrained features without any additional learning or finetuning, or working with limited data. This reveals the superior quality of the learned representation.

We thank Shashanka Venkataramanan for his valuable contribution to certain experiments. This work was supported by computational time granted from GRNET in the Greek HPC facility ARIS under projects PR009017, PR011004 and PR012047 and by the HPC resources of GENCI-IDRIS in France under the 2021 grant AD011012884. NTUA thanks NVIDIA for the support with the donation of GPU hardware. This work has been supported by RAMONES and iToBos projects, funded by the EU Horizon 2020 research and innovation programme, under grants 101017808 and 965221, respectively.

References

Appendix 0.A More Experiments

We provide more benchmarks (subsection 0.A.1), more ablations (subsection 0.A.2), and more visualizations (subsection 0.A.3).

In contrast with the DINO objective that is applied only on the output [CLS] token embeddings, the MIM objective is directly applied to the output features of the patch tokens. Table A10 shows that using global average pooling (GAP) over patch features instead of the [CLS] token embeddings, AttMask outperforms baseline iBOT by 9.0% $k$-NN accuracy. This indicates that AttMask leads to a more challenging MIM objective, which in turn forces the ViT to learn more discriminative patch features.

Does AttMask Lead to Better Exploitation of Non-Salient Parts?

We examine the performance of the models pre-trained on 100% of ImageNet-1k on a more challenging ImageNet-1k validation set. In particular, we gradually mask the salient parts using the attention maps of the official pre-trained DINO ViT-Base model and setting the corresponding masked pixel values to zero. Our assumption is that a more robust model should be less sensitive when salient parts of an object are missing. In Table A11, we observe that as more parts of the images are hidden, a larger gain occurs by using AttMask with iBOT. This indicates that AttMask leads to less sensitive models that exploit better the non-salient parts or even background context.

Downstream Tasks using Linear Probing.

We experiment on scene classification on Places205 , measuring classification accuracy, using linear probing evaluation on models pre-trained on 100% of ImageNet-1k for 100 epochs. In Table A12, we observe that AttMask improves scores by 0.8% accuracy.

Training for More Epochs.

We train iBOT with AttMask on 100% of ImageNet-1k for 300 epochs. AttMask not only accelerates the learning process and has better performance on data-limited regimes as explained in the main paper, but as we see in Table A13(a), even when trained for many epochs and with many data, it still brings an improvement of 0.4% $k$-NN and 0.1% linear probing over baseline iBOT . Also, AttMask outperforms all other state-of-the-art frameworks on linear probing evaluation on ImageNet-1k validation set. We highlight that MST employs an additional CNN decoder, while AttMask achieves improved linear probing performance with fewer learnable parameters.

We argue that the higher improvement of AttMask $k$-NN compared with linear probing indicates higher quality of learned embeddings, since linear probing amounts to supervised classification on higher-dimensional embeddingsWe remind that, following the evaluation setups of DINO for ViT-S, for linear probing we use the concatenated features from the last 4 layers of ViT while for $k$-NN the feature from only the last layer. So, linear probing uses 4 times higher-dimensional features and on the same dataset that was used for self-supervised pre-training. To validate this, we experiment with a more challenging variant of $k$-NN where only $\nu\in\{1,5,10,20\}$ examples per class of the training set are used. Table A13(b) shows that using AttMask for self-supervised pre-training and then using only simple $k$-NN classifier with only one example per class, achieves an accuracy improvement of 1.5% compared with the default iBOT. This highlights the superiority of AttMask in low-shot learning regimes, which are of great practical interest.

Full fine-tuning on ImageNet-1k.

For iBOT and iBOT+AttMask pre-trained on ImageNet-1k for 300 epochs, we also experiment with further supervised fine-tuning on ImageNet-1k, training for 100 epochs. We report results in Table A14. AttMask improves iBOT by 0.2% (81.1% $\rightarrow$ 81.3%), providing a better network initialization for supervised finetuning.

A.2 More Ablations

The overall loss of iBOT is a weighted sum of $L_{\textsc{mim}}$ (3), with weight $\lambda$, and $L_{\textsc{g}}$ (4) $+$ $L_{\textsc{lc}}$ (A13) (DINO), with weight 1. Table A15 shows that AttMask is superior to the default block-wise random masking of iBOT in all cases, while the default $\lambda=1$ works best for both and yields the greatest gain of $3\%$ $k$-NN accuracy for AttMask. In particular, increasing the weight of the MIM loss leads to a larger gain in $k$-NN accuracy. This shows that AttMask boosts the MIM task.

Masking strategy and mask ratio.

We ablate both the masking strategy (random block-wise, random or AttMask-High) and the mask ratio $r$ in Table A16. AttMask-High with 10-50 mask ratio gives the best results.

A.3 More Visualizations

In Figure A5, we utilized the pre-trained models on 20% of ImageNet and observe that, when training iBOT with the default block-wise random masking strategy, there is at least one head (in blue) that attends the background to a great extent. By contrast, with our AttMask, all heads mostly attend salient objects or object parts. It appears that by focusing on reconstructing highly-attended masked tokens, the network learns to focus more on foreground objects.

Visualization of Masking Examples.

We illustrate the effect of mask ratio $r$ (%) to various masking strategies in Figure A6 and Figure A7. While random Block-Wise and Random masking fail to consistently mask informative parts of an image, AttMask-High and AttMask-Hint make use of attention to hide salient and all but very salient parts respectively. This gives rise to a more challenging MIM task.

Appendix 0.B Experimental Setup

We provide more details on the experimental setup, including multi-crop, training details and evaluation details.

Following , we apply the multi-crop strategy to generate a set of $m$ low-resolution local crops, which cover only small parts of the image, tokenized as $Z^{c}_{1},\dots,Z^{c}_{m}$. Similar to $L_{\textsc{g}}$ (4), the loss is applied globally on the [cls] tokens, in particular between the student output for a local crop $Z^{c}_{j}$ and the teacher output for a global view $Z^{v}$, both of which are non-masked:

The overall loss is a weighted sum of $L_{\textsc{mim}}$ (3), $L_{\textsc{g}}$ (4) and $L_{\textsc{lc}}$ (A13).

Training Details.

For our analysis and ablation (subsection 4.2, subsection 4.4 and subsection 0.A.2), we pre-train models on 20% of ImageNet-1k for 100 epochs. For both iBOT and DINO we use AdamW as optimizer. Unless otherwise stated, we use the ViT-S/16 architecture and a batch size of 240. We warm-up learning rate $\eta$ for 10 epochs following the linear scaling rule $\eta=5\times 10^{-4}\times\texttt{bs}/256$ where bs is the batch size and then decay using a cosine schedule. We also use a cosine schedule from $0.04$ to $0.4$ for weight decay. We set teacher momentum to $0.99$ and student temperature to $0.1$. We use a linear warm-up for teacher temperature from $0.04$ to $0.07$ for the first 30 epochs following DINO.

All methods in subsection 4.2, subsection 4.4 and subsection 0.A.2 use the multi-crop scheme with two $224^{2}$ global crops and six $96^{2}$ local crops that approximately scale the training time by a factor of $\gamma=2+6\times(96/224)^{2}=3.10$. We use color jittering, Gaussian blur and solarization as data augmentations. Local crops scales are sampled from $(0.05,s)$ and global crop scales from $(s,1)$. We set $s$ to $0.4$ for DINO and $0.25$ for iBOT. We set the dimensionality of the head output to $65536$ for DINO, while for iBOT, we use a shared projection head for [CLS] and patch tokens, of dimensionality $8192$. We do not perform weight normalization on the last layer of the MLP heads.

For our benchmark (subsection 4.3 and subsection 0.A.1), we pre-train models on 100% of ImageNet-1k for 100 and 300 epochs. For the 100-epoch experiments, the setup is the same as on 20% of ImageNet-1k except for increasing the teacher momentum to $0.996$ and the number of local crops to ten. The scaling factor of the training time in this case is $\gamma=2+10\times(96/224)^{2}=3.84$. For the 300 epochs experiments, we increase the batch size to 800 and set $s$ to $0.32$, similar to the iBOT default scale.

Evaluation Details.

For the ImageNet-1k evaluation, we use $k$-NN and linear probing as in DINO and iBOT . We evaluate on ImageNet-1k validation set. For $k$-NN, we use the [CLS] feature from the last ViT layer and set $k$ to 20. For linear probing, we train a linear classifier using SGD with a batch size of 1024 for 100 epochs. We set learning rate to $0.003$ and do not apply weight decay. We apply random resized crops and horizontal flips as data augmentations and keep the central crop. Following DINO and iBOT , we use the concatenation of the [CLS] features from the last four layers as input to the linear classifier.

For the evaluation of downstream tasks with finetuning, we train models on CIFAR10, CIFAR100 for 500 epochs and on Oxford Flowers for 1000 epochs. We set learning rate to $7.5\times 10^{-6}$, weight decay to $0.05$ and use a batch size of 900.

On COCO , we evaluate the performance of object detection and instance segmentation downstream tasks. We consider the COCO 2017 set, which contains $118$K training images, $5$k validation and $20$ test-dev. We consider the Cascade Mask R-CNN as task layer and follow the setup from . We use the hyper-parameter configuration from : multi-scale training (resizing image with shorter size between $480$ and $800$, with the longer side no larger than $1333$). We use AdamW with initial learning rate $10^{-4}$, the $1\times$ schedule (12 epochs with the learning rate decayed by $10\times$ at epochs $9$ and $11$) and weight decay $0.05$. Unlike , where training is on $8$ GPUs with $4$ images per GPU, we use $2$ images per GPU due to hardware limitations. For a fair and direct comparison, we fine-tune iBOT baseline with the same configuration.

We evaluate on ADE20K for the semantic segmentation downstream task. It consists of 25k images in 150 classes, with 20k for training, 2k for validation and 3k for testing. We rely on UperNet as task layer and fine-tune the entire network following the setup from : $160$k iterations with $512\times 512$ images. We do not perform multi-scale training and testing. We adopt the same hyper-parameters as in . We use the AdamW optimizer with an initial learning rate of $7\times 10^{-4}$ with poly-scheduling, layer decay rate $0.65$ and weight decay $0.05$. We train on $8$ GPUs with $2$ images per GPU.

For the evaluation of downstream tasks without finetuning, we follow the protocol of DINO on $\mathcal{R}$Oxford, $\mathcal{R}$Paris and DAVIS 2017 . On Caltech-UCSD Birds (CUB200) , Cars (CARS196) , Stanford Online Products (SOP) and In-Shop Clothing Retrieval (In-Shop) , we extract features from test set images and directly apply nearest neighbor search to measure Recall@$k$ . On Places205 , we train a $205$-way linear classifier on pre-cached features, using only horizontal flip as augmentation. Training is with SGD for $50$ epochs using an initial learning rate of $0.01$ that is decreased to with cosine schedule, a batch-size of $1024$, and no weight decay.