iBOT: Image BERT Pre-Training with Online Tokenizer — part 1/4

Cleaned-text code points [0, 15494). Bibliography omitted here; source.md includes it.

# iBOT : Image BERT Pre-Training with Online Tokenizer

###### Abstract

The success of language Transformers is primarily attributed to the pretext task of masked language modeling (MLM) ( Devlin et al., 2019 ) , 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, e.g . , object detection, instance segmentation, and semantic segmentation. The code and models are publicly available at https://github.com/bytedance/ibot .

## 1 Introduction

M asked L anguage M odeling (MLM), which first randomly masks and then reconstructs a set of input tokens, is a popular pre-training paradigm for language models. The MLM pre-trained Transformers ( Devlin et al., 2019 ) have demonstrated their scalability to large-capacity models and datasets, becoming a de-facto standard for lingual tasks. However, its potential for Vision Transformer (ViT), which recently started to revolutionize computer vision research ( Touvron et al., 2021 ; Dosovitskiy et al., 2021 ) , has been largely underexplored. Most popular unsupervised pre-training schemes in vision deal with the global views ( Chen et al., 2021 ; Caron et al., 2021 ) , neglecting images’ internal structures, as opposed to MLM modeling local tokens. In this work, we seek to continue the success of MLM and explore M asked I mage M odeling (MIM) for training better Vision Transformers such that it can serve as a standard component, as it does for NLP.

One of the most crucial components in MLM is the lingual tokenizer which splits language into semantically meaningful tokens, e.g . , WordPiece ( Wu et al., 2016 ) in BERT. Similarly, the crux of MIM lies in a proper design of visual tokenizer , which transforms the masked patches to supervisory signals for the target model, as shown in Fig. 2 . However, unlike lingual semantics arising naturally from the statistical analysis of word frequency ( Sennrich et al., 2016 ) , visual semantics cannot be extracted such easily due to the continuous property of images. Empirically, visual semantics emerges progressively by bootstrapping online representation that enforces a similarity of distorted image views ( He et al., 2020 ; Grill et al., 2020 ; Caron et al., 2020 ) . This property intuitively indicates a multi-stage training pipeline, where we need to first train an off-the-shelf semantic-rich tokenizer before training the target model. However, since acquiring visual semantics is a common end for both the tokenizer and target model, a single-stage training pipeline where the tokenizer and target model can be jointly optimized awaits further exploration.

Previous works partially tackle the above challenges. Several works use identity mapping as the visual tokenizer, i.e . , predicting the raw pixel values ( Pathak et al., 2016 ; Atito et al., 2021 ) . Such paradigm struggles in semantic abstraction and wastes the capacity at modeling high-frequency details, yielding less competitive performance in semantic understanding ( Liu et al., 2021a ) . Recently, BEiT ( Bao et al., 2021 ) proposes to use a pre-trained discrete VAE ( Ramesh et al., 2021 ) as the tokenizer. Though providing some level of abstraction, the discrete VAE is still found only to capture low-level semantics within local details (as observed by Tab. 9 ). Moreover, the tokenizer needs to be offline pre-trained with fixed model architectures and extra dataset ( Ramesh et al., 2021 ) , which potentially limits its adapativity to perform MIM using data from different domains.

To this end, we present iBOT , short for i mage B ERT pre-training with O nline T okenizer, a new framework that performs MIM with a tokenizer handling above-mentioned challenges favorably. We motivate iBOT by formulating the MIM as knowledge distillation (KD), which learns to distill knowledge from the tokenizer, and further propose to perform self-distillation for MIM with the help of twin teacher as online tokenizer. The target network is fed with a masked image while the online tokenizer with the original image. The goal is to let the target network recover each masked patch token to its corresponding tokenizer output. Our online tokenizer naturally resolves two major challenges. On the one hand, our tokenizer captures high-level visual semantics progressively learned by enforcing the similarity of cross-view images on class tokens. On the other hand, our tokenizer needs no extra stages of training as pre-processing setup since it is jointly optimized with MIM via momentum update.

