Adaptive Client Sampling in Federated Learning via Online Learning with Bandit Feedback

TL;DR

This work tackles the challenge of high communication costs in Federated Learning by proposing an adaptive client-sampling strategy that cast client selection as an online learning problem with bandit feedback. The authors develop the OSMD Sampler, operating on the probability simplex with an unnormalized negative entropy regularizer, and construct unbiased gradient estimators from partial feedback to minimize sampling variance. They prove dynamic regret bounds and provide convergence guarantees for common FL algorithms such as mini-batch SGD and FedAvg when paired with OSMD sampling, showing improvements over uniform sampling especially under client heterogeneity. Extensive simulations and real-data experiments demonstrate that Adaptive-OSMD consistently outperforms baselines across different heterogeneity regimes and remains robust to hyperparameter choices, highlighting its broad applicability beyond FL to stochastic optimization.

Abstract

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.

Figures (8)

• Figure 1: The training loss (top row) and cumulative regret (bottom row) are compared for the Adaptive-OSMD Sampler, Uniform Sampler, and Optimal Sampler, under $\sigma = 1.0$, $\sigma = 3.0$, and $\sigma = 10.0$. Solid lines represent the mean values, while the shaded regions indicate $\text{mean} \pm \text{standard deviation}$ across independent runs.
• Figure 2: The training loss (top row) and cumulative regret (bottom row) are compared across the Adaptive-OSMD Sampler, MABS, VRB, and Avare methods for $\sigma = 1.0$, $\sigma = 3.0$, and $\sigma = 10.0$. Solid lines represent the mean values, while shaded regions indicate $\text{mean} \pm \text{standard deviation}$ across independent runs.
• Figure 3: The training loss (top row) and cumulative regret (bottom row) are shown for the Adaptive-OSMD Sampler with different values of $\alpha$ under $\sigma = 1.0$, $\sigma = 3.0$, and $\sigma = 10.0$. Solid lines represent the mean values, while shaded regions indicate $\text{mean} \pm \text{standard deviation}$ across independent runs.
• Figure 4: The training loss (top row) and cumulative regret (bottom row) are compared between the Adaptive-OSMD Sampler and $p^{\text{IS}}$ for $\nu = 1.0$, $\nu = 3.0$, and $\nu = 10.0$. Solid lines represent the mean values, while shaded regions illustrate $\text{mean} \pm \text{standard deviation}$ across independent runs.
• Figure 5: The sample size distribution in the training set across clients.

Theorems & Definitions (16)

• Theorem 4
• Theorem 5
• Theorem 6
• Corollary 7
• Theorem 8
• Theorem 9
• Theorem 10
• Proposition 11: Proposition 3 of el2020adaptive
• Proposition 12
• Proposition 13