[논문 리뷰] SelfFed: Self-Supervised Federated Learning for Data Heterogeneity and Label Scarcity in Medical Images

Sunder Ali Khowaja, Kapal Dev외 2명|arXiv (Cornell University)|2023. 07. 04.

Privacy-Preserving Technologies in Data인용 수 5

한 줄 요약

이 논문은 의료 사물 인터넷(IoMT) 환경에서 데이터 이질성과 레이블 부족 문제를 해결하기 위해 자기지도 학습 연합 학습 프레임워크인 SelfFed를 제안한다. 이는 두 단계의 접근 방식을 사용한다: Swin Transformer를 활용한 마스킹 자동에코딩을 통한 사전 훈련을 통해 데이터 이질성을 다루고, 새로운 집계 전략을 사용한 대비 훈련을 통해 레이블 부족 문제를 완화한다. 이로 인해 비독립 동일 분포(non-IID) 데이터에서 최대 8.8%의 정확도 향상을 달성했으며, 레이블이 10% 밖에 없는 조건에서도 기존 기준 모델들을 능가한다.

ABSTRACT

Self-supervised learning in the federated learning paradigm has been gaining a lot of interest both in industry and research due to the collaborative learning capability on unlabeled yet isolated data. However, self-supervised based federated learning strategies suffer from performance degradation due to label scarcity and diverse data distributions, i.e., data heterogeneity. In this paper, we propose the SelfFed framework for medical images to overcome data heterogeneity and label scarcity issues. The first phase of the SelfFed framework helps to overcome the data heterogeneity issue by leveraging the pre-training paradigm that performs augmentative modeling using Swin Transformer-based encoder in a decentralized manner. The label scarcity issue is addressed by fine-tuning paradigm that introduces a contrastive network and a novel aggregation strategy. We perform our experimental analysis on publicly available medical imaging datasets to show that SelfFed performs better when compared to existing baselines and works. Our method achieves a maximum improvement of 8.8% and 4.1% on Retina and COVID-FL datasets on non-IID datasets. Further, our proposed method outperforms existing baselines even when trained on a few (10%) labeled instances.

연구 동기 및 목표

• 의료 영상 연합 학습에서 데이터 이질성과 레이블 부족이라는 이중 과제를 IoMT 환경에서 해결한다.
• 왜곡된 데이터 분포로 인해 기존 연합 학습 방법이 실패하는 비독립 동일 분포(non-IID) 환경에서의 성능 저하 문제를 해결한다.
• 대규모 레이블 데이터셋에 대한 의존도를 줄이고, 분산된 방식으로 제한된 레이블 데이터를 효과적으로 학습할 수 있도록 한다.
• 다양하고 분산된 의료 영상 클라이언트 간에서 모델 성능을 유지하는 강력하고 개인정보 보호 기반의 프레임워크를 개발한다.
• 자기지도 사전 훈련과 대비 훈련, 적응형 집계를 통합하여 연합 의료 AI의 일반화 능력과 수렴 성능을 향상시킨다.

제안 방법

• 두 단계의 SelfFed 프레임워크를 제안한다: (1) Swin Transformer 인코더를 사용한 마스킹 자동에코딩(MAE)을 통한 사전 훈련으로 분산 방식의 강력한 표현을 학습한다.
• 사전 훈련 중 16x16 이미지 패치의 60%에 대해 증강된 마스킹을 적용하여 강력한 특징 학습을 유도한다.
• 미세 조정 단계에서 동일한 입력의 두 개의 시각을 기반으로 학습하는 대비 네트워크를 도입하여, 레이블 부족 상황에서 표현 품질을 향상시킨다.
• 미세 조정 동안 다수의 참여를 반복하는 클라이언트의 영향을 줄이기 위해 새로운 클라이언트 집계 전략을 설계한다. 이는 과적합된 클라이언트로부터 발생하는 편향을 감소시킨다.
• 사전 훈련과 미세 조정 모두에 AdamW 옵timizer를 사용하고, 코즈인 감쇠와 웜업 전략을 적용하여 연합 학습과 중심화된 학습 간 일관된 초모수를 유지한다.
• 클라이언트가 로컬에서 훈련하고 업데이트된 가중치를 서버에 전송하여 집계하는 글로벌 모델 집계 프로토콜을 사용하여 데이터 프라이버시를 보장한다.

실험 결과

연구 질문

• RQ1 마스킹 자동에코딩을 통한 자기지도 사전 훈련이 연합 의료 영상 학습에서 데이터 이질성으로 인한 성능 저하를 완화할 수 있는가?
• RQ2 제안된 대비 미세 조정과 새로운 집계 전략이 비독립 동일 분포(non-IID) 환경에서 레이블 부족 상황에서 성능을 어떻게 향상시키는가?
• RQ3 제한된 레이블 데이터에서 연합 환경에서 미세 조정된 경우, SelfFed가 ImageNet 사전 훈련 모델(ViT, BEiT, MAE 등)을 초월할 수 있는 정도는 어느 정도인가?
• RQ4 제안된 프레임워크는 비독립 동일 분포(non-IID) 및 동일 분포(IID) 데이터 분포에서 레이블 데이터의 10%만으로도 강건성과 일관성을 유지하는가?
• RQ5 제안된 새로운 집계 전략이 미세 조정 과정에서 높은 참여 빈도를 가진 클라이언트의 영향을 효과적으로 줄여 모델의 공정성과 수렴성을 향상시키는가?

주요 결과

• SelfFed는 비독립 동일 분포 조건에서 ImageNet 사전 훈련 기준 모델 대비 Retina 데이터셋에서 최대 8.8%의 정확도 향상을 달성했고, COVID-FL 데이터셋에서는 4.1% 향상되었다.
• 레이블 샘플의 10%만으로도 SelfFed는 IID 및 비독립 동일 분포 조건에서 100% 레이블 데이터로 훈련된 기준 모델들을 능가하는 성능을 보였다.
• 제안된 SelfFed-MAE 사전 훈련 단계는 데이터 이질성으로 인한 성능 저하를 크게 감소시켰으며, 초기 학습된 Swin Transformer 및 ImageNet 사전 훈련 모델보다 뛰어난 성능을 보였다.
• 새로운 집계 전략을 적용한 대비 미세 조정은 자주 참여하는 클라이언트의 영향을 효과적으로 완화하여 모델의 안정성과 공정성을 향상시켰다.
• CXR14 데이터셋에서 SelfFed는 사전 훈련된 ViT, BEiT, MAE 모델들과 비교해도 열등한 성능를 유지하지 않아, 다양한 의료 영상 벤치마크에서의 강건성을 확인했다.
• 다양한 데이터셋과 환경에서 일관된 성능을 보이며, 제한된 레이블과 이질적인 데이터를 가진 실세계 IoMT 시나리오에서의 유효성을 입증했다.