Probabilistic Client Sampling and Power Allocation for Wireless Federated Learning

Wen Xu, Ben Liang, Gary Boudreau, Hamza Umit Sokun

Abstract: Despite the many known benefits of Federated Learning (FL), in the wireless environment, its performance is significantly impacted by the statistical and system heterogeneities among the local data sets and local clients. Therefore, judicious sampling of clients and resource allocation among them are of vital importance in FL. In this work, we consider the online joint optimization of probabilistic client sampling and power allocation to improve the training performance of wireless FL. Our optimization is based on a new convergence bound for non-convex loss functions under probabilistic client sampling, which considers the different data ratios and gradient norms among clients. We propose a new algorithm based on the Lyapunov optimization framework, termed PCSPA, that accounts for how the statistical and system heterogeneities affect both the convergence rate and training time of FL, as well as the long-term power constraints and the expected number of sampled clients. Experiments on image classification with wireless FL show that the proposed algorithm can substantially outperform conventional separate optimization strategies and a state-of-the-art joint optimization method.