Federated Learning
Federated learning is a machine learning approach that enables multiple actors to collaborate on model training while maintaining data privacy. This tutorial covers the core concepts and provides a worked example.

Introduction to Federated Learning
Federated learning is a machine learning approach that enables multiple actors to collaborate on model training while maintaining data privacy. This is particularly useful in scenarios where data cannot be shared due to privacy concerns, such as in healthcare or finance.
Context and Why it Matters
The increasing use of machine learning in various industries has led to a growing need for data sharing and collaboration. However, data sharing can be challenging due to privacy concerns, regulatory restrictions, and data ownership issues. Federated learning addresses these challenges by enabling multiple actors to collaborate on model training while maintaining data privacy.
Core Concept
The core concept of federated learning is to train a model on decentralized data. Each actor has a local dataset, and the goal is to train a global model that performs well on all local datasets. This is achieved through the following steps:
- Local training: Each actor trains a local model on their local dataset.
- Model aggregation: The local models are aggregated to form a global model.
- Global update: The global model is updated and sent back to each actor for further local training.
Worked Example
Here's an example of federated learning using PyTorch. We'll use a simple neural network to classify handwritten digits (MNIST dataset).
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# Define the neural network model
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28 * 28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# Define the federated learning algorithm
class FederatedLearning:
def __init__(self, num_actors, num_epochs):
self.num_actors = num_actors
self.num_epochs = num_epochs
self.models = [Net() for _ in range(num_actors)]
def local_train(self, model, device, loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data.view(-1, 28 * 28))
loss = nn.CrossEntropyLoss()(output, target)
loss.backward()
optimizer.step()
def aggregate_models(self):
# Aggregate the local models
global_model = Net()
for param in global_model.parameters():
param.data.zero_()
for local_model in self.models:
for global_param, local_param in zip(global_model.parameters(), local_model.parameters()):
global_param.data += local_param.data / self.num_actors
return global_model
def run(self):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for epoch in range(self.num_epochs):
for i in range(self.num_actors):
# Local training
local_model = self.models[i]
local_optimizer = optim.SGD(local_model.parameters(), lr=0.01)
local_loader = torch.utils.data.DataLoader(datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transforms.Compose([transforms.ToTensor()])), batch_size=64, shuffle=True)
self.local_train(local_model, device, local_loader, local_optimizer, epoch)
# Aggregate the local models
global_model = self.aggregate_models()
# Update the local models
for i in range(self.num_actors):
self.models[i].load_state_dict(global_model.state_dict())
# Create a federated learning instance
federated_learning = FederatedLearning(num_actors=5, num_epochs=10)
federated_learning.run()
Pitfalls
Federated learning can be challenging due to the following pitfalls:
- Non-IID data: The local datasets may not be independent and identically distributed (IID), which can affect the performance of the global model.
- Communication overhead: The communication overhead of aggregating the local models can be significant, especially in large-scale federated learning scenarios.
- Security: Federated learning can be vulnerable to security attacks, such as model inversion attacks or data poisoning attacks.
What to Read Next
If you're interested in learning more about federated learning, here are some recommended readings: