On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein Geometry

1 Introduction . Learning on symmetric positive definite ( SPD ) matrices is a fundamental problem in various machine learning applications , including metric and kernel learning [ 69 , 24 , 38 , 13 , 66 ] , medical imaging [ 57 , 56 ] , natural language processing [ 37 , 40 ] , computer vision [ 28 , 35 , 34 ] , multi-task learning [ 39 , 54 ] , domain adaptation [ 48 , 41 ] , modeling time-varying data [ 18 ] , object detection [ 70 ] , and quantum mechanics [ 53 , 47 ] , among others . The set of SPD matrices of size n× n , defined as Sn++ : = { X : X ∈ Rn×n , X > = X , and X 0 } , has a smooth manifold structure with a richer geometry than the Euclidean space . When endowed with a metric ( inner product structure ) , the set of SPD matrices becomes a Riemannian manifold [ 11 ] . Hence , numerous existing works [ 57 , 7 , 42 , 36 , 34 , 46 , 56 ] have studied and employed the Riemannian optimization framework for learning over the space of SPD matrices [ 2 , 16 ] . Several Riemannian metrics on Sn++ have been proposed such as the Affine-Invariant [ 57 , 11 ] , the Log-Euclidean [ 7 , 58 ] , the Log-Det [ 64 , 20 ] , and the Log-Cholesky [ 46 ] , to name a few . One can additionally obtain different families of Riemannian metrics on Sn++ by appropriate parameterizations based on the principles of invariance and symmetry [ 23 , 20 , 67 , 56 ] . However , to the best of our knowledge , a systematic study comparing the different metrics for optimizing generic cost functions defined on Sn++ is missing . Practically , the Affine-Invariant ( AI ) metric seems to be the most widely used metric in Riemannian first- and second-order algorithms ( e.g. , steepest descent , conjugate gradients , and trust regions ) as it is the only Riemannian SPD metric available in manifold optimization toolboxes , such as Manopt [ 17 ] , Manopt.jl [ 10 ] , Pymanopt [ 68 ] , ROPTLIB [ 32 ] , and McTorch [ 50 ] . Moreover , many interesting problems in machine learning are found to be geodesic convex ( generalization of Euclidean convexity ) under the AI metric , which allows fast convergence of optimization algorithms [ 76 , 31 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recent works have studied the Bures-Wasserstein ( BW ) distance on SPD matrices [ 49 , 12 , 72 ] . It is a well-known result that the Wasserstein distance between two multivariate Gaussian densities is a function of the BW distance between their covariance matrices . Indeed , the BW metric is a Riemannian metric . Under this metric , the necessary tools for Riemannian optimization , including the Riemannian gradient and Hessian expressions , can be efficiently computed [ 49 ] . Hence , it is a promising candidate for Riemannian optimization on Sn++ . In this work , we theoretically and empirically analyze the quality of optimization with the BW geometry and show that it is a viable alternative to the default choice of AI geometry . Our analysis discusses the classes of cost functions ( e.g. , polynomial ) for which the BW metric has better convergence rates than the AI metric . We also discuss cases ( e.g. , log-det ) where the reverse is true . In particular , our contributions are as follows . • We observe that the BW metric has a linear dependence on SPD matrices while the AI metric has a quadratic dependence . We show this impacts the condition number of the Riemannian Hessian and makes the BW metric more suited to learning ill-conditioned SPD matrices than the AI metric . • In contrast to the non-positively curved AI geometry , the BW geometry is shown to be non-negatively curved , which leads to a tighter trigonometry distance bound and faster convergence rates for optimization algorithms . • For both metrics , we analyze the convergence rates of Riemannian steepest descent and trust region methods and highlight the issues arising from the differences in the curvature and condition number of the Riemannian Hessian . • We verify that common optimization problems that are geodesic convex under the AI metric are also geodesic convex under the BW metric . • We support our analysis with extensive experiments on applications such as weighted least squares , trace regression , metric learning , and Gaussian mixture model . 2 Preliminaries . The fundamental ingredients for Riemannian optimization are Riemannian metric , exponential map , Riemannian gradient , and Riemannian Hessian . We refer readers to [ 2 , 16 ] for a general treatment on Riemannian optimization . A Riemannian metric is a smooth , bilinear , and symmetric positive definite function on the tangent space TxM for any x ∈ M. That is , g : TxM× TxM −→ R , which is often written as an inner product 〈· , ·〉x . The induced norm of a tangent vector u ∈ TxM is given by ‖u‖x = √ 〈u , u〉 x . A geodesic on a manifold γ : [ 0 , 1 ] −→M is defined as a locally shortest curve with zero acceleration . For any x ∈ M , u ∈ TxM , the exponential map , Expx : TxM −→ M is defined such that there exists a geodesic curve γ with γ ( 0 ) = x , γ ( 1 ) = Expx ( u ) and γ ′ ( 0 ) = u. First-order geometry and Riemannian steepest descent . Riemannian gradient of a differentiable function f : M−→ R at x , denoted as gradf ( x ) , is a tangent vector that satisfies for any u ∈ TxM , 〈gradf ( x ) , u〉x = Duf ( x ) , where Duf ( x ) is the directional derivative of f ( x ) along u . The Riemannian steepest descent method [ 71 ] generalizes the standard gradient descent in the Euclidean space to Riemannian manifolds by ensuring that the updates are along the geodesic and stay on the manifolds . That is , xt+1 = Expxt ( −ηt gradf ( xt ) ) for some step size ηt . Second-order geometry and Riemannian trust region . Second-order methods such as trust region and cubic regularized Newton methods are generalized to Riemannian manifolds [ 1 , 3 ] . They make use of the Riemannian Hessian , Hessf ( x ) : TxM−→ TxM , which is a linear operator that is defined as the covariant derivative of the Riemannian gradient . Both the trust region and cubic regularized Newton methods are Hessian-free in the sense that only evaluation of the Hessian acting on a tangent vector , i.e. , Hessf ( x ) [ u ] is required . Similar to the Euclidean counterpart , the Riemannian trust region method approximates the Newton step by solving a subproblem , i.e. , min u∈TxtM : ‖u‖xt≤∆ mxt ( u ) = f ( xt ) + 〈gradf ( xt ) , u〉xt + 1 2 〈Hxt [ u ] , u〉xt , whereHxt : TxtM−→ TxtM is a symmetric and linear operator that approximates the Riemannian Hessian . ∆ is the radius of trust region , which may be increased or decreased depending on how model value mxt ( u ) changes . The subproblem is solved iteratively using a truncated conjugate gradient algorithm . The next iterate is given by xt+1 = Expxt ( u ) with the optimized u . Next , the eigenvalues and the condition number of the Riemannian Hessian are defined as follows , which we use for analysis in Section 3 . Definition 1 . The minimum and maximum eigenvalues of Hessf ( x ) are defined as λmin = min‖u‖2x=1〈Hessf ( x ) [ u ] , u〉x and λmax = max‖u‖2x=1〈Hessf ( x ) [ u ] , u〉x . The condition number of Hessf ( x ) is defined as κ ( Hessf ( x ) ) : = λmax/λmin . Function classes on Riemannian manifolds . For analyzing algorithm convergence , we require the definitions for several important function classes on Riemannian manifolds , including geodesic convexity and smoothness . Similarly , we require the definition for geodesic convex sets that generalize ( Euclidean ) convex sets to manifolds [ 65 , 74 ] . Definition 2 ( Geodesic convex set [ 65 , 74 ] ) . A set X ⊆M is geodesic convex if for any x , y ∈ X , the distance minimizing geodesic γ joining the two points lies entirely in X . Indeed , this notion is well-defined for any manifold because a sufficiently small geodesic ball is always geodesic convex . Definition 3 ( Geodesic convexity [ 65 , 74 ] ) . Consider a geodesic convex set X ⊆ M. A function f : X −→ R is called geodesic convex if for any x , y ∈ X , the distance minimizing geodesic γ joining x and y satisfies ∀ t ∈ [ 0 , 1 ] , f ( γ ( t ) ) ≤ ( 1− t ) f ( x ) + tf ( y ) . Function f is strictly geodesic convex if the equality holds only when t = 0 , 1 . Definition 4 ( Geodesic strong convexity and smoothness [ 65 , 33 ] ) . Under the same settings in Definition 3 . A twice-continuously differentiable function f : X −→ R is called geodesic µ-strongly convex if for any distance minimizing geodesic γ in X with ‖γ′ ( 0 ) ‖ = 1 , it satisfies d 2f ( γ ( t ) ) dt2 ≥ µ , for some µ > 0 . Function f is called geodesic L-smooth if d 2f ( γ ( t ) ) dt2 ≤ L , for some L > 0 . 3 Comparing BW with AI for Riemannian optimization . This section starts with an observation of a linear-versus-quadratic dependency between the two metrics . From this observation , we analyze the condition number of the Riemannian Hessian . Then , we further compare the sectional curvature of the two geometries . Together with the differences in the condition number , this allows us to compare the convergence rates of optimization algorithms on the two geometries . We conclude this section by showing geodesic convexity of several generic cost functions under the BW geometry . The proofs for this section are in our extended report [ 27 ] . AI and BW geometries on SPD matrices . When endowed with a Riemannian metric g , the set of SPD matrices of size n becomes a Riemannian manifoldM = ( Sn++ , g ) . The tangent space at X is TXM : = { U : U ∈ Rn×n and U > = U } . Under the AI and BW metrics , the Riemannian exponential map , Riemannian gradient , and Hessian are compared in Table 1 , where we denote { A } S : = ( A + A > ) /2 and exp ( A ) as the matrix exponential of A. LX [ U ] is the solution to the matrix linear system LX [ U ] X + XLX [ U ] = U and is known as the Lyapunov operator . We use ∇f ( X ) and ∇2f ( X ) to represent the first-order and second-order derivatives , i.e. , the Euclidean gradient and Hessian , respectively . The derivations in Table 1 can be found in [ 56 , 12 ] . In the rest of the paper , we useMai andMbw to denote the SPD manifolds under the two metrics . From Table 1 , the computational costs for evaluating the AI and BW ingredients are dominated by the matrix exponential/inversion operations and the Lyapunov operator L computation , respectively . Both at most cost O ( n3 ) , which implies a comparable per-iteration cost of optimization algorithms between the two metric choices . This claim is validated in Section 4 . A key observation . From Table 1 , the Affine-Invariant metric on the SPD manifold can be rewritten as for any U , V ∈ TXM , 〈U , V〉ai = tr ( X−1UX−1V ) = vec ( U ) > ( X⊗X ) −1vec ( V ) , ( 1 ) where vec ( U ) and vec ( V ) are the vectorizations of U and V , respectively . Note that we omit the subscript X for inner product 〈· , ·〉 to simplify the notation . The specific tangent space where the inner product is computed should be clear from contexts . The Bures-Wasserstein metric is rewritten as , for any U , V ∈ TXM , 〈U , V〉bw = 1 2 tr ( LX [ U ] V ) = 1 2 vec ( U ) > ( X⊕X ) −1vec ( V ) , ( 2 ) where X⊕X = X⊗ I + I⊗X is the Kronecker sum . Remark 1 . Comparing Eq . ( 1 ) and ( 2 ) reveals that the BW metric has a linear dependence on X while the AI metric has a quadratic dependence . This suggests that optimization algorithms under the AI metric should be more sensitive to the condition number of X compared to the BW metric . The above observation serves as a key motivation for further analysis .

Summary: The authors consider the Riemannian optimization with Bures-Wasserstein (BW) geometry for symmetric positive definite (SPD) matrix manifold. The authors compare the proposed approach with its counterpart with the popular Affine-Invariant (AI) geometry.  The authors illustrate that the proposed approach (with BW geometry) is more suitable than its counterpart with AI geometry when SPD matrices are ill-conditioned.