PaperlayerView exact revision on arXiv

Source-linked AI summary

Projected Riemannian Gradient Descent for the Bures-Wasserstein Barycenter: Dimension-Independent Linear Convergence at Unit Step Size

A. Afham

arXiv:2609.03762v1cs.LGmath.OCquant-ph

TL;DR

Unit-step RGD is used in practice, but existing analyses offer either dimension-dependent guarantees or dimension-independent guarantees at slower small step sizes. The paper composes unit-step RGD with an exact, non-expansive projection onto a well-conditioned spectral interval. This yields dimension-independent linear convergence for the barycenter and invariant matrix projection problems, without materially increasing per-step cost.

  • Problem

    Existing unit-step RGD guarantees can depend on dimension, whereas dimension-free guarantees require small step sizes that cost two to three orders of magnitude more iterations on typical instances.

    Source: §1 · pp. 2–3

  • Method

    Projected BW-GD composes unit-step RGD with eigenvalue clipping, the exact non-expansive BW projection onto a well-conditioned spectral interval.

    Source: §1 · p. 4

  • Results

    The method converges at rate (1 − κ^-3/2) for both problems, improving the dimension-free iteration complexity from O(κ^5/2 log(1/ϵ)) to κ^3/2 log(1/ϵ).

    Source: §1 · p. 4, §6 · p. 20

  • Takeaways & Limitations

    The guarantee extends to the invariant matrix projection problem because its fixed-point algorithm is unit-step RGD on a totally geodesic invariant submanifold.

    Source: §1 · p. 4, §6 · p. 21

    The observed convergence rates are much faster than the proved bounds, and whether the generic-case rate admits a sharper exponent remains open.

    Source: §6 · p. 21

Abstract

from arXiv · show

The computation of the Bures-Wasserstein (BW) barycenter of an ensemble of positive definite matrices arises throughout machine learning, optimal transport, and quantum information. Riemannian gradient descent (RGD) at unit step size -- the fixed-point iteration used in practice -- converges rapidly, yet existing analyses present a dichotomy: unit-step guarantees carry worst-case exponential dependence on the dimension, while dimension-independent guarantees require small step sizes that forfeit the empirical speed. We resolve this dichotomy, not by improving the guarantees for unit-step RGD, but by proposing a Projected RGD algorithm that achieves dimension-independent linear convergence at unit step size. The achieved rate, $(1 - \kappa^{-3/2})$, where $\kappa$ is the condition number of the ensemble, also polynomially improves on the best small-step guarantee ($\kappa^{3/2}$ versus $\kappa^{5/2}$ iteration complexity). The crux is a novel Projection Lemma: clipping the eigenvalues of a positive matrix to an interval $[\alpha, \beta]$ is the closed-form, non-expansive (1-Lipschitz) BW-metric projection onto the set $\{S : \alpha I \leq S \leq \beta I\}$ -- a statement which, unlike its known one-sided counterpart, does not follow from convexity. The projection is moreover free: it reuses an eigendecomposition the next iteration must perform in any case, so the projected and unprojected iterations cost the same per step. The same analysis covers the invariant matrix projection problem of Brahmachari et al. (2025), whose fixed-point algorithm we identify as unit-step RGD on a totally geodesic submanifold, thereby extending the dimension-independent guarantee to that setting verbatim.

1 Introduction

Unit-step RGD is fast in practice, but its existing guarantees either depend on dimension or require small steps that can cost substantially more iterations. Projected BW-GD restores dimension-independent convergence at unit step size by keeping iterates in a well-conditioned spectral interval.

1 Introduction

  • Existing unit-step analyses can have worst-case κ^d dependence, while dimension-free guarantees require η = O(1/κ) [ACGS21].On typical instances, those small steps cost two to three orders of magnitude more iterations.Source: §1 · pp. 2–3
  • Unit-step iterates can transiently fall below the spectral floor α, so the well-conditioned set is not invariant even when the ensemble and optimum lie inside it.The exit is temporary; the cited example returns to the set after two steps.Source: §1 · p. 3
  • The invariant matrix projection fixed-point update is unit-step BW gradient descent on a totally geodesic invariant submanifold, extending the convergence guarantee to that problem.The invariant submanifold is also closed under clipping.Source: §1 · p. 4

