Adaptive Client Sampling in Federated Learning via Online Learning with Bandit Feedback
Jan 1, 2025·
,,
,,·
0 min read
Boxin Zhao
Ziqi Liu
Chaochao Chen
Mladen Kolar
Zhiqiang Zhang
Jun Zhou
Abstract
In federated learning (FL) problems, client sampling plays a key role in the convergence speed of training algorithm. However, while being an important problem in FL, client sampling is lack of study. In this paper, we propose an online learning with bandit feedback framework to understand the client sampling problem in FL. By adapting an Online Stochastic Mirror Descent algorithm to minimize the variance of gradient estimation, we propose a new adaptive client sampling algorithm. Besides, we use online ensemble method and doubling trick to automatically choose the tuning parameters in the algorithm. Theoretically, we show dynamic regret bound with comparator as the theoretically optimal sampling sequence; we also include the total variation of this sequence in our upper bound, which is a natural measure of the intrinsic difficulty of the problem. To the best of our knowledge, these theoretical contributions are novel to existing literature. Moreover, by implementing both synthetic and real data experiments, we show empirical evidence of the advantages of our proposed algorithms over widely-used uniform sampling and also other online learning based sampling strategies in previous studies. We also examine its robustness to the choice of tuning parameters. Finally, we discuss its possible extension to sampling without replacement and personalized FL objective. While the original goal is to solve client sampling problem, this work has more general applications on stochastic gradient descent and stochastic coordinate descent methods.
Type
Publication
Journal of Machine Learning Research, 26(8), 1–67

Authors
PhD (2020-2025)
Boxin Zhao was a PhD student in Econometrics and Statistics at University of Chicago, Booth School of Business. His research interests include probabilistic graphical models, functional data analysis and distributed learning, with a focus on developing novel methodologies with both practical applications and theoretical guarantees.

Authors
Professor of Data Sciences and Operations
Mladen Kolar is a Professor of Data Sciences and Operations at the University of Southern California Marshall School of Business and a Visiting Professor of Statistics and Data Science at Mohamed bin Zayed University of Artificial Intelligence. Before joining USC, he was on the faculty of the University of Chicago Booth School of Business. His research is focused on high-dimensional statistical methods, graphical models, varying-coefficient models and data mining, driven by the need to uncover interesting and scientifically meaningful structures from observational data. He is a Fellow of the Institute of Mathematical Statistics.