Nonconvex Federated Learning on Compact Smooth Submanifolds With Heterogeneous Data

Authors: jiaojiao  zhang
jiang  hu
man-cho  anthony
  so
mikael  johansson
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 20:43:28

Abstract: 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.

Main paper: 
 introduction: 
 Federated learning (FL), which enables clients to collaboratively train models without exchanging their raw data, has gained significant traction in machine learning [1, 2] . The framework is appreciated for its capacity to leverage distributed data, accelerate the training process via parallel computation, and bolster privacy protection. The majority of existing FL algorithms address problems that are either unconstrained or have convex constraints. However, for applications such as principal component analysis (PCA) and matrix completion, where model parameters are subject to nonconvex manifold constraints, there are very few options in the federated setting.In this paper, we study FL problems over manifolds in the form ofEQUATIONHere, n is the number of clients, x is the matrix of model parameters, and M is a compact smooth submanifold embedded in R d×k . Examples of such manifolds include the Stiefel, oblique, and symplectic manifolds [3, 4, 5] . For instance, PCA-related optimization problems use the Stiefel manifold M = St(d, k) = {x ∈ R d×k : x T x = I k } to maintain orthogonality [6, 7] . In (1) , the global loss f : R d×k → R is smooth but nonconvex 1 , and the local loss function f i of each client i is the average of the losses f il on the m i data points in its local dataset D i = {D i1 , . . . , D imi }. We consider a heterogeneous data scenario where the statistical properties of D i differ across clients.Manifold optimization problems of the form (1) appear in many important machine learning tasks, such as PCA [8, 9] , low-rank matrix completion [10, 11] , multitask learning [12, 13] , and deep neural network training [14, 15] . Still, there are very few federated algorithms for machine learning on manifolds. In fact, the work [16] appears to be the only FL algorithm that can deal with manifold optimization problems of a similar generality as ours. Handling manifold constraints in an FL setting poses significant challenges: (i) Existing single-machine methods for manifold optimization [6, 5, 4] cannot be directly adapted to the federated setting. Due to the distributed framework, the server has to average the clients local models. Even if each of these models lies on the manifold, their average typically does not due to the nonconvexity of M. The current literature relies on complicated geometric operators like the exponential map, inverse exponential map, and parallel transport, to design an averaging operator for the manifold [16] . However, these mappings may lack closed-form expressions and can be computationally expensive to evaluate. For example, computing the inverse exponential map on the Stiefel manifold requires solving a nonlinear matrix equation [17] . (ii) Extending typical FL algorithms to scenarios with manifold constraints is not straightforward, either. Most existing FL algorithms either are unconstrained [18, 19] or only allow for convex constraints [20, 21, 22, 23, 24] , but manifold constraints are typically nonconvex. Moreover, compared to nonconvex optimization in Euclidean space, manifold optimization necessitates the consideration of the geometric structure of the manifold and properties of the loss functions, which poses challenges for algorithm design and analysis. (iii) Traditional methods for enhancing communication efficiency in FL, like local updates [25] , need substantial modifications to accommodate manifold constraints. The so-called client drift issue due to local updates and heterogeneous data [19] persists in the realm of manifold optimization. Directly using client-drift correcting techniques originally developed for Euclidean spaces [19, 26, 27] could lead to additional communication or computational costs due to the manifold constraints. For instance, in [16] , the correction term requires additional communication of local Riemannian gradients and involves using parallel transport to move the correction term onto some tangent space in preparation for the exponential mapping. Although some existing decentralized manifold optimization algorithms [9, 28, 29] can be simplified to an FL scenario with only one local update under the assumption of a fully connected network, these algorithms cannot be directly applied to FL scenarios with more than one local update, especially in cases of data heterogeneity. Extending the analysis of these algorithms to FL scenarios with multiple local updates is not straightforward. On the other hand, the use of local updates in FL, compared to these decentralized distributed algorithms, can more effectively reduce the number of communication rounds.We consider the nonconvex FL problem (1) with M being a compact smooth submanifold and allow for heterogeneous data distribution among clients. Our contributions are summarized as follows.We propose an FL algorithm for solving (1) that is efficient in terms of both computation and communication. We employ stochastic Riemannian gradients and a projection operator to address manifold constraints, use local updates to reduce the communication frequency between clients and the server, and design correction terms to overcome client drift. In terms of server updates, our algorithm ensures feasibility of all global model iterates and is computationally efficient since it avoids the techniques used in [16] based on the exponential mapping and inverse exponential mapping for averaging local models on manifolds. For local updates, our algorithm constructs the correction terms locally without increasing communication costs. In comparison, the approach presented in [16] requires each client to transmit an additional local stochastic Riemannian gradient for constructing correction terms. Moreover, [16] necessitates parallel transport to position the correction terms on tangent spaces so that the exponential map can be applied to ensure the feasibility of local models, thereby increasing computational costs. In contrast, our algorithm utilizes a simple projection operator, effectively eliminating the need for parallel transport of correction terms.2) Theoretically, we establish sub-linear convergence to a neighborhood of a first-order optimal solution and demonstrate how this neighborhood depends on the stochastic sampling variance and algorithm parameters. Our analysis introduces novel proof techniques that utilize the curvature of the manifolds and the properties of the loss functions to overcome the challenges posed by the nonconvexity of manifold constraints in the nonconvex FL scenario. Compared to the existing work [16] where analytical results are limited to cases where either the number of local updates is one or the number of participating clients per communication round is one, our theoretical results allow for an arbitrary number of local updates and support full client participation. The key components of our analysis are the manifold geometry and the Lipschitz continuity of the projection operator, both of which are inherent to the submanifold constraint.3) Our algorithm demonstrates superior performance over alternative methods in the numerical experiments. In particular, it produces high-accuracy results for kPCA and low-rank matrix completion at a significantly lower communication and computation cost than alternative algorithms.In this section, we first review federated learning algorithms for composite optimization with and without constraints. Then, we discuss FL algorithms with manifold constraints.Composite FL in Euclidean space. Problem (1) can be viewed as a special case of composite FL where the loss function is a composition of f and the indicator function of M. It is important to note that since the manifold is nonconvex, its indicator function is also nonconvex. Most existing composite FL methods can only handle convex constraints. The work [20] proposed a federated dual averaging method and established its convergence for a general loss function under bounded gradient assumptions, but only for quadratic losses under the bounded heterogeneity assumption that the degree of data heterogeneity among clients is bounded. In contrast, we make no assumptions about the similarity of data across clients. The fast federated dual averaging algorithm [21] extends the work in [20] by using both past gradient information and past model information in the local updates. However, the work [21] requires each client to transmit the local gradient as well as the local model, and it assumes bounded data heterogeneity. The work [22] introduces the federated Douglas-Rachford method, and the work [23] applies this algorithm to solve dual problems. Although these two methods avoid bounded data heterogeneity, they require an increasing number of local updates to ensure convergence, which reduces their practicality in FL. The recent work [24] proposes a communicationefficient FL algorithm that overcomes client drift by decoupling the proximal operator evaluation and the communication and shows that the method converges without any assumptions on data similarity.Federated learning on manifolds. The existing composite FL in Euclidean space [20] - [24] only considers scenarios where the nonsmooth term in the loss functions is convex. However, incorporating a nonconvex manifold constraint as an indicator function introduces a nonconvex nonsmooth term. Consequently, the methods in [20] - [24] are not directly applicable. A typical challenge caused by the nonconvex manifold constraint is that the average of local models, each of which lies on the manifold, may not belong to the manifold. To address this issue, the work [16] introduced Riemannian federated SVRG (RFedSVRG), where the server maps the local models onto a tangent space, calculates an average, and then retracts the average back to the manifold. This process sequentially employs inverse exponential and exponential mappings. Moreover, RFedSVRG employs a correction term to overcome client drift but requires additional communication of local Riemannian gradients to construct the correction term. In addition, the method uses parallel transport to position the correction term, which increases the computation cost even further. Note that the manifold we consider is a compact smooth submanifold embedded in R d×k , which is more restrictive than the manifolds discussed in [16] . However, this approach still encompasses many common manifolds, including the Stiefel, oblique, and symplectic manifolds [3, 4, 5] . The work [30] explores the differential privacy of RFedSVRG. The work [31] considers the specific manifold optimization problem that appears in PCA and investigates an ADMM-type method that penalizes the orthogonality constraint. However, this algorithm requires solving a subproblem to desired accuracy, which increases computational cost. The work [32] introduces a differentially private FL algorithm for solving PCA. Finally, the works [33] and [34] consider the leading eigenvector problem with homogeneous data across clients, where the loss functions are quadratic and the manifold is a sphere. In contrast, we consider a more general setting with heterogeneous data, where x lies on M and the loss functions are smooth and nonconvex.Notations. We use I k to denote a k × k identity matrix. We use (col{z i } n i=1 ) = col{P M (z i )} n i=1. Given a positive definite matrix x, we use x -1/2 to denote the inverse of the square root of x, i.e., x -1/2 x -1/2 = x -1 . We define D 2 to be the second-order differential operator. 
