Practice: Implementing Advanced Aggregation

Now that we've explored the theoretical underpinnings of advanced aggregation algorithms like FedProx and SCAFFOLD, let's put them into practice. This section guides you through implementing FedProx, demonstrating how it addresses some of FedAvg's limitations, particularly in non-IID settings. We will compare its performance against the standard FedAvg baseline in a simulated environment.

Setting the Stage: Simulation Environment

We assume you have a basic federated learning simulation setup. This typically involves:

1. Server: Orchestrates the training process, aggregates model updates.
2. Clients: Simulate individual devices holding local data. Each client performs local training.
3. Dataset: A standard dataset (like MNIST or CIFAR-10) partitioned to simulate non-IID data distributions across clients. A common way to achieve this is by assigning only a limited number of classes to each client.
4. Model: A machine learning model (e.g., a simple CNN for image classification) implemented in TensorFlow or PyTorch.

Our baseline will be a standard FedAvg implementation where clients train locally for $E$ epochs using SGD and send their updated model weights back to the server for averaging.

Implementing FedProx Client Update

Recall from the "FedProx: Addressing Statistical Heterogeneity" section that FedProx modifies the client's local objective function. Instead of just minimizing the local loss $F_k(w)$ on client $k$'s data, the client minimizes:

$$\min_w F_k(w) + \frac{\mu}{2} ||w - w^t||^2$$

Here, $w^t$ represents the global model weights received from the server at the start of round $t$, and $\mu$ is a non-negative hyperparameter controlling the strength of the proximal term. This term pulls the local solution $w$ closer to the global model $w^t$, mitigating client drift caused by diverging local data distributions.

In practice, implementing this involves modifying the client's local training loop. Specifically, when computing the gradient for SGD (or any optimizer), you need to add the gradient of the proximal term. The gradient of the proximal term with respect to $w$ is simply $\mu(w - w^t)$.

Here's a conceptual Python snippet illustrating the modification within a client's local training step (assuming PyTorch):

# Assume 'model' is the client's local model instance
# 'global_model_weights' holds the weights w^t received from the server
# 'optimizer' is the standard SGD optimizer
# 'criterion' is the loss function (e.g., CrossEntropyLoss)
# 'local_data_loader' provides batches of local data
# 'mu' is the FedProx hyperparameter

# Store the initial global weights before local training
initial_global_weights = [p.clone().detach() for p in model.parameters()]

# Standard local training loop
for epoch in range(num_local_epochs):
    for batch_idx, (data, target) in enumerate(local_data_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)

        # --- FedProx Modification Starts ---
        proximal_term = 0.0
        # Iterate over model parameters and the initial global weights
        for local_param, global_param in zip(model.parameters(), initial_global_weights):
            # Ensure parameters requiring gradients are included
            if local_param.requires_grad:
                 # Calculate the squared L2 norm difference
                proximal_term += torch.sum((local_param - global_param.to(local_param.device))**2)

        loss += (mu / 2.0) * proximal_term
        # --- FedProx Modification Ends ---

        loss.backward()
        optimizer.step()

# After local training, the 'model' contains the updated weights
# to be sent back to the server (or the delta, depending on implementation)

Key Implementation Points:

• Storing $w^t$: Before starting local training iterations, the client must store a copy of the global model weights $w^t$ it received.
• Gradient Calculation: The proximal term's gradient $\mu(w - w^t)$ is added to the gradient calculated from the local loss $F_k(w)$ before the optimizer step. Many deep learning frameworks allow you to simply add the proximal term directly to the loss function, as shown above, and automatic differentiation handles the gradient calculation.
• Hyperparameter $\mu$: The choice of $\mu > 0$ is important.
  • If $\mu = 0$, FedProx reduces to FedAvg.
  • A small $\mu$ allows more deviation from the global model, potentially leading to better personalization but risking divergence similar to FedAvg on highly non-IID data.
  • A large $\mu$ forces local models to stay very close to the global model, limiting personalization but improving stability and convergence of the global model under heterogeneity. Finding a good value often requires empirical tuning.

The server-side aggregation in FedProx remains identical to FedAvg: it collects the updated local models (or model deltas) and computes a weighted average based on the number of data points on each client.

Comparing FedAvg and FedProx

To observe the impact of FedProx, set up a simulation with significant statistical heterogeneity. For example, partition the MNIST dataset among 100 clients such that each client only has data from two digit classes. Train a simple CNN using both FedAvg ($\mu = 0$) and FedProx (e.g., $\mu = 0.01$ or $\mu = 0.1$) for a fixed number of communication rounds.

Track the global model's accuracy on a held-out, balanced test set after each communication round. You might observe results similar to those depicted below:

Global model accuracy on a balanced MNIST test set over communication rounds for FedAvg and FedProx with different $\mu$ values, simulated under a non-IID data distribution (each client has data from only two classes).

Analysis and Further Steps

The chart illustrates a typical outcome:

• FedAvg: May converge faster initially but often plateaus at a lower final accuracy due to client drift caused by non-IID data. The aggregated model struggles to generalize across all data distributions.
• FedProx ($\mu > 0$): Might converge slightly slower initially, especially with larger $\mu$, because local training is constrained. However, by mitigating client drift, it often reaches a higher final global accuracy. The choice of $\mu$ balances stability and convergence speed.

This practical exercise demonstrates the tangible benefits of using an advanced aggregation algorithm like FedProx in heterogeneous environments.

Further Exploration:

1. Experiment with $\mu$: Try different values for $\mu$. How does it affect the convergence speed and final accuracy? How does it impact the performance variance across clients' local models?
2. Implement SCAFFOLD: As a next step, try implementing SCAFFOLD. This involves modifying both client and server logic to handle control variates ($c_i$, $c$) alongside model updates. Compare its convergence against FedAvg and FedProx. Note the difference in client updates (using control variates) and the additional communication overhead (sending control variate deltas).
3. Vary Heterogeneity: Change the degree of non-IID data distribution (e.g., clients having 1, 3, or 5 classes) and observe how the performance gap between FedAvg and FedProx changes.
4. Combine Techniques: Explore combining FedProx with other techniques discussed later, such as communication efficiency methods (Chapter 5).

By implementing and experimenting with these algorithms, you gain practical insight into their behavior and the trade-offs involved, preparing you to build more effective federated learning systems for real-world applications.