#  Green Hierarchical Vision Transformer
for Masked Image Modeling

Lang Huang1, Shan You2, Mingkai Zheng3, Fei Wang2, Chen Qian2, Toshihiko
Yamasaki1
1The University of Tokyo; 2SenseTime Research; 3The University of Sydney
<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>Corresponding author. †Code and pre-trained models:
https://github.com/LayneH/GreenMIM.

###### Abstract

We present an efficient approach for Masked Image Modeling (MIM) with
hierarchical Vision Transformers (ViTs), _e.g._ , Swin Transformer [43],
allowing the hierarchical ViTs to discard masked patches and operate only on
the visible ones. Our approach consists of two key components. First, for the
window attention, we design a Group Window Attention scheme following the
Divide-and-Conquer strategy. To mitigate the quadratic complexity of the self-
attention w.r.t. the number of patches, group attention encourages a uniform
partition that visible patches within each local window of arbitrary size can
be grouped with equal size, where masked self-attention is then performed
within each group. Second, we further improve the grouping strategy via the
Dynamic Programming algorithm to minimize the overall computation cost of the
attention on the grouped patches. As a result, MIM now can work on
hierarchical ViTs in a green and efficient way. For example, we can train the
hierarchical ViTs about 2.7$\times$ faster and reduce the GPU memory usage by
70%, while still enjoying competitive performance on ImageNet classification
and the superiority on downstream COCO object detection benchmarks.†

##  1 Introduction

Driven by the great success of Masked Language Modeling (MLM) [50, 51, 13, 5]
in natural language processing (NLP) and the advancement of Vision Transformer
(ViT) [15, 43, 60, 69], Masked Image Modeling (MIM) emerged as a promising
self-supervised pre-training paradigm for computer vision (CV). MIM learns
representations from unlabelled data by masked prediction, _e.g._ , predicting
the discrete tokens [2], the latent features [73, 62, 1], or the raw pixels
[22, 66] of the randomly masked input image patches. Among them, the
representative work Masked Autoencoder (MAE) [22] exhibited competitive
performance as well as impressive efficiency. In essence, MAE proposed an
asymmetric encoder-decoder architecture for MIM, where the encoder (_e.g._ , a
standard ViT model [15]) operates only on visible patches, and the light-
weight decoder recovers all patches for mask prediction.

On the one hand, the asymmetric encoder-decoder architecture significantly
reduces the computation burden of pre-training. On the other hand, MAE only
supports the isotropic ViT [15] architecture as the encoder, while most of the
modern vision models adopt hierarchical structure [37, 25, 43], in part due to
the need of handling the scale-variations of visual elements. In fact, the
hierarchical structure and local inductive bias are crucial in various CV
tasks that require representations of different levels or scales to make
predictions, including image classification [25], and object detection [19].
Yet it is still not straightforward how the hierarchical vision transformers,
_e.g._ Swin Transformer [43], can be integrated into the MAE framework.
Moreover, though the work SimMIM [66] has explored Swin Transformer for MIM,
it operates on both visible and masked patches and suffers from heavy
computation costs compared with MAE. As a concrete example, we find that even
the base size model of SimMIM cannot be trained on a single machine with eight
32GB GPUs, let alone the larger ones. The computation burden makes it
difficult for a wider range of researchers to dive into this field of
research, not even to mention the amount of carbon emission during the model
development process.

To this end, we strive to devise a new and green approach for MIM with
hierarchical models, in the spirit of Green AI [55, 67]. Our work focuses on
extending the asymmetric encoder-decoder architecture of MAE to hierarchical
vision transformers, particularly the representative model Swin Transformer
[43], for the sake of efficient pre-training on visible patches only. We
identify that the major obstacle is the limitation of the local window
attention. Though widely used in hierarchical vision transformers, the local
window attention does not work well with the random masking as it creates
various-sized local windows that are infeasible for computing in parallel.

This paper provides the first attempt to address this drawback. Our
methodology is conceptually simple and consists of two components. First,
guided by the Divide-and-Conquer principle, we present a Group Window
Attention scheme first partitioning the local windows with uneven numbers of
visible patches into several equal-sized groups and then applying the masked
attention within each group. Second, we formulate the aforementioned group
partition as a constrained optimization problem, where the objective is to
find a group partition that minimizes the computation cost of the attention on
the grouped tokens. Inspired by the concept of Dynamic Programming [4] and the
greedy principle, we propose an Optimal Grouping algorithm that adaptively
selects the optimal group size and partitions the local windows into a minimum
number of groups. Our methodology is generic and does not make _any_
modification to the architecture of the backbone models, such that we can make
apple-to-apple comparisons with the baseline operating on both visible and
masked patches. In our experimental evaluation, we observed that our method
requires substantially less training time and consumes much less GPU memory
while performing on par with the baseline. Concretely, using Swin-B [43], our
method requires only half of the training time and about 40% of GPU memory
consumption compared with the baseline SimMIM, while achieving 83.7% top-1
fine-tuning accuracy on the ImageNet-1K [54] that is on par with SimMIM.

Figure 1: Comparison with SimMIM in terms of efficiency. All methods use a
Swin-B/Swin-L backbone and batch size of 2,048. The experiments of our method
are conducted on a single machine with eight V100 GPUs, CUDA 10, and PyTorch
1.8, while those of SimMIM require 2 or 4 machines.

Our contributions. We make the following main contributions:

  1. 1.

We design a Green Hierarchical Vision Transformer for Masked Image Modeling
that advocates a more practical method with drastically improved efficiency.

  2. 2.

As shown in Figure 2, we present a Group Window Attention scheme, gathering
the local windows with different numbers of visible patches into several
equal-sized groups and then applying the masked attention within each group
(see Figure 3).

  3. 3.

Inspired by the concept of Dynamic Programming [4] and the principle of the
greedy algorithm, we propose an Optimal Grouping algorithm (see Algorithm 1)
that adaptively selects the optimal group size and partitions the local
windows into the minimum number of groups.

  4. 4.

Extensive experimental evaluations on the ImageNet-1K [54] and the MS-COCO
[41] datasets demonstrate that our method can obtain comparable performance to
the baselines with more than $2\times$ efficiency improvement (see Figure 1,
Tables 2 and 3).

##  2 Related Works

Self-Supervised Learning. Representation learning is a long-standing and
fundamental question in CV. For a long time, representation learning had been
dominated by supervised learning. Until recent three years, self-supervised
learning (SSL) exhibited impressive performance and attained significant
attention. Generally, SSL solves a proxy task without the actual interest to
learn good representations. According to the proxy tasks, SSL methods can be
categorized into generative approaches and discriminative approaches.
Generative approaches predict the original data based on the partially
observed inputs [59, 49, 38], predict the transformation applied to the input
[46, 18], or model pixels in the input space [36, 20, 33, 27]. Masked Image
Modeling also falls into this category. Discriminative approaches received
more interest during the past few years, especially the contrastive learning
methods. Contrastive learning creates multiple views of images with a set of
random data distortions and encourages the representations to be invariant to
the distortions. A large number of contrastive learning approaches [63, 47,
23, 9, 71, 72] drive the training by maximizing the similarities between
positive samples (_i.e._ , views from the same image) while minimizing those
between the negative samples, and some works simply get rid of the negative
pairs [21, 6, 10, 31, 70, 7]. Beyond contrastive learning on global features,
several methods proposed to maintain the spatial information of
representations and use region/mask/pixel-level contrastive learning [61, 65,
64, 26, 29].

