Federated Learning

Federated learning enables multiple participants to collaboratively train a shared model without exchanging raw data. Each participant trains locally and shares only model updates, preserving data privacy.

1. Problem Formulation

Federated Optimization

Minimize a global objective across participants:

where is the weight of participant .

Communication Rounds

Each round : (1) Server broadcasts , (2) Participants compute local updates, (3) Participants send updates to server, (4) Server aggregates.

2. Federated Averaging (FedAvg)

Algorithm

Local SGD update at participant :

Global aggregation:

Convergence Analysis

Theorem (McMahan et al., 2017): Under -smoothness and bounded client drift :

Client Sampling

Select subset per round:

3. Communication Efficiency

Gradient Compression

Top-k sparsification: Send only the largest gradient components:

Error feedback: Accumulate compression error:

Scaffold: Variance Reduction

SCAFFOLD uses control variates:

4. Differential Privacy

Definitions

A mechanism satisfies -DP if:

DP-SGD

Gradient clipping:

Gaussian noise:

Rényi DP composition after steps:

Federated DP-SGD

Privacy amplification:

5. Secure Aggregation

Secret sharing: Each participant splits update into shares:

Homomorphic encryption: Aggregate on encrypted updates:

6. Data Heterogeneity

Non-IID Data

When , client drift increases.

FedProx

Add proximal regularization:

Personalization

Per-FedAvg: Meta-learn initialization for personalization:

7. Implementation

import torch
import torch.nn as nn
from typing import List, Dict

class FedAvg:
    def __init__(self, model, num_clients, fraction=0.1):
        self.model = model
        self.num_clients = num_clients
        self.fraction = fraction

    def aggregate(self, client_updates, client_weights):
        total_weight = sum(client_weights)
        aggregated = {}
        for key in client_updates[0].keys():
            aggregated[key] = sum(
                client_updates[i][key] * client_weights[i] / total_weight
                for i in range(len(client_updates))
            )
        return aggregated

    def train_round(self, clients, round_num):
        num_selected = max(1, int(self.fraction * self.num_clients))
        selected = torch.randperm(self.num_clients)[:num_selected]
        client_updates, client_weights = [], []
        for k in selected:
            update = clients[k].local_train(self.model.state_dict())
            client_updates.append(update)
            client_weights.append(clients[k].dataset_size)
        self.model.load_state_dict(self.aggregate(client_updates, client_weights))

class SCAFFOLD:
    def __init__(self, model, lr=0.01):
        self.model = model
        self.lr = lr
        self.c_global = {n: torch.zeros_like(p) for n, p in model.named_parameters()}

    def client_update(self, client_model, client_c, local_data, num_epochs=5):
        c_delta = {n: torch.zeros_like(p) for n, p in client_model.named_parameters()}
        for _ in range(num_epochs):
            for batch in local_data:
                loss = compute_loss(client_model, batch)
                grads = torch.autograd.grad(loss, client_model.parameters())
                for (n, p), g in zip(client_model.named_parameters(), grads):
                    step = self.lr * (g - client_c[n] + self.c_global[n])
                    p.data -= step
                    c_delta[n] += (self.c_global[n] - client_c[n] - step / self.lr)
        return c_delta

class DPLocalSGD:
    def __init__(self, model, noise_multiplier=1.0, max_grad_norm=1.0):
        self.model = model
        self.noise_multiplier = noise_multiplier
        self.max_grad_norm = max_grad_norm

    def private_update(self, dataloader, lr):
        for batch in dataloader:
            loss = compute_loss(self.model, batch)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
            scale = self.noise_multiplier * self.max_grad_norm
            for p in self.model.parameters():
                if p.grad is not None:
                    p.grad += torch.randn_like(p.grad) * scale
            self.model.step()
            self.model.zero_grad()

8. SVG: Federated Learning Architecture

9. SVG: FedAvg Communication Rounds

10. Comparison of Methods

11. Open Problems

• Communication compression: Balancing compression ratio with convergence
• Heterogeneous devices: Stragglers and varying compute capabilities
• Byzantine robustness: Resilient aggregation against malicious participants
• Federated unlearning: Removing data influence after training
• Cross-silo vs cross-device: Different trust and participation models

References

1. McMahan, B., et al. (2017). Communication-Efficient Learning of Deep Networks from Decentralized Data. AISTATS.
2. Karimireddy, S. P., et al. (2019). SCAFFOLD: Stochastic Controlled Averaging for Federated Learning. ICML.
3. Abadi, M., et al. (2016). Deep Learning with Differential Privacy. CCS.
4. Li, T., et al. (2018). Federated Optimization in Heterogeneous Networks. MLSys.
5. Bonawitz, K., et al. (2017). Practical Secure Aggregation for Privacy-Preserving Machine Learning. CCS.