Source-linked AI summary
Patient Clustering Improves Efficiency of Federated Machine Learning to predict mortality and hospital stay time using distributed Electronic Medical Records
Li Huang, Dianbo Liu
TL;DR
Centralized EMR learning is difficult because records are non-IID, distributed, and privacy-sensitive. CBFL clusters EMR data into communities and trains one model per community while keeping data local; it outperformed baseline FL across reported predictive metrics and communication cost, with accuracy close to centralized learning.
Problem
Federated learning may underperform on non-IID EMR data, while centralized EMR learning conflicts with records' distributed and privacy-sensitive nature.
Method
CBFL clusters EMR data into communities and simultaneously trains one model per community while keeping data local and aggregating locally computed results.
Results
CBFL outperformed baseline FL across ROC AUC, PR AUC, and communication cost, with predictive accuracy close to centralized learning.
Takeaways & Limitations
Patient communities captured clinically meaningful similarities, and community performance differences could be explained by dissimilarity between communities.
Takeaways & Limitations
Training K community models on each client transfers K−1 times more model parameters than training one model.
Abstract
from arXiv · showhide
Electronic medical records (EMRs) supports the development of machine learning algorithms for predicting disease incidence, patient response to treatment, and other healthcare events. But insofar most algorithms have been centralized, taking little account of the decentralized, non-identically independently distributed (non-IID), and privacy-sensitive characteristics of EMRs that can complicate data collection, sharing and learning. To address this challenge, we introduced a community-based federated machine learning (CBFL) algorithm and evaluated it on non-IID ICU EMRs. Our algorithm clustered the distributed data into clinically meaningful communities that captured similar diagnoses and geological locations, and learnt one model for each community. Throughout the learning process, the data was kept local on hospitals, while locally-computed results were aggregated on a server. Evaluation results show that CBFL outperformed the baseline FL algorithm in terms of Area Under the Receiver Operating Characteristic Curve (ROC AUC), Area Under the Precision-Recall Curve (PR AUC), and communication cost between hospitals and the server. Furthermore, communities' performance difference could be explained by how dissimilar one community was to others.
3. Medical School, Harvard University, 25 Shattuck St, Boston, MA 02115
The paper identifies distributed clustering, autoencoders, federated machine learning, and non-IID data as key topics.
- The keywords emphasize distributed clustering, autoencoders, federated machine learning, and non-IID data.
1. Introduction
Existing EMR machine-learning applications often assume centralized data sharing, although EMRs are distributed and privacy-sensitive. The paper proposes CBFL to address non-IID data by clustering EMRs and training one model per community.
- EMR machine-learning applications have often assumed that records could be shared across silos and stored centrally.
- EMRs are distributed and sensitive, raising concerns about security, cost-effectiveness, privacy, and data availability.
- Federated learning keeps data and computation local while aggregating computational results to train a global predictive model.
- FL may underperform on non-IID data, a characteristic commonly found in EMRs.
- CBFL clusters EMR data into communities and simultaneously trains one model per community to make learning more efficient than FL.
- The study develops and evaluates decentralized clustering with federated learning for predictions on ICU EMRs.
2. Materials and Methods
CBFL was developed and evaluated on distributed ICU EMRs, using privacy-preserving representations, hospital-level clustering, and community-specific federated models.
- The study used the eICU database and selected 50 hospitals with more than 600 patients, sampling 560 patients from each.
- The extracted data covered drugs administered during the first 48 ICU hours, discharge status, and ICU stay time.
- The final extracted dataset contained 126,490 examples.
- Each hospital trained a denoising autoencoder locally and returned only encoder weights for server-side averaging.
- Hospitals transformed data into representations, which the server clustered with k-means using K centroids representing communities.
- Each hospital trained K community models, which the server aggregated using weights based on community size.
3. Results
The cohort comprised 28,000 ICU patients from 50 hospitals, and clustering identified clinically and geographically coherent hospital and patient communities.
- Cohort analysis: 28,000 ICU patients from 50 hospitals formed the study cohort.
- Cohort analysis: 61.82% of patients were older than 60 years, while males and females represented 54.95% and 45.03% of the cohort, respectively.
- Cohort analysis: The mortality rate was 4.98%, and 6.12% experienced a prolonged ICU stay.
- Cohort analysis: A prolonged stay was defined as greater than or equal to eight days.
- Community analysis: Communities tended to contain patients with similar diagnoses and hospitals with similar geological locations.
- Community analysis: 50 hospitals clustered into five communities with recognizable separation on a PCA plot.
3.3. Mortality prediction
Mortality prediction used prescribed drug features and evaluated centralized learning, FL, and CBFL across community counts under same- and different-hospital splits. CBFL generally achieved higher predictive metrics with fewer communication rounds than FL, although more communities could reduce accuracy through overfitting.
- Same hospitals in training and test sets: CBFL consistently outperformed FL, reaching higher ROC AUCs with fewer communication rounds.In the same-hospital split, FL achieved ROC AUC 0.6895 and PR AUC 0.1107 in 101 rounds, while five-community CBFL achieved ROC AUC 0.6984 and PR AUC 0.1430 in 75 rounds.
- Same hospitals in training and test sets: 0.7368 ROC AUC and 0.1449 PR AUC from centralized learning exceeded both FL and CBFL in the same-hospital evaluation.This result was consistent with prior literature reported by the authors.
- Different hospitals in training and test sets: 0.6628 ROC AUC and 0.0912 PR AUC were obtained by 10-community CBFL in 27 rounds under different-hospital training and test distributions.FL reached ROC AUC 0.6520 and PR AUC 0.0871 in 66 rounds; more communities again reduced ROC AUC while using fewer rounds.
3.4. Stay time prediction
Stay-time prediction used prescribed drug features and the same evaluation metrics as mortality prediction. CBFL with five communities outperformed FL, while different training and test hospital distributions substantially reduced ROC AUCs and changed the comparison with centralized learning.
- Setup: Prolonged ICU stay time was predicted from prescribed drug features and assessed using ROC AUC, PR AUC, and communication rounds.Training and test data were formed from the same 50 hospitals for the initial evaluation.
- Same hospitals in training and test sets: More communities intensified overfitting in stay-time prediction, reducing performance as community count increased.In one evaluation, CBFL obtained ROC AUC 0.6400 and PR AUC 0.0822 in 23 rounds; increasing communities to 10 or 15 instead produced 31 rounds.
- Different hospitals in training and test sets: 0.7083 to 0.6189 ROC AUC was the reduction for centralized learning when training and test hospitals came from different distributions; FL fell from 0.6360 to 0.6212.CBFL also declined from approximately 0.63 to 0.62, despite faster convergence.
- Different hospitals in training and test sets: FL and CBFL outperformed centralized learning for the different-hospital stay-time task.The authors conjectured that federated regularization outweighed decentralization information loss in this setting.
3.5. Community distribution analysis
CBFL communities grouped patients with similar diagnoses and geographical locations, making separate community models easier to learn on average. Community performance differences were associated with how far each community was from the others.
- Community structure: Communities tended to contain patients with similar diagnoses and geological locations, making their individual models easier to learn on average.This contrasts with learning one model across all patients.
- Community performance: Community 1 achieved the highest ROC AUC of 0.7561, while Community 4 achieved the highest PR AUC of 0.2155.These results came from five-community CBFL mortality prediction.
- Community performance: Community 2 was farthest from the others, with an average centroid distance of 2.562.Average distance to other community centroids on a PCA plot was used to explain performance differences.
- Community performance: Community performance was compared using ROC AUC, PR AUC, and average distance to other communities.The reported analysis links differences among community models with their relative centroid distances.
4. Discussion
CBFL clustered ICU EMR data using admission diagnoses and hospital locations, then learned separate community models that achieved higher predictive accuracy with fewer communication rounds than baseline FL. The approach also preserved patient-data privacy, while its communication burden and performance portability across datasets remained important limitations.
- Community-based learning: CBFL clustered patients by admission diagnoses and hospital locations, learning separate models for individual communities.The paper describes these dimensions as the primary clustering focus.
- Predictive performance: CBFL converged to higher predictive accuracy in fewer communication rounds than baseline FL for mortality and stay-time prediction.This pattern held across both prediction tasks.
- Interpretability: Clustering made prediction results interpretable because distances between communities helped explain differences in prediction performance.The discussion connects community dissimilarity with variation in results.
- Privacy: CBFL kept patient data local and avoided transmitting patient data, unlike cited federated-learning optimizers that required sharing a fraction of client data.Encoded data sent for clustering was difficult to use for recovering original drug features because client decoders were discarded.
- Limitations: With K community models, CBFL transferred K−1 times more model parameters than FL, increasing communication load as training samples and communities grew.The paper identifies this additional communication as a limitation of the approach.
- Limitations and future work: CBFL performed best with five or at most ten communities in these experiments, but fewer communities may not be optimal on other biomedical datasets.The authors identify future work on reducing communication load and enriching clustering beyond diagnoses and locations.
5. Conclusions
CBFL addressed non-IID ICU EMR learning by clustering patients into clinically meaningful communities and optimizing a model for each community. It outperformed baseline FL across prediction tasks and metrics, while community performance differences were associated with inter-community dissimilarity.
- Approach: CBFL clustered non-IID ICU patients into clinically meaningful communities and optimized mortality and ICU stay-time prediction.The model was evaluated against baseline FL using ROC AUC, PR AUC, and communication rounds.
- Performance: CBFL achieved predictive accuracy close to centralized learning, although neither FL nor CBFL performed better than centralized learning.This positions CBFL as competitive with centralized learning while retaining the federated setting.
- Performance: CBFL outperformed FL on ROC AUC, PR AUC, and communication cost in both mortality and stay-time prediction tasks.The comparison held with or without matching training and test data distributions.
- Community structure: CBFL communities contained different overrepresented diagnoses and accommodated hospitals from diverse geographical locations.Community composition therefore reflected both clinical and geographic variation.
- Community structure: Performance differences between communities could be attributed to Euclidean distances on the PCA plot.Greater inter-community dissimilarity was linked to differences in community performance.
- Scope: The study focused on ICU EMRs, while CBFL could be extended to medical image recognition and medical planning across healthcare silos.These applications involve large, distributed, and privacy-sensitive data.