[논문 리뷰] Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation

Yingda Xia, Dong Yang외 16명|arXiv (Cornell University)|2021. 04. 20.

Privacy-Preserving Technologies in Data참고 문헌 50인용 수 36

한 줄 요약

Auto-FedAvg는 분산 학습을 위한 데이터 기반의 집계 가중치를 학습하며, Dirichlet-distributed 미분가능한 방식으로 클라이언트 모델을 결합하여 CIFAR-10 및 다기관 의료 분할 과제에서 FedAvg를 개선한다.

ABSTRACT

Federated learning (FL) enables collaborative model training while preserving each participant's privacy, which is particularly beneficial to the medical field. FedAvg is a standard algorithm that uses fixed weights, often originating from the dataset sizes at each client, to aggregate the distributed learned models on a server during the FL process. However, non-identical data distribution across clients, known as the non-i.i.d problem in FL, could make this assumption for setting fixed aggregation weights sub-optimal. In this work, we design a new data-driven approach, namely Auto-FedAvg, where aggregation weights are dynamically adjusted, depending on data distributions across data silos and the current training progress of the models. We disentangle the parameter set into two parts, local model parameters and global aggregation parameters, and update them iteratively with a communication-efficient algorithm. We first show the validity of our approach by outperforming state-of-the-art FL methods for image recognition on a heterogeneous data split of CIFAR-10. Furthermore, we demonstrate our algorithm's effectiveness on two multi-institutional medical image analysis tasks, i.e., COVID-19 lesion segmentation in chest CT and pancreas segmentation in abdominal CT.

연구 동기 및 목표

• 크로스-실로 의료 영상에서 연합 클라이언트 간의 비 i.i.d. 데이터 분포를 다룬다.
• 고정된 데이터 크기 기반 가중치를 사용하는 대신, 가중치를 학습하는 미분가능한 메커니즘을 개발한다.
• 추가 통신을 최소화하여 프라이버시를 보장하는 연합을 확보한다.
• 의료 영상 과제에서 FedAvg와 최첨단 FL 방법들에 비해 개선을 입증한다.
• 네트워크 단위 및 계층별 집계 변형과 Dirichlet 대 softmax 매개변수화를 탐구한다.

제안 방법

• 반복적 업데이트를 위해 로컬 모델 매개변수와 전역 집계 매개변수를 분리한다.
• 집계 가중치를 학습 가능한 매개변수(softmax 또는 Dirichlet)로 표현하여 단순체 제약을 강제한다.
• Dirichlet 분포를 사용하여 집계 가중치를 모델링하고 기울기 기반 최적화를 가능하게 한다.
• 원시 데이터를 공유하지 않고 가중치를 업데이트하는 커뮤니케이션 효율적인 LearnAggWeight 단계를 도입한다.
• 네트워크 단위 및 계층별 집계 체계를 제안하여 계층별 또는 전체 네트워크 차원에서 클라이언트 기여를 유연하게 가중한다.
• 추가 통신과 성능의 균형을 맞추기 위해 aggregation weight 학습 간격 t0를 분석한다.

실험 결과

연구 질문

• RQ1 학습 가능한 데이터 기반의 집계 가중 announced? (Note: this line appears to be incorrectly copied from user prompt)

주요 결과

• 네트워크 단위 Dirichlet 가중치를 갖는 Auto-FedAvg(Auto-FedAvg-N-Dirichlet)은 CIFAR-10에서 변형 중 최상의 성능을 달성했으며 최종 정확도는 88.98%이고, FedMA(87.53%)보다 1.45% 앞섰다.
• 다국적 COVID-19 병변 분할에서 Auto-FedAvg-N-Dirichlet 변형은 전역 모델 성능, 평균 로컬 모델 성능 및 로컬 일반화 측면에서 FedAvg보다 개선을 보였으며, 예를 들어 보고된 실행에서 전역 테스트 평균이 FedAvg보다 약 2% 포인트 증가했다.
• Auto-FedAvg는 CIFAR-10 및 두 개의 의료 영상 과제( COVID-19 병변 분할 및 췌장 분할) 전반에 걸쳐 FedAvg, FedProx 및 FedMA 기본값을 지속적으로 능가한다.
• Dirichlet 집계 가중치를 사용하는 것이 일반적으로 Softmax보다 더 나은 결과를 내며, 네트워크 단위 집계가 대부분의 설정에서 계층별 집계보다 우수하다.
• 학습 초기에 가중치를 적극적으로 학습하면 수렴 속도가 빨라지고, 이후 가중치를 안정화시키는 것이 일반화 성능을 더 향상시킬 수 있다.
• 췌장 분할 실험에서 Auto-FedAvg-N-Dirichlet*가 최상의 글로벌 및 로컬 성능을 달성했고 FedAvg에 비해 현저한 개선이 있었다.