[논문 리뷰] FedIIC: Towards Robust Federated Learning for Class-Imbalanced Medical Image Classification

Nannan Wu, Li Yu외 3명|arXiv (Cornell University)|2022. 06. 28.

Imbalanced Data Classification Techniques인용 수 5

한 줄 요약

FedIIC는 특징 학습과 분류기 캘리브레이션을 공동 최적화하여 의료 영상 분류에서의 클래스 불균형 문제를 해결하는 프라이버시 보장형 분산 학습 프레임워크를 제안한다. 이는 내부 클라이언트 및 상호 클라이언트 수준의 이중 대비 학습을 통해 클래스별 표현을 향상시키고, 난이도 인식 로짓 조정을 통해 결정 경계를 동적으로 균형 잡는다. 실세계 및 시뮬레이션된 의료 데이터셋에서 최신 기술 수준의 성능을 달성하며, 균형 잡힌 정확도에서 뚜렷한 향상을 보였다.

ABSTRACT

Federated learning (FL), training deep models from decentralized data without privacy leakage, has shown great potential in medical image computing recently. However, considering the ubiquitous class imbalance in medical data, FL can exhibit performance degradation, especially for minority classes (e.g. rare diseases). Existing methods towards this problem mainly focus on training a balanced classifier to eliminate class prior bias among classes, but neglect to explore better representation to facilitate classification performance. In this paper, we present a privacy-preserving FL method named FedIIC to combat class imbalance from two perspectives: feature learning and classifier learning. In feature learning, two levels of contrastive learning are designed to extract better class-specific features with imbalanced data in FL. In classifier learning, per-class margins are dynamically set according to real-time difficulty and class priors, which helps the model learn classes equally. Experimental results on publicly-available datasets demonstrate the superior performance of FedIIC in dealing with both real-world and simulated multi-source medical imaging data under class imbalance. Code is available at https://github.com/wnn2000/FedIIC.

연구 동기 및 목표

• 소수의 질병 클래스가 부족하고 잘 인식되지 않는 문제로 인해 발생하는 분산 학습에서의 의료 영상 분류의 심각한 클래스 불균형 문제를 해결하기 위해.
• 기관 간 원시 데이터 공유가 필요 없는 프라이버시 보장형 FL 프레임워크를 개발하기 위해.
• 글로벌 모델 훈련 중에 발생하는 속성 편향(불균형 배경 속성으로 인한)과 클래스 편향(사전 확률 차이로 인한)을 동시에 완화하기 위해.
• 분산형, 다중 소스 의료 영상 환경에서 주류 클래스의 정확도를 희생시키지 않고도 소수 클래스의 성능을 향상시키기 위해.

제안 방법

• 이중 수준의 대비 학습을 도입한다: 내부 클라이언트 대비 학습은 각 클라이언트의 로컬 데이터 내에서 특징의 구분 능력을 향상시키고, 상호 클라이언트 대비 학습은 클라이언트 간 특징을 정렬하여 일반화 능력을 향상시킨다.
• 클래스 사전 확률에 적응하는 온도 스케줄링 대비 손실을 사용하여 특징 추출기를 캘리브레이션함으로써, 클래스 간 유사도를 감소시키고 표현 품질을 향상시킨다.
• 난이도 인식 로짓 조정(DALA) 메커니즘은 실시간 샘플 난이도와 클래스 사전 확률에 기반하여 분류기 로짓을 동적으로 조정하여 결정 경계를 균형 잡는다.
• 이 프레임워크는 대비 특징 학습과 동적 분류기 캘리브레이션을 통합된 FL 파이프라인 내에 통합하여 속성 편향과 클래스 편향을 동시에 감소시킨다.
• 각 라운드에서 전체 클라이언트 참여를 보장하는 글로벌 모델 집합 전략을 사용하여 이질적이고 불균형한 데이터 소스 간의 강력한 수렴성을 확보한다.
• 표준 딥 러닝 백본을 사용하여 구현되며, 기존의 FL 훈련 프로토콜과 호환되어 임상 환경에서의 실용적 구현을 보장한다.

실험 결과

연구 질문

• RQ1 의료 영상에서 글로벌 클래스 불균형 상황에서 특징 학습과 분류기 캘리브레이션을 공동 최적화하면 성능 향상이 이루어지는가?
• RQ2 데이터가 글로벌하게 불균형할 경우, 이중 대비 학습(내부 및 상호 클라이언트)이 강력하고 클래스 특화된 표현을 학습하는 데 얼마나 효과적인가?
• RQ3 샘플 난이도와 클래스 사전 확률에 기반한 동적 로짓 조정이 더 나은 결정 경계와 향상된 소수 클래스 인식을 이끌어내는가?
• RQ4 FedIIC는 실세계 및 시뮬레이션된 다중 소스, 클래스 불균형 의료 영상 데이터셋을 처리할 때 최신 기술 수준의 방법보다 어떻게 비교되는가?

주요 결과

• FedIIC는 평가된 모든 데이터셋에서 가장 높은 균형 잡힌 정확도(BACC)를 달성하였으며, Real 데이터셋에서는 두 번째로 우수한 방법(CReFF)보다 3.99% 높고, ISIC에서는 7.32% 높으며, ICH에서는 2.01% 높았다.
• 제거 실험을 통해 내부 및 상호 클라이언트 대비 학습과 DALA를 결합한 조합이 가장 우수한 성능을 보였으며, ISIC에서 BACC 78.84%, ICH에서 84.22%를 기록했다.
• 온도 매개변수 t=0.5(클래스 사전 확률 기반)로 설정할 경우 t=0보다 성능 향상이著명하게 향상되어 동적 온도 스케줄링의 중요성을 입증했다.
• 프로토타입 벡터의 직교화는 클래스 간 유사도를 감소시키고 성능을 향상시키며, 적용 시 BACC가 72.64%에서 78.84%로 증가했다.
• 난이도 인식 로짓 조정(d=0.25)을 도입하면 BACC는 향상되지만 F1은 약간 감소하여, 균형 잡힌 정확도와 정밀도/재현율의 조화 평균 사이의 트레이드오���을 확인했다.
• 전체 FedIIC 프레임워크는 ISIC에서 F1 점수 78.05, ICH에서 84.73을 기록하여 모든 지표에서 일관된 향상을 보였다.