Towards elastic vision transformers via contrastive masked distillation

Vision Transformers (ViTs), which embed images into sequences by first splitting them into patches, have yielded brilliant results on various vision tasks. In addition to the demand for enormous amounts of training data, ViTs suffer from unsatisfactory performance with large patches. However, large patches typically bring about lower computational cost and faster inference. To efficiently improve the performance of ViTs with large patches, we propose CoMimViD, a two-stage, patch-size-agnostic distillation elaborately designed for ViTs, which bridges masked image modeling and contrastive learning in a straightforward manner. In the first pretraining stage, the reconstructed image, which is overlaid with visible patches in the masked image modeling, is fed back to the large-patch ViT for contrastive distillation with a teacher ViT at various patch sizes. The resulting ViT is fine-tuned at a fixed patch size in the second fine-tuning stage, which surprisingly can perform well across different patch sizes with a single set of weights. Based on our experiments on ImageNet-1 K, distilling from ViT-S/16, the proposed method yields a +7.6% absolute accuracy gain over the vanilla ViT-S/32 baseline, achieving a result that is better than the counterpart pre-trained on the significantly larger ImageNet-21 K.