[논문 리뷰] Federated Contrastive Learning for Dermatological Disease Diagnosis via On-device Learning

Yawen Wu, Dewen Zeng외 6명|arXiv (Cornell University)|2022. 02. 14.

Cutaneous Melanoma Detection and Management인용 수 7

한 줄 요약

이 논문은 제한된 레이블 데이터를 사용하여 피부질환 진단을 위한 디바이스 내 연합 대비 학습 프레임워크를 제안한다. 자기지도 학습 전훈 중에 장치 간 특징 공유를 가능하게 하고 국소 음성 예제를 제거함으로써 데이터 다양성과 표현 품질을 향상시켜, 기존 연합 및 대비 학습 기준 대비 피부질환 데이터셋에서 최신 기술 수준의 재현율과 정밀도 향상을 달성한다.

ABSTRACT

Deep learning models have been deployed in an increasing number of edge and mobile devices to provide healthcare. These models rely on training with a tremendous amount of labeled data to achieve high accuracy. However, for medical applications such as dermatological disease diagnosis, the private data collected by mobile dermatology assistants exist on distributed mobile devices of patients, and each device only has a limited amount of data. Directly learning from limited data greatly deteriorates the performance of learned models. Federated learning (FL) can train models by using data distributed on devices while keeping the data local for privacy. Existing works on FL assume all the data have ground-truth labels. However, medical data often comes without any accompanying labels since labeling requires expertise and results in prohibitively high labor costs. The recently developed self-supervised learning approach, contrastive learning (CL), can leverage the unlabeled data to pre-train a model, after which the model is fine-tuned on limited labeled data for dermatological disease diagnosis. However, simply combining CL with FL as federated contrastive learning (FCL) will result in ineffective learning since CL requires diverse data for learning but each device only has limited data. In this work, we propose an on-device FCL framework for dermatological disease diagnosis with limited labels. Features are shared in the FCL pre-training process to provide diverse and accurate contrastive information. After that, the pre-trained model is fine-tuned with local labeled data independently on each device or collaboratively with supervised federated learning on all devices. Experiments on dermatological disease datasets show that the proposed framework effectively improves the recall and precision of dermatological disease diagnosis compared with state-of-the-art methods.

연구 동기 및 목표

• 레이블이 부족하고 모바일 장치에 분산되어 있는 상황에서 정확한 피부질환 진단 모델을 훈련시키는 데 도전하는 것.
• 각 장치에서 데이터 다양성이 낮아 성능이 떨어지는 연합 환경에서 대비 학습의 한계를 극복하는 것.
• 원시 데이터 공유를 피하면서도 장치 간 특징 교환을 통해 표현 학습을 향상시켜 환자 개인정보를 보호하는 것.
• 사전 훈련된 모델을 사용해 최소한의 레이블 데이터로 효과적인 피지컬 테이닝을 가능하게 하여 레이블링 효율을 향상시키는 것.
• 연합 자기지도 전훈과 협업 또는 국지적 피지컬 테이닝을 결합한 이중 단계 프레임워크를 개발하는 것.

제안 방법

• 이중 단계 프레임워크를 도입: 레이블이 없는 데이터를 사용한 연합 자기지도 전훈, 이어서 제한된 레이블 데이터로 피지컬 테이닝.
• 전훈 중에 장치 간 특징 교환(FE)을 제안하여 데이터 다양성을 증가시키고 대비 학습 품질을 향상시키며, 원시 이미지를 공유하지 않음.
• 대비 학습 중 국소 음성 예제를 제거하여 학습된 표현의 클러스터링을 향상시키고 노이즈를 감소시킴.
• 모델 일관성을 유지하기 위해 전훈 및 피지컬 테이닝 단계 모두에 FedAvg를 집계 전략으로 사용.
• 대비 손실를 위한 양성 및 음성 쌍을 생성하기 위해 데이터 증강(예: 무작위 자르기, 색상 왜곡)을 적용한 시아미즈 네트워크 아키텍처를 활용.
• 데이터 가용성 및 개인정보 보호 제약 조건에 따라 국지적 또는 연합적 피지컬 테이닝을 모두 지원하여 배포의 유연성 확보.

실험 결과

연구 질문

• RQ1 각 장치에서 레이블 없이 제한된 데이터만 있는 상황에서 연합 대비 학습이 피부질환 진단을 위한 강력한 표현을 효과적으로 학습할 수 있는가?
• RQ2 클라이언트 당 데이터 다양성이 낮은 연합 환경에서 장치 간 특징 공유가 대비 학습 성능을 어떻게 향상시키는가?
• RQ3 국소 음성 예제를 제거함으로써 학습된 표현의 품질과 후속 진단 정확도에 어떤 영향을 미치는가?
• RQ4 제안된 프레임워크는 진단 성능을 유지하거나 향상시키면서 레이블 요구량을 어느 정도 줄일 수 있는가?
• RQ5 피부질환 데이터셋에서 재현율과 정밀도 측면에서 제안된 방법이 최신 기술 수준의 연합 및 대비 학습 기준 대비 어떻게 비교되는가?

주요 결과

• 10%의 레이블 데이터로도 제안된 방법은 연합 피지컬 테이닝에서 48.03%의 재현율과 42.87%의 정밀도를 달성하여 최고의 베이스라인인 FedSimCLR보다 재현율 2.14%p, 정밀도 1.43%p 향상.
• 40% 레이블 비율에서 제안된 방법은 재현율 55.13%, 정밀도 48.73%에 도달하여, 이중 레이블 데이터로 훈련된 베이스라인의 성능을 맞추거나 초월.
• 제거 실험에서 특징 교환은 재현율을 44.72%에서 47.45%로 향상시키고, 국소 음성 예제 제거로 추가로 48.03%까지 상승시켜 10% 레이블에서 확인.
• 모든 레이블 비율(10%, 20%, 40%, 80%)에서 국지적 및 연합 피지컬 테이닝 환경 모두에서 모든 베이스라인 대비 재현율과 정밀도가 지속적으로 향상됨.
• 다양한 피부 색소 및 질환 표현을 포함한 피부질환 데이터셋에서 평가한 결과, 다양한 피부 유형에서도 프레임워크가 뛰어난 강건성을 보임.
• 결과는 특징 공유가 낮은 데이터, 개인정보 보호 제약 조건이 있는 환경에서 표현 학습을 크게 향상시키며, 실생활 모바일 피부질환 응용 분야에 적합함을 시사한다.