Handling Non‑IID Clinical Data in Federated Learning: Algorithms, Pitfalls, and PyTorch Examples

Federated Learning (FL) is appealing in healthcare because it lets hospitals collaborate on models without pooling raw patient data. But the very reason FL is attractive, different sites, different patients, different workflow,ws creates the hardest technical problem: non‑IID data.

In clinical machine learning, “non‑IID” is not a corner case. It’s the default: patient populations vary by geography and referral patterns, measurement devices differ, and documentation practices are inconsistent.

If you’re an ML engineer, data scientist, or technical clinician trying to build models across hospitals, this matters right now. Regulatory pressure, privacy constraints, and the push to deploy models responsibly make “train centrally on everything” increasingly unrealistic.

This deep dive is for readers who already know basic ML and PyTorch and want serious depth. We’ll connect the optimization theory to what you see in real EHR pipelines: missingness, prevalence shifts, calibration failures, and site-to-site brittleness.

By the end, you’ll understand why vanilla FedAvg can become unstable on hospital data, and what “client drift” looks like in practice. You’ll also be able to implement FedAvg and FedProx in PyTorch, evaluate them with clinically meaningful metrics, and apply personalization strategies that make sense for real workflows.

Background and prerequisites

You’ll be most comfortable with this article if you can read PyTorch training loops and know what a loss function and optimizer do. Basic probability and linear algebra help, especially for understanding why gradients “fight” each other across heterogeneous sites.

You should also know the basics of supervised evaluation, including class imbalance. In healthcare, the difference between AUROC and AUPRC is not academic; it can change decisions about deployment readiness.

On the clinical side, you do not need to be a clinician. But you do need to internalize that EHR data are not “just tabular data,” they are a record of decisions, workflows, and incentives.

A lab value being missing often means something clinically meaningful happened. It might not mean the test was “unavailable”; it might mean the clinician didn’t order it because the patient looked stable.

Two hospitals can disagree about a label even when both are acting reasonably. A diagnosis code may be applied differently, an outcome may be recorded differently, or the “clock” for time-to-event may be operationalized differently.

On the technical side, federated learning is a distributed optimization problem with real system constraints. You’re not just training a model, you’re coordinating training across organizations with different infrastructure, schedules, and risk tolerances.

That interdisciplinary reality is why healthcare FL is hard and interesting.
You must design algorithms that converge under heterogeneity and systems that respect privacy, governance, and clinical safety.

What does non‑IID mean in clinical federated learning

In textbooks, i.i.d. data means each sample comes from the same distribution as every other sample. In multi-hospital settings, that assumption breaks in several overlapping ways, and each one matters.

The simplest form is covariate shift, where feature distributions differ across hospitals. One site might have older patients, different comorbidity profiles, or different measurement ranges due to device calibration.

A second common form is prevalence shift, where $P(y)$ differs across sites. For the same task, disease prevalence can be higher at a tertiary referral center than at a community hospital.

More subtle and often more damaging is concept shift, where $P(y \mid x)$ differs across sites. The same features can imply different outcomes because treatment protocols, triage policies, and resource constraints differ.

Clinical datasets also show missingness shift, which is rarely random. A lab might only be ordered when the clinician suspects deterioration, so “missing” becomes a proxy for “not concerning.”

Healthcare data can also include “schema shift” disguised as non‑IID.
Different coding systems, inconsistent units, and different feature availability mean the feature space itself is not identical across sites.

All of this connects back to the domain goal: patient care decisions happen locally. If your model performs “well on average” but fails at one hospital, you can cause harm even while reporting impressive aggregate metrics.

Core theory: the federated objective and why heterogeneity breaks FedAvg

Most cross-silo FL setups aim to minimize a weighted global objective.
With $K$ hospitals, hospital $k$ has $n_k$ samples and local loss $F_k(w)$.

$$\min_w F(w) = \sum_{k=1}^{K} p_k F_k(w), \quad p_k = \frac{n_k}{\sum_{j} n_j}$$

This looks like standard empirical risk minimization, except you can’t compute gradients on pooled data. Each hospital only sees its own distribution, and that’s where non‑IID becomes an optimization problem.

FedAvg: simple averaging that can hide unstable dynamics

In Federated Averaging (FedAvg), the server sends weights $w_t$ to clients each round. Each hospital performs several local training steps and returns updated weights $w_{t+1}^{(k)}$.

