SCAFFOLD: Variance Reduction in Federated Optimization

While FedAvg provides a simple baseline for aggregating client updates, its performance often degrades significantly in the presence of statistical heterogeneity (Non-IID data). Clients training on very different local datasets can pull the global model in conflicting directions, leading to slow convergence, oscillations, or even divergence. This phenomenon, often called "client drift," stems from the discrepancy between the local optimization objective each client pursues and the global objective.

SCAFFOLD (Stochastic Controlled Averaging for Federated Learning) offers an elegant solution to mitigate this client drift by introducing control variates. The core idea is to estimate the update direction each client would have taken if it had access to the global data distribution and then correct the actual local updates based on this estimate. This reduces the variance between client updates, aligning them more closely with the global objective and thereby accelerating convergence.

The Mechanics of SCAFFOLD

SCAFFOLD modifies the standard federated learning process by maintaining additional state information on both the server and the clients:

1. Server Control Variate ($c$): Represents an estimate of the average gradient direction across all data in the federation.
2. Client Control Variates ($c_k$): Each client $k$ maintains its own control variate $c_k$, representing an estimate of the gradient direction based on its local data.

The difference $c - c_k$ effectively captures the "drift" for client $k$. SCAFFOLD uses these control variates to adjust the local training process and the aggregation step.

Algorithm Overview

Let $w^t$ be the global model weights at communication round $t$.

1. Server Broadcast: The server sends the current global model $w^t$ and its control variate $c^t$ to the selected clients.

2. Client Computation: Each participating client $k$ performs the following steps:

   • Initializes its local model $w_k^t = w^t$.
   • Receives the server control variate $c^t$.
   • Performs $E$ local steps of stochastic gradient descent (SGD). For each local step $\tau = 0, ..., E-1$, using a local minibatch $b$:
     • Compute the local gradient: $g_k(w_{k, \tau}^t) = \nabla F_k(w_{k, \tau}^t; b)$, where $F_k$ is the local loss function for client $k$.
     • Update the local model using the corrected gradient: $$w_{k, \tau+1}^t = w_{k, \tau}^t - \eta_l (g_k(w_{k, \tau}^t) - c_k^t + c^t)$$ Here, $\eta_l$ is the local learning rate. Notice how the local gradient $g_k$ is adjusted by the difference between the server control variate $c^t$ and the client's control variate $c_k^t$. This adjustment aims to correct for the client's local drift.
   • Compute the total model update direction: $\Delta w_k^t = w_{k, E}^t - w^t$.
   • Update the client control variate $c_k$. A common way is: $$c_k^{t+1} = c_k^t - c^t + \frac{1}{E \eta_l} (w^t - w_{k, E}^t)$$ This update reflects the average local gradient direction observed during the $E$ steps.
   • Compute the change in the client control variate: $\Delta c_k^t = c_k^{t+1} - c_k^t$.
   • Send $\Delta w_k^t$ and $\Delta c_k^t$ back to the server.

3. Server Aggregation: The server receives updates $(\Delta w_k^t, \Delta c_k^t)$ from a set $S_t$ of participating clients.

   • Aggregate model updates (similar to FedAvg): $$\Delta w^t = \frac{1}{|S_t|} \sum_{k \in S_t} \Delta w_k^t$$
   • Update the global model: $$w^{t+1} = w^t + \eta_g \Delta w^t$$ ($\eta_g$ is the server learning rate, often set to 1).
   • Aggregate control variate updates: $$\Delta c^t = \frac{1}{|S_t|} \sum_{k \in S_t} \Delta c_k^t$$
   • Update the server control variate: $$c^{t+1} = c^t + \Delta c^t$$

Why Does Variance Reduction Work?

The main intuition is that the term $(g_k(w_{k, \tau}^t) - c_k^t + c^t)$ used in the client update is a better estimate of the global gradient $\nabla F(w_{k, \tau}^t)$ than the raw local gradient $g_k(w_{k, \tau}^t)$. By subtracting the estimated local direction ($c_k^t$) and adding the estimated global direction ($c^t$), SCAFFOLD effectively steers the client updates towards the global minimum, reducing the variance caused by disparate local data distributions.

This variance reduction leads to more stable and typically faster convergence compared to FedAvg, especially when dealing with significant statistical heterogeneity (Non-IID data).

Convergence behavior of FedAvg, FedProx, and SCAFFOLD under simulated non-IID conditions. SCAFFOLD often achieves faster convergence and potentially higher final accuracy due to its variance reduction mechanism.

Implementation

Implementing SCAFFOLD introduces some overhead compared to FedAvg:

• State: Both the server and clients need to store their respective control variates ($c$ and $c_k$). These are typically the same size as the model parameters.
• Communication: Clients need to send both the model update ($\Delta w_k$) and the control variate update ($\Delta c_k$) to the server in each round. The server needs to broadcast the global control variate $c$ in addition to the global model $w$. This effectively doubles the communication cost per round compared to FedAvg, although the potential reduction in the number of rounds required for convergence can often compensate for this.

Despite the increased state and communication requirements, SCAFFOLD's ability to effectively handle client drift makes it a valuable tool for improving the performance and reliability of federated learning systems operating on heterogeneous data. Its theoretical underpinnings provide stronger convergence guarantees than FedAvg under non-IID settings.