mc-BEiT: Multi-choice Discretization for Image BERT Pre-training

TL;DR

mc-BEiT introduces a multi-choice discretization approach for image BERT pre-training, improving the quality of masked image modeling by using soft supervision and inter-patch perception, leading to superior performance across vision tasks.

Contribution

It proposes a novel multi-choice training objective for masked image modeling that refines discretization using soft probabilities and inter-patch relations, enhancing image pre-training.

Findings

Achieves 84.1% top-1 accuracy on ImageNet-1K.

Outperforms counterparts in object detection and segmentation on COCO.

Excels in semantic segmentation on ADE20K.

Abstract

Image BERT pre-training with masked image modeling (MIM) becomes a popular practice to cope with self-supervised representation learning. A seminal work, BEiT, casts MIM as a classification task with a visual vocabulary, tokenizing the continuous visual signals into discrete vision tokens using a pre-learned dVAE. Despite a feasible solution, the improper discretization hinders further improvements of image pre-training. Since image discretization has no ground-truth answers, we believe that the masked patch should not be assigned with a unique token id even if a better tokenizer can be obtained. In this work, we introduce an improved BERT-style image pre-training method, namely mc-BEiT, which performs MIM proxy tasks towards eased and refined multi-choice training objectives. Specifically, the multi-choice supervision for the masked image patches is formed by the soft probability vectors of the discrete token ids, which are predicted by the off-the-shelf image tokenizer and further refined by high-level inter-patch perceptions resorting to the observation that similar patches should share their choices. Extensive experiments on classification, segmentation, and detection tasks demonstrate the superiority of our method, e.g., the pre-trained ViT-B achieves 84.1% top-1 fine-tuning accuracy on ImageNet-1K classification, 49.2% AP^b and 44.0% AP^m of object detection and instance segmentation on COCO, 50.8% mIOU on ADE20K semantic segmentation, outperforming the competitive counterparts. The code will be available at https://github.com/lixiaotong97/mc-BEiT.