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

Abstract:
In this paper, we comparatively analyze the Bures-Wasserstein (BW) geometry with the popular Affine-Invariant (AI) geometry for Riemannian optimization on the symmetric positive definite (SPD) matrix manifold. Our study begins with an observation that the BW metric has a linear dependence on SPD matrices in contrast to the quadratic dependence of the AI metric. We build on this to show that the BW metric is a more suitable and robust choice for several Riemannian optimization problems over ill-conditioned SPD matrices. We show that the BW geometry has a non-negative curvature, which further improves convergence rates of algorithms over the non-positively curved AI geometry. Finally, we verify that several popular cost functions, which are known to be geodesic convex under the AI geometry, are also geodesic convex under the BW geometry. Extensive experiments on various applications support our findings.

heading: 1 Introduction
text: 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.

heading: 2 Preliminaries
text: 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.

heading: 3 Comparing BW with AI for Riemannian optimization
text: 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.

heading: 3.1 Condition number of Riemannian Hessian at optimality
text: Throughout the rest of the paper, we make the following assumptions. Assumption 1. (a). f is at least twice continuously differentiable with a non-degenerate local minimizer X∗. (b). The subset X ⊆M (usually as a neighbourhood of a center point) we consider throughout this paper is totally normal, i.e., the exponential map is a diffeomorphism.
Assumption 1 is easy to satisfy. Particularly, Assumption 1(b) is guaranteed for the SPD manifold under the AI metric because its geodesic is unique. Under the BW metric, for a center point X, we can choose the neighbourhood such that X = {ExpX(U) : I + LX[U] ∈ Sn++} as in [49]. In other words, X is assumed to be unique-geodesic under both the metrics. We now formalize the impact of the linear-versus-quadratic dependency, highlighted in Remark 1. At a local minimizer X∗ where the Riemannian gradient vanishes, we first simplify the expression for the Riemannian Hessian in Table 1.
OnMai, 〈Hessaif(X∗)[U],U〉ai = tr(∇2f(X∗)[U]U) = vec(U)>H(X∗)vec(U), with H(X) ∈ Rn2×n2 is the matrix representation of the Euclidean Hessian ∇2f(X) and U ∈ TX∗Mai. The maximum eigenvalue of Hessaif(X∗) is then given by λ∗max = max‖U‖2ai=1 vec(U)
>H(X∗)vec(U), where ‖U‖2ai = vec(U)>(X∗ ⊗X∗)−1vec(U). This is a generalized eigenvalue problem with the solution to be the maximum eigenvalue of (X∗ ⊗X∗)H(X∗). Similarly, λ∗min corresponds to the minimum eigenvalue of (X∗ ⊗X∗)H(X∗).
On Mbw, 〈Hessbwf(X∗)[U],U〉bw = vec(U)>H(X∗)vec(U) and the norm is ‖U‖2bw = vec(U)>(X∗ ⊕X∗)−1vec(U). Hence, the minimum/maximum eigenvalue of Hessbwf(X∗) equals the minimum/maximum eigenvalue of (X∗ ⊕X∗)H(X∗). Let κ∗ai := κ(Hessaif(X
∗)) = κ((X∗ ⊗X∗)H(X∗)) and κ∗bw := κ(Hessbwf(X∗)) = κ((X∗ ⊕ X∗)H(X∗)). The following lemma bounds these two condition numbers. Lemma 1. For a local minimizer X∗ of f(X), the condition number of Hessf(X∗) satisfies
κ(X∗)2/κ(H(X∗)) ≤κ∗ai ≤ κ(X∗)2κ(H(X∗)) κ(X∗)/κ(H(X∗)) ≤κ∗bw ≤ κ(X∗)κ(H(X∗)).
It is clear that κ∗bw ≤ κ∗ai when κ(H(X∗)) ≤ √ κ(X∗). This is true for linear, quadratic, higher-order polynomial functions and in general holds for several machine learning optimization problems on the SPD matrices (discussed in Section 4). Case 1 (Condition number for linear and quadratic optimization). For a linear function f(X) = tr(XA), its Euclidean Hessian matrix is H(X) = 0. For a quadratic function f(X) = tr(XAXB) with A,B ∈ Sn++, H(X) = A ⊗ B + B ⊗ A. Therefore, κ(H(X∗)) is a constant and for ill-conditioned X∗, we have κ(H(X∗)) ≤ √ κ(X∗), which leads to κ∗ai ≥ κ∗bw. Case 2 (Condition number for higher-order polynomial optimization). For an integer α ≥ 3, consider a function f(X) = tr(Xα) with derived H(X) = α ∑α−2 l=0 (X
l ⊗ Xα−l−2). We get κ∗ai = α ∑α−1 l=1 ((X ∗)l ⊗ (X∗)α−l) and κ∗bw = α(X∗ ⊕X∗)( ∑α−2 l=0 (X
l ⊗Xα−l−2)). It is apparent that κ∗ai = O(κ(X∗)α) while κ∗bw = O(κ(X∗)α−1). Hence, for ill-conditioned X∗, κ∗ai ≥ κ∗bw.
One counter-example where κ∗bw ≥ κ∗ai is the log-det function. Case 3 (Condition number for log-det optimization). For the log-det function f(X) = − log det(X), its Euclidean Hessian is∇2f(X)[U] = X−1UX−1 and H(X) = X−1⊗X−1. At a local minimizer X∗, Hessaif(X∗)[U] = U with κ∗ai = 1. While onMbw, we have κ∗bw = κ((X∗ ⊕X∗)((X∗)−1 ⊗ (X∗)−1)) = κ((X∗)−1 ⊕ (X∗)−1) = κ(X∗). Therefore, κ∗ai ≤ κ∗bw.

