[논문 리뷰] Federated Cross Learning for Medical Image Segmentation

Xuanang Xu, Deng, Hannah H.외 7명|arXiv (Cornell University)|2022. 04. 05.

Privacy-Preserving Technologies in Data인용 수 11

한 줄 요약

이 논문은 의료 영상 분할을 위한 모델 평균화가 필요 없는 분산 학습 프레임워크인 페더레이티드 크로스 러닝(FedCross)을 제안한다. 비독립 동일 분포(non-i.i.i.d.) 데이터에서 성능 저하를 완화하기 위해 클라이언트 간에 순환 방식으로 모델을 순차적으로 훈련시킨다. FedCrossEns는 앙상블 확장 기법으로, 중심화된 학습 수준의 정확도를 달성하고 픽셀 수준의 불확실성 지도를 제공하여 임상 적용 가능성을 높이며, 모델 평균화에 기인한 불안정성을 피한다.

ABSTRACT

Federated learning (FL) can collaboratively train deep learning models using isolated patient data owned by different hospitals for various clinical applications, including medical image segmentation. However, a major problem of FL is its performance degradation when dealing with data that are not independently and identically distributed (non-iid), which is often the case in medical images. In this paper, we first conduct a theoretical analysis on the FL algorithm to reveal the problem of model aggregation during training on non-iid data. With the insights gained through the analysis, we propose a simple yet effective method, federated cross learning (FedCross), to tackle this challenging problem. Unlike the conventional FL methods that combine multiple individually trained local models on a server node, our FedCross sequentially trains the global model across different clients in a round-robin manner, and thus the entire training procedure does not involve any model aggregation steps. To further improve its performance to be comparable with the centralized learning method, we combine the FedCross with an ensemble learning mechanism to compose a federated cross ensemble learning (FedCrossEns) method. Finally, we conduct extensive experiments using a set of public datasets. The experimental results show that the proposed FedCross training strategy outperforms the mainstream FL methods on non-iid data. In addition to improving the segmentation performance, our FedCrossEns can further provide a quantitative estimation of the model uncertainty, demonstrating the effectiveness and clinical significance of our designs. Source code is publicly available at https://github.com/DIAL-RPI/FedCross.

연구 동기 및 목표

• 병원 간 데이터가 비독립 동일 분포일 경우 의료 영상 분할에서 페더레이티드 러닝(FL)의 성능 저하 문제를 해결한다.
• FedAvg 및 그 변종에서 발생하는 모델 평균화의 한계를 극복하여 기울기의 차이로 인한 최적해 이격 문제를 해결한다.
• 이질적인 의료 영상 데이터에서 안정성과 성능을 향상시키기 위해 파rameter 집계를 피하는 훈련 전략을 개발한다.
• 임상 의사결정 지원을 위해 예측의 신뢰도를 추정함으로써 모델 신뢰도를 향상시킨다.
• 데이터 프라이버시를 유지하면서도 중심화된 학습 수준의 성능을 달성한다.

제안 방법

• 모델 집계가 필요 없는 새로운 FL 전략인 FedCross를 제안하며, 글로벌 모델을 클라이언트 간 순환 방식으로 순차적으로 훈련시는 방식이다.
• 로컬 모델 파라미터의 평균화를 피하기 위해 한 클라이언트의 모델을 한 번에 하나씩 훈련시음으로써, 비독립 동일 분포 조건에서 오차의 증폭을 줄인다.
• 다수의 FedCross 모델을 동시에 훈련하고 예측 결과를 통합함으로써 정확도 향상과 불확실성 추정을 향상시키는 앙상블 방법인 FedCrossEns를 도입한다.
• 원시 데이터를 공유하지 않고도 중심화된 훈련의 성능을 근사하기 위해 앙상블 학습을 활용한다.
• 앙상블 분산에서 유도된 불확실성 지도를 활용해 예측의 신뢰도를 정량화하며, 임상 검증 워크플로우를 지원한다.
• 이론적 분석을 통해 FedAvg에서의 모델 평균화가 비독립 동일 분포 조건에서 오차를 증폭시킴을 밝혀내며, 집계 없는 설계의 필요성을 뒷받침한다.

실험 결과

연구 질문

• RQ1 비독립 동일 분포 조건에서 FedAvg의 모델 평균화가 성능을 어떻게 저하시키는가?
• RQ2 모델 집계를 피하는 페더레이티드 훈련 전략이 비독립 동일 분포 조건에서 의료 영상 분할 작업에서 더 나은 성능을 내는가?
• RQ3 FedCross와 같은 집계 없는 방법이 프라이버시 보호 환경에서 중심화된 학습의 정확도를 유지하거나 초월할 수 있는가?
• RQ4 FedCrossEns는 분할 결과에 대해 신뢰할 수 있는 불확실성 추정을 제공할 수 있는가? 이는 임상적 신뢰도와 사용성 향상에 기여하는가?
• RQ5 로컬 훈련 에포크 수가 집계 기반 및 집계 없는 FL 방법의 안정성과 성능에 미치는 영향은 어떠한가?

주요 결과

• FedCross는 PROMISE12 데이터셋에서 비독립 동일 분포 조건에서 최고의 DICE 스코어 90.77±2.47을 기록하며, FedAvg, FedProx, FedBN을 모두 앞섰다.
• FedCrossEns는 테스트 세트에서 DICE 스코어 88.73을 기록하여 중심화된 학습의 성능 90.66에 매우 가까이 도달했다.
• FedCrossEns가 생성한 불확실성 지도는 그림 4에 시각화된 바와 같이 임상적으로 유용한 픽셀 수준의 신뢰도 추정을 제공했다.
• 로컬 훈련 에포크 수가 1에서 16으로 증가함에 따라 FedCross는 집계 기반 방법(FedAvg, FedProx, FedBN)보다 성능 저하가 적었으며, 이는 안정성을 확인하는 데 기여했다.
• 작은 표본 수를 가진 데이터셋(19, 30, 48명의 환자)에서 모든 FL 방법이 국소 훈련보다 유의미하게 높은 성능을 보였으며, 지식 전이의 효율성을 입증했다.
• 제거 실험을 통해 FedCross에서 모델 집계의 부재가 특히 높은 로컬 훈련 에포크 수 조건에서 오차 증폭을 감소시킴을 확인했다.