LoLaFL: Low-Latency Federated Learning via Forward-only Propagation

LoLaFL: Low-Latency Federated Learning via Forward-only Propagation Jierui Zhang, Graduate Student Member, IEEE, Jianhao Huang, Member, IEEE, and Kaibin Huang, Fellow, IEEE Abstract --Federated learning (FL) has emerged as a widely adopted paradigm for enabling edge learning with distributed data while ensuring data privacy. However, the traditional FL with deep neural networks trained via backpropagation can hardly meet the low-latency learning requirements in the sixth generation (6G) mobile networks. This challenge mainly arises from the high-dimensional model parameters to be transmitted and the numerous rounds of communication required for convergence due to the inherent randomness of the training process. T o address this issue, we adopt the state-of-the-art principle of maximal coding rate reduction to learn linear discriminative features and extend the resultant white-box neural network into FL, yielding the novel framework of Low-Latency Federated Learning (LoLaFL) via forward-only propagation. LoLaFL enables layer-wise transmissions and aggregation with significantly fewer communication rounds, thereby considerably reducing latency. Additionally, we propose two nonlinear aggregation schemes for LoLaFL. The first scheme is based on the proof that the optimal NN parameter aggregation in LoLaFL should be harmonic-mean-like. The second scheme further exploits the low-rank structures of the features and transmits the low-rank-approximated covariance matrices of features to achieve additional latency reduction. Theoretic analysis and experiments are conducted to evaluate the performance of LoLaFL. In comparison with traditional FL, the two nonlinear aggregation schemes for LoLaFL can achieve reductions in latency of over 91% and 98%, respectively, while maintaining comparable accuracies. I NTRODUCTION With the growing volume of data and the increasing number of edge devices, the sixth generation (6G) mobile networks are envisioned to support a wide range of AI-based applications at the network edge, including augmented/mixed/virtual reality, connected robotics and autonomous systems, and smart cities and homes, among others [1], [2]. To realize this vision, researchers have been motivated to develop technologies to deploy AI models at the network edge. These technologies, collectively called edge learning, leverage the mobile-edge-computing platform to train edge-AI models among edge servers and devices [3], [4]. For its preservation of data privacy, federated learning (FL) emerges as a widely adopted solution for distributed edge learning, where local models are trained using local devices' data and sent to the server for updating the global model [5]-[8]. This collaborative training approach enables multiple devices and a server to train a global J. Zhang, J. Huang, and K. Huang are with the Department of Electrical and Electronic Engineering, The University of Hong Kong, Hong Kong. However, FL faces its own challenges.