The online tokenizer enables iBOT to achieve excellent performance for feature representation. Specifically, iBOT advances ImageNet-1K classification benchmark under k k -NN, linear probing and fine-tuning protocols to 77.1%, 79.5%, 84.0% with ViT-Base/16 respectively, which is 1.0%, 1.3%, 0.4% higher than previous best results. When pre-trained with ImageNet-22K, iBOT with ViT-L/16 achieves a linear probing accuracy of 82.3% and a fine-tuning accuracy of 87.8% , which is 1.0% and 1.8% higher than previous best results. Beyond that, the advancement is also valid when transferring to other datasets or under semi-supervised and unsupervised classification settings. Of particular interest, we have identified an emerging part-level semantics that can help the model with image recognition both on global and local scales. We identify that the semantic patterns learned in patch tokens, which sufficiently lack in the off-line tokenizer as in BEiT ( Bao et al., 2021 ) , helps the model to be advanced in linear classification and robustness against common image corruptions. When it is transferred to downstream tasks, we show that in downstream tasks related to image classification, object detection, instance segmentation, and semantic segmentation, iBOT surpasses previous methods with nontrivial margins. All of the evidence demonstrates that iBOT has largely closed the gap of masked modeling pre-training between language and vision Transformers.

## 2 Preliminaries

### 2.1 Masked Image Modeling as Knowledge Distillation

Masked image modeling (MIM), which takes a similar formulation as MLM in BERT, has been proposed in several recent works ( Bao et al., 2021 ; Tan et al., 2021 ) . Specifically, for an image token sequence 𝒙 = { 𝒙 i } i = 1 N {\bm{x}}=\{{\bm{x}}_{i}\}_{i=1}^{N} , MIM first samples a random mask 𝒎 ∈ { 0 , 1 } N {\bm{m}}\in\{0,1\}^{N} according to a prediction ratio r r , where N N is the number of tokens. The patch token 𝒙 i {\bm{x}}_{i} where m i m_{i} being 1 1 , denoted as 𝒙 ~ ≜ { 𝒙 i | m i = 1 } \tilde{\bm{x}}\triangleq\{{\bm{x}}_{i}\ |\ m_{i}=1\} , are then replaced with a mask token 𝒆 [MASK] {\bm{e}}_{\texttt{[MASK]}} , yielding a corrupted image 𝒙 ^ ≜ { 𝒙 ^ i | ( 1 − m i ) ​ 𝒙 i + m i ​ 𝒆 [MASK] } i = 1 N \hat{\bm{x}}\triangleq\{\hat{\bm{x}}_{i}\ |\ (1-m_{i}){\bm{x}}_{i}+m_{i}{\bm{e}}_{\texttt{[MASK]}}\}_{i=1}^{N} . MIM is to recover the masked tokens 𝒙 ~ \tilde{\bm{x}} from the corrupted image 𝒙 ^ \hat{\bm{x}} , i.e . , to maximize: log ​ q 𝜽 ​ ( 𝒙 ~ | 𝒙 ^ ) ≈ ∑ i = 1 N m i ⋅ log ​ q 𝜽 ​ ( 𝒙 i | 𝒙 ^ ) \mathrm{log}\ q_{\bm{\theta}}(\tilde{\bm{x}}|\hat{\bm{x}})\approx\sum_{i=1}^{N}\ m_{i}\cdot\mathrm{log}\ q_{\bm{\theta}}({\bm{x}}_{i}|\hat{\bm{x}}) , where ≈ \approx holds with an independence assumption that each masked token can be reconstructed separately. In BEiT ( Bao et al., 2021 ) , q 𝜽 q_{\bm{\theta}} is modelled as a categorical distribution and the task is to minimize − ∑ i = 1 N m i ⋅ P ϕ ( 𝒙 i ) T log P 𝜽 ( 𝒙 ^ i ) , -\sum_{i=1}^{N}\ m_{i}\cdot P_{\bm{\phi}}({\bm{x}}_{i})^{\mathrm{T}}\ \mathrm{log}\ P_{\bm{\theta}}(\hat{\bm{x}}_{i}), (1) where P ⁡ ( ⋅ ) P(\cdot) transforms the input to a probability distribution over K K dimensions, and ϕ {\bm{\phi}} is parameters of a discrete VAE ( Ramesh et al., 2021 ) that clusters image patches into K K categories and assigns each patch token a one-hot encoding identifying its category. We note this loss is formulated similarly to knowledge distillation ( Hinton et al., 2015 ) , where knowledge is distilled from a pre-fixed tokenizer parameterized by ϕ {\bm{\phi}} to current model parameterized by 𝜽 {\bm{\theta}} .

