Riemannian Low-Rank Model Compression for Federated Learning with Over-the-Air Aggregation

Abstract

Low-rank model compression is a widely used technique for reducing the computational load when training machine learning models. However, existing methods often rely on relaxing the low-rank constraint of the model weights using a regularized nuclear norm penalty, which requires an appropriate hyperparameter that can be difficult to determine in practice. Furthermore, existing compression techniques are not directly applicable to efficient over-the-air (OTA) aggregation in federated learning (FL) systems for distributed Internet-of-Things (IoT) scenarios. In this paper, we propose a novel manifold optimization formulation for low-rank model compression in FL that does not relax the low-rank constraint. Our optimization is conducted directly over the low-rank manifold, guaranteeing that the model is exactly low-rank. We also introduce a consensus penalty in the optimization formulation to support OTA aggregation. Based on our optimization formulation, we propose an alternating Riemannian optimization algorithm with a precoder that enables efficient OTA aggregation of low-rank local models without sacrificing training performance. Additionally, we provide convergence analysis in terms of key system parameters and conduct extensive experiments with real-world datasets to demonstrate the effectiveness of our proposed Riemannian low-rank model compression scheme compared to various state-of-the-art baselines.

arXiv commentThis paper has been accepted for publication as a REGULAR paper in the IEEE Transactions on Signal Processing

References in corpus (13)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Message Passing Algorithms for Compressed Sensing
• Generalized Cross Entropy Loss for Training Deep Neural Networks with Noisy Labels
• Compressing Deep Convolutional Networks using Vector Quantization
• Machine Learning at the Wireless Edge: Distributed Stochastic Gradient Descent Over-the-Air
• An overview of low-rank matrix recovery from incomplete observations
• On Analog Gradient Descent Learning over Multiple Access Fading Channels
• PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization
• Compression-aware Training of Deep Networks
• High-Dimensional Stochastic Gradient Quantization for Communication-Efficient Edge Learning
• Linear Mode Connectivity and the Lottery Ticket Hypothesis
• Intrinsic formulation of KKT conditions and constraint qualifications on smooth manifolds
• Federated Composite Optimization

Cited by in corpus (1)

• Sparse Training for Federated Learning with Regularized Error Correction