[Paper Review] Learning Rate Optimization for Federated Learning Exploiting Over-the-air Computation

Chunmei Xu, Shengheng Liu+3 more|arXiv (Cornell University)|Feb 5, 2021

Indoor and Outdoor Localization Technologies46 references151 citations

TL;DR

This paper proposes dynamic learning rate (DLR) optimization in over-the-air computation (AirComp)-assisted federated learning to combat channel-induced aggregation distortion. By adapting the learning rate based on wireless channel fading, the method reduces mean squared error (MSE) and improves test accuracy on MNIST and CIFAR10 datasets, with closed-form beamforming designs verified asymptotically and via simulations.

ABSTRACT

Federated learning (FL) as a promising edge-learning framework can effectively address the latency and privacy issues by featuring distributed learning at the devices and model aggregation in the central server. In order to enable efficient wireless data aggregation, over-the-air computation (AirComp) has recently been proposed and attracted immediate attention. However, fading of wireless channels can produce aggregate distortions in an AirComp-based FL scheme. To combat this effect, the concept of dynamic learning rate (DLR) is proposed in this work. We begin our discussion by considering multiple-input-single-output (MISO) scenario, since the underlying optimization problem is convex and has closed-form solution. We then extend our studies to more general multiple-input-multiple-output (MIMO) case and an iterative method is derived. Extensive simulation results demonstrate the effectiveness of the proposed scheme in reducing the aggregate distortion and guaranteeing the testing accuracy using the MNIST and CIFAR10 datasets. In addition, we present the asymptotic analysis and give a near-optimal receive beamforming design solution in closed form, which is verified by numerical simulations.

Motivation & Objective

• Address the challenge of wireless channel fading in over-the-air computation (AirComp)-based federated learning, which distorts model aggregation and degrades performance.
• Overcome the limitations of existing approaches that focus only on physical-layer resource optimization (e.g., power control, beamforming) without leveraging machine learning hyperparameters.
• Introduce dynamic learning rate (DLR) as a novel adaptive mechanism to compensate for channel-induced distortions in real time.
• Develop closed-form and iterative solutions for DLR optimization in MISO and MIMO scenarios, respectively.
• Verify the effectiveness of DLR through extensive simulations on MNIST and CIFAR10 datasets, with asymptotic analysis for massive antenna deployment.

Proposed method

• Formulate the DLR optimization problem in a multiple-input-single-output (MISO) scenario, where the problem is convex and admits a closed-form solution.
• Extend the framework to multiple-input-multiple-output (MIMO) systems using an iterative algorithm to solve the non-convex optimization problem.
• Introduce a dynamic learning rate that varies between predefined minimum and maximum boundaries (rmin and rmax) to adapt to time-varying channel conditions.
• Propose a near-optimal receive beamforming design that sums channel vectors in closed form, minimizing aggregate distortion.
• Apply asymptotic analysis to derive theoretical performance bounds under massive antenna deployment in MISO, SIMO, and MIMO configurations.
• Integrate the DLR mechanism with AirComp to enable simultaneous, analog aggregation of local model updates while minimizing distortion.

Experimental results

Research questions

• RQ1 Can dynamic learning rate adaptation effectively reduce model aggregation error caused by wireless fading channels in AirComp-based federated learning?
• RQ2 How does the performance of DLR compare to fixed learning rate baselines in terms of test accuracy and training loss on real-world datasets?
• RQ3 What is the impact of the number of transmit and receive antennas on the effectiveness of DLR and beamforming design?
• RQ4 Can a closed-form receive beamforming solution achieve near-optimal performance in MISO and MIMO systems under asymptotic conditions?
• RQ5 How does the DLR mechanism interact with physical-layer beamforming and power control to improve overall FL convergence and accuracy?

Key findings

• The proposed DLR scheme reduces aggregate distortion and improves test accuracy on MNIST and CIFAR10 datasets compared to fixed learning rate baselines.
• On MNIST with 20 devices, DLR achieved 97.17% test accuracy versus 97.01% with fixed learning rate, demonstrating consistent improvement.
• In MIMO scenarios with Nd=2, Nt=4, DLR achieved 73.86% accuracy on CIFAR10, outperforming the fixed learning rate baseline of 71.29%.
• Asymptotic analysis confirms that with infinite antennas, the MSE of DLR converges to the theoretical bound, validating the closed-form beamforming design.
• The proposed closed-form receive beamforming design closely matches the theoretical minimum MSE, especially in massive MIMO regimes, with performance approaching the optimal bound.
• Simulation results show that the performance gap between DLR and fixed learning rate diminishes as the number of antennas increases, confirming the robustness of the approach under high-diversity conditions.