Improving Masked Autoencoders by Learning
Where to Mask

Unlike language modeling, the information density of images is more sparse and usually dominated by the objects that appear in the image. Therefore we conduct a preliminary experiment to explore whether introducing the object location priors with high information density can help representation learning. We randomly select a subset ($10\%$) of the ImageNet dataset [7] and use the provided true object bounding boxes to indicate object locations. Instead of random masking, we manually raise the dropped probability of patches within the bounding box and use these samples for self-supervised pretraining. We experiment with different values and show the linear probing results in Fig. 2, where $\beta$ represents the additionally raised probability for patches within the bounding box against those in the background. We have two observations from Fig. 2:

• •

  Compared with random masking, slightly increasing the masking probability for patches within the bounding box significantly improves the linear probing results.

• •

  Excessively raising the masking probability of patches within the bounding box may affect the validity of the pretext task and degrades the performance.

These results indicate that different image patches provide different effects on the learning of visual representations, and learning to reconstruct more foreground objects with higher information density can lead to better pretraining results. Nonetheless, we need to strike a balance between the information gain through masked image modeling and its training difficulty. Inspired by this, a natural solution is to integrate mask generation and image reconstruction into a fully differentiable framework. For different images, it can adaptively provide masking weights that reflect prior object hints and are friendly to masked image modeling, that is, we want to find patches with lower reconstruction difficulty within the obtained areas with higher information density.

Motivated by our observations in Section 2.2, we design a novel self-supervised framework that can learn to select more informative foreground patches to mask, in which two key challenges need to be properly handled: First, how to guide the ViT model mining the informative patches without any explicit supervision? Second, how to control the difficulty of the pretext task via an adaptively learned mask strategy?

For the first challenge, we propose a differentiable mask generator $G$ to generate the mask image which contains the important weight for each image patch. The mask generator contains a pretrained ViT encoder (also in a self-supervised manner) and two trainable convolution layers. Given an input image, we use the pretrained ViT encoder to extract the multi-head attention maps from the last transformer block. The obtained attention maps contain indications of different semantic regions of the image and are further processed by the convolution layers for mask generation. Since we cannot obtain explicit supervision to highlight the location of the foreground object, we introduce an object-centered adversarial training strategy to guide the mask generator in producing higher weights on patches within possible foreground objects. Specifically, we introduce a discriminator and generate real mask images by manually raising the masking probabilities of a randomly sampled rectangle area. In this way, we minimize the distribution distance between the generated mask and the sample real mask, which helps the model find the informative foreground patches.

For the second challenge, we directly propagate the gradients from the masked autoencoder back to the mask generator and train these two modules synchronously. In other words, the reconstruction loss applied for the masked autoencoder also constrains the mask generator from generating too hard mask images. By this means, the masked autoencoder encourages the mask generator to produce high weights on patches that can easily infer the masking information, while the discriminator regularizes the mask generator and focuses more on possible foreground objects instead of the background. These two branches cooperate with each other to facilitate the masked autoencoder to learn more representative patch relations.