1.2 Our contributions

  • Projected BW-GD converges at rate (1 − κ^-3/2) for both problems, improving the best dimension-free iteration complexity from O(κ^5/2 log(1/ϵ)) to κ^3/2 log(1/ϵ).The guarantee is dimension-independent and applies at unit step size.Source: §1 · p. 4, §1.2 · p. 6

1 Introduction

  • Clipping reuses an eigendecomposition needed by the next iteration, so projected and unprojected steps cost essentially the same per step.The remaining O(d) difference is dominated by O(d^3) eigendecompositions.Source: §1 · pp. 3–4

1.2 Our contributions

  • Clipping eigenvalues to [α, β] is the exact, non-expansive BW-metric projection onto [αI, βI], including the lower-bound clipping not established by the known one-sided argument.The two-sided result uses a new argument rather than convexity.Source: §1.2 · pp. 6, 8
  • The analysis identifies BW-GD with Euclidean gradient descent on the square-root manifold, where strong convexity yields a PL inequality and the descent lemma combines with projection.The framework also identifies the projected iteration upstairs as singular-value-clipped Euclidean gradient descent.Source: §1.2 · p. 7
  • The projection keeps iterates where the dimension-free PL inequality holds and, because it is non-expansive, does not undo the gradient step’s progress.References within [αI, βI] are fixed by the projection.Source: §1.2 · pp. 6, 8

2 Mathematical Preliminaries

The preliminaries set the BW geometry on positive definite matrices and collect the quotient, distance, and invariant-subspace tools used later.

Matrix notation and geometric mean

  • The matrix geometric mean is the unique positive definite solution of a Riccati equation and is symmetric and operator monotone in each argument.It also satisfies A#(cI) = √cA for c > 0.Source: §2 · p. 8

The Bures–Wasserstein manifold

  • The BW metric on positive definite matrices is defined through the Lyapunov solution that maps a tangent vector to its metric representation.The resulting metric induces the squared BW distance.Source: §2 · p. 8
  • The square-root map identifies the positive definite cone with invertible matrices modulo unitary transformations, and the BW distance is the distance between their fibers.For full-rank matrices, the squared distance also has a closed form in terms of fidelity.Source: §2 · p. 9

Further properties of the BW distance

  • The squared BW distance to a fixed reference is strongly convex on a bounded spectral region, supporting the later convergence analysis.The preliminaries also note its data processing inequality for quantum channels.Source: §2 · p. 9

Invariant matrix projection

  • For invariant matrices, the twirl is the Hilbert–Schmidt orthogonal projection onto the commutant, whose positive definite elements form the invariant submanifold.Its tangent space is the Hermitian part of the commutant; BW distance also obeys data processing, including for pinching channels.Source: §2 · p. 10

3 The Algorithm: Projected Bures–Wasserstein Gradient Descent

The algorithm lifts BW gradient steps to square-root space, establishes eigenvalue clipping as a non-expansive BW projection, and composes that projection with unit-step RGD.

3.1 BW gradient descent and the square-root manifold

  • BW gradient descent on positive definite matrices is Euclidean gradient descent on the square-root manifold, and the BW step length equals the lifted gradient’s Euclidean length.This identity lets the iteration monitor progress through gradient norms.Source: §3.1 · p. 11

3.2 Eigenvalue clipping is the non-expansive BW projection

  • Clipping eigenvalues to [α,β] gives the unique BW projection onto [αI,βI] and is 1-Lipschitz in BW distance.The Frobenius projection on the lifted, non-convex feasible set can instead be expansive; for singular inputs, the BW projection need not be unique.Source: §3.2 · pp. 12–13

3.3 Projected Bures–Wasserstein gradient descent

  • Projected BW-GD applies the unit-step fixed-point map K and then clips the result to [αI,βI].Clipping reuses the eigendecomposition needed at the next iteration, so both variants require n + 1 spectral decompositions per step; the iteration also has an equivalent lifted singular-value-clipped form.Source: §3.3 · pp. 13–14

