README.md

---
license: apache-2.0
library_name: brainmu-tokenizer
pipeline_tag: feature-extraction
tags:
  - neuroscience
  - neural-signals
  - eeg
  - calcium-imaging
  - neuropixels
  - tokenizer
  - vector-quantization
  - mouse
---

# 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](https://github.com/BAAI-Brain-Inspired-Group/Brainmu);
the token-consuming foundation model is not part of this release.

![architecture](architecture.png)

## Usage

The package is not on PyPI. Clone the
[Brainmu repository](https://github.com/BAAI-Brain-Inspired-Group/Brainmu)
and install it from the repository root:

```bash
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
```

```python
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](https://github.com/BAAI-Brain-Inspired-Group/Brainmu).

## 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.

```bibtex
@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:

```bibtex
@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}
}
```