[Paper Review] Context-Aware Online Client Selection for Hierarchical Federated Learning

Zhe Qu, Rui Duan+4 more|arXiv (Cornell University)|Dec 1, 2021

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes Context-aware Online Client Selection (COCS), a learning-based client selection policy for Hierarchical Federated Learning (HFL) that uses Contextual Combinatorial Multi-Armed Bandits (CC-MAB) to dynamically select clients based on real-time context such as client computation capacity and wireless channel conditions. COCS achieves sublinear regret and significantly improves training performance over baselines, with 67.93% test accuracy on CIFAR-10—within 0.48% of the Oracle policy—under non-convex settings.

ABSTRACT

Federated Learning (FL) has been considered as an appealing framework to tackle data privacy issues of mobile devices compared to conventional Machine Learning (ML). Using Edge Servers (ESs) as intermediaries to perform model aggregation in proximity can reduce the transmission overhead, and it enables great potentials in low-latency FL, where the hierarchical architecture of FL (HFL) has been attracted more attention. Designing a proper client selection policy can significantly improve training performance, and it has been extensively used in FL studies. However, to the best of our knowledge, there are no studies focusing on HFL. In addition, client selection for HFL faces more challenges than conventional FL, e.g., the time-varying connection of client-ES pairs and the limited budget of the Network Operator (NO). In this paper, we investigate a client selection problem for HFL, where the NO learns the number of successful participating clients to improve the training performance (i.e., select as many clients in each round) as well as under the limited budget on each ES. An online policy, called Context-aware Online Client Selection (COCS), is developed based on Contextual Combinatorial Multi-Armed Bandit (CC-MAB). COCS observes the side-information (context) of local computing and transmission of client-ES pairs and makes client selection decisions to maximize NO's utility given a limited budget. Theoretically, COCS achieves a sublinear regret compared to an Oracle policy on both strongly convex and non-convex HFL. Simulation results also support the efficiency of the proposed COCS policy on real-world datasets.

Motivation & Objective

• Address the lack of client selection strategies tailored for Hierarchical Federated Learning (HFL), which faces unique challenges like time-varying client-ES connectivity and limited network operator (NO) budgets.
• Improve training performance in HFL by maximizing the number of successfully participating clients per round while respecting per-ES budget constraints.
• Design an online, adaptive client selection policy that learns from real-time context (e.g., transmission quality, computation availability) without prior knowledge of client performance.
• Overcome the limitations of static optimization-based policies by introducing a learning-based approach that adapts to dynamic and uncertain HFL network conditions.

Proposed method

• Develop COCS based on the Contextual Combinatorial Multi-Armed Bandit (CC-MAB) framework to model client selection as a sequential decision-making problem under uncertainty.
• Use side information (context) such as client computation resources, distance to edge server (ES), and wireless channel bandwidth to estimate the probability of successful client participation.
• Formulate the NO’s utility as a function of the number of successfully participating clients, with constraints on the total rental budget per ES.
• Implement a bandit-based learning mechanism that balances exploration (testing underperforming client-ES pairs) and exploitation (selecting high-probability successful pairs) to maximize cumulative utility.
• Apply a partitioning strategy on the context space to group similar client-ES pairs into hypercubes, enabling efficient estimation of participation probabilities.
• Theoretical analysis proves that COCS achieves sublinear regret compared to an Oracle policy in both strongly convex and non-convex HFL settings.

Experimental results

Research questions

• RQ1 How can client selection in HFL be optimized under time-varying client-ES connectivity and limited network operator budgets?
• RQ2 Can a learning-based approach outperform static or optimization-based client selection policies in HFL when network conditions are uncertain and dynamic?
• RQ3 To what extent does incorporating contextual information (e.g., computation capacity, channel quality) improve client selection performance in HFL?
• RQ4 What is the theoretical performance guarantee of an online client selection policy in HFL, and does it converge to near-optimal decisions over time?

Key findings

• COCS achieves sublinear regret in both strongly convex and non-convex HFL settings, indicating asymptotically optimal performance relative to an Oracle policy.
• On the CIFAR-10 dataset with a CNN model, COCS achieves 67.93% test accuracy, outperforming LinUCB (65.19%), CUCB (63.76%), and Random (62.25%) policies.
• COCS reduces the number of required training rounds by 10% compared to the Oracle policy in the non-convex setting, converging in 101 rounds versus 92 for the Oracle.
• Cumulative utility of COCS reaches 5,378 under a deadline of 8 seconds, showing diminishing returns beyond a certain deadline or budget increase.
• Increasing the budget from 3.5 to 5 units improves cumulative utility, but further increases yield minimal gains due to inherent bottlenecks in low-throughput client-ES pairs.
• The performance gap between COCS and the Oracle policy is smaller in non-convex settings (0.48% accuracy difference) than in convex settings, due to the use of an approximated Oracle in the former.