spd_learn.modules.SPDBatchNormLie

class spd_learn.modules.SPDBatchNormLie(num_features, metric='AIM', theta=1.0, alpha=1.0, beta=0.0, momentum=0.1, eps=1e-05, n_iter=1, congruence='cholesky', device=None, dtype=None)

Bases: Module

Lie Group Batch Normalization for SPD matrices.

Implements the LieBN framework [Chen et al., 2024] for SPD manifolds. Unlike SPDBatchNormMeanVar, which normalizes under a single Riemannian metric (AIRM), this layer exploits the Lie group structure of three classical SPD geometries to define centering, scaling, and biasing as group-theoretic operations with formal statistical guarantees.

Algorithm. Given a batch \(\{P_i\}_{i=1}^N \subset \mathcal{S}_{++}^n\), the forward pass applies three steps in the Lie algebra selected by metric:

1. Centering – translate the batch mean \(M\) to the group identity \(E\) via the inverse left translation:

   \[\bar{P}_i = L_{M_\odot^{-1}}(P_i)\]

2. Scaling – normalize the Fréchet variance \(v^2\) with a learnable shift \(s \in \mathbb{R}_{>0}\):

   \[\hat{P}_i = \operatorname{Exp}_E \!\left[\frac{s}{\sqrt{v^2 + \epsilon}}\, \operatorname{Log}_E(\bar{P}_i)\right]\]

3. Biasing – translate to the learnable SPD parameter \(B\):

   \[\tilde{P}_i = L_B(\hat{P}_i)\]

Theoretical guarantees (Proposition 4.2 of the paper):

• Mean control: after centering and biasing with \(B = E\), the Fréchet mean of the output batch equals \(E\).

• Variance control: after scaling, the output dispersion satisfies \(\sum_i w_i\,d^2(\hat{P}_i, E) = s^2\).

Supported metrics. The metric parameter selects one of three Lie group structures, each inducing a family of parameterized metrics via the power deformation \(\mathrm{P}_\theta\). The table below summarizes how each step is realized (see Table 2 in [Chen et al., 2024]):

Operation	\((\theta,\alpha,\beta)\)-AIM	\((\alpha,\beta)\)-LEM	\(\theta\)-LCM
Pullback map	\(\mathrm{P}_\theta\)	\(\operatorname{mlog}\)	\(\psi_{\mathrm{LC}} \circ \mathrm{P}_\theta\)
Left translation \(L_Q(P)\)	\(Q^{1/2} P\, Q^{1/2}\)	\(P + Q\)	\(P + Q\)
Scaling	\(\operatorname{Exp}_I[s\,\operatorname{Log}_I(P)]\)	\(s \cdot P\)	\(s \cdot P\)
Fréchet mean	Karcher flow	Arithmetic mean	Arithmetic mean
Running mean update	AIRM geodesic	Linear interpolation	Linear interpolation

Bi-invariant distance. The Fréchet variance uses the \((\alpha, \beta)\) bi-invariant metric (Definition 3 and Eq. 3 of the paper):

\[d^2(P, Q) = \alpha \lVert V \rVert_F^2 + \beta \, g(V)^2\]

where \(V\) is the tangent representation (log-map) and \(g(V) = \log\det(P)\) for AIM or \(\operatorname{tr}(V)\) for LEM/LCM. The variance is normalized by \(\theta^2\) for AIM and LCM.

Parameters:

• num_features (int) – Size of the SPD matrices (\(n \times n\)).

• metric ({"AIM", "LEM", "LCM"}, default="AIM") – Lie group invariant metric.

• theta (float, default=1.0) – Power deformation parameter \(\theta\). When \(\theta = 1\), no deformation is applied.

• alpha (float, default=1.0) – Frobenius norm weight \(\alpha\) in the bi-invariant distance.

• beta (float, default=0.0) – Trace / log-determinant weight \(\beta\) in the bi-invariant distance. Must satisfy \(\min(\alpha, \alpha + n\beta) > 0\).

• momentum (float, default=0.1) – Momentum \(\gamma\) for exponential moving average of running statistics.

• eps (float, default=1e-5) – Numerical stability constant \(\epsilon\) added to the variance before taking the square root.

• n_iter (int, default=1) – Number of Karcher flow iterations for the AIM Fréchet mean. Ignored by LEM and LCM (which use arithmetic means).

• congruence ({"cholesky", "eig"}, default="cholesky") – Implementation of the AIM congruence action (centering/biasing). "cholesky" uses the Cholesky factor \(L\) of \(P\) to compute \(L X L^\top\) (as in the original LieBN paper). "eig" uses eigendecomposition-based \(M^{-1/2} X M^{-1/2}\) (matching spd_centering()). Both are mathematically equivalent; Cholesky is typically faster, while eigendecomposition reuses the infrastructure of SPDBatchNormMeanVar. Only affects the AIM metric.

• device (torch.device or str, optional) – Device on which to create parameters and buffers.

• dtype (torch.dtype, optional) – Data type of parameters and buffers.

bias

Learnable SPD bias matrix \(B \in \mathcal{S}_{++}^n\), parametrized via SymmetricPositiveDefinite. Initialized to the identity.

Type:

nn.Parameter

shift

Learnable positive scalar \(s > 0\), parametrized via PositiveDefiniteScalar. Initialized to 1.

Type:

nn.Parameter

running_mean

Exponential moving average of the batch Fréchet mean.

Type:

torch.Tensor

running_var

Exponential moving average of the batch variance.

Type:

torch.Tensor

See also

SPDBatchNormMean

Mean-only Riemannian batch normalization (AIRM centering without variance normalization) [Brooks et al., 2019].

SPDBatchNormMeanVar

Full Riemannian batch normalization under the AIRM [Kobler et al., 2022].

frechet_mean()

Fréchet mean via Karcher flow (used internally for AIM).

lie_group_variance()

Bi-invariant Fréchet variance computation.

References

[1] (1,2)

Ziheng Chen, Yue Song, Yunmei Xu, and Nicu Sebe. A lie group approach to riemannian batch normalization. In International Conference on Learning Representations. 2024. URL: https://openreview.net/forum?id=okYdj8Ysru.

Examples

>>> import torch
>>> from spd_learn.modules import SPDBatchNormLie
>>> bn = SPDBatchNormLie(num_features=4, metric="AIM")
>>> X = torch.randn(8, 4, 4, dtype=torch.float64)
>>> X = X @ X.mT + 0.1 * torch.eye(4, dtype=torch.float64)
>>> bn = bn.to(dtype=torch.float64)
>>> Y = bn(X)
>>> Y.shape
torch.Size([8, 4, 4])

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(X)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters()