Optimizing Local Computation

As highlighted in the introduction to this chapter, communication costs, particularly the uplink from numerous clients to a central server, represent a significant performance bottleneck in many federated learning scenarios. While techniques like gradient compression directly tackle the size of the messages, another complementary approach focuses on reducing the frequency of communication rounds. This is achieved by optimizing the amount of computation performed locally on each client device before synchronization.

In the standard Federated Averaging (FedAvg) algorithm, two primary parameters govern this balance: the number of clients participating in each round ($K$), and the number of local training epochs ($E$) each selected client performs on its data before sending an update. This section examines the role of $E$ and the associated trade-offs.

The Trade-off: Local Work vs. Communication Rounds

Increasing the number of local epochs ($E$) allows each client to perform more steps of local optimization (e.g., stochastic gradient descent) using its own dataset between communication rounds. Intuitively, if clients make more progress locally, fewer communication rounds might be needed to reach a target global model accuracy. Each communication round incurs overhead (network latency, server aggregation time), so reducing the total number of rounds can lead to faster overall training in terms of wall-clock time.

Consider the basic federated optimization objective:

$$\min_{w} F(w) = \sum_{i=1}^{N} p_i F_i(w)$$

where $w$ represents the global model parameters, $N$ is the total number of clients, $p_i$ is the weight for client $i$ (often proportional to its dataset size), and $F_i(w)$ is the local objective function for client $i$.

When a client $i$ performs $E$ local epochs, it starts with the current global model $w_t$ and iteratively updates it using its local data, resulting in a local model $w_{t+1}^i$. The server then aggregates these local models (e.g., averaging) to produce the next global model $w_{t+1}$.

If $E=1$, clients perform only one pass (or potentially just one batch update) over their local data. This closely resembles a large-batch synchronous SGD approach distributed across clients, often referred to as Federated SGD (FedSGD). Communication happens frequently.

If $E > 1$, clients perform multiple updates locally. This reduces communication frequency but introduces a potential challenge: client drift. As clients optimize intensely on their local, potentially unique (Non-IID) data distributions, their local models ($w_{t+1}^i$) might diverge significantly from each other and from the optimal solution for the global objective $F(w)$. Aggregating these diverged models can slow down convergence or even prevent the global model from reaching a good solution.

Impact of Statistical Heterogeneity

The effectiveness of increasing $E$ is strongly tied to the degree of statistical heterogeneity (Non-IID data) across clients, a topic discussed extensively in Chapter 4.

• IID Data: If client data distributions are independent and identically distributed (IID), performing more local steps ($E > 1$) generally accelerates convergence compared to $E=1$, as local updates approximate updates on the global data distribution well. Client drift is less of a concern.
• Non-IID Data: With heterogeneous data, increasing $E$ becomes a delicate balancing act. More local steps can exacerbate client drift. The local minima found by clients on their specific data distributions might be far from the global minimum. Algorithms like FedProx (Chapter 2), which introduces a proximal term to keep local updates closer to the global model, $$w_{t+1}^i = \arg\min_{w} F_i(w) + \frac{\mu}{2} \|w - w_t\|^2$$ can make training more effective to larger values of $E$ even with Non-IID data. Similarly, SCAFFOLD (Chapter 2) uses control variates to correct for client drift, potentially allowing for more aggressive local computation schedules.

System Heterogeneity and Stragglers

In statistical heterogeneity, system heterogeneity (variations in client hardware, network connectivity, and availability) also influences the optimal choice of $E$. In synchronous FL, the server waits for all $K$ selected clients to complete their $E$ local epochs before proceeding. If $E$ is large, clients with slower processors or limited power will take longer, delaying the entire round. These "stragglers" can severely limit the effective training speed.

Choosing a smaller $E$ reduces the computation time per round, mitigating the impact of stragglers in synchronous settings. Asynchronous protocols, discussed later in this chapter, offer alternative ways to handle stragglers without being strictly bottlenecked by the slowest device, potentially making larger $E$ values more viable in heterogeneous systems.

Practical Tuning and Considerations

Finding the "best" value for $E$ is often an empirical process specific to the application, model architecture, dataset characteristics, and system constraints. There's typically a sweet spot:

• Too small $E$: High communication overhead dominates the training time.
• Too large $E$: Client drift slows convergence, or stragglers significantly delay synchronous rounds.

A common practical approach is to start with a small $E$ (e.g., 1 to 5) and gradually increase it while monitoring convergence. Important metrics to track include:

1. Global model accuracy vs. communication rounds.
2. Global model accuracy vs. wall-clock time.

The goal is usually to minimize wall-clock time to reach a desired accuracy level.

This scenario illustrates how increasing the number of local epochs ($E$) can initially speed up convergence in terms of wall-clock time (e.g., $E=5, E=10$ reach higher accuracy faster than $E=1$). However, excessively large $E$ (e.g., $E=20$) might lead to slower convergence or instability due to client drift, especially on heterogeneous data. The optimal $E$ balances computation and communication.

Also, consider the interaction between $E$ and the local learning rate. Performing more local updates ($E > 1$) might require using a smaller local learning rate compared to the $E=1$ case to maintain stability.

In summary, optimizing local computation by tuning $E$ is a significant lever for improving the efficiency of federated learning systems. It allows trading off local processing effort against communication frequency. However, this tuning must carefully account for the statistical properties of the data (Non-IID) and the systemic characteristics (hardware/network variability) of the participating clients to avoid issues like client drift and straggler bottlenecks.