Green Hierarchical Vision Transformer for Masked Image Modeling

TL;DR

This work tackles the high computational cost of masked image modeling (MIM) for hierarchical Vision Transformers by introducing Green MIM, which enables training on visible patches only. It combines Group Window Attention with a Dynamic Programming–driven Optimal Grouping strategy and replaces standard convolutions with Sparse Convolution to exploit sparsity, achieving substantial efficiency gains while maintaining accuracy. The method yields up to $2.7\times$ faster pretraining and up to $70\%$ memory savings and demonstrates competitive ImageNet results alongside superior downstream COCO object detection performance. The approach remains architecture-agnostic across Swin and Twins backbones, contributing a practical, greener paradigm for self-supervised learning in vision.

Abstract

We present an efficient approach for Masked Image Modeling (MIM) with hierarchical Vision Transformers (ViTs), allowing the hierarchical ViTs to discard masked patches and operate only on the visible ones. Our approach consists of three key designs. First, for window attention, we propose 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. Third, as for the convolution layers, we convert them to the Sparse Convolution that works seamlessly with the sparse data, i.e., the visible patches in MIM. As a result, MIM can now work on most, if not all, hierarchical ViTs in a green and efficient way. For example, we can train the hierarchical ViTs, e.g., Swin Transformer and Twins Transformer, 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. Code and pre-trained models have been made publicly available at https://github.com/LayneH/GreenMIM.

Figures (7)

• 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.
• 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.
• 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.
• 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.
• 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.