spd_learn.modules.SPDBatchNormMean

class spd_learn.modules.SPDBatchNormMean(num_features, momentum=0.1, rebias=True, n_iter=1, device=None, dtype=None)

Bases: Module

Riemannian Batch Normalization for SPD Matrices (Mean-only).

This class implements the Riemannian Batch Normalization (RBN) layer for the Symmetric Positive Definite (SPD) manifold [Brooks et al., 2019].

\[\tilde{P}_i = \mathcal{G}^{-\frac{1}{2}} P_i \mathcal{G}^{-\frac{1}{2}}\]

where \(\mathcal{G}\) is the Fréchet mean of the batch.

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.

• rebias (bool, default=True) – If True, the layer rebases the data.

• n_iter (int, default=1) – Number of Karcher flow iterations to estimate the batch mean.

Examples

>>> import torch
>>> from spd_learn.modules import SPDBatchNormMean
>>> bn = SPDBatchNormMean(num_features=4, momentum=0.1)
>>> X = torch.randn(8, 4, 4)
>>> X = X @ X.mT + 0.1 * torch.eye(4)  # Make SPD
>>> Y = bn(X)
>>> Y.shape
torch.Size([8, 4, 4])

import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from spd_learn.modules import SPDBatchNormMean

def spd_to_ellipse(spd_matrix, center=(0, 0), scale=1.0):
    eigvals, eigvecs = np.linalg.eigh(spd_matrix)
    width = 2 * np.sqrt(eigvals[1]) * scale
    height = 2 * np.sqrt(eigvals[0]) * scale
    angle = np.degrees(np.arctan2(eigvecs[1, 1], eigvecs[0, 1]))
    return Ellipse(center, width, height, angle=angle)

# Create batch of 2x2 SPD matrices
torch.manual_seed(42)
np.random.seed(42)
batch_size = 6

spd_batch = []
for i in range(batch_size):
    scale = np.random.uniform(0.5, 2.0)
    angle = np.random.uniform(0, np.pi)
    R = np.array([[np.cos(angle), -np.sin(angle)],
                  [np.sin(angle), np.cos(angle)]])
    D = np.diag([scale, scale * np.random.uniform(0.3, 1.0)])
    S = R @ D @ D @ R.T
    spd_batch.append(S)

X = torch.tensor(np.array(spd_batch), dtype=torch.float32)

# Apply SPDBatchNormMean
bn = SPDBatchNormMean(num_features=2, momentum=0.1, rebias=False)
bn.train()
Y = bn(X)

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

colors = plt.cm.tab10(np.linspace(0, 1, batch_size))

# Before normalization
ax1 = axes[0]
for i, S in enumerate(X.numpy()):
    ellipse = spd_to_ellipse(S, scale=0.5)
    ellipse.set_facecolor(colors[i])
    ellipse.set_alpha(0.6)
    ellipse.set_edgecolor('black')
    ax1.add_patch(ellipse)
ax1.set_xlim(-3, 3)
ax1.set_ylim(-3, 3)
ax1.set_aspect('equal')
ax1.grid(True, alpha=0.3)
ax1.axhline(y=0, color='k', linewidth=0.5)
ax1.axvline(x=0, color='k', linewidth=0.5)
ax1.set_title('Before SPDBatchNormMean', fontweight='bold')

# After normalization
ax2 = axes[1]
for i, S in enumerate(Y.detach().numpy()):
    ellipse = spd_to_ellipse(S, scale=0.5)
    ellipse.set_facecolor(colors[i])
    ellipse.set_alpha(0.6)
    ellipse.set_edgecolor('black')
    ax2.add_patch(ellipse)
identity = Ellipse((0, 0), 1, 1, facecolor='none',
                   edgecolor='red', linewidth=2, linestyle='--')
ax2.add_patch(identity)
ax2.set_xlim(-3, 3)
ax2.set_ylim(-3, 3)
ax2.set_aspect('equal')
ax2.grid(True, alpha=0.3)
ax2.axhline(y=0, color='k', linewidth=0.5)
ax2.axvline(x=0, color='k', linewidth=0.5)
ax2.set_title('After SPDBatchNormMean', fontweight='bold')

plt.suptitle('SPDBatchNormMean: Riemannian Centering', fontsize=13, fontweight='bold')
plt.tight_layout()
plt.show()

(Source code)

forward(input)

Forward pass of the Riemannian Batch Normalization layer.

Parameters:

input (torch.Tensor) – Input tensor of shape (batch_size, h, n, n), where each slice along the batch dimension is an SPD matrix.

Returns:

Normalized tensor of the same shape as the input.

Return type:

torch.Tensor

reset_parameters() → None

reset_running_stats() → None