BrainStratify

This is the official implementation of our work "A disentanglement-based self-supervised framework for versatile neural decoding".

Abstract

Invasive brain-computer interfaces (BCIs) can restore communication and movement, but existing decoders typically require extensive patient-specific calibration and are difficult to reuse across tasks and recording modalities. We introduce BrainStratify, a unified decoding framework that leverages self-supervised learning on abundant unlabeled recordings to reduce reliance on labeled data. BrainStratify performs hierarchical disentanglement by first discovering task-relevant functional channel modules using an attention-based spatial context objective and label-efficient cluster selection, and then factorizing within-module neural dynamics into discrete latent state tokens with a decouple-then-product quantization module. A codebook-guided masked modeling stage learns temporal context over these tokens, producing a pre-trained encoder that can be fine-tuned for downstream tasks, including classification, regression, and sequence decoding. Across stereo-electroencephalography (sEEG) and electrocorticography (ECoG), microelectrode array recordings, and surface electromyography (EMG), spanning speech production, perception, and motor control, BrainStratify consistently outperforms state-of-the-art baselines and maintains strong performance in low-label regimes. It improves performance by about 10% on most tested modalities and tasks. Our results suggest that neural disentanglement can provide practical, data-efficient decoders that make effective use of unlabeled neural activity, offering a promising direction for clinical BCIs.

Environment Setup

Install required packages:

conda create env -f environment-torch.yml
conda activate brain2vec-torch

Run Experiments

Stage 1: Download full repository

The full repository of BrainStratify is organized as follows:

- data                      # The data directory to organize datasets.
  |- seeg.he2023xuanwu      # Our collected sEEG dataset, including 12 subjects.
  |- ecog.wang2025xuanwu    # Our collected (epidural) ECoG dataset, including 1 subject.
  |- ...
- environment-prep.yml      # The conda environment config for preprocess.
- environment-torch.yml     # The conda environment config for training.
- figures                   # The resources to reproduce main figures and related tables.
  |- main2                  # The resources to reproduce main figure 2 and related tables.
     |- main2c              # The resources to reproduce main figure 2 (c) and related tables.
        |- expr-*           # The resources to reproduce main figure 2 (c) and related tables.
        |- main2c.svg       # The reproduced main figure 2 (c).
     |- ...
  |- ...
- models                    # The models used to reproduce results.
- params                    # The parameters used to reproduce results.
- plot                      # The plotting functions to reproduce main figures and related tables.
  |- main2                  # The plotting functions to reproduce main figure 2 and related tables.
     |- main2c              # The plotting functions to reproduce main figure 2 (c) and related tables.
  |- ...
- pretrains                 # The pre-trained checkpoints to reproduce results.
- train                     # The train scripts to reproduce results.
- utils                     # The tool functions to reproduce results.

Then, unzip directories according to the following commands:

unzip figures.zip
cd data/seeg.he2023xuanwu
unzip 001.zip && unzip 002.zip && unzip 003.zip ...
cd data/ecog.wang2025xuanwu
unzip 001.zip

Stage 2: Reproduce main results

Reproduce main results according to README.md in plot.