---
title: Bures-Wasserstein Gradient Descent
url: https://www.emergentmind.com/topics/bures-wasserstein-gradient-descent-bwgd
type: topic
---

# Bures-Wasserstein Gradient Descent

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 [2106.08502, 2310.20090, 2303.03027].

## 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 [2106.08502].

A more intrinsic formulation works on \(\mathbb P_d\), the cone of \(d\times d\) positive definite matrices, with tangent space \(T_P\mathbb P_d\cong \mathbb H_d\). The BW metric at \(P\) is
\[
\langle U,V\rangle_P^{\mathrm{BW}}=\operatorname{Tr}\!\left[\mathcal L_P(U)\,P\,\mathcal L_P(V)\right],
\]
where \(\mathcal L_P(U)\) is the unique solution \(X\) of the Lyapunov equation
\[
XP+PX=U.
\]
In this formulation, the exponential and logarithmic maps are available in closed form:
\[
\operatorname{Exp}_P^{\mathrm{BW}}[V]=(\mathbb I+\mathcal L_P(V))\,P\,(\mathbb I+\mathcal L_P(V)),
\]
\[
\operatorname{Log}_P^{\mathrm{BW}}[Q]=\mathcal L_P^{-1}\!\left(P^{-1}\#Q-\mathbb I\right), \qquad \mathcal L_P^{-1}(Y)=YP+PY.
\]
The generalized squared Bures distance at a base point \(R\),
\[
\operatorname{B}_R(P,Q)=\|\operatorname{Log}_R[P]-\operatorname{Log}_R[Q]\|_R^2,
\]
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 \(\operatorname{Exp}_R\) [2410.04937].

## 2. Canonical BWGD update rules

The most classical BWGD task is Gaussian barycenter computation. For a probability law \(P\) over centered Gaussians or covariance matrices, the objective is
\[
F(\Sigma)=\frac12\int W_2^2(\Sigma,\cdot)\,dP.
\]
In Wasserstein space, the generic gradient step is
\[
\mu^+=\exp_\mu(-\eta\nabla F(\mu))=[I-\eta\nabla F(\mu)]_\#\mu,
\]
and for Gaussian covariances this becomes the explicit matrix iteration
\[
S_t\gets (1-\eta)I_d+\eta\int \mathrm{GM}(\Sigma_{t-1}^{-1},\Sigma)\,dP(\Sigma),\qquad
\Sigma_t\gets S_t\,\Sigma_{t-1}\,S_t.
\]
Its stochastic counterpart replaces the integral by a sample \(K_t\sim P\):
\[
\hat S_t\gets (1-\eta_t)I_d+\eta_t\,\mathrm{GM}(\Sigma_{t-1}^{-1},K_t),\qquad
\Sigma_t\gets \hat S_t\,\Sigma_{t-1}\,\hat S_t.
\]
This is the algorithmic form often called Bures-Wasserstein gradient descent or Bures-SGD in the barycenter literature [2106.08502, 2001.01700].

For composite objectives arising in Gaussian variational inference, a forward-backward Euler splitting is also used on the BW manifold. Writing
\[
\mathcal F(\mu)=\mathcal E_V(\mu)+\mathscr H(\mu),
\]
with \(\mathcal E_V(\mu)=\int V(x)\,d\mu(x)\) and \(\mathscr H(\mu)=\int \log(\mu(x))\,d\mu(x)\), the scheme consists of a forward step
\[
\mu_{k+\frac12}=(I-\eta \nabla_{\BW}\mathcal E_V(\mu_k))_\#\mu_k
\]
followed by the BW proximal step
\[
\mu_{k+1}= \arg\min_{\mu\in\BW(\mathbb R^d)} \left\{ \mathscr H(\mu)+\frac{1}{2\eta}W_2^2(\mu,\mu_{k+\frac12}) \right\}.
\]
For Gaussian \(\mu_{k+\frac12}=\mathcal N(m_{k+\frac12},\Sigma_{k+\frac12})\), the backward step is closed form:
\[
m_{k+1}=m_{k+\frac12},
\]
\[
\Sigma_{k+1}=\frac12\left(\Sigma_{k+\frac12}+2\eta I+\big[\Sigma_{k+\frac12}(\Sigma_{k+\frac12}+4\eta I)\big]^{1/2}\right).
\]
This closed-form backward step is one reason BW methods are practical for Gaussian KL objectives [2410.02490].

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 \(G^{(k)}\) from the exact manifold derivative and updates
\[
\Sigma_\beta^{(k)}=G^{(k)}\Sigma_\beta^{(k-1)}G^{(k)},
\]
while the Hybrid Gradient Descent method alternates a closed-form SUOT denoising step with a standard BW barycenter update
\[
\Sigma_\beta^{(k)}=S^{(k)}\Sigma_\beta^{(k-1)}S^{(k)}.
\]
These algorithms remain intrinsic to the Bures-Wasserstein manifold and avoid Euclidean updates that would leave \(\mathbb S_{++}^d\) [2410.08117].

## 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 \(q_\phi=\mathcal N(\mu,\Sigma)\) with factorization \(\Sigma=BB^\top\), the variational objective is
\[
\mathcal L(q_\phi)=\mathrm{KL}(q_\phi\|p),\qquad
F(\phi)=\mathbb E_{q_\phi}[\log q_\phi(x)-\log p(x)],
\]
with reparameterization
\[
x=T_\phi(z)=\mu+Bz,\qquad z\sim\mathcal N(0,I).
\]
In this setting, the paper shows that the Bures-Wasserstein gradient flow can be recast as the Euclidean flow
\[
\frac{d\mu_t}{dt}=-\nabla_\mu F(\mu_t,B_t),\qquad
\frac{dB_t}{dt}=-\nabla_B F(\mu_t,B_t),
\]
and that the forward Euler discretization
\[
\phi^{k+1}=\phi^k-\eta\nabla_\phi F(\phi^k)
\]
is exactly the standard black-box variational inference update when gradients are estimated by the path-derivative estimator. In this interpretation, the pathwise gradient
\[
\nabla_\phi F(\phi)=\mathbb E_{z}\left[\nabla_\phi\big(\log q_\phi(T_\phi(z))-\log p(T_\phi(z))\big)\right]
\]
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 \(f\)-divergences and non-Gaussian variational families [2310.20090].

Subsequent BW variational-inference work retains the same Gaussian manifold but changes either the objective or the estimator. For KL minimization against \(\pi(x)\propto e^{-V(x)}\), the BW gradient of the potential energy at \(\mu=\mathcal N(m,\Sigma)\) is the affine map
\[
\nabla_{\BW}\mathcal E_V(\mu):x\mapsto \mathbb E_{\mu}\nabla V+(\mathbb E_{\mu}\nabla^2V)(x-m).
\]
Because \(\mathbb E_{\mu}\nabla V\) and \(\mathbb E_{\mu}\nabla^2V\) 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
\[
\widetilde b_k=\nabla V(X_k)-c_k\,\Sigma_k^{-1}(X_k-m_k),
\]
while keeping \(S_k=\nabla^2V(X_k)\), thereby reducing variance within the same forward-backward BW scheme [2410.02490].

The importance-weighted setting produces a different but closely related BWGD. In Gaussian VI with \(q=\mathcal N(m,\Sigma)\), the BW tangent space consists of affine maps
\[
\mathcal T_q \operatorname{BW}(R^d)=\Big\{ x\mapsto a+S(x-m)\mid a\in R^d,\ S\in \bm S^d\Big\},
\]
and the Wasserstein gradient of the IW-ELBO is projected onto this tangent space to obtain a BW gradient of the form
\[
\nabla^{\mathrm{BW}}[-\mathcal F_K(q)](z)=a_*+S_*(z-m).
\]
The induced update is
\[
m_{k+1}=m_k-\eta a_*,\qquad
\Sigma_{k+1}=(I-\eta S_*)\Sigma_k(I-\eta S_*).
\]
A distinctive result in this line is a signal-to-noise analysis: while the Euclidean gradient estimator for IW-ELBO has SNR \(O(\sqrt{M}/\sqrt{K})\), the Wasserstein gradient estimator has SNR \(\Omega(\sqrt{MK})\). The same favorable scaling is extended to the Variational Rényi Importance-Weighted Autoencoder bound [2602.04272].

## 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 \(\zeta\)-regularity assumption on the Gaussian support, these results yield a linear rate for gradient descent and an \(O(1/n)\) rate for stochastic gradient descent [2001.01700].

A later analysis strengthened this theory by proving dimension-free convergence for Riemannian gradient descent on Gaussian barycenters. Under spectral bounds \([\lambda_{\min},\lambda_{\max}]\) and condition number \(\kappa=\lambda_{\max}/\lambda_{\min}\), 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 [2106.08502].

Robust barycenters on the Bures manifold admit a similar theory. For the SUOT-based Gaussian barycenter, if the iterates remain in the spectral set
\[
\mathcal K_{[1/\rho,\rho]}=\left\{\Sigma\in\mathbb S_{++}^d:\frac1\rho\le \lambda_i(\Sigma)\le \rho,\ \forall i\right\},
\]
the Exact Geodesic Gradient Descent algorithm converges to an optimal solution with a dimension-free geometric rate
\[
\mathcal D(\Sigma_\beta^{(k)})\le \left(1-\frac{8\tau^2\eta(1-\eta/2)}{\rho(\rho^2+2\tau\rho)^{3/2}}\right)^k \mathcal D(\Sigma_\beta^{(0)}),
\]
while the Hybrid algorithm is also proven convergent under the same spectral boundedness assumptions [2410.08117].

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 [2106.08502].

## 5. Specialized objective classes

Beyond averaging, BWGD has been analyzed for generative covariance models. In a deep linear generator with latent Gaussian \(Z\sim\mathcal N(0,I_m)\) and output \(X=WZ\sim\mathcal N(0,WW^\top)\), the BW loss is
\[
L^1(W)=\operatorname{tr}\!\left(WW^\top+\Sigma_0-2(\Sigma_0^{1/2}WW^\top\Sigma_0^{1/2})^{1/2}\right).
\]
The paper characterizes rank-\(k\) critical points explicitly: if \(\Sigma_0=\Omega\Lambda\Omega^\top\) has distinct positive eigenvalues, then the critical points of \(L^1|_{M(k)}\) are
\[
W^*=\Omega_{J_k}\,\bar\Lambda_{J_k}^{1/2}\,V^\top.
\]
The global minimizer selects the top \(k\) eigenvalues, all non-minimizing critical points are strict saddles, and the optimal rank-\(k\) 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 \(L_\tau^1\) and convergence analyses under balanced initialization and a modified deficiency margin [2303.03027].

In rank-one matrix recovery and real phase retrieval, BWGD has a different interpretation. For whitened measurements satisfying
\[
\frac1n\sum_{i=1}^n a_i a_i^T=I,
\]
the amplitude loss
\[
F(u)=\frac{1}{2n}\sum_{i=1}^n (|a_i^Tu|-y_i)^2
\]
has Hessian \(\nabla^2F(u)=I\) at differentiable points, and the BWGD iteration
\[
u_{t+1}=(1-\eta)u_t+\frac{\eta}{n}\sum_{i=1}^n \frac{y_i\,a_i a_i^T u_t}{|a_i^T u_t|}
\]
coincides with Newton’s method when \(\eta=1\). This explains the empirical superlinear convergence reported earlier, but it also exposes an instability: the objective is nonsmooth when some \(a_i^Tu=0\). A smoothing framework based on
\[
F_\epsilon(u)=\frac{1}{2n}\sum_{i=1}^n \left(\sqrt{(a_i^Tu)^2+\epsilon\|a_i\|^2}-\sqrt{y_i^2+\epsilon\|a_i\|^2}\right)^2
\]
leads to BWGD with Dynamic Smoothing and local superlinear convergence guarantees [2507.23017].

Fréchet regression on the BW manifold introduces signed barycenter objectives,
\[
F(S)=\sum_{k=1}^n \lambda_k\,W_2^2(S,\Sigma_k),
\]
with affine weights that may be negative in extrapolation regimes. A sufficient existence condition is the Spectral Dominance of Positive Weights,
\[
\sum_{i\in\mathcal I}\lambda_i^+\sqrt{\lambda_{\min}(\Sigma_i)}>
\sum_{j\in\mathcal J}\lambda_j^-\sqrt{\lambda_{\max}(\Sigma_j)}.
\]
Under this condition, the objective has no local maxima, and the projection-free BWGD update
\[
\tilde S_t=(1-\eta)I+\eta\sum_{k=1}^n \lambda_k\,\mathrm{GM}(S_{t-1}^{-1},\Sigma_k),\qquad
S_t=\tilde S_t\,S_{t-1}\,\tilde S_t
\]
stays inside \(\mathbb S_{++}^d\). With \(\eta\le 1/\sum_k|\lambda_k|\), the method has an \(O(1/T)\) stationarity guarantee, and a pairwise stochastic reformulation supports off-the-shelf stochastic Riemannian optimization methods [2604.03566].

## 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
\[
\operatorname{B}_R(P,Q)=\|\operatorname{Log}_R[P]-\operatorname{Log}_R[Q]\|_R^2,
\]
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 [2410.04937].

Another extension changes the geometry itself. The generalized Bures-Wasserstein (GBW) geometry introduces a parameter \(\mathbf M\in\mathcal S_{++}^n\), with
\[
d_{\rm gbw}(\mathbf X,\mathbf Y)=d_{\rm bw}\big(\mathbf M^{-1/2}\mathbf X\mathbf M^{-1/2},\mathbf M^{-1/2}\mathbf Y\mathbf M^{-1/2}\big),
\]
and recovers standard BW when \(\mathbf M=\mathbf I\). The paper derives the associated gradient, Hessian, geodesics, and exp/log maps, and argues that choosing \(\mathbf M\) appropriately can improve conditioning relative to the standard BW metric. This is a generalization of BWGD rather than BWGD itself [2110.10464].

Acceleration methods treat BWGD as a fixed-point iteration. In the Gaussian setting, if
\[
G(\rho)=\operatorname{Exp}_\rho(-h\,\partial_W E(\rho)),
\]
then plain BWGD is the unaccelerated fixed-point iteration \(\rho\mapsto G(\rho)\). Riemannian Anderson Mixing (BWRAM) augments this with transported residual histories and a least-squares correction,
\[
x_{k+1}=\operatorname{Exp}_{x_k}\!\big(-X_k\Gamma_k+\beta_k\bar r_k\big),
\]
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 [2601.22038].

Alternative computational paradigms sometimes replace intrinsic BWGD altogether. One line reformulates the BW metric as a convex semidefinite program,
\[
\rho^2(A,B)=\min \ \operatorname{Tr}(A)+\operatorname{Tr}(B)-2\operatorname{Tr}\!\left(\sqrt{A}U\right)
\quad\text{subject to}\quad
\begin{bmatrix} B & U^\top \\ U & I \end{bmatrix}\succeq 0,
\]
and uses this to compute BW barycenters, distances between convex subsets, and BW distance constraints via standard SDP machinery [2303.03883]. Another line, motivated by covariance alignment, proposes ITSPACE, a proximal majorization-minimization method on a square-root factorization \(X=YY^\top\) for the exact BW objective
\[
F(Y)=\|Y\|_F^2+\operatorname{tr}(\Sigma_\star)-2\|\Sigma_\star^{1/2}Y\|_\ast,
\]
with closed-form update
\[
Y_{k+1}=\alpha\,\Sigma_\star^{1/2}\,\mathrm{polar}(\Sigma_\star^{1/2}Y_k)+(1-\alpha)Y_k.
\]
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 [2606.30523].

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.

Source: https://www.emergentmind.com/topics/bures-wasserstein-gradient-descent-bwgd