Masked Language/Image Modeling. Self-supervised pre-training has
revolutionized the NLP. Among them, the Masked Language Modeling (MLM)
proposed in BERT [13] and its variants [5] are the most dominant methods,
which learn representations by predicting the tokens that are randomly masked
in the input sentence. Masked Image Modeling has a similar idea of predicting
corrupted images, and some of these methods [59, 49] were proposed Even
preceded to BERT. These methods were, however, unable to perform on par with
other pre-training paradigms at that time. Until recently, aided by the
significant advancement of Vision Transformers [15], several MIM methods
presented promising results [15, 2, 22, 66, 62] and became the state-of-the-
art of self-supervised learning in CV. These methods can be roughly
differentiated in terms of the prediction target, _e.g._ , color bins [15],
discrete tokens [2, 14] from pre-trained VAEs [57, 52], raw pixels [22, 66],
and handcrafted features [62]. Among these approaches, MAE [22] exhibited
competitive performance as well as impressive efficiency as it discards the
masked tokens and operates only on the visible ones.

Isotropic and Hierarchical Vision Transformers. The seminal work Vision
Transformer (ViT) [15] revolutionized the conventional view of images. ViT,
and its variants [56], treat an image as a sequence of patches and adopt a
pure Transformer [58] backbone to model the patches for image classification,
achieving an impressive performance even compared with the Convolutional
Neural Networks. Nevertheless, while the results of ViT are promising in
classification, its performance on dense prediction tasks is less favorable,
which is largely due to its low-resolution feature maps inherited from its
isotropic structure and the quadratic complexity of self-attention [58]. To
this end, a strand of works proposed hierarchical structure [60, 43, 16, 69]
and efficient attention [30, 43, 12, 32] for ViTs, unleashing the potential of
ViTs as general-purpose vision backbones. Our work performs studies upon the
representative method Swin Transformer [43] with (shifted) window local
attention that is also widely adopted in other methods [12, 69, 32].

Green AI. Witnessing the exponential growth of computations of big AI models
[13, 5, 42], the concept of Green AI attains mounting attention in recent
years [55, 67]. Rather than being merely obsessed with accuracy, Green AI
advocates making efficiency an important measure of AI models, championing the
greener approaches that are more inclusive to the research community. This
work follows the path of Green AI and presents a greener approach for MIM with
hierarchical ViTs.

##  3 Approach

###  3.1 Preliminary

Notations. Let $\mathbf{X}\in\mathbb{R}^{C\times H\times W}$ denote the input
feature where $C$, $H$, and $W$ are the numbers of channels, height, and width
of $\mathbf{X}$; $\mathbf{M}\in\\{0,1\\}^{H\times W}$ denotes the (spatial)
mask generated randomly during training where $0$ indicates a patch is
invisible for the encoder, and vice versa.

Figure 2: Illustration of the Group Window Attention scheme. In Masked Image
Modeling (MIM), the input $\mathbf{X}$, where different colors indicate the
tokens belong to different local windows, is randomly masked, producing
$\widehat{\mathbf{X}}$ of which most tokens are invisible. Our Group Window
Attention first performs an optimal grouping to group the visible tokens of
the local windows into several equal-sized groups, forming
$\overline{\mathbf{X}}$. Finally, we perform the Masked Attention within each
group to ensure no inter-window information leakage.

Masked Image Modeling. MIM learns representations by predicting the masked
portion of a input $\mathbf{X}$ from its partial observation
$\widehat{\mathbf{X}}\leftarrow\mathrm{Mask}(\mathbf{X},\mathbf{M})$. Existing
MIM methods fall into two categories regarding the
$\mathrm{Mask}(\cdot,\cdot)$ operation. Most methods [2, 66, 62] use the
Hadamard product for masking and retain the masked patches, _i.e._ ,
$\widehat{\mathbf{X}}\leftarrow\mathrm{Mask}(\mathbf{X},\mathbf{M})=\mathbf{X}\odot\mathbf{M}$
with $\mathbf{M}$ broadcasted for $C$ times along the channel dimension. In
sharp contrast to these methods, Masked Autoencoders (MAE) [22] proposes
throwing the masked patches at the masking stage, _i.e._ ,

$\displaystyle\widehat{\mathbf{X}}\leftarrow\mathrm{Mask}(\mathbf{X},\mathbf{M})=\\{\mathbf{X}_{i,j}:\mathbf{M}_{i,j}=1\\}.$
(1)

MAE designs an asymmetric and isotropic encoder-decoder architecture to take
advantage of partial inputs: the encoder operates only on the visible patches
$\widehat{\mathbf{X}}$ without mask tokens; the decoder reconstructs the
original images from representations of visible patches and masked tokens.
This design allows MAE to achieve competitive performance as well as
impressive efficiency, _e.g._ , $3\times$ training speedup compared with the
ones operating on all patches. Nevertheless, the MAE works only with the
isotropic ViTs and it is unclear _how to translate the efficiency of MAE to
hierarchical ViTs_ , which exhibited nearly unanimous superiority over the
isotropic ones on most vision tasks [60, 43, 16, 12, 69]. In this paper, we
attempt to answer this question and propose a much greener approach for MIM
with Hierarchical ViTs.

###  3.2 Green Hierarchical Vision Transformer for Masked Image Modeling

Base architecture. We choose the representative hierarchical Vision
Transformer–Swin Transformer [43]–as our baseline, which mainly consists of
the eed-Forward Networks (FFNs) and the (shifted) window attentions. While the
FFNs are point-wise operations and can operate only on the visible patches,
the window attentions fail to do so. Given a window size $p$ (_e.g._ , 7 for
Swin), window attention first partitions the feature map $\mathbf{X}$ into
$n_{w}=\frac{H}{p}\times\frac{W}{p}$ non-overlapped local windows as
$\mathbf{X}\rightarrow\\{\mathbf{X}_{i}\\}_{i=0}^{n_{w}}$ where each
$\mathbf{X}_{i}$ contains $p\times p$ patches. Then the Multi-head Self-
Attention (MSA) [58] is performed within each window independently in parallel
because each window contains the same number of patches. Nevertheless, when
the numbers of patches within the local windows are uneven, _i.e._ , due to
the random masking as in Figure 2, it is unclear how to efficiently compute
the attention in parallel. To this end, we propose an efficient Group Window
Attention scheme and directly replace all (shifted) window attentions in Swin
to make it operate only on visible patches in a green manner.

