Skip to main content

Federated averaging, by hand

You now know that each participant uploads a list of arrays (their weights, shaped into layers) plus a single weighting. This page shows exactly what the most common way of combining them — federated averaging — does, worked out with real numbers, all the way down.

Federated averaging (often called FedAvg) is just a weighted average, an element-by-element mean. It is the default method on BranchKey and the one most federations use. It is not, however, the only way to aggregate: combining participants' updates is a general step, and different methods do it differently. This page unpacks averaging specifically, because once you can do it by hand the idea stops being mysterious.

The setup: three participants, one tiny model

Imagine three hospitals training the same small model on their own patient data. To keep the sums readable, our model has just two layers:

  • layer_0 — a vector of 3 weights
  • layer_1 — a single bias

Each hospital trains locally and ends up with slightly different numbers, because each saw different data. They also have different amounts of data, which we record as the weighting (here, the number of training samples).

Participantweighting (samples)layer_0layer_1
Hospital A100[0.2, 0.4, 0.6][1.0]
Hospital B300[0.4, 0.2, 0.0][2.0]
Hospital C100[0.6, 0.6, 0.3][6.0]

Each hospital uploads its own .npz (weighting + layer_0 + layer_1). Nobody sees anyone else's data — only these numbers.

Step 1 — the simple case: equal weighting

First, pretend every participant counts equally (as if each set weighting = 1). Then aggregation is the plain average: add up each position and divide by 3.

layer_0 — average each of the three positions independently:

position 0: (0.2 + 0.4 + 0.6) / 3 = 1.2 / 3 = 0.4
position 1: (0.4 + 0.2 + 0.6) / 3 = 1.2 / 3 = 0.4
position 2: (0.6 + 0.0 + 0.3) / 3 = 0.9 / 3 = 0.3

so layer_0 = [0.4, 0.4, 0.3].

layer_1 — one number to average:

(1.0 + 2.0 + 6.0) / 3 = 9.0 / 3 = 3.0

so layer_1 = [3.0].

Notice that averaging happens position by position, within each layer. This is exactly why everyone's layer_0 must have the same shape — you can only add position 0 to position 0 if they line up.

Step 2 — the real case: weighted by data size

Equal weighting isn't quite fair here: Hospital B trained on three times as much data as A or C. Its numbers are backed by more evidence, so it should pull the average a little harder. That's what the weighting field is for.

The weighted average is: multiply each participant's value by its weighting, add those up, then divide by the total weighting. In symbols, for each position:

result = (w_A·x_A + w_B·x_B + w_C·x_C) / (w_A + w_B + w_C)

The total weighting is 100 + 300 + 100 = 500. Now layer_0, position by position:

position 0: (100·0.2 + 300·0.4 + 100·0.6) / 500 = (20 + 120 + 60) / 500 = 200 / 500 = 0.40
position 1: (100·0.4 + 300·0.2 + 100·0.6) / 500 = (40 + 60 + 60) / 500 = 160 / 500 = 0.32
position 2: (100·0.6 + 300·0.0 + 100·0.3) / 500 = (60 + 0 + 30) / 500 = 90 / 500 = 0.18

so layer_0 = [0.40, 0.32, 0.18].

And layer_1:

(100·1.0 + 300·2.0 + 100·6.0) / 500 = (100 + 600 + 600) / 500 = 1300 / 500 = 2.6

so layer_1 = [2.6].

What the weighting actually did

Line the two results up:

layer_0layer_1
Equal weighting[0.4, 0.4, 0.3][3.0]
Weighted by samples[0.40, 0.32, 0.18][2.6]

Positions 1 and 2, and the bias, all moved towards Hospital B's numbers — because B carried three times the weight. Position 0 happened to land on 0.40 either way (the values balanced out). This is the entire effect of weighting: it tilts the shared result towards participants with more (or more trusted) data. Only relative sizes matter — 100, 300, 100 gives the same result as 1, 3, 1.

The same thing in NumPy

Here is the whole computation in code. This is essentially what BranchKey's average aggregation service runs on the server — you never write this yourself, but now you can see there's nothing hidden in it:

import numpy as np

# Each participant: (weighting, [layer_0, layer_1])
updates = [
(100, [np.array([0.2, 0.4, 0.6]), np.array([1.0])]),
(300, [np.array([0.4, 0.2, 0.0]), np.array([2.0])]),
(100, [np.array([0.6, 0.6, 0.3]), np.array([6.0])]),
]

total_weight = sum(w for w, _ in updates)
num_layers = len(updates[0][1])

aggregated = []
for layer_index in range(num_layers):
weighted_sum = sum(w * params[layer_index] for w, params in updates)
aggregated.append(weighted_sum / total_weight)

print(aggregated)
# [array([0.4 , 0.32, 0.18]), array([2.6])]

That result — layer_0 and layer_1, no weighting field — is exactly what gets packaged into the aggregated .npz and sent back to every participant to load into their model.

Why this is the whole idea of federated learning

Each hospital's numbers encode what its data taught the model. By averaging the numbers — and never the data — every hospital walks away with a model shaped by all the data, while each dataset stays exactly where it was. Repeat the round many times (train a bit more, upload, average, download) and the shared model steadily improves for everyone. That loop is federated learning; the weighted average you just did by hand is the most common way its combine step is performed.

For the narrative picture with diagrams, see How federated learning works. For the algorithms and settings a branch owner configures — including average (what you saw here) and the advanced maml variant — see Configuring a branch.