spd_learn.models.TangentGraphSPDAggregation#

class spd_learn.models.TangentGraphSPDAggregation(k_neighbors: int = 2, eps: float = 0.0001)[source]#

Bases: Module

Tangent-Space \(k\)-NN Graph-SPD Attention Aggregation.

This module performs geometric graph aggregation over a set of SPD node representations [Tibermacine et al., 2026].

  1. Tangent Space Mapping: Node covariance matrices \(S \in \mathbb{R}^{B \times V \times d \times d}\) are mapped into the Log-Euclidean tangent space:

    \[L_u = \log(S_u) \in \mathcal{S}^d\]

  2. Pairwise Tangent Distance: Computes pairwise Frobenius distances in the tangent space:

    \[D_{uv} = \|L_u - L_v\|_F = \sqrt{\sum_{i,j} (L_{u,ij} - L_{v,ij})^2}\]

  3. Top-:math:`k` Nearest-Neighbor Masking: For each node \(u\), identifies its \(k\) nearest neighbors \(\mathcal{N}_k(u)\) corresponding to the smallest distances \(D_{uv}\). Non-neighbor connections are masked with \(-\infty\).

  4. Attention Weights: Computes softmax attention scores over the neighbors:

    \[A_{uv} = \frac{\exp(-D_{uv})}{\sum_{w \in \mathcal{N}_k(u)} \exp(-D_{uw})}, \quad v \in \mathcal{N}_k(u)\]

  5. Tangent Aggregation and Manifold Retraction: Combines the node’s tangent vector with the attention-weighted neighbor tangent vectors, then projects back to the SPD manifold via matrix exponential:

    \[\tilde{L}_u = 0.5 L_u + 0.5 \sum_{v \in \mathcal{N}_k(u)} A_{uv} L_v\]
    \[\tilde{S}_u = \exp(\tilde{L}_u) \in \mathcal{S}_{++}^d\]
Parameters:

  • k_neighbors (int, default=2) – Number of nearest neighbors to aggregate for each node.

  • eps (float, default=1e-4) – Numerical stability constant for distance computation.

Examples

>>> import torch
>>> from spd_learn.models.ssgcnet import TangentGraphSPDAggregation
>>> agg = TangentGraphSPDAggregation(k_neighbors=2)
>>> S = torch.eye(6).expand(2, 4, 6, 6).clone()
>>> S_out = agg(S)
>>> S_out.shape
torch.Size([2, 4, 6, 6])
forward(S: Tensor) → Tensor[source]#

Forward pass of the tangent-space graph-SPD aggregation.

Parameters:

S (torch.Tensor) – Input node SPD matrices of shape (…, V, d, d).

Returns:

Aggregated node SPD matrices of shape (…, V, d, d).

Return type:

torch.Tensor