Joint User Association and Resource Allocation for Wireless Hierarchical Federated Learning With IID and Non-IID Data

Abstract

In this work, hierarchical federated learning (HFL) over wireless multi-cell networks is proposed for large-scale model training while preserving data privacy. However, the imbalanced data distribution has a significant impact on the convergence rate and learning accuracy. In addition, a large learning latency is incurred due to the traffic load imbalance among base stations (BSs) and limited wireless resources. To cope with these challenges, we first provide an analysis of the model error and learning latency in wireless HFL. Then, joint user association and wireless resource allocation algorithms are investigated under independent identically distributed (IID) and non-IID training data, respectively. For the IID case, a learning latency aware strategy is designed to minimize the learning latency by optimizing user association and wireless resource allocation, where a mobile device selects the BS with the maximal uplink channel signal-to-noise ratio (SNR). For the non-IID case, the total data distribution distance and learning latency are jointly minimized to achieve the optimal user association and resource allocation. The results show that both data distribution and uplink channel SNR should be taken into consideration for user association in the non-IID case. Finally, the effectiveness of the proposed algorithms are demonstrated by the simulations.