Add BrainOmni, the downstream classifier on BrainTokenizer by bruAristimunha · Pull Request #1231 · braindecode/braindecode

bruAristimunha · 2026-10-05T13:12:52Z

Split 3/3 of #1043, stacked on #1230 (its diff shrinks to the classifier once #1230 merges). Adds BrainOmni, the downstream classifier on a frozen BrainTokenizer (RoPE spatiotemporal transformer blocks + pooled head), strictly loading the authors' raw tiny and base checkpoints. reset_head goes through _set_n_outputs/_update_init_kwargs so a re-serialized model reports its real head. license="mit".

Licence: MIT.
Replication: PhysioNet-MI test balanced accuracy 0.5796 ± 0.0164 vs the paper's 0.580 ± 0.019 (BrainOmni tiny, released protocol, 30/30 cells).

BrainOmni is ported from OpenTSLab/BrainOmni (MIT); its SEANet codec from Meta's EnCodec (MIT) and its residual VQ from lucidrains/vector-quantize-pytorch (MIT). Per the model-port convention a port links to its licence rather than bundling the text, so add only the file lines for brainomni.py and the shared modules/quantization.py under the existing MIT section (the MIT link is already present) and leave pyproject.toml/MANIFEST.in untouched.

Port of the OpenTSLab BrainOmni tokenizer (MIT) with strict loading of the authors' raw BrainTokenizer.pt beside the architecture (BrainTokenizer.from_opentslab_config + load_state_dict key mapping). Sensor geometry, SEANet codec, residual VQ (EMA codebooks, K-means init, rotation trick) and the cross-attention bridges stay private in braindecode/models/brainomni.py; no public braindecode.modules or models.util API is added. Tests live in test/unit_tests/models/test_brainomni.py (Hub checks are marked network). Registration: models export, summary.csv, util registry lists, docs/api.rst, whats_new, NOTICE, TorchScript skip.

BrainOmni (MIT, OpenTSLab) wraps the frozen BrainTokenizer with the released spatial-temporal factored attention blocks and a downstream head, and strictly loads the authors' raw tiny/base Stage-2 checkpoints (pretraining head and RoPE caches dropped, RoPE recomputed in float32). The RoPE attention stays private in braindecode/models/brainomni.py; no public braindecode.modules API is added. Tests extend test_brainomni.py (Hub checks marked network). Registration: models export, summary.csv, util registry, docs/api.rst, whats_new, TorchScript skip.

Resolved NOTICE.txt by keeping both MIT entries.

Resolved docs/whats_new.rst by keeping both entries.

Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 1 · 3

Open (4)

embed_onehot is cast to dtype (the input dtype). In mixed precision this can make one_hot_sum… · New
This warning is emitted from BrainTokenizer.__init__, but the message says `BrainOmni pretrained… · New
Using assert for runtime argument validation is unsafe in library code because asserts can be… · New
The file header declares License: BSD-3, but the PR description and added NOTICE entries indicate… · New

What changed in this PR

Adds BrainOmni’s downstream classifier and its BrainTokenizer backbone, including shared residual vector quantization utilities and focused tests to ensure strict parity with the authors’ released checkpoints/configs.

Changes:

Introduce BrainTokenizer (VQ-VAE tokenizer) and BrainOmni (frozen-tokenizer classifier) with OpenTSLab config/key translation and strict checkpoint loading.
Add shared EMA-based residual vector quantization module under braindecode.modules.
Add extensive unit + network parity tests, plus documentation/API registry updates.

File	Description
`braindecode/models/brainomni.py`	New BrainTokenizer/BrainOmni implementations, config translation, strict key remapping, geometry derivation, and model internals.
`braindecode/modules/quantization.py`	New shared EMA residual vector quantization implementation used by the tokenizer.
`test/unit_tests/models/test_brainomni.py`	New focused tests (geometry, SEANet, quantization, strict-load parity) including network-gated checkpoint parity checks.
`test/unit_tests/models/test_integration.py`	Excludes BrainOmni/BrainTokenizer from TorchScript integration test.
`braindecode/models/__init__.py`	Exposes `BrainOmni` and `BrainTokenizer` in the public models API.
`braindecode/modules/__init__.py`	Exposes quantization modules in the public modules API.
`braindecode/models/util.py`	Registers model signatures and marks BrainTokenizer as non-logit-returning.
`braindecode/models/summary.csv`	Adds summary entries for BrainOmni and BrainTokenizer.
`docs/whats_new.rst`	Documents new BrainTokenizer and BrainOmni additions.
`docs/api.rst`	Adds BrainOmni/BrainTokenizer to the generated API docs.
`NOTICE.txt`	Adds attributions for BrainOmni/EnCodec/vector-quantize-pytorch derived components.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…enizer

…nosec B105

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

codecov · 2026-10-05T21:21:31Z

Codecov Report

❌ Patch coverage is 94.62500% with 43 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.27%. Comparing base (5e00a5b) to head (47945af).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1231      +/-   ##
==========================================
+ Coverage   87.99%   88.27%   +0.28%     
==========================================
  Files         151      153       +2     
  Lines       17628    18428     +800     
==========================================
+ Hits        15511    16268     +757     
- Misses       2117     2160      +43

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…via w41/brainomni-2-tokenizer)

Copilot

Copilot review overview

🟡 Changes recommended

Resetting the head on an evaluated model reactivates dropout, making inference stochastic.

Review effort: Balanced
Findings: 1 · 4

Open (5)

embed_onehot is cast to dtype (the input dtype). In mixed precision this can make one_hot_sum…
Reset head fails to preserve eval mode, causing stochastic inference · New
The file header declares License: BSD-3, but the PR description and added NOTICE entries indicate…
Using assert for runtime argument validation is unsafe in library code because asserts can be…
This warning is emitted from BrainTokenizer.__init__, but the message says `BrainOmni pretrained…

