Skip to content

[models] Add TMSA-Net - #1209

Merged
bruAristimunha merged 9 commits into
braindecode:masterfrom
lindicaphxag-tech:feat/tmsa-net-1201-clean
Oct 9, 2026
Merged

bruAristimunha merged 9 commits into
braindecode:masterfrom
lindicaphxag-tech:feat/tmsa-net-1201-clean

Conversation

@lindicaphxag-tech

@lindicaphxag-tech lindicaphxag-tech commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Part of #1201.

Summary

This adds TMSA-Net (Zhao & Zhu, 2025), one of the proposed first ports in #1201, as a Braindecode motor-imagery model.

The port keeps the released computation rather than replacing the custom attention with nn.MultiheadAttention.

In particular, the reference implementation uses

head_dim = embed_dim // num_heads
inner_dim = num_heads * head_dim

so the default BCI-IV-2a setting embed_dim=19, num_heads=4 is intentionally 19 → 16 → 19. The local-key and global-key attention branches are computed separately and summed before the output projection.

Braindecode integration

  • TMSANet follows EEGModuleMixin and the standard n_chans / n_times / n_outputs contract.
  • The released BCI-IV-2a, BCI-IV-2b, and HGD model geometries are covered by tests.
  • final_layer is exposed for the standard classifier/head integration.
  • The model is registered in the model registry, summary table, API docs, and model categorization.
  • The reference-only radix argument is not exposed: in the released source it only multiplies the input-channel count, while Braindecode's n_chans already denotes the actual number of input channels. All released configurations use radix=1.

Reference fidelity

Pinned reference:

Whit3Zhao/TMSA-Net@c60882db35eeff860a5014df7b0f54dda6601c65

A deterministic CPU parity assay maps all 45 model-state entries from the released implementation into this port. On the default BCI-IV-2a configuration and the same input:

  • attention head width: 4
  • attention inner width: 16
  • logits max absolute error: 0
  • input-gradient max absolute error: 0

Current-head parity run (82b61e569ea9dcf0f90efd3eaf07ce05d1c1b128):
https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37189419406

The parity check is against the released source implementation; this PR does not claim pretrained-checkpoint parity.

Validation

Focused regressions cover:

  • the three released dataset configurations;
  • the non-divisible 19/4 attention bottleneck;
  • the local + global attention sum;
  • classifier temporal geometry;
  • invalid architectural configurations.

The current PR head passes 13 focused TMSA-Net tests, 12 TMSANet-selected Braindecode integration cases, and the touched-file pre-commit checks:

https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37189437290

The parity and focused validation above both target the current PR head; no pretrained-checkpoint result is implied.

Copy link
Copy Markdown
Contributor Author

Exact-head validation is green on 3a62b2e6d94c498e9fa92df3767c4e98be33a4ce: all 13 TMSA-Net focused regressions pass, and 12 shared Braindecode integration cases pass on the same head. This complements the pinned reference-parity assay already in the PR body (45/45 state entries, zero logits/input-gradient error). Evidence run: https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37136260152. I’ll hold this head for maintainer review.

Copy link
Copy Markdown
Contributor Author

Current-head validation is now complete on 82b61e569ea9dcf0f90efd3eaf07ce05d1c1b128. Pinned-reference parity maps all 45/45 state entries and remains exact for both logits and input gradients (max_abs = 0.0), while preserving the released 19→16→19 attention bottleneck (head_dim=4, inner_dim=16): https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37189419406. The same current head also passes 13 focused TMSA-Net tests, 12 Braindecode integration cases, and all touched-file pre-commit hooks: https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37189437290. I’ll hold this head for maintainer review.

@codecov

codecov Bot commented Oct 4, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.83721% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 89.20%. Comparing base (6abd308) to head (10444f9).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #1209      +/-   ##
==========================================
+ Coverage   89.02%   89.20%   +0.18%     
==========================================
  Files         157      159       +2     
  Lines       19398    19727     +329     
==========================================
+ Hits        17269    17598     +329     
  Misses       2129     2129              
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@bruAristimunha

bruAristimunha commented Oct 5, 2026 •

Copy link
Copy Markdown
Collaborator

Integration gate (braindecode maintainers)

