[Paper Review] Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation

Yingda Xia, Dong Yang+16 more|arXiv (Cornell University)|Apr 20, 2021

Privacy-Preserving Technologies in Data50 references36 citations

TL;DR

Auto-FedAvg learns data-driven aggregation weights for federated learning, using a Dirichlet-distributed, differentiable approach to combine client models, improving FedAvg in CIFAR-10 and multi-institutional medical segmentation tasks.

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.

Motivation & Objective

• Address non-i.i.d. data distributions across federated clients in cross-silo medical imaging.
• Develop a differentiable mechanism to learn aggregation weights rather than using fixed, data-size-based weights.
• Ensure privacy-preserving federation with minimal extra communication.
• Demonstrate improvements over FedAvg and state-of-the-art FL methods on medical imaging任务.
• Explore network-wise and layer-wise aggregation variants and Dirichlet vs softmax parameterizations.

Proposed method

• Disentangle local model parameters and global aggregation parameters for iterative updates.
• Represent aggregation weights as learnable parameters (softmax or Dirichlet) to enforce simplex constraints.
• Use a Dirichlet distribution to model aggregation weights and enable gradient-based optimization.
• Introduce a communication-efficient LearnAggWeight step that updates aggregation weights without sharing raw data.
• Propose network-wise and layer-wise aggregation schemes to flexibly weight client contributions by layer or whole network.
• Analyze aggregation weight learning interval t0 to balance extra communication with performance.

Experimental results

Research questions

• RQ1 Can learnable, data-driven aggregation weights improve federated averaging under non-i.i.d. data in cross-silo FL?
• RQ2 Do Dirichlet-distributed aggregation weights yield better performance and generalizability than softmax-based weights?
• RQ3 What is the impact of network-wise vs layer-wise aggregation on performance and communication cost?
• RQ4 How does the aggregation-weights learning interval t0 affect performance and efficiency in cross-institutional settings?

Key findings

• Auto-FedAvg with network-wise Dirichlet weights (Auto-FedAvg-N-Dirichlet) achieves the best results among variants on CIFAR-10, with final accuracy 88.98%, outperforming FedMA (87.53%) by 1.45%.
• On multi-national COVID-19 lesion segmentation, Auto-FedAvg-N-Dirichlet variants show improvements over FedAvg in global model performance, average local model performance, and local generalization, e.g., global test average increased by about 2 percentage points over FedAvg in reported runs.
• Auto-FedAvg consistently outperforms FedAvg, FedProx, and FedMA baselines across CIFAR-10 and two medical imaging tasks (COVID-19 lesion segmentation and pancreas segmentation).
• Using Dirichlet aggregation weights generally yields better results than Softmax across tasks, and network-wise aggregation outperforms layer-wise aggregation in most settings.
• Aggressive learning of aggregation weights early in training accelerates convergence, while stabilizing weights later can further enhance generalization.
• In pancreas segmentation experiments, Auto-FedAvg-N-Dirichlet* achieved the best global and local performance, with notable improvements over FedAvg.