4 Dimension-Independent Linear Convergence at Unit Step-Size

Combining projected descent with a PL inequality, the analysis proves dimension-independent linear convergence at unit step for both objectives and supplies iteration-count and stopping certificates.

4 Dimension-Independent Linear Convergence at Unit Step-Size

  • A lifted PL inequality from strong convexity combines with projected descent, while clipping preserves spectral feasibility and, for the invariant problem, invariance.The projected descent step does not increase the objective and gives a decrease controlled by the squared step length.Source: §4 · pp. 15–17
  • At unit step, both objectives contract the function gap by (1 - κ^-3/2) per iteration, independently of dimension.The theorem covers the BW barycenter and invariant matrix projection, with initialization in the prescribed feasible set.Source: §4 · pp. 17–18
  • On the refined interval, κ′ = β′/α′ ≤ κ gives rate (1 - κ′^-3/2) and iteration count κ′3/2 log(1/ϵ), versus O(κ′5/2) for the best dimension-free guarantee.The refined-interval proof requires a separate argument because the input matrices need not lie in that interval.Source: §4 · p. 18

4.1 Choosing the iteration count

  • For absolute tolerance δ, stop when B(St,K(St)) ≤ 2κ^-3/2δ; the squared length of the computed step then certifies a function gap at most δ.The certificate uses a quantity the iteration has already computed, and its resulting count is often far smaller in practice than the a priori count.Source: §4.1 · p. 19

5 Numerical Experiments

The experiments compare Projected RGD with unprojected and small-step RGD, finding much faster unit-step convergence and iteration counts that remain nearly flat as dimension increases. The projection is active in the adversarial example but stays dormant in the other reported experiments, while the certificate is conservative yet useful for stopping.

5 Numerical Experiments

  • The certificate overestimates the true gap by roughly its κ^3/2/2 prefactor, but its certified stopping time improves on the a priori count by orders of magnitude.Because convergence is linear, the multiplicative slack costs only a few additional iterations.Source: §5 · p. 19
  • On Example 13, unit-step RGD reached relative gaps of 10^-6 and 10^-12 in 3 and 6 iterations, versus 544 and 1108 for small-step RGD.Projection was active on exactly the two exit steps, after which projected and unprojected trajectories had indistinguishable convergence rates.Source: §5 · p. 20
  • In the subsequent non-adversarial experiments, clipping never activated, so projected and unprojected unit-step iterations coincided.The authors describe projection as a dormant safeguard in those experiments.Source: §5 · p. 20
  • At κ = 10^3, unit-step RGD reached the numerical floor in 16 iterations versus 16288 for small-step RGD.Figure 2(a) compares the methods on pinned ensembles; its iteration axis is logarithmic.Source: Figure 2 · p. 19, §5 · p. 20
  • For pinned ensembles with κ fixed, unit-step iteration counts nearly coincide across d ∈ {4, 16, 64}, with residual spread tracking κ′ rather than dimension.Pinning the ensembles controls changes in effective difficulty as dimension varies.Source: §5 · p. 20

6 Conclusion

The paper establishes a dimension-independent linear guarantee for projected unit-step RGD on the BW barycenter and invariant matrix projection problems. The projection enforces well-conditioned iterates, while sharper rates for typical unprojected trajectories remain an open question.

6 Conclusion

  • Projected BW-GD converges at the dimension-independent rate (1 − κ^-3/2), improving the best dimension-free iteration complexity from O(κ′^5/2) to κ′^3/2.The conclusion also notes an a posteriori certificate for on-the-fly termination.Source: §6 · p. 20
  • Eigenvalue clipping projects onto the well-conditioned set and enforces the iterate condition needed by the analysis without undoing gradient-step progress.The conclusion notes that unprojected unit-step iterates can genuinely fail to remain well-conditioned.Source: §6 · pp. 20–21
  • The same rate extends to the invariant matrix projection problem because its invariant submanifold is totally geodesic and closed under clipping.Whether typical unprojected trajectories admit a sharper generic-case rate, or the bound is tight only on rare instances, remains open.Source: §6 · p. 21

