Federated Learning: Training Models Without Sharing Data
Federated learning concepts, implementation patterns and the practical challenges of training models across distributed devices while preserving privacy.
Traditional machine learning requires collecting all training data in one place. Federated learning flips this model: instead of bringing data to the model, you bring the model to the data. Each device trains locally, and only model updates travel to the server.
This approach solves real problems. Healthcare organizations can't share patient records. Banks can't pool transaction data. Mobile users won't upload their typing patterns. Federated learning lets all of them benefit from collective intelligence without compromising privacy.
How Federated Learning Works
The core loop is deceptively simple: distribute a model, train locally, aggregate updates, repeat. The complexity hides in the details of each step.
# ❌ Traditional centralized training
def train_centralized(all_data, model):
# All data must exist in one location
for batch in all_data:
loss = model.forward(batch)
loss.backward()
model.update()
return model
# Problem: all_data contains sensitive records from every user# ✅ Federated learning loop
from typing import List, Dict
import numpy as np
class FederatedServer:
def __init__(self, initial_model: Dict[str, np.ndarray]):
self.global_model = initial_model
self.round_number = 0
def select_clients(
self, available_clients: List[str], fraction: float = 0.1
) -> List[str]:
"""Select a random subset of clients for this round."""
num_selected = max(1, int(len(available_clients) * fraction))
indices = np.random.choice(
len(available_clients), num_selected, replace=False
)
return [available_clients[i] for i in indices]
def aggregate_updates(
self, client_updates: List[Dict[str, np.ndarray]],
client_sizes: List[int]
) -> Dict[str, np.ndarray]:
"""Weighted average of client model updates (FedAvg)."""
total_size = sum(client_sizes)
aggregated = {}
for key in self.global_model:
weighted_sum = sum(
update[key] * (size / total_size)
for update, size in zip(client_updates, client_sizes)
)
aggregated[key] = weighted_sum
self.global_model = aggregated
self.round_number += 1
return aggregatedEach client trains on its own data, computes the difference between the updated model and the received model, and sends only that delta back. The server never sees raw data—only averaged gradients or model weights.
Client-Side Training Implementation
The client component handles local training and communicates only model updates. The raw training data never leaves the device.
class FederatedClient:
def __init__(self, client_id: str, local_data: np.ndarray):
self.client_id = client_id
self.local_data = local_data
self.local_model = None
def receive_model(
self, global_model: Dict[str, np.ndarray]
) -> None:
"""Receive the latest global model from server."""
self.local_model = {
k: v.copy() for k, v in global_model.items()
}
def train_local(
self, epochs: int = 5, learning_rate: float = 0.01
) -> Dict[str, np.ndarray]:
"""Train on local data and return model update."""
for epoch in range(epochs):
for batch in self._get_batches(self.local_data):
gradients = self._compute_gradients(batch)
for key in self.local_model:
self.local_model[key] -= learning_rate * gradients[key]
return self.local_model
def _get_batches(
self, data: np.ndarray, batch_size: int = 32
):
"""Split local data into training batches."""
indices = np.random.permutation(len(data))
for i in range(0, len(data), batch_size):
batch_idx = indices[i:i + batch_size]
yield data[batch_idx]
def _compute_gradients(
self, batch: np.ndarray
) -> Dict[str, np.ndarray]:
"""Compute gradients for a single batch."""
# Simplified - real implementation uses autograd
gradients = {}
for key, weights in self.local_model.items():
gradients[key] = np.random.randn(*weights.shape) * 0.01
return gradientsThe number of local epochs matters significantly. Too few epochs and clients barely learn from their data. Too many epochs and client models diverge from each other, making aggregation less effective—a phenomenon called "client drift."
Handling Non-IID Data Distribution
In practice, data across clients is rarely identically distributed. A keyboard prediction model on a developer's phone sees different patterns than one on a teenager's phone. This non-IID (non-independent and identically distributed) nature is federated learning's biggest challenge.
# ❌ Assuming uniform data distribution
def naive_aggregate(updates: List[Dict[str, np.ndarray]]):
"""Simple average assumes all clients have similar data."""
aggregated = {}
for key in updates[0]:
aggregated[key] = np.mean(
[u[key] for u in updates], axis=0
)
return aggregated# ✅ FedProx: Adding a proximal term to handle heterogeneity
class FedProxClient(FederatedClient):
def __init__(
self, client_id: str, local_data: np.ndarray, mu: float = 0.01
):
super().__init__(client_id, local_data)
self.mu = mu # Proximal term strength
self.reference_model = None
def receive_model(
self, global_model: Dict[str, np.ndarray]
) -> None:
super().receive_model(global_model)
self.reference_model = {
k: v.copy() for k, v in global_model.items()
}
def train_local(
self, epochs: int = 5, learning_rate: float = 0.01
) -> Dict[str, np.ndarray]:
"""Train with proximal term to limit client drift."""
for epoch in range(epochs):
for batch in self._get_batches(self.local_data):
gradients = self._compute_gradients(batch)
for key in self.local_model:
# Standard gradient update
update = learning_rate * gradients[key]
# Proximal term: penalize divergence from global model
proximal_penalty = self.mu * (
self.local_model[key] - self.reference_model[key]
)
self.local_model[key] -= update + proximal_penalty
return self.local_modelFedProx adds a regularization term that penalizes local models for drifting too far from the global model. The mu hyperparameter controls the trade-off: higher values keep clients closer to the global model but may underfit local data.
Privacy Guarantees with Differential Privacy
Federated learning alone doesn't guarantee privacy. Model updates can leak information about training data through gradient inversion attacks. Adding differential privacy provides mathematical guarantees.
class DifferentiallyPrivateClient(FederatedClient):
def __init__(
self,
client_id: str,
local_data: np.ndarray,
noise_multiplier: float = 1.0,
max_grad_norm: float = 1.0,
):
super().__init__(client_id, local_data)
self.noise_multiplier = noise_multiplier
self.max_grad_norm = max_grad_norm
def _clip_gradients(
self, gradients: Dict[str, np.ndarray]
) -> Dict[str, np.ndarray]:
"""Clip gradient norm to bound sensitivity."""
total_norm = np.sqrt(
sum(np.sum(g ** 2) for g in gradients.values())
)
clip_factor = min(1.0, self.max_grad_norm / (total_norm + 1e-6))
return {
key: grad * clip_factor
for key, grad in gradients.items()
}
def _add_noise(
self, gradients: Dict[str, np.ndarray]
) -> Dict[str, np.ndarray]:
"""Add calibrated Gaussian noise for differential privacy."""
noise_scale = self.noise_multiplier * self.max_grad_norm
return {
key: grad + np.random.normal(0, noise_scale, grad.shape)
for key, grad in gradients.items()
}
def train_local(
self, epochs: int = 5, learning_rate: float = 0.01
) -> Dict[str, np.ndarray]:
"""Train with DP guarantees: clip then noise."""
for epoch in range(epochs):
for batch in self._get_batches(self.local_data):
gradients = self._compute_gradients(batch)
gradients = self._clip_gradients(gradients)
gradients = self._add_noise(gradients)
for key in self.local_model:
self.local_model[key] -= learning_rate * gradients[key]
return self.local_modelThe privacy-utility trade-off is real. More noise means stronger privacy guarantees but slower convergence. In practice, you need to track the privacy budget (epsilon) across rounds and stop training when the budget is exhausted.
Communication Efficiency Strategies
Bandwidth is the bottleneck in federated learning. Sending full model updates from thousands of devices is expensive. Compression techniques reduce communication costs dramatically.
def compress_update_top_k(
update: Dict[str, np.ndarray], k_fraction: float = 0.1
) -> Dict[str, tuple]:
"""Keep only top-k% of gradient values by magnitude."""
compressed = {}
for key, values in update.items():
flat = values.flatten()
k = max(1, int(len(flat) * k_fraction))
top_indices = np.argpartition(np.abs(flat), -k)[-k:]
top_values = flat[top_indices]
compressed[key] = (
top_indices,
top_values,
values.shape,
)
return compressed
def decompress_update(
compressed: Dict[str, tuple]
) -> Dict[str, np.ndarray]:
"""Reconstruct full update from compressed representation."""
decompressed = {}
for key, (indices, values, shape) in compressed.items():
full = np.zeros(np.prod(shape))
full[indices] = values
decompressed[key] = full.reshape(shape)
return decompressedTop-k sparsification typically retains 90%+ of the model quality while transmitting only 1-10% of the parameters. Combined with quantization (reducing float32 to int8), you can achieve 100x compression with minimal accuracy loss.
Orchestrating the Full Training Loop
Putting it all together requires careful orchestration of client selection, training, aggregation, and evaluation across multiple rounds.
def run_federated_training(
server: FederatedServer,
clients: List[FederatedClient],
num_rounds: int = 100,
clients_per_round: float = 0.1,
local_epochs: int = 5,
) -> List[float]:
"""Run complete federated training loop."""
accuracies = []
for round_num in range(num_rounds):
# Step 1: Select participating clients
client_ids = [c.client_id for c in clients]
selected_ids = server.select_clients(client_ids, clients_per_round)
selected_clients = [
c for c in clients if c.client_id in selected_ids
]
# Step 2: Distribute global model
for client in selected_clients:
client.receive_model(server.global_model)
# Step 3: Local training
updates = []
sizes = []
for client in selected_clients:
update = client.train_local(epochs=local_epochs)
updates.append(update)
sizes.append(len(client.local_data))
# Step 4: Aggregate updates
server.aggregate_updates(updates, sizes)
# Step 5: Evaluate
accuracy = evaluate_model(server.global_model)
accuracies.append(accuracy)
if round_num % 10 == 0:
print(
f"Round {round_num}: accuracy={accuracy:.4f}, "
f"clients={len(selected_clients)}"
)
return accuracies
def evaluate_model(
model: Dict[str, np.ndarray]
) -> float:
"""Evaluate model on held-out test set."""
# Simplified evaluation
return 0.0Production federated systems add fault tolerance (clients dropping mid-round), secure aggregation (the server can't see individual updates), and asynchronous updates (don't wait for slow clients). Each layer adds complexity but addresses real deployment challenges.
Key Takeaways
Federated learning represents a fundamental shift in how we think about data and model training. The core insight is that models can learn from data they never see directly—and this enables use cases that were previously impossible due to privacy regulations, competitive concerns, or simple logistics.
The technical challenges are real: non-IID data distributions degrade model quality, communication costs scale with model size and client count, and privacy guarantees come at the cost of model accuracy. But the field is maturing rapidly. If your application involves sensitive user data and you've been limited by the centralized training paradigm, federated learning might be the approach that unlocks your next breakthrough.
