You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add PopulationTransformer (PopT) iEEG/sEEG foundation model - #1105
Adds a braindecode-native port of PopulationTransformer (PopT) (Chau et al. 2024, arXiv:2406.03044), the third iEEG/sEEG model from #1097 after Brant (#1100) and BrainBERT (#1104).
What it is
Unlike a per-channel encoder, PopT is a population aggregator: its input is a set of per-electrode feature vectors (typically the frozen embeddings of a channel-level foundation model such as BrainBERT) plus each electrode's integer anatomical coordinates. A CLS token summarises the population after a stack of Transformer encoder layers.
Design
Following the braindecode convention, the per-electrode feature vector plays the role of the n_times axis, so the model keeps the standard (batch, n_chans, n_times) signature (n_chans = number of electrodes, n_times = feature dimension, 768 for BrainBERT stft features).
Electrode coordinates are read from chs_info and discretised to integer indices inside the model; without positions the electrodes fall back to distinct sequential indices.
The input embedding, spatial position encoding and Transformer are ported weight-for-weight; only the classification head is a braindecode-native addition.
Parity
The input embedding + Transformer are checked bit-exact against the upstream PtModelCustom (input embedding + spatial encoding + post-norm encoder), with the released popt_brainbert_stft weights copied over — test_popt.py::test_encoder_is_bit_exact_with_upstream (gated by POPT_SRC, skipped in CI without it). Reproduce with:
The shared model suites (config / categorization / forward / serialization / HF roundtrip) pass for PopulationTransformer.
Weights are re-hosted on braindecode/popt-pretrained (revision 50b02d6f8cb436ee81dea1dcf972ca635e9a07fd), so PopulationTransformer.from_pretrained("braindecode/popt-pretrained") works out of the box.
Released-checkpoint parity (2026-09-30, extends the synthetic-weight gate above): loaded the exact released pretrained_popt_brainbert_stft.pth (HF PopulationTransformer/popt_brainbert_stft rev b11d552, SHA-256 eae617d6...523d8ca) into upstream PtModelCustom, and the port's official mirror (braindecode/popt-pretrained, SHA-256 edba2cc8...b61ae5) into this port. The 78 tensor names shared between the two checkpoints are bit-identical (max abs diff 0.0) — the port's mirror is confirmed to be the same released b11d552 weights, renamed, not an independent retrain. CLS-token forward parity on identical synthetic input: max abs diff 8.3e-7 (float32 noise). No port-code divergence found; nothing pushed this session.
Real-data Volume replication on Brain Treebank. The released BrainBERT STFT and PopT checkpoints are used on the paper's 7 held-out subjects. Settings follow the released code: 80/10/10 split and AdamW for 2000 steps with the paper's schedule. Frozen BrainBERT features are cached per subject and shared by both codes, so the original code and this port see the same features, splits, head initialisation and seeds (42, 43, 44). Runs used Gaudi HPU in FP32.
Subject
Port AUC
Original AUC
Δ
1
0.947
0.948
−0.002
2
0.962
0.955
+0.007
3
0.940
0.939
+0.001
4
0.898
0.901
−0.003
6
0.817
0.835
−0.019
7
0.915
0.913
+0.002
10
0.942
0.951
−0.009
Mean ± SEM
0.917 ± 0.019
0.920 ± 0.016
−0.003
The port matches the original code. The largest per-subject gap is subject 6, which also has the largest spread across seeds. Both land above the paper's 0.87 ± 0.03. The setup differs from the paper in three ways:
The per-subject top-90 electrode lists are unpublished, so each subject uses the first 90 clean_laplacian electrodes.
Faithful re-implementation of the upstream PopT input embedding, spatial
sinusoidal position encoding, feature-reconstruction head and classification
head as standalone braindecode modules (no new runtime dependency; the encoder
is a stock torch.nn.TransformerEncoder).
braindecode-native port of PopT (Chau et al. 2024), a self-supervised
population aggregator over intracranial electrodes. The per-electrode feature
vector plays the role of the n_times axis so the model keeps the standard
(batch, n_chans, n_times) signature; electrode coordinates are read from
chs_info. Supports from_pretrained, reset_head and return_features.
Standard forward/feature/head/coord checks plus a bit-exact parity gate against
upstream PtModelCustom (gated by POPT_SRC, skipped in CI without it). Skip the
polymorphic-return model in the TorchScript export suite.
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
Adds a braindecode-native implementation of PopulationTransformer (PopT), an iEEG/sEEG population-level Transformer that aggregates per-electrode feature vectors plus discretized anatomical coordinates, and integrates it into the model registry, docs, and test suite.
Changes:
Implement PopT model + reusable encoder building blocks (input embedding, spatial positional encoding, heads).
Register the new model in braindecode’s model registry/categorization and update release notes.
Add PopT-specific unit tests (including an optional upstream parity gate) and update integration TorchScript skip list.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
File
Description
braindecode/models/popt.py
Adds PopulationTransformer model and coordinate discretization logic; documents intended pretrained usage.
braindecode/modules/popt_modules.py
Adds PopT-specific embedding/positional encoding + heads as standalone modules.
braindecode/models/__init__.py
Exposes PopulationTransformer in public model imports.
braindecode/models/util.py
Adds PopT to models_mandatory_parameters so shared model suites cover it.
braindecode/models/summary.csv
Adds PopT categorization/metadata entry for the model zoo.
docs/whats_new.rst
Announces PopT in the 1.7.0 development release notes.
test/unit_tests/models/test_popt.py
Adds PopT unit tests (return_features/CLS, coords-from-chs_info, optional upstream parity gate).
test/unit_tests/models/test_integration.py
Skips TorchScript test for PopT due to polymorphic return type.
Comments suppressed due to low confidence (1)
braindecode/models/popt.py:70
This "Pre-trained weights available" note claims the checkpoint loads directly from braindecode/popt-pretrained. If the re-hosting is a follow-up (per PR description), this block should instead point to the upstream checkpoint and avoid promising an immediate from_pretrained path.
.. important::
**Pre-trained weights available.** The official checkpoint is released by
the authors and loads directly::
model = PopulationTransformer.from_pretrained(
…leak
- point popt.py / popt_modules.py docstrings at the actual parity gate
(test_encoder_is_bit_exact_with_upstream) instead of a non-existent
scripts/popt_parity_check
- convert the upstream import helper into a fixture that restores sys.path
on teardown so POPT_SRC never leaks onto the import path of other tests
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
braindecode/models/popt.py:19
This docstring states the official checkpoint loads from braindecode/popt-pretrained, but the PR description says re-hosting on the braindecode Hub is a follow-up and not yet available. This is likely to mislead users (and would fail if the repo/tag is missing).
verified by the ``test_encoder_is_bit_exact_with_upstream`` parity gate in
``test/unit_tests/models/test_popt.py``); the classification head is a
braindecode-native addition. The official checkpoint loads directly via
``PopulationTransformer.from_pretrained("braindecode/popt-pretrained")``.
braindecode/models/popt.py:73
The .. important:: block claims PopulationTransformer.from_pretrained("braindecode/popt-pretrained") works now, but the PR description explicitly says weight re-hosting is a follow-up. Consider wording this as future/optional, or pointing to the upstream HF weights instead.
the authors and loads directly::
model = PopulationTransformer.from_pretrained(
"braindecode/popt-pretrained", n_outputs=2
)
braindecode/modules/popt_modules.py:70
The sinusoidal table construction breaks when pe_dim is odd (i.e., hidden_dim divisible by 4 but not by 8): pe[:, 1::2] has fewer columns than div_term, causing a shape mismatch at the cosine assignment.
If x.shape[1] differs from self.n_chans and coords/seq_id are left as None, the buffer-based defaults will have a different electrode count than x, leading to a later shape error that’s hard to diagnose. Raising an explicit ValueError here would make the failure mode clear.
batch_size, n_chans, _ = x.shape
if coords is None:
coords = self.electrode_coords.unsqueeze(0).expand(batch_size, -1, -1)
if seq_id is None:
seq_id = self.electrode_seq_id.unsqueeze(0).expand(batch_size, -1)
test/unit_tests/models/test_popt.py:130
This parity-gate test unconditionally imports omegaconf, but omegaconf is not listed in the project's core or tests optional dependencies (pyproject.toml). If a developer sets POPT_SRC locally, the test will error instead of skipping. Consider skipping the parity gate when omegaconf is unavailable.
Encoder parity. The test_encoder_is_bit_exact_with_upstream gate is skipped in CI (needs POPT_SRC, the upstream clone). Ran locally: after copying upstream weights, encoder output matches to max abs diff 9.5e-7 (atol=1e-5), spatial PE buffers to atol=1e-6. The ported input embedding + Transformer are bit-exact with upstream.
Re-hosted weights (braindecode/popt-pretrained). Re-ran the paper's Table 1 protocol (Brain Treebank, frozen BrainBERT features → PopT, 7 subjects, paper's fine-tuning recipe). Mean ROC-AUC:
task
paper
re-hosted
scratch
volume
0.84
0.82
0.55
pitch
0.69
0.75
0.53
onset
0.86
0.87
0.73
speech
0.89
0.82
0.60
All four tasks in the paper's range, pretrained≫scratch preserved everywhere (+0.14 to +0.27). Not bit-exact absolute numbers (single seed, subject subset) but a faithful functional reproduction.
The reason will be displayed to describe this comment to others. Learn more.
I reviewed this exact candidate against the PopulationTransformer paper and
its BrainBERT dependency. I am requesting changes because the mandatory
paper-result gate is not met and cannot yet be executed faithfully.
The qualifying target is the primary-paper Table 1 BrainBERT + pretrained PopT
Volume result at 90 electrodes: ROC-AUC 0.87 +/- 0.03, reported as mean +/-
SEM across subjects. Before any run, the exact electrode ranking, subject
splits, seeds/repeats, aggregation, PopT source/checkpoint lineage, and
compatible reuse rights must be frozen. The gate also depends on a faithful,
licensed BrainBERT checkpoint path, which PR #1104 does not yet provide.
The current implementation is a pooled-classification shell rather than the
complete paper pipeline: it does not integrate BrainBERT, has no padding mask
for variable electrode sets, and accepts a hidden-width configuration that
later crashes. Green synthetic shape tests therefore cannot establish the
published result.
Please keep this behind the BrainBERT gate, reproduce the selected Table 1 row
on one exact candidate SHA within a preregistered tolerance, and reduce the
library PR to the smallest faithful model/masking contract with focused tests.
Data acquisition, checkpoint conversion, and full experiments should remain
external and opt-in. No data or checkpoint bytes were downloaded because the
protocol and provenance gates are incomplete.
Paper-replication gate at exact head b717d90: hold behind #1104.
A bounded PopulationTransformer Volume run on the Iceberg NEMAR Brain Treebank adaptation completed with validation ROC-AUC 0.9167148370 and test ROC-AUC 0.9335834897 (BCE 0.5066329166). The score is not a paper reproduction: it is one subject/run, uses 91 released clean channels rather than the unpublished task-specific top 90, lacks the paper split/subject aggregation, uses a post-paper checkpoint, and has no proven NEMAR-to-original-HDF5 signal equivalence.
Merge requires BrainBERT first, strict recoverable checkpoint/source provenance and rights, the paper's electrode-selection/split protocol, and the reported aggregation. The current successful run is diagnostic evidence only.
The reason will be displayed to describe this comment to others. Learn more.
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.
❌ Patch coverage is 92.59259% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.82%. Comparing base (1edbf60) to head (6e87a8e). ⚠️ Report is 12 commits behind head on master.
- coord_units="raw" keeps the released Brain Treebank integer
(left, inferior, posterior) coordinates, as NEMAR nm000253 stores them,
instead of treating them as metres (which clamped every electrode to
4999). The default "m" path is unchanged and documented as not being
the pretrained coordinate space.
- forward accepts key_padding_mask (upstream src_key_padding_mask) and
builds the default sequence ids from the input, so padded batches with
other electrode counts work; drop the now-unused seq-id buffer.
- reset_head keeps the old head's device and dtype and updates n_outputs
through _set_n_outputs.
- hidden_dim must be divisible by 8 (each positional quarter must be even);
hidden_dim=12 now raises a clear ValueError instead of a shape error.
- Keyword-only model parameters after the signal parameters;
versionadded 1.8.2; "numerically equivalent" instead of "bit-exact";
file header.
- Tests for each item; the upstream parity gate still passes with POPT_SRC.
The parameter count does not match the constructor in this PR: with n_times=768, the default hidden_dim=128, ffn_dim=256, and two encoder layers, the model has 479,746 parameters, not 480,002. Please update the summary metadata so the published model size is accurate.
This public export is not accompanied by an entry in the explicit model autosummary in docs/api.rst (the list currently ends at ZUNA), so Sphinx will not generate API documentation or a public API link for PopulationTransformer. Add the class to that autosummary as part of registering this public model.
Correct the PopT parameter count
braindecode/models/summary.csv:51
The parameter count in this summary row does not match the model constructed by the new defaults. For n_times=768, hidden_dim=128, ffn_dim=256, n_layers=2, and n_outputs=2, the trainable modules in popt.py total 479,746 parameters, not 480,002, so the generated summary metadata is inaccurate.
Both conflicts are additive: master added BrainBERT to
models_mandatory_parameters and BaRISTA/Brant to the changelog while this
branch added PopulationTransformer to the same two places. Both sides kept.
Bruno asked for the parenthetical aside to go. It also referenced a test file
and a numerical tolerance, which contradicted the weight-for-weight claim made
in the same sentence.
The header claimed MIT because upstream is MIT, but this file is a
braindecode-native port that copies no upstream line -- the encoder is a stock
torch.nn.TransformerEncoder. BrainBERT, the same situation, carries BSD-3.
Bruno: 'We have a function for this...'. extract_channel_locations_from_chs_info
already covers the three fallback cases this reimplemented by hand -- missing,
malformed and all-zero chs_info -- and is the helper SignalJEPA, LUNA and ZUNA
use.
It stops at the first unusable channel and does not screen NaNs, so the two
guards PopT needs on top (one row per electrode, all finite) are kept and the
reason is written down.
Bruno: 'unlock this'. The skip was justified by the polymorphic return type,
and that justification was wrong -- scripting failed earlier than that, on two
real defects:
- key_padding_mask was reassigned inside its own 'is not None' branch, which
drops TorchScript's narrowing; and '.bool()' is not in TorchScript's Tensor
API, so '.to(torch.bool)' replaces it.
- the return type is then handled the way Brant already does it, with
torch.jit.is_scripting() pruning the dict branch at compile time.
test_model_torch_script[PopulationTransformer] now passes.
Bruno: 'delete this'. Brant and BrainBERT ship no per-model test file either;
the shared suites cover PopT through the registry.
PopulationTransformer is added to the explicit _MODELS list in
test_return_features.py (has_cls=True -- PopT returns a real CLS token, not
None). The reset_head cases were already picked up from the registry.
25 shared tests now run against PopT. The one remaining skip is the test
declaring PopT has no embedding parameter, which is true.
The upstream parity gate goes with the file. It stays reproducible evidence for
the PR rather than CI surface, as it needs a local clone of the upstream repo.
…d stop overclaiming parity
Two findings from the automated review that still applied.
The forward comment claimed that sizing seq_id from the input made 'padded
batches with more electrodes than at construction' work. It did not:
electrode_coords still held the configured count, so the call died inside
torch.cat with 'Sizes of tensors must match except in dimension 2'. The default
coordinates only describe the electrode set they were derived from, so the
mismatch is now refused with a message naming both counts and the way out --
passing coords explicitly, which does work and is now the documented path.
The check uses string concatenation rather than an f-string to stay scriptable.
The changelog said the encoder was 'ported bit-exact'. The parity evidence in
this PR is a max abs difference around 1e-6, not zero. Weights do map 1:1, so
'weight-for-weight' is the claim the evidence supports.
Three port deviations from czlwang/PopulationTransformer@d237755b, found by
the PR #1105 reproduction (pr-evidence-20260924/pr1105-reproduction-report.md):
- Coordinates: chs_info positions are no longer shifted so that each axis
starts at 0. Upstream feeds absolute integer (L, I, P) indices; the shift
moved the CLS output of the released weights by ~9 % (cosine 0.9975). The
shift stays available as the opt-in shift_coords=True (default False).
Indices outside [0, max_len - 1] are still clamped, now with a warning.
- Head: final_layer is upstream's fine-tuning head, a single Linear on the
CLS token (PtDownstreamModel.linear_out), instead of LayerNorm + Linear.
load_state_dict maps the earlier final_layer.fc.* keys onto it and drops
final_layer.norm.*, so braindecode/popt-pretrained@50b02d6 still loads
with strict=True.
- Defaults: hidden_dim=512, ffn_dim=2048, n_layers=6, n_heads=8, the
released popt_brainbert_stft config, so a default-constructed model and
from_pretrained("braindecode/popt-pretrained") agree. summary.csv updated
(480002 -> 19967746 parameters).
Regression tests go in the shared test_foundation_models.py suite.
Review follow-up. Restore the caveat that absolute MNE head-frame positions
are not the checkpoint's (left, inferior, posterior) space, make the clamp
warning's hint depend on shift_coords, and document that the legacy head's
LayerNorm is dropped unconditionally and only at the top level.
Preserve MIT licensing attribution for the upstream port
braindecode/models/popt.py:3
This is documented as a faithful port of the upstream PopulationTransformer code, whose repository is MIT-licensed, but this file declares BSD and neither this file nor popt_modules.py is listed in the MIT section of NOTICE.txt (unlisted package files are declared BSD there). Preserve the upstream MIT attribution in the notice and pass the corresponding non-BSD license metadata to EEGModuleMixin; otherwise the distributed source and Hub metadata misstate the port's licensing.
Shorter module/class comments, f-string error message, concise whats_new entry.
Header and class kwarg declare the upstream MIT license; NOTICE lists both files;
PopulationTransformer added to docs/api.rst.
Following @bruAristimunha's report: the paper's random split leaks (5 s windows around words that are a fraction of a second apart, so test windows overlap training ones). Same setup (7 subjects, 3 seeds), pretrained PopT, random vs time-blocked split:
Task
Random (paper)
Time-blocked
Pitch
0.79
0.51
Volume
0.89
0.64
Onset
0.86
0.84
Speech
0.90
0.84
The model and weights are not affected, only the evaluation. I added a warning in the docstring (6e87a8e).
The reason will be displayed to describe this comment to others. Learn more.
Copilot review overview
🔵 Needs a closer look
Several moderate issues remain, including missing parity coverage, invalid fallback indexing, an incorrect registry input dimension, and API documentation indentation.
The reason will be displayed to describe this comment to others. Learn more.
Thanks Adam. Replication on our side: 1,176 runs, Table 1 11/16 within 1 SE and Figure 3 reproduced; the docstring warning about the released split is the right call. Merging as agreed.
Conflict resolution:
- test/unit_tests/models/test_foundation_models.py: kept both the SleepFM/SleepFMStager imports (braindecode#1106) and master's PopulationTransformer import (braindecode#1105), alphabetized.
- docs/whats_new.rst auto-merged cleanly, keeping both the SleepFM and PopulationTransformer entries.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
bruAristimunha
left a comment
Towards #1097.
Adds a braindecode-native port of PopulationTransformer (PopT) (Chau et al. 2024, arXiv:2406.03044), the third iEEG/sEEG model from #1097 after
Brant(#1100) andBrainBERT(#1104).What it is
Unlike a per-channel encoder, PopT is a population aggregator: its input is a set of per-electrode feature vectors (typically the frozen embeddings of a channel-level foundation model such as
BrainBERT) plus each electrode's integer anatomical coordinates. ACLStoken summarises the population after a stack of Transformer encoder layers.Design
n_timesaxis, so the model keeps the standard(batch, n_chans, n_times)signature (n_chans= number of electrodes,n_times= feature dimension, 768 for BrainBERTstftfeatures).chs_infoand discretised to integer indices inside the model; without positions the electrodes fall back to distinct sequential indices.Parity
The input embedding + Transformer are checked bit-exact against the upstream
PtModelCustom(input embedding + spatial encoding + post-norm encoder), with the releasedpopt_brainbert_stftweights copied over —test_popt.py::test_encoder_is_bit_exact_with_upstream(gated byPOPT_SRC, skipped in CI without it). Reproduce with:Notes
PopulationTransformer.braindecode/popt-pretrained(revision50b02d6f8cb436ee81dea1dcf972ca635e9a07fd), soPopulationTransformer.from_pretrained("braindecode/popt-pretrained")works out of the box.cc @bruAristimunha
Replication (braindecode maintainers)
Released-checkpoint parity (2026-09-30, extends the synthetic-weight gate above): loaded the exact released
pretrained_popt_brainbert_stft.pth(HFPopulationTransformer/popt_brainbert_stftrevb11d552, SHA-256eae617d6...523d8ca) into upstreamPtModelCustom, and the port's official mirror (braindecode/popt-pretrained, SHA-256edba2cc8...b61ae5) into this port. The 78 tensor names shared between the two checkpoints are bit-identical (max abs diff0.0) — the port's mirror is confirmed to be the same releasedb11d552weights, renamed, not an independent retrain. CLS-token forward parity on identical synthetic input: max abs diff8.3e-7(float32 noise). No port-code divergence found; nothing pushed this session.Real-data Volume replication on Brain Treebank. The released BrainBERT STFT and PopT checkpoints are used on the paper's 7 held-out subjects. Settings follow the released code: 80/10/10 split and AdamW for 2000 steps with the paper's schedule. Frozen BrainBERT features are cached per subject and shared by both codes, so the original code and this port see the same features, splits, head initialisation and seeds (42, 43, 44). Runs used Gaudi HPU in FP32.
The port matches the original code. The largest per-subject gap is subject 6, which also has the largest spread across seeds. Both land above the paper's
0.87 ± 0.03. The setup differs from the paper in three ways:clean_laplacianelectrodes.