spd_learn.modules.GBWAttention#

class spd_learn.modules.GBWAttention(in_features, out_features, power=1.5, learn_power=False, n_iter=10, aggregation='fixed_point')[source]#

Bases: Module

Apply generalized BW attention to SPD tokens.

See [Jin et al., 2026].

Parameters:

  • in_features (int) – Input matrix dimension.

  • out_features (int) – Output matrix dimension.

  • power (float, default=1.5) – Positive power deformation, fixed unless learn_power=True.

  • learn_power (bool, default=False) – Enable constrained learnable power.

  • n_iter (int, default=10) – Unrolled BW barycenter iterations (approximation, not convergence).

  • aggregation ({'fixed_point', 'root'}, default='fixed_point') – Paper Eq. 22 iteration or author weighted-root approximation. Fixed-point accumulation uses float64 for float32 inputs, returning checked float32 output with differentiable casts. Root stays in input dtype.

Notes

Independent Stiefel Q/K/V projections precede the common chart. Attention aggregates tokens within each sample; no masks or padding are supported.

forward(x, *, return_weights=False)[source]#

Attend to a sequence of SPD tokens.

Parameters:

  • x (torch.Tensor) – SPD tokens of shape (batch, tokens, in_features, in_features).

  • return_weights (bool, default=False) – Also return attention weights.

Returns:

  • torch.Tensor – SPD outputs of shape (batch, tokens, out_features, out_features).

  • weights (torch.Tensor, optional) – Attention of shape (batch, tokens, tokens), returned when return_weights=True.