FedNova: Normalized Averaging for Heterogeneous Systems

"While Federated Averaging (FedAvg) provides a simple baseline for aggregating client updates, its performance can degrade significantly in federated networks characterized by systems heterogeneity. Clients often possess varying computational capabilities (CPU, memory), network bandwidths, and power constraints. This heterogeneity leads to variations in the amount of local computation each client can perform within a given communication round. Some clients might complete many local training epochs ($E$), while others only manage a few."

Standard FedAvg aggregates client models (or updates) typically weighted by the number of data points ($n_k$). However, it doesn't explicitly account for the differing amounts of local work ($\tau_k$, the number of local gradient steps or updates performed by client $k$) done to arrive at those updates. When $\tau_k$ varies greatly across clients, FedAvg can suffer from:

1. Client Drift Mismatch: Updates from clients that performed more local steps might represent models that have diverged further from the initial global model $w^t$. Simply averaging these can lead to instability or slow convergence, as the aggregate update might be dominated by clients who performed extensive, potentially diverging, local optimization.
2. Stale Updates (Implicitly): In synchronous settings where the server waits for a quorum, faster clients finish early. In asynchronous settings, updates arrive at different times reflecting different amounts of work. FedAvg doesn't inherently compensate for the staleness or computational effort discrepancies.
3. Inefficient Gradient Approximation: The goal of aggregation is often to approximate the true global gradient $\nabla F(w) = \sum p_k \nabla F_k(w)$. If clients perform different numbers of local steps ($\tau_k$), the simple weighted average of local updates $\Delta_k^t = w_k^{t+1} - w^t$ used in FedAvg might not be the most effective approximation, especially when local learning rates $\eta_l$ are involved (as local update $\Delta_k^t \approx - \eta_l \tau_k \nabla F_k(w^t)$).

The FedNova Approach: Normalizing Local Updates

Federated Normalized Averaging (FedNova) tackles systems heterogeneity directly by normalizing client updates before aggregation. The central idea is to adjust each client's contribution to counteract the effect of varying local computations, aiming for an update that better reflects the average gradient direction across clients, irrespective of how many steps each took locally.

Instead of averaging the final local models or the raw updates $\Delta_k^t$, FedNova averages the updates normalized by the amount of local work performed. Let $\tau_k$ be the number of local steps (e.g., SGD updates) performed by client $k$ in round $t$. The local update submitted by client $k$ is $\Delta_k^t = w_k^{t+1} - w^t$.

FedNova computes the global model update as follows:

$$w^{t+1} = w^t + \sum_{k \in S_t} p_k \left( \frac{\Delta_k^t}{\tau_k} \right)$$

Here, $S_t$ is the set of clients participating in round $t$, and $p_k$ is the aggregation weight for client $k$ (commonly $p_k = n_k / \sum_{j \in S_t} n_j$, but other weightings are possible).

Compare this to the implicit update direction in FedAvg when viewed through local updates:

$$w^{t+1}_{\text{FedAvg}} = w^t + \sum_{k \in S_t} p_k \Delta_k^t$$

Effectively, FedNova aims to average the average local step update ($\Delta_k^t / \tau_k$), weighted by $p_k$. This prevents clients who performed many local steps ($\tau_k \gg \tau_j$) from disproportionately influencing the magnitude and direction of the aggregated update simply because they computed more. It assumes that $\Delta_k^t / \tau_k$ provides a better estimate of the local gradient direction relevant to the global objective than the raw update $\Delta_k^t$ when $\tau_k$ varies.

Visualizing the Normalization Effect

Consider two clients with the same amount of data ($p_1 = p_2 = 0.5$) but different computational capabilities. Client 1 performs $\tau_1 = 2$ local steps, while Client 2 performs $\tau_2 = 10$ local steps. Their raw updates are $\Delta_1^t$ and $\Delta_2^t$.

FedAvg directly averages potentially disparate updates $\Delta_1^t$ and $\Delta_2^t$. If Client 2's update $\Delta_2^t$ is much larger due to $\tau_2=10$, it might dominate. FedNova first normalizes by the number of steps ($\tau_k$) before averaging, potentially leading to a different aggregate direction that better reflects the average single-step progress.

Benefits and Convergence

The primary advantage of FedNova is its robustness to systems heterogeneity. By normalizing updates based on local work:

• Improved Convergence: Theoretical analysis suggests FedNova can achieve faster convergence rates than FedAvg under significant systems heterogeneity, as it corrects for the varying contributions arising from different $\tau_k$.
• Fairness: It prevents computationally powerful clients or stragglers (who might run for longer if allowed) from having an undue influence on the global model update purely based on the volume of local computation.
• Stability: Normalization can lead to more stable training dynamics compared to FedAvg when $\tau_k$ values fluctuate widely across clients or rounds.

Practical Notes

Implementing FedNova requires a minor modification to the standard FL protocol:

• Reporting Local Work: Clients need to report the amount of local work performed ($\tau_k$) along with their model updates ($\Delta_k^t$). This is typically a single scalar value (e.g., the number of local gradient steps taken), adding minimal communication overhead.
• Server-Side Calculation: The server performs the normalization ($\Delta_k^t / \tau_k$) before executing the weighted aggregation step.

It's important to note that FedNova primarily addresses systems heterogeneity. It does not directly solve the challenges of statistical heterogeneity (Non-IID data), although it can potentially be combined with algorithms like FedProx or SCAFFOLD designed for that purpose. The effectiveness of FedNova relies on the assumption that the normalized update $\Delta_k^t / \tau_k$ is a meaningful quantity, representing an average step direction. This generally holds well for common local solvers like SGD.

In summary, FedNova provides a theoretically grounded and practically effective mechanism for mitigating the adverse effects of varying computational resources and local steps across clients in a federated network. By normalizing updates before aggregation, it promotes fairer contributions and can lead to faster, more stable convergence compared to FedAvg in heterogeneous system environments. It's a valuable tool for building more performant FL systems when client capabilities vary significantly.