The framework of AutoMAE is shown in Fig. 3. The differentiable mask generator $G$ aims to generate sample-specific mask images to guide the token selection process in the training of the masked autoencoder, which mainly consists of a pretrained ViT encoder and two trainable convolutional layers. We first exploit an MAE-pretrained [2] ViT encoder¹¹1Different self-supervised pretrained ViT models can be used for attention map extraction. Their effects are compared in our experiment. In the rest of this paper, we use the MAE-pretrained ViT-Base model without loss of generality. to embed the image $\boldsymbol{I}\in\mathbb{R}^{c\times h\times w}$ and extract the multi-head self-attention map $\boldsymbol{A}\in\mathbb{R}^{H\times\frac{h}{p}\times\frac{w}{p}}$ between the query embedding of the [CLS] token and key embeddings of other patch tokens from the last transformer block, where $H$ represents the number of heads. Specifically, the calculation of the $i$-th attention map $\boldsymbol{A}_{i}\in\mathbb{R}^{1\times\frac{h}{p}\times\frac{w}{p}}$ can be formalized as $\boldsymbol{A}_{i}=\texttt{softmax}(\boldsymbol{q}_{i}^{c}\boldsymbol{K}_{i}^{\top}/\sqrt{d^{\prime}}),$ where $\boldsymbol{q}_{i}^{c}\in\mathbb{R}^{1\times d^{\prime}}$ represents the query embedding of the [CLS] token, $\boldsymbol{K}_{i}\in\mathbb{R}^{\frac{hw}{p^{2}}\times d^{\prime}}$ represents the key embeddings of other patch tokens, and $d^{\prime}=d/H$ represents the embedding dim. Then we reshape the attention map with spatial size equal to $1\times\frac{h}{p}\times\frac{w}{p}$. As mentioned in [2], different attention maps in the last transformer block can attend to different semantic regions of the image. Therefore these attention maps can be served as informed initialization for further mask generation. Besides, the parameters of the pretrained ViT encoder are frozen during the training process. Given the multi-head self-attention map $\boldsymbol{A}$, we further exploit two convolutional layers with ReLU activation $f_{\theta}$ to embed these attention maps $\boldsymbol{F}=f_{\theta}(\boldsymbol{A}),$ where $\boldsymbol{F}\in\mathbb{R}^{1\times\frac{h}{p}\times\frac{w}{p}}$. We then sample the final mask image $\boldsymbol{M}\in\mathbb{R}^{1\times\frac{h}{p}\times\frac{w}{p}}$ given the embedding feature $\boldsymbol{F}$:

where $\log(softmax(\cdot))$ is used to stabilize the training process and $z$ represents a random noise sampled from a Gumbel distribution. The values in $\boldsymbol{M}$ are continuous and can be seen as the important weight of each image patch. We use an extra discriminator and the masked autoencoder itself to supervise the generation process.

Although the attention map $\boldsymbol{A}$ can provide indications for different semantic regions, how to supervise the mask generator to produce a proper pretext task that can focus on foreground objects is still a challenging problem. To tackle this problem, we introduce an extra discriminator and design an object-centered adversarial training strategy to regularize mask generation. The key insight of this approach is to generate pseudo mask images $\boldsymbol{M}^{p}$ as “real” samples which can imitate the difference between foreground and background patches, and use the adversarial training strategy to minimize the distribution shift between generated mask $\boldsymbol{M}$ and “real” samples $\boldsymbol{M}^{p}$.

The way to generate pseudo-mask images is quite simple. Given a zero-initialized image $\boldsymbol{M}^{p}\in\mathbb{R}^{1\times\frac{h}{p}\times\frac{w}{p}}$, we first randomly sample a rectangle area within this image whose spatial size is between $20\%\sim 80\%$ of the whole image area. This sampled rectangle area is assumed to be the bounding box of one pseudo object. Then we flatten the pseudo mask $\boldsymbol{M}^{p}=\{m_{i}^{p}\}_{i=1}^{n}$ and suppose the index set within the sample rectangle area is $\boldsymbol{X}$, the values in the pseudo mask are sampled as follows:

where $\epsilon$ represents a random noise sampled from a uniform distribution $\epsilon\sim U(0,1)$, and $\alpha$ represents the additional important weight for patches within the rectangle, which is set to $0.5$ in our experiments. In this way, we generate the pseudo mask images where foreground patches contain higher weights compared with background patches. We then use the pseudo-mask images as a real sample for adversarial training with the loss function from LSGAN [23]:

where $D$ is the discriminator and we set $a=-1,b=1,c=0$ as in previous work.

Given the sampled mask $\boldsymbol{M}$, we suggest that patches with higher weights are more likely to be foreground objects and more informative than patches with lower weights. Therefore, we raise the masking probabilities of these high-weighted patches and use low-weighted patches as visible hints in masked image modeling. Specifically, we first sort the values in $\boldsymbol{M}=\{m_{1},m_{2},\cdots,m_{n}\}$ and then get the top $K$ index set $\boldsymbol{Y}$. The unnormalized masking probability of each token can be sampled as follows:

