Low-Rank Extragradient Method for Nonsmooth and Low-Rank Matrix Optimization Problems

1 Introduction . Low-rank and nonsmooth matrix optimization problems have many important applications in statistics , machine learning , and related fields , such as sparse PCA [ 21 , 34 ] , robust PCA [ 28 , 33 , 2 , 9 , 38 ] , phase synchronization [ 41 , 6 , 29 ] , community detection and stochastic block models [ 1 ] 1 , low-rank and sparse covariance matrix recovery [ 35 ] , robust matrix completion [ 22 , 10 ] , and more . For many of these problems , convex relaxations , in which one replaces the nonconvex low-rank constraint with a trace-norm constraint , have been demonstrated in numerous papers to be highly effective both in theory ( under suitable assumptions ) and empirically ( see references above ) . These convex relaxations can be formulated as the following general nonsmooth optimization problem : min X∈Sn g ( X ) , ( 1 ) where g : Sn → R is convex but nonsmooth , and Sn = { X ∈ Sn | Tr ( X ) = 1 , X 0 } is the spectrahedron in Sn , Sn being the space of n× n real symmetric matrices . Problem ( 1 ) , despite being convex , is notoriously difficult to solve in large scale . The simplest and most general approach applicable to it is the projected subgradient method [ 3 , 7 ] , which requires on 1in [ 41 , 6 , 29 ] and [ 1 ] the authors consider SDPs with linear objective function and affine constraints of the form A ( X ) = b . By incorporating the linear constraints into the objective function via a ` 2 penalty term of the form λ‖A ( X ) − b‖2 , λ > 0 , we obtain a nonsmooth objective function . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . each iteration to compute a Euclidean projection onto the spectrahedron Sn , which in worst case amounts to O ( n3 ) runtime per iteration . In many applications g ( X ) follows a composite model , i.e. , g ( X ) = h ( X ) + w ( X ) , where h ( · ) is convex and smooth and w ( · ) is convex and nonsmooth but admits a simple structure ( e.g. , nonsmooth regularizer ) . For such composite objectives , without the spectrahedron constraint , proximal methods such as FISTA [ 4 ] or splitting methods such as ADMM [ 32 ] are often very effective . However , with the spectrahderon constraint , all such methods require on each iteration to apply a subprocedure ( e.g. , computing the proximal mapping ) which in worst case amounts to at least O ( n3 ) runtime . A third type of off-the-shelf methods include those which are based on the conditional gradient method and adapted to nonsmooth problems , see for instance [ 31 , 19 , 36 , 26 ] . The advantage of such methods is that no expensive high-rank SVD computations are needed . Instead , only a single leading eigenvector computation ( i.e. , a rank-one SVD ) per iteration is required . However , whenever the number of iterations is not small , these methods still require to store high-rank matrices in memory , even when the optimal solution is low-rank . Thus , to conclude , standard first-order methods for Problem ( 1 ) require in worst case Ω ( n3 ) runtime per iteration and/or to store high-rank matrices in memory . In the recent works [ 17 , 18 ] it was established that for smooth objective functions , the high-rank SVD computations required for Euclidean projections onto the spectrahedron in standard gradient methods , can be replaced with low-rank SVDs in the close proximity of a low-rank optimal solution . This is significant since the runtime to compute a rank-r SVD of a given n×n matrix using efficient iterative methods typically scales with rn2 ( and further improves when the matrix is sparse ) , instead of n3 for a full-rank SVD . These results depend on the existence of eigen-gaps in the gradient of the optimal solution , which we refer to as a generalized strict complementarity condition . These results also hinge on a unique property of the Euclidean projection onto the spectrahedron . The projection onto the spectrahedron of a matrix X ∈ Sn , which admits an eigen-decomposition X = ∑n i=1 λiviv > i , is given by ΠSn [ X ] = n∑ i=1 max { 0 , λi − λ } viv > i , ( 2 ) where λ ∈ R is the unique scalar satisfying ∑n i=1 max { 0 , λi−λ } = 1 . This operation thus truncates all eigenvalues that are smaller than λ , while leaving the eigenvectors unchanged , thereby returning a matrix with rank equal to the number of eigenvalues greater than λ . Importantly , when the projection of X onto Sn is of rank r , only the first r components in the eigen-decomposition of X are required to compute it in the first place , and thus , only a rank-r SVD of X is required . In other words and simplifying , [ 17 , 18 ] show that under strict complementary , at the proximity of an optimal solution of rank r , the exact Euclidean projection equals the rank-r truncated projection given by : Π̂rSn [ X ] : = ΠSn [ r∑ i=1 λiviv > i ] . ( 3 ) Extending the results of [ 17 , 18 ] to the nonsmooth setting is difficult since the smoothness assumption is critical to the analysis . Moreover , while [ 17 , 18 ] rely on certain eigen-gaps in the gradients at optimal points , for nonsmooth problems , since the subdifferential set is often not a singleton , it is not likely that a similar eigen-gap property holds for all subgradients of an optimal solution . In this paper we show that under the mild assumption that Problem ( 1 ) can be formulated as a smooth convex-concave saddle-point problem , i.e. , the nonsmooth term can be written as a maximum over ( possibly infinite number of ) smooth convex functions , we can obtain results in the spirit of [ 17 , 18 ] . Concretely , we show that if a generalized strict complementarity ( GSC ) assumption holds for a low-rank optimal solution ( see Assumption 1 in the sequel ) , the extragradient method for smooth convex-concave saddle-point problems [ 23 , 30 ] ( see Algorithm 1 below ) , when initialized in the proximity of the optimal solution , converges with its original convergence rate of O ( 1/t ) , while requiring only two low-rank SVDs per iteration2 . It is important to recall that while the extragradient method requires two SVDs per iteration , it has the benefit of a fast O ( 1/t ) convergence rate , while simpler saddle-point methods such as mirror-descent-based only achieve a O ( 1/ √ t ) rate [ 7 ] . Our contributions can be summarized as follows : 2note that the extradgradient method computes two projected-gradient steps on each iteration , and thus two SVDs are needed per iteration . • We prove that even under ( standard ) strict complementarity , the projected subgradient method , when initialized with a “ warm-start '' , may produce iterates with rank higher than that of the optimal solution . This phenomena further motivates our saddle-point approach . See Lemma 5 . • We suggest a generalized strict complementarity ( GSC ) condition for saddle-point problems and prove that when g ( · ) — the objective function in Problem ( 1 ) , admits a highly popular saddle-point structure ( one which captures all applications we mentioned in this paper ) , GSC w.r.t . an optimal solution to Problem ( 1 ) implies GSC ( with the same parameters ) w.r.t . a corresponding optimal solution of the equivalent saddle-point problem ( the other direction always holds ) . See Section 3 . • Main result : we prove that for a smooth convex-concave saddle-point problem and an optimal solution which satisfies GSC , the extragradient method , when initialized with a “ warm-start '' , converges with its original rate of O ( 1/t ) while requiring only two low-rank SVDs per iteration . Moreover , we prove GSC facilitates a precise and powerful tradeoff : increasing the rank of SVD computations ( beyond the rank of the optimal solution ) can significantly increase the radius of the ball in which the method needs to be initialized . See Theorem 1 . • We present extensive numerical evidence that demonstrate both the plausibility of the GSC assumption in various tasks , and more importantly , demonstrate that indeed the extragradient method with simple initialization converges correctly ( i.e. , produces exactly the same sequences of iterates ) when the rank of the SVDs used to compute the ( truncated ) projections matches the rank of the ( low-rank ) ground-truth matrix to be recovered , instead of naively using full-rank SVDs ( as suggested by ( 2 ) ) . See Section 5 . 1.1 Additional related work . Since , as in the works [ 17 , 18 ] mentioned before which deal with smooth objectives , strict complementarity plays a key role in our analysis , we refer the interested reader to the recent works [ 16 , 39 , 13 , 20 ] which also exploit this property for efficient smooth and convex optimization over the spectrahedron . Strict complementarity has also played an instrumental role in two recent and very influential works which used it to prove linear convergence rates for proximal gradient methods [ 42 , 14 ] . Besides convex relaxations such as Problem ( 1 ) , considerable advances have been made in the past several yeas in developing efficient nonconvex methods with global convergence guarantees for low-rank matrix problems . In [ 37 ] the authors consider semidefinite programs and prove that under a smooth manifold assumption on the constraints , such methods converge to the optimal global solution . In [ 24 ] the authors prove global convergence of factorized nonconvex gradient descent from a ” warm-start ” initialization point for non-linear smooth minimization on the positive semidefinite cone . Very recently , [ 8 ] has established , under statistical conditions , fast convergence results from “ warm-start ” initialization of nonconvex first-order methods , when applied to nonsmooth nonconvex matrix recovery problems which are based on the explicit factorization of the low-rank matrix . A result of similar flavor concerning nonsmooth and nonconvex formulation of robust recovery of lowrank matrices from random linear measurements was presented in [ 25 ] . Finally , several recent works have considered nonconvex low-rank regularizers which result in nonconvex nonsmooth optimization problems , but guarantee convergence only to a stationary point [ 27 , 40 ] . 2 Strict complementarity for nonsmooth optimization and difficulty of applying low-rank projected subgradient steps . Our analysis of the nonsmooth Problem ( 1 ) naturally depends on certain subgradients of an optimal solution which , in many aspects , behave like the gradients of smooth functions . The existence of such a subgradient is guaranteed from the first-order optimality condition for constrained convex minimization problems : Lemma 1 ( first-order optimality condition , see [ 3 ] ) . Let g : Sn → R be a convex function . Then X∗ ∈ Sn minimizes g over Sn if and only if there exists a subgradient G∗ ∈ ∂g ( X∗ ) such that 〈X−X∗ , G∗〉 ≥ 0 for all X ∈ Sn . For some G∗ ∈ ∂g ( X∗ ) which satisfies the first-order optimality condition for an optimal solution X∗ , if the multiplicity of the smallest eigenvalue equals r∗ = rank ( X∗ ) , then it can be shown that the optimal solution satisfies a strict complementarity assumption . The equivalence between a standard strict complementarity assumption on some low-rank optimal solution of a smooth optimization problem over the spectrahedron and an eigen-gap in the gradient of the optimal solution was established in [ 39 ] . We generalize this equivalence to also include nonsmooth problems . The proof follows similar arguments and is given in Appendix A.1 . Definition 1 ( strict complementarity ) . An optimal solution X∗ ∈ Sn of rank r∗ for Problem ( 1 ) satisfies the strict complementarity assumption with parameter δ > 0 , if there exists an optimal solution of the dual problem3 ( Z∗ , s∗ ) ∈ Sn ×R such that rank ( Z∗ ) = n− r∗ , and λn−r∗ ( Z∗ ) ≥ δ. Lemma 2 . Let X∗ ∈ Sn be a rank-r∗ optimal solution to Problem ( 1 ) . X∗ satisfies the ( standard ) strict complementarity assumption with parameter δ > 0 if and only if there exists a subgradient G∗ ∈ ∂g ( X∗ ) such that 〈X−X∗ , G∗〉 ≥ 0 for all X ∈ Sn and λn−r∗ ( G∗ ) − λn ( G∗ ) ≥ δ . Throughout this paper we assume a weaker and more general assumption than strict complementarity , namely generalized strict complementarity ( GSC ) , which we present now . Assumption 1 ( generalized strict complementarity ) . We say an optimal solution X∗ to Problem ( 1 ) satisfies the generalized strict complementarity assumption with parameters r , δ , if there exists a subgradient G∗ ∈ ∂g ( X∗ ) such that 〈X−X∗ , G∗〉 ≥ 0 for all X ∈ Sn and λn−r ( G∗ ) −λn ( G∗ ) ≥ δ . In [ 17 ] the author presents several characteristic properties of the gradient of the optimal solution in optimization problems over the spectrahedron . Using the existence of subgradients which satisfy the condition in Lemma 1 , we can extend these properties also to the nonsmooth setting . The following lemma shows that GSC with parameters ( r , δ ) for some δ > 0 ( Assumption 1 ) is a sufficient condition for the optimal solution to be of rank at most r. The proof follows immediately from the proof of the analogous Lemma 7 in [ 17 ] , by replacing the gradient of the optimal solution with a subgradient for which the first-order optimality condition holds . Lemma 3 . Let X∗ be an optimal solution to Problem ( 1 ) and write its eigen-decomposition as X∗ = ∑r∗ i=1 λiviv T i . Then , any subgradient G ∗ ∈ ∂g ( X∗ ) which satisfies 〈X−X∗ , G∗〉 ≥ 0 for all X ∈ Sn , admits an eigen-decomposition such that the set of vectors { vi } r ∗ i=1 is a set of leading eigenvectors of ( −G∗ ) which corresponds to the eigenvalue λ1 ( −G∗ ) = −λn ( G∗ ) . Furthermore , there exists at least one such subgradient . One motivation for assuming ( standard ) strict complementarity ( Assumption 1 with parameters r = rank ( X∗ ) and δ > 0 ) is that it guarantees a certain notion of robustness of the problem to small perturbations in the parameters . It is well known ( see for instance [ 3 ] ) that a projected subgradient step from X∗ with respect to a subgradient G∗ ∈ ∂g ( X∗ ) for which the first-order optimality condition holds , returns the optimal solution X∗ itself . This implies that rank ( ΠSn [ X ∗ − ηG∗ ] ) = rank ( X∗ ) ( here η is the step-size ) . Without ( standard ) strict complementarity however , a small change in the parameters could result in a higher rank matrix . This is captured in the following lemma which is analogous to Lemma 3 in [ 18 ] , where again the proof is straightforward from the proof in [ 18 ] by replacing the gradient of the optimal solution with a subgradient for which the first-order optimality condition holds . Lemma 4 . Let X∗ be an optimal solution of rank r∗ to Problem ( 1 ) . Let G∗ ∈ ∂g ( X∗ ) be a subgradient at X∗ such that 〈X −X∗ , G∗〉 ≥ 0 for all X ∈ Sn . Then , λn−r∗ ( G∗ ) = λn ( G∗ ) if and only if for any arbitrarily small ζ > 0 it holds that rank ( Π ( 1+ζ ) Sn [ X ∗ − ηG∗ ] ) > r∗ , where η > 0 , ( 1 + ζ ) Sn = { ( 1 + ζ ) X |X ∈ Sn } , and Π ( 1+ζ ) Sn [ · ] denotes the Euclidean projection onto the set ( 1 + ζ ) Sn .

Summary: This paper considers the problem of minimizing a non-smooth matrix function  $f(X)$ over the spectrahedron $S_n$, i.e., positive semidefinite matrices with trace equal to 1. Such optimization problems often arise as the convex relaxation of various low-rank matrix recovery problems. In general solving these problems with projected subgradient descent requires a projection onto the spectraherdon at every iteration, which is expensive because of the need to compute a full SVD at every step.  Under the assumption that the non-smooth objective function $f(X)$ can be written as the pointwise maximum of a family of smooth functions,  the authors reformulate the original optimization problem as a saddle point problem which can be solved via the projected extragradient method. The key contribution is that under a “generalized strict complementarity” condition and when close to the ground truth, the projection onto $S_n$ can be replaced by a low-rank projection. Therefore, the overall algorithm can be significantly more efficient, because only a low-rank SVD is needed.