Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling

TL;DR

This work asks how to port BERT-style masked pre-training to convolutional networks, identifying irregular masked inputs and single-scale biases as key obstacles. It introduces SparK, which sparsely gathers unmasked patches into a sparse representation encoded with sparse convolutions and uses a hierarchical UNet-like decoder with mask embeddings to reconstruct multi-scale features. The encoder is then used for downstream tasks, with a simple, scalable pre-training setup that yields consistent gains over self-supervised transformers and contrastive methods, including substantial improvements on COCO. SparK demonstrates strong transferability and favorable scaling across model sizes, suggesting that generative pre-training on convnets is a viable path for future vision systems.

Abstract

We identify and overcome two key obstacles in extending the success of BERT-style pre-training, or the masked image modeling, to convolutional networks (convnets): (i) convolution operation cannot handle irregular, random-masked input images; (ii) the single-scale nature of BERT pre-training is inconsistent with convnet's hierarchical structure. For (i), we treat unmasked pixels as sparse voxels of 3D point clouds and use sparse convolution to encode. This is the first use of sparse convolution for 2D masked modeling. For (ii), we develop a hierarchical decoder to reconstruct images from multi-scale encoded features. Our method called Sparse masKed modeling (SparK) is general: it can be used directly on any convolutional model without backbone modifications. We validate it on both classical (ResNet) and modern (ConvNeXt) models: on three downstream tasks, it surpasses both state-of-the-art contrastive learning and transformer-based masked modeling by similarly large margins (around +1.0%). Improvements on object detection and instance segmentation are more substantial (up to +3.5%), verifying the strong transferability of features learned. We also find its favorable scaling behavior by observing more gains on larger models. All this evidence reveals a promising future of generative pre-training on convnets. Codes and models are released at https://github.com/keyu-tian/SparK.

Figures (4)

• Figure 1: Different masking strategies with pixel intensity histograms plotted before (in gray) and after (blue) masking. (b) is a straightforward idea to apply masked modeling to convnets, which results in a distribution shift. (a) illustrates MAE mae that has no such side effect thanks to the transformer's ability to process variable-length input. We propose (c) to adapt convnets to irregular masked input without a distribution shift.
• Figure 2: Sparse masked modeling with hierarchy. To adapt convolution to irregular masked input, visible patches are gathered into a sparse image and encoded by sparse convolution. To pre-train a hierarchical encoder, we employ a UNet-style architecture to decode multi-scale sparse feature maps, where all empty positions are filled with mask embedding. This "densifying" is necessary to reconstruct a dense image. Only the regression loss on masked patches will be optimized. After pre-training, only the encoder is used for downstream tasks.
• Figure 3: Using sparse convolution to address "mask pattern vanishing" issue. Three mask examples are shown. As in left, when computing ordinary "dense" convolution centered at a zero (masked) position, the result would be non-zero if the filter covers any non-zero (unmasked) points. Repeating this convolution will erode masked regions (zero positions) and dilate the unmasked ones, leading to the vanishing problem. We use sparse convolution to overcome this undesired property by skipping all masked positions and keeping the mask pattern.
• Figure 4: Reconstruction examples by a pre-trained ConvNeXt-Base with a mask ratio of $60\%$. Images are randomly selected from ImageNet validation set. Several interesting regions are highlighted.