FedAvgM

The core innovation of FedAvgM is the application of a momentum term exclusively on the server. Unlike standard momentum in centralized SGD, which smooths the gradient trajectory during local client updates, FedAvgM's momentum smooths the trajectory of the aggregated model updates received from clients. The server maintains a velocity vector v_t that accumulates a fraction of past aggregated updates, which is then added to the current update. This helps dampen oscillations caused by client heterogeneity and noisy updates, leading to a more stable descent path for the global model.

A primary challenge in federated learning is client drift, where local models diverge from the global objective due to optimization on statistically heterogeneous (non-IID) data. FedAvgM's server momentum acts as a stabilizing force against this drift. By incorporating historical update directions, the algorithm reduces the variance in the global update step. This provides a form of inertia that helps correct for the biased updates from individual clients, pulling the global model update toward a consensus direction that is more representative of the overall population data distribution.

The momentum mechanism in FedAvgM is proven to accelerate convergence rates compared to vanilla FedAvg, especially under non-IID data conditions. The acceleration occurs because momentum helps navigate the optimization landscape more efficiently. In practice, this can lead to:

• Fewer communication rounds required to reach a target model accuracy.
• More stable loss curves with reduced oscillation between rounds.
• Improved final test accuracy, as the model converges to a better optimum by escaping shallow local minima that standard averaging might get stuck in.

FedAvgM is a specific instance within the broader FedOpt (Federated Optimization) framework. This framework generalizes server-side aggregation beyond simple averaging to include adaptive optimizers. FedAvgM applies a momentum-based optimizer (analogous to SGD with Momentum) on the server. Other members of the FedOpt family include FedAdam and FedYogi, which use Adam and Yogi optimizers, respectively, on the server. FedAvgM is often the simplest and most computationally efficient member of this family, providing a clear performance boost over FedAvg without the additional hyperparameters of adaptive methods.

The algorithm introduces a key hyperparameter, the momentum coefficient (β), typically a value between 0 and 1 (e.g., 0.9). This parameter controls the weight given to the historical velocity vector versus the current aggregated update.

• A high β (e.g., 0.99) gives strong weight to past updates, leading to very smooth but potentially slow adaptation to new data trends.
• A low β (e.g., 0.5) makes the algorithm more responsive to the current round's updates but provides less stabilization. Tuning β is crucial and often depends on the degree of client data heterogeneity and the desired trade-off between stability and adaptability.

It is critical to distinguish FedAvgM from applying momentum during local client training. In standard FedAvg, clients can use SGD with Momentum locally, but this momentum state is discarded after each round. FedAvgM's momentum is persistent on the server across rounds. This server-side persistence is what provides the long-term memory and stabilization effect for the global training process. The two techniques are orthogonal and can be combined: clients can use local momentum for faster local convergence, while the server uses FedAvgM momentum for stable global aggregation.

FedAvgM's most critical use case is training on statistically heterogeneous (non-IID) data across clients. The server-side momentum term acts as a low-pass filter, smoothing the aggregation of potentially conflicting client updates. This reduces the variance in the global update direction, counteracting client drift and leading to faster, more stable convergence compared to vanilla FedAvg. It is particularly effective when client data distributions diverge significantly.

In practical deployments with partial client participation, the subset of clients selected each round can introduce high variance into the aggregated update. FedAvgM's momentum accumulates update information across rounds, providing inertia against noisy updates from a small or unrepresentative client sample. This results in a more stable training trajectory and reduces the performance degradation associated with low participation rates.

In systems with high client heterogeneity in compute or connectivity, stragglers can delay synchronous rounds or, in asynchronous protocols, send stale updates. The momentum mechanism helps dampen the disruptive effect of these delayed or outdated updates. By maintaining a running average of recent update directions, the global model is less susceptible to being pulled off course by a single slow or late-arriving client update.

FedAvgM is well-suited for cross-silo settings (e.g., collaboration between a few large organizations like hospitals or banks). In these environments, each client (silo) has a large, stable dataset, and communication rounds are expensive but reliable. The momentum term efficiently utilizes the high-quality updates from each silo, accelerating progress toward a high-performance global model where standard averaging might plateau or oscillate.

The primary engineering benefit of FedAvgM is achieving a higher final test accuracy for the global model under challenging federated conditions. By navigating the optimization landscape more efficiently, it often converges to a better minimum than FedAvg. This is quantifiable in benchmarks, showing accuracy improvements of several percentage points on non-IID splits of datasets like CIFAR-10 or FEMNIST.

FedAvgM can significantly reduce the total number of communication rounds required to reach a target accuracy. The momentum-aided updates make more progress per round. This directly translates to lower overall training time and reduced bandwidth consumption, which is a major cost and latency factor in federated learning. It provides a clear return on investment for the minimal added computation on the server.

The FedOpt framework generalizes server-side aggregation beyond simple averaging. Instead of applying a static learning rate to the averaged update, the server uses adaptive optimization algorithms (like Adam, Yogi, or Adagrad) on the aggregated client gradients. This allows the server to maintain and adapt per-parameter momentum and learning rates, which can accelerate convergence, especially in heterogeneous settings. FedAvgM is a specific instantiation of this framework using a momentum optimizer on the server.

FedAdam is a direct sibling to FedAvgM within the FedOpt framework. While FedAvgM uses standard momentum (similar to SGD with Momentum), FedAdam employs the Adam optimizer on the server. This incorporates both momentum (first moment) and adaptive learning rates based on the variance of the gradients (second moment). Key differences include:

• Adaptive per-parameter learning rates that can handle sparse gradients better.
• Bias correction for initial steps.
• Often requires tuning additional hyperparameters (beta1, beta2, epsilon).

SCAFFOLD (Stochastic Controlled Averaging) tackles client drift—a primary cause of slow convergence—using control variates. Each client and the server maintain a state variable that estimates the update direction. Clients use this to correct their local descent, reducing the variance between their update and the true global direction. Unlike FedAvgM, which adds momentum on the server, SCAFFOLD modifies the client's local objective to align updates, proving particularly effective under high statistical heterogeneity (non-IID data).

FedProx addresses client drift by adding a proximal term to the local loss function on each client. This term penalizes the local model parameters for straying too far from the initial global model received at the start of the round. It effectively acts as a regularizer, constraining local updates. While FedAvgM uses server momentum to smooth the aggregation path, FedProx controls drift at the source during local training. It is especially useful when client capabilities vary, as it allows for adaptive local computation.

FedNova (Federated Normalized Averaging) solves a different convergence issue: the inconsistency in the number of local update steps taken by clients. In standard FedAvg, clients performing more local epochs contribute disproportionately larger updates. FedNova normalizes each client's update by its local steps before aggregation, ensuring all updates are on an equivalent scale. This normalization stabilizes convergence. FedAvgM can be combined with this normalization principle, applying momentum to the normalized updates for further improvement.

Momentum is a classic optimization technique that accelerates gradient descent by incorporating a fraction of the past update vector into the current update. It helps smooth the optimization path and escape shallow local minima. In FedAvgM, this concept is applied at the server level:

• The server maintains a momentum buffer.
• The aggregated client update is used to refresh this buffer.
• The global model is updated using the momentum buffer, not the raw aggregate. This provides inertia against noisy or conflicting updates from heterogeneous clients, leading to more stable convergence.