The server aggregates by a weighted average:

$$w_{t+1} = \sum_{k \in \mathcal{S}_t} \frac{n_k}{\sum_{j \in \mathcal{S}_t} n_j}\, w_{t+1}^{(k)}$$

When hospital data are similar, this can work surprisingly well. But in clinical FL, hospitals often optimize toward different local optima, especially with multiple local epochs.

Client drift: the failure mode you actually observe

Client drift happens when local training moves the model far toward a hospital’s local optimum. If you do more local epochs to save bandwidth, you often increase drift under non‑IID.

From a clinical perspective, drift corresponds to “local medicine” being encoded strongly in the weights. When you average many locally specialized models, you can end up with a global model that is not truly good anywhere.

This can show up as oscillating validation metrics across rounds, unstable calibration, or per-site performance collapse. It can also show up as a global model that favors large sites, because their gradients dominate the aggregation.

FedProx: stabilizing local training with a proximal anchor

FedProx modifies the client objective so local optimization stays closer to the global model. Hospital $k$ minimizes a proximal-regularized objective:

$$F_k(w) + \frac{\mu}{2}\lVert w - w_t \rVert^2$$

The parameter $\mu$ controls how strongly the model is pulled toward $w_t$.
In practice, $\mu$ is a knob you tune based on how heterogeneous your hospitals are.

Intuitively, FedProx says: “Learn from local data, but don’t run off too far.” That constraint often improves training stability when sites have different prevalence, missingness, or feature ranges.

FedProx does not “solve” non‑IID, but it makes the baseline more reliable. That reliability matters in healthcare, because unstable training translates into unstable clinical validation and harder governance.

Personalization: why a single global model is often not the right goal

Healthcare is local. Even if physiology is universal, care pathways, resources, and documentation patterns are not.

A global FL model is often best viewed as a shared representation. You want it to learn general relationships, like how vital signs correlate with deterioration risk.

But you also want site-specific adaptation to handle practice patterns. For example, one hospital might only order certain labs for severe patients, which changes what “normal” looks like in the dataset.

One practical personalization approach is local fine-tuning. You train a global model federatively, then adapt it locally using a small amount of on-site training with careful validation.

Another approach is shared backbone + local head.
Hospitals share the feature extractor but keep their own final layer, which maps features to a site-specific decision boundary.

This maps well onto clinical realities. The shared backbone captures general physiology, while local heads capture differences in prevalence, coding, and workflow thresholds.

A third approach is “statistics personalization,” commonly discussed as FedBN. BatchNorm running means and variances can be harmful to average across heterogeneous sites, so you keep them local.

Even in tabular models without BatchNorm, the idea is instructive.
Many failures come from mixing incompatible feature distributions, not from insufficient model capacity.

Hands-on implementation: FedAvg and FedProx on skewed multi-hospital tabular data

To make this concrete, we’ll build a realistic toy setup: multiple hospitals with EHR-like tabular features. We’ll simulate site-specific feature shifts, prevalence shifts, and missingness patterns, then train federated models in PyTorch.

This is not meant to be a perfect clinical simulator.
It’s meant to mirror the engineering reality: separate data loaders per site, local training loops, server aggregation, and per-site evaluation.

Why these choices match clinical practice

Tabular baselines remain common in clinical risk prediction because they are easier to audit and validate. They also map well to structured EHR features like vitals, labs, and prior utilization.

We’ll evaluate with AUROC and AUPRC because clinical outcomes are often imbalanced. We’ll also evaluate calibration using the Brier score, because predicted probabilities often drive resource allocation decisions.

Code: end-to-end FL skeleton with FedAvg and FedProx

import copy
import math
import numpy as np

import torch
import torch.nn as nn
from torch. utils.data import Dataset, DataLoader

from sklearn.metrics import roc_auc_score, average_precision_score, brier_score_loss


# -----------------------------
# 1) Synthetic multi-hospital data with non-IID patterns
# -----------------------------

def sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


