[论文解读] Intermediate Layers Matter in Momentum Contrastive Self Supervised Learning

Aakash Kaku, Sahana Upadhya+1 more|arXiv (Cornell University)|Oct 27, 2021

AI in cancer detection参考文献 28被引用 9

一句话总结

该论文通过使用均方误差（MSE）或Barlow Twins风格的不变性损失，在中间网络层最小化特征差异，以增强动量对比（MoCo）自监督学习。该方法通过在网络早期学习更具可重用性和信息量的特征，提升了下游性能，尤其在低标签数据设置下表现更优，在三个多样化的医学影像数据集上，相较于标准MoCo，平均性能提升高达5%（在1%标签数据下）。

ABSTRACT

We show that bringing intermediate layers' representations of two augmented versions of an image closer together in self-supervised learning helps to improve the momentum contrastive (MoCo) method. To this end, in addition to the contrastive loss, we minimize the mean squared error between the intermediate layer representations or make their cross-correlation matrix closer to an identity matrix. Both loss objectives either outperform standard MoCo, or achieve similar performances on three diverse medical imaging datasets: NIH-Chest Xrays, Breast Cancer Histopathology, and Diabetic Retinopathy. The gains of the improved MoCo are especially large in a low-labeled data regime (e.g. 1% labeled data) with an average gain of 5% across three datasets. We analyze the models trained using our novel approach via feature similarity analysis and layer-wise probing. Our analysis reveals that models trained via our approach have higher feature reuse compared to a standard MoCo and learn informative features earlier in the network. Finally, by comparing the output probability distribution of models fine-tuned on small versus large labeled data, we conclude that our proposed method of pre-training leads to lower Kolmogorov-Smirnov distance, as compared to a standard MoCo. This provides additional evidence that our proposed method learns more informative features in the pre-training phase which could be leveraged in a low-labeled data regime.

研究动机与目标

• 探究改进中间层表征对齐是否能提升医学影像中自监督表示学习的性能。
• 通过提升预训练质量，应对医学AI中标签数据有限的挑战，以优化下游微调性能。
• 评估更好的特征质量是否能带来更优的泛化能力，尤其是在低数据量设置下。
• 通过分析特征重用和逐层信息量，量化自监督特征质量，超越线性探测的评估范围。
• 比较微调后模型输出概率分布，评估其与最优性能的接近程度。

提出的方法

• 引入两种额外损失项：在数据增强后的图像视图之间，对中间层特征计算均方误差（MSE）；以及采用Barlow Twins风格的损失，最小化特征间的互相关性。
• 将MoCo框架扩展至在多个中间层应用这些损失，而不仅限于最终层。
• 沿用标准MoCo中的动量更新队列作为负样本键，但增加了中间层的对比学习目标。
• 采用特征重用分析，测量微调前后特征的相似性，以反映特征质量。
• 通过逐层探测评估网络不同深度处特征的信息量。
• 比较在1%和6%标签数据与100%标签数据上微调的模型之间，输出概率分布的柯尔莫哥洛夫-斯米尔诺夫（KS）距离，以评估泛化能力。

实验结果

研究问题

• RQ1 在MoCo中对齐中间层表征是否能提升医学影像数据集上的下游分类性能？
• RQ2 与标准MoCo相比，该方法在不同网络层上对特征可重用性和信息量的影响如何？
• RQ3 表征质量的提升在低标签数据设置下能带来多大程度的性能增益？
• RQ4 与在小样本标签数据上微调的最佳模型相比，采用该方法预训练的模型输出概率分布有多接近最优性能？
• RQ5 中间层对齐是否能在多种医学影像模态中生成更鲁棒、更具泛化能力的特征？

主要发现

• 在三种医学影像数据集（NIH胸部X光、糖尿病视网膜病变、乳腺癌组织病理学）上，采用MSE损失的该方法均优于标准MoCo，尤其在低标签比例下表现更优。
• 在1%标签数据下，该方法在所有数据集上的平均性能较标准MoCo提升5%，表明在低数据量设置下有显著改进。
• 采用中间层对齐训练的模型表现出更高的特征重用率，表明其学习到的特征质量更高、更具迁移性。
• 逐层探测结果表明，与标准MoCo相比，该方法在更早的网络层中学习到的特征更具信息量。
• 该方法在1%和6%标签数据上微调的模型与100%标签数据下最优模型之间的柯尔莫哥洛夫-斯米尔诺夫（KS）距离更低，表明其输出分布更接近最优性能。
• 该方法在多种医学影像模态中均保持一致的性能增益，表明其具有良好的泛化能力，不局限于特定数据类型。