Brain-Response Models on Apple Silicon

Problem

TRIBE v2 is a brain-encoding model from Meta. You give it a video, audio or text, and it predicts the fMRI response an average person's brain would show, across about 20,000 points on the cortical surface. Researchers use models like this for "in-silico" experiments: try a stimulus on the model before putting a person in a scanner.

That only works if trying a stimulus is quick. On a Mac, TRIBE's vision path runs Meta's V-JEPA2 ViT-g video encoder in PyTorch on the CPU, because the feature-extraction library it uses does not put this extractor on the Apple GPU. It took about one and a half to two minutes for each second of video. A ten-second clip meant a coffee break. I wanted it fast enough to try things interactively on a laptop.

How it works

TRIBE does not look at pixels itself. It reads features from three large pretrained models (video, audio and text) and learns how to map them onto the brain. For video, the slow part is the encoder, not TRIBE. So I left TRIBE alone and replaced only the encoder.

1. Cut the video into seconds. TRIBE predicts at 1 Hz. For each second it takes the one-second window that ends there.
2. Sample 16 frames from that window. They are resized to 292 pixels on the short side, centre-cropped to 256 × 256 and normalised the way the original processor does.
3. Run V-JEPA2 ViT-g. A 40-block vision transformer, 1,408 numbers wide, cuts the clip into 16-pixel patches two frames deep. Each block's output is averaged over all patches.
4. Keep the layers TRIBE was trained on. Counting the patch embedding there are 41 hidden states. TRIBE keeps 20 evenly spaced ones, then averages two groups: from the halfway point to three quarters, and from three quarters to the top. That gives two vectors of 1,408 numbers per second.
5. Predict the brain response. The unchanged TRIBE v2 checkpoint maps those vectors to a value at each of about 20,000 vertices of the fsaverage5 cortical mesh. Predictions are shifted 5 s back in time to allow for the slow blood-flow response that fMRI measures.

Converting the encoder

MLX is Apple's array framework for its own chips. It runs on the GPU and shares memory with the CPU, so a model this size (about 1.9 GB at half precision) fits without copying. I did not write the network from scratch: the open vjepa2-mlx project already had a V-JEPA2 encoder in MLX for the smaller models. I configured it for ViT-g and moved the weights across.

1. Match the architecture. 1,408 hidden size, 40 layers, 22 attention heads, patch 16, tubelet 2, crop 256. The feed-forward width is 48/11 × the hidden size, about 4.36, not the usual 4. The MLX code assumed a whole number, so I patched it to accept a fraction.
2. Expose every layer. I added an option to return all hidden states, because TRIBE needs the middle of the network, not just the end.
3. Map the weights. Load facebook/vjepa2-vitg-fpc64-256 from Hugging Face in fp32, flatten it, and convert each tensor: linear and layer-norm weights stay as they are, the 3D patch convolution is transposed into MLX's layout, and fused query-key-value weights are split where MLX expects them separately.
4. Drop what TRIBE does not use. V-JEPA2 also has a predictor network for its own training objective. TRIBE only needs the encoder, so only encoder.* weights are kept.
5. Cast to fp16 and save. The weights go out as one safetensors file with a config, and are on Hugging Face. I kept an fp32 copy locally for parity checks.

Plugging it in without forking TRIBE

I wanted the official model and checkpoint to stay exactly as Meta published them, so that any difference in the output comes from the encoder alone. TRIBE gets its video features from one method, HuggingFaceVideo._get_data, inside the neuralset library. install_mlx_video_hooks() swaps that one method for an MLX version, and TribeModel.predict carries on as normal.

model = TribeModel.from_pretrained("facebook/tribev2", device="mps", ...)
install_mlx_video_hooks(model, weights_dir="mlx_weights/V-JEPA2-vitg-fpc64-256")
preds, segments = model.predict(events=events)   # vision features now come from MLX

Three details made it work:

• The method is wrapped by the exca caching library and cannot simply be reassigned. The hook patches the wrapper's underlying function and keeps it as a module-level function called _get_data, so cache IDs and pickling still work.
• MLX is not safe across worker processes, so the hook runs feature extraction in-process with no worker pool.
• MLX features get their own cache version (for example mlx-f16), so they never mix with cached PyTorch features or a different frame count.

Then some speed work inside the hook: the token average happens on the GPU and only the 20 needed layer vectors come back to the CPU, frames for the next second load while the current one encodes, and preprocessing uses NumPy and PIL instead of the PyTorch processor. A --fast option uses 8 frames instead of 16 for about twice the encoding speed, at some cost to quality.

What I measured

The check is simple: the same one-second clip and the same TRIBE settings, vision only, run once through the official PyTorch path and once through the MLX hook. Then compare the predicted brain maps. The script is demo_data/validate_mlx_vs_torch.py, and it passes if cosine similarity and Pearson correlation are both at least 0.95.

• 11×faster for a one-second clip, about 97 s down to 8.5 s
• 0.962cosine similarity to the official brain predictions
• 0.964Pearson r against the official predictions
• 12.6 sfor a ten-second video end to end, model load included

Per second of video, the efficient script works out to roughly 1.2 s plus about 2 s to load the model and install the hook. Feature caching on disk means a second run over the same video skips the encoder entirely.

Faster speech recognition

When a video has speech, TRIBE also needs to know which word was said when. The official path transcribes with WhisperX, which is slow on a Mac. I added Parakeet, running in MLX through parakeet-mlx, as an option: asr_backend="parakeet". Parakeet returns sub-word pieces, so the backend merges them back into words in the same format WhisperX produces, including contractions like "world's" and splits like "iPhone 17".

On a ten-second clip it took about 1 s against about 17 s for WhisperX. The word error rate was about 0.09, measured against WhisperX's transcript rather than a human one, so it says the two mostly agree, not that either is right.

Limits

• Parity was validated on one short clip. The outputs are close but not bit-identical, because of fp16 and frame sampling. A larger set of videos would show whether 0.96 holds.
• The timings are single cold runs on one Apple Silicon Mac, with model loading included. I did not record repeat runs or test other chips, so treat 11× as one measurement, not a benchmark.
• Only the vision path is accelerated. Text and audio features still run through the official PyTorch stack, so a full multimodal prediction is still slow.
• The 8-frame --fast mode trades some quality for speed. I did not measure its parity.
• The Parakeet word error rate is against WhisperX on one clip, not against a human transcript.
• This is an inference port, not a new brain model, and not peer reviewed. Predictions are for an "average" subject, as in the original.

Credits and licence

TRIBE v2 and V-JEPA2 are by Meta, and TRIBE v2 is released under CC BY-NC 4.0, so this fork and the converted weights are for non-commercial use only. The MLX encoder builds on the open vjepa2-mlx project. Speech recognition uses parakeet-mlx with mlx-community/parakeet-tdt-0.6b-v3. I did the conversion, the hooks, the speed work and the validation. Not affiliated with Meta. The stimulus and cortex figure at the top is from the TRIBE v2 paper (CC BY-NC 4.0). If you use the brain model, cite the TRIBE paper.

Read more

• The fork on GitHub: hooks, conversion script, validation and efficient run scripts.
• MLX conversion write-up: every step, the gotchas and the parity table.
• MLX weights on Hugging Face: fp16 encoder, about 1.9 GB.
• TRIBE v2 by Meta and the TRIBE v2 paper.