Channel-Aware Joint AoI and Diversity Optimization for Client Scheduling in Federated Learning With Non-IID Datasets

Federated learning (FL) is a distributed learning framework where clients jointly train a global model without sharing their local datasets. In each communication round of FL, a subset of clients are scheduled to participate in training. Recent research has shown that diversity-based FL can improve the convergence performance of FL, especially when the client datasets are not independent and identically distributed (non-IID). In this paper, we show that by considering the channel state information and age of information (AoI) of each client, the convergence of FL can further be improved. We formulate a channel-aware joint AoI and diversity-based client scheduling problem as a constrained Markov decision process (CMDP). By using Lagrangian index and one-step lookahead approaches, we develop a two-stage online algorithm which is scalable and has a low computational complexity. For FL tasks with non-IID client datasets, our results show that the proposed algorithm can speed up the convergence of FL by up to 71%, through reducing the duration of uplink transmission, when compared with three state-of-the-art FL algorithms.