[论文解读] SuperCon: Supervised Contrastive Learning for Imbalanced Skin Lesion Classification

Keyu Chen, Di Zhuang+1 more|arXiv (Cornell University)|Feb 11, 2022

Cutaneous Melanoma Detection and Management被引用 5

一句话总结

本文提出 SuperCon，一种两阶段训练策略：首先使用监督对比损失学习类别感知、可分的特征表示，随后使用焦点损失微调以提升在类别不平衡的皮肤病变数据集上的分类器性能。该方法实现了最先进性能，在现有方法基础上平均提升 F1 分数 0.33 和 AUC 0.032。

ABSTRACT

Convolutional neural networks (CNNs) have achieved great success in skin lesion classification. A balanced dataset is required to train a good model. However, due to the appearance of different skin lesions in practice, severe or even deadliest skin lesion types (e.g., melanoma) naturally have quite small amount represented in a dataset. In that, classification performance degradation occurs widely, it is significantly important to have CNNs that work well on class imbalanced skin lesion image dataset. In this paper, we propose SuperCon, a two-stage training strategy to overcome the class imbalance problem on skin lesion classification. It contains two stages: (i) representation training that tries to learn a feature representation that closely aligned among intra-classes and distantly apart from inter-classes, and (ii) classifier fine-tuning that aims to learn a classifier that correctly predict the label based on the learnt representations. In the experimental evaluation, extensive comparisons have been made among our approach and other existing approaches on skin lesion benchmark datasets. The results show that our two-stage training strategy effectively addresses the class imbalance classification problem, and significantly improves existing works in terms of F1-score and AUC score, resulting in state-of-the-art performance.

研究动机与目标

• 解决皮肤病变分类中严重的类别不平衡问题，其中恶性（如黑色素瘤）病变稀少但需准确检测。
• 克服现有方法对数据分布敏感且在不同数据集间泛化能力不足的局限。
• 提升深度学习模型学习鲁棒、判别性特征表示的能力，且不受类别样本不平衡影响。
• 在不依赖数据增强或复杂集成模型的前提下，提升对少数类（如黑色素瘤）的性能。
• 开发一种训练策略，即使在极端类别不平衡下，也能实现类内特征分离、类间特征分离。

提出的方法

• 在第一阶段使用监督对比损失，学习在每个类别内紧密聚集、类别间良好分离的特征表示。
• 使用对比目标训练表示头，最大化同类别样本间的相似性，最小化不同类别样本间的相似性。
• 第一阶段完成后固定主干网络与表示头，仅使用焦点损失对分类器头进行微调，以降低多数类样本的影响。
• 应用焦点损失重新加权交叉熵损失，使训练更关注难以分类的样本，同时最小化易分类的多数类样本的影响。
• 按顺序训练两个阶段：先进行表示学习，再使用焦点损失对分类器进行微调。
• 使用 t-SNE 可视化分析学习到的特征表示，并验证类间可分性的提升。

实验结果

研究问题

• RQ1 是否两阶段训练策略（将表示学习与分类器训练解耦）能提升在类别不平衡皮肤病变分类中的性能？
• RQ2 监督对比学习是否能在少数类别样本不足的情况下，仍有效学习到判别性、类别感知的表示？
• RQ3 将监督对比损失与焦点损失结合，相较于现有数据级、损失级或模型级方法，在皮肤病变基准上能多大程度上实现性能超越？
• RQ4 SuperCon 的性能对第一阶段训练轮数（epoch 数）的敏感性如何？
• RQ5 该模型能否在数据分布不同的数据集（如 ISIC 2019 与 ISIC 2020）之间实现良好泛化？

主要发现

• 在类别不平衡的皮肤病变数据集上，SuperCon 相较于现有最先进方法，平均提升 F1 分数 0.33 和 AUC 0.032。
• 在 ISIC 2020 数据集上，SuperCon 相较于焦点损失，F1 分数提升 0.34，AUC 提升 0.30，表明其在高度不平衡数据上表现更优。
• 混淆矩阵显示，SuperCon 同时实现了高真正例率与高真负例率，而 ROS 等方法则以牺牲高真负例率为代价换取高假正例率。
• t-SNE 可视化证实，SuperCon 学习到了非重叠、良好分离的不同皮肤病变类别特征表示，而 Vanilla、Focal-Loss 或 ROS 则未实现此效果。
• SuperCon 的性能对第一阶段训练轮数敏感，表明超参数调优对实现最优表示学习至关重要。
• 在表示学习阶段引入 ISIC 2019 的分布外数据会降低性能，表明数据集间领域偏移会负面影响表示学习。