preliminaries: 
 Below, we introduce fundamental definitions and inequalities for optimization on manifolds.Manifold optimization aims to minimize a real-valued function over a manifold, i.e., min x∈M f (x). Throughout the paper, we restrict our discussion to embedded submanifolds of the Euclidean space, where the associated topology coincides with the subspace topology of the Euclidean space. We refer to these as embedded submanifolds. Some examples of such manifolds include the Stiefel manifold, oblique manifold, and symplectic manifold [3, 4, 5] . We define the tangent space of M at point x as T x M, which contains all tangent vectors to M at x, and the normal space as N x M which is orthogonal to the tangent space. With the definition of tangent space, we can define the Riemannian gradient that plays a central role in the characterization of optimality conditions and algorithm design for manifold optimization. Definition 2.1 (Riemannian gradient gradf (x)). The Riemannian gradient gradf (x) of a function f at the point x ∈ M is the unique tangent vector that satisfies⟨gradf (x), ξ⟩ x = df (x)[ξ], ∀ξ ∈ T x M,where ⟨•, •⟩ x is the Riemannian metric and df denotes the differential of function f .For a submanifold M, the Riemannian gradient gradf (x) (under the Euclidean inner product) can be computed as [5, Proposition 3.61] gradf (x) = P TxM (∇f (x)), where P TxM (∇f (x)) represents the orthogonal projection of ∇f (x) onto T x M. The Riemannian gradient gradf (x) reduces to the Euclidean gradient ∇f (x) when M is the Euclidean space R d×k .In our federated manifold learning algorithm, the server needs to fuse models that have undergone multiple rounds of local updates by the clients. Due to the nonconvexity of the manifold, the average of points on the manifold is not guaranteed to belong to the manifold. The tangent space-based exponential mapping or other retraction operations commonly used in manifold optimization are expensive in FL [16] . Specifically, the server needs to map the local models onto a tangent space using inverse exponential mapping, calculate an average on the tangent space, and then perform an exponential mapping to retract this average back onto the manifold. This exponential mapping, due to its dependency on the tangent space, also calls for parallel transport during the local updates when there are correction terms. To overcome this difficulty, we use a projection operator P M defined byP M (x) ∈ argmin u∈M ∥x -u∥ 2 (2)to ensure the feasibility of manifold constraints. The projection operator P M can be explicitly calculated for many common submanifolds, as discussed in [35] . For the Stiefel manifold, the closedform expression for P M (x) of a given matrix x with full column rank is P M (x) = x(x T x) -1/2 ; see [35, Proposition 7] . It is worth noting that P M can be regarded as a special retraction operator when restricted to the tangent space [35] . However, unlike a typical retraction operator, its domain is R d×k , not just the tangent space, which enables a more practical averaging operation across clients in FL. Despite these advantageous properties, the nonconvex nature of M means that P M (x) may be set-valued and non-Lipschitz, making the use and analysis of P M in the FL setting highly nontrivial.To tackle this, we introduce the concept of proximal smoothness that refers to a property of a closed set, including M, where the projection becomes a singleton when the point is sufficiently close to the set. Definition 2.2 (γ-proximal smoothness of M). For any γ > 0, we define the γ-tube around M asU M (γ) := {x : dist(x, M) < γ},where dist(x, M) := min u∈M ∥u -x∥ is the Eulidean distance between x and M. We say that M is γ-proximally smooth if the projection operator P M (x) is a singleton whenever x ∈ U M (γ).It is worth noting that any compact smooth submanifold M embedded in R d×k is a proximally smooth set [36, 37] . The constant γ can be calculated with the method of supporting principle for proximally smooth sets [38, 39] . For instance, the Stiefel manifold is 1-proximally smooth. Assumption 2.3. The manifold M is assumed to be a compact smooth submanifold embedded in R d×k , with the Euclidean inner product serving as its Riemannian metric. Moreover, we assume that the proximal smoothness constant of M is 2γ.With Assumption 2.3, we can ensure not only the uniqueness of the projection but also the Lipschitz continuity of the projection operator P M around M, analogous to the non-expansiveness of projections under Euclidean convex constraints.Lipschitz continuity of P M . Define U M (γ) := {x : dist(x, M) ≤ γ} as the closure of U M (γ).Following the proof in [36, Theorem 4.8] , for a 2γ-proximally smooth M, the projection operatorEQUATIONNormal inequality. In the normal space N x M, we exploit the so-called normal inequality [36, 37] . Following [36] , given a 2γ-proximally smooth M, for any x ∈ M and v ∈ N x M, it holds thatEQUATIONIntuitively, when x and y are close enough, the matrix y -x is approximately in the tangent space, thus being nearly orthogonal to the normal space.Question: Does the paper discuss the limitations of the work performed by the authors? Answer:[Yes] Justification: See Conclusions and limitations. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