where $\epsilon$ represents a random noise sampled from a uniform distribution $\epsilon\sim U(0,1)$, and $\beta$ represents the additionally raised probability for tokens with higher weights. In our experiment, we set $K=n/4$ and $\beta=0.5$ where $n$ represents the number of patch tokens. In this case, patches in set $\boldsymbol{Y}$ are supposed to be more informative than other patches, and models are supposed to pay more attention to these patches.

After that, we use the sampled probability to guide the token selection process where tokens with the top $75\%$ biggest probability are dropped and the remaining tokens with $25\%$ smallest probability are saved as visible hints for masked image modeling. In other words, we still follow the mask ratio used in the original masked autoencoder while using the sampled probabilities to enhance the learning of more important patches.

Particularly, we slightly change the formulation of the input patch tokens to allow the gradients from the masked autoencoder can be propagated back to the mask generator. Given the embedded patch tokens $\boldsymbol{Z}=\{\boldsymbol{z}_{i}\}_{i=1}^{n}$ and the generated mask image $\boldsymbol{M}=\{m_{i}\}_{i=1}^{n}$, the new embedded patch tokens $\boldsymbol{Z^{\prime}}=\{\boldsymbol{z^{\prime}}_{i}\}_{i=1}^{n}$ are obtained via $\boldsymbol{z^{\prime}}_{i}=\boldsymbol{z}_{i}\cdot m_{i}+\boldsymbol{z}_{i}\cdot\texttt{sg}(1-m_{i}),$ where $\texttt{sg}(\cdot)$ represents the stop gradient operation. We conduct token selection on these new embedded tokens $\boldsymbol{Z^{\prime}}$ and then perform the following encoding and decoding stages which are the same as the original masked autoencoder. In this way, we preserve the values of patch tokens while allowing the masked autoencoder to influence mask generation. Combined with the adversarial loss, the full objective function of the mask generator can be written as $\mathcal{L}_{\text{G}}=\mathcal{L}_{\text{recon}}+\lambda\mathcal{L}_{\text{adv}}$ where we set $\lambda=0.2$ by grid search.

When we have finished the preliminary model analyses on small subsets, we start to train AutoMAE on the full ImageNet-1K and analyze the experimental results.

To verify whether the mask generator can perceive meaningful areas, we conduct experiments on ImageNet-9 [30], a dataset with foreground objects only. As shown in Fig. 4, the most high-weighted masks in AutoMAE mainly focus on meaningful patches, such as the back and eyes of the jaguar, and gradually involve less significant patches as the masking ratio increases. More results are provided in supplementary materials.

In Table S1, we compare the use of different self-supervised warmup methods (MAE-800 vs. MAE-400-Small) for the ViT-Base feature extractor in the mask generator. The latter one is pretrained on the $30\%$ subset for $400$ epochs. It achieves a comparable Top-1 accuracy ($53.57\%$) with the one with the MAE-800 backbone ($53.58\%$), which is significantly higher ($+2.89$) than its baseline MAE ($50.50\%$) and the best MAE model with ground-truth bounding box guidance (as described in 2.2 of the paper). These results indicate that the main improvement in the performance of AutoMAE comes from the differentiable mask sampling method, as the mask generator is effective even with a weak feature extractor.

In most of the experiments, the mask generator of AutoMAE relies on a pretrained ViT. We here consider removing such a requirement by using a train-from-scratch ViT backbone. The simplest way is to reuse the same ViT model for both mask generation and mask reconstruction. As shown in Table S2, this model (“From-Scratch”) remarkably outperforms the vanilla MAE model ($26.00\%$ vs $23.07\%$), both pretrained on the $10\%$ ImageNet subset. It suggests that AutoMAE is effective even without any pretrained models. We further perform the exponential moving average (EMA) to synchronize the ViT used for mask generation ($k$) and the one used for image reconstruction ($q$). We have $\theta_{k}=m\times\theta_{k}+(1-m)\times\theta_{q}$. Results are shown in Table S2.

