Federated learning with heavy-tailed gradient noise and communication noise: a variance-reduction based algorithm

In large-scale federated learning (FL) deployments, such as those involving Internet of Things (IoT) devices or wireless networks, the training process is often disrupted by two types of noise: heavy-tailed gradient noise (arising from data heterogeneity or privacy-preserving mechanisms) and heavy-tailed communication noise (caused by channel fading or interference). The authors seek to develop a robust optimization algorithm that can maintain stable convergence despite these noise sources.

The authors propose VRA-FedSGD, an algorithm that integrates three key mechanisms to suppress noise:

1. Variance-Reduced Aggregation (VRA): A mechanism that mitigates downlink communication noise by having clients maintain a local estimate of the global model that is updated using a weighted average of the noisy broadcast.
2. Momentum Variance Reduction: A technique that uses a nonlinear mapping (such as clipping or normalization) to handle heavy-tailed gradient noise, preventing extreme gradient values from destabilizing the model update.
3. Partial Client Participation: The algorithm is designed to handle scenarios where only a subset of clients participates in each round, improving communication efficiency.

The authors provide a theoretical analysis of the convergence rates for both non-convex and strongly convex objective functions, establishing bounds in both the mean sense and the almost sure sense.

VRA-FedSGD demonstrates strong theoretical performance. For non-convex objective functions, it achieves a convergence rate of O(K^{-(p-1)/(2p-1)}), where p is the tail index of the noise. For strongly convex objectives, it establishes an almost sure convergence rate of Õ(K^{-(1-1/(p-ε))}). These results generalize previous findings from the finite-variance setting to the more challenging heavy-tailed noise regime. Empirical simulations on logistic regression tasks confirm that the algorithm effectively maintains performance where standard FedSGD would otherwise diverge.