MIM-Refiner: A Contrastive Learning Boost from Intermediate Pre-Trained Masked Image Modeling Representations

Authors: Benedikt Alkin, Lukas Miklautz, Sepp Hochreiter, Johannes Brandstetter
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: self-supervised learning, masked image modeling, instance discrimination, computer vision, contrastive learning

Abstract:
We introduce MIM (Masked Image Modeling)-Refiner, a contrastive learning boost for pre-trained MIM models. MIM-Refiner is motivated by the insight that strong representations within MIM models generally reside in intermediate layers. Accordingly, MIM-Refiner leverages multiple instance discrimination (ID) heads that are connected to different intermediate layers. In each head, a nearest neighbor ID objective constructs clusters that capture semantic information which improves performance on downstream tasks, including off-the-shelf and fine-tuning settings.

The refinement process is short and simple -  yet highly effective. Within a few epochs, we refine the features of MIM models from subpar to state-of-the-art, off-the-shelf features. Refining a ViT-H, pre-trained with data2vec 2.0 on ImageNet-1K, sets a new state-of-the-art in linear probing (84.7\%) and low-shot classification among models that are pre-trained on ImageNet-1K. MIM-Refiner efficiently combines the advantages of MIM and ID objectives, enabling scaling ID objectives to billion parameter models using relatively little compute. MIM-Refiner compares favorably against previous state-of-the-art SSL models on various benchmarks such as low-shot classification, long-tailed classification and semantic segmentation.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730523774637

thread_type:
official_review

<thread_object>
creator: reviewer_jdmo

<content>
summary:
This paper introduces MIM-Refiner, which leverages contrastive learning to boost MIM models. The proposed method is simple and has demonstrated effectiveness in few-shot image classification tasks.

strengths:
S1: This paper is well-written and easy to follow.

S2: This paper is not the first to point out that the encoder in MIM methods partially performs image encoding and representation learning. A similar conclusion is also discussed in [A], highlighting that MIM methods using a single ViT structure tend to face this issue. The reviewer previously conducted experiments on MAE-B, showing that introducing an additional decoder can effectively alleviate this problem. This paper demonstrates that, for methods like MAE that use an asymmetric encoder-decoder architecture， especially in larger models， a small decoder cannot fully decouple encoding and decoding, providing academic insights.

S3: This paper proposes a simple and effective MIM-Refiner method, refining the later blocks of MIM models to enhance MIM representations effectively.

[A] Context Autoencoder for Self-Supervised Representation Learning.

weaknesses:
W1: Existing work [A] has shown that fine-tuning MIM models can enhance their representation capability (for image classification), but the improvement under full fine-tuning is minimal. Additionally, MAE has demonstrated significant transfer performance on dense prediction tasks [B] (object detection/instance segmentation). Fine-tuning MIM models with contrastive learning methods is unlikely to bring substantial improvement and may even negatively impact performance.

W2: Current vision foundation models, such as DINOv2, exhibit strong patch-level representation learning capabilities and combine MIM and CL. Their learned representations have shown effectiveness in tasks like image classification, pixel classification, and depth estimation. Although this paper discusses the relationship between MIM-Refiner and these models, suggesting that MIM-Refiner can build on them for further improvement, I am concerned that MIM-Refiner may degrade pixel-level representation performance for tasks like semantic segmentation or depth estimation (especially when the backbone is fixed).

[A] Layer Grafted Pre-training: Bridging Contrastive Learning And Masked Image Modeling For Label Efficient Representations.

[B] Exploring Plain Vision Transformer Backbones for Object Detection.

questions:
see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730601084258

thread_type:
official_review

<thread_object>
creator: reviewer_haab

<content>
summary:
---

## **Summary**

The paper identifies the representation degradation issue in Masked Image Modeling (MIM)-pretrained large foundation models. To address this, the authors propose a simple yet effective method to prevent degradation and further improve the representation quality of MIM methods by adding auxiliary contrastive losses to the last layers of Vision Transformers (ViTs) on top of the MIM objective. The paper provides improved performances with large margins over current state-of-the-art (SOTA) methods through extensive experiments and rigorous analysis, demonstrating the success of the proposed approach.

---

strengths:
---

## **Strengths**

- The paper is well-written, with clear observations, a well-developed motivation, a straightforward idea, a clearly-stated method, extensive experiments, and comprehensive analysis.
  
- It effectively identifies the representational degradation phenomenon in large visual foundation models pre-trained with MIM self-supervised learning (SSL), providing evidence through multiple experiments.
  
- The proposed method offers a simple and effective solution to prevent this issue and improve the representation quality of MIM SSL.

- Rigorous experiments and analysis are conducted to show the success of the proposed method, with large improvements over current SOTA.

---

weaknesses:
---
### **Limitations**

1. To prevent representation quality degradation in the last layers of ViTs, the authors experiment with contrastive loss, which requires constructing a queue/pool for positive and negative samples. I noticed the proposed method uses a top 20-NN approach to retrieve positive samples in the queue, which could contribute significantly to the increased training time per step. What's the queue size used? how much does it contribute to the increased training time per step?

