Source-linked AI summary

BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, Furu Wei

arXiv:2208.06366v2cs.CV

TL;DR

BEiT v2 addresses the limited use of high-level semantics in masked image modeling by using semantic visual tokens as reconstruction targets. It trains those tokens with vector-quantized knowledge distillation and adds patch aggregation, achieving stronger results across image classification and semantic segmentation tasks.

  • Problem

    Most MIM methods reconstruct low-level image elements, leaving high-level semantics underused in visual representation learning.

  • Method

    BEiT v2 trains a semantic visual tokenizer with vector-quantized knowledge distillation, then predicts its discrete codes and uses patch aggregation during pretraining.

  • Results

    BEiT v2 outperforms compared MIM methods across image classification, linear probing, and semantic segmentation, including 85.5% top-1 accuracy for base-size ImageNet-1K fine-tuning.

  • Takeaways & Limitations

    Semantic visual tokens and explicit global representation aggregation improve BEiT pretraining and transfer performance on downstream tasks.

Abstract

from arXiv · show

Masked image modeling (MIM) has demonstrated impressive results in self-supervised representation learning by recovering corrupted image patches. However, most existing studies operate on low-level image pixels, which hinders the exploitation of high-level semantics for representation models. In this work, we propose to use a semantic-rich visual tokenizer as the reconstruction target for masked prediction, providing a systematic way to promote MIM from pixel-level to semantic-level. Specifically, we propose vector-quantized knowledge distillation to train the tokenizer, which discretizes a continuous semantic space to compact codes. We then pretrain vision Transformers by predicting the original visual tokens for the masked image patches. Furthermore, we introduce a patch aggregation strategy which associates discrete image patches to enhance global semantic representation. Experiments on image classification and semantic segmentation show that BEiT v2 outperforms all compared MIM methods. On ImageNet-1K (224 size), the base-size BEiT v2 achieves 85.5% top-1 accuracy for fine-tuning and 80.1% top-1 accuracy for linear probing. The large-size BEiT v2 obtains 87.3% top-1 accuracy for ImageNet-1K (224 size) fine-tuning, and 56.7% mIoU on ADE20K for semantic segmentation. The code and pretrained models are available at https://aka.ms/beitv2.

1 Introduction

BEiT v2 advances masked image modeling from low-level reconstruction toward semantic-aware supervision through a learned visual tokenizer and patch aggregation. It reports stronger performance across model sizes, training schedules, and downstream tasks.

  • Motivation: Existing MIM methods largely reconstruct low-level image elements, underusing high-level semantic information.This contrasts with masked language modeling, where masked words carry high-level semantics.
  • Method: BEiT v2 trains a semantic-aware visual tokenizer with vector-quantized knowledge distillation to discretize a continuous semantic space.The tokenizer’s encoder produces discrete codes used as supervision during BEiT pretraining.
  • Method: Patch aggregation encourages the [CLS] token to associate all discrete image patches, strengthening global image representations.The strategy addresses the emphasis on patch reconstruction in MIM and supports improved linear-probing performance.
  • Results: BEiT v2 improves performance across ImageNet fine-tuning, linear probing, and ADE20K semantic segmentation evaluations.The experiments cover base- and large-size vision Transformers and multiple pretraining schedules.
  • Results: BEiT v2 improves over BEiT by about two points on ImageNet fine-tuning for both ViT-B/16 and ViT-L/16.Figure 1 reports top-1 fine-tuning accuracy on ImageNet at 224-pixel size for these two model variants.

2 Methodology

BEIT V2 trains a semantic-aware visual tokenizer with vector-quantized knowledge distillation, then pretrains Transformers to predict masked visual tokens. Patch aggregation additionally trains the [CLS] token to capture global image representations.

  • 2.1 Image Representation: BEIT V2 uses a visual tokenizer that converts each image into discrete visual tokens corresponding to image patches.The tokenizer comprises a vision Transformer encoder and a quantizer with a learnable codebook.
  • 2.2 Training Visual Tokenizer: Vector-quantized knowledge distillation selects nearest codebook embeddings and trains a decoder to reconstruct teacher-model semantic features.Teacher guidance can come from models such as DINO or CLIP; gradients are copied through the non-differentiable quantization step.
  • 2.2 Training Visual Tokenizer: The tokenizer feeds normalized quantized codebook embeddings to a Transformer decoder whose outputs maximize cosine similarity with teacher feature vectors.This makes the discrete reconstruction targets semantic-feature-based rather than raw-pixel-based.
  • 2.3 Pretraining BEIT V2: Patch aggregation adds an auxiliary masked-prediction pathway that encourages the [CLS] token to aggregate patch information into global representations.The final loss sums the original MIM loss with the shallow decoder's MIM loss, while the added decoder is discarded after pretraining.
  • 2.3 Pretraining BEIT V2: During MIM pretraining, around 40% of image patches are block-wise masked and replaced by a shared learnable embedding.The model predicts the tokenizer-derived visual tokens at masked positions using a softmax MIM head.

3 Experiments

