Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation

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.

References in corpus (15)

• Federated Learning with Non-IID Data
• On the Convergence of FedAvg on Non-IID Data
• A large annotated medical image dataset for the development and evaluation of segmentation algorithms
• AutoAugment: Learning Augmentation Policies from Data
• FedMD: Heterogenous Federated Learning via Model Distillation
• Federated Optimization in Heterogeneous Networks
• Communication-Efficient On-Device Machine Learning: Federated Distillation and Augmentation under Non-IID Private Data
• Domain Generalization via Model-Agnostic Learning of Semantic Features
• Adaptive Personalized Federated Learning
• Uncertainty-aware multi-view co-training for semi-supervised medical image segmentation and domain adaptation
• Implicit Reparameterization Gradients
• Federated Learning with Matched Averaging
• Privacy-preserving Federated Brain Tumour Segmentation
• FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning
• Multi-site fMRI Analysis Using Privacy-preserving Federated Learning and Domain Adaptation: ABIDE Results

Cited by in corpus (5)

• Federated Continual Learning via Knowledge Fusion: A Survey
• Federated Multi-organ Segmentation with Inconsistent Labels
• Optimality and Stability in Federated Learning: A Game-theoretic Approach
• FedMM: Saddle Point Optimization for Federated Adversarial Domain Adaptation
• Multi-task Federated Learning for Heterogeneous Pancreas Segmentation