Nonconvex Federated Learning on Compact Smooth Submanifolds With Heterogeneous Data
Authors: Jiaojiao Zhang
Jiang Hu
Anthony Man-Cho So
Mikael Johansson
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated learning
manifold optimization
heterogeneous data
Primary area: optimization
TL;DR: This paper proposes a computation- and communication-efficient algorithm for federated learning over manifolds with heterogeneous data.

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.

Nonconvex Federated Learning on Compact Smooth Submanifolds With Heterogeneous Data
Jiaojiao Zhang', Jiang Hu?; Anthony Man-Cho So
Mikael Johansson' !
KTH Royal Institute of Technology
2 University of California, Berkeley 3 Thhe Chinese University of Hong Kong
{jiaoz, mikaeljjEkth. se, hujiangopt&@maill com, manchosodse, cuhk, edu. hk
Abstract
Many machine learning tasks, such as principal component analysis and low-rank matrix completion, give rise to manifold optimization problems. Although there is large body of work studying the design and analysis of algorithms for manifold optimization in the centralized setting, there are currently very few works address- ing the federated setting. In this paper, consider nonconvex federated learning over a compact smooth submanifold in the setting of heterogeneous client data. We propose algorithm that leverages stochastic Riemannian gradients and mani- fold projection operator to improve computational efficiency, uses local updates to show improve communication efficiency, and avoids client drift. T'heoretically, that our proposed algorithm converges sub-linearly to a neighborhood 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.
first-
1 Introduction
Federated learning (FL), which enables clients to collaboratively train models without exchanging their raw data, has gained significant traction in machine learning [I []. 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 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 manifolds in the form of
1 M f((z), f((z)
my
minimize REMC
f(z)
fa(z;
(I)
m
Here, n is the number of clients, x is the matrix of model parameters, and M is compact smooth submanifold embedded in Ridkk, Examples of such manifolds include the Stiefel, oblique, and symplectic manifolds 3
4 manifold M St,d, k) {ar E Rdxk xT x Ig maintain orthogonality 6 7. In (I. the global loss f dx + R is smooth but nonconvey and the local loss function f: of each client is the average of the losses fa on the m, data points in its local dataset D. = ......... Dim, }. We consider a heterogeneous data scenario where the statistical properties of D. differ across clients.
[5). For instance, PCA-related optimization problems use the Stiefel
"Corresponding author: Jiang Hu 'hhroughout this paper, all convexity-related concepts are defined the Euclidean space.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Manifold optimization problems of the form ①1) appear in many important machine learning tasks, such as PCA 8 99. low-rank matrix completion [O] TT]. multitask learning [2 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 FI. 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 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] [9] 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 local updates and heterogeneous data [9] 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.
1.1 Contributions
We consider the nonconvex FL problem ① with M being a compact smooth submanifold and allow for heterogeneous data distribution among clients. Our contributions are summarized as follows.
I) We propose an FL algorithm for solving |[ that is efficient in terms of both computation and communication. We employ stochastic Riemannian gradients and a projection operator address manifold constraints, use local updates to reduce the communication frequency between clients and the server, and design correction terms 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 [6 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 [66 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 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 [6] 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
2

---

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 exper- iments. 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.
1.2 Related work
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 ① can be viewed as 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 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 assumes bounded data heterogeneity. The work [22] introduces the federated Douglas-Ran 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 communication- efficient 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 considers scenarios where the nonsmooth term in the loss functions is convex. However, incorporating nonconvex manifold constraint as an indicator function introduces nonconvex nonsmooth term. Consequently, the methods in 20 24 are not directly applicable. A typical challenge caused by the nonconvex manifold constraint 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 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 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 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 sphere. In contrast, we consider a more general setting with heterogeneous data, where lies on M and the loss functions are smooth and nonconvex.
only
Notations. We use I3 to denote k x i identity matrix. We use ||- denote Frobenius norm and tr(-) to denote the trace of a matrix. For a set B, we use [ to denote the cardinality. For random variable v, we use Elu] denote the expectation and E[ denote the expectation given event F. For an integer n,, we use [nx] to denote the set [1,...,n}} n.]. For two matrices x, E gdxk, we define their Euclidean inner product (I. : > Σ T4j3ii- For matrices 21,..., An E Rdck we use z = (2 := :.:;...;:0 2n] E kdn to denote the vertical stack of all matrices. The bold notations ž, c, and A are defined similarly. Specifically, for a matrix
3

---

E R we define x := ;= == [z:...;;; : E Rudxk, We use r to denote the index of the communication round and denote the index of local updates. Given the local Riemannian gradient gradfi e ;3,) at point ,f, with the mini-batch dataset B% we define the stack of Riemannian gradients as gradf (2]; B3) = fi 1i and the stack of average local Riemannian gradients as gradf (2;; B]] == col ( grad B Given col (ol == and (M.(2s) define P. ([:col22]]]] (==) col(PM
Given positive definite matrix I, to denote the inverse of the square root of ie., x 1/2z 11/2 x=1 We define D2 to be the second-order differential operator.
x 1/2
Preliminaries
Below, we introduce fundamental definitions and inequalities for optimization manifolds.
2.1 Optimization manifolds
Manifold optimization aims to minimize a real-valued function over a manifold, i.c., minge. f(x). Throughout the paper, restrict 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 embedded submanifolds. Some examples of such manifolds include the Stiefel manifold, oblique manifold, and symplectic manifold B 4 5 We define the tangent space of M at point T,M, which contains all tangent vectors M æ; and the normal space N_M which is orthogonal to the tangent space. With the definition of tangent space, we can define the Riemannian gradient that plays central role in the characterization of optimality conditions and algorithm design for manifold optimization.
Definition 2.1 (Riemannian gradient gx)) The Riemannian gradient grad (x) of a function f at the point E M is the unique langent vector that satisfies
fgrad E) = df)) VE E T.M,
where ). is the Riemannian metric and df denotes the differential finction f.
For submanifold M, the Riemannian gradient fx) (under the Euclidean inner product) can be computed as (5) Proposition 3.61]
gradf(z) Pr..((f(z))
where Pr.M( (..M(V/J) (v)()) represents the orthogonal projection of V f() onto T. The Riemannian gradient gradf(x) reduces to the Euclidean gradient fz( when M is the Euclidean space Rd
2.2 Proximal smoothness of M
In our federated manifold learning algorithm, the server needs fuse models that have undergone multiple rounds 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 map the local models onto tangent space using inverse exponential mapping, calculate an average on the tangent space, and then perform an exponential mapping 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 correction terms. To overcome this difficulty, we use projection operator PM defined by
(M( E argmin || u||²
NEM
to ensure the feasibility of manifold constraints. The projection operator PM can be explicitly calculated for many common submanifolds, as discussed in (35). For the Stiefel manifold, the closed- form expression for (M((z)
given matrix with full column rank is (M) x) see 35 Proposition 7]. It is worth noting that PM can be regarded as a special retraction operator when restricted to the tangent space (5). However, unlike a typical retraction operator, domain d 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 (M(()) may be set-valued and non-Lipschitz, making the use and analysis of PM in the FL setting highly nontrivial. To tackle this, we introduce the concept of proximal smoothness that refers to property of closed set, including M, where the projection becomes a singleton when the point is sufficiently close the set.
Definition 2.2 (6-proximal smoothness of M). For any & > O, we define the Sy-tube around M as (+) {a: dist( M) 7)
where dist(x, M) := min I|u - æl| is the Eulidean distance between x and M. We say that M is "-proximally smooth if the projection operator (M( is singleton whenever UM(()..
It is worth noting that any compact smooth submanifold M embedded in Rd* is proximally smooth set 33 (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 Ra* with the Euclidean inner product serving its Riemannian metric. Moreover, assume that the proximal smoothness constart of M is 27-
With Assumption 2.3
ensure only the uniqueness of the projection but also the Lips- chitz continuity of the projection operator PM around M, analogous to the non-expansiveness of projections under Euclidean convex constraints.
Lipschitz continuity of P.M Define U (7) == {o dist(z, M) S ? as the closure of (M))). Following the proof in 36 Theorem 4.8], for 27-proximally smooth M, the projection operator PM is 2-Lipschitz continuous over U (7) such that
- - 2|1x vll. (e,yEUM(1)
Normal inequality. the normal space NzM, we exploit the so-called normal inequality 36 Following (36), given a 2 proxiiialy smooth M, for any E M and v E N,,, it holds that
||||| bly =1 4y
(v, x) S
Vy E M.
Intuitively, when & and are close enough, the matrix y- - x is approximately in the tangent space, thus being nearly orthogonal the normal space.
3 Proposed algorithm
In this section, we develop a novel algorithm for nonconvex federated learning on manifolds. The algorithm 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 handle nonconvex loss functions, which call for different convergence analysis.
3.1 Algorithm description
The per-client implementation of algorithm is detailed in Algorithm [ Similarly to the well- known FedAvg, it operates in a federated learning setting with one server and n clients. Each client engages in 7 steps of local updates before updating the server. We use as the index of communication rounds and f. as the index of local updates.
At any communication round r, client downloads the global model x" from the server and computes PM (x"). Each client updates two local variables, oh, and 2f,,, where 2 aggregates the Riemannian gradients from local updates, and = = (x)(2()) ensures that Riemannian gradients can be computed points on M. The update of 2,, is given in Line 8, where B, is mini-batch dataset and of is correction term to eliminate client drift. After 7 local updates, client sends ب to the server.
The server receives all T computes their average to form the global model ++1 following Line 13, and broadcasts x'+1 to each client i that uses r+1 to locally construct the correction term c
5

---

Algorithm Proposed algorithm
1: Input: R, T, 1), f1g- f) "yg, at, and ct for all i [n]
2: for for-=1 1, 1.2,.... R do
3: Client
4: Set 2r = (2))) and 2 = (M*
5: for 0, ...,,...,.......................
do
6:
Sample a mini-batch dataset 1 C Di with - Update (ff, [(,) 1 ED grad Pu)
7:
8:
Update 21,41 2,, (grad (,,) c.)
9:
Update zf,411 = P.M (2f,441)
10: end for
Send , to the server
12: Server
13:
Update x PM((2)) + n1 ( -i o0 - PM(l7))
14: Broadcast x=+1 to all the clients
Client
16:
Receive x1+1
Update of NOTT -2=0)=))=++))))))))))))))))))))) Et= gradfi B(,,) (i.,)
18:
end for
19:
Output: (2
In the proposed algorithm, each client i downloads x" at the start of local updates and uploads af.- at the end of the local updates. Therefore, each communication round involves each client and the server exchanging only single d k matrix.
3.2 Algorithm intuition and innovations
To better understand the proposed algorithm, we present its equivalent and more compact form:
1/ = i -n gradf (zi; B3) + gradf (z, ; :,=)))
Σ gradf ; :;; 1)
+1 = PM (2)+1),
(5)
x = PM(xx)) 17
gradf (z,;
1=u
For the initialization of correction term, gradfi (2h,5 B..)) for all t and that - gradf (2f; 35. - Σn gradf (zf: BB) = 0, which coincides with the initialization 0 in Algorithm The equivalence between Algorithm and (5) can be proved following the and is therefore omitted.
similar derivations
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 gradf (2];B] and 7 into the last step in (5). we have
P (PM(())
1 nt
(gradfi (zf, B(,)))
(6)
Thanks to the introduction of the variable a., during the local updates for each client in Algorithm I. the server after averaging ěr obtains an accumulation of 7 local Riemannian gradients across local updates and an average of the local Riemannian gradients across all clients. In the special case where 7 = and = m,, 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)
#++1
PM (PM(x") ngradf
(7)

---

In our analysis, we will compare the sequence (x generated by our algorithm with the virtual iterate 3r+1 to establish the convergence of our algorithm.
2) Feasibility of all iterates at a low computational cost. Our algorithm uses PM 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, cannot be directly used in our algorithm. In the local updates, is difficult to perform exponential mapping on 2[+1 because i is not on the manifold; see the first step in (5). As shown in (24], 2 is essential for the server to obtain aggregated Riemannian gradients from n clients after r local updates. Moreover, at the server, although (M((x" is on the manifold, the aggregated direction does not lie in the tangent space at P.u (x"). The algorithm suggested in [[6] 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 computationally expensive. Our use of PM on point in the Fuclidean 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 to address client drift. According the first step of 5). the correction employs the idea of "variance reduction", which involves replacing the old local Riemannian gradient
Σi= gradf (z, BT with the new one gradf (z"; B3) in the average of all client Riemannian gradients Ei= gradf (z, BB 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 conrection term can be locally generated, leading to 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.
4 Analysis
In this section, we analyze the convergence of the proposed Algorithm ! Throughout the paper, we make the following assumptions, which are common manifold optimization.
Assumption 4.1. Each fa(r;; Da): : Rdxk b TR. has i.-Lipschitz continuous gradiens V fa(;;; Dai) on the convex hull of M, denoted by conv(M) ie., for any x, y E conv(M), it holds that
[V Da) V fuly: Da)ll S Í|| vll.
(8)
With the compactness of M. there exists a constant Dj > 0 such that the Euclidean gradient [a( Da)) of sa is bounded by D,, ie., max | [a(x;; Da)ll S D,. It then follows from
28 Lemma 4.2] that there exists a constant L < L < oo such that for any 2, E M,
L a - v113,
faly; Da) S fa(a; Da) + (gradfa Da), - x) + 2
(1;; Da) - ( Da)ll S A|| - -
To address the stochasticity introduced by the random sampling B define Fj. as the o-algebra generated by the set [B | [m], E [r], E [t - 1}} and make the following assumptions regarding the stochastic Riemannian gradients, similar to (40, Assumption 2}.
Assumption 4.2. Each stochastic Riemannian gradient gradfi Bi,)) in Algorithm 7 satisfies
E [gradd((=T [[ = gr
0²
E [1 Bi.) f. 11 | I
b
Considering the nonconvexity of f and the manifold constraints, we characterize the first-order optimality of (T). point a² is defined as a first-order optimal solution of ( if x E M and
gradf(z*) = 0. We employ the norm of G (PM (x")) as suboptimality metric, defined as
1 (s(PM(x')) = (PM(=") - -
(10)
7

---

and ø++1 defined in 7 is used only for analytical purposes. In optimization on Euclidean space such that .M- = R the quantity (r(PM serves as a widely accepted metric to assess first- order optimality for nonconvex composite problems A1. In optimization on manifold, we have (²(PM(z") = 0 if and only if
= for any vj > Moreover, for a suitable ij, the use of Riemannian gradients in the update of 5 helps us establish a quasi-isometric property, specifically that
I || (PM((2)) I)ll I 2 g (PM((=))) (z"))];; see Lemmas A.1 and A.2 With (g((M((")) we have the following theorem.
Theorem 4.3. Under Assumptions
2.3
and 4.2 the step sizes satisfy
7 1
24ML' 6D," DyLp) fg Vn, max (diam(MM)7:) 2}. diam(M)
i) 1T min
(11)
where M
max vll. Dj
maxitaEM V (a;;; Da)ll, and Lp = ZEU (M))l( then the sequence (m(x")
generated by Algorithm 0 satisfies
Σ
822
64
(12)
R
V
nrb"
where Ω2 > 0 is a constant related to initialization.
In Theorem 4.3 the first term on the right hand of 12
converges at sub-linear rate, which is The second term a constant error caused
common for constrained nonconvex optimization 41 42
by the variance o² of stochastic Riemannian gradients.
Theoretical contributions. The work [6 establishes convergence rates of OL/RR) for r 1 and (()((TR) for r > 1 but only if a single client participates in the training per communication round. In contrast, our Theorem 4.3 achieves a rate of ((L((VnTR) for t > 1 and full client participation. Our theorem indicates that multiple local updates enable faster convergence, which distinguishes our algorithm from decentralized manifold optimization algorithms ( (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 B- local dissimilarity @1 Definition 3], respectively, to establish convergence. These assumptions imply some similarity between client data. Our main analytical advantages over [43) and []] 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 several novel techniques. Specifically, capitalize the strueture of M and exploit the proximal smoothness of M to guarantee the uniqueness of PM and Lipschitz continuity of PMM within 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, select appropriate first-order optimality metric (see (TO)) 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.
5 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 RFeaavg Reedprox, and RFedSVRG. RFedavg and RFedprox are direct extensions of FedAvg [25] and Fedprox [1] For RFedSVRG. there are no theoretical guarantees when we set 7 > 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 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
n M(x), f((x) = -irla" A Atx),
1
minimiie f(x)= zES((d
n
8

---

the step where St(d, k) f E Raxk | æ²xx = } denotes the Stiefel manifold, and ATA, E R is the covariance matrix of the local data A, E Rexd of client i. We conduct experiments where the matrix A, is from the Mnist dataset. The specific experiment settings can be found Appendix A.44.
10a
10a
10f
Igrad (xO)
10a
Igrad ft')
107
Igrad (tO
10a
10
10
10
1000
1500 2000
1000
2000
4000
200
Communication round
Communication quantily
Runtime (s)
RFedavg
RFedprox
RFedSVRG
Figure 1: kPCA problem with Mnist dataset: Comparison on (grrad(("')
In the first set of experiments, we compare our algorithm with RFedavg, RFedprox, and RFedSVRG. RFedS requires each client to transmit two d i matrices per communication round, while our algorithm only transmits a single matrix. We use communication quantity to count the total number of d & matrices that per client transmits the server. We the local full gradient V fz to mitigate the effects of stochastic gradient noise. In Figs. D we set T = 10 and = 1/3 for all algorithms, where
the square of the largest singular value of col < We 77 1 to facilitate comparison with other algorithms. As noted below (40) in the Appendix, all analytical It can be observed that RFedavg and RFedprox
results leading up (40) remain valid for n2 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 the impact of For all the algorithms, size = 1/3 and 7 E (10, 15, 20). For our algorithm, we set Y/2 = The experiment results are shown in Figs. For all values of 7, our algorithm achieves better convergence and requires less communication quantity.
10a
==10 10
10a
r= 15
10a
==
20
I|grad r(x']l
10a
Ilgrad [[x
10a
Ilgrad fix"
102
10-2
10-2
10-2
10-6
10-1
10-4
10 10
10
10 10
1000
2000
3000 4000
1000
2000
500
1000
1500
2000
Communication quantity
Communication quantity
Communication quantity
RFedavg
RFedprox
RFedSVRG
our
Figure 2: kPCA with Mnist dataset: The impacts of r.
In addition, we test the impact stochastic Riemannian gradients with different batch sizes. We set 7 1/(203). As shown in Figs. 3 our algorithm converges to neighborhood due to the sampling noise and larger batch size leads to faster convergence. Additional experimental results can be found in Appendix A.4.I
LRMC. LRMC aims to recover a low-rank matrix A E IRdxT' from its partial observations. Let Q be the set of indices of known entries in A, the rank-k LRMC problem can be written as minimize ( ES((d.)), VERAXT [| - A)]l where the projection operator Ps is defined in an
entry-wise manner with (Pn(A)) = Ar, (2,42) E Ω and otherwise. terms of the FL setting, consider the case where the observed data matrix (A(AA) is equally divided into clients

---

15
Tw
10a
10a
10a
Igrad [t)]
[gradf(x)]
Igra
103
10a
10a
10
10
5000
7500
10000
2000
4000
Communication round
Communication round
Communication round
Figure 3: kPCA problem with Mnist dataset: The impacts of stochastic Riemannian gradients.
10a
10a
[grad
10-1
Igrad
10
Igrad
10-1
10-3
10-3
10"5
10:1
10
200
300
400
200
150
200
Communication round
Communication quantity
Runtime (s)
RFedavg
RFedprox
RFed5VRG
our
Figure 4: LRMC: Comparison on Igradf(x"))
by columns, denoted by Ai..... Ag. Then, the FL LRMC problem is
1
minimize xd)t(
IIPn, (XV:(X) A:)11:
(13)
2n
where Qi is the subset corresponding to client i in Ω and V:(X) == argminy ||Pn, (XV - -A,) I. In the experiments, we set T = 1000, d = 100, 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 terms of communication quantity and run time. Additional experimental results can be found in Appendix A.4.2
6 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 the loss function, we prove sub-linear convergence to neighborhood of a first-order stationary point. Numerical experiments show superior performance of 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 PM for certain manifolds necessitates exploring methods to calculate projections approximately. Additionally, smoothness constant 7, underscoring the need for estimating y either off-line for specific manifolds adaptively on-line. Furthermore, designing algorithms for partial participation and devising corresponding client-drift correction mechanisms require further investigation.
step-size selection relies on the proximal
10

---

Acknowledgments
This research is supported in part by the Hong Kong Research Grant Council (RGC) through the General Research Fund (GRF) project CUHK 14205421, in part by the Knut and Alice Wallenberg Foundation through grant 2022.0050, and in part by the Swedish Research Council through grant 2023-05538.
References
[1] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 450, 2020.
[2] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings. et al. Advances and open problems in federated learning, Foundations and Trends in Machine Learning, 14(1-2):1 1-200, 2021.
[3] P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2009.
[4] Jiang Hu, Xin Liu, Zaiwen Wen, and Yaxiang Yuan. A brief introduction manifold optimiza- tion. Journal of the Operations Research Society of China, 1:199-248, 2020.
[5] Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
[6] Shixiang Chen, Shiqian Ma, Anthony Man-Cho So, and Tong Zhang. Proximal gradient method for nonsmooth optimization over the stiefel manifold. SIAM Journal on Optimization, 30(1):210-239) 2020.
[7] Lei Wang and Xin Liu, Decentralized optimization over the Stiefel manifold by an approximate augmented lagrangian function. IEEE Transactions on Signal Processing, 70:3029-3041, 2022.
[8] Haishan Ye and Tong Zhang. DeEPCA: Decentralized exact PCA with linear convergence rate. Journal of Machine Learning Research, 22(238):1-27, 2021.
(9] Shixiang Chen, Alfredo Garcia, Mingyi Hong, and Shahin Shahrampour. Decentralized Rrie- mannian gradient descent the Stiefel manifold. In International Conference Machine Learning, pages 1594-1605. PMLR, 2021.
[10] Nicolas Boumal and P-A Absil. Low-rank matrix completion via preconditioned optimization on the Grassmann manifold. Linear Algebra and its Applications, 475:200-233 2090000-22,222 2015.
[11] Hiroyuki Kasai, Pratik Jawanpuria, and Bamdev Mishra. Riemannian adaptive stochastic gradient algorithms matrix manifolds. In International Conference on Machine Learning, pages 3262-3271. PMLR, 2019.
[12] Nilesh Tripuraneni, Chi Jin, and Michael Jordan. Provable meta-lcarning of linear repre- sentations. In International Conference Machine Learning, pages 10443. PMLR, 2021.
[13] Nikolaos Dimitriadis, Pascal Frossard, and François Fleuret. Pareto manifold learning: Tackling multiple tasks via ensembles of single-task models. In International Conference Machine Learning, pages 8015-8052. PMLR, 2023.
[14] German Magai. Deep neural networks architectures from the perspective of manifold learning. In 2023 IEEE 6th International Conference on Pattern Recognition and Artificial Intelligence (PRAI), pages 1021-1031. IEEE, 2023.
[15] Thomas Yerxa, Yilun Kuang, Eero Simoncelli, and Yeon Chung. Learning efficient coding of natural images with maximum manifold capacity representations. Advances in Neural Information Processing Systems, 36:24103-24128, 2023.

---

[16] Jiaxiang Li and Shiqian Ma. Federated learning on Riemannian manifolds. arXiv preprint arXiv:2206.05663 2022.
[17] Ralf Zimmermann and Knut Huper. Computing the Riemannian logarithm on the Stiefel mani- fold: Metrics, methods, and performance. SIAM Journal on Matrix Analysis and Applications, 43(2):953-965 2022.
[18] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg non-iid data. In International Conference on Learning Representations, 2019.
[19] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132-5143, 2020.
[20] Honglin Yuan, Manzil Zaheer, and Sashank Reddi. Federated composite optimization. In International Conference on Machine Learning, pages 12253-12266, 2021.
[21] Yajie Bao, Michael Crawshaw, Shan Luo, and Mingrui Liu. Fast composite optimization and statistical recovery in federated learning. In International Conference on Machine Learning, pages 1508-1536, 2022.
[22] Quoc Tran Dinh, Nhan H Pham, Dzung Phan, and Lam Nguyen. Rachford splitting algorithms for nonconvex federated composite optimization. Advances in Neural Information Processing Systems, 11:30326-30338 2021.
Douglas-
[23] Han Wang. Siddartha Marella, and James Anderson. FedADMM: A federated primal-dual algorithm allowing partial participation. In 2022 IEEE 61 Conference on Decision and Control (CDC), pages 287-294, 2022.
[24] Jiaojiao Zhang, Jiang Hu, and Mikael Johansson. Composite federated learning with heteroge- neous data. In ICASSP 2024-2024 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages
IEEE, 2024.
[25] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effcoi learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273-1282, 2017.
[26] Sai Praneeth Karimireddy, Martin Jaggi. Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning, arXiv preprini arXiv:2008 2330:00008.03606, 2020.
[27] Aritra Mitra, Rayana Jaafar, George J Pappas, and Hamed Hassani. Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients. Advances in Neural Information Processing Systems, 34: 34:1446661444 14619, 2021.
[28] Kangkang Deng and Jiang Hu. Decentralized projected Riemannian gradient method for smooth optimization on compact submanifolds. arXiv preprint arXiv:2304.08241, 2023.
[29] Jun Chen, Haishan Ye, Mengmeng Wang, Tianxin Huang, Guang Dai, Ivor Tsang, and Yong Liu. Decentralized Riemannian conjugate gradient method on the Stiefel manifold. In The Tweifth International Conference on Learning Representations, 2024.
[30] Zhenwei Huang, Wen Huang, Pratik Jawanpuria, and Bamdev Mishra. Federated learning on Riemannian manifolds with differential privacy. arXiv preprin! arXiv:2404. 10029, 2024.
[31] Tung-Anh Nguyen, Jiayu He, Long Tan Le, Wei Bao, and Nguyen H Tran. Federated PCA on Grassmann manifold for anomaly detection in iot networks. In IEEE INFOCOM 2023-IEEE Conference on Computer Communications, pages 1-10. IEEE, 2023.
[32] Andreas Grammenos, Rodrigo Mendoza Smith, Jon Crowcroft, and Cecilia Mascolo. Federated principal component analysis. Advances Neural Information Processing Systems, 33:6453- 6464, 2020.
12

---

[33] Long- Kai Huang and Sinno Pan. Commumication-effiae distributed pea by riemannian optimization. International Conference on Machine Learning. pages 4465-4474. PMLR. 2020.
[34] Foivos Alimisis, Peter Davies, Bart Vandereycken, and Dan Alistarh. Distributed principal component analysis with limited communication. Advances in Neural Information Processing Systems, 34:2823-288 22344, 2021.
[35] P-A Absil and Jérôme Malick. Projection-like retractions matrix manifolds. SIAM Journal on Optimization, 22(1):135-158, 2012.
[36] 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.
[37] Damek Davis, Dmitriy Drusvyatskiy, and Zhan Shi. Stochastic optimization over proximally smooth sets. arXiv preprint arXiv:2002.06339, 2020.
[38] MV Balashov. Nonconvex optimization. Control theory (additional chapters): tutorial. Moscow: Lenand, 2019.
[39] MV Balashov and AA Tremba. Error bound conditions and convergence of optimization methods on smooth and proximally smooth manifolds. Optimization, 71(3)711-775, 2022.
[40] Pan Zhou, Xiao-Tong Yuan, and Jiashi Feng. Faster first-order methods for stochastic non- convex optimization on Riemannian manifolds. In The 22nd International Conference on Artificial Intelligence and Statisties, pages 138 147. PMLR, 2019.
[41] Sashank J Reddi, Suvrit Sra, Barnabas Poczos, and Alexander Smola. Proximal stochastic methods for nonsmooth nonconvex finite-sum optimization. Advances in Neural Information Processing Systems, 29, 2016.
[42] Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018.
[43] Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In P Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 5693-5700, 2019.
[44] Alan Edelman, Tomás A Arias, and Steven T Smith. The geometry of algorithms with or- thogonality constraints. SIAM journal on Matrix Analysis and Applications, 20(2):303-353,
1998.
[45] Nicolas Boumal, Bamdev Mishra, P-A Absil, and Rodolphe Sepulchre. Manopt, Matlab toolbox for optimization on manifolds. The Journal of Machine Learning Research, 15(1):1455- 1459, 2014.
[46] James Townsend, Niklas Koep, and Sebastian Weichwald. Pymanopt: A python toolbox for optimization on manifolds using automatic differentiation. Journal of Machine Learning Research, 17(137):1-5, 2016.
[47] Robert L Foote. Regularity of the distance function. Proceedings of the American Mathematical Society, 92(1):153-155, 1984.
[48] Maxence Noble, Aurélien Bellet, and Aymeric Dieuleveut. Differentially private federated learning on heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pages 101 10110-10145, 2022.
13

---

A Appendix
A.1 Notations
We analyze the proposed algorithm using the Lyapunov function !2" defined by
N [(PM(>")) A' 131
(14)
where f* is the optimal value of problem
and we define
A ((rgrad((PM(x))
gradf 2 1,
Σ grad ;B
1,,,-11)
and A col ( E2= N
The Lyapunov function consists of two parts: to bound the suboptimality of the global model P (x*)
and the reduction of "variance" among clients, respectively.
A.2 Preliminary lemmas
Let us start with the following lemma on the global- like Lipschitz- Lipschitz-continiity continuity property of P.M Lemma A.1. There exists constant M O such that for any M..
(a u) =|| S M [|]]
(15)
Proof. Let us consider two cases:
|w|| 2 y: Since PM(x u) and x belong to M, we have
[| xell diam(M) diam(M) ?
||]]
where diam(M) max max yll is the diameter of MM.
|ull S 7: By the 2-L.ipschitz continuity of PM over U (m) in we have
[P u x|1 < 2/|u||-
Setting M max diam(M) 3). we complete the proof.
In the following, we show the reasonableness of the suboptimality metric [[G Lemma A.2. Consider Gif(-) defined by (10). Then, for any E M, it holds that
gradf(#) if and only if Gz(x)) 0.
In addition, under Assumptions 2.3 and 4.1 if S min { 3. L. }
with Lp being the smoothness
constant of Df over T M(7), it holds that
Iggrdd < (x))1
(16)
Proof. If grrrd((()) 0, it follows directly from the definition of GT(+) that GT(x)) Conversely, Gq(x)= = 0. we have
PM (ar - nggrad/ (x)) argmin |y + ngrrdd/ (x))11
NEM
It follow from the optimality of that 0 (r..a((ggrdf(())) which implies that gradf(z) 0.
14

---

Lemma], PM(()) is sufficiently smooth over (M()). Let us define Lp == max M())(() MD (z)]l, where D2 denotes the second-order differential operator. Then, have
With 47
[9g(,)] = و 11 PMu( - ngradd
Iv 2V Igradf(<x)] L L
(x))l,
1 in the second inequality. This gives where we use f) S D, in the first inequality and ij S Ln (6).
To prove Theorem 4.3 we use the following lemma to establish a recursion on the second term A' in the Lyapunov function.
Lemma A.3. Under Assumptions 2.3
4.1 and 4.2 if n S min
Yo 20 T6L
}. we have
2D,
1 EI - - (M
A
PM(
(17)
n
1 T TL2 ( + - A 1 11 18 1 2)220) + n
1
n a²
n
nTb
Proof. As first step, we bound the drift error || P (x") that is caused by the local updates. I = 1, the error is zero since zi, = PM(x"") When 2 2, repeating the local updates for steps and substituting 00,0 = (m and -,,11 (M we have
E T..1 = n
(gradfi(2),;; Bi,)) G)) - -PM(=') (18)
To bound the right-hand side of (18., compare algorithm with the exact C-PRGD step given in under the step size (t + 1)n
CPGG PM (PM(x*) (1 (P
It follows from
15 that
C PM((")ll < Mrnlgrad/PM
Then from 18 we have
B B| - PM(x')1
(19)
= =PIPMMPMMeT -
+c)) -xetInD + #ctrneD - PM(x")]
< 2F([PM(PM(T))
grradf((z)) B + c)) htt CRRDD 12M"
E=0
where we use ||a + b112 2||a|12 2| in the inequality.
To bound the term (1) on the right hand of 19), from ij S 7Hh and maxi,I..M [[V]a< Da)ll Dy.
2Dy
we have
7
(g
+ ST.
15

---

Thus, by substituting definition of gr+1
"C-PRGD* we can invoke the 2-Lipschitz continuity of PM over
U (M)() given in B and get
(I) 21 PM PM((P((()) Σ (gradfi(*[ ct))
(PM(P) (L 1)ngrad (PM
(20)
41E n Σ gradff Bi,)) + + - gradf (P
Next, bound the right-hand side of (20) we rewrite in terms of ||A A A' by substituting the definition of the of given in (5)
(1) - 7
grad (E) - (Padf
(21)
+ gradf:(PM((1)
grad f: ;
t=
gradfi : gPad |
(A ^²) )| ||
4F 17 grad (1 grads:(P
ggra fe(2 B gPad +8E & ++ A7
1 A
$ 8E m
Next, for the term (II), substituting Assumption 4.2 yields
(II) = 8E
gradft Bi,)) gradfi + gradfi(2 - gradf((PM(z")) (22)
8=0
< 16( + 1E
gradff gradf((2[.)))
16E 17
- (rad
For the first term can be handled by the fact
48 Corollary C.1] that
fgradf ( grad |
16(t - 1)
(23)
+
16(t + +
1 o²
E || (gradfi (=T,e: Bi,,) gradSi (zf,e) [1937]
<
1)2
t+1 b
0=
Combining
22
and
we have
+ 1
20² 0²
(1) <16(t + 1)m² 2
Bl PM( )12 +8
EIAA X'1 +
16(
1)m7
(24)
b
E=0
where we use Assumption
Next we substitute (24 into 19 to get
E
21,41 (=)(2"))11
(25)
16

---

S 16(t + 1)n
B1 Pa)(2')1?
+
8E| - X' 1677)
==A-
The following proof is similar to that in 24 We define A" the sum of the last three terms on the right hand of (25) and ST = E
Ell (M By 21ll PMM("")
St - S and 25 we have
, S (1+ 1/67)) )) S + A",
(26)
where the inequality from 7 S /// (16L) and thus 16(l + 1)r L2 N 1/(16r). With (26). we get
- SA'
Σ (1+ 1/(66))) & 1.17A"
(27)
where use Ee (1+ 1/()))) < Ee exp (1/( s Exd cxp(1/16 I 1.1r. Sum- ming (27) over all the clients i, we get
| - PM(x')]²]
(28)
< 3nnM2 7/
+ A' 1' + 18n7 1) b
E 71 - A 11a. 112. By the definition of A r+ and A
Now we are ready to bound
we have
n
E 11 ||2
(29)
= / rgradf (PM(x*+1)
gradf (x); (PM +
gradf (zf: BB) |?
1n B rgradf (PM((**)) gradf (x]; ||
= B rgradf (PM(x*+1)) rgradf (PM(x")) + rgradf (Pm((x))
Σ
gradf (z]:B) 2
gradf (z,)
gradf (zf)
1=0
1=0
o²
< 22 L1 aaaae||PPa PM PM(x'||2
+4 4n2² E B| PM
b
1=1 t=0
Here, the first inequality is due to ||
A
1|2 < || the last inequality is due to
||a + b112 < 2||a||2 2||b| Assumption 4.1
and following similar derivations as (3).
By substituting (28 into 229 and reorganizing the results, we complete the proof of Lemma A.3
A.3 Proof of Theorem
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.
Lemma A.4. Given M, v E T,M, n 0, x nuu E U M())), and x PM(2 -no), it holds that
f(xt) < f2) + (grad
- u. - z)
(11o x=12 - || - xm|²) (1 1
3||||||
||2 - x
2n
2n 47
(30)
1
11
z=|²
+
L
|| z 211,Vz V2 M.
2
2
17

---

Proof. For any p-strongly convex function h, we have for any 2 E M
h(=) 2 hly) + (Vh(y), - u) + -1= - yl1
= h(y) + (gradh(y) + Vh(y) - gradh(y), 2 - y/ + 1 2112- y11
2
(31)
2 hyy) + (gradh(y), v)
|| yll?,
2
47
where the second inequality is from the normal inequality 4 and Vh(y) - gy)d < [h(y)]l- Setting h(y) = 2n
1 fly (x - muull in (31 with μ = 1/n, = x², and noting the optimality of x
(i.e., gxadh(x*) = 0). we have
1 ||2 (x mu)ll 2
1
( 2n
11 (xc (x- - ||| xt1
11 (x -m0-n0)112
+
21
2m
107
1
3|1ull
47 ) || - x+11 x=11²
1x=
(xr n1=)))
2m
27
where the second inequality is from x - un E U. (m) and IP - mul) - (xr - -nnill S 1PM
vyo - x|l n||le|| s 3y/lu|. Rearranging the above inequality leads
1 (1|=* -11--1z 1|2 xx1 +
1
3|10|| -|2- - 1t
(1,z-x+)> N
(32)
21
2n
47
It follows from the L-smoothness of f that
S 5()) + (gradf(z), x+ -x) +
| x1|
2
L
L
< {()) + (gradd(())
x+ - 2) +
= x]|²
+
|| - ==1??
2
2
where we use f(x) + (grad ((r (x). x)
4 b|1= 2 z1|2 < f() in the last inequality. Combining the
above inequality and 32) gives 30
In the following, we use Lemma A.4 to and respectively. First, to apply Lemma A.4 to substitute x = +r+1, PM ), x PM (x') ) and g (x")) and get
ff ( ) s [ ( + (
L 1 | PM(x')) 1- hp
|r
PM
(x {1²]
27
27
(33)
where we use 7 S D. to guarantee ar+1 E EUM(7)) and p: 3Dg,
27
Next, to use Lemma A.4 6 we set a+ = PM(x ) x = PM(x") == a v = v and get
E [[(PM(z*)))]
(34)
<
E
[( + (grad[(PM(=")) - o², PP -1+11
1 (IPm)())t)) - (x1 -|art1 PMn(=")133) -
1 1- ||rt1 t1)
27
2n
I
I ! 11 - (M)
+ ( PM(("))1
E f (art1) + (grads - PPM(x*+) - - ( L. 11) (Partt))PMmm))
-
L +
1
|| PMM('") |1² 21= -np -
+
27
27j
where we use ;
<
77 to guarantee (l) 2D,
E M.?)).
18

---

Recall that
) (Pu(tE) -
1 TCT
n
fr (2f,e: ((t)
1mp"
Combining the above inequality,
33 and 34 yields
E [[(PM(Z*)))]
(35)
+ P
1 -MM(=))1 (1-))++(()))
21j 1 [P - (x)(1))
s E [ [f (PM(2*)) +
- 27j
2
1 Ap
|Pn((=*)))
xr+ 11 ( gradf (Pm(x"))-e"))
2ñ
(IV)
According to l1a b112 2 1||13 + 11611 have
(IV) + (V)
(36)
$ (TV) + 1 7- 7p
- +
(P 112
"
(Prad((PM<z") u=1
2îj
2(1 Ap)
(gMa - 11
2(1 7i))
To bound the above inequality, following similar derivations as in
23 we obtain
E |
(P'"))
gradfi (zf,6i Bl,,) gradfi(=[,)) + gr)d)( gradf((Px(z")))
IIT
(37)
E| PM(x")| )|2 + 2 o²
< 2L
tn b
VLT
Substituting (36) and 337 into 35 we have
(( ))]
(38)
< (Pm(x*))
T.
1
())
-PM(=))))
(
1
-PM))
(x"))
21j
27
212 n
--1
tn )
2 220²)
+
| - PM
2(1 011-1p) 7p)
nt
The final term the drift-error that be bounded in 28 Thus, 38
becomes
E[((PM(**t)]]
(39)
E f (PM(x")) (L- (-
1 +b) || -PM(3')] +
(5
L 1 [P PM(("")1
2nj
27
212
12 (PM( - A'I22 1817 2 20
2 0²
112
2(1 πpp)
nt
2(1 1p))
Tm
b
where we use (ggrad/(PMtan i S v we have 2(1 (p)))pp S nj Combining the recursions given by Lemma A.3 and (39]. we have for the Lyapunov function that
< 2| from Lemma A.2 3. By substituting 7 S
1 as 4p
IE (F(PM(=*+))) f*) +
- - A I??
(40)
19

---

8ij o² " (P (1))]
<[(U(M((()))
A' A A
nr b
where substitute conditions Substituting the definition of the yapunov function Q and repeating the above inequality, we complete the proof of Theorem 4.3
the step sizes and omit straightforward algebraic calculations.
Discussion on step sizes. our analysis, the condition f/g - vn applied only in the final step, when combining Lemma A.3 A..3andd and (39) to derive the recursion on the Lyapunov function in (40). Up to this point, we retain f2 as a variable without assigning a specific value, ensuring that preceding results remain valid for 7l 1 well. We choose 7g vni in (40) to conveniently cancel out the number clients n, yielding a more concise expression in the final result.
TT I min 24ME 6D,; DyLe }. The term 6Dj
controls the client drift, (39); the term D,LD used to derive (T6). which ensures that when the metric
In we require that
(M,(()(()))) approaches zero, the first-order optimality condition is met. The term 544L is used to establish the Lyapunov function recursion (40) by combining the recursions in (17) and (39). Additional conditions the step sizes ensure that the iterates remain within the 2-ytube during
analysis, such as in 33 and (34) but they are implied by the three terms in (TT).
A.4
Additional results for numerical experiments
A.4.1 kPCA
The settings for kPCA problem with Mnist Dataset. The Mnist dataset consists of 60,000 handwritten digit images ranging from to 9, each with dimensions 28 x 28. We reshape these images into data matrix A 8,0000xx88 784 To construct the heterogeneous A,, 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") is shown in Figs. 5
Our algorithm performs much better than alternative methods in terms of communication quantity and run time.
10a
10a
10a
10-1_
10-1
10-1
9
10
200
400
600
800
1000
500
1000
1500
2000
20
40
60
100
Communication round
Communication quantity
Runtime (s)
RFedavg
RFedprox
RFedSVRG
Figure 5: kPCA problem with Mnist dataset: Comparison on f(x") - f
Synthetic Dataset. We also solve kPCA with synthetic datasets on larger networks with n = 30. We generate each entry of A, from Gaussian distribution / (0, such that A; are heterogeneous among clients. We set (d, k) (20, 5) and p 15. We use the local full gradient V to remove the influence stochastic Riemannian gradient noise.
In the first set of experiments, we compare with existing algorithms, including RFedavg. RFedprox, and RRG. For all algorithms, we set the number of local steps as and the step size as 4e-3. For our algorithm, we set ny 1. The experimental results are shown in Figs. 6 The y-axis represents Iggrad and (f(x") ) f** respectively, while the r-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 not converge accurately. Both
FddSVRG and our algorithm can overcome the client drift issue, but algorithm is slightly faster
20

---

in terms of communication rounds and is much faster in terms of both communication quantity and run time.
10a
10a
10a
Igrad fix')
1
Igrra f(x')]
10-3
Igrad ffx'
10
10-7
10--
10
10
10-11
10 11
10-15
10-15
10-15
1000
2000
3000
4000
2000
4000
6000
8000
10
20
30
Communication round
Communication quantity
Runtime (5)
RFedavg
RFedprox
RFedSVRG
our
109
10a
10a
L.
L
10
10-12
10-12
10-12
10-16
10-16
10-16
2000
1000
2000
3000
4000
Communication round
Communication quantity
Runtime (s)
RFedaag
Rfedprox
RFedSVRG
our
Figure 6: kPCA with synthetic dataset:
Comparison on ||grad (x")) and f(x") f
In the second set of experiments, we test the impact of the number of local updates т. For all the algorithms, set R = 4000, the step size n 0.7e-3, 3, and {10, 15, 20). For our algorithm, set f/5 = 1. The results are shown Figs. with the y-axis representing ggadd fx=)) and x-axis representing the communication quantity. When increases, the convergence becomes faster. For all values of r, our algorithm achieves high accuracy and requires less time.
10a
T 10
104
T 15
10a
T 20
Igrad ||
10
Igrad fix"""") '|l
10
Lgrad '11
10
10-7
10-7
10-1
10-15
10-11
10-12
10-15
10-15
10-15
2000
6000
8000
2000
4000
6000
8000
2000 Communication quantity
4000
6000
Communication quantity
Communication quantity
RFedavg
RFedorox
RFed5VRG
our
Figure 7: kPCA with synthetic dataset: The impacts of T.
A.4.2 Low-rank matrix completion
For numerical tests, consider random generated A. To be specific, first generate two random matrices L e pa and Å k where each entry obeys the standard Gaussian distribution. For the indices set Q, we generate a random matrix B with each entry following from the uniform distribution, then set Ogj if Bij and 0 otherwise. The parameter is set 10k(d T kd)(dT).
As shown in Figs. 8 our algorithm is faster than existing algorithms in terms of communication quantity and run time.
We also show the impact of and the impact of n. As shown in Figs. larger yields less communication quantity achieve the same accuracy. As shown in Figs. for n 10, n 20, and = 30, our algorithm consistently outperforms the compared algorithms.
21

---

10a
10a
10
10
10
10
10
10-5
10
10 8
10
10
100
200
300
400
200
400
600
800
100
150
200
Communication round
Communication quantity
Runtime (s)
RFeddag
RFedprox
RFedSVRG
our
Figure 8: LRMC: Comparison on f(
109
109
109
grad [Ix]
Igrad fIx
grad f(x]
10-1
10-1
10-1
10-6
10-4
10
10
10-12
10-12
200
400
600
200
400
600
200
400
600
Communication quantity
Communication quantity
Communication quantity
RFedavg
RFedorox
RFedSVRG
Figure 9: LRMC: The impacts of r.
[gradf(x(1
109
Igrad [())
109
Igrad Rxx))
10"
10
10
10
10-4
10
10
100
200
300
400
200
400
100
150
Communication round
Communication quantity
Runtime (s)
RFedavg .....
RFedprox
RFedsVRG
our
Iggaad(x" fix')l
109
\|I
Igrad fix"")) I||
10a
Iggad Fixe")
10-2
10-2
10-2
10-4
10-4
300
400
800
150
200
Communication round
Communication quantity
Runtime (s)
RFedavg
RFedorox
RFedSVRG
our
Igradf(x]]
10a
Igrad [(x]]
10a
Igrad ftx]]
10-2
10-2
10-2
10-4
10-4
10-4
200 300
400
200
400
600
800
200
Communication round
Communication quantity
Runtime (s)
RFedavg
RFedprax
Figure 10: LRMC: The impacts of n. (First row: 10, second row: 20, and third row:
n 30. We set T = 5.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
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.
Justification:
Guidelines:
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 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 is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer:[Yes)
Justification:
See Conclusions and limitations.
Guidelines:
The answer NA means that the paper has 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 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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
23

---

Answer: [Yes]
Justification: See Assumptions we made and Appendix A.1-A.3 for a complete (and correct) proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See Numerical experiments and Appendix A.4 for details of experiment settings and additional results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
24

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipssccc/pill wttp:ccppppcccceeeeeeeeee..c.ett httpa////gueeee/edeSSuuiiiissss Coooeooomisssssseeeeeeee www.iissssiiiicciico for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Our algorithm uses stochastic gradients. We fixed the random seed to better compare the impact 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, confi- dence intervals, statistical significance tests, 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 with given experimental conditions).
25

---

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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: Our experiments can be completed on laptop.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our research conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
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.
26

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 open datasets.
Guidelines:
The answer NA means that the paper does 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
27

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
28


Reviews:
[{"review_id": "3yesugsedt", "reviewer": "Reviewer_q9zL", "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.", "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?\n- 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 and open problems are adequately discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written, with its novelties and contributions clearly articulated. The development flows naturally, providing sufficient intuitions and implications throughout.\n- The paper identifies the computational challenges in federated learning on manifolds, specifically involving exponential mappings, logarithmic mappings, and parallel transports, and effectively addresses these challenges using a simple projection mapping.\n- The idea behind the methodology is both intuitive and effective.", "weaknesses": "I am overall satisfied with the originality and quality of the paper, with some suggestions:\n\n- Nonconvexity\n    - 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?\n    - 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.\n    - 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.\n    - 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.\n- Numerical experiments\n    - Most experiments use full gradients instead of stochastic gradients.\n    - In the kPCA experiment, the number of clients is $1 = n = \\eta _{g}^{2}$ according to Theorem 4.3. Could the authors clarify this?\n- Notation and typos\n    - I would suggest moving some necessary notation from Appendix A.1 to the main text to make the analysis more self-contained.\n    - Is $D$ in the definition of $L_{\\mathcal{P}}$ in Theorem 4.3 missing a subscript $f$?\n    - The third line of Paragraph *Theoretical contributions* is missing a closing parenthesis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and positive assessment of our work!\n\nSuggestions\n\nNonconvexity\n1)  All the concepts related to convexity and nonconvexity are defined in the ambient space, which is the Euclidean space for the considered submanifolds. We will state this clearly in the revised manuscript. \n2)  Yes, the nonconvexity involves both the nonconvexity of the manifold constraint and the nonconvexity of the objective function.  When the manifold constraint is incorporated into the objective function as an indicator function, the resulting objective function becomes composite, with both the smooth and nonsmooth components being nonconvex. We will state this clearly in the final version of the paper.\n3) It is due to the nonconvexity of the manifold constraint. \nThe \"existing composite FL\" mentioned here refers to the related work discussed in the previous paragraph, specifically Composite FL in Euclidean space ([19]-[23]). When we refer to composite FL in Euclidean space ([19]-[23]), the composite objective function consists of two parts: a smooth part and a nonsmooth part. All of these works ([19]-[23]) require the nonsmooth part to be convex. For the smooth objective function, except for [21] and the follow-up [22], which allow the smooth objective function to be nonconvex, the others require the smooth objective function to be convex or strongly convex. When we incorporate the manifold constraint as an indicator function into the objective function, it means that the nonsmooth function in the composite objective also becomes nonconvex. This results in the methods in [19]-[23] not being directly applicable. We will clarify in the revised version what specific nonconvexity we are referring to.\n4) Thank you very much for your suggestion. We should indeed mention that the related work [13] can also handle geodesically nonconvex objective functions. Our contribution concerning the nonconvexity of the smooth objective function is to compare against existing composite FL work in Euclidean space. Among these, apart from [21] and its follow-up work [22], the other works [19], [20], and [23] all require the smooth objective function in the composite objectives to be convex or strongly convex.\n\nNumerical experiments \n1) Thank you for your comments. \nCompared to other related algorithms, our algorithm demonstrates significant advantages in overcoming client drift, reducing communication quantity, and saving computation time. Therefore, in the comparative experiments, we have chosen to focus on these aspects. We employ full gradients to conduct ablation studies to eliminate the interference of stochastic gradients, thereby better demonstrating our algorithm's ability to handle client drift, while also showing advantages in terms of communication quantity and computation time. The impact of stochastic gradients is predictable and aligns with classical conclusions from single-machine SGD algorithms. As shown in Figure 3, the error caused by stochastic gradients can be reduced by increasing the mini-batch size.\n2)  Thank you very much for allowing us the opportunity to clarify this point. \nTo facilitate comparison with other algorithms, we fixed $\\eta_g=1$ and then manually adjusted the step size $\\eta$ to its optimal value. In our analysis, the condition $\\eta_g = \\sqrt{n}$ is used only in the last step of our analysis, when combining Lemma A.3 and (39) to derive the recursion on the Lyapunov function in (40). For all analysis results prior to inequality (40), we retained $\\eta_g$ without substituting a specific value, and they still hold for $\\eta_g=1$. We set $\\eta_g = \\sqrt{n}$ when deriving (40) because this setting conveniently cancels out the number of clients $n$, leading to a more concise result in (40). We acknowledge that the upper bound in (40) obtained with $\\eta_g = \\sqrt{n}$ may not be optimal. Retaining $\\eta_g$ allows the server and clients to use different step sizes, increasing flexibility, which aligns with existing well-known FL literature, such as [18], [19], and [23].\n\nNotation and typos\n1) Agree.\n2) In the notation $D^2 \\mathcal{P}_{\\mathcal{M}}(x) $, $D^2$ represents the second-order derivative. We will introduce this notation more carefully in the revised manuscript to avoid confusion.\n3) Agree.\n\nQuestions\n1) The ambient Euclidean structure of the considered submanifold is crucial in the analysis of proximal smoothness and the proposed method. Note that such manifolds cover many of the most widely used manifolds, including the Stiefel manifold, the oblique manifold, and the symplectic manifold.\n2) For a general manifold where the ambient geometry is not Euclidean, it will be difficult to define the\nprojection operator and the associated proximal smoothness. This will complicate the design and\nanalysis of the algorithm. We will comment on this in\nthe Conclusions and Limitations of the revised manuscript.\n3) Thank you very much for your comments. \nIn Eq.(11), we require $\\tilde{\\eta}:=\\eta_g \\eta\\tau\\le \\min \\\\{ {\\frac{1}{24ML}}, \\frac{\\gamma}{6D_f}, {\\frac{1}{D_f L_{\\mathcal{P}}}} \\\\}$. The term $\\frac{\\gamma}{6D_f}$ is to control the client drift, see (39); the third term $\\frac{1}{D_f L_{\\mathcal{P}}}$ is to derive (16), which indicates that the metric $\\\\| \\mathcal{G}_{\\tilde{\\eta}}( \\cdot ) \\\\| $ to zero implies that the first-order optimality condition is met; the first term $\\frac{1}{24ML}$ is used to establish the Lyapunov function recursion (40) by combing the recursions of (17) and (39). There are some other conditions on the step sizes for restricting the iterates to remain within the $2\\gamma$-tube during our analysis, such as in (17), (33), and (34), but they are implied by these three terms in Eq.(11).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WpjCqGZBFs", "reviewer": "Reviewer_qvSZ", "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.", "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": "The author(s) have discussed the limitations of the proposed algorithm in Section 6.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Optimization on manifolds in the federated learning setting is a challenging topic, and this article seems to make visible contributions on this direction. The theoretical analysis also appears to be novel.", "weaknesses": "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our sincere gratitude for your careful review and positive assessment of our work!\n\nWeaknesses\n\n1) The projection operator can be explicitly calculated for many common submanifolds, as discussed in [SIOPT 2012]. For the Stiefel manifold, the closed-form expression for the projection operator of a given matrix $X$ is $\\mathcal{P}_{\\mathcal{M}}(X) = X(X^\\top X)^{-1/2}$; see Proposition 7 in [SIOPT 2012]. \n\n[SIOPT 2012] P-A. Absil, and Jerome Malick. ``Projection-like retractions on matrix manifolds.\" SIAM Journal on Optimization 22, no. 1 (2012): 135-158.\n\nQuestions\n\n1) Thank you very much for your careful review. Although the kPCA problem is nonconvex, it has certain properties\nsimilar to those of a strongly convex function, e.g., the Lojasiewicz property [MP 2019]. Similar\nphenomena have also been observed in manifold optimization on low-rank matrices [NeurIPS 2021,\nAAAI 2022]. While challenging, it would be interesting to consider local regularity conditions such\nas the Lojasiewicz property and establish convergence rates faster than sublinear.\n\n[MP 2019] Huikang Liu, Anthony Man-Cho So, and Weijie Wu. “Quadratic optimization with\northogonality constraint: explicit Lojasiewicz exponent and linear convergence of retraction-based\nline-search and stochastic variance-reduced gradient methods.” Mathematical Programming 178,\nno. 1 (2019): 215-262.\n\n[Neurips 2021] Ye T, Du S S. Global convergence of gradient descent for asymmetric low-rank matrix\nfactorization[J]. Advances in Neural Information Processing Systems, 2021, 34: 1429-1439.\n\n[AAAI 2022] Bi Y, Zhang H, Lavaei J. Local and global linear convergence of general low-rank\nmatrix recovery problems[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2022,\n36(9): 10129-10137.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F7dIA62yGo", "reviewer": "Reviewer_GhoW", "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.", "questions": "-  Please clarify the theoretical results in terms of assumptions on the manifold needed.\n-  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.\n-  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.\n\n[13] \"Jiaxiang Li and Shiqian Ma. Federated learning on Riemannian manifolds. arXiv preprint arXiv:2206.05668, 2022.\"", "limitations": "Limitations are properly addressed.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The computational cost is advantageous, which could make new problems feasible.\n- The federated learning aspects seem excellent.", "weaknesses": "- 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.\n- 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.\n\n- 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.\n\n[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\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your careful review and valuable comments!\n\nWeaknesses\n\n1)  Thank you for your comment and the opportunity to clarify this point. In our algorithm design, we not only utilize projection but also exploit the manifold geometry, specifically the Riemannian gradient, to create a more efficient algorithm. Theoretically,  we use the manifold geometry to define the optimality metric $\\mathcal{G}_{\\tilde{\\eta}}(x)$ and establish a quasi-isometric property, i.e.,  $1/2\\\\|\\operatorname{grad} f(x)\\\\|\\leq\\\\|\\mathcal{G}{\\tiny \\tilde{\\eta}} ( x)\\\\|\\leq 2\\\\|\\operatorname{grad} f( x)\\\\|$   (see Line 259, Lemmas A.1 and A.2), which is crucial to establishing the convergence of our algorithm. It is challenging to verify this inequality using the Euclidean gradient, which ignores the manifold geometry, leading to the possibility that an algorithm that employs the Euclidean gradient instead of the Riemannian gradient may fail to converge. Therefore, both algorithm analysis and design rely on the effective utilization of manifold geometry.\n\n2) All results are derived and proven for Euclidean immersions and the Euclidean metric. In the\nsubmitted version of the manuscript, we stated this on Lines 132 and 135: “Throughout the paper,\nwe restrict our discussion to embedded submanifolds of the Euclidean space, where the associated\ntopology coincides with the subspace topology of the Euclidean space. · · · In addition, we always\ntake the Euclidean metric as the Riemannian metric.” We will revise the manuscript to ensure that\nthere are no misunderstanding that our results are restricted to this case. \n\n3) The proof of Theorem 4.8 in [33] relies on the uniqueness of the projection and the application of the Cauchy-Schwarz inequality.  As we mentioned in Lines 161-162 in the submitted version of the manuscript, \"we introduce the concept of proximal smoothness that refers to a property of a closed set, including $\\mathcal{M}$, where the projection becomes a singleton when the point is sufficiently close to the set.\" Because the projection is unique, Theorem 4.8 in [33] still holds in our case.  However,  Theorem 4.8 in [33] does not necessarily extend to general Riemannian manifolds. We will emphasize this point in the revised version.\n\nQuestions\n\n1) We assume that the manifold is a compact and smooth submanifold embedded in $\\mathbb{R}^{d \\times k}$, as we mentioned below Eq.(1). We will state the assumption clearly in our theorems in the revised manuscript.\n\n2) These are good questions. We will address them one by one.\n\ni) The proximal (or projection) operator can be calculated explicitly for many common submanifolds, see [SIOPT 2012].   For the Stiefel manifold, for example, the closed-form expression for the projection operator of a matrix $X$ is $\\mathcal{P}_{\\mathcal{M}}(X) = X(X^\\top X)^{-1/2}$; see Proposition 7 in [SIOPT 2012]. \n\n[SIOPT 2012] P-A. Absil, and Jerome Malick. \"Projection-like retractions on matrix manifolds.\" SIAM Journal on Optimization 22, no. 1 (2012): 135-158.\n\nii) Under Assumption 2.3 and the step size condition (11), we guarantee that the projection used in our algorithm is a singleton. In our algorithm, we calculate \n$\\mathcal{P}\\_\\mathcal{M}(\\hat{z}\\_{i, t+1}^r)$ \nand $\\mathcal{P}\\_\\mathcal{M}(x^r)$ which both lie on the manifold. Neither $\\hat{z}\\_{i, t+1}^r$ nor $x^r$ are necessary on the manifold, but we require that $\\hat{z}\\_{i, t+1}^r$ and $x^r$ are sufficiently close to the manifold, specifically within a tube $\\overline{U}\\_\\mathcal{M}(\\gamma):=\\\\{x: \\operatorname{dist}(x, \\mathcal{M}) \\le \\gamma \\\\}$. We ensure that this condition is met by restricting the step size to be small enough in our analysis. In our algorithm, $\\hat{z}\\_{i,\\tau}^r$, which is the last local model during the local updates, is uploaded, and $x^r$ is downloaded. The variable $\\hat{z}\\_{i,\\tau}^r$ is essential for the server to extract the average local gradient information by simply averaging all $\\hat{z}\\_{i,\\tau}^r$ (since the nonlinearity of $\\mathcal{P}\\_\\mathcal{M}(\\cdot)$ prevents this from being done directly through $\\mathcal{P}\\_\\mathcal{M}(\\hat{z}_{i,\\tau}^r)$). Meanwhile, $x^r$ is used locally by each client to construct the correction term $c_i^r$ without requiring additional communication.\n\n3) The work [13] appears to be the only FL algorithm that can deal with manifold optimization problems of a similar generality as ours. \nThe manifold we consider is a compact, smooth submanifold embedded in $\\mathbb{R}^{d \\times k}$, which is more restrictive than the manifolds discussed in [13] but it covers many common manifolds such as the Stiefel manifold, oblique manifold, and symplectic manifold. \nWe will make this point noted clearly in our revised manuscript. \n\nIn [13], the inverse of the exponential map is used to pull back the local updates to the same tangent space, allowing for the computation of the tangent mean.  However, the convergence in [13] is shown only for two specific settings: (1) $\\tau=1$, and (2) $\\tau>1$ but with only a single client participating in each training round. By leveraging the structure of the compact smooth submanifold, we demonstrate convergence for $\\tau \\geq 1$ with full client participation.\n\nWe genuinely appreciate your thoughtful and constructive feedback. Each of your comments has been carefully considered, resulting in significant improvements and a stronger overall paper. We believe that our responses have thoroughly addressed all the concerns raised. We would be most grateful if you could kindly consider raising the score!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qjYyT7GTTL", "reviewer": "Reviewer_3dV9", "summary": "This paper studies federated learning on a compact smooth submanifold using distributed Riemannian optimization methods.\n\nThe proposed methods were examined on PCA and low-rank matrix completion tasks.", "questions": "There are several other works on distributed Riemannian optimization methods for similar tasks, such as\nDistributed Principal Component Analysis with Limited Communication, Neurips 2021.\nCommunication-Efficient Distributed PCA by Riemannian Optimization, ICML 2020.\n\nCould you please compare your proposed methods with these related work theoretically and/or experimentally?\n\nThe 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?\n\nHow did you measure CPU time?\n\nCould you please provide comparative analyses with Euclidean baselines, such as FedAvg and FedProx?\n\nIn 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?\n\nOne 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 were partially addressed.", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed proximal smoothness of a manifold is intuitive and interesting. The notation can be improved but the theoretical analyses of the algorithm are clear in general.", "weaknesses": "In the theoretical analyses, a comparison with the other related Euclidean and Riemannian methods can be given.\n\nExperimental analyses are limited and should be improved as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your careful review and valuable comments. \n\nWeaknesses\n\nWe have addressed all the comments accordingly. Please see the response to Questions for details. \n\nQuestions\n\n1) Thank you very much for your comments. We have provided a comparison with [ICML 2020] and [Neurips 2021] in the general response. \n\n2) Thank you very much for your suggestion.  In our kPCA experiments, we presented results for the MNIST dataset with $n=10$ clients in the main text and for synthetic datasets with $n=30$ clients in the appendix. The results for kPCA on MNIST with $n=30$ clients, reported in Figs.6-7 in the Appendix, are qualitatively similar to those observed for $n=10$, demonstrating that our algorithm effectively reduces communication quantity and runtime. In the MNIST experiment, when we set $n=10$, data from each of the 10 classes was distributed across 10 clients, ensuring that no two clients shared data with the same label. This setup introduced significant data heterogeneity, allowing us to better demonstrate our algorithm's ability to overcome client drift issues caused by such a high level of heterogeneity.\n\nBased on your feedback, we have conducted additional experiments on low-rank matrix completion for $n=10, 20, 30$. In all cases, our algorithm consistently outperforms the compared algorithms. The results are in the uploaded PDF in the general response. \n\n3) Thank you for raising this issue. We should indeed use \"runtime\" instead of  \"CPU time.\" Our code was implemented in Python, and to measure runtime we recorded timestamps at the start and end of each communication round using time.time( ). The difference between these two timestamps represents the runtime for one communication round. \n\n4) Thank you very much for your suggestion. We have provided the comparison with FedAvg and FedProx.  Please see the general response. \n\n5) This is an excellent question. Our algorithm can indeed be extended to achieve differential privacy. Below, we outline how this can be accomplished.\n\nAlgorithmically, each client can add Gaussian noise to the local gradient information to ensure differential privacy protection. Let $\\xi \\in \\mathbb{R}^{d\\times k}$ be Gaussian noise, where each entry is independently and identically distributed, following $\\mathcal{N}(0, \\sigma^2)$. We modify Algorithm 1 at Line 8 and Line 17 as follows: \n$$\n\\begin{aligned}\n\\widehat{z}_{i,{t+1}}^r&=\\widehat{z}\\_{i,t}^r   -\\eta (\\widetilde{\\operatorname{grad}}f_i (z\\_{i,t}^r;\\mathcal{B}\\_{i,t}^r)+c_i^r )\\\\\\\\\nc_i^{r+1}&=\\frac{1}{\\eta_g \\eta\\tau} (\\mathcal{P}\\_\\mathcal{M} (x^r)-x^{r+1})-\\frac{1}{\\tau} \\sum\\_{t=0}^{\\tau-1} \\widetilde{\\operatorname{grad}} f_i(z\\_{i,t}^r;\\mathcal{B}\\_{i,t}^r)\n\\end{aligned}\n$$\nwhere $\\widetilde{\\operatorname{grad}} f_i(z\\_{i,t}^r;\\mathcal{B}\\_{i,t}^r):=\\operatorname{grad} f_i(z\\_{i,t}^r;\\mathcal{B}\\_{i,t}^r) +\\xi $. \n\nFor the analysis, we can establish a quantitative relationship between privacy loss and the variance $ \\sigma^2$ as follows:  \n\ni) First, we determine the privacy loss for a single local update. According to the post-processing theorem [Lemma 1.8, TCC 2016] and the definition of $\\rho$-zCDP [Definition 1.1, TCC 2016], our algorithm, after a single local update satisfies $\\frac{D_f^2}{2\\sigma^2}$-zCDP, where we use the fact that $ \\\\| \\operatorname{grad} f_i(z_{i,t}^r;\\mathcal{B}_{i,t}^r)\\\\| \\leq D_f$.\n\nii) Then, according to the composition theorem [Lemma 2.3, TCC 2016], the total privacy loss after $R$ rounds of communication is $\\frac{\\tau R  D_f^2}{2\\sigma^2}$-zCDP.\n\niii) Finally, we convert $\\rho$-zCDP to standard $(\\epsilon, \\delta)$-DP [Proposition 1.3, TCC 2016], which includes the quantitative relationship between the total privacy loss and the variance $\\sigma^2$. Given a privacy budget, we can calculate the corresponding $\\sigma^2$ to ensure that the algorithm satisfies differential privacy throughout its entire execution.\n\nThe above differential privacy mechanism can defend against scenarios where the server is honest but curious, as well as where eavesdroppers can intercept the upload and download channels. Our algorithm can also be extended to other threat scenarios, with further efforts on the algorithm design and privacy analysis.\n\n[TCC 2016] Bun M, Steinke T. Concentrated differential privacy: Simplifications, extensions, and lower bounds[C]//Theory of cryptography conference. Springer Berlin Heidelberg, 2016.\n\n6) Thank you very much for your valuable feedback, which allows us to clarify and improve our work.\nOur theoretical result, Theorem 4.3, is applicable to varying levels of data heterogeneity because our algorithm completely overcomes client drift caused by heterogeneous data and multiple local updates. First, we do not make any assumptions about the similarity of data across clients; second, we fully address client drift, as reflected in our convergence error, which eliminates the error associated with the degree of data heterogeneity. Please also refer to the response to Q4.\n\nWe sincerely appreciate your thoughtful and constructive feedback! We believe our responses have thoroughly addressed all concerns, leading to significant improvements and a stronger overall paper. We would be truly grateful if you could kindly reassess our work and consider raising the score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely appreciate the careful reviews and valuable comments from all the reviewers. We have provided point-by-point responses to each comment. In particular, for Reviewer 3dV9's Q1 regarding the comparison with related Riemannian methods and Q4 regarding the comparison with Euclidean baselines FedAvg and FedProx, our responses are as follows.

1) Reviewer 3dV9 Q1: Comparison with related Riemannian methods

i) Both [ICML 2020] and [Neurips 2021] consider the leading eigenvector problem, formulated as  $$\min_{x \in \mathbb{R}^d, \\|x\\| = 1} (-x^T A x) = \sum_{i=1}^n f_i(x), $$ 
where $A=\sum_{i=1}^n A_i $ and $f_i(x)=-x^T A_i x$. Let $ v_1$ be the leading eigenvector of $A$. Both papers consider the case when $A$ satisfies conditions that ensure that the problem has exactly two global minima: $v_1$ and $-v_1$. Both  [ICML 2020] and [Neurips 2021] establish linear convergence rates but rely heavily on this property of the global minima. 

In contrast, we consider a more general manifold optimization problem, formulated as 
$$ \min_{x \in \mathcal{M} \subset \mathbb{R}^{d\times k}} \frac{1}{n} \sum_{i=1}^n f_i(x),$$ where $x$ lies on $\mathcal{M}$ rather than on a sphere, and the objective functions $f_i$ are general nonconvex functions rather than only quadratic. Due to its generality, our analysis does not rely on any prior information about the global minima. 

ii) Both [ICML 2020] and [Neurips 2021] assume that data across clients are drawn from the same distribution. In contrast, our work addresses scenarios where the data distribution varies across clients, introducing heterogeneity.

iii) Similar to us, the paper [ICML 2020] also considers using local updates to reduce communication frequency. However, their algorithm requires the transmission of both local models and local gradient information, which increases the communication overhead. Additionally, during local updates, parallel transportation is needed to translate the global gradient information to a tangent space in preparation for using the retraction operator, which increases the computational cost. In contrast, with our novel algorithm design, we only transmit local models and avoid parallel transportation.

iv) [NeurIPS 2021] uses quantization to reduce the amount of transmitted information per communication round, while our algorithm reduces communication frequency through multiple local updates. These two communication-saving strategies are distinct and lead to significant differences in the algorithms. [NeurIPS 2021] encodes $\operatorname{grad} f_i$ at each client and decodes it at the server. For decoding, previous local information must be translated via parallel transport. In contrast, our algorithm is designed to efficiently overcome client drift caused by data heterogeneity and multiple local updates while addressing more general manifold constraints. Through an innovative algorithm design, we can achieve both communication and computation efficiency.

In the revised manuscript, we will compare our method with [ICML 2020] and [Neurips 2021]. 

2) Reviewer 3dV9 Q4: Comparison with related  Euclidean methods

 i) For FedAvg, the paper [AAAI 2019] studied the unconstrained nonconvex problem. 

 [AAAI 2019] Yu H et al. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. AAAI 2019. 

Under the assumption of bounded second moments $ \mathbb{E}\_{\zeta_i \sim \mathcal{D}_i} \\| \nabla f_i(x;\zeta_i) \\|^2 \le G^2$ (Assumption 1), they established the following convergence bound (Theorem 1): 
$$\mathcal{O}\left( \frac{\Delta^0}{\eta R\tau}  +  \eta^2 \tau^2 G^2 L^2 + \frac{\eta\sigma^2L}{n}\right).$$ 
Here, $\Delta^0$ represents the initialization error, and the other notation is consistent with the one used in our paper. The three terms in the error bound are attributed to initialization error, client drift, and stochastic variance, respectively. When data across clients are highly heterogeneous, $G$ can be very large, which causes the accuracy guarantee to deteriorate. In comparison, our algorithm, which overcomes the client drift issue, only includes the first term related to the initialization error and the third term related to the stochastic variance. Moreover, our algorithm does not require bounded second moments, i.e., that $ \mathbb{E}\_{\zeta_i \sim \mathcal{D}_i} \\| \nabla f_i(x;\zeta_i) \\|^2 \le G^2$. 

ii) For FedProx, Theorem 6 in [1] proves that after $R=\mathcal{O}(\frac{\Delta^0}{\rho \epsilon})$ communication rounds where $\rho$ is some parameter given in [1, Theorem 4], it holds that $\frac{1}{R} \sum_{r=1}^R \\|\nabla f(x^r)\\|^2 \le \epsilon $.  However, the validity of this result requires the following conditions:

(1) A $B$-local dissimilarity condition must be satisfied, i.e.,$\frac{1}{n} \sum_{i=1}^n \\|\nabla f_i (x) \\|^2 \le \\| \nabla f(x) \\|^2 B^2$.
This condition implies that the data across clients are somewhat similar. Please refer to Theorem 4 and Definition 3 in [1].

(2) The local updates need to solve the subproblem inexactly until a $\gamma$-inexactness condition is met, where $\gamma$ must satisfy  $\gamma B < 1$ (Remark 5). This implies that when the degree of data heterogeneity is high and, thus, $B$ is large, $\gamma$ becomes very small, requiring a large number of local updates. This, in turn, leads FedProx to be inefficient in practice.

In summary, compared to FedAvg and FedProx, despite our objective function being generally non-convex and subject to non-convex manifold constraints, our algorithm still overcomes the issue of client drift caused by heterogeneous data and local updates. Our main analytical advantages over FedAvg and FedProx are twofold: first, we do not make any assumptions about the similarity of data across clients; second, we completely overcome client drift, as reflected in our convergence error, which eliminates the error caused by the degree of data heterogeneity.

Meta-review: 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.
Decision: Accept (poster)