BEIT V2 is evaluated across image classification, linear probing, robustness, semantic segmentation, and ablation settings, showing strong performance and benefits from its design choices.

  • Experimental setup: Experiments evaluate ImageNet classification, linear probing, semantic segmentation, robustness, and ablations using base- and large-size Vision Transformers.ImageNet-1K is used for pretraining and classification; ADE20K is used for semantic segmentation.
  • Image classification: 85.0% top-1 accuracy is achieved by base-size BEIT V2 after 300 epochs, exceeding BEIT, CAE, SplitMask, and PeCo by 2.1%, 1.4%, 1.4%, and 0.9%.The result is reported for ImageNet-1K fine-tuning.
  • Image classification: 87.3% top-1 accuracy is reached by ViT-L/16 with a longer pretraining schedule, while 85.5% is achieved by BEIT V2 under a longer schedule.The 85.5% result is reported as a new state of the art among self-supervised methods on ImageNet-1K.
  • Linear probing: BEIT V2 respectively outperforms BEIT, CAE, and MVP by 23.4%, 16.0%, and 4.7% in ImageNet-1K linear probing.All compared methods use ViT-B/16 and 300 pretraining epochs except MAE, which uses 1600 epochs.
  • Robustness: BEIT V2 achieves gains across ImageNet-Adversarial, ImageNet-Rendition, and ImageNet-Sketch compared with MAE, indicating stronger model generalization.The evaluation uses three ImageNet validation variants.
  • Semantic segmentation: 56.7 mIoU is achieved by ViT-L/16 on ADE20K, establishing a new state of the art for masked image modeling on that benchmark.Semantic segmentation uses UperNet with 512 × 512 input resolution.
  • Ablation analysis: Deeper VQ-KD decoders improve reconstruction but reduce codebook usage and downstream task performance, while reducing lookup dimension improves codebook utilization.The ablation uses ViT-B/16 without patch aggregation as the baseline and 300 epochs of pretraining.
  • Ablation analysis: A shallower patch-aggregation head performs better than a deeper head, and patch aggregation outperforms the baseline without it.The findings suggest that the shallower head emphasizes the input [CLS] token more strongly.

4 Related Work

Related work on visual tokenizers and masked image modeling motivates BEIT V2’s semantic visual-token reconstruction approach.

  • Visual tokenizer: Visual tokenizers convert images into discrete codes and reconstruct inputs; later systems use Transformer blocks, adversarial losses, or perceptual losses to preserve details.The related methods include VQ-VAE, DALL-E, VQGAN, and ViT-VQGAN.
  • Masked image modeling: Masked image modeling has been adapted from language tasks to vision, with methods exploring discrete visual tokens, pixel reconstruction, and distilled features as prediction targets.BEIT recovers discrete visual tokens, whereas MAE performs pixel-level denoising reconstruction.

5 Conclusion

BEIT V2 uses vector-quantized knowledge distillation to train a semantic visual tokenizer, replacing pixel supervision with discrete semantic targets for masked image modeling. A patch aggregation mechanism further encourages global image representations and improves downstream transfer.

  • VQ-KD discretizes a continuous semantic space into supervision targets for masked image modeling instead of relying on image pixels.
  • Patch aggregation explicitly encourages global image representations and narrows the gap between patch-level pretraining and image-level aggregation.
  • The learned semantic tokenizer and patch aggregation significantly boost transfer performance on image classification and semantic segmentation.

A Visualization of Codebook

Discrete codes tend to capture explicit semantic concepts while ignoring image details such as color, illumination, rotation, and scale. However, VQ-KD can mismatch concepts when local structures are similar.

  • Discrete codes tend to represent explicit semantic concepts while ignoring color, illumination, rotation, and scale.
  • VQ-KD misclassifies a spot on a fish body as an eye because of local structural similarity.
  • Figure 5 contrasts patches matching learned semantic concepts with patches that mismatch them.

B Comparison with Large-scale Supervised Pretraining

BEIT V2 shows strong ImageNet-1K fine-tuning performance at large scale and resolution. With ViT-L/16 at 384×384, it reaches 89.0% top-1 accuracy, exceeding ViT-H/14 trained on Google JFT-3B by 0.5%.

  • 89.0% top-1 accuracy is achieved by BEIT V2 with ViT-L/16 at 384×384 on ImageNet-1K fine-tuning.
  • BEIT V2 exceeds ViT-H/14 using the Google JFT-3B labeled dataset by 0.5% in this comparison.
  • The reported gain is presented as evidence of BEIT V2's data efficiency and superiority.
  • Table 7 reports top-1 accuracy on ImageNet-1K fine-tuning and distinguishes 224^2 and 384^2 model resolutions.

C Overall Framework for BEIT V2

The overall BEIT V2 pretraining process is presented as a framework spanning tokenizer training and BEIT V2 pretraining. Figure 6 summarizes this complete process.

  • The paper presents tokenizer training and BEIT V2 pretraining separately before showing the whole pretraining process.
  • Figure 6 presents the overall framework for BEIT V2 pretraining.

D Hyperparameters for VQ-KD Training

Table 8 presents the hyperparameters used to train VQ-KD on ImageNet-1K.

  • Table 8 lists hyperparameters for VQ-KD training.
  • The listed training configuration targets ImageNet-1K.
  • This table documents the setup for the VQ-KD training stage.

E Hyperparameters for BEIT V2 pretraining

Table 9 presents BEIT V2 pretraining hyperparameters on ImageNet-1K, including a condition for 300-epoch schedules.

  • Table 9 lists hyperparameters for BEIT V2 pretraining on ImageNet-1K.
  • The configuration is specified for ImageNet-1K pretraining.
  • An asterisk marks hyperparameters adopted when the pretraining schedule is 300 epochs.

F Hyperparameters for Image Classification Fine-tuning

Tables 10 and 11 provide BEIT V2 fine-tuning hyperparameters for ImageNet-1K and ADE20K.

  • ImageNet-1K: Table 10 lists hyperparameters for fine-tuning BEIT V2 on ImageNet-1K.
  • Together, the tables document BEIT V2 fine-tuning configurations for image classification and semantic segmentation datasets.
  • ADE20K: Table 11 lists hyperparameters for fine-tuning BEIT V2 on ADE20K.
Loading 2208.06366v2…