Bures-Wasserstein Gradient Descent

• BWGD is a gradient-based optimization method on the manifold of Gaussian covariances using the Bures-Wasserstein metric and closed-form transport maps.
• It supports various applications such as barycenter computation, variational inference, and phase retrieval by leveraging intrinsic manifold retractions.
• BWGD avoids Euclidean limitations by ensuring structure-preserving updates with convergence guarantees under smoothness and Polyak-Łojasiewicz conditions.

Bures-Wasserstein Gradient Descent (BWGD) denotes gradient-based optimization on the manifold of Gaussian measures or positive definite covariance matrices endowed with the Bures-Wasserstein geometry, i.e. the Gaussian specialization of the $2$-Wasserstein metric. In the literature, the term is used in several closely related senses: as Riemannian gradient descent for Gaussian barycenters, as the explicit Euler discretization of a Bures-Wasserstein gradient flow in Gaussian variational inference, and as a baseline geometry-aware optimizer for symmetric positive definite (SPD) objectives such as covariance matching, phase retrieval, and conditional barycenters (Altschuler et al., 2021, Yi et al., 2023, Bréchet et al., 2023).

1. Geometric setting on Gaussian and SPD manifolds

For centered Gaussians, a distribution $\mathcal N(0,\Sigma)$ is identified with its covariance matrix $\Sigma$, and the relevant state space is the cone of positive semidefinite or positive definite matrices. The squared Bures-Wasserstein distance between covariances is

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$

which is also the closed form of the Gaussian $2$-Wasserstein distance. In barycenter and transport formulations, the optimal transport map from $\Sigma$ to $\Sigma'$ is

$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$

and the Bures-Wasserstein geodesic is

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$

These formulas make BWGD a transport-based method: descent follows the geodesics induced by optimal transport rather than Euclidean straight lines (Altschuler et al., 2021).

A more intrinsic formulation works on $\mathbb P_d$, the cone of $\mathcal N(0,\Sigma)$0 positive definite matrices, with tangent space $\mathcal N(0,\Sigma)$1. The BW metric at $\mathcal N(0,\Sigma)$2 is

$\mathcal N(0,\Sigma)$3

where $\mathcal N(0,\Sigma)$4 is the unique solution $\mathcal N(0,\Sigma)$5 of the Lyapunov equation

$\mathcal N(0,\Sigma)$6

In this formulation, the exponential and logarithmic maps are available in closed form: $\mathcal N(0,\Sigma)$7

$\mathcal N(0,\Sigma)$8

The generalized squared Bures distance at a base point $\mathcal N(0,\Sigma)$9,

$\Sigma$0

shows that BW geometry can be linearized in tangent space, so descent on the manifold can be interpreted as Euclidean steps after log-mapping, followed by retraction with $\Sigma$1 (Afham et al., 2024).

2. Canonical BWGD update rules

The most classical BWGD task is Gaussian barycenter computation. For a probability law $\Sigma$2 over centered Gaussians or covariance matrices, the objective is

$\Sigma$3

In Wasserstein space, the generic gradient step is

$\Sigma$4

and for Gaussian covariances this becomes the explicit matrix iteration

$\Sigma$5

Its stochastic counterpart replaces the integral by a sample $\Sigma$6: $\Sigma$7 This is the algorithmic form often called Bures-Wasserstein gradient descent or Bures-SGD in the barycenter literature (Altschuler et al., 2021, Chewi et al., 2020).

For composite objectives arising in Gaussian variational inference, a forward-backward Euler splitting is also used on the BW manifold. Writing

$\Sigma$8

with $\Sigma$9 and $$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$0, the scheme consists of a forward step

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$1

followed by the BW proximal step

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$2

For Gaussian $$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$3, the backward step is closed form: $$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$4

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$5

This closed-form backward step is one reason BW methods are practical for Gaussian KL objectives (Luu et al., 2024).

Robust barycenter problems on the Bures manifold use the same multiplicative transport structure. In the Semi-Unbalanced Optimal Transport (SUOT) Gaussian setting, the Exact Geodesic Gradient Descent method forms a matrix $$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$6 from the exact manifold derivative and updates

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$7

while the Hybrid Gradient Descent method alternates a closed-form SUOT denoising step with a standard BW barycenter update

$$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$8