### 2.2 Self-Distillation

Self-distillation, proposed recently in DINO ( Caron et al., 2021 ) , distills knowledge not from posterior distributions P ϕ ​ ( 𝒙 ) P_{\bm{\phi}}({\bm{x}}) but past iterations of model itself P 𝜽 ′ ​ ( 𝒙 ) P_{{\bm{\theta^{\prime}}}}({\bm{x}}) and is cast as a discriminative self-supervised objective. Given the training set ℐ \mathcal{I} , an image 𝒙 ∼ ℐ {\bm{x}}\sim\mathcal{I} is sampled uniformly, over which two random augmentations are applied, yielding two distorted views 𝒖 {\bm{u}} and 𝒗 {\bm{v}} . The two distorted views are then put through a teacher-student framework to get the predictive categorical distributions from the [CLS] token: 𝒗 t [CLS] = P 𝜽 ′ [CLS] ​ ( 𝒗 ) {\bm{v}}_{t}^{\texttt{[CLS]}}=P_{{\bm{\theta}}^{\prime}}^{\texttt{[CLS]}}({\bm{v}}) and 𝒖 s [CLS] = P 𝜽 [CLS] ​ ( 𝒖 ) {\bm{u}}_{s}^{\texttt{[CLS]}}=P_{{\bm{\theta}}}^{\texttt{[CLS]}}({\bm{u}}) . The knowledge is distilled from teacher to student by minimizing their cross-entropy, formulated as ℒ [CLS] = − P 𝜽 ′ [CLS] ​ ( 𝒗 ) T ​ log ​ P 𝜽 [CLS] ​ ( 𝒖 ) . \mathcal{L}_{\texttt{[CLS]}}=-P_{\bm{\theta^{\prime}}}^{\texttt{[CLS]}}({\bm{v}})^{\mathrm{T}}\ \mathrm{log}\ P_{\bm{\theta}}^{\texttt{[CLS]}}({\bm{u}}). (2) The teacher and the student share the same architecture consisting of a backbone f f ( e.g . , ViT) and a projection head h [CLS] h^{\texttt{[CLS]}} . The parameters of the student network 𝜽 {\bm{\theta}} are Exponentially Moving Averaged (EMA) to the parameters of teacher network 𝜽 ′ {\bm{\theta^{\prime}}} . The loss is symmetrized by averaging with another cross-entropy term between 𝒗 s [CLS] {\bm{v}}_{s}^{\texttt{[CLS]}} and 𝒖 t [CLS] {\bm{u}}_{t}^{\texttt{[CLS]}} .

## 3 iBOT

We motivate our method by identifying the similar formulation of Eq. ( 1 ) and Eq. ( 2 ). A visual tokenizer parameterized by online 𝜽 ′ {\bm{\theta^{\prime}}} instead of pre-fixed ϕ {\bm{\phi}} thus arises naturally. In this section, we present iBOT, casting self-distillation as a token-generation self-supervised objective and perform MIM via self-distillation. We illustrate the framework of iBOT in Fig. 3 and demonstrate the pseudo-code in Appendix A . In Sec. 3.2 , we briefly introduce the architecture and pre-training setup.

### 3.1 Framework