whats_new: keep both; refs for NeuroRVQ (braindecode#1218), BrainTokenizer (braindecode#1230), BrainOmni (braindecode#1231) point to the merged PRs.

bruAristimunha added 6 commits October 4, 2026 15:51

Merge branch 'master' into w41/brainomni-1-license

0f5b321

Resolved NOTICE.txt by keeping both MIT entries.

Merge branch 'w41/brainomni-1-license' into w41/brainomni-2-tokenizer

e9ca464

Resolved docs/whats_new.rst by keeping both entries.

Merge branch 'w41/brainomni-2-tokenizer' into w41/brainomni-3-classifier

328698e

Copilot AI balanced review requested due to automatic review settings October 5, 2026 13:12

bruAristimunha added needs-replication Model PR: paper number must be replicated (NeuralBench) before merge model Adds a new model labels Oct 5, 2026

Copilot AI reviewed Oct 5, 2026

View reviewed changes

Comment thread braindecode/modules/quantization.py

Comment thread braindecode/models/brainomni.py

Comment thread braindecode/models/brainomni.py

Comment thread test/unit_tests/models/test_brainomni.py

Copilot started reviewing on behalf of bruAristimunha October 5, 2026 13:36 View session

bruAristimunha added 4 commits October 5, 2026 22:32

FIX chunk the codebook K-means distances to bound memory

42cb8bf

Merge remote-tracking branch 'origin/master' into w41/brainomni-2-tok…

b1cc58e

…enizer

Merge branch 'w41/brainomni-2-tokenizer' into w41/brainomni-3-classifier

ac04123

FIX anchor BrainOmni tokenizer.decoder key rename and mark cls_token …

47945af

…nosec B105

Copilot AI balanced review requested due to automatic review settings October 5, 2026 20:46

Copilot AI reviewed Oct 5, 2026

Copilot started reviewing on behalf of bruAristimunha October 5, 2026 21:24 View session

bruAristimunha added 2 commits October 6, 2026 02:50

Merge braindecode#1230 squash commit (content already on this branch …

5923ee9

…via w41/brainomni-2-tokenizer)

Merge remote-tracking branch 'origin/master' into HEAD

76aa410

Copilot AI balanced review requested due to automatic review settings October 6, 2026 00:52

Copilot started reviewing on behalf of bruAristimunha October 6, 2026 00:52 View session

Copilot AI reviewed Oct 6, 2026

View reviewed changes

Comment thread braindecode/models/brainomni.py

bruAristimunha merged commit 6d9800c into braindecode:master Oct 6, 2026
19 of 21 checks passed

This was referenced Oct 6, 2026

BrainOmni: type CTF/KIT axial MEG gradiometers as gradiometers, as the release does #1238

Merged

Follow-up fixes from the review threads of the BrainOmni, NeuroRVQ and MAPA PRs #1243

Merged

Navigation Menu

Repository navigation

Add BrainOmni, the downstream classifier on BrainTokenizer - #1231

Add BrainOmni, the downstream classifier on BrainTokenizer#1231
bruAristimunha merged 12 commits into
braindecode:masterfrom
bruAristimunha:w41/brainomni-3-classifier

bruAristimunha commented Oct 5, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

codecov Bot commented Oct 5, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

bruAristimunha commented Oct 5, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 5, 2026

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants