🤖 AI

A Robust Pipeline for Differentially Private Federated Learning on Imbalanced Clinical Data using SMOTETomek and FedProx

This paper presents a robust framework for differentially private federated learning on imbalanced clinical data that combines client-level SMOTETomek resampling and an optimized FedProx algorithm to achieve high recall (>77%) while maintaining strong privacy guarantees (epsilon 9.0) for cardiovascular risk prediction.

Original authors: Rodrigo Tertulino

Published 2026-07-27

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Rodrigo Tertulino

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where hospitals could team up to build a super-smart doctor's assistant, one that learns from millions of patients to spot diseases early. The problem is, patient records are like precious, secret treasures locked in separate vaults. Laws like GDPR and HIPAA mean these vaults can't be opened and merged into one giant pile; the data must stay where it lives. Enter Federated Learning, a clever trick where the "brain" of the computer travels to the vaults instead of the data traveling to the brain. Each hospital trains a piece of the model on its own local data and only sends back the lessons learned, not the secrets.

But there's a catch. To make sure no one can reverse-engineer those lessons to steal a specific patient's identity, scientists add a layer of Differential Privacy. Think of this as adding a little bit of "static" or "noise" to the lessons before they are sent. It's like whispering a secret to a friend while someone else is shouting nearby; the friend hears the message, but the eavesdropper can't be sure exactly what was said. The big question is: how much static can we add before the message becomes garbled and useless? This is the "privacy-utility trade-off." If we add too much noise to be super safe, the model might get so confused it stops working. If we add too little, the secrets might be at risk. This is especially tricky when the data is "imbalanced," meaning there are thousands of healthy patients but very few sick ones, making it easy for a model to get lazy and just guess "everyone is healthy."

This paper tackles that exact mess. The researchers tried to build a Federated Learning system to predict cardiovascular risks (like strokes) using data that was both private and heavily imbalanced. They found that if you just use the standard methods, the system fails spectacularly: it becomes so confused by the lack of sick patients that it predicts "no one is sick" for everyone, achieving zero ability to catch actual cases. To fix this, they built a robust pipeline with two main upgrades. First, they used a technique called SMOTETomek at each hospital to artificially create more examples of the rare sick patients, balancing the scales before training. Second, they swapped the standard training algorithm for FedProx, which acts like a gentle guide to keep the local models from drifting too far apart when the data looks different at each hospital.

The result is a working system that successfully navigates the privacy-utility tightrope. The authors measured how well the model worked at different levels of privacy "noise." They discovered a "sweet spot" where the privacy budget (a number called $\epsilon$ ) is around 9.0. At this level, the model maintains strong privacy guarantees while still being clinically useful. In their tests, the model successfully identified about 77% of the actual at-risk patients (Recall) and achieved a discriminative power score (ROC-AUC) of roughly 0.82. The paper explicitly rules out the idea that standard Federated Learning works out of the box for this type of messy, real-world data, showing that without their specific balancing and stabilization steps, the model collapses. While the privacy budget of 9.0 is more relaxed than the ultra-tight numbers often discussed in pure math theory, the authors argue it is a practical, safe, and effective operating point for real healthcare applications, proving that you can have both security and a model that actually saves lives.

Technical Summary: A Robust Pipeline for Differentially Private Federated Learning on Imbalanced Clinical Data

Problem Formulation
The paper addresses the critical intersection of three challenges in healthcare machine learning: the need for collaborative model training across data silos, the presence of severe class imbalance in clinical datasets, and the requirement for formal privacy guarantees. Specifically, the authors investigate the prediction of cardiovascular events (focusing on stroke risk) using Federated Learning (FL). They highlight that standard FL approaches, when combined with Differential Privacy (DP), often fail in realistic clinical settings characterized by non-IID (statistically heterogeneous) data and extreme class imbalance (approximately 5% positive cases). Initial experiments demonstrated that applying standard DP-FL to such data resulted in a model with misleadingly high accuracy but zero clinical Recall, rendering it useless for screening high-risk patients.

