PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Armed Bandit Based Client Scheduling for Federated Learning

Wenchao Xia, Tony Q. S. Quek, Kun Guo, Wanli Wen, Howard H. Yang, Hongbo Zhu

arXiv:2007.02315v1cs.ITcs.LG

TL;DR

Federated learning faces substantial latency from repeated communication rounds and client-side computation. This paper develops online multi-armed-bandit client scheduling without wireless CSI or client statistical characteristics, using CS-UCB for ideal conditions and CS-UCB-Q for non-ideal conditions. The algorithms have logarithmic or conditionally sub-linear regret bounds, and the paper analyzes convergence and validates efficiency through simulations.

  • Problem

    Federated learning requires hundreds to thousands of communication rounds, creating a latency bottleneck that motivates minimizing total training time.

  • Method

    The paper formulates online client scheduling as a multi-armed bandit problem and proposes CS-UCB for ideal settings and CS-UCB-Q, combining UCB with virtual queues, for non-ideal settings.

  • Results

    CS-UCB has an expected performance regret bound growing logarithmically over communication rounds, while CS-UCB-Q can have sub-linear regret growth under certain conditions.

  • Takeaways & Limitations

    The framework provides a client-scheduling approach for ideal and non-ideal federated-learning scenarios, with analyzed regret and convergence performance and simulation validation.

Abstract

from arXiv · show

By exploiting the computing power and local data of distributed clients, federated learning (FL) features ubiquitous properties such as reduction of communication overhead and preserving data privacy. In each communication round of FL, the clients update local models based on their own data and upload their local updates via wireless channels. However, latency caused by hundreds to thousands of communication rounds remains a bottleneck in FL. To minimize the training latency, this work provides a multi-armed bandit-based framework for online client scheduling (CS) in FL without knowing wireless channel state information and statistical characteristics of clients. Firstly, we propose a CS algorithm based on the upper confidence bound policy (CS-UCB) for ideal scenarios where local datasets of clients are independent and identically distributed (i.i.d.) and balanced. An upper bound of the expected performance regret of the proposed CS-UCB algorithm is provided, which indicates that the regret grows logarithmically over communication rounds. Then, to address non-ideal scenarios with non-i.i.d. and unbalanced properties of local datasets and varying availability of clients, we further propose a CS algorithm based on the UCB policy and virtual queue technique (CS-UCB-Q). An upper bound is also derived, which shows that the expected performance regret of the proposed CS-UCB-Q algorithm can have a sub-linear growth over communication rounds under certain conditions. Besides, the convergence performance of FL training is also analyzed. Finally, simulation results validate the efficiency of the proposed algorithms.

I. INTRODUCTION

The paper frames client scheduling as a way to reduce federated-learning latency when wireless channels are limited, client conditions vary, and prior system information is unavailable. It formulates online scheduling as a multi-armed bandit problem and develops UCB-based algorithms for ideal and non-ideal settings, with regret and convergence analyses.

  • Motivation: Federated learning latency remains a bottleneck because training requires hundreds to thousands of communication rounds, each involving wireless transmission and local computation.Dynamic wireless propagation and limited client computing power can degrade model performance under a limited training-time budget.
  • Limitations of Existing Work: Existing scheduling approaches commonly rely on prior wireless CSI or client computing-resource information, which is difficult to obtain at large scale.Many prior works also focus on transmission time without accounting for local-training computation time.
  • Motivation: Client scheduling is needed to select appropriate clients when many clients compete for a limited number of wireless channels.The objective is to minimize total wall-clock training time, including transmission and local computation.
  • Proposed Framework: The paper reformulates client scheduling as a multi-armed bandit problem to learn clients’ statistical information online without knowing wireless CSI or computing-resource dynamics.This is presented as the first application of the MAB tool to client scheduling in federated learning.
  • Ideal Scenario: For balanced i.i.d. datasets and always-available clients, CS-UCB uses the UCB policy and has expected performance regret growing logarithmically over communication rounds.The policy balances exploiting actions that performed well with exploring actions that may yield higher future rewards.
  • Non-Ideal Scenario: For non-i.i.d. and unbalanced data with varying client availability, CS-UCB-Q combines UCB with virtual queues and can achieve sub-linear expected regret under certain conditions.The non-ideal formulation also introduces a fairness constraint requiring each client to participate in a certain proportion of rounds.
  • Convergence Analysis: The convergence analysis finds that more participating clients can improve the ideal-scenario convergence rate, while non-ideal convergence depends only weakly on that number.The analysis also reveals a tradeoff between performance regret and the speed of convergence toward satisfying fairness constraints.

II. SYSTEM MODEL AND PROBLEM FORMULATION

The system models federated learning over a wireless network where an access point selects clients, distributes the global model, and aggregates locally computed updates. Training latency combines communication and local-computation delays, whose client-side dynamics are unknown and observed through total round time.

  • The network has a single-antenna access point, K clients, and N orthogonal channels, with N < K and no inter-client interference.
  • FL training minimizes a weighted global loss formed from the clients’ local losses and private datasets.
  • Each client’s local loss is computed over its sample points using a parameterized loss function.
  • Each communication round selects a client subset, distributes the global model, receives private-data updates, and aggregates the uploaded gradients.
  • A. Training Latency: A round’s client latency includes distribution time, local update time, and transmission time, while aggregation delay is treated as negligible.
  • A. Training Latency: The round interval is capped by τmax to prevent indefinite waiting caused by stragglers, and distribution time depends on model size and wireless channel conditions.
  • A. Training Latency: Client computing power varies with concurrent tasks; its available rate φk(t) is modeled as i.i.d. over time with unknown expectation.
  • A. Training Latency: Rather than separately estimating latency components, the access point directly observes each client’s total latency τk(t), reducing measurement difficulty and overhead.

B. Availability Constraint

The availability constraint restricts scheduling to clients currently able to participate, while fairness constraints preserve participation for clients with heterogeneous data importance and activity.

  • Clients may become temporarily unavailable because of poor wireless channels or exhausted computing power.
  • The availability indicator ak(t) equals one when client k is available, defining A(t) as the available-client set revealed at each round’s start.
  • Availability sets are assumed i.i.d. over time, but their distribution is unknown in advance.
  • Because clients outnumber channels, the access point must select a subset from the available clients.
  • Non-i.i.d. and unbalanced local datasets, differing dataset sizes, and heterogeneous activity make client importance uneven.
  • Larger datasets with distributions more similar to the global distribution are treated as more important and should participate more often.
  • The fairness indicator bk(t) records whether client k is selected, with required participation fraction ck constrained to [0, 1).

D. Problem Formulation

The paper formulates latency-aware client scheduling as sequential decision-making under availability and fairness constraints. It then recasts the ideal case as a combinatorial multi-armed bandit problem evaluated by cumulative regret.

  • D. Problem Formulation: The objective is to minimize total FL training time, with each round’s duration determined after selected clients finish uploading.
  • D. Problem Formulation: In the ideal scenario, all clients are always available and datasets are i.i.d. and balanced, so availability and fairness constraints are unnecessary.
  • D. Problem Formulation: The non-ideal formulation incorporates temporary unavailability, non-i.i.d. unbalanced datasets, availability constraints, and fairness constraints.
  • D. Problem Formulation: The formulation selects as many clients as possible per round because participating-client count has weak convergence-rate dependence in the non-i.i.d. setting.
  • D. Problem Formulation: Optimal scheduling is difficult because computing-power dynamics and wireless CSI are unknown, requiring information learned from feedback.
  • A. MAB Problems: The ideal problem is reformulated as an MAB problem in which the access point chooses client subsets as sequential actions.
  • B. Problem Reformulation: The selected subset is a super arm with normalized reward r(S(t)) = −τ̄(S(t)) + 1, so maximizing reward corresponds to reducing latency.
  • B. Problem Reformulation: Policy quality is measured by regret, the difference between optimal expected reward and policy reward; minimizing regret is equivalent to maximizing total reward.

C. Proposed Algorithm

The proposed CS-UCB algorithm learns client rewards through exploration and exploitation, using arm-level feedback to select client subsets. Its analysis establishes logarithmic regret and relates participation to convergence.

  • CS-UCB is based on the UCB policy and selects a super arm using empirical rewards plus an exploration term.
  • Arm-level feedback can be reused across different super arms, allowing information from multiple actions to improve future decisions.
  • The vectors y(t) and z(t) store each arm’s observed-reward sample mean and play count, respectively, initialized at zero.
  • The algorithm initializes by ensuring every arm is played at least once before entering its main selection loop.
  • The combinatorial action space grows exponentially with K, making the generic optimization potentially expensive despite the specialized linear structure.
  • Because rewards are linear and arm rewards are independent, sorting individual rewards and choosing the top N clients finds the best super arm.
  • Theorem 1 bounds CS-UCB’s expected regret, with the bound expressed through N, K, and ln T.
  • The regret grows as O(N^3K ln T), polynomially in channels, linearly in clients, and logarithmically in communication rounds.

IV. SCHEDULING IN NON-IDEAL SCENARIO

The non-ideal scheduling problem incorporates client availability and fairness constraints into a multi-armed bandit formulation. Because direct UCB application cannot satisfy fairness, the problem is reformulated using randomized availability-aware policies and regret minimization.

  • Client availability and fairness constraints are jointly considered in the non-ideal scheduling problem.
  • Direct application of Algorithm 1 is insufficient because it cannot handle the fairness constraint.Availability can be handled by checking clients before action selection, whereas fairness requires an additional approach.
  • Virtual queues are introduced to enforce fairness while scheduling available clients.The approach follows prior virtual-queue techniques for constrained bandit problems.
  • An A-only policy selects super arms as a function of the observed available-client set, with q_S(e) defining selection probabilities.The policy is represented by probability distributions over feasible super arms for each availability set.
  • A feasible A-policy exists that can meet the fairness constraint when the requirement is feasible.The feasibility condition requires c to lie strictly inside the maximal feasibility region C.
  • The resulting objective minimizes cumulative regret relative to the optimal policy after training convergence.The optimal reward is defined under known mean rewards, while online scheduling must estimate arm rewards through exploration and exploitation.

B. Proposed Algorithm

CS-UCB-Q combines UCB reward estimation with virtual queues to schedule available clients while enforcing long-term fairness. Its analysis establishes feasibility and sub-linear regret under stated conditions, alongside convergence implications for federated learning.

  • CS-UCB-Q combines the UCB policy with virtual queues to address client scheduling under fairness constraints.The algorithm is presented as Algorithm 2 for the constrained scheduling problem.
  • The scheduler chooses a feasible super arm by maximizing a compound value built from truncated UCB estimates and virtual-queue lengths.A weighting factor β ∈ [0, 1] balances reward optimization against fairness enforcement.
  • Virtual queue D_k tracks each client’s fairness deficit, decreasing after selection and increasing by c_k each round.The queue is initialized at zero and uses [x]+ = max{x, 0}.
  • Theorem 3 states that CS-UCB-Q is feasibility-optimal and satisfies fairness constraints for any c strictly inside C.Thus, long-term fairness is achieved whenever the requested constraint is feasible.
  • The expected performance gap becomes sub-linear over communication rounds under the theorem’s conditions.The regret bound contains terms whose balance is controlled by β; larger β prioritizes fairness.
  • Convergence analysis shows that smaller δ3 improves convergence performance, while increasing participating clients has only weak rate dependence.Fairness can prioritize clients with important local datasets to reduce δ3.

V. NUMERICAL RESULTS

The experiments evaluate the proposed scheduling algorithms using wireless communication, computation, and federated-learning settings. The simulation uses MNIST logistic regression with explicitly specified channel, power, bandwidth, and optimization parameters.

  • Experiments evaluate CS-UCB and CS-UCB-Q using client scheduling costs derived from wireless transmission and local-update times.The study uses average rather than instantaneous distribution rates to model long-term learning acceleration.
  • Each client’s round cost combines distribution, upload, and local-update components.Distribution uses model size, channel bandwidth, and average downlink rate; upload uses data size and average uplink rate.
  • Wireless channels vary during distribution when a round spans multiple coherence-time slots, so average distribution rates are used.The model includes both small-scale and large-scale fading for uplink and downlink links.
  • The default system assumes equal 5 × 10^3-bit distribution and upload payloads and a maximum interval τmax = 5 seconds.The network uses a single-antenna access point and clients distributed within a 500 m disc.
  • The simulations use MNIST handwritten-digit classification with multinomial logistic regression.The dataset contains 60,000 training samples and 10,000 test samples, with SGD mini-batches of size 2 and learning rate 0.001.

B. Performance in Ideal Scenario

In ideal and non-ideal simulations, the proposed algorithms improve scheduling performance while addressing convergence, availability, and fairness. CS-UCB reduces wall-clock training time, whereas CS-UCB-Q preserves learning performance and satisfies feasible fairness constraints.

  • B. Performance in Ideal Scenario: CS-UCB outperforms two baseline algorithms after 20 communication rounds, with the performance gap increasing thereafter.The authors attribute this to learning client statistics while balancing exploitation and exploration.
  • B. Performance in Ideal Scenario: CS-UCB produces fewer failed clients, allowing more clients to participate within a given time interval.A client is counted as failed when its round time reaches τmax = 5 seconds.
  • B. Performance in Ideal Scenario: CS-UCB has a convergence rate close to the baselines in communication rounds under i.i.d. client datasets.It nevertheless completes the required rounds for a target accuracy in less wall-clock time.
  • B. Performance in Ideal Scenario: CS-UCB performs better than random scheduling and round robin as the number of available channels varies.The cumulative performance gap increases with more channels because high-latency clients become more likely to be selected.
  • C. Performance in Non-Ideal Scenario: CS-UCB-Q achieves better test accuracy than three baselines in the heterogeneous setting by selecting clients with important datasets more often.It reduces training latency without degrading learning performance and can finish at a target accuracy in less wall-clock time.

VI. CONCLUSIONS

The paper addresses client scheduling in ideal and non-ideal federated-learning scenarios to reduce training time. It derives regret bounds, analyzes convergence, and validates the algorithms through simulations.

  • Contributions: The work formulates client scheduling for both ideal and non-ideal federated-learning scenarios.The ideal and non-ideal settings differ in client-data properties and availability.
  • Ideal scenario: The CS-UCB algorithm receives an upper regret bound that grows logarithmically over communication rounds.This result applies to the ideal scenario.
  • Non-ideal scenario: The non-ideal scenario includes non-i.i.d., unbalanced client datasets and dynamic client availability.The conclusion motivates an explicit fairness constraint for client participation.
  • Non-ideal scenario: The CS-UCB-Q algorithm combines the UCB policy with a virtual queue technique and has an upper regret bound.The virtual-queue mechanism is used with a fairness constraint during training.
  • Results: The expected regret of CS-UCB-Q can grow sub-linearly, while convergence performance is analyzed in both scenarios.The paper also reports a tradeoff between performance regret and convergence speed.
  • Results: Simulation results validate the efficiency of the proposed CS-UCB and CS-UCB-Q algorithms.The reported simulations also examine test accuracy and performance comparisons involving communication rounds and channel counts.

APPENDIX A

The appendix develops the technical analyses underlying the scheduling and convergence results. It uses counters, concentration bounds, contraction arguments, and virtual-queue stability analysis.

  • Regret analysis: The regret proof tracks client-selection counters and distinguishes optimal from non-optimal super-arm selections.Counters are updated when non-optimal super arms are selected, with random tie handling when needed.
  • Regret analysis: The proof bounds selection errors using sample means, confidence terms, and Chernoff-Hoeffding concentration arguments.The analysis bounds probabilities of the relevant inequalities and chooses a logarithmic counter threshold.
  • Convergence analysis: The convergence analysis constructs a local linearized function and obtains a contraction inequality under 0 < γΦ ≤ 1.The optimal point and minimum value of the constructed function are used in the argument.

2. Define S′(t) as the super

This appendix section bounds the drift-plus-regret terms used in the non-ideal scheduling analysis. The resulting inequalities support the expected-regret and virtual-queue stability results.

  • Regret bound: The analysis defines an auxiliary comparison involving selected and optimal super arms to relate scheduling decisions to regret.Indicator variables identify the super arms selected in each round.
  • Regret bound: A lemma bounds the cumulative expectation of Υ1(t) through a drift-plus-regret construction.The proof sums the conditional bounds over rounds and uses the Lyapunov function at the initial and final times.
  • Regret bound: The appendix combines the intermediate inequalities to obtain an upper bound on expected regret.The bound is assembled by substituting the preceding lemmas and inequalities into the regret expression.
  • Stability analysis: The stability proof uses assumptions on the system and bounds on the second moment of the decision vector.The resulting Lyapunov-drift argument completes the technical proof under the stated assumptions.
Loading 2007.02315v1…