Model Aggregation in Federated Learning: FedAvg Guide

The Story That Explains Model Aggregation

Federated Learning lets many devices or organizations train one shared model without ever centralizing their raw data. The clients do the learning; the server does the merging. This tutorial zooms into that merge — the single most important operation on the server side.

What Actually Happens on the Server

In one federated round, the server runs a tight loop: it broadcasts the current global model, each client trains locally on its private data, every client uploads only its model update, and the server aggregates those updates into the next global model. Watch the data flow below — notice that only model weights move, never raw data.

FedAvg — The Aggregation Formula

The standard server-side aggregation is Federated Averaging (FedAvg), introduced by McMahan et al. It is a weighted average of every client's model, where each client's weight is the fraction of total data it holds. After clients train locally, the server computes one new global model for the next round.

A Worked Numerical Example

Let's aggregate a single weight value from our three hospitals. Each trained the model locally and arrived at a different value for one parameter. (Real models have millions of these — the math is applied element-by-element.)

Compare that to a plain average: (0.80 + 0.50 + 0.20) / 3 = 0.50. FedAvg lands at 0.65 instead, pulled toward Client A's value — correctly, because Client A backed its number with six times more data than Client C. That 0.15 difference is the data-weighting at work.

Aggregation Strategies at a Glance

FedAvg is the default, but the server can blend updates in several ways depending on whether you care about data heterogeneity, privacy, or defending against malicious clients.

Implementing Server Aggregation in Python

Here is FedAvg as the server would run it. Each client returns its trained weights plus how many samples it trained on. The server blends them into one global model.

import numpy as np

# Each client sends back: (model_weights, num_samples)
# Here we use the 3-hospital values from Section 04.
w_A, w_B, w_C = np.array([0.80]), np.array([0.50]), np.array([0.20])

client_updates = [
    (w_A, 600),
    (w_B, 300),
    (w_C, 100),
]

def federated_average(updates):
    """Weighted average of client weights by sample count (FedAvg)."""
    total_samples = sum(n for _, n in updates)

    # accumulator shaped exactly like the model weights
    global_w = np.zeros_like(updates[0][0])

    for weights, n in updates:
        alpha = n / total_samples      # this client's data share
        global_w += alpha * weights   # weighted contribution

    return global_w

new_global = federated_average(client_updates)
print("Aggregated global weight:", new_global)

Real neural networks store weights as a state_dict — a dictionary of named tensors (one per layer). The server simply applies the same weighted average key-by-key across every layer:

def aggregate_state_dicts(client_states, client_sizes):
    """FedAvg over PyTorch-style state_dicts (layer by layer)."""
    total = sum(client_sizes)
    avg_state = {}

    # every client shares the same layer names (same architecture)
    for key in client_states[0].keys():
        avg_state[key] = sum(
            (client_sizes[i] / total) * client_states[i][key]
            for i in range(len(client_states))
        )

    return avg_state   # load this back into the global model

Challenges the Server Must Handle