the-enigma-project/ntt-tokenizers

NTT tokenizers

Pretrained tokenizers for neural spike data from the NTT (NeuroTicTokenization) benchmark: 46 models, one folder each (config.json + model.safetensors), trained on public brainsets only. Load them with ntt.

Paper link: TBD

1. Use a tokenizer

import torch
from ntt.hub import load_pretrained

tok = load_pretrained("conv-vae-100to1", revision="v1.0")  # downloads once, then cached
spec = tok.spec
print(spec.n_neurons, spec.window_length_seconds, spec.bin_size_ms, spec.compress_n)

load_pretrained also accepts a local export folder (one that contains config.json) in place of a name.

A batch holds binned spike counts in spike_matrix, shaped [batch, n_bins, n_units] with n_bins = window_length_seconds * 1000 / bin_size_ms. tok.spec tells you the rest:

• bin_size_ms is the bin width in milliseconds.
• window_length_seconds is a number for most models. For some (the Transformer-VQ models) it is a distribution such as UniformDist(low=0.1, high=2.0), the range of window lengths used in training; choose a length in that range whose n_bins is a multiple of spec.patch_length.
• n_units: with n_neurons=False (the LDNS-Perceiver) the model takes any number of units. With an integer n_neurons and compress_n=True the model reads exactly n_neurons units. With compress_n=False the evaluation passes all units of a recording to the model directly.
• If input_type is binary_spikes and binarize_spikes is true, ntt's data pipeline clips counts to 0/1 before the model sees them; do the same.

A batch can be built and encoded with

from ntt.core.model.tokenizer.validation import make_dummy_input

# Pass window_length_seconds=... when the spec's window is a distribution, and n_neurons=...
# for the number of units you feed (needed when spec.n_neurons is False).
batch = make_dummy_input(spec, batch_size=2)
batch = batch._replace(spike_matrix=my_counts)  # int32, same shape as batch.spike_matrix
with torch.no_grad():
    latents = tok.encode(batch)

