Dynamic Hyperparameter Adjustment via Reinforcement Learning in Asynchronous Federated Learning for Medical Image Analysis
摘要
Federated learning is gaining prominence in the healthcare domain due to its ability to protect data privacy and facilitate collaboration among multiple institutions. However, the heterogeneity among clients poses a significant challenge due to differences in client data and computing facilities. Conventional synchronized federated learning, e.g., FedAvg, often leads to inefficient system utilization in training. Thus, Asynchronous Federated Learning (AFL) paradigms have been proposed to address the straggler problem. However, it also introduces the problem of model staleness, leading to significant accuracy drops. Moreover, in the context of heterogeneous client data, fixed hyperparameters (e.g., learning rate, aggregation weights, and local epochs) can detrimentally impact the convergence speed and overall model performance. To tackle these challenges, we propose a novel asynchronous federated learning architecture, AFedRL, that employs reinforcement learning to dynamically adjust hyperparameters in AFL, accelerating the convergence of the global model. It aggregates client models as soon as any client finishes its local training, while the REINFORCE algorithm is utilized to adjust hyperparameters for each client iteratively, with a reward function that considers the relative reduction in the loss of the global models across iterations and evaluates aleatoric uncertainty in the local model and epistemic uncertainty in the global model. Experiments on two multi-center medical image segmentation tasks demonstrate that AFedRL achieves competitive accuracy while significantly reducing training time. This showcases its effectiveness and practical advantages for real-world applications. The code is available at https://github.com/shuanggu0815/afedrl .