baai-brain/brainmu-tokenizer

Brainmu Tokenizer (Brainmu0)

A pretrained tokenizer that maps three kinds of mouse neural recordings — skull EEG/EMG field potentials, two-photon calcium ΔF/F and deconvolved events, and binned Neuropixels spike counts — onto one shared discrete token interface. Modality-specific convolutional stems feed a shared 24-block SigLIP2-ViT-L transformer trunk; a product quantizer splits each 64-dimensional token feature into eight 8-dimensional blocks and assigns each block by cosine similarity to one of 4096 codes in its own codebook; an independent decoder trunk reconstructs the selected modality from the discrete IDs. This is the tokenizer stage of Brainmu1.0; the token-consuming foundation model is not part of this release.

Usage

The package is not on PyPI. Clone the Brainmu repository and install it from the repository root:

git clone https://github.com/BAAI-Brain-Inspired-Group/Brainmu.git
cd Brainmu
pip install -e .                      # model and core API
pip install -e '.[data,notebook]'     # + public-data demo notebooks

import torch
from brainmu_tokenizer import UniTok

device = "cuda" if torch.cuda.is_available() else "cpu"
model = UniTok.from_pretrained("baai-brain/brainmu-tokenizer").eval().to(device)

# API smoke test: one canonical 4 s mouse EEG/EMG window at 250 Hz.
# Random values verify shapes only; they are not a meaningful neural example.
x = torch.randn(1, 2, 1000, device=device)
encoding = model.encode(x, modality="eeg")

encoding.indices      # (1, 8, 250)  discrete code IDs, one per codebook
encoding.features     # (1, 250, 64) quantized token features
encoding.token_mask   # (1, 250)     valid (non-padding) positions
encoding.grid_shape   # (2, 125)     channel-group × time token grid

reconstruction = model.decode(
    encoding.indices, modality="eeg",
    token_mask=encoding.token_mask, grid_shape=encoding.grid_shape,
)

The first from_pretrained call downloads the BF16 checkpoint (~1.48 GB) into the standard Hugging Face cache. The repository is gated: requesting access on this page shares your Hugging Face username and email with us, and approval is automatic. The download then needs an authenticated session — run hf auth login (huggingface-cli login on older huggingface_hub) or set HF_TOKEN. Set BRAINMU_TOKENIZER_CKPT=/path/to/brainmu-tokenizer to load a local release directory instead.

For source code, demo notebooks, preprocessing utilities, and the data guide, see the Brainmu repository.

Intended use

Extracting discrete tokens or frozen features from mouse neural recordings, reconstructing signals from tokens, and training small supervised heads on pooled frozen features. The tokenizer itself is not meant to be fine-tuned in these workflows.

Training domain — mouse only. The checkpoint was trained exclusively on mouse recordings. Other species and other recording technologies are out-of-domain transfer and require independent validation; a successful API call is not evidence of scientific validity after a domain shift.

modality	Expected signal	Training-time sampled ranges	Release chunk canvas	Decoder output
eeg	mouse EEG/EMG field potentials	128–250 Hz; 1–2 channels; 800 or 1000 samples	2 × 1000	signal values
dff	calcium ΔF/F traces	1–30 Hz; 8–256 cells; 10 or 15 frames	128 × 15	signal values
event	deconvolved calcium events	1–30 Hz; 8–256 cells; 10 or 15 frames	128 × 15	Poisson log-rates
spikes	per-unit binned spike counts	5–60 Hz; 16–1024 units; 40 or 60 bins	128 × 120	Poisson log-rates

Normalize each channel over the full recording before cutting windows (brainmu_tokenizer.data.normalize_traces implements the training-matched rule), and convert event / spikes reconstructions out of log-rate space with reconstruction.clamp(max=6).exp() before comparing them to counts. Inputs that do not divide the modality stride must be zero-padded with aligned channel_mask / time_mask; oversized inputs are chunked automatically.

Weight precision and reproducibility

The uploaded model.safetensors stores all tensors in BF16, cast from the FP32 training weights. from_pretrained loads them into an FP32 model, which is the recommended way to run it.

Code assignment is a cosine argmax whose top-1/top-2 margins reach down to ~1e-4, so the discrete IDs are sensitive to numerical precision:

• fixed model dtype and device → repeated runs agree exactly. encode ignores any ambient torch.autocast and keeps the quantization path in FP32, the same invariant training held.
• changing the precision — model.to(torch.bfloat16), a different GPU, other kernels — shifts a few percent of code IDs.
• token features and reconstructions are far more stable: pooled features move far less than the discrete IDs, and reconstructions are unaffected in practice.

Practical consequence: extract every feature set you intend to compare in one pass, and do not mix features cached from different dtypes or devices.

Results

Numbers below come from this checkpoint on the documented public mouse-data subsets used by the demo notebooks. They are reference points for checking your own run of those notebooks, not a standardized benchmark.

Reconstruction (masked metrics over the demo windows):

Modality	Signal	Dataset	Pearson r	SNR (dB)
eeg	mouse EEG/EMG @ 250 Hz	AccuSleep (Mouse01, day 1)	0.995	19.5
dff	calcium ΔF/F traces	Allen Visual Coding ophys	0.970	10.9
event	calcium event proxy derived from ΔF/F	Allen Visual Coding ophys	0.986	13.5
spikes	Neuropixels binned spike counts	Allen Visual Coding Neuropixels (DANDI:000021)	0.970	8.8

Frozen-feature heads — a small MLP (LayerNorm → Linear → GELU → Dropout → Linear, 17k–134k parameters) trained on pooled frozen features; the tokenizer is never fine-tuned:

Task	Dataset	Modality	Split rule	Example score (chance)
Sleep staging, 3-class	AccuSleep (2 mice × 2 days)	eeg	subject-isolated mouse	92.9% balanced accuracy (33%)
Grating direction, 8-way / axis, 4-way	Allen Visual Coding ophys (experiment 510517131)	dff	held-out presentations	71.5% / 92.4% balanced accuracy (12.5% / 25%)
Running speed, regression / 2-state	Allen Visual Coding Neuropixels (DANDI:000021)	spikes	contiguous time blocks	Pearson r 0.974 / 96.7% balanced accuracy (50%)

Token interface

indices are parallel codebook assignments, not a flat language-model vocabulary: eight IDs, each indexing its own 4096-entry codebook, at every token-grid position. A downstream model must embed and merge the eight assignments explicitly, or consume features instead. The modalities share the codebooks, which gives them a uniform interface; that alone should not be read as proof that one ID means the same thing across modalities.

License and citation

Apache-2.0.

@software{brainmu,
  title   = {Brainmu Project},
  author  = {{The Brainmu Authors}},
  year    = {2026},
  version = {1.0.0},
  url     = {https://github.com/BAAI-Brain-Inspired-Group/Brainmu}
}

Associated study, where the mouse sleep work is relevant:

@article{yu2026memory,
  title   = {Memory reactivation underlies experience-dependent adaptive
             regulation of sleep},
  author  = {Yu, Menghan and Wang, Junjie and Zhai, Zihan and Huang, Ruoyi and
             Fan, Guofan and Su, Xiaoya and Niu, Yijun and Zhu, Haochen and
             Chen, Jiayi and Jiang, Grace and Zhang, Tian and Zhong, Yi and
             Lei, Bo},
  journal = {Science},
  volume  = {392},
  number  = {6802},
  pages   = {eaed8630},
  year    = {2026},
  doi     = {10.1126/science.aed8630}
}