DEV Community

Cover image for Federated learning in 40 lines: training a model without ever moving the data
Spicy
Spicy

Posted on

Federated learning in 40 lines: training a model without ever moving the data

#machinelearning #privacy #ai #python

The usual way to train a model on sensitive data is to pile it all into one place first. Dump every user's records into a central store, train, done. That central pile is also a single, searchable honeypot, which is exactly why security teams have started pushing back on it.

Federated learning flips the idea: instead of bringing the data to the model, you send the model to the data. Each device or site trains on its own local data and sends back only the resulting weights. The raw data never moves.

That sounds abstract until you see it in code. Here's a complete, runnable example of the core algorithm (FedAvg) in plain NumPy.

The whole thing 

import numpy as np

# Each client has its OWN data. This never leaves the client.
def make_client(n, true_w, noise=0.1):
    X = np.random.randn(n, len(true_w))
    y = X @ true_w + noise * np.random.randn(n)
    return X, y

true_w = np.array([2.0, -3.0, 0.5])
clients = [make_client(100, true_w) for _ in range(5)]

def local_train(X, y, w, lr=0.05, steps=20):
    # One client runs gradient descent on ITS data only.
    for _ in range(steps):
        grad = X.T @ (X @ w - y) / len(y)
        w = w - lr * grad
    return w

# Federated averaging
global_w = np.zeros(3)
for rnd in range(10):
    local_weights = []
    for X, y in clients:
        # Server sends global_w out; client trains locally;
        # only the trained weights come back.
        w = local_train(X, y, global_w.copy())
        local_weights.append(w)
    # Server averages the weights. It never touched any client's X or y.
    global_w = np.mean(local_weights, axis=0)
    print(f"round {rnd+1}: w = {np.round(global_w, 3)}")

print("true w:", true_w)
Enter fullscreen mode Exit fullscreen mode

Run it and you'll watch global_w converge toward [2, -3, 0.5] without the server ever seeing a single row of X or y.

The one line that is the whole point

This one:

w = local_train(X, y, global_w.copy())
local_weights.append(w)
Enter fullscreen mode Exit fullscreen mode

The client gets the current global model, trains it on data that stays on the client, and returns only the weights. The server averages those weights into a new global model and sends it back out. Repeat. At no point does raw data cross the network. That's the entire privacy pitch in two lines.

Where the real world gets harder

This toy version hides three things you'd deal with in production:

  • Weights can leak too. Gradients and weight updates can be reverse-engineered to reveal training data. The fix is differential privacy: add calibrated noise to the updates so no single record can be recovered.
  • Clients disagree. Real clients have different data distributions (non-IID), so naive averaging can stall or wobble. Weighting by dataset size and smarter aggregation help.
  • Stragglers and scale. Thousands of phones that drop mid-round need handling. That's what frameworks like Flower, TensorFlow Federated, and NVIDIA FLARE exist for.

But the mental model is exactly the 40 lines above: move the model, not the data, and only ever share what you learned, never what you saw.

I wrote up why this is quietly replacing "shove everything into a vector DB" for sensitive enterprise data, and what that means for data sovereignty, here: The Great Data Recall: Why Federated Learning Beats Vector DB Leakage

Anyone running FL in production yet? Curious what you're using for the differential-privacy layer.

Top comments (0)