def make_site_dataset(
    n: int,
    d: int,
    site_id: int,
    rng: np.random.Generator,
):
    """
    Simulate non-IID hospital data:
      - covariate shift (feature mean shifts by site)
      - prevalence shift (site-specific intercept changes)
      - missingness shift (site-specific missing rates per feature)
    We also add missingness indicators, which is a common clinical baseline trick.
    """

    # Covariate shift: measurement/population differences across sites
    mean_shift = rng.normal(loc=0.0, scale=0.8, size=d) * (site_id + 1) / 3.0
    X = rng.normal(size=(n, d)) + mean_shift

    # Shared "true" relationship + site-specific intercept = prevalence shift
    w_true = rng.normal(size=d) / math.sqrt(d)
    intercept = -0.6 + 0.6 * site_id  # later sites have higher baseline risk
    logits = X @ w_true + intercept
    p = sigmoid(logits)
    y = rng.binomial(n=1, p=p).astype(np.int64)

    # Missingness shift: ordering practices differ by site
    base_miss = 0.05 + 0.05 * site_id
    miss_probs = np.clip(rng.uniform(base_miss, base_miss + 0.15, size=d), 0.0, 0.6)
    missing_mask = rng.binomial(n=1, p=miss_probs, size=(n, d)).astype(bool)

    X_obs = X.copy()
    X_obs[missing_mask] = np.nan

    # Clinical-style baseline preprocessing:
    # - simple imputation (0)
    # - add missingness indicators as extra features
    miss_ind = np.isnan(X_obs).astype(np.float32)
    X_imp = np.nan_to_num(X_obs, nan=0.0).astype(np.float32)

    X_final = np.concatenate([X_imp, miss_ind], axis=1).astype(np.float32)
    return X_final, y


class TabularClinicalDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        self.X = torch.from_numpy(X)
        self.y = torch.from_numpy(y)

    def __len__(self):
        return self.X.shape[0]

    def __getitem__(self, idx: int):
        return self.X[idx], self.y[idx]


def make_federated_loaders(
    num_sites: int = 5,
    n_per_site: int = 4000,
    d: int = 20,
    batch_size: int = 128,
    seed: int = 42,
):
    """
    Create per-hospital train/val/test loaders.
    In real clinical FL, each hospital owns its split and validates locally.
    """
    rng = np.random.default_rng(seed)
    sites = []

    for k in range(num_sites):
        X, y = make_site_dataset(n=n_per_site, d=d, site_id=k, rng=rng)

        idx = rng.permutation(n_per_site)
        n_train = int(0.7 * n_per_site)
        n_val = int(0.1 * n_per_site)

        train_idx = idx[:n_train]
        val_idx = idx[n_train:n_train + n_val]
        test_idx = idx[n_train + n_val:]

        ds_train = TabularClinicalDataset(X[train_idx], y[train_idx])
        ds_val = TabularClinicalDataset(X[val_idx], y[val_idx])
        ds_test = TabularClinicalDataset(X[test_idx], y[test_idx])

        sites.append({
            "train": DataLoader(ds_train, batch_size=batch_size, shuffle=True),
            "val": DataLoader(ds_val, batch_size=batch_size, shuffle=False),
            "test": DataLoader(ds_test, batch_size=batch_size, shuffle=False),
            "n_train": len(ds_train),
        })

    input_dim = 2 * d  # original features + missingness indicators
    return sites, input_dim


# -----------------------------
# 2) Model: a simple clinical risk MLP
# -----------------------------

class RiskMLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=0.1),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)  # logits


# -----------------------------
# 3) Evaluation: AUROC, AUPRC, calibration proxy (Brier)
# -----------------------------

@torch.no_grad()
def predict_proba(model: nn.Module, loader: DataLoader, device: str):
    model.eval()
    probs, ys = [], []
    for x, y in loader:
        x = x.to(device)
        logits = model(x)
        p = torch.sigmoid(logits).cpu().numpy()
        probs.append(p)
        ys.append(y.numpy())
    return np.concatenate(probs), np.concatenate(ys)


def compute_metrics(probs: np.ndarray, y: np.ndarray):
    # Clinical note: Some sites can have extreme prevalence. Guard AUROC.
    auroc = float("nan")
    if len(np.unique(y)) == 2:
        auroc = float(roc_auc_score(y, probs))

    auprc = float(average_precision_score(y, probs))
    brier = float(brier_score_loss(y, probs))
    prevalence = float(y.mean())
    return {"auroc": auroc, "auprc": auprc, "brier": brier, "prevalence": prevalence}


def evaluate_per_site(model: nn.Module, sites, device: str, split: str = "test"):
    rows = []
    for k, site in enumerate(sites):
        probs, y = predict_proba(model, site[split], device)
        rows.append((k, compute_metrics(probs, y)))
    return rows


