Deep Compression for Over-the-Air Federated Learning in Wireless Systems

Abstract

Federated learning (FL) is a powerful, privacy-preserving machine learning (ML) technique that trains models across edge devices without centralizing data, ensuring enhanced privacy and optimizing bandwidth. Over-the-air federated learning (OTA-FL) builds on this by utilizing over-the-air aggregation (OTA) of wireless channels to transmit deep neural network (DNN) parameters to a central parameter server (PS). The PS receives an aggregated version of these parameters, simplifying computations and reducing overhead. However, OTA transmissions face significant challenges due to wireless perturbations, which can degrade model performance, creating a need for more resilient techniques in such environments. This thesis addresses these challenges by first exploring the trade-off between the num ber of channels in a convolutional neural network (CNN) layer and the number of partici pating edge devices, both of which directly impact model accuracy in noisy channels. In creasing the number of CNN channels improves FL performance with fewer participating devices but comes at the cost of higher computational load. On the other hand, increas ing the number of edge devices proves to be even more crucial in boosting overall model performance. With edge devices constrained by limited memory, energy, and bandwidth, achieving an optimal balance between model size, accuracy, and system efficiency becomes a key focus of this work. A major contribution of this thesis is the in-depth analysis and implementation of various pruning techniques alongside quantization-aware training (QAT) to ensure that model performance and size are balanced in OTA-FL networks. While traditional one-shot pruning (OSP) reduces model size significantly, it often sacrifices performance. To mitigate these losses, the thesis introduces advanced pruning methods like iterative magnitude pruning (IMP), which, through a cyclic process of pruning and retraining, allows the model to re tain its essential connections while progressively reducing complexity. This method proves particularly effective in OTA-FL networks, where channel perturbations and signal-to-noise ratios (SNRs) can vary dramatically. Furthermore, the thesis delves into the impact of non-independent and identically distributed (non-IID) data among edge devices, which further exacerbates the accuracy-size trade-off, especially with OSP. To address this, the research proposes PS-side retraining using representative data from the network, combined with the federated proximal (Fed Prox) algorithm, which has been shown to significantly improve model performance and robustness in non-IID environments. By integrating these techniques, this thesis paves the way for more efficient and reliable OTA-FL systems, offering a refined approach to model compression, communication, and computation in FL, and ultimately transforming how wireless networks handle distributed ML.