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

📄 arXiv: 2402.10093v4 📥 PDF

作者: Benedikt Alkin, Lukas Miklautz, Sepp Hochreiter, Johannes Brandstetter

分类: cs.CV, cs.AI, cs.LG

发布日期: 2024-02-15 (更新: 2025-02-20)

备注: Published as a conference paper at ICLR 2025. Github: https://github.com/ml-jku/MIM-Refiner

💡 一句话要点

提出MIM-Refiner以提升预训练MIM模型的对比学习效果

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 对比学习 特征精炼 中间层表示 图像分类 语义分割 聚类 预训练模型

📋 核心要点

  1. 现有的MIM模型在下游任务中的表现受限于特征提取的有效性，尤其是在中间层的表示能力不足。
  2. MIM-Refiner通过引入多个对比头，针对不同中间层进行特征精炼，利用语义聚类提升模型性能。
  3. 实验结果显示，MIM-Refiner在多个任务上超越了现有的最先进模型，尤其是在低样本分类和线性探测中表现优异。

📝 摘要（中文）

我们介绍了MIM-Refiner，这是一种针对预训练的MIM模型的对比学习增强方法。MIM-Refiner的动机在于，MIM模型中的强表示通常存在于中间层。因此，MIM-Refiner利用多个与不同中间层连接的对比头。在每个头中，修改后的最近邻目标构建了捕捉语义信息的语义聚类，从而提升下游任务的性能，包括现成和微调设置。精炼过程简短而简单，但效果显著。在几个训练周期内，我们将MIM模型的特征从较差提升至最先进的现成特征。对使用data2vec 2.0在ImageNet-1K上预训练的ViT-H进行精炼，在线性探测中创下84.7%的新记录，并在低样本分类中表现优异。MIM-Refiner有效结合了MIM和ID目标的优势，并在低样本分类、长尾分类、聚类和语义分割等多个基准上优于之前的最先进SSL模型。

🔬 方法详解

问题定义：本论文旨在解决现有MIM模型在下游任务中表现不佳的问题，尤其是中间层特征提取能力不足导致的性能瓶颈。

核心思路：MIM-Refiner的核心思路是利用多个对比学习头，分别连接到不同的中间层，以此提取更强的语义特征，从而提升模型在下游任务中的表现。

技术框架：整体架构包括多个对比头，每个头连接到MIM模型的不同中间层。通过修改的最近邻目标，构建语义聚类，进而优化特征表示。

关键创新：MIM-Refiner的主要创新在于其对中间层特征的利用，通过对比学习增强了MIM模型的表示能力，与传统方法相比，显著提升了下游任务的性能。

关键设计：在设计中，采用了修改后的最近邻损失函数，确保语义聚类的有效性；同时，选择了适合的网络结构和参数设置，以保证精炼过程的高效性和准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，MIM-Refiner对使用data2vec 2.0在ImageNet-1K上预训练的ViT-H模型进行了精炼，在线性探测任务中达到了84.7%的新记录。此外，在低样本分类等多个基准上，MIM-Refiner的表现均优于现有的最先进SSL模型，显示出显著的性能提升。

🎯 应用场景

MIM-Refiner的研究成果可广泛应用于计算机视觉领域，尤其是在图像分类、语义分割和聚类等任务中。通过提升预训练模型的特征提取能力，该方法能够为实际应用提供更高的准确性和效率，具有重要的商业价值和研究意义。

📄 摘要（原文）

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 contrastive heads that are connected to different intermediate layers. In each head, a modified nearest neighbor objective constructs semantic 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 and compares favorably against previous state-of-the-art SSL models on a variety of benchmarks such as low-shot classification, long-tailed classification, clustering and semantic segmentation.