proposed algorithm: 
 In this section, we develop a novel algorithm for nonconvex federated learning on manifolds. The algorithm is inspired by the proximal FL algorithm for strongly convex problems in Euclidean space recently proposed in [24] but includes several non-trivial extensions. These include the use of Riemannian gradients and manifold projection operators and the ability to handle nonconvex loss functions, which call for a different convergence analysis.The per-client implementation of our algorithm is detailed in Algorithm 1. Similarly to the wellknown FedAvg, it operates in a federated learning setting with one server and n clients. Each client i engages in τ steps of local updates before updating the server. We use r as the index of communication rounds and t as the index of local updates.At any communication round r, client i downloads the global model x r from the server and computes P M (x r ). Each client i updates two local variables, ẑr i,t and z r i,t , where ẑr i,t aggregates the Riemannian gradients from local updates, and z r i,t = P M (ẑ r i,t ) ensures that Riemannian gradients can be computed at points on M. The update of ẑr i,t is given in Line 8, where B r i,t is a mini-batch dataset and c r i is a correction term to eliminate client drift. After τ local updates, client i sends ẑr i,τ to the server. The server receives all ẑr i,τ , computes their average to form the global model x r+1 following Line 13, and broadcasts x r+1 to each client i that uses x r+1 to locally construct the correction term c r+1 i . Algorithm 1 Proposed algorithm 1: Input: R, τ , η, η g , η = ηη g τ , x 1 , andc 1 i = 0 for all i ∈ [n] 2: for r = 1, 2, . . . , R do 3: Client i 4:Set z r i,0 = P M (x r ) and z r i,0 = P M (x r ) 5:for t = 0, 1, . . . , τ -1 do 6: Sample a mini-batch dataset B r i,t ⊆ D i with |B r i,t | = b 7: Update gradf i (z r i,t ; B r i,t ) = 1 b D il ∈B r i,t gradf il (z r i,t ; D il ) 8: Update z r i,t+1 = z r i,t -η gradf i (z r i,t ; B r i,t ) + c r i 9:Update z r i,t+1 = P M z r i,t+110:end for 11:Send z r i,τ to the server 12:Server 13:Updatex r+1 = P M (x r ) + η g 1 n n i=1 z r i,τ -P M (x r ) 14:Broadcast x r+1 to all the clients 15:Client i 16:Receive x r+1 17:Updatec r+1 i = 1 ηgητ (P M (x r ) -x r+1 ) -1 τ τ -1 t=0 gradf i (z r i,t ; B r i,t ) 18: end for 19: Output: P M (x R+1 )In the proposed algorithm, each client i downloads x r at the start of local updates and uploads ẑr i,τ at the end of the local updates. Therefore, each communication round involves each client and the server exchanging only a single d × k matrix.To better understand the proposed algorithm, we present its equivalent and more compact form:EQUATIONFor the initialization of correction term, we set gradf i z 0 i,t ; B 0 i,t= 0 for all t and i so that1 τ τ -1 t=0 gradf z 0 t ; B 0 t -1 τ τ -1 t=0 gradf z 0 t ; B 0 t = 0, which coincides with the initialization c 1 i = 0 in Algorithm 1.The equivalence between Algorithm 1 and ( 5) can be proved following the similar derivations in [24] and is therefore omitted. With (5), we highlight the key properties and innovations of the proposed algorithm.1) Recovery of the centralized algorithm in special cases. Substituting the definitions of x, gradf (z r t ; B r t ), and η into the last step in (5), we haveEQUATIONThanks to the introduction of the variable ẑr i,t during the local updates for each client i in Algorithm 1, the server after averaging ẑr i,τ obtains an accumulation of τ local Riemannian gradients across local updates and an average of the local Riemannian gradients across all clients. In the special case where τ = 1 and b = m i , i.e., with the local full Riemannian gradient for each client i, the update of (6) recovers the centralized projected Riemannian gradient descent (C-PRGD)EQUATIONIn our analysis, we will compare the sequence P M (x r+1 ) generated by our algorithm with the virtual iterate xr+1 to establish the convergence of our algorithm.2) Feasibility of all iterates at a low computational cost. Our algorithm uses P M to obtain feasible solutions on the manifold, which is computationally more efficient than the commonly used exponential mapping. In fact, since the exponential mapping relies on a point on the manifold and the tangent space at that point, it cannot be directly used in our algorithm. In the local updates, it is difficult to perform exponential mapping on ẑr t+1 because ẑr t is not on the manifold; see the first step in (5) . As shown in [24] , ẑr t+1 is essential for the server to obtain aggregated Riemannian gradients from n clients after τ local updates. Moreover, at the server, although P M (x r ) is on the manifold, the aggregated direction does not lie in the tangent space at P M (x r ). The algorithm suggested in [16] uses an exponential mapping to fuse local models. It needs to map the local models to a tangent space using the inverse exponential mapping and then retract the result back to the manifold, which is computationally expensive. Our use of P M on a point in the Euclidean space close to the manifold avoids these high computational costs, but creates new challenges for the analysis.3) Overcoming client drift. Inspired by [24] , we use a correction term c r i to address client drift. According to the first step of ( 5), the correction employs the idea of "variance reduction", which involves replacing the old local Riemannian gradient1 τ τ -1 t=0 gradf z r-1 t ; B r-1 twith the new one gradf (z r t ; B r t ) in the average of all client Riemannian gradients1 τ τ -1 t=0 gradf z r-1 t ; B r-1 t, where the "variance" refers to the differences in Riemannian gradients among clients caused by data heterogeneity. Compared to [16] , our correction improves communication and computation. The correction approach in [16] necessitates extra transmissions of local Riemannian gradients, while our correction term can be locally generated, leading to a significantly reduced communication overhead. Furthermore, [16] employs parallel transport to position the correction term with a specific tangent space for the exponential mapping to ensure local model feasibility. Our approach, which utilizes P M , eliminates the need for parallel transport and reduces the computations per iteration even further.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: We did not provide the code during the submission stage.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
analysis: 
 In this section, we analyze the convergence of the proposed Algorithm 1. Throughout the paper, we make the following assumptions, which are common in manifold optimization. EQUATION)With the compactness of M, there exists a constant D f > 0 such that the Euclidean gradient ∇f il (x; D il ) of f il is bounded by D f , i.e., max i,l,x∈M ∥∇f il (x; D il )∥ ≤ D f . It then follows from [28, Lemma 4.2] that there exists a constant L ≤ L < ∞ such that for any x, y ∈ M,f il (y; D il ) ≤ f il (x; D il ) + ⟨gradf il (x; D il ), y -x⟩ + L 2 ∥x -y∥ 2 , ∥gradf il (x; D il ) -gradf il (y; D il )∥ ≤ L∥x -y∥.To address the stochasticity introduced by the random sampling B r i,t , we define F r t as the σ-algebra generated by the set {B r i, t | i ∈ [n], r ∈ [r], t ∈ [t -1]}EQUATIONConsidering the nonconvexity of f and the manifold constraints, we characterize the first-order optimality of (1). A point x ⋆ is defined as a first-order optimal solution of (1) if x ⋆ ∈ M and gradf (x ⋆ ) = 0. We employ the norm of G η (P M (x r )) as a suboptimality metric, defined asEQUATIONand xr+1 defined in ( 7) is used only for analytical purposes. In optimization on Euclidean space such that M = R d×k , the quantity G η (P M (x r )) serves as a widely accepted metric to assess firstorder optimality for nonconvex composite problems [41] . In optimization on manifold, we have G η (P M (x r )) = 0 if and only if gradf (P M (x r )) = 0 for any η > 0. Moreover, for a suitable η, the use of Riemannian gradients in the update of xr+1 helps us to establish a quasi-isometric property, specifically that 1/2∥gradf (P M (x r ))∥ ≤ ∥G η (P M (x r ))∥ ≤ 2∥gradf (P M (x r ))∥; see Lemmas A. EQUATIONwhereEQUATIONwhere Ω 1 > 0 is a constant related to initialization.In Theorem 4.3, the first term on the right hand of ( 12) converges at a sub-linear rate, which is common for constrained nonconvex optimization [41, 42] . The second term is a constant error caused by the variance σ 2 of stochastic Riemannian gradients.Theoretical contributions. The work [16] establishes convergence rates of O(1/R) for τ = 1 and O(1/(τ R)) for τ > 1 but only if a single client participates in the training per communication round. In contrast, our Theorem 4.3 achieves a rate of O(1/( √ nτ R)) for τ > 1 and full client participation. Our theorem indicates that multiple local updates enable faster convergence, which distinguishes our algorithm from decentralized manifold optimization algorithms [9, 28] that limit clients to do a single local update. The works [43] and [1] study smooth nonconvex FL with heterogeneous data in Euclidean space, relying on the assumptions of bounded second moments [43, Assumption 1] and Blocal dissimilarity [1, Definition 3] , respectively, to establish convergence. These assumptions imply some similarity between client data. Our main analytical advantages over [43] and [1] are twofold: We do not assume data similarity, and we completely eliminate client drift. This is reflected in our convergence error, which eliminates the error introduced by data heterogeneity. Our convergence analysis relies on several novel techniques. Specifically, we capitalize on the structure of M and exploit the proximal smoothness of M to guarantee the uniqueness of P M and Lipschitz continuity of P M within a tube around M. Additionally, we carefully select the step sizes to ensure that the iterates remain close to M, thus preserving the established properties throughout the iterations. Last but not least, we select an appropriate first-order optimality metric (see (10) ) and jointly consider the properties of M and the loss functions to establish some new inequalities for the convergence of this metric, given in Appendix. 