Group Window Attention. Regarding the aforementioned problem, we propose a
Group Window Attention scheme that significantly improves the computation
efficiency of window attention on masked features. Given the masked feature
$\widehat{\mathbf{X}}=\mathrm{Mask}(\mathbf{X},\mathbf{M})$ following Equation
(1), we collect a set of uneven local windows
$\widehat{\mathbf{X}}\rightarrow\\{\widehat{\mathbf{X}}_{i}\\}_{i=0}^{n_{w}}$
of which each element contains the visible tokens only, with sizes
$\\{w_{i}\\}_{i=0}^{n_{w}}$ accordingly. As shown in Figure 2, our Group
Window Attention first uses an Optimal Grouping algorithm to partition the
uneven windows into several equal-sized groups and then perform Masked
Attention within each group to avoid information leakage. In the next two
subsections, we will elaborate on these two components, respectively.

###  3.3 Optimal Grouping with Dynamic Programming

General formulation. The first step of the optimal grouping is to find an
indexes partition $\Pi$ with respect to the group size $g_{s}$:

$\displaystyle\Pi=\\{\pi_{j}\\}_{j=1}^{n_{g}},$ (2) $\displaystyle
s.t.~{}~{}\cup_{j}~{}\pi_{j}=\\{1,\cdots,\sum_{i}w_{i}\\},~{}~{}\sum_{j}|\pi_{j}|$
$\displaystyle=\sum_{i}w_{i},\mathrm{and}~{}\forall_{j}~{}\sum_{k\in\pi_{j}}w_{k}\leq
g_{s}.$ (3)

where $n_{g}$ is the number of resulting groups. The conditions in Equation
(3) constrain the partition to contain all local windows with no duplicate and
enforce the actual size of each group to be smaller than $g_{s}$. Based on the
partition $\Pi$, we obtain a set of grouped tokens
$\\{\bar{\mathbf{X}}_{j}\\}_{j=1}^{n_{g}}$ that

$\displaystyle\overline{\mathbf{X}}_{j}=\mathrm{Concat}(\\{\widehat{\mathbf{X}}_{k}:\forall
k\in\pi_{j}\\}),~{}~{}\overline{\mathbf{X}}_{j}\in\mathbb{R}^{C\times g_{s}},$
(4)

upon which the Masked Attention is performed. Finally, we apply the inverse
operation of the partition $\Pi$ to recover the positions of output tokens.

With the formulation above, there remain two questions unresolved: 1) how to
choose the optimal group size $g_{s}^{*}$, and 2) how to obtain the optimal
partition $\Pi^{*}$ given $g_{s}^{*}$. To this end, we formulate our
objectives as the below min-min optimization problem,

$\displaystyle
g_{s}^{*}=\mathrm{argmin}_{g_{s}}~{}\mathcal{C}(g_{s},\mathrm{Partition}(g_{s},\\{w_{i}\\}_{i})),$
(5)
$\displaystyle\mathrm{Partition}(g_{s},\\{w_{i}\\}_{n_{w}})=\mathrm{argmin}_{\Pi}~{}|\Pi|,~{}s.t.~{}~{}\mathrm{Equation}~{}\eqref{eq:partition_cond},$
(6)

where $\mathcal{C}(\cdot)$ is a cost function measuring the computation cost
of the attention with the grouped tokens. Intuitively, Equation (5) aims to
find the optimal group size $g_{s}^{*}$ that the computation cost of the
optimal partition w.r.t.$g_{s}^{*}$ is the minimum. Equation (6) searches for
the optimal partition, with the constraints of Equation (3). Having the
optimal group size, we can directly obtain the optimal partition
$\Pi^{*}=\mathrm{Partition}(g_{s}^{*},\\{w_{i}\\}_{n_{w}})$. Next, we will
present in detail how we solve the above optimization problem.

Algorithm 1 Optimal Grouping

0: The number of visible patches within each local window
$\\{w_{i}\\}_{i=0}^{n_{w}}$,

1: Minimum computational cost $c^{*}\leftarrow+\infty$

2: for $g_{s}=\max_{i}\\{w_{i}\\}_{i=1}^{n_{w}}$ to $\sum_{i=1}^{n_{w}}w_{i}$
do

3: Remaining windows $\Phi\leftarrow\\{w_{i}\\}_{i=1}^{n_{w}}$; partition
$\Pi\leftarrow\emptyset$; the number of group $n_{g}\leftarrow 0$

4: repeat

5: $\pi_{n_{g}}\leftarrow\mathrm{Knapsack}(g_{s},\Phi)$, as in Equation (7)

6: $\Pi\leftarrow\Pi\cup\pi_{n_{g}}$; $\Phi\leftarrow\Phi\setminus\pi_{n_{g}}$

7: $n_{g}\leftarrow n_{g}+1$

8: until $\Phi=\emptyset$

9: $c\leftarrow\mathcal{C}(g_{s},\Pi)$, as in Equation (8)

10: if $c<c^{*}$ then

11: $c^{*}\leftarrow c$; $\Pi^{*}\leftarrow\Pi$

12: end if

13: end for

14: return Optimal group partition $\Pi^{*}$

Group partition with Dynamic Programming. We find that the optimization
problem in Equations (6) and (3) is a special case of the multiple subset sum
problem with identical capacities (MSSP-I), a variant of the well-known 0-1
multiple knapsack problem with identical capacities (MKP-I) [34, Chapter 10].
In our case, the group size is analogous to the capacity of knapsacks, the
numbers of visible tokens $\\{w_{i}\\}_{n_{w}}$ are analogous to the values of
goods, the weights of goods are the same as their values, and the number of
knapsacks is unbounded. Although the general multiple knapsacks problem is NP-
complete, its variant MSSP-I can be solved in pseudo-polynomial time using the
dynamic programming (DP) algorithms [4], owing to the identical capacities.
Specifically, we make use of the DP algorithm for the single knapsack problem
(or the subset sum problem):

$\displaystyle\pi\leftarrow\mathrm{Knapsack}(g_{s},\Phi),$ (7)

which selects a subset $\pi$ that $\sum_{u\in\pi}|u|\leq g_{s}$ out of the
full set $\Phi$ (the pseudo-code of this algorithm is given in Appendix). We
alternatively apply this algorithm to the remaining full set $\Phi$ and
exclude the selected subset $\pi$ from $\Phi$, until $\Phi$ is empty. In
practice, we found that our algorithm is very fast and the time cost is
negligible because the number of local windows is often small, _e.g._ , less
than 100 in our pre-training stage.

Cost function. Because we mainly care about the efficiency, we use the FLOPs
to measure the computation complexity of the multi-head attention on grouped
tokens, _i.e._ ,

$\displaystyle\mathcal{C}(g_{s},\Pi)=|\Pi|\times(4g_{s}C^{2}+2g_{s}^{2}C)=n_{g}\times(4g_{s}C^{2}+2g_{s}^{2}C),$
(8)