LDNS-Perceiver: make_dummy_input sets a placeholder session_id, for which the perceiver uses all-zero unit embeddings: the outputs are meaningless and no error is raised. Set session_id to a session in the uploaded embeddings table and unit_ids to the units' indices within that session's units table (as ntt's data pipeline numbers them):

import json
from safetensors import safe_open

with safe_open(tok.spec.embeddings_path, framework="pt") as f:
    sessions = json.loads(f.metadata()["sessions"])  # the sessions the table covers

batch = make_dummy_input(tok.spec, batch_size=2, n_neurons=n_units)
batch = batch._replace(
    spike_matrix=my_counts,  # [2, n_bins, n_units]
    session_id=(my_session,) * 2,  # one of `sessions`
    unit_ids=torch.as_tensor(my_unit_indices, dtype=torch.int32).repeat(2, 1),  # [2, n_units]
)

For recordings with more units than a fixed-n_neurons model reads, the evaluation (job.evaluate_all_neurons=true, the default) first caps each session at max_neurons units with a seeded random subset (select_eval_units in ntt.core.data.components.transforms; the setting is max_neurons in ntt.evaluation.config.schema.EvaluationConfig), then, for compress_n=True models, splits the units into groups of n_neurons and runs the model on each group (NeuronCoveringWrapper in ntt.evaluation.utils.neuron_covering).

2. Re-run our evaluation

An hf://<owner>/<repo>/<name>[@<revision>] path works as job.ckpt_path for job=eval, job=eval_all (all its passes) and both robustness jobs, and as job.ckpt for the downstream jobs (job=downstream_encoding, and job=downstream_decoding in its single-run and calibrate modes). The model is downloaded into the Hugging Face cache and converted to a local checkpoint under $NTT_HUB_CACHE (default ~/.cache/ntt/hub).

Requirements:

• The published scores are on the test split.
• The public brainsets prepared with brainsets; see public_dataset_creation.md. Set paths.data_dir to the folder holding your prepared brainsets.

Reconstruction axis, every model:

uv run ntt job=eval_all \
    paths.run_name=conv-vae-100to1 \
    paths.data_dir=/absolute/path/to/brainsets/processed \
    job.ckpt_path=hf://the-enigma-project/ntt-tokenizers/conv-vae-100to1@v1.0 \
    job.split=test \
    job.main.enabled=true job.snr.enabled=false \
    job.input_robustness.enabled=false job.latent_robustness.enabled=false \
    job.compare_poisson=true job.poisson_rate_bin_ms=10.0 job.poisson_n_realisations=5 \
    job.output_path=results/conv-vae-100to1

paths.run_name is required by every ntt job; it names the run's log directory. The main pass uses eval_all's default data config, ntt_eval_light. For ldns-perceiver, also add job.compression_n_neurons=256 (it only affects the compression score).

conv-vae-1000to1 is the exception; it is scored with the standalone eval job:

uv run ntt job=eval \
    paths.run_name=conv-vae-1000to1 \
    paths.data_dir=/absolute/path/to/brainsets/processed \
    job.ckpt_path=hf://the-enigma-project/ntt-tokenizers/conv-vae-1000to1@v1.0 \
    evaluation@job.eval_config=hill_climb data@job.data_config=ntt_eval_light \
    job.split=test \
    ++job.eval_config.eval_window_seconds=5.0 \
    ++job.eval_config.max_eval_seconds=600 \
    ++job.eval_config.max_neurons=200 \
    job.compare_poisson=true job.poisson_rate_bin_ms=10.0 job.poisson_n_realisations=5 \
    job.output_path=results/conv-vae-1000to1

Without a precomputed job.poisson_reference, the Poisson baseline is computed in the same run on the same units as the model.

Scores are comparable to the published ones when they are produced with the same settings. If you aggregate results with ntt's aggregation job (job=aggregation, see eval_aggregation.md), its output metadata records a comparability_version; the published scores have comparability_version 6. The remaining axes (SNR, input and latent robustness, encoding, decoding) follow eval_battery_runbook.md and downstream_eval.md; use the same hf:// path as the checkpoint.

3. LDNS-Perceiver and NuCLR embeddings

The LDNS-Perceiver reads a NuCLR embedding for every unit it sees, from the embeddings table uploaded next to it. The uploaded perceiver supports only the sessions in that table, which are sessions of the public brainsets; on any other session it raises a KeyError naming the session.

It cannot be extended to other recordings by adding embeddings: a newly trained NuCLR encoder embeds units in a different space from the one the uploaded perceiver was trained on, and the NuCLR encoder behind the uploaded table is not released. For other recordings, train NuCLR, compute embeddings and train your own perceiver on them. ntt has the code for each step:

1. Train NuCLR: uv run ntt job=nuclr_contrastive paths.run_name=<name> (code in ntt.contrib.nuclr).

2. Compute the per-unit embeddings: uv run python -m ntt.contrib.nuclr.eval.extract_embeddings --ckpt-path <nuclr.ckpt> --data-root <brainsets root> --data-config <config> --split <split> --output-path embeddings.pt.

3. Optionally, save them as safetensors (the perceiver reads either format):

   import torch
   from ntt.contrib.tokenizer.perceiver.perceiver_tokenizer import save_embeddings_safetensors
   
   result = torch.load("embeddings.pt", weights_only=False)  # your own file (holds numpy arrays); keys: embs, unit_ids, session_ids
   save_embeddings_safetensors("my_embeddings.safetensors", result["embs"], result["unit_ids"], result["session_ids"])
   

4. Train a perceiver with uv run ntt job=train paths.run_name=<name> tokenizer@job.tokenizer_config=<your config> loss@job.loss_config=contrib/poisson_nll. Your tokenizer config (a YAML file under config/tokenizer/) has _target_: ntt.contrib.tokenizer.perceiver.perceiver_tokenizer.PerceiverTokenizerSpec and embeddings_path set to your embeddings file; the spec in the uploaded model's config.json lists the settings we used.

License

Weights, embeddings and cards: CC-BY-NC-4.0. The models were trained on public datasets, including data released under CC-BY-NC.

Citation

TBD