# -----------------------------
# 4) Local training: FedAvg (mu=0) and FedProx (mu>0)
# -----------------------------

def local_train(
    model: nn.Module,
    loader: DataLoader,
    device: str,
    epochs: int = 1,
    lr: float = 1e-3,
    mu: float = 0.0,
    global_params: dict | None = None,
):
    model = model.to(device)
    model.train()

    opt = torch.optim.Adam(model.parameters(), lr=lr)
    bce = nn.BCEWithLogitsLoss()

    if mu > 0.0:
        assert global_params is not None
        global_params = {k: v.to(device) for k, v in global_params.items()}

    for _ in range(epochs):
        for x, y in loader:
            x = x.to(device)
            y = y.float()to(device)

            logits = model(x)
            loss = bce(logits, y)

            # FedProx: proximal penalty discourages drifting far from global weights
            if mu > 0.0:
                prox = 0.0
                for name, p in model.named_parameters():
                    prox += torch.sum((p - global_params[name]) ** 2)
                loss = loss + (mu / 2.0) * prox

            opt.zero_grad()
            loss.backward()
            opt.step()

    return model


def weighted_average_state_dicts(state_dicts: list[dict], weights: list[float]):
    avg = copy.deepcopy(state_dicts[0])
    for key in avg.keys():
        avg[key] = avg[key] * 0.0
        for sd, w in zip(state_dicts, weights):
            avg[key] += sd[key] * w
    return avg


def federated_train(
    sites,
    input_dim: int,
    rounds: int = 20,
    clients_per_round: int | None = None,
    local_epochs: int = 2,
    lr: float = 1e-3,
    mu: float = 0.0,   # mu=0 => FedAvg, mu>0 => FedProx
    device: str = "cpu",
):
    global_model = RiskMLP(input_dim=input_dim).to(device)

    num_sites = len(sites)
    if clients_per_round is None:
        clients_per_round = num_sites

    for r in range(rounds):
        selected = np.random.choice(num_sites, size=clients_per_round, replace=False)

        global_state = copy.deepcopy(global_model.state_dict())
        client_states = []
        client_weights = []

        for k in selected:
            local_model = RiskMLP(input_dim=input_dim).to(device)
            local_model.load_state_dict(global_state)

            local_model = local_train(
                model=local_model,
                loader=sites[k]["train"],
                device=device,
                epochs=local_epochs,
                lr=lr,
                mu=mu,
                global_params=global_state if mu > 0.0 else None,
            )

            client_states.append(copy.deepcopy(local_model.state_dict()))
            client_weights.append(sites[k]["n_train"])

        total = float(sum(client_weights))
        w = [cw / total for cw in client_weights]
        new_global_state = weighted_average_state_dicts(client_states, w)
        global_model.load_state_dict(new_global_state)

        # Minimal monitoring: in practice, report per-site every N rounds
        if (r + 1) % 5 == 0:
            site0_probs, site0_y = predict_proba(global_model, sites[0]["val"], device)
            m = compute_metrics(site0_probs, site0_y)
            tag = "FedAvg" if mu == 0.0 else f"FedProx(mu={mu})"
            print(f"[{tag}] round={r+1} site0_val auprc={m['auprc']:.3f} brier={m['brier']:.3f}")

    return global_model


def personalize_by_finetuning(
    global_model: nn.Module,
    sites,
    input_dim: int,
    device: str,
    finetune_epochs: int = 2,
    lr: float = 5e-4,
):
    """
    Personalization step: start fromthe global model, fine-tune locally.
    In clinical deployments, this can improve per-site calibration and decision fit.
    """
    personalized = []
    global_state = copy.deepcopy(global_model.state_dict())

    for k in range(len(sites)):
        local_model = RiskMLP(input_dim=input_dim).to(device)
        local_model.load_state_dict(global_state)

        local_model = local_train(
            model=local_model,
            loader=sites[k]["train"],
            device=device,
            epochs=finetune_epochs,
            lr=lr,
            mu=0.0,
            global_params=None,
        )
        personalized.append(local_model)

    return personalized


