iBOT: Image BERT Pre-Training with Online Tokenizer

In plain words

Language models are often pre-trained by hiding some words and asking the model to guess them. The same idea works for images: hide some patches of a picture and ask the network to predict what was there. The hard part is deciding what "what was there" means. Pixel values are a poor target, because they spend effort on fine detail and say little about what the patch shows. Earlier image methods fixed this by first training a separate model that turns patches into labels, which is an extra stage and ties the result to that model's training data.

iBOT removes the separate stage. It keeps two copies of the same network. The student sees a picture with some patches hidden. The teacher, a slowly updated copy of the student, sees the full picture. The teacher's output for each patch becomes the target the student must match at the hidden positions. Because the teacher is learned at the same time as the student, the targets improve as training goes on. A second loss, the same kind used in earlier self-distillation work, compares the summary token for two differently cropped views of one picture, which pushes the teacher's outputs toward meaningful categories.

The point is a single training run in which the image model learns both a useful summary of each picture and useful descriptions of its parts. The card's reported results show strong classification accuracy and gains on dense tasks such as detection and segmentation, with costs in memory and training time compared with the strongest self-distillation baseline.

TL;DR

iBOT is masked image modeling for Vision Transformers in which the prediction targets come from an EMA teacher instead of a pre-trained tokenizer. The teacher and student are trained jointly with a cross-view [CLS] loss and a masked-patch loss. Author-reported: 82.3% linear probing with ViT-L/16 pre-trained on ImageNet-22K, 84.0% fine-tuning with ViT-B/16 on ImageNet-1K, and gains over DINO on COCO and ADE20K.

Problem

• Masked language modeling works with a discrete vocabulary, but image patches have no natural discrete units §1.
• Pixel regression spends capacity on high-frequency detail and gives weak semantics §1.
• Discrete VAE tokenizers (BEiT) are pre-trained offline on a fixed dataset and architecture, which adds a stage and limits adaptation to new domains §1 §2.1.
• Self-distillation on the [CLS] token gives semantics but does not supply patch-level targets by itself §2.2 §4.4, tab. 9.

Method

iBOT uses a student network g_s and a teacher network g_t, both a ViT backbone followed by a projection head. The teacher parameters are an exponential moving average of the student parameters §2.2 main_ibot.py.

Steps per training iteration:

1. Views. Sample an image and apply two random augmentations to get views u and v §2.2 app. B. The official pre-training uses 2 global crops at 224 px and 10 local crops at 96 px app. B main_ibot.py.
2. Masks. Apply blockwise masking §3.1 eq. 3 app. B. With multi-crop, the paper's chosen pipeline masks only the global crops, and the ratio is sometimes 0 (random MIM) app. B fig. 7. The prediction ratio r is drawn per crop, equal to 0 with probability 0.5 and uniformly from [0.1, 0.5] otherwise §3.2 loader.py.
3. Teacher pass. The teacher reads the unmasked views and outputs a [CLS] distribution and one distribution per patch token §3.1 alg. 1.
4. Student pass. The student reads the masked views and outputs the same quantities §3.1 alg. 1. Masked patch embeddings are replaced by one learned mask vector vision_transformer.py.
5. [CLS] loss. Cross-entropy between the teacher [CLS] distribution of one view and the student [CLS] distribution of the other view, averaged over both orderings §2.2 eq. 2 alg. 1.
6. Patch (MIM) loss. For each view, cross-entropy between the teacher patch distribution of the unmasked view and the student patch distribution of the masked view, summed over masked positions only and averaged over both orderings §3.1 eq. 3 alg. 1.
7. Total loss. Add the two losses with weight 1 each, without rescaling §3.2 main_ibot.py.
8. Updates. Update the student with AdamW. Update the teacher as an EMA of the student. Update both centers as exponential moving averages of the teacher outputs §3.2 alg. 1 main_ibot.py.

The teacher output is a softmax distribution, not a hard one-hot label. Teacher outputs are centred (subtract a running mean) and sharpened (low temperature) before the cross-entropy. Centring and sharpening are applied to [CLS] and patch outputs with separate centres and temperatures §3.1 alg. 1.

The projection head is a three-layer MLP with a normalized bottleneck, following DINO §3.2. The [CLS] and patch heads are shared in the configuration the paper reports as the default §3.2 app. E.

