spd_learn.functional.karcher_mean_iteration#

spd_learn.functional.karcher_mean_iteration(X: Tensor, current_mean: Tensor, detach: bool = True, return_tangent: bool = False) → Tensor | Tuple[Tensor, Tensor][source]#

Perform one iteration of the Karcher mean algorithm.

The Karcher (Fréchet) mean on the SPD manifold is the minimizer of the sum of squared geodesic distances. This function performs one iteration of the iterative algorithm to compute it.

Given a current estimate \(M\) of the mean, the update is:

\[M_{\text{new}} = M^{1/2} \exp\left(\frac{1}{N} \sum_{i=1}^N \log(M^{-1/2} X_i M^{-1/2})\right) M^{1/2}\]

Parameters:

X (torch.Tensor) – Batch of SPD matrices with shape (batch_size, …, n, n).
current_mean (torch.Tensor) – Current estimate of the Karcher mean with shape (1, …, n, n).
detach (bool, default=True) – If True, detaches current_mean from the computational graph before computing the update. Set to False when gradients with respect to the mean are needed.
return_tangent (bool, default=False) – If True, also returns the mean tangent update used in this Karcher step.

Returns:

Updated Karcher mean estimate with shape (1, …, n, n). When return_tangent=True, also returns the mean tangent update with the same shape.

Return type:

torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

Notes

For well-conditioned data, a single iteration often suffices. The algorithm converges quadratically near the solution.

When detach=False, gradients flow through the entire computation, including the matrix square root and inverse square root of the current mean.