Federated Optimization

Federated Averaging (FedAvg) is the foundational and most widely cited algorithm for Federated Learning. It operates in synchronous communication rounds:

• The server broadcasts the current global model to a subset of selected clients.
• Each client performs multiple steps of local Stochastic Gradient Descent (SGD) on its private data.
• Clients send their updated local model weights (or gradients) back to the server.
• The server computes a weighted average of these updates, typically based on the number of data samples per client, to form a new global model.

Its simplicity is its strength, but performance degrades significantly under high statistical heterogeneity (non-IID data) and systems heterogeneity (varied client compute power).

FedProx is a robust generalization of FedAvg designed to handle the challenges of systems and statistical heterogeneity. It modifies the local optimization problem on each client by adding a proximal term to the objective function.

• This term penalizes the local model updates that stray too far from the current global model.
• The proximal term acts as a regularizer, stabilizing training when clients have non-IID data or perform vastly different amounts of local work.
• It provides theoretical convergence guarantees under more realistic assumptions of partial device participation and data heterogeneity than FedAvg.

FedProx demonstrates improved convergence stability and final model accuracy in heterogeneous environments.

SCAFFOLD (Stochastic Controlled Averaging for Federated Learning) addresses the client drift problem caused by non-IID data. Client drift occurs when local SGD updates on heterogeneous data pull client models in different directions, hindering global convergence.

SCAFFOLD's key innovation is the use of control variates (correction terms):

• Each client and the server maintain a control variate that estimates the update direction.
• Clients use their local control variate to correct their local gradient, reducing the variance introduced by data heterogeneity.
• This correction aligns local updates closer to the global objective, leading to faster convergence and often superior final accuracy compared to FedAvg, especially under high heterogeneity.

It requires storing and communicating additional state (the control variates).

FedOpt refers to a family of algorithms that generalize the server-side aggregation step of FedAvg. Instead of using a simple weighted average, FedOpt employs adaptive optimization algorithms on the server, treating client updates as pseudo-gradients.

• The server aggregates client updates (e.g., computes an average delta).
• It then applies an optimizer like Adam, Adagrad, or Yogi to this aggregated update before updating the global model.
• This adaptive server-side optimization can accelerate convergence and improve performance, particularly for complex models.

FedAvgM, which uses server-side momentum, is a common and effective example. FedOpt provides a flexible framework for incorporating advanced optimization techniques into the federated setting.

q-FedAvg is an optimization algorithm designed for fairness across clients in Federated Learning. Standard FedAvg aims to minimize the average loss across all clients, which can result in a model that performs well on average but poorly for devices with rare data distributions.

q-FedAvg reformulates the objective to minimize the aggregated loss variance across clients:

• It uses a q-parameter to trade-off between fairness (minimizing variance) and average accuracy.
• A higher q places more emphasis on improving the performance of the worst-performing clients.
• The algorithm involves weighting client updates based on their local loss, favoring clients with higher loss in the aggregation step.

This approach is crucial for applications where equitable performance across a diverse user base is a primary requirement.

Federated Multi-Task Learning (MTL) frameworks treat the federated optimization problem through the lens of multi-task learning, where each client's data distribution is considered a related but distinct task.

Instead of learning a single global model, these methods learn a model per client that is regularized to be close to a shared global model or to other related client models.

Key approaches include:

• MOCHA: Solves a dual formulation of the MTL problem, providing robustness to stragglers and theoretical convergence guarantees.
• Personalization Layers: Learning a global feature extractor with personalized final layers for each client.
• Model Interpolation: Training a global model alongside local models and interpolating them for inference.

This paradigm is particularly well-suited for highly non-IID data, as it explicitly models and leverages the relationships between client tasks.