spd_learn.modules.SPMultiheadAttention#

class spd_learn.modules.SPMultiheadAttention(embed_dim=6, num_heads=2, dropout=0.0, bias=True, *, device=None, dtype=None)[source]#

Bases: Module

Apply SP-MHA to symmetric log coordinates.

See [Seraphim et al., 2024].

Parameters:

  • embed_dim (int, default=6) – Positive triangular coordinate dimension, divisible by num_heads.

  • num_heads (int, default=2) – Number of independently normalized attention maps.

  • dropout (float, default=0.0) – Dropout of each head’s post-softmax probabilities during training.

  • bias (bool, default=True) – Include biases in query/key projections (head mixing is bias-free).

  • device (torch.device or str, optional) – Device for parameters and buffers.

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

Notes

Q/K are projected per head. A signed linear combination of head maps multiplies the unprojected values. Raw covariance entries are not valid value tokens. Masks, padding and causal operation are unsupported.

forward(query: Tensor, key: Tensor | None = None, value: Tensor | None = None, *, key_padding_mask=None, attn_mask=None, is_causal=False, need_weights=False) → tuple[Tensor, Tensor | None][source]#

Apply structure-preserving multihead attention.

Parameters:

  • query (torch.Tensor) – Log tokens of shape (batch, tokens, embed_dim).

  • key (torch.Tensor or None) – Keys with the query shape; None uses query.

  • value (torch.Tensor or None) – Unprojected values with the query shape; None uses query.

  • key_padding_mask (None) – Reserved argument; padding masks are not supported.

  • attn_mask (None) – Reserved argument; masks are not supported.

  • is_causal (bool, default=False) – Must be False; causal attention is not supported.

  • need_weights (bool, default=False) – Return the combined attention map.

Returns:

  • output (torch.Tensor) – Attended log tokens with the query shape.

  • weights (torch.Tensor or None) – Combined map of shape (batch, tokens, tokens) when need_weights=True.