Use of Large Language Models

The conclusion describes both routine LLM assistance and specific technical contributions made with LLM assistance. It identifies Gemini’s proof of part of Lemma 6 and other results obtained with Claude Fable.

Use of Large Language Models

  • The author reports standard LLM uses including proofreading, restructuring, basic questions, and coding.The paper distinguishes these routine uses from technical contributions.Source: Use of Large Language Models · p. 21
  • Google Gemini 3.1 Pro was used to prove the contraction result in the latter half of Lemma 6, whose proof was then reworked for presentation.The prompt concerned two-sided eigenvalue clipping as a contraction in BW distance.Source: Use of Large Language Models · p. 21
  • The paper also credits Claude Fable with assistance on Proposition 15, Corollary 18, and the Example 13 counterexample.These results concern refined-interval clipping, Frobenius projection, and the counterexample instance, respectively.Source: Use of Large Language Models · p. 21

A Unit-Step RGD Exits the Well-Conditioned Set

A two-dimensional, three-matrix example shows that unit-step RGD can leave the well-conditioned interval even when the ensemble, initialization, and barycenter are strictly interior. The excursion lasts two iterations before the trajectory re-enters.

A Unit-Step RGD Exits the Well-Conditioned Set

  • With d = 2, n = 3, and κ = 300, interior ensemble matrices, initialization, and barycenter do not prevent unit-step RGD from leaving [αI, βI].The example uses a strictly interior initialization and a barycenter certified to lie strictly inside the interval.Source: A Unit-Step RGD Exits the… · pp. 24–25
  • At t = 1 and 2, λmin falls below α by 2.98 × 10−1 and 4.43 × 10−3, then re-enters the interval at t = 3.The limiting barycenter remains strictly above the floor, with λmin(S⋆) −α ≈+2.52 × 10−3.Source: Figure 3 · p. 25, A Unit-Step RGD Exits the… · p. 25
  • Both excursion iterates fall below the certified floor α′ of the optimum, so the exit is not explained by convergence toward a near-boundary optimum.Source: A Unit-Step RGD Exits the… · p. 25
  • The construction extends to every n ≥ 3 by splitting one target’s weight and to every d ≥ 2 by block-diagonal padding, preserving the violations.The authors choose n = 3 because the n = 2 barycenter has a closed form.Source: A Unit-Step RGD Exits the… · pp. 25–26

B A Variance Inequality Proof of the Descent Lemma

The descent lemma is proved by lifting the barycenter problem to a flat square-root space, where a Euclidean variance identity applies. Passing back to the BW manifold contracts distances and preserves a per-step decrease.

A Variance Inequality Proof of the Descent Lemma

  • In the flat total space, the descent statement reduces to the Euclidean variance identity for the mean of optimally paired square roots.The paper gives a self-contained proof of the descent lemma previously proved using optimal-transport tools [BJL19, Theorem 10].Source: B A Variance Inequality Proof of the… · p. 26
  • The quotient map from square roots to positive definite matrices contracts distances because BW distance minimizes Frobenius distance over the fibres.Source: B A Variance Inequality Proof of the… · p. 26
  • The averaged square roots project to the next RGD iterate, so the contraction gives a genuine per-step decrease with the gradient-norm term retained.The argument uses non-negative curvature of the BW manifold, inherited from the flat total space under the Riemannian submersion.Source: B A Variance Inequality Proof of the… · pp. 26–27

C Descent under Projection onto the Refined Interval

For the barycenter problem, clipping to a refined spectral interval yields convergence with refined condition-number constants. Although ensemble matrices may lie outside that interval, clipping still decreases the objective, allowing the convergence analysis to proceed.

