Pretrained OmniMouse: a multi-modal, multi-task transformer for the mouse visual cortex. Trained on 2.3M neurons from 73 mice across 323 sessions (150B+ neural tokens), unified across neural prediction, behavioral decoding, and neural forecasting.

Scaling data and models has transformed AI. Does the same hold for brain modeling? We train multi-modal, multi-task models on 2.3 million neurons from 73 mice (150B+ neural tokens), flexibly supporting neural prediction, behavioral decoding, and neural forecasting. OmniMouse achieves state-of-the-art performance, outperforming specialized baselines across virtually all regimes. Yet performance scales with more data while gains from larger models saturate — inverting the standard AI scaling story: brain models remain data-limited even with vast recordings.

This repo stores the training checkpoint as per-rank shards (rank_*.ckpt), loaded through omnimouse.utils.ModelCheckpoint. The loader handles multi-rank → single-rank consolidation transparently.

Pass ckpt_dir as ckpt_path to the standard eval pipeline:

Or wire it into a notebook via notebooks/inference.ipynb — see that notebook for the end-to-end flow (config composition, dataloader setup, checkpoint load, evaluation).

The matching Hydra experiment config is configs/experiment/omnimouse_20M.yaml. The session-specific readouts (sess_params) in the checkpoint are tied to the training corpus; when fine-tuning to new sessions, load with strict=False so unmatched readouts initialize fresh.

OmniMouse follows an encode–fuse–decode design. A cross-attention encoder with learned latent queries compresses unmasked neural and behavioral tokens into a fixed-length latent sequence, using local sliding-window attention plus a few global "register" tokens. A fusion stack then integrates these latents with video features through transformer layers that interleave sliding-window and global self-attention at a 5:1 ratio. All attention blocks use 1D RoPE derived from each token's timestamp, and a shared cross-attention decoder reads out neuronal and behavioral targets from the fused representation.

2.3M+ single neurons from the visual cortex of 73 mice across 323 sessions, totaling 150B+ neural tokens. Mice viewed naturalistic movies (cinematic clips, Sports-1M), ImageNet images, and parametric stimuli (Gabors, random dot kinematograms, pink noise, model-generated stimuli) while running on a wheel. Pupil position, pupil dilation and its derivative, and running speed were recorded throughout.

The dataset is available at the-enigma-project/omnimouse-dataset.

Single-trial Pearson correlation on the held-out test set (full 323-session corpus).