Joint Power Control and Data Size Selection for Over-the-Air Computation-Aided Federated Learning

Federated learning (FL) has emerged as an appealing machine learning approach to deal with massive raw data generated at multiple mobile devices, which needs to aggregate the training parameter of every mobile device at one base station (BS) iteratively. For parameter aggregating in FL, over-the-air computation is a spectrum-efficient solution, which allows all mobile devices to transmit their parameter-mapped signals concurrently to a BS. Due to heterogeneous channel fading and noise, there exists difference between the BS’s received signal and its desired signal, measured as the mean-squared error (MSE). To minimize the MSE, we propose to jointly optimize the signal amplification factors at the BS and the mobile devices as well as the data size (the number of data samples involved in local training) at every mobile device. The formulated problem is difficult to address due to its nonconvexity. To find the optimal solution, we perform cost function simplification and variable transformation, and solve the transformed problem in a two-level structure. Optimal solution of the lower level problem is found by analyzing every candidate solution from the Karush–Kuhn–Tucker (KKT) condition. Optimal solution of the upper level problem is found by exploring its piecewise convexity. Numerical results show that our proposed method can greatly reduce the MSE and can help to enhance the training performance of FL compared with benchmark methods.