[논문 리뷰] Robust Split Federated Learning for U-shaped Medical Image Networks

Ziyuan Yang, Yingyu Chen외 5명|arXiv (Cornell University)|2022. 12. 13.

Privacy-Preserving Technologies in Data인용 수 9

한 줄 요약

이 논문은 U자형 의료 영상 네트워크의 개인정보 보호 및 계산 효율적인 훈련을 가능하게 하는 하이브리드 프레임워크인 강건한 분할 연합 학습(RoS-FL)을 제안한다. 입력, 레이블, 출력 및 모델 개인정보를 보호하면서 병렬 훈련을 가능하게 하기 위해 모델을 머리, 몸통, 꼬리의 세 부분으로 분할하여 클라이언트, 계산 서버, 집계 서버에 분산 배치한다. 주요 기여는 데이터 이질성 하에서 훈련을 안정화시키기 위한 동적 가중치 보정 전략(DWCS)으로, 모델 이탈을 감소시키고 클라이언트 간 성능 일관성을 향상시킨다.

ABSTRACT

U-shaped networks are widely used in various medical image tasks, such as segmentation, restoration and reconstruction, but most of them usually rely on centralized learning and thus ignore privacy issues. To address the privacy concerns, federated learning (FL) and split learning (SL) have attracted increasing attention. However, it is hard for both FL and SL to balance the local computational cost, model privacy and parallel training simultaneously. To achieve this goal, in this paper, we propose Robust Split Federated Learning (RoS-FL) for U-shaped medical image networks, which is a novel hybrid learning paradigm of FL and SL. Previous works cannot preserve the data privacy, including the input, model parameters, label and output simultaneously. To effectively deal with all of them, we design a novel splitting method for U-shaped medical image networks, which splits the network into three parts hosted by different parties. Besides, the distributed learning methods usually suffer from a drift between local and global models caused by data heterogeneity. Based on this consideration, we propose a dynamic weight correction strategy ( extbf{DWCS}) to stabilize the training process and avoid model drift. Specifically, a weight correction loss is designed to quantify the drift between the models from two adjacent communication rounds. By minimizing this loss, a correction model is obtained. Then we treat the weighted sum of correction model and final round models as the result. The effectiveness of the proposed RoS-FL is supported by extensive experimental results on different tasks. Related codes will be released at https://github.com/Zi-YuanYang/RoS-FL.

연구 동기 및 목표

• 원시 데이터, 레이블, 출력이 민감한 U자형 의료 영상 네트워크에서의 개인정보 우려를 해결하기 위해.
• 클라이언트의 계산 부담을 줄이기 위해 무거운 모델 구성 요소를 서버로 오프로드함으로써 연합 학습의 클라이언트 부담을 감소시키기 위해.
• 분산 훈련 환경에서 데이터 이질성으로 인한 모델 이탈을 완화하기 위해.
• 클라이언트와 서버 간의 데이터 및 모델 개인정보를 유지하면서 병렬 훈련을 가능하게 하기 위해.
• 비독립 동일 분포(non-IID) 데이터를 가진 클라이언트 간의 훈련 안정성과 성능 일관성을 향상시키기 위한 새로운 동적 가중치 보정 메커니즘을 통해.

제안 방법

• U자형 네트워크를 머리(클라이언트 측), 몸통(서버 측), 꼬리(클라이언트 측)의 세 구성 요소로 분할하여 계산을 분산하고 개인정보 보호를 향상시키기 위해.
• 연합 학습과 분할 학습을 융합한 하이브리드 아키텍처를 사용하여 병렬 훈련을 가능하게 하고 클라이언트 측 계산 부담을 감소시키기 위해.
• 원시 입력, 레이블, 최종 출력의 전송을 피하기 위해 클라이언트와 서버 간에 중간 특징 맵만 전송하기 위해.
• 현재 및 이전 전역 모델의 가중치 합을 계산하여 이탈을 보정하는 동적 가중치 보정 전략(DWCS)을 구현하기 위해.
• 연속적인 통신 라운드 간의 모델 이탈을 정량화하기 위해 가중치 보정 손실을 정의하고, 이 손실을 최소화하도록 보정 모델을 최적화하기 위해.
• 클라이언트 측 기울기와 서버 측 모델 업데이트를 중앙 집중식 집계 서버에서 집계하여 전역 모델을 업데이트하기 위해.

실험 결과

연구 질문

• RQ1 하이브리드 분할 연합 학습 프레임워크가 U자형 의료 영상 네트워크에서 입력, 레이블, 출력 및 모델 파라미터를 효과적으로 보호할 수 있는가?
• RQ2 의료 영상 모델의 분산 훈련에서 데이터 이질성으로 인한 모델 이탈은 어떻게 완화할 수 있는가?
• RQ3 제안된 동적 가중치 보정 전략(DWCS)이 비독립 동일 분포 데이터를 가진 클라이언트 간의 훈련 안정성과 성능 일관성에 얼마나 기여하는가?
• RQ4 제안된 RoS-FL 프레임워크는 표준 연합 학습에 비해 클라이언트 측 계산 비용을 줄이면서도 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 수정 가중치 μ와 같은 하이퍼파rameter 선택에 대해 RoS-FL의 성능는 얼마나 민감한가?

주요 결과

• RoS-FL는 모델 이탈로 인한 기존 방법에서 관찰된 심각한 성능 격차 없이 클라이언트 간 일관된 성능을 달성한다.
• 복원 작업에서 최고의 PSNR 42.81을 1회의 통신 라운드와 500회의 로컬 에포크로 달성했으며, 다른 설정보다 뛰어난 성능을 보였다.
• 분할 작업에서 최고의 Dice 스코어 88.21을 2회의 통신 라운드와 150회의 로컬 에포크로 달성했으며, 데이터 이질성에 대한 강건성을 입증했다.
• 하이퍼파rameter μ에 대해 성능이 크게 민감하지 않으며, 1×10⁻⁶에서 100의 넓은 범위에서 안정적인 성능을 보였고, μ = 1×10⁻⁴일 때 최고의 성능를 기록했다.
• 이미징 작업에서 임계값을 초과해 로컬 에포크를 늘일수록 성능이 악화되는 것으로 나타났으며, 이는 모델 이탈 증가로 인한 것으로, DWCS가 훈련 안정화에 중요한 역할을 한다는 점을 강조한다.
• 시각적 비교 결과 RoS-FL는 다른 방법에 비해 더 선명한 윤곽선과 더 잘 구조화된 출력을 생성하며 더 적은 잡음 요소를 보였다.