Riemannian Low-Rank Model Compression for Federated Learning with Over-the-Air Aggregation
This work addresses computational and communication bottlenecks in distributed IoT systems, offering a novel solution for efficient federated learning, though it is incremental in improving existing compression methods.
The paper tackles the problem of low-rank model compression in federated learning for IoT scenarios by proposing a manifold optimization formulation that enforces exact low-rank constraints and supports over-the-air aggregation, resulting in efficient training without performance loss as demonstrated in experiments with real-world datasets.
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.