spd_learn.modules.GaussianSPD#

class spd_learn.modules.GaussianSPD(embedding_type: str = 'quotient', alpha: float = 1.0, beta: float = 1.0, eps: float = 1e-05, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Gaussian SPD Descriptor.

Embeds empirical first-order (mean) and second-order (covariance) statistics of a sequence of token representations into an extended Symmetric Positive Definite (SPD) matrix using the quotient / Lovrić embedding [Hu et al., 2026].

Given a sequence of \(T\) token representations \(X \in \mathbb{R}^{\dots \times T \times D}\), this layer computes the empirical mean \(\mu \in \mathbb{R}^{\dots \times D \times 1}\):

\[\mu = \frac{1}{T} \sum_{t=1}^T X_t\]

and centered token covariance matrix \(\Sigma \in \mathcal{S}_{++}^D\):

\[\Sigma = \frac{1}{\max(T - 1, 1)} \sum_{t=1}^T (X_t - \mu)(X_t - \mu)^* + \epsilon I_D\]

where \(\epsilon > 0\) guarantees strict positive definiteness.

The mean and covariance are embedded into a \((D + 1) \times (D + 1)\) SPD matrix \(\mathcal{G}(\mu, \Sigma) \in \mathcal{S}_{++}^{D+1}\):

\[\begin{split}\mathcal{G}(\mu, \Sigma) = \begin{pmatrix} \Sigma + \beta \mu \mu^* & \alpha \mu \\ \alpha \mu^* & \frac{\alpha^2}{\beta} \end{pmatrix}\end{split}\]

By the Schur complement condition, \(\mathcal{G}(\mu, \Sigma)\) is strictly positive definite whenever \(\Sigma \succ 0\), \(\beta > 0\), and \(\alpha \neq 0\).

Parameters:

  • embedding_type (str, default="quotient") – The type of Gaussian embedding into the SPD manifold. Supported: "quotient" (or "lovric").

  • alpha (float, default=1.0) – Cross-term coupling coefficient between mean and covariance.

  • beta (float, default=1.0) – Mean rank-one perturbation weight. Must be positive.

  • eps (float, default=1e-5) – Regularization added to the covariance diagonal to ensure strict positive definiteness.

  • device (torch.device, optional) – Device on which tensors will be allocated.

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

See also

RHOP

End-to-end pooling layer using this descriptor.

CovLayer

Standard sample covariance layer.

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 GaussianSPD
>>> layer = GaussianSPD()
>>> tokens = torch.randn(2, 10, 8)  # batch=2, tokens=10, features=8
>>> G = layer(tokens)
>>> G.shape
torch.Size([2, 9, 9])
forward(X: Tensor) → Tensor[source]#

Compute Gaussian SPD descriptor from sequence tensor.

Parameters:

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

Returns:

Gaussian SPD matrix of shape (…, n_features + 1, n_features + 1).

Return type:

torch.Tensor