Many machine learning tasks, such as principal component analysis and low-rank matrix completion, give rise to manifold optimization problems. Although there is a large body of work studying the design and analysis of algorithms for manifold optimization in the centralized setting, there are currently very few works addressing the federated setting. In this paper, we consider nonconvex federated learning over a compact smooth submanifold in the setting of heterogeneous client data. We propose an algorithm that leverages stochastic Riemannian gradients and a manifold projection operator to improve computational efficiency, uses local updates to improve communication efficiency, and avoids client drift. Theoretically, we show that our proposed algorithm converges sub-linearly to a neighborhood of a first-order optimal solution by using a novel analysis that jointly exploits the manifold structure and properties of the loss functions. Numerical experiments demonstrate that our algorithm has significantly smaller computational and communication overhead than existing methods.

翻译：许多机器学习任务，如主成分分析和低秩矩阵补全，均涉及流形优化问题。尽管已有大量工作研究集中式环境下流形优化算法的设计与分析，但针对联邦环境的研究目前仍十分有限。本文考虑异构客户端数据场景下紧致光滑子流形上的非凸联邦学习问题。我们提出一种算法，通过利用随机黎曼梯度与流形投影算子提升计算效率，采用本地更新机制改善通信效率，并避免客户端漂移现象。理论上，我们通过联合利用流形结构与损失函数性质的新颖分析，证明该算法能以次线性速率收敛至一阶最优解的邻域。数值实验表明，相较于现有方法，本算法显著降低了计算与通信开销。