Keypoint Aware Masked Image Modelling

TL;DR

KAMIM enhances masked image modeling by incorporating keypoint-based patch weighting, significantly improving linear probing accuracy and providing more effective local context understanding in vision transformers.

Contribution

The paper introduces KAMIM, a novel patch-wise weighting method based on keypoint features that improves masked image modeling performance, especially for linear probing tasks.

Findings

Linear probing accuracy increased from 16.12% to 33.97%.

Finetuning accuracy improved slightly from 76.78% to 77.3%.

Patch-wise weighting benefits larger pretraining datasets.

Abstract

SimMIM is a widely used method for pretraining vision transformers using masked image modeling. However, despite its success in fine-tuning performance, it has been shown to perform sub-optimally when used for linear probing. We propose an efficient patch-wise weighting derived from keypoint features which captures the local information and provides better context during SimMIM's reconstruction phase. Our method, KAMIM, improves the top-1 linear probing accuracy from 16.12% to 33.97%, and finetuning accuracy from 76.78% to 77.3% when tested on the ImageNet-1K dataset with a ViT-B when trained for the same number of epochs. We conduct extensive testing on different datasets, keypoint extractors, and model architectures and observe that patch-wise weighting augments linear probing performance for larger pretraining datasets. We also analyze the learned representations of a ViT-B trained using KAMIM and observe that they behave similar to contrastive learning with regard to its behavior, with longer attention distances and homogenous self-attention across layers. Our code is publicly available at https://github.com/madhava20217/KAMIM.