Skip to content

Add BrainOmni + BrainTokenizer (unified EEG/MEG foundation model) - #1043

Closed
bruAristimunha wants to merge 70 commits into
masterfrom
feat/brainomni-model
Closed

bruAristimunha wants to merge 70 commits into
masterfrom
feat/brainomni-model

Conversation

@bruAristimunha

@bruAristimunha bruAristimunha commented Jun 10, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Adds two models derived from BrainOmni: A Brain Foundation Model for Unified
EEG and MEG Signals
(Xiao et al., NeurIPS 2025;
arXiv:2505.18185; upstream
OpenTSLab/BrainOmni, MIT):

  • BrainTokenizer — the SEANet/channel-compression/residual-VQ tokenizer,
    registered as a non-classification model.
  • BrainOmni — the downstream classifier: frozen tokenizer inference,
    factored spatio-temporal transformer, released mean-pooling path, and
    classification head.

Key contracts:

  • Strict public loaders accept the authors' raw tokenizer, tiny, and base
    checkpoints. Immutable artifact revisions and SHA-256 values are pinned in
    the integration tests; no checkpoint is copied or re-hosted by this PR.
  • Pinned-source comparisons cover tokenizer initialization/quantization and
    released numerical signatures. The deterministic RoPE cache is regenerated
    from its authoritative frequency buffer because the DeepSpeed export lost
    its complex phase.
  • Geometry is derived once from chs_info for EEG, MAG, planar GRAD, and axial
    GRAD. Real MNE types use public mne.channel_type; lightweight ch_type
    metadata and real-MNE metadata both round-trip through config and local Hub
    serialization with exact state/output parity.
  • Fresh codebooks use the released first-batch K-means initialization.
    Distributed EMA initialization and sufficient statistics are synchronized
    across ranks.
  • Token extraction temporarily enters evaluation mode and restores the
    tokenizer's prior mode; projection, transformer blocks, and head remain
    trainable during downstream fine-tuning.
  • The historical positional PatchTokenizer API remains valid, short and odd
    SEANet cases follow the source padding policy, invalid public configurations
    fail early, and the implementation remains compatible with the declared
    PyTorch 2.4 floor, using native RMSNorm directly.

The released source preprocessing differs from the paper's prose at its final
normalization step. Model documentation records that discrepancy and tells
users to reproduce the released source pipeline when using released weights.

Provenance and licensing

  • Audited upstream source commit: 340d6b5aba886af76b217272cdb3251651e9bf16.
  • Audited OpenTSLab artifact revision:
    9a4d3c70495370397ccfbfd6d2496f25647545a5.
  • LICENSES/BrainOmni-MIT.txt preserves the OpenTSLab and Meta/EnCodec MIT
    grants; LICENSES/vector-quantize-pytorch-MIT.txt preserves the additional
    derived quantizer grant; NOTICE.txt records both components.
  • Wheel and sdist gates verify that the project BSD license, NOTICE, and both
    MIT files are packaged. No external weight is bundled.

The official OpenTSLab repositories contain raw .pt state dictionaries, not
Braindecode config.json plus safetensors repositories. The documentation
therefore shows from_opentslab_config(...) followed by strict
load_state_dict(...); it does not claim the upstream repository works with
Braindecode from_pretrained(...).

Verification

  • All three immutable released artifacts strict-load through the public
    models; tokenizer/tiny numerical signatures match the pinned source.
  • PyTorch 2.4.0: 89 model, geometry and quantizer tests passed (3 skips).
  • Removed obsolete RMSNorm and compiler-namespace compatibility fallbacks.
  • Two-rank distributed EMA synchronization and bounded first training
    update.
  • Real-MNE and lightweight EEG/MAG/GRAD geometry, config, and strict local
    Hub round trips for both public models.
  • Foundation/modules, model integration, config/Hub, return-features,
    compile/export, reset-head, and invalid-boundary suites.
  • Ruff, mypy, isort, codespell, Sphinx lint, build, Twine, archive, and
    packaged-license gates.
  • Independent review at exact c8b01ab478220a2ea8460bbc064c6f1bd61753c1
    found no actionable issue.

The PR remains a draft while fresh exact-SHA CI runs. Publishing converted
Braindecode-format weights, if desired later, is a separate maintainer action
and is not part of this PR.

Adds _BrainSensorModule, _ForwardSolution, _BackwardSolution,
_BrainTokenizerEncoder, and _BrainTokenizerDecoder to brainomni.py,
faithfully ported from upstream BrainOmni/model_utils/module.py with
_ prefixes and upstream state-dict key alignment preserved.
Adds scripts/convert_brainomni_weights.py to convert OpenTSLab/BrainOmni
upstream .pt checkpoints into braindecode hub format (config.json +
safetensors).  Includes exhaustive key-coverage assertions, geometry-buffer
awareness (pos/sensor_type), round-trip validation via save_pretrained /
from_pretrained, optional numerical parity mode, and an offline --self-test
that validates the full remap logic without network access.
Copilot AI review requested due to automatic review settings June 10, 2026 14:50

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b8ad7899dd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread braindecode/models/brainomni.py Outdated
Comment on lines +64 to +66
coil = str(ch.get("coil_type", ""))
if "PLANAR" in coil or "GRAD" in coil:
return _SENSOR_GRAD

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Classify integer-coded gradiometers correctly

