Federated Learning

Federated learning was a large part of my bachelor thesis. I was interested in training a shared model when the data stays with the people or devices that collected it. Building an implementation for the thesis also led me to study unreliable clients, malicious updates, and privacy.

What federated learning is

A centralized training setup brings the training data together. In federated learning, clients train locally and send model updates to a server, which combines them into a shared model. The raw data stays on the clients, although the updates can still reveal information about it.

For a dataset of \(N\) examples, a centralized training objective can be written as:

\[\min_{w \in \mathbb{R}^d} F(w) = \frac{1}{N} \sum_{i=1}^N \ell(x_{i}, y_{i}; w)\]

When the data is divided among clients, we can express that objective as a weighted sum of local losses:

\[F(w) = \sum_{k=1}^K \frac{n_{k}}{n_{\text{total}}} F_{k}(w)\]

where each client \(k\) holds its own local loss \(F_k\) over \(n_k\) data points.

A training round follows these steps:

1. The server sends the current model to all clients.
2. Each client trains locally on its own data for a bit.
3. Clients send their updates back, never their data.
4. The server aggregates the updates into a new global model.

Repeating these rounds lets the model learn from several clients without collecting their raw training examples on the server.

Aggregation methods

FedSGD

In FedSGD, clients compute gradients at the current global weights and send them to the server. The server combines the gradients and takes an update step:

\[w^{t+1} = w^t - \eta \sum_{k=1}^K \frac{n_{k}}{n_{\text{total}}} \nabla F_{k}(w^t)\]

This requires communication for every global update, which can be expensive when clients have slow or unreliable connections.

FedAvg

FedAvg allows each client to take several local training steps before sending its model back. The server then averages the resulting weights:

\[w^{t+1} = \sum_{k=1}^K \frac{n_{k}}{n_{\text{total}}} w_{k}^t\]

Taking more steps locally reduces communication, but clients may move in different directions if their data distributions differ. The number of local steps becomes an important part of tuning the training process.

Robust aggregation

A client may send an incorrect update because of a software bug or an attempt to damage the model. A few extreme updates can pull an arithmetic mean far from the other clients’ results.

Geometric median

The geometric median is less sensitive to extreme updates. It finds a point that minimizes the weighted sum of distances to the client updates:

\[\min_{z} \sum_{k=1}^m \alpha_{k} \|w_{k} - z\|\]

A smoothed form of the Weiszfeld algorithm approximates it iteratively:

\[z^{(i+1)} = \frac{\sum_{k=1}^m \beta_{k}^{(i)} w_{k}}{\sum_{k=1}^m \beta_{k}^{(i)}}, \quad \beta_{k}^{(i)} = \frac{\alpha_{k}}{\max\{\nu, \|w_{k} - z^{(i)}\|\}}\]

Updates farther from the current estimate receive less weight in the next iteration. This reduces their influence, although it doesn’t make every kind of malicious update harmless.

Differential privacy

Keeping raw data on a device doesn’t by itself guarantee privacy. A server may be able to infer information about training examples from the updates it receives. Differential privacy gives us a way to bound how much the output can reveal.

The definition