where $C$ is the number of channels. Although the complexity is quadratic
w.r.t.the group size $g_{s}$, using smaller $g_{s}$ might produce more groups
(and more padding) and suffers from suboptimal efficiency. Therefore the
optimal group size is determined adaptively during training.

Putting everything together. We sweep over the possible values of group size,
from $\max_{i}~{}w_{i}$ to $\sum_{i}w_{i}$, to find the optimal group size.
For each selected group size, we firstly use the DP algorithm in Equation (7)
to partition the windows and then calculate the computation cost of the
attention this partition. The one with minimum cost is chosen as the optimal
group size. The pseudo-code of the optimal grouping is summarized in Algorithm
1.

Figure 3: Illustration of the Masked Attention scheme. Given a group of
tokens, we first compute their pairwise attention weights and then set the
attention weights between tokens from different local windows to $-\infty$
(indicated by the gray cells). The final attention output is then computed
with the masked attention map.

###  3.4 Masked Attention

Because non-adjacent local windows are partitioned into the same groups,
masking the attention weights is needed to avoid the information exchange
between these local windows. As illustrated in Figure 3, having computed the
attention map, we retain only the intra-window attention weights (_i.e._ , the
block-diagonal elements) and discard the inter-window ones. A similar masking
scheme is also applied to the retrieving of relative position bias [43], where
we store the original absolute position of each token and compute the relative
positions on-the-fly to retrieve the corresponding biases.

###  3.5 Batch-wise Random Masking

We observed that the per-sample random masking strategy would deteriorate the
efficiency of our method: 1) it might produce different numbers of groups of
local windows for each sample, which is intractable for the parallel
computation of the Masked Attention; 2) when the mask patch size is smaller
than the largest patch size of the hierarchical models, some patches might
contain both visible and masked inputs. In this case, we can not discard such
patches during training and fail to fully take advantage of the
sparsification. Therefore, we propose to set the mask patch size to the same
value as the largest patch size of the encoder (_e.g._ , 32 for most
hierarchical models, which is also the default choice of [66]) and use the
same random mask for all samples in the same GPU device (a.k.a., a micro-
batch).

##  4 Experiments

###  4.1 Implementation Details and Experimental Setups

We conduct experiments on the ImageNet-1K [54] (BSD 3-Clause License) image
classification dataset and MS-COCO [41] (CC BY 4.0 License) object
detection/instance segmentation dataset. The Swin-Base and Swin-Large [43]
models, which consist of four stages with features of stride 4/8/16/32, are
used as the encoder throughout this paper for direct comparisons with the
baseline [66]. The model is first pre-trained on the ImageNet-1K dataset
without label and then fine-tuned on downstream tasks. All the experiments of
our method are performed on a single machine with eight 32G Tesla V100 GPUs,
CUDA 10.1, PyTorch [48] 1.8, and automatic mixed-precision training [45].

Pre-training setup. We patchify images of size 224$\times$224 with a patch
size of 4$\times$4 and randomly mask the patches with a ratio $r$ ($r=0.75$ by
default) following the scheme in Section 3.5. The input images are transformed
by a set of simple data augmentations, including random cropping and
horizontal flip, and standardization. Following the prior work MAE [22], we
use a lightweight decoder that consists of $n_{d}$ (by default $n_{d}=1$)
transformer blocks with an embedding dimension of 512. The decoder takes the
representations of visible patches and the mask token as input and is appended
after the final stage of the encoder to learning representations for the
masked patches. It is followed by a linear layer to predict the normalized
pixel values of the masked patches. The models are trained for 100/200/400/800
epochs with a batch size of 2,048. We use the AdamW optimizer [35] with the
cosine annealing schedule [44]. We set the base learning rate to
1.5$\mathrm{e}^{-4}$, weight decay to 0.05, the hyper-parameters of Adam
$\beta_{1}=0.9$, $beta_{2}=0.999$, the number of warmup epochs to 40 with an
initial base learning rate 1.5$\mathrm{e}^{-7}$. The effective learning rate
is scaled linearly by $\mathrm{batch\\_size}/256$.

Figure 4: The optimal group size $\bm{g_{s}}$ at each stage. The figure of the
fourth stage is omitted here because there is only one local window in this
stage, so the grouping is not necessary. The simulation is repeated 100 times,
of which the mean and standard deviation (the shaded regions) are reported.

Fine-tuning on the ImageNet-1K dataset. For fine-tuning, we drop the decoder
and directly append a 1,000-way fully-connected layer to the average-pooled
output of the encoder as the classifier. The models are also optimized by the
AdamW optimizer [35] with 100 training epochs in total, 20 warmup epochs, a
base/warmup learning rate of 1.25$\mathrm{e}^{-4}$/2.5$\mathrm{e}^{-7}$, the
cosine annealing schedule [44], a weight decay of 0.05, a layerwise learning
rate decay [2] of 0.9, and a stochastic depth [28] ratio of 0.1. The data
augmentations are the same as [2, 66].

Fine-tuning on the MS-COCO dataset. We adopt the Mask R-CNN [24] architecture
with the FPN [40] as the detector. All models are fine-tuned on the MS-COCO
[41] 2017 $\operatorname{train}$ split ($\sim$118k images) and finally
evaluated on the $\operatorname{val}$ split ($\sim$5k images). We use a batch
size of 16, AdamW optimizer [35] with a learning rate of 1$\mathrm{e}^{-4}$, a
weight decay of 0.05. The $1\times$/$3\times$ schedule in the mmdetection [8]
is adopted, which uses 12/36 training epochs in total and decays the learning
rate at the $\frac{3}{4}$ and $\frac{11}{12}$ of the total epochs by a factor
of 10. The standard COCO metrics, including AP, AP50, and AP75 for both object
detection and instance segmentation are used for evaluation.

###  4.2 Ablations studies

Efficiency comparison with SimMIM. We compare the efficiency of our method to
the baseline SimMIM in Figure 1. The evaluations are performed on a single
machine with eight 32GB V100 GPUs for our method and on 2 or 4 machines for
the SimMIM because it fails to fit into a single machine with the default
batch size of the original paper [66] (_i.e._ , 2,048). As we can see from the
figures, the training of SimMIM with images of size $224^{2}$ is very slow and
memory-hungry. Although training with smaller images considerably reduces the
training time and memory consumption, it still far lags behind our method with
images of size $224^{2}$. Concretely, with the same amount of training epochs,
our method performs on par with baseline with $\sim$2$\times$ speedup and
$\sim$60% of memory reduction using Swin-B. We also observe that the
efficiency improvements become larger with the larger Swin-L, _e.g._ ,
2.7$\times$ speedup compared with SimMIM192, highlighting the efficiency of
our method with larger models.

