PaperlayerView exact revision on arXiv

Source-linked AI summary

Tackling System and Statistical Heterogeneity for Federated Learning with Adaptive Client Sampling

Bing Luo, Wenli Xiao, Shiqiang Wang, Jianwei Huang, Leandros Tassiulas

arXiv:2112.11256v1cs.LGcs.AIcs.DCcs.NImath.OC

TL;DR

Federated learning’s conventional client sampling can converge slowly in wall-clock time under system and statistical heterogeneity. This paper derives a convergence-based adaptive sampling framework, solves its resulting optimization approximately, and reports 73% less hardware-prototype time than uniform sampling for the same target loss.

  • Problem

    Existing convergence analyses mainly use uniform or data-size-proportional sampling, which can suffer slow wall-clock convergence under system and statistical heterogeneity.

  • Method

    The paper derives a tractable convergence bound for arbitrary sampling probabilities, learns unknown parameters, and approximately solves the resulting non-convex training-time problem.

  • Results

    73% less time than uniform sampling was required by the hardware-prototype scheme to reach the same target loss.

  • Takeaways & Limitations

    The sampling design characterizes how communication delays and data importance jointly shape client-selection probabilities for minimizing wall-clock convergence time.

Abstract

from arXiv · show

Federated learning (FL) algorithms usually sample a fraction of clients in each round (partial participation) when the number of participants is large and the server's communication bandwidth is limited. Recent works on the convergence analysis of FL have focused on unbiased client sampling, e.g., sampling uniformly at random, which suffers from slow wall-clock time for convergence due to high degrees of system heterogeneity and statistical heterogeneity. This paper aims to design an adaptive client sampling algorithm that tackles both system and statistical heterogeneity to minimize the wall-clock convergence time. We obtain a new tractable convergence bound for FL algorithms with arbitrary client sampling probabilities. Based on the bound, we analytically establish the relationship between the total learning time and sampling probabilities, which results in a non-convex optimization problem for training time minimization. We design an efficient algorithm for learning the unknown parameters in the convergence bound and develop a low-complexity algorithm to approximately solve the non-convex problem. Experimental results from both hardware prototype and simulation demonstrate that our proposed sampling scheme significantly reduces the convergence time compared to several baseline sampling schemes. Notably, our scheme in hardware prototype spends 73% less time than the uniform sampling baseline for reaching the same target loss.

I. INTRODUCTION

The paper targets slow wall-clock convergence in federated learning by jointly adapting client sampling to system and statistical heterogeneity. It derives a tractable optimization framework, develops practical algorithms, and reports faster convergence in simulation and hardware experiments.

  • Motivation: System heterogeneity lets stragglers lengthen physical training time, while statistical heterogeneity reflects unbalanced, non-i.i.d. client data.Both factors affect federated learning convergence and are illustrated as distinct challenges.
  • Motivation: Prior sampling methods largely ignore the joint impact of communication capability and data value on client selection.The paper frames this mismatch as the central limitation motivating adaptive sampling.
  • Research question: The paper asks how to design client sampling probabilities that address both heterogeneities while minimizing wall-clock convergence time.The problem is difficult because arbitrary-probability convergence analysis, straggler-aware time optimization, and unknown bound parameters must be handled together.
  • Approach: A tractable convergence upper bound relates arbitrary sampling probabilities to the rounds required for target precision, enabling a non-convex training-time minimization problem.The bound exposes sampling variance as an additional term relative to full client participation.
  • Approach: The proposed low-cost parameter-learning and approximate optimization algorithms account for communication time, data quantity, and data quality in sampling design.The method addresses unknown convergence-related parameters during learning.
  • Experiments: 73% less time than uniform sampling was required by the hardware-prototype scheme to reach the same target loss.Across simulation and prototype experiments, the proposed scheme reduced convergence time against several baselines.

III. PRELIMINARIES AND SYSTEM MODEL

