spd_learn.modules.SPDBatchNormBW#

class spd_learn.modules.SPDBatchNormBW(num_features: int, momentum: float = 0.1, theta: float = 1.0, n_iter: int = 3, rebias: bool = True, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: SPDBatchNormMean

Bures-Wasserstein Riemannian Batch Normalization for SPD Matrices.

Implements Generalized Bures-Wasserstein Batch Normalization (GBWBN) [Wang et al., 2025].

Parameters:

  • num_features (int) – The size of the SPD matrices (number of features).

  • momentum (float, default=0.1) – Momentum factor for updating the running mean.

  • theta (float, default=1.0) – Power deformation parameter \(\theta > 0\).

  • n_iter (int, default=3) – Number of fixed-point iterations for Bures-Wasserstein barycenter.

  • rebias (bool, default=True) – If True, re-biases the centered data with a learnable SPD matrix.

  • device (Optional[torch.device], default=None) – Device on which to allocate parameters and buffers.

  • dtype (Optional[torch.dtype], default=None) – Data type of parameters and buffers.

forward(input: Tensor) → Tensor[source]#

Forward pass of Bures-Wasserstein Riemannian Batch Normalization.