The optimal group size $\bm{g_{s}}$ on each stage. Because hierarchical models
have multiple stages with different scales of features, the optimal group size
of each stage might also be different. Regarding this, we design a simulation
experiment to analyze the optimal $g_{s}$ at different stages. In the
simulation, we randomly generate 100 masks following Section 3.5, compute the
costs w.r.t.different choices of $g_{s}$, and report the mean/standard
deviation of the costs in Figure 4. Note that the analysis of the 4th stage is
omitted here because it only has one local window. In general, we observed
that the cost increases quadratically w.r.t.the group size, except for some
cases where group size is exactly equal to the sum of a subset of local
windows. Another intriguing observation is that the cost at each stage seems
to be the minimum around the point that $g_{s}=49$, which is equal to the
window size of the window attention. This observation indicates that we may
not need to sweep over all possible group sizes but simply set $g_{s}=p\times
p$ in practice.

Figure 5: Ablation studies of (a) the choice of the mask ratio $r$, (b) the
number of transformer blocks $n_{d}$ in the decoder, and (c) the number of
pre-training epochs.

The influence of the mask ratio, the number of decoder blocks, and the pre-
training epochs. From Figure 5(a), we can see that the performance of our
method is quite stable with the mask ratio $r$ varying from 0.5 to 0.85, which
conforms with the observation of [22]. In Figure 5(b), we also study the
influence of the depth of the decoder. Intriguingly, the results suggest that
fewer decoder blocks produce better results. This study favors the simple
prediction head design of SimMIM [66] with hierarchical models and is in
contrast to the observation of MAE [22] with isotropic ones. For simplicity
and efficiency, we fix $r=0.75$ and the number of decoder blocks to 1
throughout the paper. Furthermore, we study the impact of pre-training budgets
on our method. As shown in Figure 5(c), the fine-tuning accuracy increases
steadily w.r.t.the number of training epochs and does not seem to stagnate,
suggesting its potential for a further performance boost.

Table 1: A larger window size $p\times p$. $p\times p$ | Time | Mem. | Acc.
---|---|---|---
$7\times 7$ | 1.1h | 121.6G | 83.2%
$14\times 14$ | 1.2h | 148.9G | 83.4%

Pre-training with larger window size. The work of [42] puts forth that using a
larger window size is beneficial for fine-tuning. In practice, however, it
might be less practical because of the quadratic complexity of self-attention
w.r.t.the window size. Fortunately, operating only on the visible patches
permits the training with a larger window size with little extra cost. As
displayed in Table 1, pre-training with doubled window size only marginally
increases the training time/GPU memory by less than 10%/20% yet brings a
moderate performance improvement despite $p=7$ in the fine-tuning.

###  4.3 ImageNet-1K Classification

We fine-tune the pre-trained models on the ImageNet validation set and report
the results in Table 2. Here, we make direct comparisons with the models 1)
trained from scratch with a longer training schedule, 2) trained with
contrastive learning, and 3) trained with MIM. Our approach achieves 83.7%
top-1 fine-tuning accuracy with the Swin-Base backbone, which is superior to
the supervised learning/contrastive learning methods and on par with other MIM
methods using backbones with similar capacities. The results demonstrate the
effectiveness of our method, in addition to the substantial efficiency
improvements over both MAE and SimMIM. The experiments with the Swin-L
backbone are given in Appendix to have a further inspection on our method.

Table 2: Top1 accuracy on the ImageNet-1K validation set with the Swin-B or ViT-B models. All methods are trained with images of size $224\times 224$ in both the pre-training and fine-tuning except for SimMIM192 using $192\times 192$ in the pre-training.  Method | Model | #Params | PT Ep. | Ep. Hours | Total Hours | FT Ep. | Acc. (%)
---|---|---|---|---|---|---|---
Training from scratch |  |  |
Scratch, DeiT [56] | ViT-B | 86M | 0 | - | - | 300 | 81.8
Scratch, MAE [22] | ViT-B | 86M | 0 | - | - | 300 | 82.3
Scratch, Swin [43] | Swin-B | 88M | 0 | - | - | 300 | 83.5
Supervised Pre-training |  |  |
Supervised, SimMIM [66] | Swin-B | 88M | 300 | - | - | 100 | 83.3
Supervised, SimMIM [66] | Swin-L | 197M | 300 | - | - | 100 | 83.5
Pre-training with Contrastive Learning |  |  |
MoCov3 [11] | ViT-B | 86M | 800 | - | - | 100 | 83.2
DINO [7] | ViT-B | 86M | 800 | - | - | 100 | 82.8
Pre-training with Masked Image Modeling |  |  |
BEiT [2] | ViT-B | 86M | 800 | - | - | 100 | 83.2
MaskFeat [62] | ViT-B | 86M | 800 | - | - | 100 | 84.0
MAE [22] | ViT-B | 86M | 1600 | 1.3 | 2069 | 100 | 83.6
SimMIM224 [66] | ViT-B | 86M | 800 | 4.1 | 3307 | 100 | 83.8
SimMIM192 [66] | Swin-B | 88M | 800 | 2.0 | 1609 | 100 | 84.0
SimMIM192 [66] | Swin-L | 197M | 800 | 3.5 | 2821 | 100 | 85.4
Ours | Swin-B | 88M | 800 | 1.1 | 887 | 100 | 83.7
Ours | Swin-L | 197M | 800 | 1.3 | 1067 | 100 | 85.1

###  4.4 MS-COCO Object Detection and Instance Segmentation

Finally, we evaluate the transfer learning performance of our pre-trained
models to the MS-COCO object detection and instance segmentation dataset.
Here, we directly use the code base of the supervised Swin Transformer without
any modification to the fine-tuning strategy. For direct comparisons, we reran
the experiments for the supervised Swin-B and SimMIM using their public
checkpoints. The experiment results are summarized in Table 3. Compared with
the supervised pre-trained Swin-B, our approach performs prominently better in
terms of all metrics, _e.g._ , 1.5% absolute improvement in APb. In addition,
we also observed that our approach still performs comparably to the SimMIM on
dense prediction tasks. More significantly, our approach outperforms most of
the baselines from [39] using 3$\times$ or 10$\times$ more fine-tuning epochs
and advanced data augmentations [17]. These experiments, combined with the
results in Table 2, verify that our approach can achieve outstanding
performance with impressive pre-training efficiency.