numerical experiments: 
 In this section, we conduct numerical experiments on two applications on the Stiefel manifold: kPCA and the low rank matrix completion (LRMC). We compare with existing algorithms, including RFedavg, RFedprox, and RFedSVRG. RFedavg and RFedprox are direct extensions of FedAvg [25] and Fedprox [1] . For RFedSVRG, there are no theoretical guarantees when we set τ > 1 and make all clients participate. In all alternative algorithms, the calculations of the exponential mapping, its inverse, and the parallel transport on the Stiefel manifold are needed. The exponential mapping has a closed-form expression but involves a matrix exponential [6] , the inverse exponential mapping needs to solve a nonlinear matrix equation [17] , and the parallel transport needs to solve a linear differential equation [44] , all of which are computationally challenging. In their implementations, approximate versions of these mappings are used [16, 45, 46] .kPCA. Consider the kPCA problem minimize x∈St(d,k) f (x) = 1 n n i=1 f i (x), f i (x) = - 1 2 tr(x T A T i A i x),where St(d, k) = {x ∈ R d×k | x T x = I k } denotes the Stiefel manifold, and A T i A i ∈ R d×d is the covariance matrix of the local data A i ∈ R p×d of client i. We conduct experiments where the matrix A i is from the Mnist dataset. The specific experiment settings can be found in Appendix A. 4 In the first set of experiments, we compare our algorithm with RFedavg, RFedprox, and RFedSVRG. RFedSVRG requires each client to transmit two d × k matrices per communication round, while our algorithm only transmits a single matrix. We use communication quantity to count the total number of d × k matrices that per client transmits to the server. We use the local full gradient ∇f i to mitigate the effects of stochastic gradient noise. In Figs. 1, we set τ = 10 and η = 1/β for all algorithms, where β is the square of the largest singular value of col{A i } n i=1 . We set η g = 1 to facilitate comparison with other algorithms. As noted below (40) in the Appendix, all analytical results leading up to (40) remain valid for η g = 1. It can be observed that RFedavg and RFedprox face the issue of client drift and have low accuracy. Both RFedSVRG and our algorithm overcome the client drift, but our algorithm, though being similar in terms of communication rounds, is much faster in terms of communication quantity and run time.In the second set of experiments, we test the impact of τ . For all the algorithms, we set the step size η = 1/β and τ ∈ {10, 15, 20}. For our algorithm, we set η g = 1. The experiment results are shown in Figs. 2 . For all values of τ , our algorithm achieves better convergence and requires less communication quantity. LRMC. LRMC aims to recover a low-rank matrix A ∈ R d×T from its partial observations. Let Ω be the set of indices of known entries in A, the rank-k LRMC problem can be written as minimize X∈St(d,k),V ∈R k×T 1 2 ∥P Ω (XV -A)∥ 2 , where the projection operator P Ω is defined in an entry-wise manner with (P Ω (A)) l1l2 = A l1l2 if (l 1 , l 2 ) ∈ Ω and 0 otherwise. In terms of the FL setting, we consider the case where the observed data matrix P Ω (A) is equally divided into n clients 0 2500 5000 7500 10000 Communication round EQUATIONwhere Ω i is the subset corresponding to client i in Ω and V i (X) := argmin V ∥P Ωi (XV -A i )∥. In the experiments, we set T = 1000, d = 100, k = 2, n = 10, and use the local full gradients. The other settings can be found in Appendix A.4.2.The numerical comparisons with RFedavg, RFedprox, and RFedSVRG are presented in Figs. 4 . Our algorithm and RFedSVRG achieve similar convergence for communication rounds, but our algorithm converges faster than RFedSVRG in terms of communication quantity and run time. Additional experimental results can be found in Appendix A.4.2. 
