RAntonello/vjepa2.1-vitl-braintuned

V-JEPA 2.1 ViT-L — brain-tuned (3 epochs)

V-JEPA 2.1 ViT-L whose upper attention blocks have been brain-tuned: fine-tuned to predict human fMRI responses to naturalistic movies, then merged back into a standalone checkpoint.

The motivating question is whether optimising a video encoder against brain data yields better video representations. On held-out data it does: the tuned tower predicts unseen brain responses substantially better than the frozen one.

	OOD s6 mean per-parcel r
frozen V-JEPA 2.1 ViT-L	0.1677
brain-tuned (this model)	0.1838
	+0.0161

This is the largest brain-tuning gain we have measured, exceeding the same procedure applied to Whisper-large-v3 (+0.0108) — notable because video is the weaker modality for whole-brain prediction in this dataset (a video-only encoding model scores ~0.25 vs ~0.28 for a text/semantic one).

Method — frozen-ridge brain-tuning

The readout is pinned at its closed-form ridge solution on the frozen features and never updated, so the LoRA adapter is the only trainable thing. This removes the head-masking escape route: any drift that hurts a linear decoder shows up directly in the loss, and the gain cannot come from the readout adapting to the encoder.

• Adapter: LoRA r=8, alpha=16 on attention qkv + proj of blocks 15–20 (the blocks feeding the layer-20 readout; blocks above it receive no gradient and are left untouched)
• Target: whole brain, 1000 Schaefer parcels, subjects 1/2/3/5
• Data: CNeuroMod Friends seasons 1–5 (~98k 4-second clips), season 6 held out
• Schedule: 3 epochs / 75k steps, batch 4, AdamW lr 2e-5
• Readout: RidgeCV fit on 40,000 clips (39 samples per input dimension — an underfit readout inflates the apparent gain, so this matters)
• Merged as W += (alpha/r) · B @ A, i.e. scale 2.0

One epoch gives only +0.0074; the gain roughly doubles by 3 epochs. Video benefits from far more training than Whisper, which plateaus within one epoch.

Files

• vjepa2_1_vitl_braintuned.pt — full 24-block checkpoint under the target_encoder key, a drop-in replacement for the original V-JEPA 2.1 ViT-L checkpoint
• braintune_provenance.json — base checkpoint, adapter hyperparameters, merged module list, scores

Usage

Loads exactly like the base checkpoint via the V-JEPA 2 repo:

from src.models import vision_transformer as vit
import torch

model = vit.vit_large(patch_size=16, img_size=(256, 256), num_frames=64, tubelet_size=2,
                      use_sdpa=True, use_SiLU=False, wide_SiLU=True, uniform_power=False,
                      use_rope=True)
ck = torch.load("vjepa2_1_vitl_braintuned.pt", map_location="cpu")
sd = {k.replace("module.", "").replace("backbone.", ""): v for k, v in ck["target_encoder"].items()}
model.load_state_dict(sd, strict=False)          # pos_embed keys absent by design (RoPE)

Features are read from block 20 (out_layers=[20]), spatially mean-pooled to 32 temporal tokens per 4-second clip, matching how the model was tuned.

Caveats

• Only blocks 15–20 differ from the base model; 0–14 and 21–23 are bit-identical.
• Tuned on 4 subjects watching one TV series. Whether the gain generalises to other video tasks is not established here.
• Evaluated on held-out episodes, not held-out subjects.

License & attribution

CC-BY-NC-4.0, inherited from Meta's V-JEPA 2.1 base checkpoint, which this derives from. Base model © Meta AI. Non-commercial use only.

fMRI data: CNeuroMod Friends dataset, via the Algonauts 2025 Challenge.