[논문 리뷰] Optimizing Federated Learning for Medical Image Classification on Distributed Non-iid Datasets with Partial Labels

Pranav Kulkarni, Adway Kanhere외 2명|arXiv (Cornell University)|2023. 03. 10.

COVID-19 diagnosis using AIMedicine인용 수 3

한 줄 요약

이 논문은 의료 영상 분류에서 데이터 이질성과 부분 레이블로 인한 수렴 문제를 완화하기 위해 사전 학습된 모델의 배치 정규화 레이어를 동결하는 FedFBN이라는 플랫폼을 제안한다. 전이 학습과 동결된 배치 정규화 레이어를 결합함으로써 FedFBN은 다양한 비독립identically distributed (non-iid) 및 부분적으로 레이블이 부여된 데이터셋에서 FedAvg와 FedBN을 능가하는 성능을 달성하며, 실제 의료 영상 환경에서의 일반화 능력과 강건성을 향상시킨다.

ABSTRACT

Numerous large-scale chest x-ray datasets have spearheaded expert-level detection of abnormalities using deep learning. However, these datasets focus on detecting a subset of disease labels that could be present, thus making them distributed and non-iid with partial labels. Recent literature has indicated the impact of batch normalization layers on the convergence of federated learning due to domain shift associated with non-iid data with partial labels. To that end, we propose FedFBN, a federated learning framework that draws inspiration from transfer learning by using pretrained networks as the model backend and freezing the batch normalization layers throughout the training process. We evaluate FedFBN with current FL strategies using synthetic iid toy datasets and large-scale non-iid datasets across scenarios with partial and complete labels. Our results demonstrate that FedFBN outperforms current aggregation strategies for training global models using distributed and non-iid data with partial labels.

연구 동기 및 목표

• 데이터가 분산되어 있고 비독립identically distributed이며 부분 레이블이 부여된 상황에서 정확한 글로벌 모델을 훈련하는 데 도전하는 것.
• 의료 영상에서 기관 간 도메인 이동과 레이블 겹침을 다루는 데 있어 FedAvg와 FedBN의 한계를 극복하는 것.
• 전이 학습을 활용하여 사전 학습된 모델의 배치 정규화 레이어를 동결함으로써 훈련을 안정화하고 수렴을 향상시키는 것.
• 다양한 레이블 완전성과 데이터 이질성을 가진 합성 및 실세계 데이터셋에서 FedFBN의 성능을 평가하는 것.
• 동결된 배치 정규화 레이어가 연합 설정에서 적응형 또는 집계된 배치 정규화 통계보다 더 나은 일반화를 이끌 수 있음을 보여주는 것.

제안 방법

• FedFBN은 훈련 과정 전반에 걸쳐 배치 정규화 레이어 통계를 동결함으로써 FedBN을 수정하여 로컬 및 글로벌 집계 동안 업데이트를 방지한다.
• 프레임워크는 전이 학습을 활용하기 위해 사전 학습된 딥 네ural 네트워크를 모델 백본으로 사용하여 가중치를 초기화하고 수렴을 향상시킨다.
• 로컬 훈련 및 글로벌 집계 중에 배치 정규화 레이어가 업데이트되지 않아 도메인 이동으로 인한 성능 저하를 방지한다.
• 글로벌 모델은 비-BN 레이어 가중치를 FedAvg 방식으로 집계하여 업데이트하며, 배치 정규화 통계는 모든 클라이언트에서 고정된다.
• 해당 방법은 NIH 흉부 X-ray 14 데이터셋에서 유도된 합성 iid 및 non-iid 데이터셋, 그리고 외부 MIMIC-CXR 테스트 세트에서 평가된다.
• 성능은 공유 레이블, 모든 레이블, 부분 레이블 겹침을 포함한 여러 레이블 구성에서 AUROC를 측정하여 평가된다.

(a) 4 Shared Disease Labels

실험 결과

연구 질문

• RQ1 연합 학습 환경에서 배치 정규화 레이어를 동결하는 것이 비독립identically distributed, 부분적으로 레이블이 부여된 의료 영상 분류 데이터셋에서 모델 수렴과 성능을 향상시킬 수 있는가?
• RQ2 레이블 완전성과 데이터 이질성의 변화에 따라 FedFBN은 일반화 및 강건성 측면에서 FedAvg와 FedBN에 비해 어떻게 비교되는가?
• RQ3 사전 학습된 모델을 사용하고 배치 정규화 레이어를 동결하는 것이 연합 의료 영상 분류에서 통계를 업데이트하거나 집계하는 전략을 능가하는가?
• RQ4 FedFBN은 데이터 공유 없이도 중앙에서 훈련된 모델과 유사한 성능을 달성할 수 있는가?
• RQ5 MIMIC-CXR와 같은 도메인 이동이 있는 외부 테스트 세트에서 FedFBN은 베이스라인 및 기존의 FL 방법에 비해 어떻게 성능을 내는가?

주요 결과

• 합성 iid 데이터셋에서 4개의 공통 질환 레이블에 대해 FedFBN은 평균 AUROC 0.89를 기록하여 기준 모델보다 유의미하게 높은 성능을 보였다 (p<0.001).
• MIMIC 외부 테스트 세트에서 FedFBN은 평균 AUROC 0.75를 기록하여 비교된 모든 모델을 능가했다 (p<0.001).
• 완전한 레이블이 있는 non-iid 데이터셋에서 FedFBN은 기준 NIH 모델의 성능을 따라잡았으며 (평균 AUROC 0.84, p=0.1), FedAvg와 로컬 FedBN을 능가했다 (p<0.001).
• 부분 레이블이 있는 non-iid 데이터셋에서 FedFBN은 7개의 공통 레이블에 대해 평균 AUROC 0.84 (p=0.03)와 모든 레이블에 대해 평균 AUROC 0.81 (p=0.06)을 기록하여 FedAvg와 로컬 FedBN을 능가했다 (p<0.001).
• MIMIC-CXR를 포함한 모든 테스트 세트에서 FedFBN은 FedAvg와 FedBN보다 더 우수한 일반화 능력을 보였으며, 평균 AUROC 0.75 (p<0.001)를 기록했다.
• FedFBN은 데이터 공유 없이도 중앙에서 훈련된 NIH 기준 모델과 유사한 성능 (평균 AUROC 0.81, p=0.06)을 보였으며, 임상적 유용성을 입증했다.

(b) All Disease Labels