We conduct an ablation study of AutoMAE by separating it into two independent learning processes, i.e., mask generation and mask-guided image modeling. In Table S1, the “Two-stage” model outperforms MAE ($52.57\%$ vs $50.50\%$) but underperforms the joint training version ($53.58\%$), which validates the effectiveness of joint training in AutoMAE. In Table 1 of the paper (full-set pretraining), we compare AutoMAE with SemMAE, a two-stage pretraining pipeline with an independent mask learning stage. It further supports the superiority of joint training over two-stage training ($66.7\%$ vs. $65.0\%$).

As DINO [2] suggests that the last self-attention layer in ViT blocks can provide meaningful masks, we remove the adversarial training loss and replace the two-layer CNN in the mask generator (after the feature extractor) with a simple $\mathrm{maximum}$ operator over the channel axis. From Table S1, although these models (“MAE-Max & DINO-Max”) still outperform the original MAE, they obtain lower accuracy than corresponding AutoMAE models with MAE-800 and DINO-B. It indicates that the mask generator does learn a better masking function than simply taking the maximum attention values.

In Fig. 2 in the main manuscript, we show how the margin value of the masking probability $\beta$ affects the linear probing performance with bounding box guidance. Here, as shown in Table S3, we conduct a grid search over $\beta=\{0.2,0.5,1.0\}$ on the ImageNet subset and find that $\beta=0.5$ performs best.

In adversarial training, we use a three-layer CNN as the discriminator. The first two convolutional layers have a $4\times 4$ kernel size, a stride of $2$, and a zero-padding length of $1$. They are followed by a Spectral Normalization [33] layer and a LeakyReLU activation. The LeakyReLU has a negative slope of $0.2$. The last convolutional layer has a filter size of $3\times 3$ and is also followed by a spectral normalization layer. As shown in Fig. S1, the spatial size of feature maps is $14\rightarrow 7\rightarrow 3\rightarrow 1$ and the number of channels is $1\rightarrow 16\rightarrow 32\rightarrow 1$.

We adopt the AdamW optimizer [22] for pretraining the MAE model and the mask generator in our AutoMAE. Following previous work [12], the base learning rate is set to $1.5\times 10^{-4}$, the momentum is configured as $\beta_{1}=0.9,\beta_{2}=0.95$, and the batch size is $4{,}096$. We use the linear learning rate schedule: lr = base_lr $\times$ batchsize $/256$. The weight decay is $0.05$. The learning rate is scheduled by cosine decay with a warm-up phase of $40$ epochs and the total number of pretraining epochs is set to $800$. We use random resized cropping and random horizontal flipping for data augmentation. For the discriminator, we use the Adam optimizer without the weight decay. Other hyper-parameters stay the same. The mask generator uses a margin of $\alpha=0.5$ to highlight the randomly sampled foreground area, and the loss weight $\lambda$ is set to $0.2$ to balance generator loss with reconstruction loss.

In Eq.(4) and Eq.(6), we set $\alpha=\beta$ as they have similar effects. In Eq.(5), the values of $a,b,c$ are directly taken from LSGAN [23]. We only tune $\beta$ and the loss weight $\lambda$ on the ImageNet subset ($30\%$) and reuse the obtained values in other datasets/setups.

We use a patch size of $16\times 16$ for all experiments. It takes $26$ hours to pretrain an AutoMAE model on $30\%$ of ImageNet with $8$ NVIDIA RTX 3080Ti cards for $400$ epochs, and $7$ days to pretrain an AutoMAE model on the full ImageNet dataset with $8$ NVIDIA A100 cards for $800$ epochs. Admittedly, MAEs commonly benefit from a smaller $8\times 8$ patch size. However, we cannot afford its training cost, e.g., it takes about $4$ weeks for $800$ epochs on the full ImageNet-1K. We want this paper to focus more on the proposed training scheme than on hyper-parameter tuning. Therefore, we follow the standard $16\times 16$ setup [13, 37], which does not affect the main conclusions when compared to SemMAE. As shown in Table S4, AutoMAE has comparable time cost to recent advances in self-supervised pre-training [16, 18] based on the same hardware ($8\times$A100 GPUs) for $800$ epochs. The full model has a similar number of FLOPs to existing work.