OmniMouse: Scaling properties of multi-modal, multi-task Brain Models on 150B Neural Tokens

1. Problem Statement

The paper addresses the question of whether the scaling laws observed in Large Language Models (LLMs) and Computer Vision (CV)—where increasing model size and dataset size yield predictable performance gains—apply to modeling brain activity.

• Current Limitations: Existing neural models are often restricted to single modalities (e.g., only neural data or only stimuli), single tasks (e.g., only forecasting), or specific datasets. They lack the flexibility to handle arbitrary combinations of inputs (neural history, visual stimuli, behavior) and outputs (predicting future neural activity, decoding behavior, or encoding stimuli).
• Data Scarcity vs. Complexity: While neuroscience has large datasets (e.g., MICrONS), they are fragmented and smaller than internet-scale corpora. It is unclear if current neural models are limited by compute/parameters or by the availability of diverse, high-quality data.

2. Methodology: OmniMouse Architecture

The authors introduce OmniMouse, a unified, multi-modal, multi-task transformer architecture designed to model the mouse visual cortex.

A. Dataset

• Scale: 3.1 million neurons recorded from 73 awake, head-fixed mice across 323 sessions.
• Total Tokens: >150 billion neural tokens.
• Modalities:
  • Neural: Wide-field two-photon calcium imaging (6–14 Hz, upsampled to 30 Hz) from the visual cortex.
  • Visual Stimuli: Naturalistic movies, images (ImageNet), and parametric stimuli (Gabor patches, noise, etc.).
  • Behavior: Running speed, pupil position, pupil dilation, and its derivative.

B. Tokenization & Masking

OmniMouse uses a flexible masking strategy similar to masked language modeling but adapted for multi-modal time-series data:

• Neural Tokenization: Uses strided 1D convolutions to convert calcium traces into tokens. Each token is augmented with learned identity embeddings for the specific neuron, session, and animal.
• Visual Tokenization: Uses the first 10 layers of a Hiera vision transformer to extract spatiotemporal embeddings from video frames.
• Behavior Tokenization: Projects behavioral traces via linear layers with channel-specific embeddings.
• Flexible Masking: The model supports arbitrary combinations of unmasked inputs and masked targets. It can be trained to:
  • Forecast future neural activity (causal context).
  • Decode behavior from neural activity.
  • Encode stimuli to predict neural responses.
  • Predict sub-populations based on other sub-populations.

C. Model Architecture

1. Encoder: Compresses unmasked neural and behavioral tokens into a fixed-length latent sequence using cross-attention with learned queries.
2. Multi-modal Fusion: Concatenates latent sequences with visual features (from Hiera) and passes them through a stack of transformer layers.
   • Attention Mechanism: Uses a hybrid of local sliding-window attention (for efficiency) and global register tokens (to maintain global context).
   • Positional Encoding: Uses 1D-RoPE (Rotary Position Embedding) based on timestamps to handle relative timing across modalities.
3. Decoder: A shared cross-attention decoder reconstructs target neural responses and behavioral traces from the fused representation. It uses modality-specific linear readouts to project back to the target dimensions.

D. Training Strategy

• Loss Functions: Poisson loss for neural activity (encoding) and Mean Squared Error (MSE) for behavior (decoding), with behavior loss down-weighted by 0.1.
• Scaling Experiments: Trained on nested datasets (8 to 323 sessions) and scaled model sizes from ~1M to ~300M parameters.
• Optimization: Warmup-stable-decay training schedule to generate dense checkpoints for scaling analysis.

3. Key Contributions

1. Unified Multi-Modal Framework: OmniMouse is the first model to jointly integrate single-neuron tokenization, video encoding, and behavioral decoding into a single architecture capable of handling arbitrary task combinations at test time.
2. Systematic Scaling Analysis: The paper provides a rigorous analysis of scaling laws in neuroscience, contrasting them with the "Chinchilla" scaling laws of LLMs.
3. State-of-the-Art Performance: The model outperforms specialized baselines (e.g., MtM, POYO+, CEBRA, Latent Models) across nearly all evaluation regimes, even when trained on the same amount of data.
4. Reversal of Scaling Paradigm: The authors demonstrate that in brain modeling, data is the bottleneck, not model size.

4. Key Results

A. Scaling Laws: Data vs. Model Size

• Model Size Saturation: Unlike LLMs where performance scales with parameter count, OmniMouse performance saturates at moderate model sizes (~80M parameters). Increasing model size beyond this point yields diminishing returns or overfitting.
• Data Scaling: Performance improves systematically and reliably with more data across all model sizes. Larger models benefit significantly more from additional data than smaller ones.
• Conclusion: Current neural models are data-limited, not compute-limited. The mouse visual cortex, despite being a "simple" system, requires vast amounts of diverse data to unlock further gains.

B. Task Performance

• Neural Prediction: OmniMouse achieves state-of-the-art results in forecasting, population prediction, and stimulus encoding.
• Behavior Decoding: This task showed the most promising scaling dynamics, with performance continuing to improve with model size, suggesting that behavioral prediction may benefit from even larger capacities.
• Generalization: The model generalizes to unseen masking configurations (e.g., varying context sizes) not seen during training, indicating it learns robust representations of neural variability.

C. Benchmarking

• Sensorium Competitions: OmniMouse-80M set new state-of-the-art records on the Sensorium 2022 and 2023 competitions, outperforming previous winners even when using a frozen backbone with only neuron-specific parameters trained.

5. Significance and Implications

• Foundation Models for Neuroscience: The work suggests that foundation models for neuroscience are feasible but require a shift in strategy: prioritize data diversity and volume over massive parameter counts.
• Phase Transitions: The observation of systematic scaling raises the possibility of "phase transitions" in neural modeling, where richer datasets might unlock qualitatively new capabilities (e.g., emergent reasoning about brain states), similar to LLMs.
• Unified Understanding: By unifying stimulus, behavior, and neural activity, OmniMouse enables the study of how brain state shapes sensory processing and how different sources of context (temporal history, population dynamics, visual input) contribute to neural variability.
• Limitations: The model scales linearly with the number of neurons (due to per-neuron embeddings), which may become computationally expensive for whole-brain recordings. Additionally, the current data is limited to spontaneous activity and simple behaviors, leaving the transfer to complex behaviors an open question.

In summary, OmniMouse demonstrates that while the principles of scaling apply to brain modeling, the optimal path differs from AI: more data is the primary driver of progress, and current models have not yet reached the data limits required to fully exploit their capacity.