When chs_info comes from real MNE objects, MEG channels usually have integer kind and integer coil_type values; planar gradiometers are not passed as kind="grad". Converting coil_type to a string here means values like 3012 never contain "PLANAR" or "GRAD", so those gradiometers are classified as magnetometers and _orientation_of uses loc[9:12] instead of the grad orientation in loc[3:6], giving incorrect sensor embeddings for MEG data.

Useful? React with 👍 / 👎.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Adds BrainOmni and BrainTokenizer foundation models to braindecode, including model registration, documentation, tests, and a maintainer conversion script to remap upstream BrainOmni checkpoints into braindecode’s Hub format.

Changes:

  • Introduces braindecode.models.BrainTokenizer (VQ-VAE tokenizer) and braindecode.models.BrainOmni (downstream classifier) with EEG/MEG geometry handling.
  • Adds extensive unit tests for geometry derivation, core blocks, tokenizer behavior, and BrainOmni classifier behavior.
  • Adds a conversion script for upstream checkpoints and updates docs/model registry/CI integration expectations (compile/TorchScript skips).

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
braindecode/models/brainomni.py New BrainOmni + BrainTokenizer implementations, geometry derivation, attention/SEANet/VQ components
braindecode/models/__init__.py Exposes new models at braindecode.models.*
braindecode/models/util.py Registers init-parameter expectations + marks BrainTokenizer as non-classification output
braindecode/models/summary.csv Adds model-zoo metadata entries for BrainOmni/BrainTokenizer
test/unit_tests/models/test_foundation_models.py Adds unit tests for BrainOmni/BrainTokenizer + helpers
test/unit_tests/models/test_integration.py Skips torch.compile/TorchScript for BrainOmni(/Tokenizer) with rationale
scripts/convert_brainomni_weights.py New conversion/validation script for upstream checkpoints
docs/api.rst Adds BrainOmni/BrainTokenizer to API docs listings
docs/whats_new.rst Changelog entry announcing the new models

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread braindecode/models/brainomni.py Outdated
Comment on lines +64 to +67
coil = str(ch.get("coil_type", ""))
if "PLANAR" in coil or "GRAD" in coil:
return _SENSOR_GRAD
return _SENSOR_MAG
Comment thread braindecode/models/brainomni.py Outdated
Comment thread braindecode/models/brainomni.py Outdated
Comment on lines +1724 to +1727
Pre-trained weights for ``BrainTokenizer`` are available on
`HuggingFace <https://huggingface.co/OpenTSLab/BrainOmni>`_
(see the repository for the *tiny* and *base* tokenizer checkpoints).
Load with ``BrainTokenizer.from_pretrained("OpenTSLab/BrainOmni")``.
Comment thread braindecode/models/brainomni.py Outdated
@bruAristimunha
bruAristimunha marked this pull request as draft June 10, 2026 15:00
bruAristimunha and others added 15 commits June 10, 2026 22:04
Resolve conflicts by keeping both sides' additions (BrainOmni/BrainTokenizer
and EEGDINO) in summary.csv, util.py, whats_new.rst, and test_modules.py.

Generalize PatchTokenizer with overlap/stride support and dynamic forward-time
padding, and reuse it for BrainTokenizer._unfold.
- EEGModuleMixin: render a braindecode model card (description, usage,
  citation) on save_pretrained/push_to_hub, and warn when a strict=False load
  leaves missing/unexpected keys (the silent mis-load on the from_pretrained
  path).
- RotaryPositionalEmbedding: clarify that n_dim is the full attention dim
  (load-bearing for checkpoint parity), not the per-head dim.
- quantization: credit lucidrains vector-quantize-pytorch (MIT) and note the
  upstream-checkpoint state-dict parity constraint.
…onventions)

Apply the braindecode model-convention review to BrainOmni/BrainTokenizer:

- Docstrings (B8): canonical "<Name> from Xiao et al. (2025) [brainomni]_"
  headers + ".. versionadded:: 1.6.1".
- Architecture docs (B8b): add Architecture Overview / Macro Components
  (Model.attr + Operations/Role) / Temporal-Spatial-Spectral / Additional
  Mechanisms rubric sections to both classes.
- Weight init (B9): shared _init_weights applied via self.apply, faithful to
  upstream (Linear/Embedding -> trunc_normal_ std=0.02, RMSNorm -> 1.0);
  weight_norm convolutions left untouched.
- Pretrained weights (B10): point from_pretrained at the converted
  braindecode/brainomni-pretrained and braindecode/braintokenizer-pretrained
  repos (raw upstream OpenTSLab keys, decoder.* vs final_layer.*, would
  silently mis-load under strict=False).
- Signature (B3/A2): "# braindecode parameters" / "# model-specific parameters"
  comment separators + "*" keyword-only model params; del all six signal args.