conclusions and limitations: 
 This paper addresses the challenges of FL on compact smooth submanifolds. We introduce a novel algorithm that enables full client participation, local updates, and heterogeneous data distributions. By leveraging stochastic Riemannian gradients and a manifold projection operator, our method enhances computational and communication efficiency while mitigating client drift. By exploiting the manifold structure and properties of the loss function, we prove sub-linear convergence to a neighborhood of a first-order stationary point. Numerical experiments show a superior performance of our algorithm in terms of computational and communication costs compared to the state-of-the-art.Our paper motivates several questions for further investigation. First, the absence of closed-form solutions for the projection operator P M for certain manifolds necessitates exploring methods to calculate projections approximately. Additionally, our step-size selection relies on the proximal smoothness constant γ, underscoring the need for estimating γ either off-line for specific manifolds or adaptively on-line. Furthermore, designing algorithms for partial participation and devising corresponding client-drift correction mechanisms require further investigation.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: See numerical experiments and Appendix A.4.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: Our algorithm uses stochastic gradients. We fixed the random seed to better compare the impact of algorithm parameters on the accuracy of the algorithm.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We use open datasets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: AWe analyze the proposed algorithm using the Lyapunov function Ω r defined byEQUATION)where f ⋆ is the optimal value of problem (1) and we defineΛ r := η(τ gradf (P M (x r )) + τ -1 t=0 gradf (z r-1 t ; B r-1 t ) - τ -1 t=0 gradf (z r-1 t ; B r-1 t ))and Λ r := col1 n n i=1 Λ r i n i=1. The Lyapunov function consists of two parts: to bound the suboptimality of the global model P M (x r ) and the reduction of "variance" among clients, respectively.Let us start with the following lemma on the global-like Lipschitz-continuity property of P M .Lemma A.1. There exists a constant M > 0 such that for any x ∈ M,EQUATIONProof. Let us consider two cases:• ∥u∥ ≥ γ: Since P M (x + u) and x belong to M, we have∥P M (x + u) -x∥ ≤ diam(M) ≤ diam(M) γ ∥u∥,where diam(M) := max x,y∈M ∥x -y∥ is the diameter of M.• ∥u∥ ≤ γ: By the 2-Lipschitz continuity of P M over U M (γ) in (3), we have∥P M (x + u) -x∥ ≤ 2∥u∥.Setting M := max diam(M) γ , 2 , we complete the proof.In the following, we show the reasonableness of the suboptimality metric ∥G η (•)∥.Lemma A.2. Consider G η (•) defined by (10) . Then, for any x ∈ M, it holds that gradf (x) = 0 if and only if G η (x) = 0.In addition, under Assumptions 2.3 and 4.1EQUATIONProof. If gradf (x) = 0, it follows directly from the definition of G η (•) that G η (x) = 0. Conversely, if G η (x) = 0, we havex = P M (x -ηgradf (x)) := argmin y∈M ∥y -x + ηgradf (x)∥ 2 .It follow from the optimality of x that 0 = P TxM (ηgradf (x)), which implies that gradf (x) = 0.With [47, Lemma] , P M (•) is sufficiently smooth over U M (γ). Let us define L P := max x∈U M (γ) ∥D 2 P M (x)∥, where D 2 denotes the second-order differential operator. Then, we have∥G η (x)∥ = 1 η ∥x -P M (x -ηgradf (x))∥ ≥ ∥gradf (x)∥ - 1 2 L P η∥gradf (x)∥ 2 ≥ 1 2 ∥gradf (x)∥,where we use η ≤ γ D f in the first inequality and η ≤ 1 L P D f in the second inequality. This gives (16).To prove Theorem 4.3, we use the following lemma to establish a recursion on the second term on Λ r in the Lyapunov function. EQUATION≤ 1 n 4η 2 τ L 2 3nM 2 τ 3 η 2 ∥gradf (P M (x r ))∥ 2 + 9τ E∥Λ r -Λ r ∥ 2 + 18nτ 2 η 2 σ 2 b + 1 n 4η 2 n 2 τ 2 σ 2 nτ b .Proof. As a first step, we bound the drift error z r i,t+1 -P M (x r ) 2 that is caused by the local updates. If τ = 1, the error is zero since z r i,t = P M (x r ). When τ ≥ 2, repeating the local updates for t steps and substituting z r i,0 = P M (x r ) and z r i,t+1 = P M (ẑ r i,t+1 ), we haveEQUATION)To bound the right-hand side of (18), we compare our algorithm with the exact C-PRGD step given in (7) under the step size (t + 1)η xr+1 C-PRGD := P M (P M (x r ) -(t + 1)ηgradf (P M (x r ))) .It follows from ( 15) that ∥x r+1 C-PRGD -P M (x r )∥ ≤ M τ η∥gradf (P M (x r ))∥. Then from (18) we haveEQUATION= E P M P M (x r ) -η t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) + c r i -xr+1 C-PRGD + xr+1 C-PRGD -P M (x r ) 2 ≤ 2E P M P M (x r ) -η t ℓ=0 (gradf i (z r i,ℓ ; B r i,ℓ ) + c r i ) -xr+1 C-PRGD 2 (I) +2M 2 τ 2 η 2 ∥gradf (P M (x r ))∥ 2 ,where we use ∥a + b∥ 2 ≤ 2∥a∥ 2 + 2∥b∥ 2 in the inequality.To bound the term (I) on the right hand of ( 19), from η ≤ γηg 2D f and max i,l,x∈M ∥∇f il (x; D il )∥ ≤ D f , we haveη t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) + c r i ≤ γ.C-PRGD , we can invoke the 2-Lipschitz continuity of P M over U M (γ) given in (3) and getEQUATIONNext, to bound the right-hand side of ( 20) we rewrite it in terms of ∥Λ r -Λ r ∥ 2 by substituting the definition of the c r i given in ( 5)EQUATION+ gradf i (P M (x r )) + 1 τ τ -1 t=0 1 n n i=1 gradf i (z r-1 i,t ; B r-1 i,t ) - 1 τ τ -1 t=0 gradf i z r-1 i,t ; B r-1 i,t -gradf (P M (x r )) 2 = 4E η t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) -gradf i (P M (x r )) + 1 ητ Λ r i -Λ r 2 ≤ 8E t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) -gradf i (P M (x r )) 2 (II) +8E t + 1 τ Λ r i - t + 1 τ Λ r 2.Next, for the term (II), substituting Assumption 4.2 yieldsEQUATION≤ 16(t + 1) 2 E η t + 1 t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) -gradf i (z r i,ℓ ) 2 + 16E η t ℓ=0 gradf i (z r i,ℓ ) -gradf i (P M (x r )) 2 .For the first term can be handled by the fact [48, Corollary C.1] that16(t + 1) 2 η 2 E 1 t + 1 t ℓ=0 gradf i (z r i,ℓ ; B r i,ℓ ) -gradf i (z r i,ℓ ) 2 (23) = 16(t + 1) 2 η 2 (t + 1) 2 t ℓ=0 E E ∥ gradf i z r i,ℓ ; B r i,ℓ -gradf i z r i,ℓ ∥ 2 |F r t ≤ 1 t + 1 σ 2 b .Combining ( 22), (23) , and ( 21), we haveEQUATIONwhere we use Assumption 4.1. Next we substitute ( 24) into (19) to getEQUATION) ≤ 16(t + 1)η 2 L 2 t ℓ=0 E∥z r i,ℓ -P M (x r )∥ 2 + 2M 2 τ 2 η 2 ∥gradf (P M (x r ))∥ 2 + 8E∥Λ r i -Λ r ∥ 2 + 16τ η 2 σ 2 b :=A r .The following proof is similar to that in [24] . We define A r as the sum of the last three terms on the right hand of ( 25) and S r i,t :=t ℓ=0 E∥z r i,ℓ -P M (x r )∥ 2 . By E[ z r i,t+1 -P M (x r ) 2 ] =S r i,t+1 -S r i,t and ( 25), we haveEQUATION)where inequality is from η ≤ η g /(16L) and thus 16(t + 1)η 2 L 2 ≤ 1/(16τ ). With ( 26), we getEQUATION)where we use 27) over all the clients i, we getτ -2 ℓ=0 (1 + 1/(16τ )) ℓ ≤ τ -2 ℓ=0 exp (ℓ/(16τ )) ≤ τ -2 ℓ=0 exp(1/16) ≤ 1.1τ . Sum- ming (EQUATION)Now we are ready to bound 1 n E∥Λ r+1 -Λ r+1 ∥ 2 . By the definition of Λ r+1 and Λ r+1 we haveE∥Λ r+1 -Λ r+1 ∥ 2 (29) = η 2 E τ gradf P M (x r+1 ) - τ -1 t=0 gradf (z r t ; B r t ) -τ gradf (P M (x r+1 )) + τ -1 t=0 gradf (z r t ; B r t ) 2 ≤ η 2 E τ gradf P M (x r+1 ) - τ -1 t=0 gradf (z r t ; B r t ) 2 = η 2 E τ gradf P M (x r+1 ) -τ gradf (P M (x r )) + τ gradf (P M (x r )) - τ -1 t=0 gradf (z r t ) + τ -1 t=0 gradf (z r t ) - τ -1 t=0 gradf (z r t ; B r t ) 2 ≤ 2η 2 τ 2 L 2 nE P M (x r+1 ) -P M (x r ) 2 + 4η 2 τ L 2 n i=1 τ -1 t=0 E z r i,t -P M (x r ) 2 + 4η 2 τ n σ 2 b .Here, the first inequality is due to ∥Λ r+1 -Λ r+1 ∥ 2 ≤ ∥Λ r+1 ∥ 2 , the last inequality is due to ∥a + b∥ 2 ≤ 2∥a∥ 2 + 2∥b∥ 2 , Assumption 4.1, and following similar derivations as in (23) .By substituting ( 28) into ( 29) and reorganizing the results, we complete the proof of Lemma A.3.To bound the first term in the Lyapunov function, we focus on the server-side update. We begin with the following lemma over the manifolds.EQUATION)Proof. For any µ-strongly convex function h, we have for any y, z ∈ MEQUATIONwhere the second inequality is from the normal inequality (4) and ∥∇h(y) -gradh(y)∥ ≤ ∥∇h(y)∥.Setting h(y) = 1 2η ∥y -(x -ηv)∥ 2 in (31) with µ = 1/η, y = x + , and noting the optimality of x + (i.e., gradh(x + ) = 0), we have1 ∥z -(x -ηv)∥ 2 ≥ 1 2η ∥x + -(x -ηv)∥ 2 + 1 2η - ∥y -(x -ηv)∥ 4ηγ ∥z -x + ∥ 2 ≥ 1 2η ∥x + -(x -ηv)∥ 2 + 1 2η - 3∥v∥ 4γ ∥z -x + ∥ 2 ,where the second inequality is from x -ηv ∈ U M (γ) and∥P M (x -ηv) -(x -ηv)∥ ≤ ∥P M (x - ηv) -x∥ + η∥v∥ ≤ 3η∥v∥.Rearranging the above inequality leads toEQUATION)It follows from the L-smoothness of f thatf (x + ) ≤ f (x) + gradf (x), x + -x + L 2 ∥x + -x∥ 2 ≤ f (z) + gradf (x), x + -z + L 2 ∥z -x∥ 2 + L 2 ∥x + -x∥ 2 ,where we use f (x) + ⟨gradf (x), z -x⟩ -L 2 ∥z -x∥ 2 ≤ f (z) in the last inequality. Combining the above inequality and (32) gives (30) .In the following, we use Lemma A.4 to ( 7) and ( 6), respectively. First, to apply Lemma A.4 to ( 7), we substitute x + = xr+1 , z = P M (x r ), x = P M (x r ), and v = gradf (P M (x r )) and getEQUATIONwhere we use η ≤ γ D f to guarantee xr+1 ∈ U M (γ) and ρ := 3D f 2γ . Next, to use Lemma A.4 to (6), we set x + = P M (x r+1 ), x = P M (x r ), z = xr+1 , and v = v r and getEQUATION≤ E f (x r+1 ) + gradf (P M (x r )) -v r , P M (x r+1 ) -xr+1 - 1 2η ∥P M (x r+1 ) -P M (x r )∥ 2 -∥x r+1 -P M (x r )∥ 2 - 1 -ηρ 2η ∥x r+1 -P M (x r+1 )∥ 2 + L 2 P M (x r+1 ) -P M (x r ) 2 + L 2 ∥x r+1 -P M (x r )∥ 2 = E f xr+1 + gradf (P M (x r )) -v r , P M (x r+1 ) -xr+1 + L 2 - 1 2η P M (x r+1 ) -P M (x r ) 2 + L 2 + 1 2η xr+1 -P M (x r ) 2 - 1 -ηρ 2η P M (x r+1 ) -xr+1 2 ,where we use η ≤ γηg2D f to guarantee P M (x r+1 ) ∈ U M (γ).Recall thatP M (x r+1 ) = P M P M (x r ) - η 1 nτ n i=1 τ -1 t=0 gradf i z r i,t ; B r i,t :=v r .Combining the above inequality, (33) , and ( 34) yieldsEQUATION)≤ E f (P M (x r )) + L - 1 2η + ρ 2 ∥x r+1 -P M (x r )∥ 2 + L 2 - 1 2η ∥P M (x r+1 ) -P M (x r )∥ 2 - 1 -ηρ 2η P M (x r+1 ) -xr+1 2 (IV) + P M (x r+1 ) -xr+1 , gradf (P M (x r )) -v r (V).According to ∥a + b∥ 2 ≤ 1 2η ∥a∥ 2 + 2 ∥b∥ 2 , we haveEQUATION≤ (IV) + 1 -ηρ 2η ∥P M (x r+1 ) -xr+1 ∥ 2 + η 2(1 -ηρ) ∥gradf (P M (x r )) -v r ∥ 2 = η 2(1 -ηρ) ∥gradf (P M (x r )) -v r ∥ 2 .To bound the above inequality, following similar derivations as in (23) we obtainEQUATIONSubstituting ( 36) and ( 37) into (35) , we haveEQUATION) ≤E f (P M (x r )) + L - 1 2η + ρ 2 xr+1 -P M (x r ) 2 + L 2 - 1 2η P M (x r+1 ) -P M (x r ) 2 + η 2(1 -ηρ) 2L 2 nτ n i=1 τ -1 t=0 ∥z r i,t -P M (x r )∥ 2 + 2 τ n σ 2 b .The final term is the drift-error that can be bounded in (28) . Thus, (38) becomesEQUATION≤ E f (P M (x r )) + L - 1 2η + ρ 2 xr+1 -P M (x r ) 2 + L 2 - 1 2η P M (x r+1 ) -P M (x r ) 2 + η 2(1 -ηρ) 2L 2 nτ 12nM 2 τ 3 η 2 ∥G ηg (P M (x r ))∥ 2 + 9τ E∥Λ r -Λ r ∥ 2 + 18nτ 2 η 2 σ 2 b + η 2(1 -ηρ) 2 τ n σ 2 b ,where we use ∥gradf(P M(x r ) )∥ ≤ 2∥G η (P M(x r ) )∥ from Lemma (A.2). By substituting η ≤ 1 4ρ as η ≤ γ 6D f , we have η 2(1-ηρ) ≤ η.Combining the recursions given by Lemma A.3 and (39), we have for the Lyapunov function thatEQUATION≤E (f (P M (x r )) -f ⋆ ) + 1 ηn ∥Λ r -Λ r ∥ 2 - η 8 ∥G(P M (x r ))∥ 2 + 8η nτ σ 2 b ,where we substitute conditions (11) on the step sizes and omit straightforward algebraic calculations. Substituting the definition of the Lyapunov function Ω r and repeating the above inequality, we complete the proof of Theorem 4.3.Discussion on step sizes. In our analysis, the condition η g = √ n is applied only in the final step, when combining Lemma A.3 and (39) to derive the recursion on the Lyapunov function in (40) . Up to this point, we retain η g as a variable without assigning a specific value, ensuring that all preceding results remain valid for η g = 1 as well. We choose η g = √ n in (40) to conveniently cancel out the number of clients n, yielding a more concise expression in the final result.In (11) , we require that η := η g ητ ≤ min (39) ; the term 1 D f L P is used to derive (16) , which ensures that when the metric ∥G η (P M (x r ))∥ approaches zero, the first-order optimality condition is met. The term 1 24M L is used to establish the Lyapunov function recursion (40) by combining the recursions in ( 17) and (39) . Additional conditions on the step sizes ensure that the iterates remain within the 2γ-tube during our analysis, such as in ( 17), (33) , and ( 34), but they are implied by the three terms in (11) .1 24M L , γ 6D f , 1 D f L P . The term γ 6D f controls the client drift, seeThe settings for kPCA problem with Mnist Dataset. The Mnist dataset consists of 60,000 handwritten digit images ranging from 0 to 9, each with dimensions of 28 × 28. We reshape these images into a data matrix A ∈ R 60000×784 . To construct the heterogeneous A i , we sort the rows in increasing order of their associated digits and then split every 60000/n rows, with n = 10 as the number of clients, among each client. This data partition introduces significant data heterogeneity. In our setup, d = 784, p = 6000, and k = 2.For kPCA problem with Mnist dataset, when τ = 10, the comparison on f (x r ) -f ⋆ is shown in Figs. 5. Our algorithm performs much better than alternative methods in terms of communication quantity and run time. Synthetic Dataset. We also solve kPCA with synthetic datasets on larger networks with n = 30.We generate each entry of A i from Gaussian distribution N 0, 2i n such that A i are heterogeneous among clients. We set (d, k) = (20, 5) and p = 15. We use the local full gradient ∇f i to remove the influence of stochastic Riemannian gradient noise.In the first set of experiments, we compare with existing algorithms, including RFedavg, RFedprox, and RFedSVRG. For all algorithms, we set the number of local steps as τ = 5 and the step size as η = 4e-3. For our algorithm, we set η g = 1. The experimental results are shown in Figs. 6 . The y-axis represents ∥gradf (x r )∥ and (f (x r ) -f ⋆ ) respectively, while the x-axis represents the number of communication rounds, communication quantity, and run time, respectively. It can be observed that RFedavg and RFedprox face the issue of client drift, hence they do not converge accurately. Both FedSVRG and our algorithm can overcome the client drift issue, but our algorithm is slightly faster in terms of communication rounds and is much faster in terms of both communication quantity and run time. For numerical tests, we consider random generated A. To be specific, we first generate two random matrices L ∈ R d×k and R ∈ R k×T , where each entry obeys the standard Gaussian distribution.For the indices set Ω, we generate a random matrix B with each entry following from the uniform distribution, then set Justification: The abstract outlines the topic of the research, the methods used, and the results obtained. In the introduction section, we summarize in detail the three contributions of this paper: algorithm contributions, theoretical contributions, and experimental results, all of which are explained in detail in the subsequent main body of the paper. Guidelines:Ω ij = 1 if B ij ≤ ν• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The reviewers were in disagreement about this work. While two reviewers were quite positive about this work, one reviewer found the applicability of the work to more general classes of manifolds limited and another reviewer (who gave a superficial review) briefly commented that experiments and discussion of previous work could be improved. After going over the paper myself I tend to also support its acceptance and think that the contribution is above bar, despite the applicability to somewhat restricted (but important) subclass of manifolds. I urge the authors to take the reviewer's comments into attention.

