spd_learn.models.GBWAtt#

class spd_learn.models.GBWAtt(n_chans=22, n_outputs=4, n_patches=3, temporal_out_channels=20, attention_out_features=18, temporal_kernel_size=12, power=1.5, learn_power=False, n_iter=10, aggregation='fixed_point', trace_normalize=True, covariance_epsilon=1e-05, reeig_threshold=1e-05)[source]#

Bases: MAtt

Generalized Bures-Wasserstein attention network for EEG.

See [Jin et al., 2026].

Parameters:

n_chans (int, default=22) – Ordered input electrode channels. Input is (batch,n_chans,n_times).
n_outputs (int, default=4) – Output classes; returns unnormalized (batch,n_outputs) logits for cross_entropy, with no softmax inside the model.
n_patches (int, default=3) – Temporal subparts in acquisition order; each needs >=2 feature samples.
temporal_out_channels (int, default=20) – Covariance dimension (learned convolution features, not electrode axes).
attention_out_features (int, default=18) – SPD dimension after Stiefel Q/K/V projections.
temporal_kernel_size (int, default=12) – Temporal kernel in samples; symmetric padding of kernel_size//2. Even kernels produce n_times+1 outputs, as in the paper MI architecture.
power (float, default=1.5) – Fixed positive deformation exponent.
learn_power (bool, default=False) – Optional constrained learning of the scalar power (not paper default).
n_iter (int, default=10) – Fixed count of differentiable BW barycenter iterations.
aggregation ({'fixed_point', 'root'}, default='fixed_point') – Paper iteration or author weighted-root approximation. The fixed-point mean uses float64 internally for float32 inputs to protect cross-product spectra, returning checked float32 output.
trace_normalize (bool, default=True) – Normalize each covariance before diagonal loading, as stated in paper. Author Signal2Spd default disables this.
covariance_epsilon (float, default=1e-05) – Absolute positive diagonal loading after optional trace normalization.
reeig_threshold (float, default=1e-05) – Paper rectification threshold (author code instead uses 1e-4).

Notes

Accepts (trial, channel, sample) EEG and returns class logits. The default frontend is the MI architecture. Trace normalization and registered metric parameters follow the stated geometry; the released code differs. Fixed-point aggregation is approximate. Train BatchNorm within the training partition and keep input units, sampling and channel order consistent.