2. Since the paper emphasizes preserving the richness of representations, evaluation on dense prediction tasks such as object detection and instance segmentation (OD/IS) would be valuable, in addition to the provided segmentation probing on ADE20K.

   - It would be meaningful to compare the performance of MIM-refiner-pretrained ViT-L on COCO object detection against MAE-pretrained ViT-L following the ViTDet framework [1].

   - [1] Li, Y., Mao, H., Girshick, R., & He, K. (2022, October). *Exploring plain vision transformer backbones for object detection.* In European Conference on Computer Vision (pp. 280-296). Cham: Springer Nature Switzerland.

---

### **Recommendation**

Considering the strengths and weaknesses discussed above, my recommendation for this paper is **ACCEPT**. This is a strong paper with a clear contribution.

questions:
---
Since D2V2 is used as a baseline, does the representation degradation issue also appear in the audio and language domains?

---
</content>
</thread_object>

</thread>
<thread>
creation date:
1730605944612

thread_type:
official_review

<thread_object>
creator: reviewer_bbvs

<content>
summary:
This paper presents a contrastive learning boosting method called MIM-Refiner to refine the features of pre-trained MIM models. MIM-Refiner leverages multiple instance discrimination heads (ID) which are connected to different immediate layers. Each ID head contains a contrastive loss that captures semantic information to improve the quality of learned representations. By training a few epochs, the features of MIM-Refiner surpass the current MIM models on multiple experiments: on ViT-H with data2vec 2.0 on ImageNet-1K, the accuracy of the proposed method reaches state-of-the-art performance on linear probing and low-shot classification.

strengths:
1. This paper proposes a detailed analysis of the blocks of MIM models in which different blocks extract features with a specific focus and the most efficient features learned by MIM are from the middle blocks.

2. A contrastive learning-based method called MIM-Refiner is proposed to refine the representation of current MIM models by attaching the middle layers with ID objective. 

3. Experimental results show the effectiveness and generalization ability of MIM-Refiner on downstream tasks.

weaknesses:
1. As the discussion of end-to-end training, the proposed method MIM-Refiner seems to be a two-step training method, with first step training MIM models and fine-tuning the updated models by incorporating ID heads to middle layers. Practically, this might increase the complexity of the training paradigm and deployment. Is it possible to improve the proposed method with end-to-end training on MIM and ID? If not, what are the potential bottlenecks to circumvent this goal?

2. There is no overview diagram that shows the detailed architecture of MIM-Refiner or how the training diagram goes. The diagram in Figure 4 provides partial information but does not clearly illustrate these points.

questions:
Please refer to weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730727427706

thread_type:
official_review

<thread_object>
creator: reviewer_auz1

<content>
summary:
The paper focuses on bridging the gap between large MIM pre-trained models and SOTA methods. The paper first discovers that  MIM models have different types of blocks: those that mainly improve the pre-training objective and others that are responsible for abstraction. Then, the paper proposes a method MIM-Refiner, which adds Instance Discriminator heads on pre-trained MIM models for refinement. The ID heads exploit the intermediate representations to consistently improve the performance of MIM pretrained models. While the performance gains on large dataset full-finetuning are small, the proposed methods show remarkable gains on few-shot settings.

strengths:
1. The paper first points out the influence of the lightweight decoder on the feature learning of the encoder in MIM methods.
2. The analyzing part is well-written.

weaknesses:
1. The description of the method and experimental setup needs to be clarified. (a) Which blocks need to be fine-tuned during refinement, or do all blocks need to be fine-tuned? (b) How many epochs are needed to refine different models? (c) What is the structure of the ID head? Answers to all these questions should be contained in the manuscript.
2. Unfair comparison. The paper misses an important baseline - train the original model with 0 heads with the same epochs to demonstrate the importance of refinement (instead of just training more epochs).
3. Some typos. L267-269, see Table 1 instead of Figure 3b.

questions:
Please refer to the weakness. I believe a clear description of the method and experimental setup is one of the most important things when writing a paper (weakness 1).

Additional question: what does the “relative” in Figure 3(d) mean? Does the value calculated by the performance of ( the i+1 th layer - the i-th layer)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731631385755

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all reviewers for their positive feedback, constructive comments and suggestions.
We are pleased to see that the reviewers highlighted the clarity of our paper and appreciated our analysis as well as our extensive experiments. Several reviewers recognized the clear motivation of our approach from the thorough analysis.

We updated the paper to incorporate the feedback of the reviewers. To summarize, we made the following changes:

- Added more details to the experimental setup in the main paper (first paragraph in Section 4) such as training duration and included a reference to the extensive implementation details and hyperparameters in the appendix (reviewer Auz1).
- Included methodology to calculate the relative improvement in Figure 3d also in the caption instead of only in the appendix (reviewer Auz1).
- Added results for COCO object detection and instance segmentation in Appendix B.10 (reviewer HaAb and JDMo)
- Added preliminary analysis of feature degradation of masked pre-training in other modalities in Appendix E (reviewer HaAb).
- corrected various typos

Additionally, we respond to each review individually, addressing the raised questions and concerns.
</content>
</thread_object>

</thread>


Metareview:
metareview:
After discussion, this submission received 4 positive scores . The major concerns about the experimental details,  preliminary analysis  and method clarity were comprehensively solved. After reading the paper, the review comments and the rebuttal, the AC thinks the remaining issue is to include all the revised content to the camera-ready version and correct typographical errors.