Riemannian Low-Rank Model Compression for Federated Learning with
  Over-the-Air Aggregation
Ye Xue Vincent Lau
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.