- Correct the frozen-backbone description: there is no train() override; the
  tokenizer is hard-frozen via torch.no_grad in BrainTokenizer.tokenize.

Tests: 163 passed, 9 skipped (unchanged from baseline).
# Conflicts:
#	braindecode/models/summary.csv
#	braindecode/models/util.py
#	braindecode/modules/blocks.py
#	test/unit_tests/models/test_foundation_models.py
#	test/unit_tests/models/test_integration.py
#	test/unit_tests/models/test_modules.py
@bruAristimunha

Copy link
Copy Markdown
Collaborator Author

I completed a maximum-scope and paper-gate audit of this draft. The exact live
head should remain provenance evidence only; it should not be rebased or merged
as one 63-commit, 21-file, +4,127-line candidate.

The smallest defensible replacement stack is:

  1. one shared license/archive prerequisite, only if still needed;
  2. one BrainTokenizer architecture/raw-checkpoint PR, keeping BrainOmni-only
    geometry, RoPE, VQ, and SEANet helpers private;
  3. one classifier-only BrainOmni PR stacked on the exact tokenizer commit;
  4. an external exact-final-SHA PhysioNet-MI gate;
  5. tokenizer then unchanged-classifier merges only after that gate returns a
    schema-valid PASS.

The qualifying target is BrainOmni-tiny, PhysioNet-MI Table 2 balanced accuracy
0.580 +/- 0.019, with 30 fold/seed/LR cells, the official padded-sampler
metric, and a clean unique-example sensitivity. This is a prespecified target,
not a reproduced result. No Margaret/Iceberg data or checkpoint was downloaded:
the experiment must not start until the final replacement classifier SHA is
reviewed, published as a draft, and frozen.

Please drop the global base ch_type change, shared PatchTokenizer stride
change, public RoPE/VQ APIs, generic geometry utilities, broad README/docs, and
the 1,342-line shared-test expansion. Keep raw checkpoint loading beside each
architecture, use one focused test_brainomni.py, and keep conversion,
training, rehosting, and reproduction code outside Braindecode. Scope remains
frozen until the external result gate passes.

Audited head: c8b01ab478220a2ea8460bbc064c6f1bd61753c1.

@codecov

codecov Bot commented Sep 21, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 91.63803% with 73 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.96%. Comparing base (55afc12) to head (bf22a1b).
⚠️ Report is 15 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #1043      +/-   ##
==========================================
+ Coverage   87.78%   87.96%   +0.17%     
==========================================
  Files         149      151       +2     
  Lines       17461    18324     +863     
==========================================
+ Hits        15329    16119     +790     
- Misses       2132     2205      +73     
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Resolve the six conflicts with master 55afc12 by keeping both sides:
- models/__init__.py: import and export BrainOmni/BrainTokenizer next to
  BrainBERT/Brant.
- models/util.py: keep the BrainOmni/BrainTokenizer mandatory-parameter
  entries next to DIVER1, keep the BrainOmni geometry helpers next to
  master's resolve_channel_indices, and take master's positions_from_chs_info
  signature.
- docs/api.rst: list BrainOmni after BrainBERT.
- docs/whats_new.rst: master's entries followed by the BrainOmni entry.
- test_return_features.py: sorted imports (BaRISTA, BrainBERT, BrainOmni,
  Brant).
- test_foundation_models.py: union of imports (master dropped reve.RMSNorm,
  unused here), master's new DIVER1/ZUNA/STEEGFormer tests followed by the
  BrainOmni block.
…dBlock)

- RotaryPositionalEmbedding: rotate pairs with functional.rotate_pairs in
  float32 instead of a cached complex buffer. No buffers, so no cast or
  state-dict repair logic; the per-head frequency split is unchanged.
- The loader drops the released RoPE keys: the export stores freqs rounded
  to bfloat16 and the complex cache without its sine part. The port now uses
  the exact float32 rotation the upstream model computes at construction.
  The tiny-checkpoint signature is updated to that upstream path.
- Replace the private _FeedForward with modules.FeedForwardBlock (same
  Linear-SELU-Linear-Dropout maths; key remap ff.layer.{0,2} -> ff.{0,3}).
- license="mit" on both classes, _set_n_outputs in reset_head, and the
  per-file MIT entries in NOTICE.txt; LICENSES/ is kept.

Random-weight and released-checkpoint parity against OpenTSLab/BrainOmni
340d6b5a: encoder max |diff| <= 5.1e-7, tokenizer 0, 0 index mismatches.
….decoder rename

Review (REVIEW.md) MINOR: the ff/aggregate_mlp renaming and the dropped RoPE
and pretraining keys were only covered by network-marked tests. NIT: the
tokenizer.decoder. rename matched anywhere in a key.
The first forward of a fresh BrainOmni runs K-means over 4096 x 512 x 256
broadcast differences (~2 GB per intermediate); test_local_push_and_pull_roundtrip
peaked at 5.2 GB RSS and the ubuntu CI runners died with a shutdown signal in
that test twice. Chunking the samples keeps every distance bit-identical
(checked) and peaks at 1.5 GB.
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