Image BERT Pre-training with Online Tokenizer

1 INTRODUCTION Masked Language Modeling ( 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 largecapacity models and datasets , becoming a defacto 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 pretraining 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 Masked Image Modeling ( MIM ) for training better Vision Transform- ers 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 can not be ex- tracted 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 image BERT pre-training with Online Tokenizer , 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 highlevel 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-NN , linear probing and fine-tuning protocols to 77.1 % , 79.5 % , 83.8 % with ViT-Base/16 respectively , which is 1.0 % , 1.3 % , 0.2 % higher than previous best results . When pre-trained with ImageNet-22K , iBOT with ViT-L/16 achieves a linear probing accuracy of 81.7 % and a fine-tuning accuracy of 86.3 % , which is 0.4 % and 0.3 % 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 x = { xi } Ni=1 , MIM first samples a random mask m ∈ { 0 , 1 } N according to a prediction ratio r , where N is the number of tokens . The patch token xi where mi being 1 , denoted as x̃ , { xi | mi = 1 } , are then replaced with a mask token e [ MASK ] , yielding a corrupted image x̂ , { x̂i | ( 1−mi ) xi+mie [ MASK ] } Ni=1 . MIM is to recover the masked tokens x̃ from the corrupted image x̂ , i.e. , to maximize : log qθ ( x̃|x̂ ) ≈ ∑N i=1 mi · log qθ ( xi|x̂ ) , where ≈ holds with an independence assumption that each masked token can be reconstructed separately . In BEiT ( Bao et al. , 2021 ) , qθ is modelled as a categorical distribution and the task is to minimize − N∑ i=1 mi · Pφ ( xi ) T log Pθ ( x̂i ) , ( 1 ) where P ( · ) transforms the input to a probability distribution overK dimensions , andφ is parameters of a discrete VAE ( Ramesh et al. , 2021 ) that clusters image patches into 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 φ to current model parameterized by θ . 2.2 SELF-DISTILLATION . Self-distillation , proposed recently in DINO ( Caron et al. , 2021 ) , distills knowledge not from posterior distributions Pφ ( x ) but past iterations of model itself Pθ′ ( x ) and is cast as a discriminative self-supervised objective . Given the training set I , an image x ∼ I is sampled uniformly , over which two random augmentations are applied , yielding two distorted views u and v. The two distorted views are then put through a teacher-student framework to get the predictive categorical distributions from the [ CLS ] token : v [ CLS ] t = P [ CLS ] θ′ ( v ) and u [ CLS ] s = P [ CLS ] θ ( u ) . The knowledge is distilled from teacher to student by minimizing their cross-entropy , formulated as L [ CLS ] = −P [ CLS ] θ′ ( v ) T log P [ CLS ] θ ( u ) . ( 2 ) The teacher and the student share the same architecture consisting of a backbone f ( e.g. , ViT ) and a projection head h [ CLS ] . The parameters of the student network θ are Exponentially Moving Averaged ( EMA ) to the parameters of teacher network θ′ . The loss is symmetrized by averaging with another cross-entropy term between v [ CLS ] s and u [ CLS ] t . 3 IBOT . We motivate our method by identifying the similar formulation of Eq . ( 1 ) and Eq . ( 2 ) . A visual tokenizer parameterized by online θ′ instead of pre-fixed φ 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 pseudocode 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 u and v and obtain their masked views û and v̂ . Taking û as an example for simplicity , the student network outputs for the masked view û projections of its patch tokens ûpatchs = P patch θ ( û ) and the teacher network outputs for the non-masked view u projections of its patch tokens upatcht = P patch θ′ ( u ) . We here define the training objective of MIM in iBOT as LMIM = − N∑ i=1 mi · P patchθ′ ( ui ) T log P patch θ ( ûi ) . ( 3 ) We symmetrize the loss by averaging with another CE term between v̂patchs and v patch t . The backbone together with the projection head of teacher network hpatcht ◦ ft 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 L [ CLS ] in Eq . ( 2 ) proposed in DINO ( Caron et al. , 2021 ) , except that now we have û [ CLS ] s instead of u [ CLS ] s 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 [ CLS ] s = h patch s , h [ CLS ] t = h patch t . 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 .

Summary: This paper presents a new strategy for self-supervised training of vision transformer models by combining the self-distillation-based method DINO with mask image modeling. The momentum teacher model in DINO is taken as an online tokenizer. A mask image modeling loss that computes self-distillation loss between the reconstructed tokens and the output tokens of the momentum teacher is proposed. Extensive experiments are conducted to show the effectiveness of the proposed self-supervised training framework. The result on downstream tasks including classification, object detection, and segmentation also shows the benefit of the mask image modeling loss.