heading: 3.2 Sectional curvature and trigonometry distance bound
text: To study the curvature ofMbw, we first show in Lemma 2, the existence of a matching geodesic between the Wasserstein geometry of zero-centered non-degenerate Gaussian measures and the BW geometry of SPD matrices. Denote the manifold of such Gaussian measures under the L2-Wasserstein distance as (N0(Σ),W2) with Σ ∈ Sn++. Lemma 2. For any X,Y ∈ Sn++, a geodesic betweenN0(X) andN0(Y) on (N0(Σ),W2) is given by N0(γ(t)), where γ(t) is the geodesic between X and Y onMbw.
The following lemma builds on a result from the Wasserstein geometry [6] and uses Lemma 2 to analyze the sectional curvature ofMbw. Lemma 3. Mbw is an Alexandrov space with non-negative sectional curvature.
It is well-known thatMai is a non-positively curved space [21, 56] while, in Lemma 3, we show that Mbw is non-negatively curved. The difference affects the curvature constant in the trigonometry distance bound of Alexandrov space [77]. This bound is crucial in analyzing convergence for optimization algorithms on Riemannian manifolds [77, 76]. In Section 3.3, only local convergence to a minimizer X∗ is analyzed. Therefore, it suffices to consider a neighbourhood Ω around X∗. In such a compact set, the sectional curvature is known to be bounded and we denote the lower bound as K−.
The following lemma compares the trigonometry distance bounds under the AI and BW geometries. This bound was originally introduced for Alexandrov space with lower bounded sectional curvature [77]. The result for non-negatively curved spaces has been applied in many work [76, 60, 25] though without a formal proof. We show the proof in [27], where it follows from the Toponogov comparison theorem [52] on the unit hypersphere and Assumption 1. Lemma 4. Let X,Y,Z ∈ Ω, which forms a geodesic triangle onM. Denote x = d(Y,Z), y = d(X,Z), z = d(X,Y) as the geodesic side lengths and let θ be the angle between sides y and z such that cos(θ) = 〈Exp−1X (Y),Exp −1 X (Z)〉/(yz). Then, we have
x2 ≤ ζy2 + z2 − 2yz cos(θ),
where ζ is a curvature constant. Under the AI metric, ζ = ζai = √ |K−ai |D
tanh( √ |K−ai |D) with D as the
diameter bound of Ω, i.e. maxX1,X2∈Ω d(X,Y) ≤ D. Under the BW metric, ζ = ζbw = 1.
It is clear that ζai > ζbw = 1, which leads to a tighter bound under the BW metric.

