Federated Learning for Healthcare

Module: Healthcare AI | Difficulty: Advanced

Federated Averaging

Differential Privacy

Secure Aggregation

Federated Learning Approaches

| Method | Privacy | Communication | Robustness | Scalability | |--------|---------|---------------|------------|-------------| | FedAvg | Low | Low | Medium | High | | FedProx | Low | Low | High | High | | DP-FedAvg | High | Medium | Medium | High | | Secure Agg | High | High | High | Medium | | Split Learning | Medium | Low | Medium | High |

import torch
import torch.nn as nn
import copy

class FederatedServer:
    def __init__(self, global_model):
        self.global_model = global_model

    def aggregate(self, client_models, client_sizes):
        total = sum(client_sizes)
        global_dict = self.global_model.state_dict()
        for key in global_dict:
            global_dict[key] = torch.zeros_like(global_dict[key], dtype=torch.float32)
        for model, size in zip(client_models, client_sizes):
            for key in global_dict:
                global_dict[key] += model.state_dict()[key].float() * (size / total)
        self.global_model.load_state_dict(global_dict)

class FedProxClient:
    def __init__(self, model, mu=0.01):
        self.model = model
        self.mu = mu
        self.global_params = None

    def train(self, dataloader, epochs=5, lr=0.01):
        optimizer = torch.optim.SGD(self.model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()
        for epoch in range(epochs):
            for batch_x, batch_y in dataloader:
                optimizer.zero_grad()
                output = self.model(batch_x)
                loss = criterion(output, batch_y)
                if self.global_params is not None:
                    prox_term = sum(((p - gp) ** 2).sum()
                                   for p, gp in zip(self.model.parameters(),
                                                  self.global_params))
                    loss += (self.mu / 2) * prox_term
                loss.backward()
                optimizer.step()
        return self.model

def add_differential_privacy(model, noise_multiplier=1.0, max_grad_norm=1.0):
    total_norm = torch.sqrt(sum(p.grad.norm() ** 2 for p in model.parameters()))
    clip_coef = max_grad_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        for p in model.parameters():
            p.grad.data.mul_(clip_coef)
    for p in model.parameters():
        noise = torch.randn_like(p.grad) * noise_multiplier * max_grad_norm
        p.grad.data.add_(noise)

pipeline = FederatedServer(nn.Linear(10, 2))
print(f'Federated pipeline initialized')

Research Insight: Federated learning in healthcare faces unique challenges: non-IID data distributions across hospitals, variable data quality, and strict privacy regulations. Personalized federated methods that keep batch normalization or model heads local while sharing lower layers achieve better performance than naive federated averaging.