Projected Bures Covariance Pooling
Implementation & benchmark of arXiv:2609.03762 — Projected Riemannian Gradient Descent for the Bures-Wasserstein Barycenter: Dimension-Independent Linear Convergence at Unit Step Size
Source paper: Projected Riemannian Gradient Descent for the Bures-Wasserstein Barycenter: Dimension-Independent Linear Convergence at Unit Step Size arXiv:2609.03762 ⓘ · analyzed Sep 4, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace Euclidean or unprojected covariance averaging with a projected Bures-Wasserstein barycenter layer. Each unit-step barycenter update is followed by eigenvalue clipping into \([\alpha,\beta]\), preserving positive definiteness and preventing ill-conditioning without an additional eigendecomposition.
Formulas
Mathematical statement
For positive semidefinite matrices \(P,Q\in\mathrm{PSD}(d)\), the squared Bures-Wasserstein distance is \(\mathrm{B}(P,Q)=\operatorname{Tr}(P+Q)-2\mathrm{F}(P,Q)\), where \(\mathrm{F}(P,Q)=\operatorname{Tr}[(P^{1/2}QP^{1/2})^{1/2}]=\|P^{1/2}Q^{1/2}\|_1\). Given SPD inputs \(A_1,\ldots,A_m\) and weights \(w_i\ge0\) with \(\sum_iw_i=1\), the BW barycenter minimizes \(\Phi(S)=\frac12\sum_iw_i\mathrm{B}(S,A_i)\). The paper's projection lemma states that for \(X=U\operatorname{diag}(\lambda_1,\ldots,\lambda_d)U^\top\), \(\Pi_{[\alpha,\beta]}(X)=U\operatorname{diag}(\operatorname{clip}(\lambda_j,\alpha,\beta))U^\top\) is the closed-form BW projection onto \(\mathcal C_{\alpha,\beta}=\{S:\alpha I\preceq S\preceq\beta I\}\), and it is non-expansive: \(d_{\mathrm{BW}}(\Pi(X),\Pi(Y))\le d_{\mathrm{BW}}(X,Y)\). Applying this projection to the unit-step BW fixed-point map \(T\), \(S_{t+1}=\Pi_{[\alpha,\beta]}(T(S_t))\), gives the claimed rate \((1-\kappa^{-3/2})\) when the ensemble condition number is \(\kappa\) and inputs lie in the spectral interval.
Implementation notes
(1) Integration point: add this as a differentiable covariance-aggregation module in a CNN, graph neural network, or set encoder. For each sample, form SPD feature covariances \(A_i=H_i^\top H_i/n+\epsilon I\), where \(H_i\) is a patch, token-neighborhood, or augmentation feature matrix. Replace arithmetic pooling with \(K\) projected BW barycenter iterations. (2) Pseudocode: initialize \(S\leftarrow\sum_iw_iA_i\). For \(t=1,\ldots,K\), compute \(R_i=(S^{1/2}A_iS^{1/2})^{1/2}\), form the unit-step proposal \(X\leftarrow(\sum_iw_iR_i)^2\), eigendecompose \(X=U\operatorname{diag}(\lambda)U^\top\), and update \(S\leftarrow U\operatorname{diag}(\operatorname{clip}(\lambda,\alpha,\beta))U^\top\). Feed \(S\), or \(\log S\), to the next layer and backpropagate through the matrix operations with stabilized eigendecomposition gradients. (3) Compute the BW update and clipping exactly; choose \(\alpha\) and \(\beta\) from running lower and upper eigenvalue quantiles of the inputs, with gradients stopped through these bounds. Estimate the empirical contraction by measuring successive \(d_{\mathrm{BW}}(S_t,S_{t-1})\); the theorem supplies the non-expansiveness and claimed rate, while task improvement is empirical. (4) First experiment: compare arithmetic pooling, log-Euclidean pooling, unprojected BW iterations, and projected BW iterations on CIFAR-100 using a small ResNet and covariance features, with \(K\in\{3,5,10\}\). Measure validation accuracy, barycenter objective, condition numbers, NaN or gradient-failure frequency, and loss versus eigendecomposition FLOPs. Success is lower instability and faster objective decrease at equal compute, with no extra eigendecomposition relative to the unprojected BW baseline.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.