# One iBOT training step  [alg. 1] [eq. 2] [eq. 3]
# Input:  images (B, 3, H, W) with random views per image
# Output: scalar loss; teacher updated by EMA after the student step

def ibot_step(images):
    # Two random augmented views per image  [app. B]
    view_u = augment(images)
    view_v = augment(images)

    # Blockwise masks on the student views only  [§3.1] [code: loader.py#L67]
    masked_u, mask_u = blockwise_mask(view_u)
    masked_v, mask_v = blockwise_mask(view_v)

    # Teacher reads unmasked views, with no gradient  [§3.1] [alg. 1]
    cls_t_u, patch_t_u = teacher(view_u)  # (B, K), (B, N, K)
    cls_t_v, patch_t_v = teacher(view_v)

    # Student reads masked views  [§3.1] [alg. 1]
    cls_s_u, patch_s_u = student(masked_u)
    cls_s_v, patch_s_v = student(masked_v)

    # [CLS] loss across the two views  [§2.2] [eq. 2]
    loss_cls = (
        cross_entropy(cls_s_v, cls_t_u, center_cls, tau_s, tau_t)
        + cross_entropy(cls_s_u, cls_t_v, center_cls, tau_s, tau_t)
    ) / 2

    # Patch loss only at masked positions, per image  [§3.1] [eq. 3]
    loss_mim = (
        masked_loss(patch_s_u, patch_t_u, mask_u, center_patch)
        + masked_loss(patch_s_v, patch_t_v, mask_v, center_patch)
    ) / 2

    # Sum without scaling  [§3.2]
    return loss_cls + loss_mim

# Cross-entropy with teacher centring and sharpening  [alg. 1]
# Shapes: logits (..., K); center (K,) or (1, K)

def cross_entropy(student_logits, teacher_logits, center, tau_s, tau_t):
    # Teacher: stop gradient, subtract centre, sharpen  [alg. 1]
    teacher_probs = softmax((teacher_logits - center) / tau_t, dim=-1)
    # Student: soften with the student temperature  [alg. 1]
    student_logp = log_softmax(student_logits / tau_s, dim=-1)
    return -(teacher_probs.detach() * student_logp).sum(dim=-1)

# Per-image mean over masked tokens, then batch mean  [code: main_ibot.py#L545]
# Images with no masked token contribute zero but still count  [Pitfalls]

def masked_loss(student_logits, teacher_logits, mask, center_patch):
    # Per-token cross-entropy with patch temperatures  [alg. 1]
    per_token = cross_entropy(
        student_logits, teacher_logits, center_patch, tau_s, tau_t_patch)
    # Keep only masked tokens; divide by their count per image  [alg. 1]
    masked = (per_token * mask).sum(dim=1)
    return (masked / mask.sum(dim=1).clamp(min=1)).mean()

# Centre update: EMA of teacher batch means  [code: main_ibot.py#L564-L572]
# m_c, m_p are centre momenta (0.9 by default)

def update_centers(teacher_cls, teacher_patch):
    # Batch mean of [CLS] outputs  [alg. 1]
    center_cls = m_c * center_cls + (1 - m_c) * mean(teacher_cls, dim=0)
    # Mean over patch tokens, then batch  [alg. 1]
    center_patch = m_p * center_patch + (1 - m_p) * mean(teacher_patch)

Key equations, as written in the paper:

• BEiT-style masked-token loss with a fixed tokenizer eq. 1.
• Self-distillation [CLS] loss, teacher to student eq. 2.
• Masked-patch self-distillation loss, Eq. 3. The paper writes it as a plain sum over masked tokens; see Paper vs code for how the code normalizes it.

Hyperparameters

Values are from the paper unless the "Where" column names the code. The "Notes" column flags differences.

Results

Author-reported. All numbers are on ImageNet-1K validation unless stated.

Classification with frozen features, k-NN and linear probing tab. 1:

Fine-tuning on ImageNet-1K tab. 2:

Fine-tuning with ImageNet-22K pre-training tab. 3: iBOT ViT-B/16 84.4 (BEiT 83.7), iBOT ViT-L/16 86.6 (BEiT 86.0), iBOT ViT-L/16 at 512 px 87.8.