Table 3: MS-COCO object detection and instance segmentation. All methods are based on the Mask R-CNN [24] architecture with the FPN [40] neck. The methods in gray are cited from [39]. Most of them use much longer training schedules and advanced data augmentations. Method | Backbone | PT Ep. | PT Hours | FT Epochs | APb | AP${}_{50}^{\mathrm{b}}$ | AP${}_{75}^{\mathrm{b}}$ | APm | AP${}_{50}^{\mathrm{m}}$ | AP${}_{75}^{\mathrm{m}}$
---|---|---|---|---|---|---|---|---|---|---
Training from scratch |  |  |  |  |  |
Benchmarking [39] | ViT-B | 0 | 0 | 400 | 48.9 | - | - | 43.6 | - | -
Supervised Pretraining |  |  |  |  |  |
Benchmarking [39] | ViT-B | 300 | 992 | 100 | 47.9 | - | - | 42.9 | - | -
PVT [60] | PVT-L | 300 | - | 36 | 44.5 | 66.0 | 48.3 | 40.7 | 63.4 | 43.7
Swin [43] | Swin-B | 300 | 840 | 36 | 48.5 | 69.8 | 53.2 | 43.2 | 66.9 | 46.7
Self-Supervised Pre-training |  |  |  |  |  |
MoCov3 [11] | ViT-B | 800 | - | 100 | 47.9 | - | - | 42.7 | - | -
BEiT [2] | ViT-B | 800 | - | 100 | 49.8 | - | - | 44.4 | - | -
MAE [22] | ViT-B | 1600 | 2069 | 25 | 48.1 | - | - | - | - | -
MAE [22] | ViT-B | 1600 | 2069 | 100 | 50.3 | - | - | 44.9 | - | -
SimMIM [66] | Swin-B | 800 | 1609 | 36 | 50.4 | 70.9 | 55.5 | 44.4 | 68.2 | 47.9
Ours | Swin-B | 800 | 887 | 36 | 50.0 | 70.7 | 55.4 | 44.1 | 67.9 | 47.5

##  5 Conclusion

In this paper, we present a green approach for Masked Image Modeling (MIM)
with hierarchical Vision Transformers, _e.g._ , Swin Transformer [43],
allowing the hierarchical models to discard masked patches and operate only on
the visible ones. Coupling the efficient Group Window Attention scheme and the
DP-algorithm-based Optimal Grouping strategy, our approach can train the
hierarchical models $\sim$2.7$\times$ faster and reduce the GPU memory
consumption by 70%, while still enjoying a competitive performance on ImageNet
classification and the superiority of downstream MS-COCO object detection
benchmarks. We hope that this work will facilitate future self-supervised
learning methods targeting both effectiveness and efficiency.

Limitations. One of the limitations of our algorithm is that it requires the
batch-wise masking scheme (as in Section 3.5) to achieve the best efficiency.
Although this limitation only has little impact on the MIM pre-training, it
restrains the application of our method on a broader range of settings, e.g.,
training ViTs with token sparification [53, 68] that requires instance-wise
sparsification. These applications are beyond the scope of this work and we
will leave them for the future study.

Broader Impact. This work proposed a green approach for MIM with hierarchical
ViTs, prominently alleviating the heavy computation burden of MIM. On the one
hand, this work provokes the efficiency as well as the effectiveness of MIM,
which may inspire new algorithms and investigations on this direction. On the
other hand, since the pre-training datasets might contain biases, our
approach, in the same way as other unsupervised/self-supervised learning
methods, might also be susceptible to replicate these biases. This concern can
be mitigated by combining the FairML methods [3].

