Protocol Choices Matter More Than the Aggregation Algorithm: Federated Learning for In-Hospital Mortality Prediction in Care-Unit and Multi-Hospital Federations

AI summary

Overview

Research area: Clinical machine learning and federated learning — specifically, benchmarking aggregation algorithms for in-hospital mortality prediction on electronic health records.

Technical level: Intermediate. The paper motivates everything in plain clinical-ML terms, but reading it comfortably assumes familiarity with AUC-ROC/AUC-PR, federated averaging, hyperparameter tuning and calibration.

Scope: A controlled, seed-repeated comparison of seven federated aggregation strategies for mortality prediction across two federations (12 care units in MIMIC-IV and 49 hospitals in eICU-CRD), designed to separate the effect of the aggregation algorithm from the effect of evaluation-protocol choices.

What This Paper Is About

Hospitals cannot pool patient records, so federated learning is used to train a shared mortality model from data that stays on site — but the server-side algorithm that combines the local models is chosen from published comparisons that were never designed to isolate that algorithm's effect. This paper rebuilds the comparison under a single protocol: one pipeline, one feature set, per-strategy tuning on validation data, the same model-selection rule for every method, ten random seeds, and correction for multiple comparisons. It then deliberately perturbs common protocol choices to measure how much they, rather than the choice of algorithm, move the final numbers.

Key Contributions

1. A controlled comparison of seven aggregation strategies (FedAvg, FedProx, SCAFFOLD, FedAdagrad, FedAdam, FedYogi, FedCluster) on an intra-hospital federation (12 MIMIC-IV care units) and an inter-hospital federation (49 eICU-CRD hospitals) built with one shared pipeline, so the effect of federation structure is not confounded with data preparation.
2. Quantification of four protocol choices — discharge-coded predictors, client-level SMOTE oversampling, evaluating the last instead of the best validation round, and hyperparameter tuning — showing that three of them change discrimination by more than the choice of strategy does in the care-unit federation.
3. Identification of SCAFFOLD (absent from prior clinical comparisons) as the top strategy in both federations, plus a dedicated control, FedAvg with tuned local SGD, that attributes 57%–63% of SCAFFOLD's AUC-ROC margin over FedAvg to its control variates rather than to its change of local optimizer.
4. Evidence that federated training matches centralized training on the multi-hospital federation, and that Platt scaling fitted on validation data repairs the risk distortions caused by imbalance corrections without losing discrimination.

Main Findings

• SCAFFOLD ranked first in both federations: AUC-ROC 0.886 ± 0.002 on MIMIC-IV care units and 0.819 ± 0.004 on eICU-CRD hospitals, first in all ten seeds (Friedman χ² = 44.2 and 50.6, p < 0.001), and first in 17 of 18 experiment variants.
• The winning margins were small: SCAFFOLD's lead over the runner-up never exceeded 0.008; on MIMIC-IV it ranged from +0.008 (over FedAdagrad) to +0.019 (over FedCluster), on eICU-CRD from 0.004 to 0.009 in AUC-ROC and 0.011 to 0.014 in AUC-PR. The full best-to-worst spread was 0.019 in AUC-ROC for MIMIC-IV and 0.009 for eICU-CRD.
• The control variates, not the optimizer, do most of the work: FedAvg with the same tuned local SGD (FedAvg-SGD) reached 0.878 versus 0.874 for FedAvg on MIMIC-IV and 0.816 versus 0.813 on eICU-CRD, but SCAFFOLD still beat FedAvg-SGD in all ten seeds of both databases (+0.008, 95% CI 0.006 to 0.010 on MIMIC-IV; +0.004, 0.002 to 0.005 on eICU-CRD; Wilcoxon p = 0.002).
• Federated matched centralized on the multi-hospital federation: SCAFFOLD versus centralized training without imbalance correction on eICU-CRD differed by +0.001 in AUC-ROC (95% CI 0.000 to +0.003, p = 0.11) and +0.001 in AUC-PR (−0.002 to +0.005, p = 0.56), and parity held after adding day-1 physiology.
• Federated fell short on the care-unit federation: SCAFFOLD trailed the uncorrected centralized model by 0.0015 in AUC-ROC (−0.002 to −0.001) and 0.011 in AUC-PR (−0.015 to −0.007) on MIMIC-IV, but did not differ from the centralized model with class weights, the treatment it shares (+0.0005 in AUC-ROC, −0.002 in AUC-PR, p > 0.2).
• Protocol beat algorithm: leakage from discharge-coded diagnosis and procedure counts raised AUC-ROC from 0.888 to 0.925 (+0.038) in a centralized gradient-boosting reference; client-level SMOTE cost up to −0.074 in AUC-ROC; evaluating the last round instead of the best validation round cost up to −0.084.
• Global SMOTE was the weakest centralized baseline: 0.838 ± 0.016 AUC-ROC on MIMIC-IV and 0.769 ± 0.009 on eICU-CRD, well below both the uncorrected (0.888, 0.818) and class-weighted (0.886, 0.817) centralized models.
• Calibration can be repaired cheaply: Platt scaling fitted on validation data corrected the distorted risk estimates caused by class weighting and resampling, with Brier scores such as 0.0194 for SCAFFOLD on MIMIC-IV and 0.0729 on eICU-CRD, without loss of discrimination; isotonic regression was evaluated as an alternative.
• SCAFFOLD requires its own local optimizer: run with the Adam optimizer used by the other strategies, the per-coordinate normalization turned its constant correction into a fixed drift and the model collapsed within three rounds on MIMIC-IV, where the Emergency Department performs more than 1,500 local steps per round.

