Federated Learning Setup

"A global model, visiting each hospital in turn, learning what it could without ever taking a single chart home."

A Coordinator That Never Saw a Patient

The previous sections of this chapter established what the system must learn and the constraints under which it may learn it: a shared clinical risk model, trained across several hospitals, with the absolute requirement that raw records never leave the institution that holds them. That constraint is what forces distribution here. It is not a throughput ceiling or a model-size ceiling of the kind that drove data parallelism in Chapter 1; it is a governance ceiling. The data physically cannot be co-located, so the computation must come to the data. Federated learning is the discipline that makes this work, and the cross-silo setting (a small number of organizational data silos rather than millions of phones) is the variant that fits a hospital consortium. We now turn the abstract FedAvg of Chapter 14 into a deployable configuration.

1. The Federated Round, Concretely Beginner

Federated training proceeds in synchronized rounds, each one a complete cycle of broadcast, local work, upload, and aggregation. The coordinator holds the authoritative global model $w^{(t)}$. At the start of round $t$ it broadcasts that model to every participating hospital. Each hospital $k$ then trains locally on its own records for $E$ epochs, producing an updated local model $w_k^{(t+1)}$ that it never shares as data, only as parameters. The hospitals upload those parameter vectors (or, equivalently, the deltas $w_k^{(t+1)} - w^{(t)}$) to the coordinator, which combines them into the next global model $w^{(t+1)}$ and begins the following round. Figure 37.4.1 traces one such round across the consortium.

The aggregation rule is federated averaging, FedAvg, exactly as introduced in Chapter 14. With $K$ hospitals, where hospital $k$ holds $n_k$ records and $n = \sum_k n_k$ is the consortium total, the new global model is the size-weighted mean of the local models,

The size weighting matters because hospitals differ enormously in patient volume; a large academic center should count for more than a small rural clinic, in proportion to the evidence each contributes. This weighted combine is a reduce in the exact sense of Chapter 4: each worker holds a vector, and the system folds those vectors into one with an associative, commutative operation. The inner update, $E$ epochs of gradient steps on the local objective $F_k$, is the local SGD of Chapter 10, run inside each hospital between communications.

2. What "Cross-Silo" Changes Beginner

The federated literature is dominated by the cross-device setting: millions of phones, each with a few records, most of them offline at any moment, joined by flaky cellular links. A hospital consortium is the opposite regime on almost every axis, and recognizing this lets us simplify the design considerably. There are few participants, perhaps five to fifty institutions, and they are known and trusted parties bound by a data-use agreement rather than anonymous clients. Every site participates in every round, so there is no client sampling and no need to reason about which subset showed up. Connectivity is reliable, since these are data-center-grade links between known endpoints, not consumer radios. Each site holds a large local dataset, thousands to millions of records, which makes many local epochs both feasible and statistically meaningful. Because each round already does substantial work and the links are good, the system converges in tens of rounds rather than the thousands a cross-device deployment may need.

These properties compound favorably. Full participation removes the variance that client sampling injects into cross-device FedAvg. Reliable links mean a straggler is a rare degraded node, not the norm, so the straggler mitigation of Chapter 18 is a safety net rather than a load-bearing mechanism. Large local datasets mean the dominant statistical problem is not data scarcity per client but data heterogeneity across clients, because each hospital's patient population, equipment, and coding practices differ. That heterogeneity is the central technical risk of this case study, and Section 37.5 is devoted to it; here it shapes one decision, the choice of server optimizer, which we take up next.

3. Choosing the Server Optimizer Intermediate

FedAvg is the baseline, and on identically distributed data it is hard to beat. The complication in a hospital consortium is that the data is not identically distributed: site $k$ minimizes its own local objective $F_k$, and when those objectives disagree, many local epochs pull each $w_k$ toward a different local minimum. The weighted average of models that have drifted apart can land in a poor place, an effect called client drift that slows or destabilizes convergence as heterogeneity grows. Two well-established remedies address this, and they intervene at different ends of the round.

FedProx changes the local objective. Each hospital adds a proximal term that penalizes drifting too far from the broadcast global model $w^{(t)}$, replacing $F_k(w)$ with

The coefficient $\mu$ tunes a leash: $\mu = 0$ recovers FedAvg exactly, while larger $\mu$ keeps local models close to the global one, trading a little local progress per round for stability under skew. This is a minimal, one-line change to the client and nothing on the server, which is why it is the natural first defense against heterogeneity. We adopt it conditionally and revisit the tuning of $\mu$ in Section 37.5.

FedAdam changes the server instead. Rather than averaging models directly, the coordinator treats the aggregated client delta $\Delta^{(t)} = \sum_k \tfrac{n_k}{n}\,(w_k^{(t+1)} - w^{(t)})$ as a pseudo-gradient and applies an adaptive optimizer to it,

Server momentum and per-coordinate scaling smooth out noisy rounds and accelerate convergence, which is valuable precisely because cross-silo training affords only a few rounds. FedProx and FedAdam are not mutually exclusive; the proximal term tames the clients while the adaptive server optimizer stabilizes the aggregate, and serious deployments often combine them. For this case study we begin with FedAvg, keep FedProx ready as the heterogeneity remedy, and reach for FedAdam when round budget is tight. Table 37.4.1 summarizes the choice.