Descent under Projection onto the Refined Interval

  • The refined bounds α′I ≤ S⋆ ≤ β′I enclose the optimum and satisfy α ≤ α′ ≤ β′ ≤ β.The bounds are derived from the ensemble’s spectral information.Source: C Descent under Projection onto the… · p. 27
  • Clipping onto [α′I, β′I] decreases the barycenter objective even though individual ensemble matrices may lie outside this refined interval.The fixed-point argument used for the wider interval cannot establish this; the proof works directly with the objective.Source: C Descent under Projection onto the… · p. 27
  • The refined strong-convexity modulus µ′ applies on S ≤ β′I, despite references Ri potentially lying outside the refined interval.This specialization supplies the modulus used with the PL inequality.Source: C Descent under Projection onto the… · p. 28
  • Projected unit-step BW-GD on [α′I, β′I] converges with iteration count κ′3/2 log(1/ϵ).The refined-interval analysis follows the earlier convergence argument with refined constants.Source: C Descent under Projection onto the… · p. 29
  • The refined count is no worse than the small-step O(κ′5/2 log(1/ϵ)) guarantee, but refinement gives no gain for invariant matrix projection or worst-case ensembles.For the invariant problem α′ = α and β′ = β; in the worst case, the barycenter bounds also attain α and β.Source: C Descent under Projection onto the… · p. 29

D The Projected Iteration on the Total Manifold

The paper relates projected BW-GD to projected Euclidean gradient descent on the square-root total manifold. Spectral clipping there intertwines with eigenvalue clipping in the BW manifold, while the invariant-projection update follows the same correspondence with twirled gradients.

The Projected Iteration on the Total Manifold

  • Singular-value clipping is the unique Frobenius projection onto square roots of the well-conditioned set and intertwines with eigenvalue clipping under the quotient map.The construction transports the projection step as well as the gradient step to the total manifold.Source: D The Projected Iteration on the Total… · pp. 29–30
  • Unlike BW projection, Frobenius projection on the total manifold can be expansive because the lower singular-value floor makes its feasible set non-convex.The expansion occurs along fibre directions collapsed by the quotient metric.Source: D The Projected Iteration on the Total… · p. 30
  • Projected BW-GD iterates are projections of the corresponding projected Euclidean GD iterates on the total manifold.Theorem 20 states that π[At] = St at every iteration.Source: D The Projected Iteration on the Total… · pp. 30–31
  • For invariant matrix projection, the correspondence holds verbatim when the gradients are twirled.The proof relies on the gradient step’s form and the clipping identity, both unaffected by twirling.Source: D The Projected Iteration on the Total… · p. 31
  • The square-root gradient update is a weighted average of optimally paired square roots of the reference matrices.This gives the unit-step descent lemma its variance-identity form.Source: D The Projected Iteration on the Total… · pp. 31–32

E The BRT Fixed-Point Algorithm as BW Gradient Descent on the Invariant Submanifold

The appendix identifies BRT’s fixed-point method as BW gradient descent on the invariant submanifold. Its flat square-root geometry enables the application of PL inequalities.

E The BRT Fixed-Point Algorithm as BW Gradient Descent on the Invariant Submanifold

  • The twirl is the BW-orthogonal projection of an ambient tangent vector onto the invariant submanifold’s tangent space, giving its restricted gradient.The ambient Riemannian gradient need not lie in that tangent space, so it must be projected to remain on the submanifold.Source: E The BRT Fixed-Point Algorithm as BW… · p. 32
  • Because the invariant submanifold is totally geodesic, its RGD step can use the ambient BW exponential map.Total geodesy means ambient geodesics starting tangent to the submanifold remain in it.Source: E The BRT Fixed-Point Algorithm as BW… · pp. 32–33
  • At step size 1, RGD on the invariant submanifold recovers the fixed-point map of Brahmachari, Rubboli, and Tomamichel.The objective is to Bures-project a non-invariant positive definite matrix onto the positive definite matrices commuting with a unitary group.Source: E The BRT Fixed-Point Algorithm as BW… · pp. 32–33
  • In the square-root representation, the update is projected Euclidean gradient descent on the invariant subspace G.The projected Euclidean step maps under π to the update on the BW manifold.Source: E The BRT Fixed-Point Algorithm as BW… · p. 33
Loading 2609.03762v1…