Semi-supervised, ViT-S/16, 1% and 10% labels tab. 4: iBOT 61.9 and 75.1; DINO 60.3 and 74.3.

Unsupervised clustering on ImageNet-1K, ViT-S/16 tab. 5: iBOT ACC 43.4, ARI 32.8, NMI 78.6, FMI 35.6; DINO ACC 41.4, ARI 29.8, NMI 76.8, FMI 32.8. The text reports "32.8% NMI", which does not match the table; see Pitfalls.

Dense tasks tab. 6 (Cascade Mask R-CNN on COCO; ADE20K mIoU):

Transfer by fine-tuning, ViT-S/16 tab. 7: iBOT 99.1 (CIFAR-10), 90.7 (CIFAR-100), 73.7 (iNat18), 78.5 (iNat19), 98.6 (Flowers), 94.0 (Cars); DINO 99.0, 90.5, 72.0, 78.2, 98.5, 93.0.

Ablation, ViT-S/16, 300 epochs, r=0.3, no multi-crop §4.4 tab. 9: full iBOT k-NN 69.1, linear 74.2, fine-tuning 81.5. Without the [CLS] loss: 9.5, 29.8, 79.4. Standalone DINO as tokenizer: 44.3, 60.0, 81.7. DINO with no MIM: 67.9, 72.5, 80.6.

Properties, qualitative (no new numbers beyond those below):

• Patch-token pattern layouts from an 800-epoch ViT-S/16 group patches with shared semantic parts, such as a vehicle headlight or a dog's ear, and also low-level textures §4.3.1 fig. 4. BEiT's offline tokenizer layouts show mostly colour and texture clusters fig. 16.
• Part-wise linear probing averages the top-k patch tokens by attention. Top-56 gives the best result, 5.9% above DINO §4.3.1 fig. 5. Using only the single top patch, iBOT advances over DINO by 17.9% §4.3.1 fig. 5.
• Self-attention maps of the [CLS] query separate objects and parts, such as a bird from a branch or a car wheel §4.3.2 fig. 6.
• Under occlusion (random, salient and non-salient patch dropping), iBOT loses less accuracy than its supervised baseline §4.3.3 app. D fig. 9. Under grid shuffle, it keeps more accuracy than the supervised baseline and ResNet-50 §4.3.3 app. D fig. 10.

Cost, ViT-S/16 tab. 19: with multi-crop, 800 epochs take 193.4 h for iBOT and 180.0 h for DINO, on two 8-GPU V100 machines at batch size 1024. Peak memory is 19.5 GB for iBOT and 15.4 GB for DINO.

Paper vs code

Code: official repository at code_commit (Apache-2.0). The repository is linked from the abstract.

• Patch loss normalization. The paper's Eq. 3 sums the masked-token loss with no normalization eq. 3. The paper's pseudocode divides by the number of masked tokens per image alg. 1. The code divides by mask.sum().clamp(min=1) per image and averages over images main_ibot.py. The code follows the pseudocode.
• Head sharing. The paper says the [CLS] and patch heads are shared and this is better §3.2 app. E. The student default is not shared main_ibot.py, while the teacher default is shared main_ibot.py. The published ViT-S command passes --shared_head true README.md. Use shared heads to reproduce the paper.
• Global and local crop scales. The paper states s=0.32, with global crops in (s, 1) and local crops in (0.05, s) app. B. The published ViT-S command uses global (0.25, 1) and local (0.05, 0.25) README.md. The argparse defaults are global (0.14, 1) and local (0.05, 0.4) main_ibot.py main_ibot.py. Use the README values for the released recipe. Unresolved which one produced the headline numbers.
• ViT-L/16 pre-training epochs. The text says 250 epochs §3.2. The ImageNet-22K table lists 300 epochs for ViT-L/16 on ImageNet-1K tab. 15, and the effective-epoch column is 1200 = 300 × 4 tab. 1. The 1K epoch count is therefore unresolved (250 vs 300). Use 300 only if you have a source for it.
• Prediction ratio. Paper: r=0 with probability 0.5 and U[0.1, 0.5] otherwise §3.2. Code: --pred_ratio 0 0.3 --pred_ratio_var 0 0.2, which picks 0 or 0.3 at random and then samples ±0.2 around 0.3 loader.py README.md. These agree.
• Teacher temperature for [CLS]. The paper keeps [CLS] parameters the same as DINO app. E. The official command uses final teacher temperature 0.07 README.md, while the argparse default is 0.04 main_ibot.py. Use 0.07 in the published ViT-S recipe.
• Augmentations, optimizer details, mixed precision, gradient clipping, weight decay schedule, drop path, momentum. These appear only in the code (see Hyperparameters). The paper text does not state them.
• No other discrepancies found between Algorithm 1 and main_ibot.py for the [CLS] loss, the centring and sharpening, the EMA update, or the centre update.

