Distributed Principal Component Analysis with Limited Communication

1 Introduction . The need to process ever increasing quantities of data has motivated significant research interest into efficient distributed algorithms for standard tasks in optimization and data analysis ; see , e.g. , [ 16 , 27 , 11 ] . In this context , one of the fundamental metrics of interest is communication cost , measured as the number of bits sent and received by the processors throughout computation . Communicationefficient variants are known for several classical tasks and algorithms , from basic ones such as mean estimation [ 28 , 20 , 8 ] to variants of gradient descent [ 19 , 12 ] or even higher-order optimization [ 2 , 15 ] . By contrast , much less is known concerning the communication cost for classical problems in data analysis , such as Principal Component Analysis ( PCA ) [ 13 , 5 ] , or spectral clustering [ 23 ] . In this paper , we will focus on the closely-related problem of computing the leading eigenvector of the data covariance matrix , i.e. , determining the first principal component , in the setting where d-dimensional data comes from an unknown distribution , and is split among n machines . Given the fundamental nature of PCA , there is a rich body of work on solving sequential variants , both in the Euclidean domain [ 24 , 25 , 26 ] and , more recently , based on Riemannian optimization . This is because maximizing the Rayleigh quotient can be re-interpreted in terms of a hyper-sphere , leading to an unconstrained Riemannian optimization problem [ 1 , 32 , 31 ] . These references have exclusively focused on the sequential setting , where the data can be stored and processed by a single node . However , in many settings the large volume of data requires processing by multiple nodes , which has inspired significant work on reducing the distribution overheads of PCA , e.g . [ 4 , 11 ] , specifically on obtaining efficient algorithms in terms of round complexity . More precisely , an impressive line of work considered the deterministic setting [ 17 , 4 , 6 , 10 ] , where the input is partitioned arbitrarily between machines , for which both round-efficient and communication-efficient solutions are now known . The general approach is to first perform a careful low-rank approximation of the local data and then communicate iteratively to reconstruct a faithful 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . estimate of the covariance matrix . Due to the worst-case setting , the number of communication rounds scales polynomially in the eigengap δ and in the precision ε , making them costly in practice . Thus , on the practical side , it is common to assume a setting with stochastic rather than worst-case data partitioning . Specifically , Garber et al . [ 11 ] consider this setting , and note that standard solutions based on distributed Power Method or Lanczos would require a number of communication rounds which increases with the sample size , which is impractical at scale . They circumvent this issue by replacing explicit Power Method iterations with a set of convex optimization problems , which can be solved in a round-efficient manner . Huang and Pan [ 14 ] proposed a round-efficient solution , by leveraging the connection to Riemannian optimization . Their proposal relies on a procedure which allows the update of node-local variables via a surrogate gradient , whose deviation from the global gradient should be bounded . Thus , communication could only be used to reach consensus on the local variables among the nodes . Unfortunately , upon close inspection , we have noted significant technical gaps in their analysis , which we discuss in detail in the next section . The line of work in the stochastic setting focused primarily on latency cost , i.e. , the number of communication rounds for solving this problem . For highly-dimensional data and/or highly-distributed settings , it is known that bandwidth cost , i.e. , the number of bits which need to be transmitted per round , can be a significant bottleneck , and an important amount of work deals with it for other classic problems , e.g . [ 28 , 12 , 15 ] . By contrast , much less is known concerning the communication complexity of distributed PCA . Contribution . In this paper , we provide a new algorithm for distributed leading eigenvector computation , which specifically minimizes the total number of bits sent and received by the nodes . Our approach leverages the connection between PCA and Riemannian optimization ; specifically , we propose the first distributed variant of Riemannian gradient descent , which supports quantized communication . We prove convergence of this algorithm utilizing a set of necessary convexity-type properties , and bound its communication complexity under different initializations . At the technical level , our result is based on a new sequential variant of Riemannian gradient descent on the sphere . This may be of general interest , and the algorithm can be generalized to Riemannian manifolds of bounded sectional curvatures using more sophisticated geometric bounds . We strengthen the gradient dominance property proved in [ 31 ] , by proving that our objective is weakly-quasi-convex in a ball of some minimizer and combining that with a known quadratic-growth condition . We leverage these properties by adapting results of [ 7 ] in the Riemannian setting and choosing carefully the learning-rate in order to guarantee convergence for the single-node version . This algorithm is specifically-designed to be amenable to distribution : specifically , we port it to the distributed setting , and add communication-compression in the tangent spaces of the sphere , using a variant of the lattice quantization technique of [ 8 ] , which we port from the context of standard gradient descent . We provide a non-trivial analysis for the communication complexity of the resulting algorithm under different initializations . We show that , under random initialization , a solution can be reached with total communication that is quadratic in the dimension d , and linear in 1/δ , where δ is the gap between the two biggest eigenvalues ; when the initialization is done to the leading eigenvector of one of the local covariance matrices , this dependency can be reduced to linear in d . 2 Setting and Related Work . Setting . We assume to be given m total samples coming from some distribution D , organized as a global m× d data matrix M , which is partitioned row-wise among n processors , with node i being assigned the matrix Mi , consisting of mi consecutive rows , such that ∑n i=1 mi = m. As is common , let A : = MTM = ∑d i=1 M T i Mi be the global covariance matrix , and Ai : = M T i Mi the local covariance matrix owned by node i . We denote by λ1 , λ2 , ... , λd the eigenvalues of A in descending order and by v1 , v2 , ... , vd the corresponding eigenvectors . Then , it is well-known that , if the gap δ = λ1 − λ2 between the two leading eigenvalues of A is positive , we can approximate the leading eigenvector by solving the following empirical risk minimization problem up to accuracy ε : x⋆ = argminx∈Rd ( −x TAx ∥x∥2 ) = argminx∈Rd , ∥x∥2=1 ( −xTAx ) = argminx∈Sd−1 ( −xTAx ) , ( 1 ) where Sd−1 is the ( d− 1 ) -dimensional sphere . We define f : Sd−1 → R , with f ( x ) = −xTAx and fi : Sd−1 → R , with fi ( x ) = −xTAix . Since the inner product is bilinear , we can write the global cost as the sum of the local costs : f ( x ) = ∑n i=1 fi ( x ) . In our analysis we use notions from the geometry of the sphere , which we recall in Appendix A . Related Work . As discussed , there has been a significant amount of research on efficient variants of PCA and related problems [ 5 , 24 , 25 , 26 , 1 , 32 , 31 ] . Due to space constraints , we focus on related work on communication-efficient algorithms . In particular , we discuss the relationship to previous round-efficient algorithms ; to our knowledge , this is the first work to specifically focus on the bit complexity of this problem in the setting where data is randomly partitioned . More precisely , previous work on this variant implicitly assumes that algorithms are able to transmit real numbers at unit cost . The straightforward approach to solve the minimization problem ( 1 ) in a distributed setting , where the data rows are partitioned , would be to use a distributed version of the Power Method , Riemannian gradient descent ( RGD ) , or the Lanczos algorithm . In order to achieve an ε-approximation of the minimizer x⋆ , the latter two algorithms require Õ ( λ1 log ( 1/ε ) /δ ) rounds , where the Õ notation hides poly-logarithmic factors in d. Distributed Lanczos and accelerated RGD would improve this by an O ( √ λ1/δ ) factor . However , Garber et al . [ 11 ] point out that , when δ is small , e.g . δ = Θ ( 1/ √ m ) , then unfortunately the number of communication rounds would increase with the sample size , which renders these algorithms non-scalable . Standard distributed convex approaches , e.g . [ 16 , 27 ] , do not directly extend to our setting due to non-convexity and the unit-norm constraint . Garber et al . [ 11 ] proposed a variant of the Power Method called Distributed Shift-and-Invert ( DSI ) which converges in roughly O ( √ b δ √ m log2 ( 1/ε ) log ( 1/δ ) ) rounds , where b is a bound on the squared ℓ2-norm of the data . Huang and Pan [ 14 ] aimed to improve the dependency of the algorithm on ε and δ , and proposed an algorithm called Communication-Efficient Distributed Riemannian Eigensolver ( CEDRE ) . This algorithm is shown to have round complexity O ( b δ √ m log ( 1/ε ) ) , which does not scale with the sample size for δ = Ω ( 1/ √ m ) , and has only logarithmic dependency on the accuracy ε . Technical issues in [ 14 ] . Huang and Pan [ 14 ] proposed an interesting approach , which could provide the most round-efficient distributed algorithm to date . Despite the fact that we find the main idea of this paper very creative , we have unfortunately identified a significant gap in their analysis , which we now outline . Specifically , one of their main results , Theorem 3 , uses the local gradient dominance shown in [ 31 ] ; yet , the application of this result is invalid , as it is done without knowing in advance that the iterates of the algorithms continue to remain in the ball of initialization . This is compounded by another error on the constant of gradient dominance ( Lemma 2 ) , which we believe is caused by a typo in the last part of the proof of Theorem 4 in [ 31 ] . This typo is unfortunately propagated into their proof . More precisely , the objective is indeed gradient dominated , but with a constant which vanishes when we approach the equator , in contrast with the 2/δ which is claimed globally . ( We reprove this result independently in Proposition 4 ) . Thus , starting from a point lying in some ball of the minimizer can lead to a new point where the objective is more poorly gradient dominated . This is a non-trivial technical problem which , in the case of gradient descent , can be addressed by a choice of the learning rate depending on the initialization . Given these issues , we perform a new and formally-correct analysis for gradient descent based on new convexity properties which guarantee convergence directly in terms of iterates , and not just function values . We would like however to note that our focus in this paper is on bit and not round complexity .

Summary: This paper studies the problem of finding the leading eigenvector in a distributed setting, where the $i$-th worker holds $m_i$ pieces of local data and the goal is to find the leading eigenvector of the global covariance data matrix while minimizing the total communication cost. Unlike most of the previous results that largely focus on round complexity, this paper considers bit complexity.   The proposed algorithm is a combination of the Riemannian gradient descent over an $\ell_2$ sphere and the vector quantization method based on Davies et al. (2021). The main technical contribution of this paper is showing a "strong-weak-convexity" result of the quadratic function $-x^TAx$ over a Riemannian manifold, which resolves a technical issue in Huang et al. (2020).