[논문 리뷰] Closing the Generalization Gap of Cross-silo Federated Medical Image Segmentation

An Xu, Wenqi Li외 8명|arXiv (Cornell University)|2022. 03. 18.

Privacy-Preserving Technologies in Data인용 수 7

한 줄 요약

이 논문은 개인화된 모델과 모델 선택기 도입을 통해 의료 영상 분할에서 크로스실 페더레이티드 러닝과 중심집합 학습 간의 일반화 갭을 해소하는 새로운 페더레이티드 러닝 프레임워크 FedSM을 제안한다. 새로운 최적화 방법인 SoftPull을 사용하여 국지적 지식과 전역 지식의 균형을 이루며, 클라이언트 드리프트를 방지하면서도 중심집합 학습 수준의 성능을 달성한다.

ABSTRACT

Cross-silo federated learning (FL) has attracted much attention in medical imaging analysis with deep learning in recent years as it can resolve the critical issues of insufficient data, data privacy, and training efficiency. However, there can be a generalization gap between the model trained from FL and the one from centralized training. This important issue comes from the non-iid data distribution of the local data in the participating clients and is well-known as client drift. In this work, we propose a novel training framework FedSM to avoid the client drift issue and successfully close the generalization gap compared with the centralized training for medical image segmentation tasks for the first time. We also propose a novel personalized FL objective formulation and a new method SoftPull to solve it in our proposed framework FedSM. We conduct rigorous theoretical analysis to guarantee its convergence for optimizing the non-convex smooth objective function. Real-world medical image segmentation experiments using deep FL validate the motivations and effectiveness of our proposed method.

연구 동기 및 목표

• 비독립 동일 분포(non-iid) 데이터와 클라이언트 드리프트로 인해 발생하는 크로스실 페더레이티드 러닝에서 의료 영상 분할의 일반화 갭을 해소하기 위해.
• 모든 클라이언트에 맞는 단일 전역 모델이 필요 없도록, 각 클라이언트별로 개인화된 모델을 학습함으로써 이를 제거하기 위해.
• 데이터 분포 유사도 기반으로 주어진 테스트 입력에 가장 적합한 개인화된 모델을 동적으로 선택할 수 있는 모델 선택기를 설계하기 위해.
• 새로운 개인화된 페더레이티드 러닝 최적화 공식과, 국지적 및 전역 모델 업데이트를 균형 있게 조절하여 일반화를 향상시키는 SoftPull이라는 방법을 개발하기 위해.
• SoftPull이 최적화하는 비볼록 스무쓰 목적 함수에 대해 엄밀한 이론적 수렴 분석을 수행하고, 실험적으로 FedSM이 중심집합 학습 성능과 동등한 성능을 달성함을 검증하기 위해.

제안 방법

• 각 클라이언트별로 개인화된 모델의 집합과 테스트 입력을 가장 적합한 개인화된 모델로 라우팅하는 공유 모델 선택기를 학습하는 프레임워크인 FedSM을 제안한다.
• 국지적 및 전역 모델 업데이트 간의 가중치 조합으로 구성된 새로운 개인화된 페더레이티드 러닝 목적 함수를 도입하며, 이는 학습 가능한 계수 λ로 매개변수화된다.
• 국지 데이터와 전역 모델 양쪽에 대한 강한 유사도를 유도하는 부드러운 제약 조건을 적용함으로써 개인화된 목적 함수를 최적화하는 SoftPull이라는 방법을 개발한다.
• 분할 헤드의 최종 레이어에서 유도된 특징 임베딩 기반의 모델 선택기를 사용하여, 테스트 입력에 대해 가장 적합한 개인화된 모델을 동적으로 선택할 수 있도록 한다.
• 특정 클라이언트에 대해 개인화된 모델이 충분히 확신이 없을 경우 전역 모델로의 패스를 허용하기 위해 모델 선택기에서 임계값 γ를 사용함으로써, 미리 보지 못한 클라이언트에 대한 강건성을 향상시킨다.
• SoftPull이 최적화하는 비볼록 스무쉬 목적 함수에 대해 표준 페더레이티드 러닝 가정 하에 엄밀한 이론적 수렴 분석을 제공한다.

Figure 1 : The proposed FedSM framework with “super model”.

실험 결과

연구 질문

• RQ1 페더레이티드 러닝 프레임워크가 크로스실 의료 영상 분할에서 페더레이티드 러닝과 중심집합 학습 간의 일반화 갭을 해소할 수 있는가?
• RQ2 개인화된 모델을 통해 클라이언트 드리프트를 방지하면 기존의 FedAvg와 같은 전통적 페더레이티드 러닝 방법보다 더 나은 일반화 성능을 달성할 수 있는가?
• RQ3 모델 선택기가 데이터 분포 유사도 기반으로 주어진 테스트 입력에 가장 적합한 개인화된 모델을 효과적으로 식별할 수 있는가?
• RQ4 SoftPull 최적화 방법이 기존의 개인화된 페더레이티드 러닝 베이스라인보다 더 나은 일반화 성능을 달성할 수 있는가?
• RQ5 개인화된 목적 함수에 통합된 조절 가능한 보간 계수 λ가 작은 국지 데이터셋에서의 과적합을 방지함으로써 일반화 성능을 향상시킬 수 있는가?

주요 결과

• FedSM는 실제 의료 영상 분할 작업에서 중심집합 학습과 비교할 만한 테스트 Dice 점수를 달성하여, 일반화 갭을 효과적으로 해소한다.
• 망막 분할 작업에서 FedSM은 FedAvg와 로컬 학습을 모두 능가하며, 특히 데이터셋이 작은 클라이언트에서 성능 향상이 두드러지며, 최적의 임계값 γ를 사용할 경우 성능이 최대 3% 향상된다.
• SoftPull은 APFL, Per-FedAvg, FT와 같은 다른 개인화 방법보다 뛰어난 성능을 보이며, 다음으로 우수한 베이스라인보다 0.62% 높은 전역 Dice 계수를 달성한다.
• 낮은 데이터 유사도 작업(예: 망막)의 경우 최적의 λ 값은 0.7이며, 높은 유사도 작업(예: 전립선)의 경우 0.3로, SoftPull 공식의 적응 가능성은 검증된다.
• 손실 표면 분석 결과, SoftPull은 FedAvg보다 더 평탄한 손실 지형을 달성함으로써, 약간의 수렴 오차가 존재하더라도 더 나은 일반화를 보여준다.
• 모델 선택기는 새로운 클라이언트의 데이터 분포와 가장 유사한 클라이언트를 성공적으로 식별하였으며, 특징 시각화(t-SNE)와 실험적 성능 모두와 일치하는 선택 빈도를 보였다.

Figure 2 : Training curves comparison. The curves are non-decreasing because we record the best result during training.