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: This paper provides several results about the Bures-Wasserstein  (BW) geometry on the positive definite matrices. A thorough comparison with the Affine Invariant (AI) metric is made. The main focus of the paper is on optimization: is it faster to optimize a function under the AI or under the BW metric? The authors argue that when the solution of the problem $X^*$ is ill-conditioned, then algorithms derived under BW metric have faster convergence theoretical convergence. This is due to the conditioning of the Riemannian Hessian at the optimum, which is better for the BW metric when $X^*$ has poor conditioning. The authors also demonstrate that a variety of functions that are convex for the AI geometry are also convex for the BW geometry. Experiments on an array of synthetic problems validate that both Riemannian gradient descent and Riemannian trust region methods are faster in the BW geometry.