CoolFace
Modelpublic

braindecode/popt-pretrained

sourceHugging Faceunknownupdated 2mo agoView on Hugging Face
1likes11downloads
Model Card

PopulationTransformer (PopT) — braindecode weights

braindecode-native re-host of the official PopulationTransformer (PopT) popt_brainbert_stft checkpoint, remapped so it loads directly with `braindecode.models.PopulationTransformer`.

PopT is a self-supervised population model for intracranial recordings (iEEG/sEEG). It does not encode a raw time signal: each electrode is represented by a feature vector — typically the frozen embedding of a per-channel foundation model such as BrainBERT (stft features, 768-d) — and PopT aggregates across the electrode population. Each electrode feature is linearly projected and given a fixed sinusoidal spatial position encoding built from its integer anatomical coordinates; a CLS token is prepended, a stack of Transformer encoder layers mixes the population, and the CLS output is the pooled representation used downstream.

Provenance

Original codehttps://github.com/czlwang/PopulationTransformer
Original weightshttps://huggingface.co/PopulationTransformer/poptbrainbertstft
PaperChau et al. (2024), Population Transformer: Learning Population-level Representations of Neural Activity, arXiv:2406.03044

These weights are a format conversion only of the authors' released checkpoint — no re-training. The input embedding, the spatial positional encoding and the 6-layer Transformer encoder are carried over bit-for-bit (verified: all 82 mapped tensors are identical to the source, and the ported encoder reproduces the upstream PtModelCustom output to < 1e-5). The upstream masked-modelling heads (cls_head, token_cls_head) are not carried; the braindecode classification head (final_layer) is randomly initialised and must be trained/fine-tuned for your task.

Configuration

paramvalue
hidden_dim512
ffn_dim2048
n_layers6
n_heads8
n_times (feature dim)768
max_len (coord table)5000
activationGELU
parameters (encoder + spec head)~20.6M

Usage

python
import torch
from braindecode.models import PopulationTransformer

# n_chans = number of electrodes; n_outputs = your task's classes.
model = PopulationTransformer.from_pretrained(
    "braindecode/popt-pretrained", n_outputs=2
)

# input = per-electrode features (e.g. frozen BrainBERT stft embeddings),
# shape (batch, n_electrodes, 768). Electrode coordinates are read from
# chs_info when available, otherwise fall back to sequential indices.
x = torch.randn(4, 64, 768)
logits = model(x)                       # (4, n_outputs)
cls = model(x, return_features=True)    # {"features": ..., "cls_token": ...}

The classification head is task-specific: pass your own n_outputs (the head is re-initialised) and fine-tune. Electrode positions can be provided through chs_info (their loc) so the spatial encoding reflects the real montage.

Licensing

The upstream PopulationTransformer repository ships no explicit license file, so the license of these weights is marked `unknown`. They are re-hosted here for research use with attribution; if you use them, cite the original work and respect any terms the authors may later publish. The braindecode code is BSD-3-Clause.

Citation

bibtex
@article{chau2024population,
  title={Population Transformer: Learning Population-level Representations of Neural Activity},
  author={Chau, Geeling and Wang, Christopher and Talukder, Sabera and Subramaniam, Vighnesh and Soedarmadji, Saraswati and Yue, Yisong and Katz, Boris and Barbu, Andrei},
  journal={arXiv preprint arXiv:2406.03044},
  year={2024}
}