[논문 리뷰] SuperCon: Supervised Contrastive Learning for Imbalanced Skin Lesion Classification

Keyu Chen, Di Zhuang외 1명|arXiv (Cornell University)|2022. 02. 11.

Cutaneous Melanoma Detection and Management인용 수 5

한 줄 요약

이 논문은 먼저 감독 대비 손실을 사용하여 클래스 인식 가능하고 분리 가능한 특징 표현을 학습한 후, 불균형한 피부 병변 데이터셋에서 분류기 성능을 향상시키기 위해 포칼 손실을 사용한 피니어싱을 수행하는 두 단계 학습 전략인 SuperCon을 제안한다. 이 방법은 최신 기술 수준의 성능을 달성하여 기존 방법 대비 평균 F1 점수와 AUC 각각 0.33과 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의 성능는 표현 학습 단계에서의 에포크 수에 얼마나 민감한가?
• 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의 도메인 외 데이터를 추가하면 성능이 악화되며, 이는 데이터셋 간 도메인 이동이 표현 학습에 부정적인 영향을 미칠 수 있음을 시사한다.