Accelerated Federated Learning Over Wireless Fading Channels With Adaptive Stochastic Momentum

Federated learning, as a well-known framework for collaborative training among distributed local sensors and devices, has been widely used in practical learning applications. To reduce communication resource consumption and training delay, acceleration training algorithms, especially momentum-based methods, are further developed for the training process. However, it is observed that under the influence of transmission noise, existing momentum methods exhibit poor training performance due to the noise accumulation along with the momentum term. This motivates us to propose a novel acceleration algorithm to achieve an efficient tradeoff between the training acceleration and noise smoothing. Specifically, to obtain clearer insights into the model update dynamics, we utilize a stochastic differential equation (SDE) model to mimic the discrete-time (DT) training trajectory. Through high-order drift approximation analysis on a general momentum-based SDE model, we propose a dynamic momentum weight and gradient stepsize design for the update rule, which is adaptive to both the training state and gradient quality. Such adaptation ensures that the training algorithm can seize good update opportunities and avoid noise explosion. The corresponding DT training algorithm is then derived via discretization of the proposed SDE model, which shows a superior training performance compared to state-of-the-art baselines.