Federated Learning: Local Training on Device — A Complete Visual Guide

The Story That Explains Federated Learning

What Is Federated Learning & On-Device Training?

Federated Learning (FL) is a way to train one shared machine-learning model across many devices — phones, cars, hospitals, IoT sensors — without ever collecting their raw data in one place. The classic phrase is: bring the model to the data, not the data to the model.

In traditional machine learning, all data is copied to a central server and the model trains there. In federated learning, the model is shipped down to each device, trained locally on data that physically stays put, and only the resulting weight updates travel back. This on-device step — the part where each phone quietly improves the model using your own typing, photos, or sensor readings — is called local training on device, and it is the heart of this tutorial.

The Federated Training Round — Step by Step

Federated learning runs in repeated communication rounds. Each round has five stages. Stages 2 and 3 — highlighted below — are the on-device local training that this tutorial focuses on.

Inside a Single Device — What "Local Training" Actually Does

When the global model lands on a device, the device behaves like a tiny, self-contained training loop. It does not train forever — that would drain the battery and overfit to one user. Instead it runs a small, fixed amount of work: typically E local epochs over the device's data in mini-batches, using standard gradient descent.

The Math Behind It — Federated Averaging (FedAvg)

The algorithm that ties all the local updates together is FedAvg, introduced by Google in 2017. The intuition is simple: after local training, the server takes a weighted average of every device's model, where each device is weighted by how much data it holds.

Python Implementation — Local Training + FedAvg from Scratch

Below is a minimal, dependency-light simulation. It defines what happens on a single device (local_train), then the server loop that broadcasts, collects updates, and aggregates with FedAvg.

import numpy as np

# ── A tiny linear model: predict y from x with weight w and bias b ──
def local_train(global_w, global_b, X, y, epochs=3, lr=0.05):
    """Runs ON the device. Raw data X, y never leave this function."""
    w, b = global_w, global_b              # start from the shared global model
    n = len(X)
    for _ in range(epochs):           # a few LOCAL epochs only
        y_pred = w * X + b
        error  = y_pred - y
        grad_w = (2/n) * np.dot(error, X)
        grad_b = (2/n) * np.sum(error)
        w     -= lr * grad_w               # local SGD step
        b     -= lr * grad_b
    return w, b, n                       # send back weights + sample count ONLY


def fed_avg(updates):
    """Server side: weighted average by each client's data size."""
    total = sum(n for _, _, n in updates)
    new_w = sum(w * n for w, _, n in updates) / total
    new_b = sum(b * n for _, b, n in updates) / total
    return new_w, new_b


# ── Simulate 3 devices, each with its own private (non-IID) data ──
np.random.seed(0)
clients = [
    (np.random.rand(600),  # Device A — lots of data
     lambda x: 3.0*x + 1.0),
    (np.random.rand(300),  # Device B
     lambda x: 3.0*x + 1.0),
    (np.random.rand(100),  # Device C — little data
     lambda x: 3.0*x + 1.0),
]
client_data = [(X, f(X)) for X, f in clients]

# ── The federated training loop (server orchestration) ──
global_w, global_b = 0.0, 0.0            # start untrained
for rnd in range(20):                   # 20 communication rounds
    updates = []
    for X, y in client_data:           # BROADCAST + LOCAL TRAIN
        upd = local_train(global_w, global_b, X, y)
        updates.append(upd)            # only Δ-weights come back
    global_w, global_b = fed_avg(updates)   # AGGREGATE

print(f"Learned w={global_w:.3f}, b={global_b:.3f}")
print("True target was   w=3.000, b=1.000")

Notice the design: local_train receives X, y but never returns them — it returns only weights and a sample count. The server in fed_avg never touches a single raw data point, yet the model converges to the true target. That is on-device federated learning in 30 lines.

Where On-Device Federated Learning Is Used Today

The Hard Parts — Challenges of Local On-Device Training

Key Takeaways