The preliminaries describe FedAvg with partial participation and introduce adaptive sampling probabilities for heterogeneous federated learning. The model includes client data weights, local computation, communication time, and sampling with replacement.

  • A. Federated Learning (FL): The global objective weights each client according to its share of the total training data.Client i has n_i samples and weight p_i = n_i/n_tot.
  • A. Federated Learning (FL): FedAvg repeatedly samples clients, performs local SGD iterations, and aggregates their updated models at the server.The server broadcasts the current model, sampled clients compute local updates, and aggregation continues until convergence.
  • A. Federated Learning (FL): Partial participation is used because communication bandwidth is limited and devices are geographically dispersed.Existing analyses commonly use uniform or data-size-proportional sampling, which may suffer from stragglers and non-i.i.d. data.
  • B. Sampling Model: The proposed sampling distribution q assigns client-selection probabilities optimized for wall-clock convergence.The probabilities are positive, sum to one, and are intended to address both system and statistical heterogeneity.
  • B. Sampling Model: Each round samples K clients with replacement, so a client may appear multiple times and receive aggregation weight multiplied by its occurrence count.The sampled client set is therefore a multiset.

3) System Heterogeneity Model:

The system heterogeneity model represents each client's round time and the resulting synchronized wall-clock training objective. The formulation combines expected total learning time with a target expected-loss constraint, but requires an analytical convergence relation and faces straggler-induced non-convexity.

  • System Heterogeneity Model: Client i's round time t_i includes local computation and global communication, remains constant across rounds, and may differ across clients.The model sorts clients by ascending t_i; extending it to time-varying round times is left for future work.
  • System Heterogeneity Model: Synchronous FL round time is determined by the slowest sampled client, so stragglers directly affect each round's physical duration.The total time accumulates these per-round maxima across R rounds.
  • Problem Formulation: The optimization minimizes expected total learning time E[Ttot(q, R)] subject to expected global loss reaching F* with ϵ precision.Randomness comes from client sampling q and local SGD.
  • Problem Formulation: Problem P1 is difficult because q and R cannot generally predict the final loss before training, motivating an analytical expression in q and R.The resulting learning-time objective is also complicated by straggling and may be non-convex.
  • Problem Formulation: The paper addresses the first challenge by deriving a tractable convergence bound for arbitrary client sampling probabilities.This bound supports approximate algorithms for the broader training-time optimization problem.

A. Machine Learning Model Assumptions

The paper states convex-learning assumptions, then develops unbiased aggregation and a convergence analysis for arbitrary sampling probabilities. The resulting bound connects sampling probabilities to rounds needed for a target precision and generalizes uniform-sampling results.

  • Machine Learning Model Assumptions: The analysis assumes L-smoothness, μ-strong convexity, bounded local stochastic-gradient variance σ^2, and bounded expected squared stochastic-gradient norm G^2.The first three assumptions are common in convex FL analyses; the fourth permits client-specific G_i values.
  • Machine Learning Model Assumptions: Unlike prior assumptions using a universal gradient bound G, Assumption 4 allows each client to have a unique G_i, which informs the sampling design.Experiments also evaluate the approach with non-convex loss functions.
  • Aggregation with Arbitrary Client Sampling: Different sampling probabilities require inverse reweighting of client updates so the aggregated global model remains unbiased relative to full participation.The algorithm reduces to FedAvg under uniform sampling q_i = 1/N and to weighted sampling when q_i = p_i.
  • Main Convergence Result: The paper derives a convergence upper bound for arbitrary sampling under the stated assumptions and a decaying learning rate η_r = 2/[μ(γ+r)].The bound applies to the corresponding unbiased aggregation procedure.
  • Main Convergence Result: Sampling variance contributes an additional term relative to full participation, producing an explicit relationship between q and rounds R for reaching E[F] − F* ≤ ϵ.The bound generalizes convergence results for uniformly sampled clients.
  • Optimal Sampling: With homogeneous communication time t_i = t_0, minimizing total learning time is equivalent to minimizing communication rounds, yielding an optimal sampling problem solvable in closed form.The paper then extends the bound-based design to general heterogeneous systems.

V. OPTIMAL ADAPTIVE CLIENT SAMPLING ALGORITHM