Target: BCI-IV-2a (BNCI2014_001) within-subject accuracy (session-1 train / session-2 test), 9-subject mean. The paper reports 83.22% on 2a. Gate: within 5% of that (0.791-0.874). The released code's 2a preprocessing is private (newbcicIV2a/ .mat, no script), and the reported number selects on the test set over 2000 epochs with no fixed seed. So we replicate under a declared, labelled substitute protocol: MOABB 4-38 Hz/250 Hz/4 s windows, AdamW lr=2**-12, cosine LR (original's T_max-per-epoch quirk reproduced), batch 16, 2000 epochs, the original's S&R augmentation, and checkpoint = best-validation accuracy on a 20 % hold-out of session 1 (primary), plus a pre-registered last-epoch control. The paper's own best-test rule is shown for context only.

Already verified: numerical parity is exact (logits/input-gradient max-abs diff 0.0, 45/45 state entries, 20303 params), so the port is the same function as Whit3Zhao/TMSA-Net@c60882db. The current head only merges master; tmsanet.py is unchanged. One disclosed training-time note: the port's FeedForwardBlock reuse drops the original's second FFN dropout (eval-neutral).

TMSA-Net (#1209), BCIC-IV-2a, 9 subjects x 5 seeds, port vs original code (matched init, same cached windows):

selection port (braindecode #1209) original code paper gap port vs paper
best-val (PRIMARY) 0.6749 (45/45) 0.6781 (45/45) 0.8322 -18.9%
last-epoch (CONTROL) 0.6819 (45/45) 0.6744 (45/45) 0.8322 -18.1%
best-test (CONTEXT, paper's own rule) 0.7431 (45/45) 0.7416 (45/45) 0.8322 -10.7%

Best-val, like-for-like on the 45 cells where both codes finished: port 0.6749 vs original 0.6781 (diff -0.0032).

Verdict (5% gate vs paper): PARTIAL: BCIC-IV-2a 9-subject mean (best-val, primary; 45/45 port, 45/45 original cells) port 0.6749 vs paper 0.8322 = -18.9% (outside 5%); like-for-like on the 45 cells where both codes finished: port 0.6749 vs original 0.6781 (diff -0.0032). Documented cause: declared substitute protocol. The paper's 2a preprocessing is a private .mat set (4-38 Hz MOABB windows substituted); the paper selects its number as the max TEST accuracy over 2000 epochs (we select on a 20% validation hold-out of session 1, which also removes 20% of the training data); its seed is random and unpublished.

No action needed from you. If the exact 2a preprocessing script or a fixed seed exists outside the published repo, it would let us tighten the substitute protocol.

@bruAristimunha bruAristimunha added model Adds a new model needs-replication Model PR: paper number must be replicated (NeuralBench) before merge labels Oct 5, 2026
Resolved NOTICE.txt
docs/whats_new.rst by keeping both entries.
NOTICE.txt: master's MIT file list plus braindecode/models/tmsanet.py.
Fold the feature extractor, transformer wrapper and local-key conv module into
their callers and use einops layers for the head split/merge and transposes.
Drop guards torch already enforces; keep embed_dim >= num_heads, which otherwise
builds an empty attention. The attention stays an explicit softmax so outputs,
gradients and the dropout RNG order match the reference exactly.
Tests: one test of the three released configurations, including the 19 -> 16
head bottleneck; the generic suites cover the rest. Shorter whats_new entry.

@bruAristimunha bruAristimunha left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified: TMSANet is bit-identical to the authors' code, passes master's model contract and integration suites including the new CPU accelerator and low-precision checks, and with online Euclidean alignment (a declared deviation) the BCI IV 2a result lands within 5% of the paper. Thanks @lindicaphxag-tech for TMSA-Net!

@bruAristimunha
bruAristimunha merged commit 1c5f213 into braindecode:master Oct 9, 2026
16 checks passed
bruAristimunha added a commit to lindicaphxag-tech/braindecode that referenced this pull request Oct 9, 2026
…raindecode#1253) into EEG-CLIP; EEGCLIP text side in _UNUSED_IN_FORWARD
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

model Adds a new model needs-replication Model PR: paper number must be replicated (NeuralBench) before merge

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants