Auto-FedRL: Federated Hyperparameter Optimization for Multi-institutional Medical Image Segmentation

TL;DR

Auto-FedRL introduces an RL-based method for dynamic hyperparameter tuning in federated learning, improving medical image segmentation across heterogeneous datasets without extensive trial-and-error.

Contribution

It presents a novel reinforcement learning approach for efficient, adaptive hyperparameter optimization in federated learning, specifically tailored for medical imaging applications.

Findings

Effective hyperparameter adjustment improves model stability.

Reduces need for extensive hyperparameter tuning trials.

Demonstrates superior performance on medical image segmentation tasks.

Abstract

Federated learning (FL) is a distributed machine learning technique that enables collaborative model training while avoiding explicit data sharing. The inherent privacy-preserving property of FL algorithms makes them especially attractive to the medical field. However, in case of heterogeneous client data distributions, standard FL methods are unstable and require intensive hyperparameter tuning to achieve optimal performance. Conventional hyperparameter optimization algorithms are impractical in real-world FL applications as they involve numerous training trials, which are often not affordable with limited compute budgets. In this work, we propose an efficient reinforcement learning (RL)-based federated hyperparameter optimization algorithm, termed Auto-FedRL, in which an online RL agent can dynamically adjust hyperparameters of each client based on the current training progress. Extensive experiments are conducted to investigate different search strategies and RL agents. The effectiveness of the proposed method is validated on a heterogeneous data split of the CIFAR-10 dataset as well as two real-world medical image segmentation datasets for COVID-19 lesion segmentation in chest CT and pancreas segmentation in abdominal CT.