The proposed algorithm first obtains an analytical expected-learning-time expression and an approximate problem from the convergence bound, then solves it efficiently using parameterized optimization and search.

  • Algorithm Design: The algorithm formulates an approximate optimization problem from the convergence upper bound and develops an efficient solver with insightful sampling principles.It first derives the expected total learning time as a function of sampling probabilities q and rounds R.

A. Analytical Expression for E[Ttot(q)]

The paper derives an expected total-learning-time expression for arbitrary client-sampling probabilities and transforms the resulting constrained problem into a more tractable approximation. The approximation supports an optimization problem whose solution is empirically reported to outperform baseline sampling schemes.

  • Convergence-time formulation: Theorem 2 expresses expected total learning time E[Ttot(q, R)] using arbitrary sampling probabilities q and training rounds R.The proof connects client straggler probabilities, expected per-round time, and total time over R rounds.
  • Convergence-time formulation: The convergence constraint is combined with the expected-time expression to formulate total-learning-time minimization over q and R.Problem P2 is more constrained than Problem P1, so every feasible solution of P2 is feasible for P1.
  • Tractable approximation: The original objective remains difficult because q appears in a polynomial sum of order K, motivating the approximation ˜E[T^(r)(q)].The approximation is exactly equal to the expected per-round time in two specified cases, including homogeneous client times.
  • Tractable approximation: The approximated problem P3 has a more straightforward objective, but solving it requires the parameters Gi, α, and β.These parameters are not directly known when the sampling probabilities must be optimized.
  • Tractable approximation: Empirical results show that solving P3 achieves superior total wall-clock-time performance compared with baseline client-sampling schemes.The paper uses P3 as an approximation of the original problem P1.

C. Solving Problem P3

The paper estimates unknown convergence-bound parameters during FL training using substitute sampling schemes and lightweight gradient-norm information. It then summarizes the resulting approximate client-sampling procedure in Algorithm 2.

  • Parameter estimation: Problem P3 is challenging because Gi, α, and β can only be obtained during the FL training process.This creates a chicken-and-egg problem: the sampling solution needs parameters that are learned while training.
  • Parameter estimation: Algorithm 2 estimates αβ using uniform sampling q1 and weighted sampling q2, where q_i=1/N and q_i=p_i, respectively.The schemes run until a predefined loss Fs rather than until the final precision target.
  • Assumptions and procedure: The method assumes that clients’ heterogeneous times ti and dataset sizes pi can be measured offline.The optimization only needs αβ rather than α and β separately because α can be divided out without changing the optimal sampling solution.
  • Parameter estimation: The server estimates Gi by receiving each sampled client’s local SGD gradient norm alongside its updated model.Only the gradient-norm value is additionally transmitted instead of full gradient information, adding little communication overhead.
  • Parameter estimation: Several predefined Fs values may be used to average parameter estimates and reduce the effect of sampling variance.The estimation process corresponds to Lines 1–6 of Algorithm 2.

2) Optimization Algorithm for q∗:

Because Problem P3 is non-convex, the paper introduces a control variable that yields a convex subproblem for fixed values and then searches over that variable.

  • Optimization strategy: Theorem 3 establishes that Problem P3 is non-convex because its objective Hessian is not positive semidefinite.The paper gives the N=2 case as an example in the proof sketch.
  • Optimization strategy: The paper defines a new control variable to transform the optimization into Problem P4.This variable enables separate treatment of the sampling probabilities and the control parameter.
  • Optimization strategy: For any fixed feasible M in [t1, tN], Problem P4 is convex in q because its objective is strictly convex and its constraints are linear.The optimal q*(M) can therefore be computed with a convex optimization tool such as CVX.
  • Optimization strategy: The algorithm performs a fixed-step linear search over [t1, tN] and uses the best M and corresponding q*(M) to approximate the global solution.The search domain is independent of the problem scale, including the number of clients N.

3) Property of Optimal Client Sampling:

The optimal sampling strategy links communication time and statistical importance: clients with shorter times and larger p_iG_i values receive higher sampling probabilities. A closed-form characterization is available in the β/α→0 case.

  • Optimal-sampling properties: Theorem 4 states that if ti≤tj and piGi≥pjGj, the optimal strategy assigns client i a higher sampling probability than client j.Thus, both lower communication time and larger statistical-importance product favor client sampling.
  • Optimal-sampling properties: Theorem 4 characterizes the interplay between system heterogeneity ti and statistical heterogeneity piGi in the optimal sampling strategy.The exact relationship is difficult to derive because Problem P3 is non-convex.
  • Closed-form case: When β/α→0, Corollary 2 gives a global optimal solution for Problem P3 in closed form.The solution follows from rewriting the objective and applying the Cauchy–Schwarz inequality.
  • Closed-form case: For homogeneous client times ti=t0, the closed-form solution reduces to the optimal sampling solution for homogeneous systems.This special case isolates the statistical-heterogeneity component of the sampling rule.

VI. EXPERIMENTAL EVALUATION

The evaluation uses hardware and simulated FL environments spanning logistic regression and CNN models, with multiple datasets and heterogeneous client settings. Results are compared across several client-sampling baselines.

  • Baselines: The study compares proposed sampling with full participation, uniform, weighted, and statistical sampling.The statistical baseline samples clients according to Corollary 1, while weighted sampling follows client data sizes in the benchmark description.
  • Evaluation Setup: Experiments use both a networked hardware prototype and a simulated environment.The prototype contains 40 Raspberry Pis and a laptop server; simulation uses 100 virtual devices.
  • Datasets and Models: Evaluations cover MNIST, EMNIST, and a synthetic dataset with logistic regression and CNN models.The CNN uses the LeNet-5 architecture, while logistic regression represents the convex setting.
  • Prototype Setup: The prototype uses unbalanced, non-i.i.d. EMNIST data across 40 devices, with device class counts ranging from 1 to 10.The 33,036 lower-case character samples follow a power-law distribution across devices.
  • Simulation Setup 1: Simulation Setup 1 uses 20,509 synthetic samples across 100 clients in a non-i.i.d. power-law distribution.It uses logistic regression and the Synthetic (1, 1) setting.
  • Simulation Setup 2: Simulation Setup 2 uses 15,129 MNIST samples across 100 clients with unbalanced, non-i.i.d. class distributions.It evaluates a CNN model where each device has 1–6 classes.

4) Training Parameters:

The experiments use fixed optimization and participation settings while varying communication delays across prototype and simulation environments. The proposed sampler reaches target loss and accuracy faster in wall-clock time, despite sometimes requiring more rounds.

  • Training Parameters: All experiments initialize w0 = 0, use batch size b = 24, and sample 10% of clients per round.Each client performs E = 50 local iterations, with K = 4 for the prototype and K = 10 for simulations.
  • Communication Settings: The prototype communication delay follows U(0.187, 7.159) seconds, while simulation delays follow exp(1) seconds.Prototype delays have mean 3.648 seconds and standard deviation 2.071 seconds; simulation delays each have mean and standard deviation of 1 second.
  • Evaluation Procedure: Wall-clock global loss and test accuracy are averaged over 50 independent runs using matched random seeds within each run.Random seeds vary across runs.
  • Loss with Wall-clock Time: The proposed sampler achieves the same target loss with significantly less time than baseline sampling schemes.This pattern is reported across Prototype Setup and both Simulation Setups.
  • Robustness Checks: The experiments vary E and K and report similar performance to the settings shown.These additional experiments are omitted from the paper because of page limitations.
  • Loss with Wall-clock Time: 73% less time is required than full and uniform sampling in Prototype Setup, while weighted and statistical sampling require around 66% more time.The comparison concerns reaching the same target loss in Fig. 3(a).
  • Accuracy with Wall-clock Time: The proposed sampler reaches target test accuracy faster than other benchmarks, including 70% less time than full sampling in Simulation Setup 1.For the 75.3% target accuracy, it takes around 46% less time than the other sampling schemes.
  • Loss with Wall-clock Time: The proposed sampler requires more training rounds than statistical sampling and full participation to reach target loss.Its objective is wall-clock time rather than minimizing the number of rounds; statistical sampling performs better in homogeneous systems.
Loading 2112.11256v1…