Add CSBrain cross-scale spatiotemporal brain foundation model by qinxwew · Pull Request #1196 · braindecode/braindecode

qinxwew · 2026-09-30T07:59:25Z

Implements CSBrain [zhou2025csbrain] (NeurIPS 2025 Spotlight) as an EEGModuleMixin model, as claimed in #1077.

What

models/csbrain.py: CBraMod-style 200-sample patching + three stacked mechanisms per encoder layer (12 layers by default):
- cross-scale temporal embedding — multi-scale convolutions (kernels 1/3/5 over the patch axis) with an exponentially decaying channel split, added residually;
- region embedding — per-region convolutions (kernels 1/3/5, circular padding over the electrode axis) mixing each anatomical region locally;
- structured sparse attention — inter-window attention over sliding windows of 5 patches plus inter-region attention restricted by a round-robin group mask (each electrode attends to at most one electrode per other region), with per-region mean-pooled global features.
Channel names (chs_info) are mapped to five anatomical regions (frontal/parietal/temporal/occipital/central, the authors' canonical full-montage rule) and reordered to be contiguous; an explicit brain_regions parameter overrides the name derivation — this is how the authors' per-dataset layouts (e.g. their BCI-IV-2a config) are reproduced exactly. Without channel information the model degenerates to a single region and keeps full attention instead of failing.
Standard registrations: models/__init__.py, summary.csv (10,962,402 params at CSBrain(n_outputs=2)), util.py, api.rst, architecture figure.
Targeted unit tests: region-name mapping, contiguity of the reordering, the structure of the group attention mask, degenerate (no-channel-info) construction, masked-patch forward, batch-size-1 train mode.

Correctness verification

Bit-exact against the reference implementation. Loading the authors' released pretrained checkpoint (CSBrain.pth, Google Drive) into both the reference implementation (yuchen2199/CSBrain, with their BCI-IV-2a region layout) and this port — key-remapped (encoder.layers.*→encoder.*, TemEmbedEEGLayer.*→temporal_embed.*, BrainEmbedEEGLayer.*→region_embed.*) and shape-filtered, 289/295 tensors — identical inputs produce max abs output diff 0.0. The six skipped tensors are the temporal-region blocks with no matching channel group in the 2a layout; the authors' own fine-tuning code filters the same way.

Benchmark (BCI IV-2a, same harness as #1186)

Subject-specific 5-fold CV over the 288 training trials of each subject's T session, evaluation on the 288 E-session trials — identical folds/seeds/test set as the MSCFormer/CTNet rows below. CSBrain is fine-tuned with the authors' recipe (pretrained init, AdamW lr 1e-4 wd 0.01, cosine to 1e-6, batch 64, label smoothing 0.1, drop 0.3, ~600 updates, selection on best validation accuracy; their BCI-IV-2a region layout via brain_regions).

model	reference acc	this harness	harness gap
MSCFormer (paper 0.8295)	0.8295	0.7434	−8.6 pp
CTNet (paper 0.8252)	0.8252	0.7068	−11.8 pp
CSBrain (authors' released 2a ckpt 0.5773)	0.5773	0.4384	−13.9 pp

Per-subject (CSBrain, mean±std over 5 folds): s1 .442±.033 · s2 .330±.006 · s3 .590±.028 · s4 .399±.045 · s5 .326±.026 · s6 .354±.026 · s7 .551±.020 · s8 .421±.016 · s9 .533±.024 (kappa .38 mean).

Reading the numbers honestly:

The cross-session harness (train on T, evaluate on E) is harder than the within-split protocols behind all three reference numbers; every model drops 9–14 pp, and CSBrain's −13.9 pp is in line with that systematic gap. Its own released fine-tuned checkpoint embeds acc_0.57726 in the filename, so ~0.58 is the authors' own level on 2a — the port adds nothing on top of that (bit-exactness) and subtracts nothing.
Task-specific compact models dominate this small dataset: 288 training trials per subject are little fine-tuning signal for a 23M-parameter foundation model (14M of which are the from-scratch task head); MSCFormer wins 45/45 paired folds. Best-epoch median was 27/150, so the gap is not an under-training artifact. A band-passed + per-epoch-z-scored input variant (closer to a LaBraM-style pipeline) was also tried and scored lower (µV amplitudes match the pretraining distribution).
The comparable-row design (same folds/seeds/test set as the MSCFormer/CTNet rows) is what makes the comparison meaningful; scripts and per-fold CSV are linked below for reproduction.

License note

The upstream repository ships no LICENSE file at the time of writing (I opened yuchen2199/CSBrain#7 asking the authors to add one). This file is an independent reimplementation in the EEGModuleMixin idiom following the paper (and the already-ported CBraMod for the shared patch-embedding design); the reference implementation was used only as a numerical cross-check (the bit-exactness section above). If the authors add a permissive license, a follow-up can switch the attribution to "adapted from" with a NOTICE.txt entry.

Reproduction

harness + recipes: bench_mscformer_2a.py / bench_csbrain_2a.py (same 5-fold subject-specific protocol, CSV-resumable; fine-tuning settings documented in the script headers)
per-fold results: results_2a.csv (model, subject, fold, acc, kappa, best_epoch, seconds)

Checklist

final_layer naming, activation class-default parameter, drop_prob
registrations (__init__, summary.csv, util.py, api.rst, architecture figure)
docs/whats_new.rst entry (added with the PR number)
bit-exactness verification against the reference implementation
benchmark under the shared harness with gap analysis

Closes #1077

Reimplementation of CSBrain (Zhou et al., NeurIPS 2025 Spotlight, arXiv:2506.23075) as an EEGModuleMixin model, verified bit-exact against the authors' reference implementation with their released pretrained checkpoint (max abs output diff 0.0 on matched inputs). - cross-scale temporal embedding (k=1/3/5), per-region circular-padded embedding and structured sparse attention (inter-window + masked inter-region) - channel names mapped to five anatomical regions, or an explicit brain_regions sequence (e.g. to reproduce the authors' per-dataset layouts); without channel info the model degrades to a single region - standard registrations: __init__, summary.csv, util.py, api.rst, architecture figure; targeted unit tests for region derivation, attention-mask structure and degenerate construction Closes braindecode#1077

codecov · 2026-09-30T09:45:43Z

Codecov Report

❌ Patch coverage is 98.13433% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.46%. Comparing base (bc80538) to head (19a61c6).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1196      +/-   ##
==========================================
+ Coverage   88.32%   88.46%   +0.13%     
==========================================
  Files         155      156       +1     
  Lines       18617    18885     +268     
==========================================
+ Hits        16443    16706     +263     
- Misses       2174     2179       +5

🚀 New features to boost your workflow:

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

# Conflicts: # docs/whats_new.rst

qinxwew · 2026-10-01T17:01:08Z

Merged the latest master (which brought the LaBraM #1155 fix and the two new tutorials) to clear the docs/whats_new.rst conflict. Both changelog entries are kept — #1155 first, the CSBrain entry right below it. The PR diff is unchanged: 9 files, +806. CI is re-running.

…model_compiled The port also Kaiming-initialised every Conv2d (temporal/region embedding, patch stem, positional conv), which the reference _weights_init does not. The residual stream then grows ~3x per layer and the default 12-layer model emits logits of ~1e4-1e5 at init, so Inductor's float32 reassociation error (~0.1 absolute) exceeds allclose(atol=1e-4) on some seeds/platforms (test_model_compiled[CSBrain] on macOS/3.13). No graph breaks are involved. Also mark the cls_token=None return with nosec B105 (Codacy/Bandit false positive, same as BrainBERT/Brant/MIRepNet).

The reference fine-tuning heads are flatten -> n_patch * 200 -> 200 -> n_outputs (e.g. 2000 hidden units for the 10 s FACED/CHB-MIT/Siena windows); the port fixed the first hidden layer at 4 * emb_dim, which matches only 4 s windows. The lazy head (no n_times) keeps 4 * emb_dim. Also correct two comments: the name-based region rule reproduces the reference PhysioNet-MI/FACED/SHU-MI/... layouts, while its BCI-IV-2a, SEED-V, SEED-VIG and Siena layouts put FC in central and PO in parietal; and without chs_info the region embedding is skipped (unmasked attention).

Keep every case of the per-model file in a CSBrain section of the shared suite, and add regression tests for the reference init (bounded residual stream), the window-dependent head width, the PhysioNet-MI name layout and the BCI-IV-2a brain_regions order. The two behaviour tests fail on c561f8a.

The reference fine-tuning models (CHB-MIT, Siena, SEED-V, ...) reorder the channels with a hand-made sorted_indices that sets the electrode ring inside each region; the stable grouping of the port differs, which changes the circular region convolution and the attention groups (random-weight parity vs the official code: 90.9 abs logit diff on Siena, 7025 on the CHB-MIT backbone). channel_order takes that permutation (validated: a permutation that groups the regions ascending) and gives 0.0 difference on both.

Review follow-up. The window-dependent head only applied when n_chans and n_times were passed explicitly; with chs_info or input_window_seconds + sfreq it fell back to the lazy 4 * emb_dim head (800 wide for 10 s windows, so the reference checkpoints could not load). Resolve both through the mixin. head_hidden_dim overrides the first hidden width for reference heads that are not n_patch * 200 wide (SEED-V: 800 for 1 s windows) and bounds the head for long windows. brain_regions is checked against n_chans; reset_head uses _set_n_outputs so the saved config follows the head.

bruAristimunha · 2026-10-05T09:34:56Z

Integration gate (braindecode maintainers)

Target: paper Table 2, PhysioNet-MI (0.6304 ± 0.0090) and BCIC-IV-2a (0.5657 ± 0.0071, test bal. acc., subjects 8-9), reproduced from the released weights + public MOABB data via NeuralBench, head 19a61c6 (add-csbrain-model). Gate: within 5% of the paper's number (or the paper's spread if larger).

Protocol (2a): the data preparation the paper used (CBraMod preprocessing_bciciv2a.py @ 92787b3, the version that yields the paper's 5,088 samples: causal order-5 Butterworth 0.3-50 Hz, [2, 6] s after trial start, per-trial CAR, 200 Hz), the reference fine-tuning model (model_for_bciciv2a.py: backbone dropout 0.1, head dropout 0.3, proj_out → Identity), finetune_CSBrain_BCI2a.sh hyperparameters, 5 seeds, best-validation epoch.

Verified: forward parity port-vs-original is exact (max-abs diff 0.0) on both datasets; 19a61c6 is forward-identical to the heads measured before. The released 2a fine-tuned checkpoint scores 0.577 on the full test set (A08+A09), matching its file name; it was trained on subjects 1-5 and its epoch was picked on validation, so there is no leak. (An earlier "0.667-0.719" came from a 96-trial debug subset and should be ignored.)

Results (updated 2026-10-06):

cell	port (5 seeds)	authors' code, authors' data prep (5 seeds)	paper	port vs paper
PhysioNet-MI	0.6135 ± 0.0045	—	0.6304 ± 0.0090	−2.7% ✅
BCIC-IV-2a	0.5146 ± 0.0356	0.5062 ± 0.0073	0.5657 ± 0.0071	−9.0% ❌

Verdict: the port matches the reference implementation; the 2a paper number does not reproduce upstream. Running the authors' released finetune_main.py unmodified on the authors' own preprocessing of the same raw files gives 0.506, and the port is within 1.7% of that. The remaining gap to 0.5657 is outside braindecode and NeuralBench, so it goes to the maintainers to decide. Nothing needed from you for this.

Small API note, not blocking: braindecode's single drop_prob sets both the backbone and the head dropout. The reference fine-tuning scripts use 0.1 in the backbone and --dropout in the head (now documented in the docstring at 19a61c6, thanks).

Licence: CSBrain ships no LICENSE file upstream, so no checkpoint mirror on the Hub until that is settled.

qinxwew · 2026-10-05T12:26:06Z

Thanks for running the integration gate — the exact parity check (max-abs diff 0.0 on both datasets, including after remapping the FeedForwardBlock keys) and the PhysioNet-MI number (0.6135 ± 0.0045 vs paper 0.6304) are exactly the kind of external validation this needed. Noted on the 2a epoch-alignment fix too.

On the LICENSE question — agreed, and here is where it stands:

This PR is an independent re-implementation: it ships no weights, no checkpoint files, and no from_pretrained loader. So the code itself is not redistribution of upstream material and isn't gated on upstream's terms.
Upstream still ships no LICENSE. I opened yuchen2199/CSBrain#7 on 2026-09-30 asking them to add an open-source license (MIT / Apache-2.0 / BSD) — no reply yet.
So I'll hold your line: no checkpoint mirror on the Hub until redistribution terms are confirmed. That supersedes my earlier "assume BSD or unlicense" framing.

Once #7 lands, I have the original CSBrain.pth plus a key-mapped state_dict and both SHA-256 hashes ready to hand over for re-hosting.

One concrete gap worth flagging on the checkpoint side. On f31f25f the encoder params are now encoder.N.ff_block.0.* / encoder.N.ff_block.3.* (48 tensors, zero linear1/linear2 left), whereas the released CSBrain.pth carries linear1/linear2. Since this PR ships no loader, anyone dropping the released weights in by hand will hit a key mismatch. Happy to add either:

a small from_pretrained-style helper with the linear1→ff_block.0 / linear2→ff_block.3 remap baked in, or
a short doc note recording the mapping.

Which would you prefer? (1) is a bit more code but makes the released weights drop-in; (2) keeps the PR lean. I'll go either way.

# Conflicts: # test/unit_tests/models/test_models.py

…SBrain Review follow-up on the integration gate note: the reference fine-tuning scripts keep the backbone at 0.1 and raise only the head dropout, which a single drop_prob cannot express; document the wrapper/subclass recipe.

qinxwew · 2026-10-06T04:16:03Z

Thanks @bruAristimunha for the updated gate — good to see 2a move once the fine-tuning config matches models/model_for_bciciv2a.py (backbone dropout 0.1, proj_out swapped for nn.Identity()), and for pinning the remaining gap to the fine-tuning side rather than the model.

Two follow-ups on the new head:

Merge: master's model additions (Add iEEG MAPA model #1178, Add BrainTokenizer, the BrainOmni EEG/MEG VQ-VAE tokenizer #1230, [models] Add SeizureTransformer #1236) conflicted with the CSBrain test section in test/unit_tests/models/test_models.py; merged origin/master in and kept both sections (SeizureTransformer first, CSBrain after). whats_new.rst and the other registration points merged cleanly.
API note: added a docstring note on drop_prob documenting that it sets both the backbone and the head dropout, and that the reference fine-tuning recipe (backbone 0.1, head --dropout 0.3 for 2a) needs a thin subclass/wrapper for the head. I kept the single-parameter API unchanged for now — happy to add a separate head_drop_prob argument instead if you prefer that direction.

csbrain.py logic is unchanged in this push (merge + docstring only), so the parity checks on f31f25f remain valid for the model itself.

qinxwew · 2026-10-06T04:46:01Z

CI note on test (windows-latest, 3.12) (head 19a61c6e): the job aborted with a runner-side ##[error]There is not enough space on the disk at ~62% of the suite. The tests it marked FAILED just before dying are all save/load round-trip tests (test_return_features.py for BENDR/CBraMod/EEGDINO/LaBraM, test_signal_jepa.py HF round-trips) — none of them touched by this PR, and all of them disk-write-heavy, which matches the disk-exhaustion timing. The same head passed on windows-latest, 3.13, pointing at the runner rather than the code. Fork accounts can't re-run workflow jobs, so documenting here instead.

bruAristimunha

Thanks for the port and the patience! Replication on Voyager via NeuralBench: on BCIC-IV-2a this port gives 0.515, matching the authors' own released code run with their own data preparation (0.506 +/- 0.007); the paper's 0.566 is not reached by the released code itself (no train/test leak, no test-set selection), so the port is faithful to the released system. PhysioNet-MI 0.6135 vs 0.6304 (-2.7%). CI green.

qinxwew added 3 commits September 30, 2026 12:49

Exempt CSBrain from torch.jit.script test like CBraMod/CodeBrain

8eb7a0c

Add whats_new entry for CSBrain PR

8a472cd

qinxwew force-pushed the add-csbrain-model branch from c75645a to 8a472cd Compare September 30, 2026 08:01

qinxwew mentioned this pull request Sep 30, 2026

Include the CSBrain #1077

Closed

Merge remote-tracking branch 'origin/master' into add-csbrain-model

c561f8a

# Conflicts: # docs/whats_new.rst

bruAristimunha added 5 commits October 1, 2026 21:09

bruAristimunha mentioned this pull request Oct 2, 2026

Integrate new EEG foundation models and supervised baselines #1201

Open

bruAristimunha added 2 commits October 4, 2026 15:52

Merge remote-tracking branch 'origin/master' into add-csbrain-model

af910a0

FIX reuse FeedForwardBlock for CSBrain encoder's ff_block

f31f25f

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

qinxwew added 2 commits October 6, 2026 12:06

Merge remote-tracking branch 'origin/master' into add-csbrain-model

7f02458

# Conflicts: # test/unit_tests/models/test_models.py

bruAristimunha approved these changes Oct 6, 2026

View reviewed changes

bruAristimunha merged commit c282258 into braindecode:master Oct 6, 2026
19 of 20 checks passed

bruAristimunha mentioned this pull request Oct 6, 2026

CSBrain: reuse CBraMod's patch embedding; fix CBraMod patch_size != 200 #1240

Merged

qinxwew mentioned this pull request Oct 6, 2026

[No Code] More models for braindecode #375

Open

This was referenced Oct 7, 2026

CSBrain: add head_drop_prob; document checkpoint key names #1247

Merged

MAINT CI: delete passing tests' tmp_path (Windows out of disk); acceptance on 3.13 #1251

Open

Navigation Menu

Repository navigation

Add CSBrain cross-scale spatiotemporal brain foundation model - #1196

Add CSBrain cross-scale spatiotemporal brain foundation model#1196
bruAristimunha merged 13 commits into
braindecode:masterfrom
qinxwew:add-csbrain-model

qinxwew commented Sep 30, 2026

Uh oh!

codecov Bot commented Sep 30, 2026 •

edited

Loading

Uh oh!

qinxwew commented Oct 1, 2026

Uh oh!

bruAristimunha commented Oct 5, 2026 •

edited

Loading

Uh oh!

qinxwew commented Oct 5, 2026

Uh oh!

qinxwew commented Oct 6, 2026

Uh oh!

qinxwew commented Oct 6, 2026

Uh oh!

bruAristimunha left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

qinxwew commented Sep 30, 2026

What

Correctness verification

Benchmark (BCI IV-2a, same harness as #1186)

License note

Reproduction

Checklist

Uh oh!

codecov Bot commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

qinxwew commented Oct 1, 2026

Uh oh!

bruAristimunha commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

qinxwew commented Oct 5, 2026

Uh oh!

qinxwew commented Oct 6, 2026

Uh oh!

qinxwew commented Oct 6, 2026

Uh oh!

bruAristimunha left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Sep 30, 2026 •

edited

Loading

bruAristimunha commented Oct 5, 2026 •

edited

Loading