## References

  * [1] A. Baevski, W.-N. Hsu, Q. Xu, A. Babu, J. Gu, and M. Auli.  Data2vec: A general framework for self-supervised learning in speech, vision and language.  arXiv preprint arXiv:2202.03555, 2022.
  * [2] H. Bao, L. Dong, and F. Wei.  Beit: Bert pre-training of image transformers.  arXiv preprint arXiv:2106.08254, 2021.
  * [3] S. Barocas, M. Hardt, and A. Narayanan.  Fairness in machine learning.  Nips tutorial, 1:2, 2017.
  * [4] R. Bellman.  Dynamic programming.  Science, 153(3731):34–37, 1966.
  * [5] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al.  Language models are few-shot learners.  In Advances in Neural Information Processing Systems, volume 33, 2020.
  * [6] M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  Advances in Neural Information Processing Systems, 33, 2020.
  * [7] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin.  Emerging properties in self-supervised vision transformers.  arXiv preprint arXiv:2104.14294, 2021.
  * [8] K. Chen, J. Wang, J. Pang, Y. Cao, Y. Xiong, X. Li, S. Sun, W. Feng, Z. Liu, J. Xu, et al.  Mmdetection: Open mmlab detection toolbox and benchmark.  arXiv preprint arXiv:1906.07155, 2019.
  * [9] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton.  A simple framework for contrastive learning of visual representations.  In International Conference on Machine Learning, 2020.
  * [10] X. Chen and K. He.  Exploring simple Siamese representation learning.  arXiv preprint arXiv:2011.10566, 2020.
  * [11] X. Chen, S. Xie, and K. He.  An empirical study of training self-supervised vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9640–9649, 2021.
  * [12] X. Chu, Z. Tian, Y. Wang, B. Zhang, H. Ren, X. Wei, H. Xia, and C. Shen.  Twins: Revisiting the design of spatial attention in vision transformers.  Advances in Neural Information Processing Systems, 34, 2021.
  * [13] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4171–4186, 2019.
  * [14] X. Dong, J. Bao, T. Zhang, D. Chen, W. Zhang, L. Yuan, D. Chen, F. Wen, and N. Yu.  Peco: Perceptual codebook for bert pre-training of vision transformers.  arXiv preprint arXiv:2111.12710, 2021.
  * [15] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [16] H. Fan, B. Xiong, K. Mangalam, Y. Li, Z. Yan, J. Malik, and C. Feichtenhofer.  Multiscale vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6824–6835, 2021.
  * [17] G. Ghiasi, Y. Cui, A. Srinivas, R. Qian, T.-Y. Lin, E. D. Cubuk, Q. V. Le, and B. Zoph.  Simple copy-paste is a strong data augmentation method for instance segmentation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2918–2928, 2021.
  * [18] S. Gidaris, P. Singh, and N. Komodakis.  Unsupervised representation learning by predicting image rotations.  In International Conference on Learning Representation, 2018.
  * [19] R. Girshick, J. Donahue, T. Darrell, and J. Malik.  Rich feature hierarchies for accurate object detection and semantic segmentation.  In IEEE Conference on Computer Vision and Pattern Recognition, pages 580–587, 2014.
  * [20] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio.  Generative adversarial nets.  In Advances in Neural Information Processing Systems, pages 2672–2680, 2014.
  * [21] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  Advances in Neural Information Processing Systems, 33, 2020.
  * [22] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick.  Masked autoencoders are scalable vision learners.  arXiv preprint arXiv:2111.06377, 2021.
  * [23] K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick.  Momentum contrast for unsupervised visual representation learning.  In IEEE Conference on Computer Vision and Pattern Recognition, pages 9729–9738, 2020.
  * [24] K. He, G. Gkioxari, P. Dollár, and R. Girshick.  Mask r-cnn.  In Proceedings of the IEEE international conference on computer vision, pages 2961–2969, 2017.
  * [25] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition, pages 770–778, 2016.
  * [26] O. J. Hénaff, S. Koppula, J.-B. Alayrac, A. v. d. Oord, O. Vinyals, and J. Carreira.  Efficient visual pretraining with contrastive detection.  arXiv preprint arXiv:2103.10957, 2021.
  * [27] J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  * [28] G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger.  Deep networks with stochastic depth.  In European Conference on Computer Vision, pages 646–661. Springer, 2016.
  * [29] L. Huang, S. You, M. Zheng, F. Wang, C. Qian, and T. Yamasaki.  Learning where to learn in cross-view self-supervised learning.  arXiv preprint arXiv:2203.14898, 2022.
  * [30] L. Huang, Y. Yuan, J. Guo, C. Zhang, X. Chen, and J. Wang.  Interlaced sparse self-attention for semantic segmentation.  arXiv preprint arXiv:1907.12273, 2019.
  * [31] L. Huang, C. Zhang, and H. Zhang.  Self-adaptive training: Bridging supervised and self-supervised learning.  arXiv preprint arXiv:2101.08732, 2021.
  * [32] Z. Huang, Y. Ben, G. Luo, P. Cheng, G. Yu, and B. Fu.  Shuffle transformer: Rethinking spatial shuffle for vision transformer.  arXiv preprint arXiv:2106.03650, 2021.
  * [33] T. Karras, S. Laine, and T. Aila.  A style-based generator architecture for generative adversarial networks.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401–4410, 2019.
  * [34] H. Kellerer, U. Pferschy, and D. Pisinger.  Knapsack problems.  Springer, 2004.
  * [35] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In International Conference on Learning Representations, 2015.
  * [36] D. P. Kingma and M. Welling.  Auto-encoding variational Bayes.  In International Conference on Learning Representations, 2014.
  * [37] A. Krizhevsky, I. Sutskever, and G. E. Hinton.  Imagenet classification with deep convolutional neural networks.  In Advances in Neural Information Processing Systems, volume 25, pages 1097–1105, 2012.
  * [38] G. Larsson, M. Maire, and G. Shakhnarovich.  Colorization as a proxy task for visual understanding.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6874–6883, 2017.
  * [39] Y. Li, S. Xie, X. Chen, P. Dollar, K. He, and R. Girshick.  Benchmarking detection transfer learning with vision transformers.  arXiv preprint arXiv:2111.11429, 2021.
  * [40] T.-Y. Lin, P. Dollár, R. Girshick, K. He, B. Hariharan, and S. Belongie.  Feature pyramid networks for object detection.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2117–2125, 2017.
  * [41] T.-Y. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitnick.  Microsoft coco: Common objects in context.  In European conference on computer vision, pages 740–755. Springer, 2014.
  * [42] Z. Liu, H. Hu, Y. Lin, Z. Yao, Z. Xie, Y. Wei, J. Ning, Y. Cao, Z. Zhang, L. Dong, et al.  Swin transformer v2: Scaling up capacity and resolution.  arXiv preprint arXiv:2111.09883, 2021.
  * [43] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012–10022, 2021.
  * [44] I. Loshchilov and F. Hutter.  SGDR: Stochastic gradient descent with warm restarts.  In International Conference on Learning Representations, 2017.
  * [45] P. Micikevicius, S. Narang, J. Alben, G. Diamos, E. Elsen, D. Garcia, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, et al.  Mixed precision training.  arXiv preprint arXiv:1710.03740, 2017.
  * [46] M. Noroozi and P. Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In European Conference on Computer Vision, pages 69–84. Springer, 2016.
  * [47] A. v. d. Oord, Y. Li, and O. Vinyals.  Representation learning with contrastive predictive coding.  arXiv preprint arXiv:1807.03748, 2018.
  * [48] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  In Advances in Neural Information Processing Systems, pages 8024–8035, 2019.
  * [49] D. Pathak, P. Krahenbuhl, J. Donahue, T. Darrell, and A. A. Efros.  Context encoders: Feature learning by inpainting.  In IEEE Conference on Computer Vision and Pattern Recognition, pages 2536–2544, 2016.
  * [50] A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever.  Improving language understanding by generative pre-training, 2018.
  * [51] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever.  Language models are unsupervised multitask learners.  OpenAI blog, 1(8):9, 2019.
  * [52] A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever.  Zero-shot text-to-image generation.  In International Conference on Machine Learning, pages 8821–8831. PMLR, 2021.
  * [53] Y. Rao, W. Zhao, B. Liu, J. Lu, J. Zhou, and C.-J. Hsieh.  Dynamicvit: Efficient vision transformers with dynamic token sparsification.  Advances in neural information processing systems, 34, 2021.
  * [54] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, et al.  Imagenet large scale visual recognition challenge.  International Journal of Computer Vision, 115(3):211–252, 2015\.
  * [55] R. Schwartz, J. Dodge, N. A. Smith, and O. Etzioni.  Green ai.  Communications of the ACM, 63(12):54–63, 2020.
  * [56] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou.  Training data-efficient image transformers & distillation through attention.  In International Conference on Machine Learning, pages 10347–10357. PMLR, 2021.
  * [57] A. Van Den Oord, O. Vinyals, et al.  Neural discrete representation learning.  Advances in neural information processing systems, 30, 2017.
  * [58] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [59] P. Vincent, H. Larochelle, Y. Bengio, and P.-A. Manzagol.  Extracting and composing robust features with denoising autoencoders.  In International Conference on Machine Learning, pages 1096–1103, 2008.
  * [60] W. Wang, E. Xie, X. Li, D.-P. Fan, K. Song, D. Liang, T. Lu, P. Luo, and L. Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 568–578, 2021.
  * [61] X. Wang, R. Zhang, C. Shen, T. Kong, and L. Li.  Dense contrastive learning for self-supervised visual pre-training.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3024–3033, 2021.
  * [62] C. Wei, H. Fan, S. Xie, C.-Y. Wu, A. Yuille, and C. Feichtenhofer.  Masked feature prediction for self-supervised visual pre-training.  arXiv preprint arXiv:2112.09133, 2021.
  * [63] Z. Wu, Y. Xiong, S. X. Yu, and D. Lin.  Unsupervised feature learning via non-parametric instance discrimination.  In IEEE Conference on Computer Vision and Pattern Recognition, pages 3733–3742, 2018.
  * [64] T. Xiao, C. J. Reed, X. Wang, K. Keutzer, and T. Darrell.  Region similarity representation learning.  arXiv preprint arXiv:2103.12902, 2021.
  * [65] Z. Xie, Y. Lin, Z. Zhang, Y. Cao, S. Lin, and H. Hu.  Propagate yourself: Exploring pixel-level consistency for unsupervised visual representation learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16684–16693, 2021.
  * [66] Z. Xie, Z. Zhang, Y. Cao, Y. Lin, J. Bao, Z. Yao, Q. Dai, and H. Hu.  Simmim: A simple framework for masked image modeling.  arXiv preprint arXiv:2111.09886, 2021.
  * [67] J. Xu, W. Zhou, Z. Fu, H. Zhou, and L. Li.  A survey on green deep learning.  arXiv preprint arXiv:2111.05193, 2021.
  * [68] H. Yin, A. Vahdat, J. Alvarez, A. Mallya, J. Kautz, and P. Molchanov.  Adavit: Adaptive tokens for efficient vision transformer.  arXiv preprint arXiv:2112.07658, 2021.
  * [69] Y. Yuan, R. Fu, L. Huang, W. Lin, C. Zhang, X. Chen, and J. Wang.  Hrformer: High-resolution vision transformer for dense predict.  Advances in Neural Information Processing Systems, 34, 2021.
  * [70] J. Zbontar, L. Jing, I. Misra, Y. LeCun, and S. Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  arXiv preprint arXiv:2103.03230, 2021.
  * [71] M. Zheng, F. Wang, S. You, C. Qian, C. Zhang, X. Wang, and C. Xu.  Weakly supervised contrastive learning.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10042–10051, October 2021.
  * [72] M. Zheng, S. You, F. Wang, C. Qian, C. Zhang, X. Wang, and C. Xu.  Ressl: Relational self-supervised learning with weak augmentation.  Advances in Neural Information Processing Systems, 34, 2021.
  * [73] J. Zhou, C. Wei, H. Wang, W. Shen, C. Xie, A. Yuille, and T. Kong.  ibot: Image bert pre-training with online tokenizer.  arXiv preprint arXiv:2111.07832, 2021.

