[论文解读] MIM-Refiner: A Contrastive Learning Boost from Intermediate Pre-Trained Representations

Benedikt Alkin, Lukas Miklautz+2 more|arXiv (Cornell University)|Feb 15, 2024

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

MIM-Refiner 通过在高质量编码器块上附加多个对比学习头，对预训练的掩码图像建模（MIM）模型的中间表征进行精炼，采用一种新颖的最近邻对齐（NNA）目标。这种简短而高效的精炼方法显著提升了模型在 ImageNet-1K 上的性能，线性探测准确率达到 84.7%，1-shot 分类准确率达到 64.2%，超越了在远超其数据量的训练数据上训练的更大模型。

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

研究动机与目标

• 为解决 MIM 模型尽管预训练效果强大，但在下游任务中表现不佳的问题，尤其是在低数据场景下。
• 识别出 MIM 模型的表征质量在中间编码器块中达到峰值，而非最终层，这是由于解码器的局限性所致。
• 开发一种轻量级、高效的精炼方法，将 MIM 特征转化为语义聚类，而无需从头开始训练。
• 提升 MIM 特征在下游任务（如线性探测、聚类和少样本分类）中的泛化能力和可迁移性。
• 通过精炼一个使用 data2vec 2.0 在 ImageNet-1K 上预训练的 ViT-H 模型，使其在极端低样本场景下超越在数量级更多数据上训练的更大模型。

提出的方法

• MIM-Refiner 在 MIM 编码器的中间块上引入多个对比学习头，尤其聚焦于表征质量最高的层。
• 每个对比学习头应用一种改进的最近邻对齐（NNA）目标，将每个样本与其最近邻对齐，同时排斥批次中其他所有样本。
• NNA 目标通过选择最近邻作为正样本对，而非单纯拉远负样本，从而提升信号质量。
• 精炼过程是顺序且简短的，仅需在与预训练相同的 ImageNet-1K 数据上进行几个训练周期。
• 批量归一化层在对比学习头中至关重要，可稳定训练过程并防止性能下降。
• 该方法保留原始 MIM 模型权重，仅微调对比学习头和归一化层。

Figure 1 : Linear probing state-of-the-art on ImageNet-1K over the last four years.

实验结果

研究问题

• RQ1 为何 MIM 模型在预训练表现强大后，仍会在下游任务中表现不佳，尤其是在低样本设置下？
• RQ2 MIM 编码器中最具语义意义的表征位于何处？如何有效利用这些表征？
• RQ3 能否通过一种轻量级、短周期的精炼过程，显著提升 MIM 模型特征质量，而无需重新训练整个模型？
• RQ4 在中间层使用多个对比学习头是否优于在最终层使用单个头的对比学习？
• RQ5 在极端低样本场景下，经精炼的 MIM 模型能否超越在远超其数据量的训练数据上训练的更大模型？

主要发现

• MIM-Refiner 使用经过 data2vec 2.0 预训练的 ViT-H 模型，在 ImageNet-1K 上实现了 84.7% 的新 SOTA 线性探测准确率。
• 该方法在 1-shot ImageNet-1K 分类任务中达到 64.2% 的新 SOTA 准确率，超越了在最多多出 2000 倍数据上训练的更大模型，如 DINOv2-g、OpenCLIP-G 和 MAWS-6.5B。
• 精炼后的模型在 k-NN 分类中达到 82.3% 的准确率，在 k-means 聚类中达到 67.3% 的准确率，表明潜在空间中具有强语义聚类能力。
• 精炼过程极为高效，仅需几个训练周期即可将特征质量从次优水平提升至 SOTA 水准。
• 在使用 100% 标签进行微调时，精炼后的模型略优于原始模型，证实其在高数据场景下无性能退化。
• 在 ADE20K 和 VTAB 等其他数据集上的迁移实验表明，该方法具有强大泛化能力，MAE-Refined 在多个基准上均取得高 mIoU 和平均排名。