Reviews: review: paper_summary: paper_summary: This paper addresses several computational challenges in federated learning on manifolds, by using a simple projection mapping to bypass the need for exponential mappings, logarithmic mappings, and parallel transports needed in prior work. The proposed technique also enables multiple local updates and full client participation. Theoretical iteration complexity is provided and various numerical experiments are conducted to demonstrate the sample-communication efficiency of the proposed method.
strength_weakness: strength_weakness: I am overall satisfied with the originality and quality of the paper, with some suggestions:

- Nonconvexity
    - First of all, the authors should define the convexity of functions on manifolds, especially since submanifolds are considered. Is the convexity defined in the ambient space or on manifolds?
    - Based on my reading, the nonconvexity of the problem involves both the nonconvexity of the manifold constraint and the nonconvexity of the objective function. As mentioned in the paper, the manifold constraint can be incorporated into the objective function as an indicator function. The authors should clearly state which aspect of nonconvexity they refer to.
    - For example, in the first sentence of Paragraph *Federated learning on manifolds*, "Since existing composite FL [...] the loss functions is convex, these methods and their analyses cannot be directly applied to FL on manifolds." It is not crystal clear whether the non-applicability is due to the nonconvexity of the manifold constraint or the objective function considered.
    - Also, it is mentioned that the algorithm is inspired by proximal FL for strongly convex problems in Euclidean spaces. Then based on my understanding, the contribution lies in not only handling the manifold constraint, but also the nonconvexity of the objective function. However, the latter part is not discussed in Section Introduction. Could the authors clarify this? The authors should also mention that related work [13] can also handle geodesically nonconvex objective functions.