Methodology
To overcome these limitations, the authors propose a multi-stage, robust pipeline implemented using the Flower framework and PyTorch. The methodology involves three core components:

Client-Side Data Balancing (SMOTETomek): Before local training, each client applies the hybrid Synthetic Minority Over-sampling Technique with Tomek Links (SMOTETomek) to their private, imbalanced dataset. This step is performed locally to prevent data leakage between institutions, generating synthetic minority samples to balance the class distribution (targeting a 30% minority ratio) without sharing raw data.
Robust Aggregation (FedProx): To address statistical heterogeneity and "client drift" caused by non-IID data distributions, the standard FedAvg algorithm is replaced with FedProx. This algorithm introduces a proximal term ( $\frac{\mu}{2} \|w - w_t\|^2$ ) to the local objective function, penalizing local updates that deviate significantly from the global model, thereby stabilizing convergence.
Differential Privacy (DP-SGD): The pipeline integrates Differentially Private Stochastic Gradient Descent (DP-SGD) via the Opacus library. This involves per-sample gradient clipping and the addition of Gaussian noise to model updates before transmission to the server. The privacy budget ( $\epsilon$ ) is tracked using a Rényi Differential Privacy (RDP) accountant.

The experimental design involves a systematic parameter sweep across noise multipliers ( $\sigma$ ), clipping norms ( $C$ ), and proximal hyperparameters ( $\mu$ ) to map the privacy-utility frontier. The study utilizes the publicly available "Stroke Prediction Dataset" (approx. 5,110 instances) partitioned among 10 clients.

Key Contributions
The paper outlines three primary contributions:

Identification of a Critical Failure Mode: The authors demonstrate that standard DP-FL applied to severely imbalanced clinical data results in a model with zero Recall, serving as a cautionary baseline for practitioners.
A Validated Methodological Pipeline: The paper proposes and validates a pipeline combining client-side SMOTETomek and FedProx. This approach successfully transforms a failing model into one with significant predictive power, effectively handling both class imbalance and statistical heterogeneity.
Granular Privacy-Utility Analysis: Through robust statistical analysis (aggregated over five independent runs), the study identifies an optimal operational region. It quantifies the trade-off, showing that a single-digit privacy budget can yield clinically viable results.

Results
The experimental results reveal a clear, non-linear trade-off between privacy and utility. The optimized pipeline achieved the following performance metrics at a privacy budget of $\epsilon \approx 9.0$ :

Recall: Approximately 77% (Mean), with a 95% confidence interval of [0.6374, 0.9026]. This indicates the model successfully identifies high-risk patients, addressing the primary clinical concern of false negatives.
ROC-AUC: Approximately 0.82 (Mean), demonstrating strong discriminative power independent of the decision threshold.
F1-Score: Approximately 0.76, indicating a balanced performance between precision and recall.
Stability: The FedProx algorithm demonstrated superior stability compared to FedAvg, maintaining consistent performance across varying privacy budgets and reducing variance caused by skewed data distributions.

The study notes that while $\epsilon \approx 9.0$ is more relaxed than the tight budgets ( $\epsilon < 1.0$ ) often cited in theoretical DP literature, it represents a pragmatic operating point that offers mathematical protection against reconstruction and membership inference attacks while preserving the high sensitivity required for effective disease screening.

Significance and Claims
The authors claim that this work provides a "practical methodological blueprint" for creating effective, secure, and accurate diagnostic tools for real-world, heterogeneous healthcare data. The significance lies in demonstrating that Privacy-Enhancing Technologies (PETs) cannot operate in isolation; they must be integrated into a data-centric pipeline that actively addresses local data sparsity and statistical heterogeneity. The paper asserts that it is possible to maintain high clinical sensitivity and precision simultaneously under formal DP guarantees, refuting the assumption that privacy noise inevitably renders models imprecise on imbalanced data. The work emphasizes that achieving a clinically viable model requires a holistic approach that balances privacy constraints with rigorous data balancing and robust aggregation strategies.

Technical Summary: A Robust Pipeline for Differentially Private Federated Learning on Imbalanced Clinical Data

More like this