spd_learn.modules.RHOP

class spd_learn.modules.RHOP(in_features: int, reduced_dim: int = 16, metric: str = 'LEM', power: float = 0.5, alpha: float = 1.0, beta: float = 1.0, eps: float = 1e-05, device: device | None = None, dtype: dtype | None = None)

Bases: Module

Riemannian High-Order Pooling (RHOP) for Brain Foundation Models.

Implements the Riemannian High-Order Pooling classification head proposed by Hu et al. [Hu et al., 2026] (ICLR 2026). RHOP augments transformer backbones (such as EEG foundation models BIOT and LaBraM) by capturing second-order Riemannian geometric statistics across sequence tokens, rather than relying exclusively on first-order average pooling or a single class token.

Algorithm Formulation

1. Dimension Reduction: Sequence tokens \(X \in \mathbb{R}^{\dots \times T \times D_{\text{in}}}\) are projected to a reduced dimension \(d = \text{reduced\_dim}\):

   \[Z = X W_{\text{down}}^\top + b_{\text{down}} \in \mathbb{R}^{\dots \times T \times d}\]

2. Gaussian SPD Descriptor: The empirical mean and covariance of \(Z\) are embedded into a \((d + 1) \times (d + 1)\) SPD matrix:

   \[G = \mathcal{G}(\mu_Z, \Sigma_Z) \in \mathcal{S}_{++}^{d+1}\]

3. Riemannian Geometry Mapping: The matrix \(G\) is mapped to Euclidean / tangent space via one of three Riemannian metrics:

   • LEM (Log-Euclidean Metric):

     \[S = \log(G)\]

   • PEM (Power-Euclidean Metric):

     \[S = \frac{G^p - I}{p}\]

   • LCM (Log-Cholesky Metric):

     \[L = \text{chol}(G), \quad S = \lfloor L \rfloor + \log(\mathbb{D}(L))\]

4. Norm-Preserving Vectorization: Vectorizes the non-redundant symmetric coefficients via Vech:

   \[v = \operatorname{vech}(S) \in \mathbb{R}^{\dots \times \frac{(d+1)(d+2)}{2}}\]

5. Up-projection & Fusion: Vector \(v\) is projected back to the backbone dimension \(D_{\text{in}}\) and fused with the summary token:

   \[f_{\text{RHOP}} = v W_{\text{up}}^\top + b_{\text{up}}\]
   \[\begin{split}y = \begin{cases} x_{\text{cls}} + f_{\text{RHOP}} & \text{if } x_{\text{cls}} \text{ is provided} \\ \bar{X} + f_{\text{RHOP}} & \text{otherwise} \end{cases}\end{split}\]

Parameters:

• in_features (int) – Input token feature dimension \(D_{\text{in}}\).

• reduced_dim (int, default=16) – Reduced dimension \(d\) for token projection before forming the SPD descriptor.

• metric (str, default="LEM") – Riemannian metric mapping. Options: "LEM" (Log-Euclidean Metric), "PEM" (Power-Euclidean Metric), "LCM" (Log-Cholesky Metric).

• power (float, default=0.5) – Power parameter \(p\) used when metric="PEM".

• alpha (float, default=1.0) – Coupling scale factor for the Gaussian SPD descriptor.

• beta (float, default=1.0) – Mean rank-one perturbation weight for the Gaussian SPD descriptor.

• eps (float, default=1e-5) – Covariance regularization diagonal loading.

• device (torch.device, optional) – Device on which module parameters and buffers are created.

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

See also

GaussianSPD

Forms the quotient Gaussian SPD matrix.

Vech

Vectorization of symmetric matrix.

References

[1]

Chen Hu, Rui Wang, Ziheng Chen, Yefeng Zheng, and Nicu Sebe. Riemannian high-order pooling for brain foundation models. In The Fourteenth International Conference on Learning Representations. 2026. URL: ChenHu-ML/RHOP.

Examples

>>> import torch
>>> from spd_learn.modules import RHOP
>>> rhop = RHOP(in_features=32, reduced_dim=8, metric="LEM")
>>> tokens = torch.randn(4, 16, 32)  # batch=4, tokens=16, dim=32
>>> out = rhop(tokens)
>>> out.shape
torch.Size([4, 32])

forward(X: Tensor, cls_token: Tensor | None = None) → Tensor

Forward pass of Riemannian High-Order Pooling.

Parameters:

• X (torch.Tensor) – Sequence tensor of shape (…, n_tokens, in_features).

• cls_token (torch.Tensor, optional) – Optional class token of shape (…, in_features) or (…, 1, in_features). If not provided, empirical token mean X.mean(dim=-2) is used for fusion.

Returns:

Pooled representation of shape (…, in_features).

Return type:

torch.Tensor