spd_learn.models.CorrAttention#

class spd_learn.models.CorrAttention(dim: int, metric: str = 'mix', device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Correlation Manifold Attention Module (CorrAttention).

This module implements a self-attention mechanism on the Riemannian manifold of full-rank correlation matrices \(\mathcal{C}_{++}^d\) [Hu et al., 2025].

Given a sequence of correlation matrices \(X \in \mathbb{R}^{B \times P \times d \times d}\), the attention mechanism:

  1. Converts inputs to correlation matrices via cov_to_corr().

  2. Maps to the hollow symmetric tangent space \(\text{Hol}(\mathcal{S}^d)\) via log_map_olm() (for metric in ("OLM", "mix")) or log_map_lsm() (for metric == "LSM").

  3. Transforms queries, keys, and values via Cayley-parameterized orthogonal mappings via BiMap, projected onto the hollow tangent space: \(Q = \text{off\_diag}(W_Q^\top L W_Q)\), \(K = \text{off\_diag}(W_K^\top L W_K)\), \(V = \text{off\_diag}(W_V^\top L W_V)\).

  4. Computes pairwise Frobenius distances \(D_{ij} = \|Q_i - K_j\|_F\) and attention weights \(A_{ij} = \text{softmax}_j(1 / (1 + \log(1 + D_{ij})))\).

  5. Computes the weighted sum \(\sum_j A_{ij} V_j + \text{off\_diag}(B + B^\top)\), followed by geodesic residual interpolation \(0.5 L + 0.5 \text{attn\_tangent}\).

  6. Maps back to \(\mathcal{C}_{++}^d\) via exp_map_lsm() (for “mix”/”LSM”) or exp_map_olm() (for “OLM”).

Parameters:

  • dim (int) – Dimension \(d\) of the correlation matrices.

  • metric ({"mix", "OLM", "LSM"}, default="mix") – Riemannian metric on the correlation manifold. In “mix” mode, the logarithmic map uses OLM (fast, closed-form) and the exponential map uses LSM (fast, closed-form).

  • device (torch.device, optional) – Device on which to allocate parameters.

  • dtype (torch.dtype, optional) – Data type for model parameters.

forward(x: Tensor) → Tensor[source]#

Forward pass of correlation attention.

Parameters:

x (torch.Tensor) – Input tensor of patch matrices with shape (…, n_patches, dim, dim).

Returns:

Attended correlation matrices of shape (…, n_patches, dim, dim).

Return type:

torch.Tensor

reset_parameters() → None[source]#

Reset layer parameters.