FedAvg across hospital nodes with non-IID patient populations — local SGD, weighted parameter averaging, and the convergence penalty of heterogeneous data
Cross-institutional federated learning begins by distributing a shared, agreed-upon model architecture to every participating hospital or research site, while patient-level data remains entirely within each institution's own infrastructure. This structural choice is what makes federated learning legally and organizationally viable for health data at all: institutional data-use agreements, IRB approvals, HIPAA business-associate constraints, and increasingly GDPR/EU data-residency rules make centralized pooling of raw multi-institutional patient data slow, costly, or outright infeasible — while sharing only model parameters sidesteps most of that friction.
The centralized-pooling alternative and why it usually fails in practice:
• Legal: multi-site IRB approval for a single pooled dataset can take 12–24 months per additional site; each institution's data-use agreement (DUA) must independently authorize the specific pooled-data use case • Regulatory: HIPAA business-associate agreements, GDPR data-residency requirements (especially for EU sites), and increasingly national health-data sovereignty rules restrict cross-border or cross-institution raw data transfer • Practical: transferring large imaging datasets (radiology, pathology whole-slide images) across institutions is bandwidth- and storage-prohibitive at scale, and creates a single centralized breach target
What federated learning changes: • Patient data physically never leaves the originating institution's firewall • Only model artifacts cross institutional boundaries: initial architecture (once), and iteratively, locally-updated parameters or gradients • Each institution retains full audit control over what computation runs against its data, and can withdraw from a training round without exposing any patient record
Setup requirements before round 1: • Common feature schema: all sites must map their local data to a shared representation (e.g., a common radiology preprocessing pipeline producing identically-sized image tensors, or a common tabular schema conformant to OMOP CDM — see the companion data-quality-scoring simulation) • Common model architecture and hyperparameters agreed upon by the consortium before training starts (architecture changes mid-training break weight-averaging compatibility) • A coordinating server (or in fully decentralized peer-to-peer schemes, a consensus protocol) that will perform aggregation — hosted by a neutral consortium body, a lead institution, or a cloud-based FL platform (e.g., NVIDIA FLARE, Owkin Substra) • Governance agreement on how many local epochs, how many total communication rounds, and what stopping/convergence criteria apply — negotiated once, applied uniformly across all sites for the whole run
In each communication round, every participating hospital initializes from the current global model parameters and runs several epochs of ordinary stochastic gradient descent on its own local cohort. Because no two hospitals see the same patient population — case-mix severity, scanner/equipment vendor, local coding conventions, demographic composition all differ — their locally-updated weights inevitably diverge from one another even though they started from the identical global starting point. This divergence, formalized later as "client drift," is the central technical challenge federated learning must manage.
Local training step, per site k, per communication round t:
1. Receive global parameters w_t from the server 2. Initialize local model: w_k ← w_t 3. Run E epochs of SGD over local data D_k: for epoch in range(E): for batch in D_k: w_k ← w_k − η · ∇L(w_k; batch) 4. Upload resulting local parameters w_k (or the delta Δw_k = w_k − w_t) back to the server
Why hospitals see genuinely different data distributions (non-IID sources in healthcare specifically):
• Case-mix / severity skew: a tertiary academic medical center systematically sees sicker, more complex patients than a community clinic — P(outcome | features) itself differs, not just feature marginals • Equipment/acquisition differences: imaging federated learning across sites with different MRI/CT scanner vendors, field strengths, and acquisition protocols introduces a covariate-shift component independent of the underlying disease biology • Coding practice variation: diagnosis coding thoroughness and specificity varies by institution, billing incentive structure, and EHR vendor — the same clinical presentation can map to different structured codes at different sites • Demographic skew: catchment-area demographics (age, race/ethnicity, socioeconomic status, comorbidity prevalence) differ systematically by geography and institution type • Label/prevalence skew: rare-disease specialty centers see a wildly different base rate of the target outcome than general hospitals — a classifier's optimal decision threshold locally can differ meaningfully from the population-wide optimum
Consequence for the local update: • Each site's local SGD trajectory pulls its weights toward that site's local loss minimum — which, under non-IID data, is a genuinely different point in parameter space than the global (all-sites-pooled) minimum • More local epochs E per round means each site drifts further toward its own local optimum before the next aggregation "pulls it back" — a direct tuning knob (exposed in this simulation's local-epochs control) trading off communication efficiency against divergence severity
FedAvg (McMahan et al., 2017) is the foundational federated learning algorithm: after each round of local training, every participating site's updated weights are combined at the coordinating server via a simple weighted average — weighted by each site's local dataset size — and the resulting averaged global model is broadcast back out to begin the next round. Despite its simplicity relative to more sophisticated federated optimization variants, FedAvg remains the default baseline against which every improvement is measured, and is what production platforms like MELLODDY and NVIDIA FLARE run by default.
FedAvg aggregation, per communication round t:
w_{t+1} = Σ_{k=1}^{K} (n_k / n) · w_k^{t+1}
where: • K = number of participating sites in this round • n_k = number of training examples at site k • n = Σ n_k, total examples across all participating sites • w_k^{t+1} = site k's locally-updated weights after its local epochs this round
Intuition: a site with more patient data gets proportionally more influence over the merged global model than a site with a smaller cohort — treating the federation as an approximation to training on the (hypothetical) union of all sites' data, without ever materializing that union.
Why simple averaging works reasonably well under IID data: • If all sites' local data were drawn from the same underlying distribution, each site's local SGD trajectory is (in expectation) an unbiased estimate of the direction centralized SGD would take on the pooled data — averaging across sites reduces variance similarly to how averaging gradients across a larger batch does in standard mini-batch SGD • This equivalence breaks down precisely under non-IID data (Stage 4), which is why FedAvg's convergence guarantees are notably weaker in the heterogeneous regime that real hospital consortia actually operate in
Secure aggregation (privacy hardening layer): • Basic FedAvg still exposes each site's individual weight update to the coordinating server — which can leak information about that site's local data distribution (and in adversarial settings, be used for gradient-inversion style reconstruction attacks) • Secure aggregation protocols (Bonawitz et al. 2017) use cryptographic additive masking or secure multi-party computation so the server can only ever see the SUM of all sites' updates, never any individual site's update in isolation — each site's mask cancels out only when all masks are combined • Production healthcare FL platforms (NVIDIA FLARE, Owkin Substra) support secure aggregation and/or differential-privacy-noised updates as an additional layer on top of the base FedAvg averaging step, particularly when the consortium includes competitively sensitive participants (e.g., pharma companies in MELLODDY)
The theoretical convergence guarantees for FedAvg were originally proven under an IID assumption — that every site's local data is drawn from the same distribution as every other site's. Real hospital consortia violate this assumption substantially, and the resulting "client drift" is the dominant practical obstacle in deployed cross-institutional federated learning: convergence is measurably slower, the final model can plateau at a worse loss than centralized training would reach, and more communication rounds are needed to close the gap.
The mechanics of client drift:
• Under non-IID data, each site k's local loss function L_k has a different minimum θ*_k than the global objective's minimum θ* • During local training, gradient steps pull site k's weights toward θ*_k, not θ* • After E local epochs, site k's weights have moved measurably toward its own local optimum — this "drift" is proportional to both the degree of distributional heterogeneity and the number of local epochs E • Averaging drifted weights across sites does not simply cancel the drift out — because the drift directions are systematically different (not just noisy in a zero-mean sense) across sites with different underlying distributions, the FedAvg average can land at a point worse than any individual site's local optimum, and worse than the true global optimum • More local epochs per round (larger E) means MORE drift accumulates before the next correcting aggregation step — creating a direct tradeoff: fewer, larger local steps reduce communication overhead but worsen drift; frequent, small local steps limit drift but increase communication cost
Quantifying the penalty (illustrative, consistent with published FL benchmarks): • Near-IID data: FedAvg typically reaches within 1–2% accuracy of a centralized-training ceiling within 20–30 communication rounds • Highly non-IID data (e.g., pathological label-skew splits, or realistic hospital case-mix splits): the same task can require 3–5× more rounds to reach a comparable point, and may plateau 2–8 percentage points below the centralized ceiling entirely
Algorithmic mitigations beyond vanilla FedAvg: • FedProx (Li et al. 2020): adds a proximal regularization term to the local objective, penalizing local weights from moving too far from the last global model — directly limits how far drift can accumulate per round • SCAFFOLD (Karimireddy et al. 2020): introduces control variates that explicitly estimate and correct for the difference between each client's update direction and the true global update direction, giving provably better convergence rates under heterogeneity • FedNova: normalizes each client's contribution by the number of local steps actually taken, correcting for a subtler bias FedAvg has when different sites run different numbers of local steps (common when sites have heterogeneous compute) • Personalization layers: rather than forcing one global model to fit all sites equally well, some deployments keep a shared global backbone but allow small site-specific "head" layers to adapt locally — accepting that a single global optimum may not exist for genuinely heterogeneous clinical populations
Client drift is not a bug to be eliminated but a structural consequence of genuine clinical heterogeneity across hospitals — the practical goal of algorithms like FedProx and SCAFFOLD is not to force convergence to a single "true" global optimum that may not meaningfully exist, but to reach a global model that generalizes reasonably well across sites while converging in a tractable number of communication rounds.
Beyond the statistical convergence challenge, cross-institutional federated learning faces a practical engineering constraint: hospitals and research sites often have limited, variable, or metered network bandwidth, and uploading full model weight updates every round does not scale to large models or many participants. Communication-efficient federated learning — gradient/weight compression via quantization and sparsification — addresses this directly, and has been proven at real production scale by consortia including MELLODDY in pharmaceutical drug discovery and NVIDIA Clara/FLARE in multi-hospital radiology.
Communication-efficient federated learning techniques:
1. Quantization: • Reduce numeric precision of transmitted weight updates from 32-bit floats to 8-bit integers (or lower) before upload • Stochastic quantization: rounds values probabilistically to preserve an unbiased expectation, avoiding systematic bias accumulation over many rounds • Typical compression: 4× reduction (32-bit → 8-bit) with minimal accuracy impact on most architectures
2. Sparsification (top-k): • Only the k largest-magnitude weight updates (by absolute value) are transmitted per round; the rest are zeroed and accumulated locally into the next round's update ("error feedback") so no information is permanently discarded, only delayed • Can achieve 10–100× reduction in transmitted parameters per round with careful error-feedback accounting
3. Structured/low-rank updates: • Constrain local updates to a low-rank factorization (analogous to LoRA-style efficient fine-tuning), transmitting only the low-rank factors rather than the full dense weight delta • Particularly effective for large architectures where most of the parameter count contributes little to the actual per-round update
4. Reduced communication frequency: • Increasing local epochs E per round directly trades communication frequency against convergence quality (see Stage 4) — a lever consortia tune based on each site's bandwidth constraints
MELLODDY (Machine Learning Ledger Orchestration for Drug Discovery): • IMI (Innovative Medicines Initiative) consortium of 10 major pharmaceutical companies (2019–2022) that federated-trained shared QSAR (quantitative structure-activity relationship) models across each company's proprietary, competitively sensitive compound-activity datasets • Scale: approximately 2.6 billion data points across ~20 million unique compounds — without any single company ever seeing another's proprietary compound library, addressing a purely competitive (not just regulatory) data-sharing barrier • Combined blockchain-based orchestration for auditability with federated learning for the actual model training, plus differential privacy safeguards against model-based data leakage between competitors
NVIDIA Clara / NVIDIA FLARE (Federated Learning Application Runtime Environment): • Open-source SDK explicitly designed for multi-institutional healthcare federated learning, widely used in academic and industry radiology/pathology consortia • Documented deployments span federated tumor segmentation, COVID-19 chest CT severity scoring across multiple international hospital networks, and breast-density classification models trained across geographically distributed sites without centralizing any patient imaging • Provides built-in support for secure aggregation, differential privacy noise injection, and several non-IID mitigation algorithms (including FedProx-style proximal terms) as configurable components rather than requiring each consortium to reimplement them from published papers