Communication-Efficient Federated Learning over Wireless Channels via Gradient Sketching

TL;DR

This paper introduces Federated Proximal Sketching (FPS), a novel method for federated learning over wireless channels that efficiently compresses data, handles heterogeneity, and maintains convergence despite noisy communications.

Contribution

The paper proposes FPS, a new federated learning algorithm that uses gradient sketching and modified loss functions to address bandwidth, noise, and heterogeneity challenges.

Findings

FPS achieves stable convergence under wireless channel noise.

FPS outperforms existing methods in accuracy and efficiency.

Numerical experiments validate FPS's robustness on real-world data.

Abstract

Large-scale federated learning (FL) over wireless multiple access channels (MACs) has emerged as a crucial learning paradigm with a wide range of applications. However, its widespread adoption is hindered by several major challenges, including limited bandwidth shared by many edge devices, noisy and erroneous wireless communications, and heterogeneous datasets with different distributions across edge devices. To overcome these fundamental challenges, we propose Federated Proximal Sketching (FPS), tailored towards band-limited wireless channels and handling data heterogeneity across edge devices. FPS uses a count sketch data structure to address the bandwidth bottleneck and enable efficient compression while maintaining accurate estimation of significant coordinates. Additionally, we modify the loss function in FPS such that it is equipped to deal with varying degrees of data heterogeneity. We establish the convergence of the FPS algorithm under mild technical conditions and characterize how the bias induced due to factors like data heterogeneity and noisy wireless channels play a role in the overall result. We complement the proposed theoretical framework with numerical experiments that demonstrate the stability, accuracy, and efficiency of FPS in comparison to state-of-the-art methods on both synthetic and real-world datasets. Overall, our results show that FPS is a promising solution to tackling the above challenges of FL over wireless MACs.