Is Partial Model Aggregation Energy-Efficient for Federated Learning Enabled Wireless Networks?

Abstract

This work aims to address two of the main challenges for federated learning (FL), i.e., the limited communication resources and the data heterogeneity across devices. To this end, we first devise a novel FL framework with partial model aggregation (PMA), which only aggregates the lower layers of neural networks responsible for feature extraction while the upper layers corresponding to complex pattern recognition remain at devices for personalization. This design is able to address the data heterogeneity and reduce the transmitted information in wireless channels. Then, we maximize the scheduled data sample volume by joint optimizing the device scheduling, bandwidth allocation, computation and communication time division. Specifically, our analysis reveals that the optimal time division is achieved when the communication and computation parts of PMA-FL have the same power. We also develop a bisection method to solve the optimal bandwidth allocation policy and use the set expansion algorithm to address the optimal device scheduling. Experimental results on the CIFAR-10 dataset show that the proposed PMA-FL improves 11.6% accuracy compared with the state-of-art benchmarks, and the proposed joint dynamic device scheduling and resource optimization approach achieves slightly higher accuracy than the considered benchmarks but reduced 25% energy or 12.5% time budgets.