Joint Online Optimization of Model Training and Analog Aggregation for Wireless Edge Learning

We consider federated learning in a wireless edge network, where multiple power-limited mobile devices collaboratively train a global model, using their local data with the assistance of an edge server. Exploiting over-the-air computation, the edge server updates the global model via analog aggregation of the local models over noisy wireless fading channels. Unlike existing works that separately optimize computation and communication at each step of the learning algorithm, in this work, we jointly optimize the training of the global model and the analog aggregation of the local models over time. Our objective is to minimize the accumulated training loss at the edge server, subject to individual long-term transmit power constraints at the mobile devices. We propose an efficient algorithm, termed Online Model Updating with Analog Aggregation (OMUAA), to adaptively update the local and global models based on the time-varying communication environment. The trained model of OMUAA is channel-and power-aware, and it is in closed form incurring low computational complexity. We study the mutual impact between model training and analog aggregation over time, to derive performance bounds on the computation and communication performance metrics. Furthermore, we consider a variant of OMUAA with double regularization on both the local and global models, termed OMUAA-DR, and show that it can significantly reduce the convergence time to reach long-term transmit power constraints. In addition, we extend both OMUAA and OMUAA-DR to enable analog gradient aggregation, while preserving their performance bounds. Simulation results based on real-world image classification datasets and typical wireless network settings demonstrate substantial performance gain of OMUAA and OMUAA-DR over the known best alternatives.