4. Local Epochs Versus Communication Intermediate

The single most consequential knob is $E$, the number of local epochs each hospital runs before uploading. It sets the ratio of computation to communication, and the two effects pull in opposite directions. Suppose the consortium needs a total of roughly $S$ local gradient passes to reach the target accuracy. The number of communication rounds is then approximately

so doubling $E$ roughly halves the number of rounds, and therefore halves the number of times the (slow, firewalled, audited) cross-institution links must carry a full model. Communication is the expensive resource in this deployment, which argues for large $E$. The opposing force is heterogeneity: with non-identical site data, each extra local epoch lets $w_k$ drift further toward its own local optimum, so very large $E$ degrades the quality of the averaged model, and beyond a point the loss per round stops improving or worsens. The convergence theory of Chapter 14 makes this dependence explicit; informally, the optimality gap after $R$ rounds behaves like

where $\sigma^2$ measures how much the site objectives disagree. The first term rewards large $E$ (it buys more rounds-worth of progress per communication); the second punishes it, scaled by the heterogeneity $\sigma^2$. The right $E$ is therefore not a constant but a function of how skewed the consortium's data is, which is exactly why $E$ and the FedProx coefficient $\mu$ are tuned together in Section 37.5. For the homogeneous-to-mild regime of the demo below, a modest $E$ in the single digits sits comfortably in the sweet spot.

5. Model Architecture and Privacy Wrapping Intermediate

The architecture must be agreed before the first round, because every site trains the identical model and the coordinator averages parameters position by position; a mismatch in shape makes the reduce meaningless. This constrains the choice toward a single, fixed, consortium-wide architecture rather than per-site customization. For structured electronic-health-record features (labs, vitals, demographics, coded diagnoses) a compact model such as regularized logistic regression or a small multilayer perceptron is the right starting point: it is cheap to communicate, since fewer parameters mean smaller uploads, it is data-efficient enough to train well even at the smaller sites, and its behavior is auditable, which matters for clinical sign-off. For imaging tasks the same federated round wraps a convolutional or transformer backbone, at the cost of much larger parameter uploads, which raises the value of the large-$E$ communication savings of Section 4. The federated machinery is indifferent to the architecture; what changes with model size is only the bytes moved per round.

Critically, the parameter upload is not sent in the clear. Each hospital's update is wrapped by the privacy mechanisms developed in Section 37.3: secure aggregation so the coordinator learns only the summed update and never any single hospital's contribution, and optionally differential-privacy noise added to each update before it leaves the firewall. These protections sit between the local-train and upload steps of every round; the FedAvg arithmetic above is unchanged, but each $w_k^{(t+1)}$ that reaches the aggregator has been encrypted or noised first. The privacy layer and the federated layer are orthogonal, which is what lets this section configure the training dynamics while Section 37.3 owns the confidentiality guarantees.

6. The Runnable Federated Loop Intermediate

The code below builds the whole round structure on synthetic patient cohorts so the dynamics are visible end to end. Five hospitals each hold a different number of records drawn with a different covariate shift, modeling the heterogeneity of real sites, while sharing the same underlying risk model. Each round broadcasts the global model, runs $E$ local epochs of logistic-regression gradient descent at every site, and combines the results with the size-weighted FedAvg reduce. The script runs both plain FedAvg and FedProx ($\mu = 0.1$) and compares the federated result against a centralized model trained on the pooled data, the privacy-violating baseline the consortium is forbidden from building.

import numpy as np

rng = np.random.default_rng(7)
H, d = 5, 8                       # hospitals, patient features
w_star = rng.standard_normal(d)   # the shared clinical risk model
b_star = 0.3

def make_hospital(n, shift):
    # Site distributions differ (covariate shift); the risk model is shared.
    X = rng.standard_normal((n, d)) + shift
    p = 1.0 / (1.0 + np.exp(-(X @ w_star + b_star)))
    y = (rng.random(n) < p).astype(np.float64)
    return X, y

sizes  = [4000, 2500, 3500, 1500, 5000]          # heterogeneous volumes
shifts = [0.0, 0.8, -0.6, 1.2, -1.0]             # heterogeneous distributions
sites  = [make_hospital(n, s) for n, s in zip(sizes, shifts)]
n_total = sum(sizes)

def sigmoid(z): return 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30)))

def local_loss(W, X, y):
    p = sigmoid(X @ W[:-1] + W[-1]); eps = 1e-9
    return -np.mean(y*np.log(p+eps) + (1-y)*np.log(1-p+eps))

def global_loss(W):
    return sum(local_loss(W, X, y)*len(y) for X, y in sites) / n_total

def local_train(W_global, X, y, epochs, lr, mu=0.0):
    # E local epochs of GD (Ch 10 local SGD); mu>0 is the FedProx proximal pull.
    W = W_global.copy(); n = len(y)
    Xb = np.hstack([X, np.ones((n, 1))])
    for _ in range(epochs):
        grad = Xb.T @ (sigmoid(Xb @ W) - y) / n + mu * (W - W_global)
        W -= lr * grad
    return W