First, we perform blockwise masking ( Bao et al., 2021 ) on the two augmented views 𝒖 {\bm{u}} and 𝒗 {\bm{v}} and obtain their masked views 𝒖 ^ \hat{\bm{u}} and 𝒗 ^ \hat{\bm{v}} . Taking 𝒖 ^ \hat{\bm{u}} as an example for simplicity, the student network outputs for the masked view 𝒖 ^ \hat{\bm{u}} projections of its patch tokens 𝒖 ^ s patch = P 𝜽 patch ​ ( 𝒖 ^ ) \hat{\bm{u}}_{s}^{\mathrm{patch}}=P_{{\bm{\theta}}}^{\mathrm{patch}}(\hat{\bm{u}}) and the teacher network outputs for the non-masked view 𝒖 {\bm{u}} projections of its patch tokens 𝒖 t patch = P 𝜽 ′ patch ​ ( 𝒖 ) {\bm{u}}_{t}^{\mathrm{patch}}=P_{{\bm{\theta^{\prime}}}}^{\mathrm{patch}}({\bm{u}}) . We here define the training objective of MIM in iBOT as ℒ MIM = − ∑ i = 1 N m i ⋅ P 𝜽 ′ patch ( 𝒖 i ) T log P 𝜽 patch ( 𝒖 ^ i ) . \mathcal{L}_{\mathrm{MIM}}=-\sum_{i=1}^{N}\ m_{i}\cdot P_{\bm{\theta^{\prime}}}^{\mathrm{patch}}({\bm{u}}_{i})^{\mathrm{T}}\ \mathrm{log}\ P_{\bm{\theta}}^{\mathrm{patch}}(\hat{\bm{u}}_{i}). (3) We symmetrize the loss by averaging with another CE term between 𝒗 ^ s patch \hat{\bm{v}}_{s}^{\mathrm{patch}} and 𝒗 t patch {\bm{v}}_{t}^{\mathrm{patch}} .

The backbone together with the projection head of teacher network h t patch ∘ f t h_{t}^{\mathrm{patch}}\circ f_{t} is, therefore, a visual tokenizer that generates online token distributions for each masked patch token. The tokenizer used in iBOT is jointly learnable to MIM objective without a need of being pre-trained in an extra stage, a bonus feature of which is now its domain knowledge can be distilled from the current dataset rather than fixed to the specified dataset.

To ensure that the online tokenizer is semantically-meaningful, we perform self-distillation on [CLS] token of cross-view images such that visual semantics can be obtained via bootstrapping, as achieved by the majority of the self-supervised methods ( He et al., 2020 ; Grill et al., 2020 ; Caron et al., 2021 ) . In practice, iBOT works with ℒ [CLS] \mathcal{L}_{\texttt{[CLS]}} in Eq. ( 2 ) proposed in DINO ( Caron et al., 2021 ) , except that now we have 𝒖 ^ s [CLS] \hat{\bm{u}}_{s}^{\texttt{[CLS]}} instead of 𝒖 s [CLS] {\bm{u}}_{s}^{\texttt{[CLS]}} as input for the student network. To further borrow the capability of semantics abstraction acquired from self-distillatin on [CLS] token, we share the parameters of projection heads for [CLS] token and patch tokens, i.e . , h s [CLS] = h s patch h_{s}^{\texttt{[CLS]}}=h_{s}^{\mathrm{patch}} , h t [CLS] = h t patch h_{t}^{\texttt{[CLS]}}=h_{t}^{\mathrm{patch}} . We empirically find that it produces better results than using separate heads.

Unlike tokenized words whose semantics are almost certain, image patch is ambiguous in its semantic meaning. Therefore, tokenization as one-hot discretization can be sub-optimal for images. In iBOT, we use the token distribution after softmax instead of the one-hot token id as a supervisory signal, which plays an important role in iBOT pre-training as shown in Tab. 18 .

### 3.2 Implementation

##### Architecture.

We use the Vision Transformers ( Dosovitskiy et al., 2021 ) and Swin Transformers ( Liu et al., 2021b ) with different amounts of parameters, ViT-S/16, ViT-B/16, ViT-L/16, and Swin-T/{7,14} as the backbone f f . For ViTs, /16 denotes the patch size being 16 16 . For Swins, / { 7 , 14 } /\{7,14\} denotes the window size being 7 7 or 14 14 . We pre-train and fine-tune the Transformers with 224 224 -size images, so the total number of patch tokens is 196 196 . The projection head h h is a 3 3 -layer MLPs with l 2 l_{2} -normalized bottleneck following DINO ( Caron et al., 2021 ) . Towards a better design to acquire visual semantics, we studied different sharing strategies between projection heads h [CLS] h^{\texttt{[CLS]}} and h patch h^{\mathrm{patch}} , considering that semantics obtained in distillation on [CLS] token helps the training of MIM on patch tokens. We empirically find that sharing the entire head prompts the best performance. We set the output dimension of the shared head to 8192 8192 .

##### Pre-Training Setup.