← Back to home
Jul 15, 2026MR. ERROR 3132 min read

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.

machine-learningfederated-learningpytorch
Federated Learning

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: