Ranking-based Client Selection with Imitation Learning for Efficient Federated Learning

TL;DR

This paper addresses the challenge of heterogeneity in federated learning by introducing FedRank, a ranking-based device selection framework that uses offline imitation learning against analytical baselines to bootstrap a pairwise ranking model. The method casts client selection as an episodic imitation learning problem and employs a RankNet-based pairwise loss during online RL to prioritize the relative quality of devices, while an early-exit probing scheme reduces state-gathering overhead. Empirical results show FedRank achieves substantial gains in accuracy, convergence speed, and energy efficiency across IID and non-IID settings, outperforming heuristic and learning-based baselines. The approach offers a scalable, energy-conscious path to deploying FL in real-world, heterogeneous device environments. It also demonstrates generalization to unseen deployments and highlights the value of combining imitation learning with pairwise ranking for robust client selection.

Abstract

Federated Learning (FL) enables multiple devices to collaboratively train a shared model while ensuring data privacy. The selection of participating devices in each training round critically affects both the model performance and training efficiency, especially given the vast heterogeneity in training capabilities and data distribution across devices. To address these challenges, we introduce a novel device selection solution called FedRank, which is an end-to-end, ranking-based approach that is pre-trained by imitation learning against state-of-the-art analytical approaches. It not only considers data and system heterogeneity at runtime but also adaptively and efficiently chooses the most suitable clients for model training. Specifically, FedRank views client selection in FL as a ranking problem and employs a pairwise training strategy for the smart selection process. Additionally, an imitation learning-based approach is designed to counteract the cold-start issues often seen in state-of-the-art learning-based approaches. Experimental results reveal that \model~ boosts model accuracy by 5.2\% to 56.9\%, accelerates the training convergence up to $2.01 \times$ and saves the energy consumption up to $40.1\%$.

Figures (9)

• Figure 1: Illustration of FedRank. (a) Learning-based approaches: tackle device selection with data-driven processes that holistically trade-off demands. (b) Heuristic-based approaches: hand-developed heuristics perform well in the specific, straightforward deployment configurations they were designed for. (c) FedRank: utilizes imitation learning and a ranking approach to optimize device selection.
• Figure 2: Workflow and real-world deployment barriers of Federated Learning.
• Figure 3: Efficiency comparison of various schemes to train the ResNet18 model on CIFAR10 (IID). We compared the FedRank with the SOTA in the three baselines, respectively.
• Figure 4: Generalization ability of FedRank to train the ResNet18 model on CIFAR10 (IID).
• Figure 5: Ablation study explores the effects of imitation learning and pairwise loss on model performance across several datasets (MNIST, CIFAR10, CINIC, and TinyImageNet) under Non-IID settings. $\text{Model}^{-I}$ (without imitation learning), $\text{Model}^{-P}$ (without pairwise loss), and $\text{Model}^{-IP}$ (lacking both features).