- Numerical experiments
    - Most experiments use full gradients instead of stochastic gradients.
    - In the kPCA experiment, the number of clients is $1 = n = \eta _{g}^{2}$ according to Theorem 4.3. Could the authors clarify this?
- Notation and typos
    - I would suggest moving some necessary notation from Appendix A.1 to the main text to make the analysis more self-contained.
    - Is $D$ in the definition of $L_{\mathcal{P}}$ in Theorem 4.3 missing a subscript $f$?
    - The third line of Paragraph *Theoretical contributions* is missing a closing parenthesis.
questions: questions: - The proposed methodology and underlying ideas rely heavily on the geometry of the ambient space. Could the authors comment on the difficulty of extending the methodology to more general manifolds where the ambient geometry is not available, or where the Riemannian geometry is not induced by the ambient space?
- Could you provide intuitive explanations for the step size constraints in Eq. (11)? Are these constraints purely for restricting the iterates to remain within the $2\gamma$-tube, or do they also serve to control the client drift?
limitations: limitations: Limitations and open problems are adequately discussed.
paper_summary: paper_summary: This article proposes a new algorithm for non-convex optimization on smooth manifolds in the federated learning setting. The algorithm is proved to have a sub-linear convergence guarantee, and numerical experiments on kPCA and low-rank matrix completion demonstrate the convergence and efficiency of the proposed algorithm.
strength_weakness: strength_weakness: I do not see obvious weaknesses of the proposed algorithm, but one potential issue is how difficult the projection operator $\mathcal{P}\_{\mathcal{M}}$ is for common manifolds. Can the author(s) provide some known examples in which the expression of $\mathcal{P}_{\mathcal{M}}$ has closed forms? How is the projection operator for the Stiefel manifold computed in the numerical experiments?
questions: questions: Besides the question given in the "Weakness" section, there is one more regarding the experiments. Although the theoretical analysis in Section 4 gives a sub-linear convergence rate, the numerical experiments mostly demonstrate a linear-like convergence pattern. Do you have any insights or conjectures for this phenomenon?
limitations: limitations: The author(s) have discussed the limitations of the proposed algorithm in Section 6.
paper_summary: paper_summary: The authors introduce a submanifold into the setting of federated learning. Utilising Riemannian gradients for optimisation, they show different properties of the proposed method, such as avoiding client drift and being computationally cheaper.
strength_weakness: strength_weakness: - A central point of the paper is to not work on the manifold directly, but to utilise a projection in order to guarantee proximity. This moves the problem away from working on manifolds directly - This could be advantageous, but the presentation as working on a manifold is unnecessary and misleading.
- The presentation is unclear - while restricting themselves to only Euclidean immersions using the Euclidean metric, they state several results for which it is unclear wether these results are for general manifolds or only for this case. For the general case, the are certain problems (i.e. injectivity radius) which would need attention.

- The Lipschitz continuity in eq(3), needed for Lemma 1.2, is proven in [33] in theorem 4.8 using the Cauchy-Schwarz inequality. The generalisation to manifolds is not apparent.

[33] "Francis H Clarke, Ronald J Stern, and Peter R Wolenski. Proximal smoothness and the lower-C2 property. Journal of Convex Analysis, 2(1-2):117–144, 1995"
questions: questions: -  Please clarify the theoretical results in terms of assumptions on the manifold needed.
-  How exactly is the proximal operator carried out? If this operator is restricted to points on the manifold, is the knowledge of which points belong to the manifold known? If yes, why not use that directly? As the proof of Lemma A.1 depends on belonging to the manifold, this appears pivotal to the concept. In [33], the proximal operator is only used in real Hilbert spaces, so belonging to a manifold is a problem arising only now.
-  Line 269 ff.: The comparison of the theoretical result with [13] seems odd, as the Li et al. do not use a proximal operator. As this operator massively eases the problem, this should at least be noted clearly - these results improve computational efficiency on another problem setting. Also, Li et al introduce a "tangent mean", where they appear to ignore the different bases of tangent spaces, more specifically that tangent spaces will differ in density in the presence of curvature. Both results are likely to not generalise to any given smooth manifold.

[13] "Jiaxiang Li and Shiqian Ma. Federated learning on Riemannian manifolds. arXiv preprint arXiv:2206.05668, 2022."
limitations: limitations: Limitations are properly addressed.
paper_summary: paper_summary: This paper studies federated learning on a compact smooth submanifold using distributed Riemannian optimization methods.

The proposed methods were examined on PCA and low-rank matrix completion tasks.
strength_weakness: strength_weakness: In the theoretical analyses, a comparison with the other related Euclidean and Riemannian methods can be given.

Experimental analyses are limited and should be improved as well.
questions: questions: There are several other works on distributed Riemannian optimization methods for similar tasks, such as
Distributed Principal Component Analysis with Limited Communication, Neurips 2021.
Communication-Efficient Distributed PCA by Riemannian Optimization, ICML 2020.

Could you please compare your proposed methods with these related work theoretically and/or experimentally?

The number of clients considered in the experimental analyses is pretty low. Could you please examine how the results change with different larger number of clients?

How did you measure CPU time?

Could you please provide comparative analyses with Euclidean baselines, such as FedAvg and FedProx?

In FL, a major issue is privacy, and therefore, several mechanisms, such as differential privacy, are utilized. Could you please provide analyses with different privacy mechanisms?

One of the major claims of the paper is improvement of the results compared to the baseline for optimization with heterogeneous data. However, there are not detailed analyses justifying this claim. Could you please provide additional analyses exploring the methods for different data heterogeneity?
limitations: limitations: Limitations were partially addressed.

Comments: comment: Dear Reviewer GhoW,

Thank you for engaging in the discussion with us, and for your valuable feedback on our manuscript. We regret that our introduction misled you. It was never our intention, and we will certainly update it to ensure that the scope and limitations of our work are clear from the start.

