MaskDistill-ViT-Base

MaskDistill ViT-Base/16

The first open-source PyTorch implementation of MaskDistill with pre-trained weights.

This model was trained using the MaskDistill-PyTorch codebase, reproducing the method from "A Unified View of Masked Image Modeling".

Model Description

MaskDistill learns visual representations by distilling knowledge from a frozen CLIP ViT-B/16 teacher into a ViT-Base student through masked image modeling. The student learns to predict the teacher's features for masked patches using Smooth L1 loss.

• Architecture: ViT-Base/16 (86M params)
• Teacher: CLIP ViT-B/16 (frozen)
• Pretraining: 300 epochs on ImageNet-1K
• Masking: Block masking at 40%, dense encoding with shared relative position bias

Results

Available Checkpoints

Usage

import torch
from src.models.vision_transformer import VisionTransformerMIM

# Load pretrained model
model = VisionTransformerMIM(
    img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12,
    use_shared_rel_pos_bias=True, use_mask_tokens=True,
)
ckpt = torch.load("pretrain_vit_base_ep290.pth", map_location="cpu")
state = {k.replace("module.student.", ""): v for k, v in ckpt["model"].items() if "student" in k}
model.load_state_dict(state, strict=False)

See the GitHub repo for full training and evaluation code.

Citation

@article{hou2022unified,
  title={A Unified View of Masked Image Modeling},
  author={Hou, Zhenda and Sun, Fei and Chen, Yun-Hao and Yuan, Jia-Hong and Yu, Jia-Mu},
  journal={arXiv preprint arXiv:2210.10615},
  year={2022}
}