Nonconvex Federated Learning on Compact Smooth Submanifolds With Heterogeneous Data

Thank you very much for your careful review and valuable comments.

Weaknesses

We have addressed all the comments accordingly. Please see the response to Questions for details.

Questions

1. Thank you very much for your comments. We have provided a comparison with [ICML 2020] and [Neurips 2021] in the general response.

2. 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.

Based 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.

3. 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.

4. Thank you very much for your suggestion. We have provided the comparison with FedAvg and FedProx. Please see the general response.

5. This is an excellent question. Our algorithm can indeed be extended to achieve differential privacy. Below, we outline how this can be accomplished.

Algorithmically, 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:

$$\begin{aligned} \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 )\\\\ c_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) \end{aligned}$$

where $\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$.

For the analysis, we can establish a quantitative relationship between privacy loss and the variance $\sigma^2$ as follows:

i) 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$.

ii) 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.

iii) 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.

The 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.

[TCC 2016] Bun M, Steinke T. Concentrated differential privacy: Simplifications, extensions, and lower bounds[C]//Theory of cryptography conference. Springer Berlin Heidelberg, 2016.

6. Thank you very much for your valuable feedback, which allows us to clarify and improve our work. Our 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.

We 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.

We would like to express our sincere gratitude for your careful review and positive assessment of our work!

Weaknesses

1. 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].

[SIOPT 2012] P-A. Absil, and Jerome Malick. ``Projection-like retractions on matrix manifolds." SIAM Journal on Optimization 22, no. 1 (2012): 135-158.

Questions

1. Thank you very much for your careful review. Although the kPCA problem is nonconvex, it has certain properties similar to those of a strongly convex function, e.g., the Lojasiewicz property [MP 2019]. Similar phenomena have also been observed in manifold optimization on low-rank matrices [NeurIPS 2021, AAAI 2022]. While challenging, it would be interesting to consider local regularity conditions such as the Lojasiewicz property and establish convergence rates faster than sublinear.

[MP 2019] Huikang Liu, Anthony Man-Cho So, and Weijie Wu. “Quadratic optimization with orthogonality constraint: explicit Lojasiewicz exponent and linear convergence of retraction-based line-search and stochastic variance-reduced gradient methods.” Mathematical Programming 178, no. 1 (2019): 215-262.

[Neurips 2021] Ye T, Du S S. Global convergence of gradient descent for asymmetric low-rank matrix factorization[J]. Advances in Neural Information Processing Systems, 2021, 34: 1429-1439.

[AAAI 2022] Bi Y, Zhang H, Lavaei J. Local and global linear convergence of general low-rank matrix recovery problems[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2022, 36(9): 10129-10137.

Thank you very much for your careful review and valuable comments!

Weaknesses

1. 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.

2. All results are derived and proven for Euclidean immersions and the Euclidean metric. In the submitted version of the manuscript, we stated this on Lines 132 and 135: “Throughout the paper, we restrict our discussion to embedded submanifolds of the Euclidean space, where the associated topology coincides with the subspace topology of the Euclidean space. · · · In addition, we always take the Euclidean metric as the Riemannian metric.” We will revise the manuscript to ensure that there are no misunderstanding that our results are restricted to this case.

3. 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.

Questions

1. 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.

2. These are good questions. We will address them one by one.

i) 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].

[SIOPT 2012] P-A. Absil, and Jerome Malick. "Projection-like retractions on matrix manifolds." SIAM Journal on Optimization 22, no. 1 (2012): 135-158.

ii) 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 $\mathcal{P}\_\mathcal{M}(\hat{z}\_{i, t+1}^r)$ and $\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.

3. The work [13] appears to be the only FL algorithm that can deal with manifold optimization problems of a similar generality as ours. The 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. We will make this point noted clearly in our revised manuscript.

In [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.

We 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!

Thank you for your careful review and positive assessment of our work!

Suggestions

Nonconvexity

1. 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.
2. 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.
3. It is due to the nonconvexity of the manifold constraint. The "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.
4. 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.

Numerical experiments

1. Thank you for your comments. Compared 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.
2. Thank you very much for allowing us the opportunity to clarify this point. To 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].

Notation and typos

1. Agree.
2. 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.
3. Agree.

Questions

1. 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.
2. For a general manifold where the ambient geometry is not Euclidean, it will be difficult to define the projection operator and the associated proximal smoothness. This will complicate the design and analysis of the algorithm. We will comment on this in the Conclusions and Limitations of the revised manuscript.
3. Thank you very much for your comments. In 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).