A Comparative Benchmark of Federated Learning Strategies for Mortality Prediction on Heterogeneous and Imbalanced Clinical Data

Technical Summary: A Comparative Benchmark of Federated Learning Strategies for Mortality Prediction

Problem Statement
Predicting in-hospital mortality using Electronic Health Records (EHRs) offers significant potential for improving clinical decision-making. However, the development of high-performance models is hindered by data privacy regulations (e.g., GDPR, HIPAA) that prevent the centralization of multi-institutional data. Federated Learning (FL) addresses this by enabling collaborative training without sharing raw data. Yet, the practical application of FL in healthcare faces two critical, compounding challenges: statistical heterogeneity (non-IID data distributions across institutions due to varying patient demographics and clinical practices) and severe class imbalance (mortality rates are typically low, often <2%). Standard FL algorithms, particularly Federated Averaging (FedAvg), often suffer from performance degradation, slow convergence, or divergence under these conditions. There is a lack of comprehensive, empirical benchmarks comparing modern FL aggregation strategies specifically on large-scale, heterogeneous, and imbalanced clinical datasets.

Methodology
The study conducts a rigorous comparative analysis of five distinct FL strategies using the MIMIC-IV dataset (version 3.1), a large-scale critical care database.

• Data Simulation: To simulate a realistic non-IID environment, the authors partitioned 466,351 hospital admissions across five distinct clinical care units (Emergency Department, Discharge Lounge, Labor & Delivery, Medicine/Cardiology, and Medicine). This partitioning creates extreme heterogeneity, with mortality rates ranging from 0.00% (Labor & Delivery, 1 death in 20,337 admissions) to 2.99% (Medicine).
• Feature Engineering: The model utilizes a feature set comprising demographics, admission-level diagnosis/procedure counts, and an 11-item laboratory panel (e.g., creatinine, lactate) averaged over the first 24 hours. Missing values are imputed with training-set medians, and missingness is explicitly indicated.
• Class Imbalance Handling: To address imbalance at the client level without centralizing data, the SMOTE-Tomek resampling technique is applied locally to each client's training set. A specific fallback mechanism (class-weighted loss) is used for the "Labor & Delivery" client, which contains only a single positive case, rendering SMOTE inapplicable.
• Model Architecture: All strategies utilize a shared Deep Neural Network (DNN) architecture (Multi-Layer Perceptron with two hidden layers of 128 and 64 neurons) to ensure that performance differences are attributable solely to the aggregation strategy.
• Evaluated Strategies:
  1. FedAvg: Standard weighted averaging.
  2. FedProx: Adds a proximal term to the local objective to constrain local updates from diverging too far from the global model.
  3. FedAdagrad & FedAdam: Server-side adaptive optimizers that maintain momentum and adaptive learning rates for the global model.
  4. FedCluster: Clusters client updates based on direction (using DBSCAN) to filter outliers before aggregation.
• Evaluation Metrics: Primary metrics are AUC-ROC and AUC-PR (threshold-independent), chosen for their reliability in highly imbalanced settings. Secondary metrics include F1-score, Precision, Recall, and Accuracy at a validation-calibrated threshold. Experiments run over 50 communication rounds with five random seeds.

Key Results

• Performance Hierarchy: FedProx emerged as the most robust strategy, achieving the highest mean AUC-ROC (0.897) and AUC-PR (0.230). Paired t-tests confirmed its AUC-ROC advantage is statistically significant against all other strategies.
• Strategy Nuances:
  • FedAdagrad showed a steady improvement in AUC-ROC (reaching 0.890) but suffered from poor probability calibration, resulting in the lowest precision (0.149) and F1 (0.232).
  • FedAdam exhibited high instability, with a sharp loss spike and AUC-ROC collapse around round 10-15 before recovering.
  • FedCluster achieved the highest F1-score (0.280) and precision (0.252) but lagged in AUC-ROC (0.873), suggesting it filters out informative but noisy updates.
  • FedAvg performed the worst among FL strategies (AUC-ROC 0.857), with its ranking quality degrading over time despite improving F1 scores.
• Centralized vs. Federated: A centralized baseline (naive, no explicit imbalance handling) significantly outperformed all federated strategies (AUC-ROC 0.929, AUC-PR 0.312). This indicates that while FL preserves privacy, it incurs a modest but real predictive performance cost compared to centralized training on the same data.
• Per-Client Heterogeneity: The global federated model did not serve all clients equally. AUC-ROC varied from 0.902 (Emergency Department) to 0.809 (Medicine), demonstrating that aggregation dominated by large clients can underperform on smaller, clinically distinct populations.
• Differential Privacy (DP): Applying per-example DP-SGD to FedProx preserved AUC-ROC (0.896) but reduced AUC-PR and F1 significantly (by ~15-22%) and increased training time by 4.1x.

Key Contributions

1. Comprehensive Benchmark: Provides the first head-to-head empirical comparison of baseline, regularized, adaptive, and clustering-based FL strategies on a large-scale, non-IID, and imbalanced clinical mortality prediction task.
2. Methodological Rigor: Demonstrates the critical importance of using threshold-independent metrics (AUC-ROC/AUC-PR) over fixed-threshold metrics (F1) in highly imbalanced clinical settings, revealing that strategies can improve in ranking quality while degrading in calibrated decision-making (and vice versa).
3. Per-Client Analysis: Highlights that aggregate metrics can mask significant performance disparities across different care units, emphasizing the need to evaluate FL systems on individual client distributions.
4. Privacy-Utility Trade-off: Quantifies the specific costs of implementing Differential Privacy in this context, showing that while ranking ability is preserved, operational metrics and computational efficiency suffer.

Significance and Claims
The paper claims that FedProx is the most robust choice among the evaluated aggregation strategies for heterogeneous and imbalanced clinical prediction tasks due to its ability to manage client drift. However, the authors modestly conclude that the decision to federate should be driven primarily by privacy and data-governance requirements, not by an expectation of superior predictive performance, as centralized training retains a measurable edge.

The study serves as a methodological cautionary note: constructing evaluation sets that preserve natural class distributions and selecting appropriate metrics are essential for accurately assessing FL strategies in healthcare. The work underscores that while FL is a viable solution for privacy-preserving collaboration, it introduces specific challenges (heterogeneity, imbalance, and aggregation costs) that require careful algorithmic selection and realistic expectations regarding performance parity with centralized models.