neuroprobe 0.1.6

Neuroprobe: Benchmark for Evaluating iEEG Foundation Models. Evaluating intracranial brain responses to naturalistic stimuli.

Neuroprobe

Neuroprobe: Benchmark for Evaluating iEEG Foundation Models.

Evaluating Intracranial Brain Responses to Naturalistic Stimuli

🌐 Website | 📄 Paper | 🚀 Example Usage | 📤 Submit

---

By Andrii Zahorodnii¹²*, Christopher Wang¹*, Bennett Stankovits¹*, Charikleia Moraitaki¹, Geeling Chau³, Andrei Barbu¹, Boris Katz¹, Ila R Fiete¹²,

¹MIT CSAIL, CBMM | ²MIT McGovern Institute | ³Caltech | *Equal contribution

Overview

Neuroprobe is a benchmark for evaluating EEG/iEEG/sEEG/ECoG foundation models and understanding how the brain processes information across multiple tasks. It analyzes intracranial recordings during naturalistic stimuli using techniques from modern natural language processing. By probing neural responses across many tasks simultaneously, Neuroprobe aims to reveal the functional organization of the brain and relationships between different cognitive processes. The benchmark includes tools for decoding neural signals using both simple linear models and advanced neural networks, enabling researchers to better understand how the brain processes information across vision, language, and audio domains.

Please see the full technical paper for more details.

Getting Started

Prerequisites

1. Install the package:

pip install neuroprobe

2. If you haven't yet, download the BrainTreebank dataset from the official release webpage, or using the following script (located here):

python braintreebank_download_extract.py --lite

(lite is an optional flag; if only using Neuroprobe as a benchmark, this flag will reduce the number of downloaded files by >50% by removing unnecessary files.)

Code Example

Start experimenting with quickstart.ipynb to create datasets and evaluate models. For example:

import os, torch
os.environ['ROOT_DIR_BRAINTREEBANK'] = '/path/to/braintreebank/'  # NOTE: Change this to your own path, or define an environment variable elsewhere

from neuroprobe import BrainTreebankSubject, BrainTreebankSubjectTrialBenchmarkDataset
subject = BrainTreebankSubject(subject_id=1, cache=True, 
                               dtype=torch.float32, coordinates_type="cortical")
dataset = BrainTreebankSubjectTrialBenchmarkDataset(subject, trial_id=2, 
                                                    dtype=torch.float32, 
                                                    eval_name="gpt2_surprisal") 

data_electrode_labels = dataset.electrode_labels 
data_electrode_coordinates = dataset.electrode_coordinates 

dataset.output_dict = True # Optionally, you can request the output_dict=True to get the data as a dictionary with a bunch of metadata.
dataset.output_indices = False # Optionally, you can request to output indices into the original BrainTreebank h5 files of the sessions, instead of raw data.
print(dataset[0])

will give the following output:

{
	'data': torch.tensor, # shape: (n_electrodes, 2048), where 2048 = 1 second at 2048 Hz
	'label': int, # index of the class to be predicted: 0, 1, etc.
	'electrode_labels': list[str], # length: (n_electrodes, )
	'electrode_coordinates': torch.tensor, # shape: (n_electrodes, 3)
	'metadata': {'dataset_identifier': 'braintreebank', 'subject_id': 1, 'trial_id': 2, 'sampling_rate': 2048}
}

In case you'd like to use your own pipeline for extracting and preprocessing data, feel free to set dataset.output_indices = True, in which case the output will look like:

{
    'data': (index_from, index_to), # tuple of indices: indices into the session's h5 file in the BrainTreebank
    ... # the same as above
}

Leaderboard Requirements

To submit to the Neuroprobe leaderboard, you MUST use the exact train/val/test splits that are provided by the Neuroprobe package:

from neuroprobe import generate_splits_cross_session
# options: generate_splits_within_session, generate_splits_cross_session, generate_splits_cross_subject
splits = generate_splits_cross_session(test_subject=subject, test_trial_id=2, 
                                       eval_name="gpt2_surprisal", output_indices=False)
print(splits[0])

will give the following output:

{
    "train_dataset": BrainTreebankSubjectTrialBenchmarkDataset,
    "val_dataset": BrainTreebankSubjectTrialBenchmarkDataset,
    "test_dataset": BrainTreebankSubjectTrialBenchmarkDataset
}

Evaluation Example

Run the linear regression model evaluation using the following example script (located here):

python eval_population.py --subject_id SUBJECT_ID --trial_id TRIAL_ID --verbose --eval_name gpt2_surprisal --split_type CrossSession

Results will be saved in the eval_results directory according to leaderboard_schema.json.

Citation

If you use Neuroprobe in your work, please cite our paper:

@misc{neuroprobe,
      title={Neuroprobe: Evaluating Intracranial Brain Responses to Naturalistic Stimuli}, 
      author={Andrii Zahorodnii and Christopher Wang and Bennett Stankovits and Charikleia Moraitaki and Geeling Chau and Andrei Barbu and Boris Katz and Ila R Fiete},
      year={2025},
      eprint={2509.21671},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2509.21671}, 
}