# -----------------------------
# 5) Example run (uncomment in your environment)
# -----------------------------
# sites, input_dim = make_federated_loaders(num_sites=5, n_per_site=4000, d=20)
# device = "cuda" if torch.cuda.is_available() else "cpu"
#
# fedavg_model = federated_train(sites, input_dim, rounds=20, local_epochs=2, lr=1e-3, mu=0.0, device=device)
# fedprox_model = federated_train(sites, input_dim, rounds=20, local_epochs=2, lr=1e-3, mu=1e-2, device=device)
#
# print("FedAvg per-site test:", evaluate_per_site(fedavg_model, sites, device, split="test"))
# print("FedProx per-site test:", evaluate_per_site(fedprox_model, sites, device, split="test"))
#
# personalized = personalize_by_finetuning(fedprox_model, sites, input_dim, device, finetune_epochs=2, lr=5e-4)
# for k, pm in enumerate(personalized):
#     print(f"Personalized site {k} test:", evaluate_per_site(pm, [sites[k]], device, split="test"))

How to read results like a clinical ML practitioner

When you run this, do not stop at “global average performance.” Look atthe per-site AUPRC and Brier score and compare them to each site’s prevalence.

A model can have a decent AUROC while being practically unusable for alerting.
If prevalence is low, a small change in calibration can flood a workflow with false positives.

If FedAvg looks unstable across rounds, it’s not necessarily your code. It may be the algorithm amplifying heterogeneity because local training is too aggressive or sites are too different.

If FedProx improves stability but reduces adaptation, that’s expected. You may need fewer local epochs, a different $\mu$, or a personalization step after convergence.

Practical pitfalls: where clinical FL projects fail even with correct code

A common failure is assuming that federated learning automatically produces a clinically deployable model. FL changes the training topology, but it does not fix label ambiguity, measurement inconsistency, or missingness bias.

Another common failure is reporting only pooled metrics. In healthcare, the unit of deployment is often a hospital or clinic, so per-site evaluation is a safety requirement, not a “nice to have.”

Teams also underestimate how much “data harmonization” still happens in FL. Even if raw data never leaves the site, you still need consistent feature definitions and outcome definitions across sites.

Clinical heterogeneity can make hyperparameter tuning deceptively hard. A learning rate or local epoch count that works for one site can cause drift or overfitting at another site.

Finally, governance is a real constraint. Hospitals may require approval cycles for training runs, and “move fast and break things” is not an acceptable posture in patient-care settings.

Systems and production: what it takes to run FL across hospitals

In cross-silo healthcare FL, the hardest part is often not the model. It’s coordinating secure computation, consistent pipelines, and monitoring across organizations that do not share infrastructure.

Data pipelines: federated training still needs consistent semantics

Each hospital has its own ETL pipeline from EHR systems into analytics tables. Those pipelines differ in schema, missingness handling, unit conversions, and coding normalization.

Even with FL, you need a shared feature contract. That contract defines what each column means, how it is computed, and what time window it represents.

Time matters more than many teams expect. A readmission model trained with post-discharge data accidentally included will look great locally and fail when deployed.

A practical pattern is to standardize feature computation locally, then federate on the resulting feature matrix. This keeps protected health information on-site, while making training behavior across sites interpretable and auditable.

Infrastructure: secure, scheduled, and heterogeneous by default

Hospitals often run training on-prem or in tightly controlled cloud environments. That means networking constraints, strict identity management, and uneven access to GPUs.

Client availability is constrained by clinical operations. Training windows might be limited to nights or weekends, and failures must be handled gracefully.

Bandwidth and compute costs matter. If you send large models frequently, aggregation becomes slow and expensive, which can reduce stakeholder patience and governance support.

Observability: per-site monitoring is the minimum standard

Clinical FL needs monitoring at multiple levels. You need to know whether training is converging, whether any site is drifting, and whether performance is degrading over time.

Performance should be tracked per site and, when allowed, per clinically meaningful cohort. If a model is systematically worse for a hospital serving a different population, that is both a technical and an ethical issue.

Calibration monitoring is especially important in clinical risk prediction. A stable AUROC can hide a calibration collapse that changes how many patients are flagged at a fixed threshold.

In mature systems, sites compute metrics locally and share only aggregated summaries. That preserves privacy while still enabling centralized governance and troubleshooting.

Risk, ethics, safety, and governance in clinical federated learning

Federated learning reduces the need to centralize raw data, but it does not eliminate privacy risk. Model updates can leak information in certain threat models, and clinical data are high-stakes.

Privacy and security: FL is not automatically safe

An attacker can sometimes infer sensitive attributes from gradients or weight deltas. This risk is workload-dependent, but it must be considered when the training population is sensitive.