A mechanism \(\mathcal{M}\) is \((\epsilon, \delta)\)-differentially private if, for any two datasets \(D\) and \(D'\) that differ by one record:

\[\Pr(\mathcal{M}(D) \in S) \leq e^\epsilon \Pr(\mathcal{M}(D') \in S) + \delta\]

For the definition above, changing one record has a bounded effect on the distribution of outputs. Smaller \(\epsilon\) means a stronger bound when the other conditions are held fixed. The noise level, sampling procedure, and number of training rounds all affect the overall guarantee.

Adding DP to FedAvg

The example here applies clipping and noise to each client’s model update. These are the two operations, written for an update vector \(g_k\):

1. Clip the update to a maximum norm: \[\tilde{g}_{k} = \frac{g_{k}}{\max(1, \|g_{k}\|_{2} / C)}\]

2. Add Gaussian noise: \[\hat{g}_k = \tilde{g}_k + \mathcal{N}(0, \sigma^2 C^2 \mathbf{I})\]

The clipping and noise operations can be written in NumPy:

1
2
3
4
5
6
7
8
9

import numpy as np

def dp_clip_and_noise(grad: np.ndarray, clip_norm: float = 1.0, noise_multiplier: float = 1.0) -> np.ndarray:
    norm = np.linalg.norm(grad)
    if norm > clip_norm:
        grad = (grad / norm) * clip_norm
    noise_std = noise_multiplier * clip_norm
    noise = np.random.normal(loc=0.0, scale=noise_std, size=grad.shape)
    return grad + noise

This function illustrates the mechanism, but it doesn’t calculate a privacy guarantee. That requires choosing what is protected, calibrating the noise, and accounting for repeated rounds. The TensorFlow Federated tutorial works through training with user-level differential privacy and privacy accounting.

Homomorphic encryption

Homomorphic encryption addresses a different question: can the server combine updates without reading each one? Clients encrypt their updates, and the server performs the aggregation on ciphertexts.

How it works

For an additive homomorphic scheme, an operation on ciphertexts corresponds to addition of the underlying values. We can write that relationship schematically as:

\[\text{Enc}(a) \oplus \text{Enc}(b) = \text{Enc}(a + b)\]

For federated learning the workflow becomes:

1. Each client encrypts its update: \(c_k = \text{Enc}(w_k)\).
2. The server sums the ciphertexts: \(c_{\text{sum}} = \sum_{k=1}^K c_k\).
3. Whoever holds the decryption key recovers \(\sum_k w_k\).

With an abstract encryption backend, summing encrypted updates might look like this:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

from typing import Any, Protocol

class HEContext(Protocol):
    def encrypt_zeros(self) -> Any: ...
    def homomorphic_add(self, a: Any, b: Any) -> Any: ...
    def decrypt(self, ciphertext: Any) -> Any: ...

def homomorphic_aggregate(encrypted_updates: list[Any], he_context: HEContext) -> Any:
    c_sum = he_context.encrypt_zeros()
    for c_u in encrypted_updates:
        c_sum = he_context.homomorphic_add(c_sum, c_u)
    return c_sum

def decrypt_sum(c_sum: Any, he_context: HEContext) -> Any:
    return he_context.decrypt(c_sum)

Putting it all together

The following example combines local training, clipping, noise, and averaging. It is a small demonstration of those operations, without the privacy accounting needed for a complete private training system:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

import numpy as np

def local_train(model_params, local_data, epochs: int = 1, lr: float = 0.01):
    w = np.copy(model_params)
    for _ in range(epochs):
        grad = compute_gradient(w, local_data)
        w = w - lr * grad
    return w

def compute_gradient(w, data):
    X, y = data
    preds = X @ w
    errs = preds - y
    return (X.T @ errs) / len(X)

def federated_round(global_params, clients_data, clip_norm: float, noise_multiplier: float):
    updated_params = []
    for local_data in clients_data:
        w_k = local_train(global_params, local_data)
        grad_k = w_k - global_params
        grad_k_dp = dp_clip_and_noise(grad_k, clip_norm, noise_multiplier)
        updated_params.append(global_params + grad_k_dp)
    return np.mean(updated_params, axis=0)

if __name__ == "__main__":
    w_global = np.zeros(2)
    
    clients_data = [
        (np.array([[1, 2], [0, 1]]), np.array([1.0, 0.0])),
        (np.array([[3, 2], [4, 1]]), np.array([2.0, 3.0])),
        (np.array([[10, 2], [8, 2]]), np.array([4.0, 5.0])),
    ]
    
    for t in range(5):
        w_global = federated_round(w_global, clients_data, clip_norm=2.0, noise_multiplier=0.5)
        print(f"Round {t}, global params = {w_global}")

Robust aggregation, differential privacy, and homomorphic encryption address different threats. Combining them takes care: an aggregation method that needs to inspect individual updates may be difficult to apply when those updates are encrypted.

What broke when I built this

The implementation brought up several issues that were easy to overlook in the equations.

Non-IID data was the biggest difficulty. Clients don’t all collect examples from the same distribution, so training on one client’s data can move the model in a different direction from training on another’s.

Client dropout meant the server couldn’t assume everyone would finish a round. Phones go offline and apps close, so the aggregation process needed a policy for incomplete participation.

Privacy and accuracy needed to be considered together. Noise changes the updates used for training, while encryption adds computation and communication costs. Choosing those settings affected both the protection and the usefulness of the model.

What stuck with me

Federated learning gave me a useful example of how an optimization method changes when it has to run across separate devices. The averaging step is small, but the system also needs to account for different data distributions, missing clients, and information that can leak through updates.

I found the interaction between privacy and robustness especially interesting. Preventing a server from seeing an update and helping it identify a bad update can pull the design in different directions. The right combination depends on which parties the system trusts and what attacks it is meant to resist.

My thesis has the fuller discussion, and the repository contains the implementation.