Pitfalls

• Teacher is not a separate tokenizer. Targets are the teacher's softmax over K=8192 outputs, not one-hot labels. The paper's ablation shows hard labels (hardmax) reduce linear probing from 74.2 to 73.3 with centring, and to 71.9 without tab. 18. Keep the softmax.
• Temperatures are sensitive. The code comment says teacher temperatures above 0.07 are unstable main_ibot.py. Keep the warm-up from 0.04 over 30 epochs main_ibot.py.
• Zero-mask images dilute the patch loss. Random MIM sets r=0 for half of the crops. Those images have no masked token, so their patch loss is zero, but they still count in the batch average main_ibot.py. The effective patch loss is smaller than a per-masked-token average would be. The paper does not discuss this.
• Ratio is per crop, not per image. The sampling happens in the dataset __getitem__, once for each crop loader.py. Two global crops from the same image can have different ratios.
• Local crops are not masked. The student is run on local crops with masked modeling switched off, and only the [CLS] output is used main_ibot.py. Local crops contribute only to the [CLS] loss.
• Mask token starts at zero. The learned mask embedding is initialised to zero and replaces the patch embedding before [CLS] and positional embeddings are added vision_transformer.py vision_transformer.py.
• Center for patches is a token mean. The patch centre is averaged over tokens and then over the batch, not per position main_ibot.py.
• Frozen last layer in epoch 0. Last-layer gradients are cancelled during the first epoch, so the weight-normalized last layer is effectively fixed then utils.py main_ibot.py.
• NMI inconsistency in the paper. The unsupervised paragraph says NMI is 32.8% §4.1, but the table reports NMI 78.6 and ARI 32.8 tab. 5. Use the table.
• Memory and time. The text says 25% more memory than DINO app. E. The table gives 19.5 GB vs 15.4 GB for ViT-S with multi-crop tab. 19, which is about 27%. Plan for the table figure.
• Figures and ablations use a reduced setting. Ablations use 300 epochs, r=0.3 and no multi-crop §4.4 app. E. Do not compare them with the 800-epoch headline numbers.
• ImageNet-22K pre-training helps some metrics and hurts others. ViT-L/16 pre-trained on 22K gets lower k-NN accuracy than on 1K (72.9 vs 78.0) but higher linear probing (82.3 vs 81.0) tab. 1 tab. 15. The paper suggests the data distribution matters more under frozen-feature protocols app. D.
• Dense results use a linear head or UperNet. Segmentation with a linear head is much weaker for BEiT (27.4 mIoU) than for iBOT (38.3) tab. 6. Choose the head to match your comparison.
• Compute. Pre-training ViT-S for 800 epochs with multi-crop took 193.4 h on two 8-GPU V100 machines tab. 19. Plan for the same scale.
• Labels are not used for training. ImageNet labels are passed through the loader and used only to compute NMI in the logs main_ibot.py main_ibot.py.
• Robustness is mixed. iBOT is more robust to background change except for O.BT. §4.3.3 tab. 8.
• Fine-tuning recipe matters. The BEiT-style recipe with layer decay is used for the reported numbers. The traditional 300-epoch SGD recipe gives lower accuracy for the same model app. C.

Lineage

• Builds on: masked image modeling with a visual tokenizer, BEiT, arXiv:2106.08254 §2.1.
• Builds on: self-distillation with a teacher and centring, DINO (Caron et al., ICCV 2021, no arXiv id in the paper's bibliography) §2.2.
• Compared against: BEiT, DINO, MoCov3, SwAV, EsViT, BYOL, SimCLRv2 §4 tab. 1.
• Related work on masked prediction with pixels or contrastive losses §5.