Enhancing Federated Learning Through Layer-Wise Aggregation Over Non-IID Data

Yang Xu, Ying Zhu, Zhiyuan Wang, Hongli Xu, Yunming Liao

Abstract: Nowadays, federated learning (FL) has been widely adopted to train deep neural networks (DNNs) among massive devices without revealing their local data in edge computing (EC). To relieve the communication bottleneck of the central server in FL, hierarchical federated learning (HFL), which leverages edge servers as intermediaries to perform model aggregation among devices in proximity, comes into being. Nevertheless, the existing HFL systems may not perform training effectively due to bandwidth constraints and non-IID issues on devices. To conquer these challenges, we introduce an HFL system with device-edge assignment and layer selection, namely Heal. Specifically, Heal organizes all the devices into a hierarchical structure (i.e., device-edge assignment) and enables each device to forward only a sub-model with several valuable layers for aggregation (i.e., layer selection). This processing procedure is called layer-wise aggregation. To further save communication resource and improve the convergence performance, we then design an iteration-based algorithm to optimize the development of our layer-wise aggregation strategy by considering the data distribution as well as resource constraints among devices. Extensive experiments on both the physical platform and the simulated environment show that Heal accelerates DNN training by about 1.4–12.5×, and reduces the network traffic consumption by about 31.9–64.1%, compared with the existing HFL systems.