These algorithms remain intrinsic to the Bures-Wasserstein manifold and avoid Euclidean updates that would leave $$W_2^2(\Sigma,\Sigma')=\operatorname{tr}\!\Bigl[\Sigma+\Sigma'-2(\Sigma^{1/2}\Sigma'\Sigma^{1/2})^{1/2}\Bigr],$$9 (Nguyen et al., 2024).

3. Variational inference and gradient-flow interpretations

A central development is the identification of Gaussian black-box variational inference with a discretized Bures-Wasserstein gradient flow. For a Gaussian variational family $2$0 with factorization $2$1, the variational objective is

$2$2

with reparameterization

$2$3

In this setting, the paper shows that the Bures-Wasserstein gradient flow can be recast as the Euclidean flow

$2$4

and that the forward Euler discretization

$2$5

is exactly the standard black-box variational inference update when gradients are estimated by the path-derivative estimator. In this interpretation, the pathwise gradient

$2$6

is not merely a Monte Carlo device but the tangent vector field of the Bures-Wasserstein flow. The paper explicitly frames this as a distillation of the Wasserstein gradient flow into parameter space and states that the viewpoint extends to $2$7-divergences and non-Gaussian variational families (Yi et al., 2023).

Subsequent BW variational-inference work retains the same Gaussian manifold but changes either the objective or the estimator. For KL minimization against $2$8, the BW gradient of the potential energy at $2$9 is the affine map

$\Sigma$0

Because $\Sigma$1 and $\Sigma$2 are usually intractable, prior BW methods used a single-sample Monte Carlo estimator. A later paper replaces the mean-gradient estimator by the control-variate estimator

$\Sigma$3

while keeping $\Sigma$4, thereby reducing variance within the same forward-backward BW scheme (Luu et al., 2024).

The importance-weighted setting produces a different but closely related BWGD. In Gaussian VI with $\Sigma$5, the BW tangent space consists of affine maps

$\Sigma$6

and the Wasserstein gradient of the IW-ELBO is projected onto this tangent space to obtain a BW gradient of the form

$\Sigma$7

The induced update is

$\Sigma$8

A distinctive result in this line is a signal-to-noise analysis: while the Euclidean gradient estimator for IW-ELBO has SNR $\Sigma$9, the Wasserstein gradient estimator has SNR $\Sigma'$0. The same favorable scaling is extended to the Variational Rényi Importance-Weighted Autoencoder bound (Jiang et al., 4 Feb 2026).

4. Nonconvexity and convergence theory

A recurring theme in BWGD is that good convergence behavior does not require geodesic convexity. The Gaussian barycenter functional is not geodesically convex on the Bures-Wasserstein manifold, and can even be concave along geodesics. The first global rates for BWGD and BW-SGD in this setting were obtained by replacing convexity arguments with a smoothness inequality and a Polyak-Łojasiewicz inequality. Under a uniform $\Sigma'$1-regularity assumption on the Gaussian support, these results yield a linear rate for gradient descent and an $\Sigma'$2 rate for stochastic gradient descent (Chewi et al., 2020).

A later analysis strengthened this theory by proving dimension-free convergence for Riemannian gradient descent on Gaussian barycenters. Under spectral bounds $\Sigma'$3 and condition number $\Sigma'$4, BWGD and BW-SGD attain rates independent of the ambient dimension. The same framework gives dimension-free guarantees for the entropically regularized barycenter and the first stationarity guarantee for the geometric median. This work also emphasizes a practical contrast: Euclidean GD/SGD on the same barycenter objective requires projection onto spectral constraints and has substantially worse condition-number dependence (Altschuler et al., 2021).

Robust barycenters on the Bures manifold admit a similar theory. For the SUOT-based Gaussian barycenter, if the iterates remain in the spectral set

$\Sigma'$5

the Exact Geodesic Gradient Descent algorithm converges to an optimal solution with a dimension-free geometric rate

$\Sigma'$6

while the Hybrid algorithm is also proven convergent under the same spectral boundedness assumptions (Nguyen et al., 2024).

These results coexist with explicit caveats. The barycenter functional is not geodesically convex; the geometric median remains globally nonconvex even after smoothing; and some earlier stronger convexity claims were corrected in subsequent analysis. BWGD theory therefore relies on more specific structural properties—PL inequalities, smoothness on spectrally bounded sets, or trapping arguments—rather than on generic convex manifold optimization (Altschuler et al., 2021).

5. Specialized objective classes

Beyond averaging, BWGD has been analyzed for generative covariance models. In a deep linear generator with latent Gaussian $\Sigma'$7 and output $\Sigma'$8, the BW loss is

$\Sigma'$9

The paper characterizes rank-$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$0 critical points explicitly: if $$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$1 has distinct positive eigenvalues, then the critical points of $$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$2 are

$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$3

The global minimizer selects the top $$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$4 eigenvalues, all non-minimizing critical points are strict saddles, and the optimal rank-$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$5 approximation error is the sum of the discarded eigenvalues. At the same time, the Hessian can blow up near rank-deficient matrices, motivating a smoothed loss $$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$6 and convergence analyses under balanced initialization and a modified deficiency margin (Bréchet et al., 2023).

In rank-one matrix recovery and real phase retrieval, BWGD has a different interpretation. For whitened measurements satisfying

$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$7

the amplitude loss

$$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$8

has Hessian $$T_{\Sigma\to\Sigma'}=\Sigma^{-1/2}\bigl(\Sigma^{1/2}\Sigma'\Sigma^{1/2}\bigr)^{1/2}\Sigma^{-1/2}=\mathrm{GM}(\Sigma^{-1},\Sigma'),$$9 at differentiable points, and the BWGD iteration

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$0

coincides with Newton’s method when $$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$1. This explains the empirical superlinear convergence reported earlier, but it also exposes an instability: the objective is nonsmooth when some $$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$2. A smoothing framework based on

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$3

leads to BWGD with Dynamic Smoothing and local superlinear convergence guarantees (Maunu et al., 30 Jul 2025).

Fréchet regression on the BW manifold introduces signed barycenter objectives,

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$4

with affine weights that may be negative in extrapolation regimes. A sufficient existence condition is the Spectral Dominance of Positive Weights,

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$5

Under this condition, the objective has no local maxima, and the projection-free BWGD update

$$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$6

stays inside $$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$7. With $$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$8, the method has an $$\Sigma_t=\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr)\Sigma\bigl((1-t)I_d+tT_{\Sigma\to\Sigma'}\bigr),\qquad t\in[0,1].$$9 stationarity guarantee, and a pairwise stochastic reformulation supports off-the-shelf stochastic Riemannian optimization methods (Nguyen et al., 4 Apr 2026).

6. Extensions, accelerations, and alternative formulations

Some work surrounding BWGD develops geometric machinery rather than a standalone descent algorithm. A 2024 paper on generalized quantum fidelities reinterprets Uhlmann, Holevo, and Matsumoto fidelities as base-dependent linearizations of the Bures-Wasserstein manifold. Its most direct optimization consequence is the identity

$\mathbb P_d$0

which provides exactly the tangent-space norm, log map, and exponential map needed for BWGD implementations on density matrices or covariance matrices. The paper explicitly notes that it does not derive the BW gradient of a particular loss, but it supplies the geometric primitives required to do so (Afham et al., 2024).

Another extension changes the geometry itself. The generalized Bures-Wasserstein (GBW) geometry introduces a parameter $\mathbb P_d$1, with

$\mathbb P_d$2

and recovers standard BW when $\mathbb P_d$3. The paper derives the associated gradient, Hessian, geodesics, and exp/log maps, and argues that choosing $\mathbb P_d$4 appropriately can improve conditioning relative to the standard BW metric. This is a generalization of BWGD rather than BWGD itself (Han et al., 2021).

Acceleration methods treat BWGD as a fixed-point iteration. In the Gaussian setting, if

$\mathbb P_d$5

then plain BWGD is the unaccelerated fixed-point iteration $\mathbb P_d$6. Riemannian Anderson Mixing (BWRAM) augments this with transported residual histories and a least-squares correction,

$\mathbb P_d$7

and local convergence is proven in sufficiently small Bures-Wasserstein balls around a nondegenerate covariance. Numerically, BWRAM is reported to provide significant acceleration over Picard iterations and performance on par with or better than Riemannian gradient descent and conjugate gradient baselines (Aksenov et al., 29 Jan 2026).

Alternative computational paradigms sometimes replace intrinsic BWGD altogether. One line reformulates the BW metric as a convex semidefinite program,

$\mathbb P_d$8

and uses this to compute BW barycenters, distances between convex subsets, and BW distance constraints via standard SDP machinery (Mohan, 2023). Another line, motivated by covariance alignment, proposes ITSPACE, a proximal majorization-minimization method on a square-root factorization $\mathbb P_d$9 for the exact BW objective

$\mathcal N(0,\Sigma)$00

with closed-form update

$\mathcal N(0,\Sigma)$01

ITSPACE is explicitly contrasted with BW-gradient descent: it has a sufficient-decrease inequality in exact arithmetic, handles inexact polar computation through a certificate-gap bound, and preserves PSD structure and rank by construction (Na et al., 29 Jun 2026).

Taken together, these developments show that BWGD is best understood not as a single fixed algorithm but as a family of geometry-aware first-order methods defined by the Bures-Wasserstein metric, explicit transport maps, and manifold retractions. What changes across the literature is the objective—barycenter, KL, IW-ELBO, SUOT, BW loss, signed Fréchet objective—or the surrounding numerical machinery, while the common principle is optimization in a geometry aligned with Gaussian optimal transport.