Riemannian Federated Learning via Averaging Gradient Streams
Zhenwei Huang, Wen Huang, Pratik Jawanpuria, Bamdev Mishra
TL;DR
RFedAGS addresses federated learning on Riemannian manifolds under partial participation and non-IID data by introducing averaging gradient streams as the server aggregation, enabling unbiased global updates via vector transport on curved spaces. The authors prove global convergence with decaying step sizes and convergence to a neighborhood with fixed step sizes, under standard FL and Riemannian optimization assumptions plus a nontrivial probability-assignment condition whose validity is supported by frequency-based probability estimates. Empirical results on PCA (Stiefel), HSP (hyperbolic), and FMC (SPD) demonstrate that RFedAGS consistently outperforms existing Riemannian FL methods under arbitrary participation and data heterogeneity, with robust performance even when true participation probabilities are unknown. The approach broadens the applicability of federated learning to manifold-valued problems, offering a practical and theoretically sound framework for distributed optimization on curved spaces.
Abstract
Federated learning (FL) as a distributed learning paradigm has a significant advantage in addressing large-scale machine learning tasks. In the Euclidean setting, FL algorithms have been extensively studied with both theoretical and empirical success. However, there exist few works that investigate federated learning algorithms in the Riemannian setting. In particular, critical challenges such as partial participation and data heterogeneity among agents are not explored in the Riemannian federated setting. This paper presents and analyzes a Riemannian FL algorithm, called RFedAGS, based on a new efficient server aggregation -- averaging gradient streams, which can simultaneously handle partial participation and data heterogeneity. We theoretically show that the proposed RFedAGS has global convergence and sublinear convergence rate under decaying step sizes cases; and converges sublinearly/linearly to a neighborhood of a stationary point/solution under fixed step sizes cases. These analyses are based on a vital and non-trivial assumption induced by partial participation, which is shown to hold with high probability. Extensive experiments conducted on synthetic and real-world data demonstrate the good performance of RFedAGS.