ube-2afc: weights and data

Pretrained weights for ube-2afc, which tests whether a person's fMRI response can tell two near-identical images apart (pairmate discrimination), offline or live during a scan.

The encoder is the Universal Brain Encoder from the Irani lab: paper Beliy et al., code WeizmannVision/brainit-fmri. This repo holds weights for it plus the data needed to retrain the base.

You do not need to download anything by hand. From the code repo:

./download_checkpoints.sh              # encoder weights + the dinov2 backbone
WITH_WARMSTART=1 ./download_checkpoints.sh   # also the nsd voxel embeddings
WITH_NSD=1 ./download_checkpoints.sh         # also the nsd training data, 17 GB

files

file	size	what it is
ube_base_infonce.pt	9.6 MB	base encoder pretrained with a contrastive (InfoNCE + entropy) objective on top of the reconstruction loss. this is the default
ube_base_recon.pt	9.6 MB	same architecture and data, reconstruction loss only. the matched control
ube_base_original.pt	9.6 MB	the earlier base this work started from
nsd_voxel_embed_infonce.pt	324 MB	learned voxel embeddings for all 315,997 NSD voxels, for warm starting a new subject
dinov2_vitl14_reg4_pretrain.pth	1.2 GB	stock DINOv2 ViT-L/14 with registers, mirrored here so clusters without internet on compute nodes still work
nsd_data/	17 GB	preprocessed NSD, only needed to pretrain a base from scratch

The .pt files are plain tensors plus a config dict, so they load with weights_only=True and execute no pickled code:

import torch
ck = torch.load("ube_base_infonce.pt", weights_only=True)   # keys: config, shared

Turning them into a working encoder needs the code repo, which builds the model and loads these weights into it. The DINOv2 backbone is stock and is not inside the .pt files, which is why they are small.

the NSD data

nsd_data/ holds the Natural Scenes Dataset, preprocessed for training this encoder: betas_fithrf_GLMdenoise_RR betas in fsaverage surface space, restricted to the Algonauts mask, z-scored per voxel within each session, repeats averaged per image, plus the 73k stimulus images at 224x224.

• fmri_v2.npz (5.8 GB): single_sub_fmri, single_sub, multi_sub_fmri, type_sample, val_single_ind
• num_voxels_all_subjects.npy (256 B): voxels per hemisphere per subject
• all_images_v2_224.npy (11 GB): images indexed by NSD id

NSD is not ours. It comes from the Natural Scenes Dataset (naturalscenesdataset.org, Allen et al., Nature Neuroscience 2021), and the stimulus images come from COCO. If you use this data, follow NSD's terms and cite their paper. This copy is here to save people the preprocessing, not to replace agreeing to their terms.

credit

• Universal Brain Encoder: Beliy, Wasserman, Zalcher, Irani. arXiv:2406.12179
• DINOv2: Meta AI, Apache 2.0
• NSD: Allen et al., Nature Neuroscience 2021