Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Election of Collaborators via Reinforcement Learning for Federated Brain Tumor Segmentation

Muhammad Irfan Khan, Elina Kontio, Suleiman A. Khan, Mojtaba Jafaritadi

TL;DR

This study addresses efficient collaborator selection in federated brain tumor segmentation by framing collaborator choice as a multi-armed bandit problem and evaluating RL strategies (Epsilon-Greedy and Upper Confidence Bound). It introduces RL-HSimAgg, a similarity-weighted, harmonic-mean aggregation method that mitigates outlier influence in non-IID FL settings. Empirical results on FeTS/BraTS-derived mpMRI data show that UCB-based collaborator selection yields higher Dice scores across enhancing tumor, tumor core, and whole tumor regions, with improved convergence and reduced Hausdorff distance, both on internal and external validation. The work demonstrates that RL-guided collaboration and robust aggregation can enhance model robustness and generalization in distributed medical imaging tasks, with potential applicability to dynamic FL environments.

Abstract

Federated learning (FL) enables collaborative model training across decentralized datasets while preserving data privacy. However, optimally selecting participating collaborators in dynamic FL environments remains challenging. We present RL-HSimAgg, a novel reinforcement learning (RL) and similarity-weighted aggregation (simAgg) algorithm using harmonic mean to manage outlier data points. This paper proposes applying multi-armed bandit algorithms to improve collaborator selection and model generalization. By balancing exploration-exploitation trade-offs, these RL methods can promote resource-efficient training with diverse datasets. We demonstrate the effectiveness of Epsilon-greedy (EG) and upper confidence bound (UCB) algorithms for federated brain lesion segmentation. In simulation experiments on internal and external validation sets, RL-HSimAgg with UCB collaborator outperformed the EG method across all metrics, achieving higher Dice scores for Enhancing Tumor (0.7334 vs 0.6797), Tumor Core (0.7432 vs 0.6821), and Whole Tumor (0.8252 vs 0.7931) segmentation. Therefore, for the Federated Tumor Segmentation Challenge (FeTS 2024), we consider UCB as our primary client selection approach in federated Glioblastoma lesion segmentation of multi-modal MRIs. In conclusion, our research demonstrates that RL-based collaborator management, e.g. using UCB, can potentially improve model robustness and flexibility in distributed learning environments, particularly in domains like brain tumor segmentation.

Election of Collaborators via Reinforcement Learning for Federated Brain Tumor Segmentation

TL;DR

This study addresses efficient collaborator selection in federated brain tumor segmentation by framing collaborator choice as a multi-armed bandit problem and evaluating RL strategies (Epsilon-Greedy and Upper Confidence Bound). It introduces RL-HSimAgg, a similarity-weighted, harmonic-mean aggregation method that mitigates outlier influence in non-IID FL settings. Empirical results on FeTS/BraTS-derived mpMRI data show that UCB-based collaborator selection yields higher Dice scores across enhancing tumor, tumor core, and whole tumor regions, with improved convergence and reduced Hausdorff distance, both on internal and external validation. The work demonstrates that RL-guided collaboration and robust aggregation can enhance model robustness and generalization in distributed medical imaging tasks, with potential applicability to dynamic FL environments.

Abstract

Federated learning (FL) enables collaborative model training across decentralized datasets while preserving data privacy. However, optimally selecting participating collaborators in dynamic FL environments remains challenging. We present RL-HSimAgg, a novel reinforcement learning (RL) and similarity-weighted aggregation (simAgg) algorithm using harmonic mean to manage outlier data points. This paper proposes applying multi-armed bandit algorithms to improve collaborator selection and model generalization. By balancing exploration-exploitation trade-offs, these RL methods can promote resource-efficient training with diverse datasets. We demonstrate the effectiveness of Epsilon-greedy (EG) and upper confidence bound (UCB) algorithms for federated brain lesion segmentation. In simulation experiments on internal and external validation sets, RL-HSimAgg with UCB collaborator outperformed the EG method across all metrics, achieving higher Dice scores for Enhancing Tumor (0.7334 vs 0.6797), Tumor Core (0.7432 vs 0.6821), and Whole Tumor (0.8252 vs 0.7931) segmentation. Therefore, for the Federated Tumor Segmentation Challenge (FeTS 2024), we consider UCB as our primary client selection approach in federated Glioblastoma lesion segmentation of multi-modal MRIs. In conclusion, our research demonstrates that RL-based collaborator management, e.g. using UCB, can potentially improve model robustness and flexibility in distributed learning environments, particularly in domains like brain tumor segmentation.
Paper Structure (17 sections, 6 equations, 1 figure, 3 tables, 3 algorithms)

This paper contains 17 sections, 6 equations, 1 figure, 3 tables, 3 algorithms.

Table of Contents

  1. Methods
  2. Data
  3. Multi-Armed Bandit Problem
  4. Epsilon-Greedy Strategy
  5. Upper Confidence Bounds (UCB) Strategy
  6. Election of Collaborators
  7. Epsilon-Greedy Approach
  8. Upper Confidence Bounds Approach
  9. Similarity Weighted Aggregation
  10. Weight Aggregation Policy
  11. Deep Learning Experiments
  12. Training Setup
  13. Results
  14. Model training and performance using internal validation data
  15. Model performance using external validation data
  16. ...and 2 more sections

Figures (1)

  • Figure 1: Performance metrics for model training of RL-HSimAgg. The horizontal axis refers to the number of rounds and the vertical axis to the metrics.