heading: 3.3 Convergence analysis
text: We now analyze the local convergence properties of the Riemannian steepest descent and trust region methods under the two Riemannian geometries. Convergence is established in terms of the
Riemannian distance induced from the geodesics. We begin by presenting a lemma that shows in a neighbourhood of X∗, the second-order derivatives of f ◦ ExpX are both lower and upper bounded. Lemma 5. In a totally normal neighbourhood Ω around a non-degenerate local minimizer X∗, for any X ∈ Ω, it satisfies that λ∗min/α ≤ d 2 dt2 f(ExpX(tU)) ≤ αλ ∗ max, for some α ≥ 1 and ‖U‖ = 1. λ∗max > λ ∗ min > 0 are the largest and smallest eigenvalues of Hessf(X ∗).
For simplicity of the analysis, we assume such an α is universal under both the Riemannian geometries. We, therefore, can work with a neighbourhood Ω with diameter uniformly bounded by D, where we can choose D := min{Dai, Dbw} such that α is universal. One can readily check that under Lemma 5 the function f is both µ-geodesic strongly convex and L-geodesic smooth in Ω where µ = λ∗min/α and L = αλ ∗ max. We now present the local convergence analysis of the two algorithms, which are based on results in [77, 1]. Theorem 1 (Local convergence of Riemannian steepest descent). Under Assumption 1 and consider a non-degenerate local minimizer X∗. For a neighbourhood Ω 3 X∗ with diameter bounded by D on two Riemannian geometriesMai,Mbw, running Riemannian steepest descent from X0 ∈ Ω with a fixed step size η = 1αλ∗max yields for t ≥ 2,
d2(Xt,X ∗) ≤ α2D2κ∗ (1−min{ 1
ζ ,
1
α2κ∗ })t−2 .
Theorem 2 (Local convergence of Riemannian trust region). Under the same settings as in Theorem 1, assume further in Ω, it holds that (1) ‖HXt − Hessf(Xt)‖ ≤ `‖gradf(Xt)‖ and (2) ‖∇2(f ◦ ExpXt)(U)−∇
2(f ◦ ExpXt)(0)‖ ≤ ρ‖U‖ for some `, ρ universal onMai,Mbw. Then running Riemannian trust region from X0 ∈ Ω yields,
d(Xt,X ∗) ≤ (2√ρ+ `)(κ∗)2d2(Xt−1,X∗).
Theorems 1 and 2 show thatMbw has a clear advantage compared toMai for learning ill-conditioned SPD matrices where κ∗bw ≤ κ∗ai. For first-order algorithms,Mbw has an additional benefit due to its non-negative sectional curvature. As ζai > ζbw = 1, the convergence rate degrades on Mai. Although the convergence is presented in Riemannian distance, it can be readily converted to function value gap by noticing µ2 d 2(Xt,X ∗) ≤ f(Xt)− f(X∗) ≤ L2 d 2(Xt,X ∗). Additionally, we note that these local convergence results hold regardless of whether the function is geodesic convex or not, and similar comparisons also exist for other Riemannian optimization methods.

heading: 3.4 Geodesic convexity under BW metric for cost functions of interest
text: Finally we show geodesic convexity of common optimization problems onMbw. Particularly, we verify that linear, quadratic, log-det optimization, and also certain geometric optimization problems, that are geodesic convex under the AI metric, are also geodesic convex under the BW metric. Proposition 1. For any A ∈ Sn+, where Sn+ := {Z : Z ∈ Rn×n,Z> = Z, and Z 0}, the functions f1(X) = tr(XA), f2(X) = tr(XAX), and f3(X) = − log det(X) are geodesic convex onMbw.
Based on the result in Proposition 1, we also prove geodesic convexity of a reparameterized version of the Gaussian density estimation and mixture model onMbw (discussed in Section 4). Similar claims onMai can be found in [31]. We further show that monotonic functions on sorted eigenvalues are geodesic convex onMbw. This is an analogue of [65, Theorem 2.3] onMai. Proposition 2. Let λ↓ : Sn++ −→ Rn+ be the decreasingly sorted eigenvalue map and h : R+ −→ R be an increasing and convex function. Then f(X) = ∑k j=1 h(λ ↓ j (X)) for 1 ≤ k ≤ n is geodesic convex onMbw. Examples of such functions include f1(X) = tr(exp(X)) and f2(X) = tr(Xα), α ≥ 1.

heading: 4 Experiments
text: In this section, we compare the empirical performance of optimization algorithms under different Riemannian geometries for various problems. In addition to AI and BW, we also include the Log-Euclidean (LE) geometry [7] in our experiments.
The LE geometry explores the the linear space of symmetric matrices where the matrix exponential acts as a global diffeomorphism from the space to Sn++. The LE metric is defined as
〈U,V〉le = tr(DU log(X)DV log(X)) (3)
for any U,V ∈ TXM, where DU log(X) is the directional derivative of matrix logarithm at X along U. Following [69, 51, 58], for deriving various Riemannian optimization ingredients under the LE metric (3), we consider the parameterization X = exp(S), where S ∈ Sn, i.e., the space of n× n symmetric matrices. Equivalently, optimization on the SPD manifold with the LE metric is identified with optimization on Sn and the function of interest becomes f(exp(S)) for S ∈ Sn. While the Riemannian gradient can be computed efficiently by exploiting the directional derivative of the matrix exponential [4], deriving the Riemannian Hessian is tricky and we rely on finite-difference Hessian approximations [15].
We present convergence mainly in terms of the distance to the solution X∗ whenever applicable. The distance is measured in the Frobenius norm, i.e., ‖Xt−X∗‖F. When X∗ is not known, convergence is shown in the modified Euclidean gradient norm ‖Xt∇f(Xt)‖F. This is comparable across different metrics as the optimality condition X∗∇f(X∗) = 0 arises from problem structure itself [43]. We initialize the algorithms with the identity matrix for the AI and BW metrics and zero matrix for the LE metric (i.e., the matrix logarithm of the identity).
We mainly present the results on the Riemannian trust region (RTR) method, which is the method of choice for Riemannian optimization. Note for RTR, the results are shown against the cumulative sum of inner iterations (which are required to solve the trust region subproblem at every iteration). We also include the Riemannian steepest descent (RSD) and Riemannian stochastic gradient (RSGD) [14] methods for some examples. The experiments are conducted in Matlab using the Manopt toolbox [17] on a i5-10500 3.1GHz CPU processor.
In our extended report [27], we include additional experiments comparing convergence in objective function values for the three geometries. We also present results for the Riemannian conjugate gradient method, and results with different initializations (other than the identity and zero matrices) to further support our claims.
The code can be found at https://github.com/andyjm3/AI-vs-BW.
Weighted least squares. We first consider the weighted least squares problem with the symmetric positive definite constraint. The optimization problem is minX∈Sn++ f(X) = 1 2‖A X−B‖ 2 F, which is encountered in for example, SPD matrix completion [63] where A is a sparse matrix. The Euclidean gradient and Hessian are∇f(X) = (A X−B) A and∇2f(X)[U] = A U A, respectively. Hence, at optimal X∗, the Euclidean Hessian in matrix representation is H(X∗) = diag(vec(A A)). We experiment with two choices of A, i.e. A = 1n1>n (Dense) and A as a random sparse matrix (Sparse). The former choice for A leads to well-conditioned H(X∗) while the latter choice leads to an ill-conditioned H(X∗). Also note that when A = 1n1>n , κ ∗ ai = κ(X ∗)2 and κ∗bw = κ(X ∗). We generate X∗ as a SPD matrix with size n = 50 and exponentially decaying eigenvalues. We consider two cases with condition numbers κ(X∗) = 10 (LowCN) and 103 (HighCN). The matrix B is generated as B = A X∗. Figure 1 compares both RSD and RTR for different metrics. When A is either dense or sparse, convergence is significantly faster onMbw than bothMai andMle. The advantage of usingMbw becomes more prominent in the setting when condition number of X∗ is high. Figure 1(e) shows thatMbw is also superior in terms of runtime.
Lyapunov equations. Continuous Lyapunov matrix equation, AX + XA = C with X ∈ Sn++, is commonly employed in analyzing optimal control systems and differential equations [59, 45]. When A is stable, i.e., λi(A) > 0 and C ∈ Sn++, the solution X∗ 0 and is unique [44]. When C ∈ Sn+ and is low rank, X∗ ∈ Sn+ is also low rank. We optimize the following problem for solving the Lyapunov equation [73], i.e., minX∈Sn++ f(X) = tr(XAX)− tr(XC). The Euclidean gradient and Hessian are respectively ∇f(X) = AX + XA − C and ∇2f(X)[U] = AU + UA with H(X) = A⊕A. At optimal X∗, the condition number κ(H(X∗)) = κ(A). We experiment with two settings for the matrix A, i.e. A as the Laplace operator on the unit square where we generate 7 interior points so that n = 49 (Ex1), and A is a particular Toeplitz matrix with n = 50 (Ex2). The generated A matrices are ill-conditioned. The above settings correspond to Examples 7.1 and 7.3 in [45]. Under each setting, X∗ is set to be either full or low rank. The matrix C is generated as C = AX∗ + X∗A. The full rank X∗ is generated from the full-rank Wishart distribution while the low rank X∗ is a diagonal matrix with r = 10 ones and n− r zeros in the diagonal. We label the four cases as Ex1Full, Ex1Low, Ex2Full, and Ex2Low. The results are shown in Figures 2(a)-(d), where we observe that in all four cases, the BW geometry outperforms both AI and LE geometries.
Trace regression. We consider the regularization-free trace regression model [62] for estimating covariance and kernel matrices [61, 19]. The optimization problem is minX∈Sd++ f(X) =
1 2m ∑m i=1(yi − tr(A>i X))2, where Ai = aia>i , i = 1, ...,m are some rank-one measurement
matrices. Thus, we have∇f(X) = ∑m i=1(a > i Xai − yi)Ai and∇2f(X)[U] = ∑m i=1(a > i Uai)Ai. We create X∗ as a rank-r Wishart matrix and {Ai} as rank-one Wishart matrices and generate yi = tr(AiX
∗) + σ i with i ∼ N (0, 1), σ = 0.1. We consider two choices, (m, d, r) = (1000, 50, 50) and (1000, 50, 10), which are respectively labelled as SynFull and SynLow. From Figures 2(e)&(f), we also observe that convergence to the optimal solution is faster for the BW geometry.
Metric learning. Distance metric learning (DML) aims to learn a distance function from samples and a popular family of such distances is the Mahalanobis distance, i.e. dM(x,y) =√
(x− y)>M(x− y) for any x,y ∈ Rd. The distance is parameterized by a symmetric positive semi-definite matrix M. We refer readers to this survey [66] for more discussions on this topic. We particularly consider a logistic discriminant learning formulation [24]. Given a training sample {xi, yi}Ni=1, denote the link tij = 1 if yi = yj and tij = 0 otherwise. The objective is given by minM∈Sd++ f(M) = − ∑ i,j (tij log pij + (1− tij) log(1− pij)) , with pij = (1 + exp(dM(xi,xj))) −1. We can derive the matrix Hessian as H(M) = ∑ i,j pij(1 − pij)(xi − xj)(xi−xj)>⊗(xi−xj)(xi−xj)>. Notice κ(H(M∗)) depends on M∗ only through the constants pij . Thus, the condition number will not be much affected by κ(M∗).
We consider two real datasets, glass and phoneme, from the Keel database [5]. The number of classes is denoted as c. The statistics of these two datasets are (N, d, c) = (241, 9, 7) for glass
(5404, 5, 2) for phoneme. In Figures 2(g)&(h), we similarly see the advantage of using the BW metric compared to the other two metrics that behave similarly.
Log-det maximization. As discussed in Section 3.1, log-det optimization is one instance where κ∗bw ≥ κ∗ai. We first consider minimizing negative log-determinant along with a linear function as studied in [75]. That is, for some C ∈ Sn++, the objective is minX∈Sn++ f(X) = tr(XC)− log det(X). The Euclidean gradient and Hessian are given by ∇f(X) = C − X−1 and ∇2f(X)[U] = X−1UX−1. This problem is geodesic convex under both AI and BW metrics. We generate X∗ the same way as in the example of weighted least square with n = 50 and set C = (X∗)−1. We consider two cases with condition number cn = 10 (LowCN) and 103 (HighCN). As expeted, we observe faster convergence of AI and LE metrics over the BW metric in Figures 3(a)&(b). This is even more evident when the condition number increases.
Gaussian mixture model. Another notable example of log-det optimization is the Gaussian density estimation and mixture model problem. Following [31], we consider a reformulated problem on augmented samples y>i = [x > i ; 1], i = 1, ..., N where xi ∈ Rd are the original samples. The density is parameterized by the augmented covariance matrix Σ ∈ Rd+1. Notice that the log-likelihood of Gaussian is geodesic convex onMai, but not onMbw. We, therefore, define S = Σ−1 and the reparameterized log-likelihood is pN (Y; S) = ∑N i=1 log ( (2π)1−d/2 exp(1/2) det(S)1/2 exp(− 12y > i Syi) ) , which is now geodesic convex onMbw due to Proposition 1. Hence, we can solve the problem of Gaussian mixture model similar as in [31].
Here, we test on a dataset included in the MixEst package [30]. The dataset has 1580 samples in R2 with 3 Gaussian components. In Figure 3(c), we observe a similar pattern with RTR as in the log-det example. We also include performance of RSGD, which is often preferred for large scale problems. We set the batch size to be 50 and consider a decaying step size, with the best initialized step size shown in Figures 3(d)&(e). Following [8], the algorithms are initialized with kmeans++. We find that the AI geometry still maintains its advantage under the stochastic setting.

heading: 5 Conclusion and discussion
text: In this paper, we show that the less explored Bures-Wasserstein (BW) geometry for SPD matrices should often be the preferable choice than the Affine-Invariant geometry for optimization, particularly for learning ill-conditioned matrices. A systematic analysis shows that the BW metric preserves geodesic convexity of some popular cost functions and leads to better convergence rates in optimization.
We also theoretically discuss a ‘counter-example’ of log-det optimization where the Affine-Invariant (AI) geometry enjoys a better second-order conditioning and validate our findings empirically. This issue is addressed in our recent work [26], where we propose a generalized Bures-Wasserstein (GBW) geometry that is built on a generalization of the Lyapunov operator in the metric:
〈U,V〉gbw = 1
2 vec(U)>(X⊗M + M⊗X)−1vec(V), (4)
where U and V are symmetric matrices and M is a given SPD matrix. When M = I, the metric (4) reduces to the special BW metric (2). The use of the parameter M in (4) allows great flexibility for optimization algorithms. For one, choosing a particular M allows preconditioning the Hessian by locally approximating the AI geometry. This leads to improved convergence for log-det optimization.
Our comparisons between AI and BW geometries are based on optimization over generic cost functions. For specific problems, however, there may exist other alternative metrics that potentially work better. This is an interesting research direction to pursue. We also remark that optimization is not the only area where the AI and BW geometries can be compared. It would be useful to qualitatively compare the two metrics for other learning problems on SPD matrices, such as barycenter learning. In addition, kernel methods have been studied on the SPD manifold [42, 29, 78] that embed SPD matrices to a high dimensional feature space, known as the Reproducing Kernel Hilbert Space (RKHS). Such representations are used for subsequent learning tasks such as clustering or classification. But only a positive definite kernel provides a valid RKHS. We show in appendix Section A that the induced Gaussian kernel based on the BW distance is a positive definite kernel unlike the case for the AI metric. This difference highlights a potential advantage of the BW metric for representation learning on SPD matrices.

heading: A Positive definite kernel from BW geometry
text: In this section, we show the existence of a positive definite Gaussian kernel on Mbw. Existing works [55, 22] have studied the Wasserstein distance kernel. First, we present the definition of a positive (resp. negative) definite function as in [9]. Definition 5. Consider X to be a nonempty set. A function f : X ×X −→ R is called positive definite if and only if k is symmetric and for all integers m ≥ 2, {x1, ..., xm} ⊆ X and {c1, ..., cm} ⊆ R, it satisfies ∑m i,j=1 cicjf(xi, xj) ≥ 0. A function f is called negative definite if and only if under the
same conditions, it satisfies ∑m i,j=1 cicjf(xi, xj) ≤ 0 with ∑m i=1 ci = 0.
The following theorem shows the Gaussian kernel induced from the BW distance is positive definite on the SPD manifold. Theorem 3. The induced Gaussian kernel k(·, ·) := exp(−d2bw(·, ·)/(2σ2)) is positive definite.
Proof of Theorem 3. From Theorem 4.3 in [42], it suffices to prove the BW distance d2bw is negative definite. Indeed for any m ≥ 2, {X1, ...,Xm} ⊆ Mbw, {c1, ..., cm} ⊆ R with ∑m i=1 ci = 0, we have m∑
i,j=1
cicjd 2 bw(Xi,Xj)
= m∑ j=1 cj m∑ i=1 citr(Xi) + m∑ i=1 ci m∑ j=1 cjtr(Xj)− 2 m∑ i,j=1 cicjtr(X 1/2 i XjX 1/2 i ) 1/2
=− 2 m∑
i,j=1
cicjtr(X 1/2 i XjX 1/2 i ) 1/2 ≤ 0.
This shows d2bw is negative definite, and therefore, the exponentiated Gaussian kernel is positive definite.
Broader Impact
This paper is of a theoretical nature and does not foresee any immediate societal impacts.

Review summaries:
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. 

This paper studies the Bures-Wasserstein (BW) geometry of symmetric positive-definite (SPD) matrices and compares it to the commonly-used affine-invariant (AI) geometry in the context of Riemannian optimization. It is shown theoretically that the BW geometry has several advantages. Particularly, it is shown that (i) the BW metric is linearly dependent on the SPD matrix, (ii) the condition number of the Riemannian Hessian of the BW metric is smaller than that of the AI metric, and (iii) the curvature constant of the BW metric is smaller than that of the AI metric. Based on these properties, it is shown that the convergence of Riemannian steepest descent and Riemannian trust-region based on the BW metric is faster than based on the AI metric. The theoretical results are demonstrated on simulations of six problems: weighted least squares, Lyapunov equations, trace regression, metric learning, log-det maximization, and Gaussian mixture model.

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.

This paper investigates the advantages of the Bures-Wasserstein (BW) metric over the widely used Affine Invariant (AI) one, for the manifold of symmetric positive definite (SPD) matrices. Rewriting the two metrics in a compatible manner reveals that for a wide variety of objectives, the BW metric is a linear function of the input matrix whereas the AI has a quadratic dependence. One exception is the log-det function which is included in the analysis. The paper experimentally validates that for ill-conditioned problems BW is a better choice compared to AI.