epfl-neuroai/vjepa2-encoder-basic

V-JEPA2 Offline Encoder for Video-Evoked BOLD Responses

This repository contains a PyTorch checkpoint for a basic V-JEPA2-based offline encoder trained to predict video-evoked BOLD responses. The encoder is intended for research workflows involving neural response prediction and neural response-guided visual synthesis.

The checkpoint stores decoder weights and metadata for an offline encoder. This repository includes a custom transformers.AutoModel wrapper and does not require the original training codebase.

Input/Output Contract

The intended input is a short video clip corresponding to the training stimulus duration:

• Input: one 3-second RGB video clip, represented as a float tensor shaped [B, T, C, H, W] with values in [0, 1].
• Output: one vector of predicted z-scored fMRI beta responses per video, shaped [B, 20484].
• Temporal dimension: the output has no time dimension. Each 3-second video maps to a single predicted response vector.

This makes the encoder suitable for scoring or optimizing short generated videos against static target neural-response patterns.

The video-input path resizes frames to 224 x 224 and applies the ImageNet normalization used by the V-JEPA2 training pipeline. If you pass already-normalized V-JEPA2 inputs, call model.predict_fmri(video, normalize=False).

Loading

This checkpoint can be loaded with transformers.AutoModel and trust_remote_code=True.

Example:

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "epfl-neuroai/vjepa2-encoder-basic",
    trust_remote_code=True,
)
model.eval()

# Replace this with a preprocessed 3-second video tensor.
# Shape: [batch, frames, channels, height, width].
video = torch.zeros(1, 16, 3, 224, 224)

with torch.no_grad():
    prediction = model.predict_fmri(video)

print(prediction.shape)  # [1, 20484]

For decoder-only debugging, the model can also run from precomputed V-JEPA2 layer features:

model = AutoModel.from_pretrained(
    "epfl-neuroai/vjepa2-encoder-basic",
    trust_remote_code=True,
    load_vjepa=False,
)

features = [
    torch.zeros(1, decoder.mean.shape[1])
    for decoder in model.decoders
]

with torch.no_grad():
    prediction = model.forward_features(features)

Data

This checkpoint was trained using data from:

• BOLD Moments Dataset (BMD): whole-brain fMRI responses to short naturalistic videos.
• Social interaction video fMRI dataset from Emalie McMahon and collaborators: fMRI responses to naturalistic two-person social action videos.

This repository does not include the underlying fMRI datasets or stimulus videos.

Files

• vjepa2_offline_encoder.pth: PyTorch checkpoint containing decoder weights, decoding-unit selection metadata, feature-extractor configuration, and registered attributes.
• config.json, configuration_vjepa2_fmri_encoder.py, modeling_vjepa2_fmri_encoder.py: custom Transformers files for AutoModel loading.
• requirements.txt: minimal Python dependencies.

Backbone Source

The V-JEPA2 backbone weights are shipped in this repository as:

vitl.pt

The loader uses the V-JEPA2 Torch Hub architecture with pretrained=False, then loads the local vitl.pt weights directly. This avoids relying on the moving facebookresearch/vjepa2 Torch Hub checkpoint URL while preserving compatibility with the original decoder features. The decoder checkpoint uses canonical extractor_config["layer_names"] metadata.

Citations

If you use this checkpoint, please cite the V-JEPA/V-JEPA 2 backbone papers and source datasets:

@article{bardes2024revisiting,
  title={Revisiting Feature Prediction for Learning Visual Representations from Video},
  author={Bardes, Adrien and Garrido, Quentin and Ponce, Jean and Chen, Xinlei and Rabbat, Michael and LeCun, Yann and Assran, Mahmoud and Ballas, Nicolas},
  journal={arXiv preprint arXiv:2404.08471},
  year={2024}
}

@article{assran2025vjepa2,
  title={V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning},
  author={Assran, Mido and Bardes, Adrien and Fan, David and Garrido, Quentin and Howes, Russell and Komeili, Mojtaba and Muckley, Matthew and Rizvi, Ammar and Roberts, Claire and Sinha, Koustuv and others},
  journal={arXiv preprint arXiv:2506.09985},
  year={2025}
}

@article{tang2025diverse,
  title={Diverse perceptual representations across visual pathways emerge from a single objective},
  author={Tang, Yingtian and Gokce, Abdulkadir and Al-Karkari, Khaled Jedoui and Yamins, Daniel and Schrimpf, Martin},
  journal={bioRxiv},
  pages={2025--07},
  year={2025},
  publisher={Cold Spring Harbor Laboratory}
}

@article{lahner2024modeling,
  title={Modeling short visual events through the BOLD moments video fMRI dataset and metadata},
  author={Lahner, Benjamin and Dwivedi, Kshitij and Iamshchinina, Polina and Graumann, Monika and Lascelles, Alex and Roig, Gemma and Gifford, Alessandro Thomas and Pan, Bowen and Jin, SouYoung and Ratan Murty, N Apurva and others},
  journal={Nature communications},
  volume={15},
  number={1},
  pages={6241},
  year={2024},
  publisher={Nature Publishing Group UK London}
}

@article{mcmahon2023hierarchical,
  title={Hierarchical organization of social action features along the lateral visual pathway},
  author={McMahon, Emalie and Bonner, Michael F and Isik, Leyla},
  journal={Current Biology},
  volume={33},
  number={23},
  pages={5035--5047},
  year={2023},
  publisher={Elsevier}
}