##  Appendix A More Details of Our Method

###  A.1 Method Overview.

We provide a diagram in Figure 6 for an intuitive illustration of our method.
The input image is first randomly masked according to masking scheme in
Section 3.5, then fed into our green hierarchical ViT to obtain
representations for each visible patches. Finally, we concatenate the
representations of visible patches with the mask tokens and feed them into a
transformer decoder, yieding the representations for the masked patches.
Finally, we predict the raw pixel values of masked patches from the
corresponding representations and implement the training by minimizing the
Mean Square Error (MSE) between the predictions and ground-truth.

###  A.2 Group Window Attention scheme with shifted windows.

In addition to Figure 2, we also illustrate how our method works with the
irregular window partition in Figure 7. We can observe that, owing to the
optimal grouping scheme, our method dynamically finds out the best group
partition despite the numbers of visible patches within each local window is
highly uneven. This figure further demonstrates that our approach is agnostic
to the window partition and works impressively well.

Figure 6: Overview of our method. The input image is randomly masked and then
fed into a 4-stage hierarchical ViTs. Finally, a lightweight decoder takes the
representations of visible patches and mask tokens to reconstruct the missing
patches. Figure 7: Illustration of the Group Window Attention scheme with
shifted windows. It shows that our approach is agnostic to the window
partition.

⬇

1def Knapsack(g_s, Phi):

2 # g_s (int): Group size

3 # Phi (list[int]): The numbers of visible patches within each local window

4

5 n_w = len(Phi) # the number of windows

6 K = [[0 for w in range(g_s + 1)] for i in range(n_w + 1)] # a buffer for the
DP algorithm

7

8 # Build table K[][] in a bottom up manner

9 for i in range(n_w + 1):

10 for w in range(g_s + 1):

11 if i == 0 or w == 0:

12 K[i][w] = 0

13 elif Phi[i - 1] <= w:

14 K[i][w] = max(Phi[i - 1] + K[i - 1][w - Phi[i - 1]], K[i - 1][w])

15 else:

16 K[i][w] = K[i - 1][w]

17

18 # Store the result of Knapsack

19 res = K[n_w][g_s]

20

21 # Store the selected indexes

22 w = g_s

23 Pi = []

24

25 for i in range(n_w, 0, -1):

26 if res <= 0:

27 break

28

29 if res == K[i - 1][w]: # This window is not included.

30 continue

31 else: # This window is included.

32 Pi.append(i - 1)

33 # Since this window is included, its value is deducted

34 res = res - Phi[i - 1]

35 w = w - Phi[i - 1]

36

37 return Pi[::-1] # Optional: make Pi in an increasing order

38

39

40def GroupPartition(g_s, Phi):

41 # g_s (int): Group size

42 # Phi (list[int]): The numbers of visible patches within each local window

43

44 win_szs = Phi.copy()

45 ori_win_idxs = list(range(len(win_szs)))

46 win_idxs = []

47

48 while len(win_szs) > 0:

49 idx = knapsack(group_size, win_szs)

50

51 # Append the selected idx

52 win_idxs.append([ori_win_idxs[i] for i in idx])

53

54 # The remaining windows and indexes

55 win_szs = [win_szs[i] for i in range(len(ori_win_idxs)) if i not in idx]

56 ori_win_idxs = [ori_win_idxs[i] for i in range(len(ori_win_idxs)) if i not
in idx]

57

58 return win_idxs

Algorithm 2 Dynamic Programming-based algorithm for the Optimal Grouping using
Python.

###  A.3 A Python Implementation of the Optimal Grouping algorithm

We provide a Python implementation of the Dynamic-Programming-based Optimal
Grouping algorithm in Algorithm 2. As we can see, the two components of the
Optimal Grouping algorithm are both easy to implement. For the DP algorithm
for the single Knapsack problem, its time/space complexity is
$\mathcal{O}(g_{s}n_{w})$ where $g_{s}$ is the group size and $n_{w}$ is the
number of windows. In practice, because $g_{s}$ and $n_{w}$ are generally
small (_i.e._ , smaller than 100) the running time of Algorithm 2 is
negligible (_i.e._ , $<$1ms).

###  A.4 A PyTorch Implementation of the Group Attention scheme

With the group partition on the indexes, we can then permute visible patches
according to the partition and obtain several groups of patches with a equal
size $g_{s}$, upon which the Masked Attention is performed. We also provide a
PyTorch implementation of the Group Attention scheme in Algorithm 3 to
facilitate future research. Note the the padding operations is omitted here
for simplicity.

⬇

1def GroupAttention(x, g_s, Phi):

2 # x (3-d tensor): Features the visible patches

3 # g_s (int): Group size

4 # Phi (list[int]): The numbers of visible patches within each local window

5

6 # B is the batch size, L is the number of visible patches, C is the number
of channels

7 B, L, C = x.shape

8

9 # Prepare for the group attention

10 win_idxs = GroupPartition(g_s, Phi)

11 patch_idxs = torch.arrange(sum(Phi))

12 patch_idxs = torch.split(patch_idxs, Phi)

13 shuffle_idxs = torch.cat([patch_idxs[wi] for wi in win_idxs])

14 unshuffle_idxs = torch.argsort(shuffle_idxs)

15

16 # Group partition. For simplicity, assume that the partition is even

17 x = torch.index_select(x, 1, shuffle_idxs) # (B, n_g * g_s, C)

18 x = x.reshape(-1, g_s, C) # (B * n_g, g_s, C)

19

20 # Attention with relative position bias as in Figure 3

21 x = MaskedAttention(x)

22

23 # Reverse the group partition

24 x = x.reshape(B, L, C).index_select(x, 1, unshuffle_idxs)

25

26 return x

Algorithm 3 Group Attention using the PyTorch framework.