With this said, the manifold geometry is critical to our work, and it is not clear if and how similar results could be derived by instead focusing on general nonconvex constraints in Euclidean space.


**For the weakness:**
* We use the manifold geometry to compute the Riemannian gradient, i.e., ${\rm grad} f_{il}(z_{i,t}^r;\mathcal{D}\_{il})$ in our Algorithm 1. Specifically, the geometry of the manifold is essential for computing both the tangent space of the manifold and the Riemannian gradient (which is the projection of the Euclidean gradient onto the tangent space). In contrast, for a general nonconvex constraint set in the Euclidean space, determining the tangent space and its projections can be challenging. The Riemannian gradient is critical for both our algorithm design and theoretical analysis.

* Furthermore, the projection operator onto the compact smooth submanifold embedded in the Euclidean space has the property of proximal smoothness. The geometry of the considered manifold helps to  establish 
   a quasi-isometric property between $\mathcal{G}\_{\tilde{\eta}}(x)$ and   ${\rm grad} f(x)$, i.e., $1/2 \\| {\rm grad} f(x) \\|\le \\|\mathcal{G}\_{\tilde{\eta}}(x)\\| \le 2\\| {\rm grad} f(x) \\|$ (see Line 259, Lemmas A.1 and A.2). 

These benefits are not always available for a nonconvex constraint set in the Euclidean space. We do not currently have a clear path for how to extend our results to general manifolds or nonconvex constraint sets, but this would be an interesting topic for future studies.

**For the Question:**

In our Algorithm 1, we calculate $z_{i,t}^r:= \mathcal{P}\_{\mathcal{M}}(\hat{z}\_{i, t}^r)$ and $\mathcal{P}\_{\mathcal{M}}(x^r)$,  where $z_{i,t}^r:=\mathcal{P}\_{\mathcal{M}}(\hat{z}\_{i, t}^r)$ ensures that our algorithm calculates the Riemannian gradient $\operatorname{grad} f_{il}(z_{i,t}^r; \mathcal{D}\_{il})$ at a point on the manifold, and $\mathcal{P}\_{\mathcal{M}}(x^r)$ ensures that the global model $\mathcal{P}\_{\mathcal{M}}(x^r)$ on the server satisfies the manifold constraint. In the case of the Stiefel manifold, we take advantage of the explicit expressions of the Riemannian gradient ${\rm grad} f_{il}(z) = \nabla f_{il}(z) - z \\, {\rm sym}(z^\top \nabla f_{il}(z))$ and the projection operator $\mathcal{P}\_{\mathcal{M}}(x) = x(x^\top x)^{-1/2}$ to obtain superior performance of our algorithm compared with the state-of-the-art, as shown in the numerical experiments. As stated in Eq.(1), the constraint set $\mathcal{M}$ can be any compact smooth submanifold embedded in the Euclidean space, which includes nonconvex compact sets in the Euclidean space with a smooth submanifold structure.

We appreciate that you view our paper as a valuable contribution to federated learning, and we hope that our additional explanations have addressed your remaining technical concerns. 
We will revise our manuscript to avoid potential misunderstandings about the scope and limitations of our work, and would appreciate it if you could kindly consider raising your rating.
Thank you for the clarifications!

While in general the answers improved my understanding of your point of view, I am still convinced that the presentation is misleading. The proposed method seems like a good improvement in the area of federated learning, but the writeup is currently letting the reader believe the method is working on any smooth submanifold. In the paper and the rebuttal (Weaknesses, 2.) it becomes clear that the results only hold for the choice of Euclidean space, so the presentation as manifolds is unnecessary. 

Some specific follow-ups:

**Weaknesses**

1. If the setting is only Euclidean space, what is the manifold geometry you are considering? 

**Question**

2. i) The existence of the solution for the Stiefel manifold is obviously nice, but it is not used in the work. As it is noted by now that the given work does not generalise to manifolds, this does not appear necessary. 

I want to reiterate, I believe this is a valuable contribution to federated learning and improves on current difficulties. To reflect this, I am updating my scores. The impression of working on manifolds is massively disadvantageous for clarity of reading and precisely understanding the contribution. With a major revision changing the narrative away from manifolds I'd be happy to accept this paper.
Thank you for your quick response. The explanation is clear. If I understand correctly, there are two sources of client drift: drift from data heterogeneity, which is corrected using variance reduction, and drift from multiple local updates, which is controlled by the step-size constraint.
Thank you so much for your reply! We will incorporate the promised revisions and your feedback in the revised version of the paper.

For your additional question on the step size, we give an explanation of why we need the step size condition in our convergence proof. The condition $\tilde{\eta} \le \frac{\gamma}{6D_f}$ comes from Eq.(39) and has two purposes: i) to control the error of SGD caused by noise and ii) to limit the client drift. To see this, note that Eq.(39) gives
$$
\mathbb{E}[f^{r+1} - f^{\star}] \lesssim \mathbb{E}\left[f^{r} - f^{\star} + A \frac{2L^2}{n\tau} 9\tau \\|\mathbf{\Lambda}^r-\overline{\mathbf{\Lambda}}^r\\|^2 + A \frac{2}{\tau n} \frac{\sigma^2}{b} \right],
$$
where $\lesssim$ means approximately smaller than or equal to, in the sense that we have ignored other terms that are not important for explaining the step size condition, and $A:= \frac{\tilde{\eta} }{2(1-\tilde{\eta}\rho)} $.  We require $A\le \tilde{\eta}$, which leads to the condition $\tilde{\eta} \le \frac{\gamma}{6D_f}$. 


To clarify why it is necessary to use a small step size to control the client drift, let us 
 assume $\sigma^2=0$. In the following, we will demonstrate that if $A$ is very large, we cannot establish the recursion of the Lyapunov function. Indeed, from Eq.(17), we know that
$$
\frac{1}{\tilde{\eta} n} \mathbb{E} \\| \mathbf{\Lambda}^{r+1} -\overline{\mathbf{\Lambda}}^{r+1}\\|^2 \lesssim \left( 4\eta^2 \tau L^2 9 \tau \right) \frac{1}{\tilde{\eta} n} \mathbb{E} \\| \mathbf{\Lambda}^{r} -\overline{\mathbf{\Lambda}}^{r}\\|^2,
$$
where we omit the other terms for simplicity. 
Adding the two inequalities, we get
$$
\mathbb{E}[f^{r+1} - f^{\star}] 
+\frac{1}{\tilde{\eta} n} \mathbb{E} \\| \mathbf{\Lambda}^{r+1} -\overline{\mathbf{\Lambda}}^{r+1}\\|^2
\lesssim \mathbb{E}\left[f^{r} - f^{\star}\right]  + \left(\tilde{\eta} A 2L^2 9 +  4\eta^2 \tau L^2 9 \tau \right) \frac{1}{\tilde{\eta} n} \mathbb{E}\\|\mathbf{\Lambda}^r-\overline{\mathbf{\Lambda}}^r\\|^2.
$$

If $A$ is very large, then we cannot establish the recursion of the Lyapunov function. Instead, by substituting $A\le \tilde{\eta}$ and step size conditions (11), we have $\left(\tilde{\eta} A 2L^2 9 +  4\eta^2 \tau L^2 9 \tau \right)\le 1$. 
Thus, we need a small step size to control client drift.
Thank you for the detailed rebuttal. I am overall satisfied with the clarifications and the promised revisions, and I am happy to maintain my scores. Specifically, I believe that the authors' clarifications on the second point of the Numerical Experiments and the third point of the Questions section would be helpful in enhancing readers' comprehension of the paper; I recommend that these clarifications be included in the main text of the revised version.

I do have one further question: Given the usage of the client drift correction mechanism, why is a step-size constraint still needed to control client drift?
Thank you very much for the reply and the review!
Thanks for the response. My questions have mostly been addressed.