---
license: apache-2.0
tags:
  - vision
  - self-supervised-learning
  - masked-image-modeling
  - knowledge-distillation
  - vit
datasets:
  - ILSVRC/imagenet-1k
  - 1aurent/ADE20K
  - detection-datasets/coco
metrics:
  - accuracy
  - mIoU
  - mAP
pipeline_tag: image-classification
---

# MaskDistill ViT-Base/16

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

This model was trained using the [MaskDistill-PyTorch](https://github.com/drkostas/MaskDistill-PyTorch) codebase, reproducing the method from ["A Unified View of Masked Image Modeling"](https://arxiv.org/abs/2210.10615).

## 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

| Evaluation | Result |
|-----------|--------|
| Finetuning (ImageNet-1K) | **84.8%** top-1 |
| k-NN (k=10) | **75.6%** top-1 |
| Linear Probe | **76.3%** top-1 |
| Sem. Seg. (ADE20K, UPerNet) | **52.6** mIoU |
| Obj. Det. (COCO, Mask R-CNN) | **44.4** bbox mAP |
| Inst. Seg. (COCO, Mask R-CNN) | **40.1** segm mAP |

## Available Checkpoints

| File | Description |
|------|------------|
| `pretrain_vit_base_ep290.pth` | Pretrained ViT-Base (300 epochs) |
| `finetune_vit_base_ep100.pth` | Finetuned on ImageNet-1K (84.8% top-1) |
| `linprobe_vit_base_ep90.pth.tar` | Linear probe (90 epochs, 76.3% top-1) |
| `semseg_upernet_ade20k_160k.pth` | UPerNet on ADE20K (52.6 mIoU) |
| `detection_maskrcnn_coco_12ep.pth` | Mask R-CNN on COCO (44.4 mAP) |

## Usage

```python
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](https://github.com/drkostas/MaskDistill-PyTorch) for full training and evaluation code.

## Citation

```bibtex
@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}
}
```