[Paper Review] BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

Zhiliang Peng, Dong Li+3 more|arXiv (Cornell University)|Aug 12, 2022

Advanced Neural Network Applications114 citations

TL;DR

BEiT v2 introduces a vector-quantized visual tokenizer trained with knowledge distillation to provide semantic targets for masked image modeling, plus a patch aggregation strategy to enhance global representation; it achieves state-of-the-art results on ImageNet classification and ADE20K segmentation.

ABSTRACT

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.

Motivation & Objective

• Promote masked image modeling from pixel-level to semantic-level using a semantic visual tokenizer.
• Train a VQ-KD tokenizer to discretize semantic space with a learnable codebook.
• Enhance BEiT pretraining with a patch aggregation strategy to improve global representations.
• Empirically validate BEiT v2 on ImageNet classification (fine-tune and linear probe) and ADE20K semantic segmentation.
• Demonstrate robustness and transferability across model sizes and pretraining schedules.

Proposed method

• Train a visual tokenizer with vector-quantized knowledge distillation (VQ-KD) where a ViT encoder maps an image to patch features, a codebook selects nearest codes, and a decoder learns to reconstruct teacher features using cosine similarity as supervision.
• Use an OpenAI CLIP-B/16 as the teacher for semantic guidance in VQ-KD and employ stop-gradient to backpropagate through the encoder via the quantizer.
• Pretrain ViT models with BEiT v2 by predicting the discrete visual tokens for masked patches, following a 40% mask ratio and patch-wise token prediction.
• Introduce a patch aggregation mechanism that encourages the [CLS] token to aggregate patch information, aligning local patch representations with global image representations.
• Pretraining objective combines standard MIM loss (reconstruct visual tokens) with a CLS-focused loss that promotes global representation quality.
• Optionally include a shallow Transformer decoder during pretraining to reinforce CLS-based global representation without increasing inference cost.

Experimental results

Research questions

• RQ1 Can a semantic, vector-quantized visual tokenizer improve masked image modeling over pixel-level targets?
• RQ2 Does a patch aggregation strategy foster better global image representations in MIM pretraining?
• RQ3 How does BEiT v2 perform on ImageNet fine-tuning, linear probing, and ADE20K semantic segmentation across model sizes and pretraining lengths?
• RQ4 What is the impact of VQ-KD target choices (DINO/CLIP) on downstream performance?
• RQ5 How robust is BEiT v2 to different downstream evaluation settings and robustness benchmarks?

Key findings

• BEiT v2 with ViT-B/16 at 300 epochs achieves 85.0% top-1 fine-tuning accuracy on ImageNet-1K and 52.7% ADE20k mIoU.
• BEiT v2 with ViT-B/16 at 1600 epochs achieves 85.5% top-1 (ImageNet) and 53.1% ADE20k mIoU; with intermediate ImageNet-21k fine-tuning, precision improves to 86.5% and 53.5% respectively.
• BEiT v2 with ViT-L/16 at 300 epochs reaches 86.6% top-1 (ImageNet) and 55.0% ADE20k mIoU; 1600 epochs yields 87.3% top-1 and 56.7% ADE20k mIoU.
• Linear probing on ImageNet-1K shows BEiT v2 achieving 80.1% top-1, outperforming BEiT, CAE, MVP, and MoCo v3 in reported comparisons.
• BEiT v2 demonstrates strong robustness across ImageNet variants (Adversarial, Rendition, Sketch) compared to MAE, with substantial gains in several settings.
• Ablation studies show deeper VQ-KD decoders reduce codebook usage and downstream performance, while smaller codebooks improve usage and accuracy.