def run_federated(rounds, epochs, lr, mu=0.0):
    W = np.zeros(d + 1); history = []
    for _ in range(rounds):
        updates = [local_train(W, X, y, epochs, lr, mu) for X, y in sites]  # broadcast + local train
        W = sum(u*len(y) for u, (X, y) in zip(updates, sites)) / n_total    # FedAvg reduce
        history.append(global_loss(W))
    return W, history

ROUNDS, EPOCHS, LR = 15, 5, 0.5
print("hospitals H            :", H)
print("total patients         :", n_total)
print("rounds x local epochs  :", f"{ROUNDS} x {EPOCHS}\n")

_, hist_fa = run_federated(ROUNDS, EPOCHS, LR, mu=0.0)
_, hist_fp = run_federated(ROUNDS, EPOCHS, LR, mu=0.1)
print("round   FedAvg loss   FedProx loss")
for r in [0, 1, 2, 4, 7, 11, 14]:
    print(f"{r+1:>4}      {hist_fa[r]:.4f}        {hist_fp[r]:.4f}")

# Centralized reference: pool everything on one machine (forbidden in practice).
Xc = np.vstack([X for X, y in sites]); yc = np.concatenate([y for X, y in sites])
Wc = np.zeros(d + 1); Xcb = np.hstack([Xc, np.ones((len(yc), 1))])
for _ in range(400):
    Wc -= LR * (Xcb.T @ (sigmoid(Xcb @ Wc) - yc) / len(yc))
print("\nfinal FedAvg loss      :", f"{hist_fa[-1]:.4f}")
print("final FedProx loss     :", f"{hist_fp[-1]:.4f}")
print("centralized (pooled)   :", f"{global_loss(Wc):.4f}")
print("FedAvg vs pooled gap   :", f"{hist_fa[-1] - global_loss(Wc):.4f}")

hospitals H            : 5
total patients         : 16500
rounds x local epochs  : 15 x 5

round   FedAvg loss   FedProx loss
   1      0.5393        0.5490
   2      0.4937        0.5013
   3      0.4736        0.4792
   5      0.4570        0.4600
   8      0.4494        0.4507
  12      0.4467        0.4472
  15      0.4462        0.4464

final FedAvg loss      : 0.4462
final FedProx loss     : 0.4464
centralized (pooled)   : 0.4458
FedAvg vs pooled gap   : 0.0004

# pip install flwr ; one coordinator process + one client process per hospital
import flwr as fl

class HospitalClient(fl.client.NumPyClient):
    def get_parameters(self, config):           # send local weights up
        return get_model_weights()
    def fit(self, parameters, config):          # E local epochs behind the firewall
        set_model_weights(parameters)
        train_locally(epochs=config["local_epochs"])
        return get_model_weights(), num_local_records(), {}

# On each hospital node, behind its own firewall:
fl.client.start_client(server_address="coordinator:8080",
                       client=HospitalClient().to_client())

# On the coordinator (FedAvg is the default strategy; FedAdam is one line away):
fl.server.start_server(
    config=fl.server.ServerConfig(num_rounds=15),
    strategy=fl.server.strategy.FedAvg(min_fit_clients=5, min_available_clients=5),
)

7. Coordinator and On-Premise Clients Advanced

The deployment topology follows the algorithm directly. A single coordinator process, hosted in a neutral environment that no single hospital controls (a shared cloud account or a consortium-operated server), holds the global model and orchestrates rounds; it is the parameter-server pattern of Chapter 11, specialized to the case where its clients are whole institutions. Inside each hospital, a client process runs on-premise, with read access to that hospital's training data and outbound connectivity to the coordinator but no inbound exposure: the hospital firewall stays closed, the client dials out, and only model parameters cross the boundary. This dial-out design is what makes the system deployable in real clinical IT, where opening an inbound port into a hospital network is effectively impossible.

Because the parties are few, known, and trusted, the coordinator's role is orchestration rather than defense against adversarial clients; the Byzantine-robust aggregation of Chapter 35 is available if the threat model later demands it, but a trusted consortium does not start there. NVIDIA FLARE targets exactly this cross-silo, on-premise, institution-as-client deployment and adds production concerns Flower leaves to the operator: client authentication, audit logging of every round for regulatory review, and admin controls for starting and stopping jobs across sites. The training algorithm is identical to Code 37.4.1; what these frameworks contribute is the operational shell that lets a real hospital run a client behind its real firewall.

We now have the training system fully specified: synchronized rounds of broadcast, $E$ local epochs, and a size-weighted reduce; FedAvg as the baseline with FedProx and FedAdam held ready; a fixed shared architecture whose uploads are privacy-wrapped per Section 37.3; and a coordinator-with-dial-out-clients topology that respects every hospital firewall. The one risk this configuration has only deferred, not solved, is the disagreement between site data distributions, which the demo's mild skew let us ignore. Section 37.5 confronts that heterogeneity head on and shows why it is the force that turns FedProx from an option into a necessity.