Adaptive Client Sampling in Federated Learning via Online Learning with Bandit Feedback
This work addresses the communication cost issue in federated learning systems, offering an adaptive sampling method that is incremental but applicable to broader stochastic optimization procedures.
The paper tackles the problem of client sampling in federated learning by framing it as an online learning task with bandit feedback, using an online stochastic mirror descent algorithm to minimize sampling variance, and shows improved convergence speed over uniform sampling in experiments.
Due to the high cost of communication, federated learning (FL) systems need to sample a subset of clients that are involved in each round of training. As a result, client sampling plays an important role in FL systems as it affects the convergence rate of optimization algorithms used to train machine learning models. Despite its importance, there is limited work on how to sample clients effectively. In this paper, we cast client sampling as an online learning task with bandit feedback, which we solve with an online stochastic mirror descent (OSMD) algorithm designed to minimize the sampling variance. We then theoretically show how our sampling method can improve the convergence speed of federated optimization algorithms over the widely used uniform sampling. Through both simulated and real data experiments, we empirically illustrate the advantages of the proposed client sampling algorithm over uniform sampling and existing online learning-based sampling strategies. The proposed adaptive sampling procedure is applicable beyond the FL problem studied here and can be used to improve the performance of stochastic optimization procedures such as stochastic gradient descent and stochastic coordinate descent.