Methodology in Plain English

The authors took two public PhysioNet databases and turned each into a federation. In MIMIC-IV, the "clients" were care units of one hospital; in eICU-CRD, they were independent hospitals. A site became a client only if it recorded at least 100 in-hospital deaths, which yielded 12 care units covering 87% of stays and 49 hospitals. Records were split at the patient level into 70% training, 10% validation and 20% test, repeated over ten seeds (42 to 51), so every strategy faced exactly the same split and test set within a seed.

All clients trained the same small neural network (two hidden layers of 128 and 64 units) for 50 communication rounds of three local epochs each. Only information available in the first 24 hours was used as predictors — demographics, insurance or ICU type, and an identical panel of 11 laboratory values — because diagnosis codes are assigned at discharge and would leak outcome information. Missing labs were median-imputed with missingness flags.

Each strategy's own hyperparameters were tuned by grid search on validation data (from a single configuration for FedAvg to 12 configurations on MIMIC-IV and 14 on eICU-CRD for SCAFFOLD), and every model was selected by the same rule: the round with the best validation AUC-PR. Comparisons used a Friedman test followed by Holm-corrected paired tests against the best-ranked strategy, with the seed as the unit of replication. A purpose-built control, FedAvg with tuned local SGD, separated the effect of SCAFFOLD's drift correction from the effect of its optimizer. Centralized models trained on pooled data served as the ceiling. Finally, the authors re-ran the pipeline while changing one protocol element at a time to see how much each one moved the answer.

Why This Matters

Impact on research. Rankings of federated algorithms are only as trustworthy as the protocol behind them. This paper shows that in a realistic clinical federation, a leaky feature, an oversampling step, or the choice of which round to report can each shift AUC-ROC by more than the entire spread between the seven algorithms. It also documents specific weaknesses in prior comparisons — untuned shared hyperparameters, mismatched selection rules between federated and centralized models, single runs, and no calibration reporting — and offers a protocol that fixes them.

Real-world applications.

• Hospital consortia choosing an aggregation strategy can treat FedAvg-family methods as near-equivalent and spend effort on protocol hygiene instead.
• Clinical deployment of risk scores requires calibrated probabilities; the paper shows validation-fitted Platt scaling recovers usable risks after class weighting or SMOTE.
• Multi-hospital research networks can use the eICU-CRD result — federated ≈ centralized — as evidence that pooling is not required for this task.
• Teams that simulate federations by splitting one center into care units now have a measured caveat: that setting behaves differently from a true multi-hospital federation.

Industry relevance. Federated learning vendors and health-system data platforms can use the finding that protocol choices dominate algorithm choices as a quality argument for auditing evaluation pipelines, and the SCAFFOLD result constrains engineering: the method only works with plain SGD locally, not with the Adam optimizer most production stacks default to.

Future Directions

• Establish whether the protocol-dominates-algorithm conclusion holds for outcomes beyond in-hospital mortality, such as acute kidney injury, sepsis or length of stay, which appear in the compared literature but not in this study.
• Determine which of the four protocol choices is the one that exceeds the algorithm spread in both federations; the reported design measures them separately but the generalizability across federation types remains open.
• Extend the comparison to drift-correction methods deliberately excluded here — FedNova and FedDyn — and to personalized federated learning, which this study only probes as one epoch of local fine-tuning.
• Work out how the validation-based steps (hyperparameter tuning, round selection, Platt scaling) would be carried out in a real deployment where a pooled validation set does not exist, for example through secure aggregation of client summaries.

Target Audience

Clinical machine learning researchers and health-informatics groups who benchmark prediction models on MIMIC-IV or eICU-CRD; federated learning practitioners moving from benchmarks to hospital deployments; and methodologists and reviewers interested in evaluation protocol, leakage, imbalance handling and calibration reporting in multi-site prediction studies.

Authors’ abstract

Federated learning allows hospitals to train a shared mortality model without pooling their records, but published comparisons of aggregation strategies rarely control the evaluation protocol. This study compared seven strategies (FedAvg, FedProx, SCAFFOLD, FedAdagrad, FedAdam, FedYogi and FedCluster) on two federations built with one pipeline: 12 care units of a hospital in MIMIC-IV and 49 hospitals in eICU-CRD. Each strategy was tuned and selected on validation data by the same rule, and every comparison was repeated across 10 seeds with Friedman and Holm-corrected tests. SCAFFOLD ranked first in both federations (AUC-ROC 0.886 and 0.819) and in 17 of 18 variants, but its margin over the runner-up never exceeded 0.008, and a control sharing its local optimizer attributed most of its margin over FedAvg to drift correction. On the multi-hospital federation, federated training matched centralized training. In the care-unit federation, three protocol choices changed discrimination by more than the choice of strategy did: discharge-coded predictors (+0.038 in AUC-ROC), client-level SMOTE (up to -0.074), and evaluating the last instead of the best validation round (up to -0.084). Platt scaling on validation data repaired the risks distorted by imbalance corrections without loss of discrimination. Studies should therefore fix the protocol before ranking algorithms.