[Paper Review] iBOT: Image BERT Pre-Training with Online Tokenizer

Jinghao Zhou, Wei Chen+5 more|arXiv (Cornell University)|Nov 15, 2021

Multimodal Machine Learning ApplicationsComputer Science209 citations

TL;DR

iBOT introduces an online visual tokenizer for masked image modeling via self-distillation, achieving state-of-the-art results on ImageNet and strong performance on dense downstream tasks.

ABSTRACT

The success of language Transformers is primarily attributed to the pretext task of masked language modeling (MLM), where texts are first tokenized into semantically meaningful pieces. In this work, we study masked image modeling (MIM) and indicate the advantages and challenges of using a semantically meaningful visual tokenizer. We present a self-supervised framework iBOT that can perform masked prediction with an online tokenizer. Specifically, we perform self-distillation on masked patch tokens and take the teacher network as the online tokenizer, along with self-distillation on the class token to acquire visual semantics. The online tokenizer is jointly learnable with the MIM objective and dispenses with a multi-stage training pipeline where the tokenizer needs to be pre-trained beforehand. We show the prominence of iBOT by achieving an 82.3% linear probing accuracy and an 87.8% fine-tuning accuracy evaluated on ImageNet-1K. Beyond the state-of-the-art image classification results, we underline emerging local semantic patterns, which helps the models to obtain strong robustness against common corruptions and achieve leading results on dense downstream tasks, eg., object detection, instance segmentation, and semantic segmentation.

Motivation & Objective

• Motivate a BERT-like pre-training paradigm for vision by leveraging semantically meaningful visual tokens.
• Eliminate the need for a pre-trained offline tokenizer by learning the tokenizer online together with the model.
• Improve masked image modeling (MIM) through knowledge distillation using a teacher tokenizer derived from the online tokenizer.
• Explore how jointly learning token semantics enhances robustness and downstream performance across classification and dense tasks.

Proposed method

• Formulate masked image modeling as knowledge distillation between a teacher (online tokenizer) and a student (Vision Transformer).
• Use two losses: a cross-view [CLS] self-distillation loss to acquire visual semantics and an MIM loss to reconstruct masked patch tokens using teacher outputs.
• Share projection heads between [CLS] and patch tokens to propagate semantic information.
• Implement an online tokenizer that is updated via momentum with the student, removing the need for offline pre-training of a tokenizer.
• Employ self-distillation on cross-view [CLS] tokens to bootstrap meaningful visual semantics and use softmax token distributions instead of hard one-hot tokens as supervision.
• Evaluate with ViT and Swin backbones, pre-training on ImageNet-1K and ImageNet-22K, followed by linear probing, k-NN, and fine-tuning.

Experimental results

Research questions

• RQ1 Can masked image modeling with an online, jointly learned visual tokenizer outperform offline-tokenizer approaches in self-supervised pre-training of Vision Transformers?
• RQ2 Do self-distillation mechanisms on [CLS] tokens and patch-level MIM signals yield stronger semantic representations and robustness for downstream tasks?
• RQ3 How does sharing projection heads between [CLS] and patch tokens affect learned semantics and performance?
• RQ4 What is the impact of online tokenizer semantics on linear probing, fine-tuning, and transfer to dense vision tasks?

Key findings

• iBOT achieves state-of-the-art results in multiple evaluation setups on ImageNet-1K, including 82.3% linear probing and 87.8% fine-tuning accuracy when pre-trained on ImageNet-22K with ViT-L/16.
• On standard ImageNet-1K, iBOT with ViT-S/16, ViT-B/16, and ViT-L/16 reaches top-1 accuracies up to 84.8% in fine-tuning and 82.3% in linear evaluation on larger pre-training data, surpassing prior SSL methods.
• iBOT discovers emergent part-level semantics in patch tokens, contributing to improved robustness to corruptions and better performance on dense tasks such as object detection, instance segmentation, and semantic segmentation.
• Compared to DINO, iBOT shows larger gains with bigger models, indicating stronger scalability of the online-tokenizer approach.
• Across transfer learning benchmarks on small datasets (e.g., CIFAR, Flowers, Cars) and larger domain datasets (iNaturalist 18/19), iBOT consistently outperforms BEiT and DINO baselines, particularly with larger backbone models.
• iBOT demonstrates improved robustness to background changes, occlusion, and out-of-distribution data relative to strong baselines.