spd_learn.modules.SPDBatchNormHarmonic#
- class spd_learn.modules.SPDBatchNormHarmonic(num_features: int, momentum: float = 0.1, alpha: float = 0.5, learnable_alpha: bool = True, rebias: bool = True, eps: float = 1e-05, device: device | None = None, dtype: dtype | None = None)[source]#
Bases:
SPDBatchNormMeanHarmonic-Arithmetic Riemannian Batch Normalization for SPD Matrices.
Implements the ARMAGNAC parametric batch normalization layer for Symmetric Positive Definite (SPD) manifolds [Gallet et al., 2026, Gallet et al., 2026].
Instead of iterative Karcher flow under the AIRM (as in
SPDBatchNormMean), ARMAGNAC replaces the Riemannian batch mean with a closed-form parametric geodesic interpolation between the batch harmonic mean and the batch arithmetic mean:\[\mathcal{M}_\alpha = \mathcal{H} \#_\alpha \mathcal{A} = \mathcal{H}^{1/2} \left(\mathcal{H}^{-1/2} \mathcal{A} \mathcal{H}^{-1/2}\right)^\alpha \mathcal{H}^{1/2}\]where \(\mathcal{A}\) is the arithmetic mean across the batch and \(\mathcal{H}\) is the harmonic mean:
\[\mathcal{A} = \frac{1}{N} \sum_{i=1}^N X_i, \quad \mathcal{H} = \left( \frac{1}{N} \sum_{i=1}^N (X_i + \epsilon I)^{-1} \right)^{-1}\]- 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.
alpha (float, default=0.5) – Initial value for the harmonic-arithmetic interpolation parameter \(\alpha \in (0, 1)\).
learnable_alpha (bool, default=True) – If True, \(\alpha\) is optimized via backpropagation.
rebias (bool, default=True) – If True, re-biases the centered data with a learnable SPD matrix.
eps (float, default=1e-5) – Small regularization constant added to matrices before inversion.
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.