Faster proximal algorithms for matrix optimization using Jacobi-based eigenvalue methods

1 Introduction . Many problems in statistics and learning can be formulated as convex optimization over matrices . This includes kernel learning [ LCB+04 ] , low-rank matrix recovery [ RFP10 , CR09 ] , sparse covariance selection [ BGdN06 , FHT08 ] , community detection [ BBV16 ] , and many others . Proximal splitting algorithms such as ( accelerated ) proximal gradient methods [ Nes03 , BT09 ] , or the alternating direction method of multipliers ( ADMM ) [ EB92 , BPC11 ] , are among the most popular methods to deal with large-scale instances of such problems . These methods solve generic convex optimization problems of the form min X∈E F ( X ) +G ( X ) ( 1 ) via calls to the proximal operators of F and/or G. When the variable X is a symmetric or rectangular matrix , it is often the case that the functions F and/or G are symmetric functions of the eigenvalues or singular values of X ; i.e. , so-called spectral or othogonally-invariant functions . Typical examples are F ( X ) = I ( X 0 ) = ∑n i=1 I ( λi ( X ) ≥ 0 ) the convex indicator function for the positive semidefinite cone , or F ( X ) = ‖X‖∗ = ∑n i=1 σi ( X ) the nuclear norm . Computing the proximal operator of such spectral functions F requires a full eigenvalue/singular value decomposition which contributes a large computational cost of at least n3 floating-point 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . operations per iteration . The most common algorithms for such decompositions are the QR method , and the divide-and-conquer algorithm [ GVL13 ] which are implemented in linear algebra packages ( LAPACK ) that are used in standard numerical computing software such as MATLAB or NumPy . Contributions In this paper we propose to use a simple and lesser known method due to Jacobi from 1846 [ Jac46 ] for the computation of these decompositions . Several features of this method make it suitable for its use within iterative proximal methods : ( i ) It can effectively exploit an initial approximate decomposition of the target matrix . In iterative proximal methods , the eigenvectors/singular vectors from the previous iterate are often a very good initial point that almost diagonalize the matrix of the current iterate . ( ii ) The method can be very easily parallelized , and is a very good fit for parallel hardware accelerators , notably GPUs , which are now very common in machine learning . ( iii ) It has a simple termination criterion that allows us to trade-off computation time with the accuracy of the desired decomposition . The present paper has two main contributions . First , we establish new theoretical results providing guarantees on the approximation quality of the proximal operator of a convex spectral function when computed with an approximate eigenvalue/singular value decomposition . When combined with existing results on the convergence of proximal splitting methods with inexact proximal oracles , we are able to appropriately set the tolerance of Jacobi-based methods and gain computation time , while guaranteeing global convergence of the optimization method . Second , we demonstrate using the observations above , and on a variety of problems and splitting algorithms , that the method of Jacobi on a GPU can yield speed-ups of 5-10x in the evaluation of proximal operators of spectral functions , compared to standard approaches on CPU and GPU . Related work Several approaches have been considered to avoid computing a full eigenvalue or singular value decomposition at each iteration , when the solution is known to be low-rank . Popular approaches include the Lanczos algorithm , see e.g. , [ CCS10 ] , randomized SVD [ HMT11 ] , or nonconvex low-rank formulations [ BM03 ] . However these approaches usually rely on an a priori knowledge of the rank of the solution . The recent paper [ RGN19 ] focusses on the ADMM method for SDPs and proposes the use of iterative eigensolvers ( LOBPCG [ Kny01 ] ) to approximately project a symmetric matrix on the positive semidefinite cone . The proposed method can achieve 10x speedups , especially when the primal or dual solution has low-rank . From the theoretical point of view , the approach relies on error bounds derived in [ GNR20 ] concerning approximate projections on the positive semidefinite cone . Our theoretical results in Section 4 constitute a significant generalization of these error bounds to proximal operators of convex spectral functions , and to different notions of approximations . The papers [ CV18 ] and [ JV21 ] consider proximal splitting algorithms with a generalized Bregman divergence ( i.e. , different than the squared Euclidean distance ) , and show that these proximal operators can sometimes be evaluated without explicitly computing a full eigenvalue decomposition . If φ is a strongly-convex function , the φ-proximal operator of F is defined by proxφF ( y ; a ) = argminx { F ( x ) + 〈a , x〉+ φ ( x ) − 〈∇φ ( y ) , x〉 } . The paper [ CV18 ] considers optimization problems over symmetric positive semidefinite Toeplitz matrices and shows that by choosing φ to be the kernel for the so-called Itakura-Saito distance one can compute the generalized proximal operator in O ( n2 ) time , instead of O ( n3 ) . In [ JV21 ] it is shown that for sparse semidefinite programs , an adequate choice of φ leads to a proximal operator that can be computed at the cost of a Cholesky decomposition of a sparse matrix which in general is much smaller than the cost of a full eigenvalue decomposition . Concerning sparse semidefinite programs , there has been a lot of work on devising fast ADMM algorithms using the clique decomposition method [ FKMN01 ] , e.g. , [ ZFP+20 , MKL15 ] . We note that these contributions are specific to sparse semidefinite programs and are complementary to the present paper because they still require an eigenvalue decomposition to be computed at each iteration . Finally let us mention that there are first-order methods that do not require a full eigenvalue decomposition at each iteration but only the computation of the smallest eigenvalue ( s ) /eigenvector ( s ) at each iteration such as e.g. , [ HR00 , Ren14 , YTF+21 ] and for which Lanczos methods are more suited . Organization In Section 2 we recall some necessary background material about convex spectral functions and their proximal operators . In Section 3 we present the Jacobi method to diagonalize symmetric matrices , and explain how it can be used with proximal algorithms . In Section 4 we present the main theoretical result of the paper and show that approximate eigenvalue/singular value decompositions yield guaranteed approximations of proximal operators for spectral functions . Finally we present numerical experiments in Section 5 . 2 Background . In this section we present some notations and background results concerning spectral functions . If f : E → R is a convex function on Euclidean space E we let ∂f ( x ) be the subdifferential of f at x , namely ∂f ( x ) = { g ∈ E : f ( y ) ≥ f ( x ) + 〈g , y − x〉 } . An element g ∈ ∂f ( x ) is called a subgradient of f at x . The proximal operator of f at y is defined as proxf ( y ) = argminx∈E { f ( x ) + ( 1/2 ) ‖x− y‖2 } . We let Sn be the Euclidean space of n×n symmetric matrices with the trace inner product 〈X , Y 〉 = tr ( XY ) . We let Sn+ ( resp . S n ++ ) be the cone of positive semidefinite ( resp . definite ) matrices in S n. Also O ( n ) denotes the set of n× n orthogonal matrices . Spectral functions A function F : Sn → R of symmetric matrices is called spectral if there is a permutation-invariant function f : Rn → R such that F ( X ) = f ( λ ( X ) ) where λ ( X ) = ( λ1 ( X ) , . . . , λn ( X ) ) is the vector of eigenvalues of X ( in any order ) . It is clear that spectral functions are orthogonally-invariant , i.e. , F ( UTXU ) = F ( X ) , because UTXU and X have the same eigenvalues ; conversely any orthogonally-invariant function F is the spectral function associated to the permutation-invariant function f ( x ) = F ( diag ( x ) ) . The main remarkable fact that we will use about spectral functions is a result due to Davis [ Dav57 ] , which asserts that F is convex as soon as f is , and gives a variational formula for F certifying this convexity . Theorem 1 ( [ Dav57 ] , see also [ BTN01 , p.282 , claim D ] ) . Let f : Rn → R be a permutationinvariant function , and let F : Sn → R , F ( X ) = f ( λ ( X ) ) be the associated spectral function . If f is convex on Rn , then F is convex on Sn , and we have the following variational formulation for F : F ( X ) = max U∈O ( n ) f ( diag ( UTXU ) ) . ( 2 ) An important , and well-known consequence of the theorem above , is that the proximal operator of F can be obtained from that of f , namely for any Y ∈ Sn we have proxF ( Y ) = U diag ( proxf ( y ) ) U T where Y = U diag ( y ) UT is an eigenvalue decomposition of Y . Functions of singular values Similar results hold for functions F : Rm×n → R of rectangular matrices that are orthogonally-invariant , i.e. , F ( UXV ) = F ( X ) for all U ∈ O ( m ) and V ∈ O ( n ) . Any such function can be written as F ( X ) = f ( σ ( X ) ) where σ ( X ) ∈ Rmin ( m , n ) is the vector of singular values and f : Rmin ( m , n ) → R is an absolutely symmetric function , i.e. , it is invariant under signed permutations : f ( η1xπ ( 1 ) , . . . , ηkxπ ( k ) ) = f ( x1 , . . . , xk ) for ηi = ±1 and π a permutation of { 1 , . . . , k } where k = min ( m , n ) [ Lew95 ] . In this case too , it holds that F is convex as soon as f is , and we have : F ( X ) = max U∈O ( m ) , V ∈O ( n ) f ( diag ( UTXV ) ) . ( 3 ) The variational formulas ( 2 ) and ( 3 ) are a consequence of well-known results concerning majorization relations between the diagonal of a matrix and its eigenvalues or singular values [ MOA79 , Chapter 9 ] . For completeness we have included a proof of these formulas in the supplementary material . We refer the reader to [ PB14 , Section 6.7 ] and references therein for more details on the proximal operator of spectral and orthogonally-invariant functions .

Summary: This paper proposes a new method to accelerate proximal-mapping based algorithms  where the proximal mapping usually needs computing the eigen/singular value decomposition.   The authors discusses how to use the Jacobi method to get the eigen/singular value decomposition approximately and parallelly,  and the convergence can be guaranteed by easily controlling the norm of the off diagonal elements.   Numerical experiments show the significant improvement over the state-of-art algorithms.