A realistic mitigation is secure aggregation, where the server only sees the sum of updates, not each site’s update. This reduces the risk of per-site update inspection and makes certain attacks harder.

Differential privacy can further reduce leakage by clipping updates and adding noise. In healthcare, the trade-off is serious: too much noise can reduce utility and harm downstream clinical decisions.

You also need operational security basics. Mutual authentication, encryption in transit, code signing, and auditable logs are part of a clinical-grade FL deployment.

Bias and representativeness: heterogeneity can become inequity

When hospitals serve different populations, performance disparities can mirror structural inequities. If a global model is dominated by a large site, smaller sites can be harmed even if they “participate.”

This is why per-site reporting is not enough by itself. You often need to examine performance relative to each site’s baseline and the clinical consequences of errors.

Clinical stakeholders also need guardrails against over-interpretation. A risk score is not a diagnosis, and a probability is not a certainty, even if the model is accurate on average.

Governance: decisions must be traceable and defensible

Healthcare organizations often operate under regulations and internal policies that require traceability. You need to document dataset versions, feature definitions, label definitions, and model versions for each training cycle.

A practical governance approach treats FL training like a software release process. You define acceptance criteria, run validation checks, review changes, and keep rollback options if monitoring detects harm.

Even if your jurisdiction differs, the principles are stable. Minimize data exposure, enforce access control, audit decisions, validate clinically, and monitor continuously.

Domain-specific scenario: federated readmission risk across hospitals

Consider a collaboration where multiple hospitals want to predict 30-day readmission risk at discharge. The goal is to allocate follow-up calls, care coordination, and home visits to patients who need them most.

Each hospital has structured EHR features like prior admissions, diagnoses, medications, and selected labs. The hospitals cannot pool raw data due to legal constraints, patient privacy expectations, and governance policies.

Non‑IID appears immediately. A tertiary center sees more complex cases and has a higher baseline readmission, while a community hospital may have different post-discharge resources.

Missingness patterns differ, too. One hospital orders labs routinely on discharge, while another orders selectively, making “missing labs” a meaningful signal.

A reasonable modeling path starts with FedAvg because it is simple and widely understood. If you observe unstable training or strong per-site disparities, you move to FedProx and reduce local epochs to mitigate drift.

You then evaluate per site with AUPRC and calibration measures. If some hospitals show miscalibration, youcan add personalization by local fine-tuning or by using local heads.

Finally, you align thresholds to the workflow. Hospitals may choose different alert thresholds because prevalence differs, and the acceptable false alert burden differs by staffing.

This scenario shows why the interdisciplinary lens matters. The “best” model is not the one with the best pooled metric, but the one that supports safe, consistent clinical decisions across sites.

Skills mapping and learning path

If you’re building a career in applied ML, this topic builds skills that map directly to real jobs. You’re learning how to make models work when data cannot be centralized and when evaluation is constrained by domain risk.

On the programming side, you gain comfort with modular training code. You learn how to structure datasets per site, implement aggregation, and keep evaluation reproducible and interpretable.

On the ML side, you develop intuition for optimization under distribution shift. You learn to diagnose client drift, tune local training schedules, and compare global versus personalized performance.

On the systems side, you learn to think about privacy and operations as first-class requirements.
Secure aggregation, metric reporting, monitoring, and governance become part of the design, not afterthoughts.

On the healthcare side, you learn what makes EHR modeling different from generic tabular ML. Missingness is informative, labels are operational definitions, and calibration matters because real decisions follow.

A practical next step is to extend the code in one targeted direction. For example, implement a shared backbone with local heads and compare per-site calibration before and after personalization.

Another strong next step is to treat monitoring as a feature. Add code that reports per-site metrics every few rounds and flags when a site’s calibration or AUPRC changes sharply.

Conclusion

Non‑IID is the default in clinical data, and it directly drives model brittleness across hospitals. FedAvg can struggle under heterogeneity because local training causes client drift and unstable aggregation.

FedProx stabilizes training by anchoring local optimization to the global model, often improving reliability. Personalization is frequently necessary in healthcare because workflows and measurement patterns are local.

Clinical-grade FL is as much about systems, monitoring, and governance as it is about algorithms. If you build per-site evaluation, calibration monitoring, and privacy-aware operations into the workflow, you can create models that support real clinical work without centralizing sensitive data.