braindecode/brainbert-pretrained

BrainBERT — Foundation Model for Intracranial Neural Signals

Pretrained weights for braindecode.models.BrainBERT, a faithful braindecode port of BrainBERT (Wang et al., ICLR 2023), a self-supervised foundation model for intracranial (sEEG/iEEG) recordings.

• Paper: BrainBERT: Self-supervised representation learning for intracranial recordings (ICLR 2023)
• Original code & weights: czlwang/BrainBERT
• braindecode docs: https://braindecode.org

Provenance & license

These weights are the official pretrained "large" checkpoint (stft variant) released by the original authors. The Transformer encoder and input encoding are mapped 1:1 into the braindecode BrainBERT module.

This repository ships the encoder only — 82 tensors, bit-identical to the authors' release. No downstream head is included: BrainBERT is pretrained by masked spectrogram reconstruction and has no classification head to inherit, so final_layer is freshly initialized by from_pretrained for whatever n_outputs you ask for.

The upstream repository ships no explicit license file, so this repository is labelled unknown: the original authors retain all rights, and these weights are re-hosted for convenience to make from_pretrained work out of the box. If you use them, please cite the original BrainBERT paper and refer to the authors' repository for terms of use.

Model configuration

This checkpoint uses the released "large" configuration (~43M parameters):

hidden_dim	768
ffn_dim	3072
n_layers	6
n_heads	12
idx_freq_cutoff (input_dim)	40
nperseg	400
noverlap	350
stft_clip	10
stft_zscore_before_clip	True
activation	GELU
sfreq	2048 Hz

Upstream ships two contradictory STFT recipes. preprocessors/stft.py — reached from conf/preprocessor/stft_pretrained.yaml and therefore the one behind the released checkpoint — z-scores the spectrogram and then trims 10 frequency bins. notebooks/demo.ipynb trims 5 bins first and z-scores after. The order matters, because it decides whether the boundary artifacts enter the z-score statistics. This repository is configured for the production recipe (stft_clip=10, stft_zscore_before_clip=True); pass stft_clip=5, stft_zscore_before_clip=False to reproduce the notebook instead.

The signal is expected at 2048 Hz (Laplacian-re-referenced, as in the paper). The short-time Fourier transform front-end is computed inside the model, so it consumes raw (batch, n_chans, n_times) signal directly (upstream fed a pre-computed spectrogram). Frames are pooled over time and channels, so you may freely change n_chans and n_outputs (the classification head is task-specific and randomly initialized — fine-tune it on your downstream task). n_times only needs to be long enough to yield at least one spectrogram frame.

Usage

from braindecode.models import BrainBERT

# The encoder loads the pretrained weights; the classification head is
# (re)initialized for your task via n_outputs.
model = BrainBERT.from_pretrained("braindecode/brainbert-pretrained", n_outputs=2)

Fidelity

The ported encoder reproduces the upstream reference (MaskedTFModel, intermediate_rep=True) to within ~6e-5 max absolute difference on the authors' demo signal — the residual is cross-device float32 rounding in the fixed sinusoidal positional table (computed on GPU at training time, regenerated on CPU here), which grows only at positions far beyond any realistic sequence length.

Citation

@inproceedings{wang2023brainbert,
  title     = {{BrainBERT}: Self-supervised representation learning for intracranial recordings},
  author    = {Wang, Christopher and Subramaniam, Vishwaas and Yaari, Adam Uri and
               Kreiman, Gabriel and Katz, Boris and Cases, Ignacio and Barbu, Andrei},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2023}
}