brain-uce-pilot-mix-v2

A pilot Universal Cell Embedding (UCE)-style model for brain single-cell RNA-seq, trained from scratch on "data mix v2" (the v2026-09 corpus): 26 caches, 147,675,214 cells, 13 species, 43.1 % brain tissue. It maps a cell's raw counts to a 512-d, L2-normalised embedding through a chromosome-ordered "cell sentence" of frozen ESM2 gene tokens.

Benchmark name	multi_v09_262k
Training run	brain_multi_v2026_09_noeval_262144steps/2026-09-20_23-48-44
Code	uce-training-suite (remote in training/git_info.json), branch ucsc-brain, commit d9df996eb62ccebb3a22cdf89d72c4a678f91d35 (the only uncommitted change was a notebook execution-count diff)
Weights	model.safetensors, 5,575,765,828 bytes, 118 fp32 tensors, final model at step 262,144
Embedding	512-d, unit L2 norm (dot product = cosine similarity)
Status	pilot research model; license not yet chosen (see Caveats)

---

1. Training data: "data mix v2"

Aggregated from the per-(cache, species, tissue) cell counts of the corpus (data_mix_composition.csv, rows mix == "data mix v2"). Cell counts are the cells in each cache; the trainer split them 0.995 / 0.005 into 146,936,824 train / 738,364 held-out cells (training/dataset_indices_summary.json; that file totals 147,675,188, i.e. one cell per cache lost to integer rounding of the split).

Brain vs non-brain: 63,658,291 brain cells (43.11 %) vs 84,016,923 non-brain (56.89 %), where "brain" is the tissue_general == "brain" label only (spinal cord, "central nervous system" and "nervous system" rows count as non-brain in this tally).

species	cells	share	of which brain	brain share within species
homo_sapiens	109,664,883	74.26 %	42,040,766	38.3 %
mus_musculus	27,910,477	18.90 %	13,002,253	46.6 %
macaca_mulatta	4,979,990	3.37 %	4,859,948	97.6 %
callithrix_jacchus	2,740,657	1.86 %	2,740,657	100 %
danio_rerio	1,220,178	0.83 %	149,763	12.3 %
microcebus_murinus	264,361	0.18 %	26,224	9.9 %
rattus_norvegicus	247,698	0.17 %	247,698	100 %
pan_troglodytes	237,072	0.16 %	237,072	100 %
sus_scrofa	137,147	0.09 %	81,159	59.2 %
macaca_nemestrina	103,183	0.07 %	103,183	100 %
monodelphis_domestica	83,509	0.06 %	83,509	100 %
aotus_nancymaae	52,084	0.04 %	52,084	100 %
tupaia_chinensis	33,975	0.02 %	33,975	100 %
total (13 species, 26 caches)	147,675,214	100 %	63,658,291	43.1 %

By source:

source	caches	cells	share
CELLxGENE Census (2025-11-08 build): human, mouse, rhesus macaque, marmoset, chimpanzee	5	119,775,282	81.1 %
BICAN NeMO (multi-species brain)	12	17,847,041	12.1 %
BICAN BrainGenome (human, macaque, marmoset brain)	3	1,844,273	1.2 %
CELLxGENE Discover (human, mouse, rat, mouse lemur, pig, zebrafish)	6	8,208,618	5.6 %

Every cache was used whole (no metadata selectors, samplers or quality filters: dataset.selectors/sampler/quality_filters: null in training/config.yaml). Sampling was uniform over the concatenated train split, so the per-species shares above are also the shares of training examples.

---

2. Architecture and training recipe

Verified against config.json (the saved UCEConfig), training/config.yaml (the resolved Hydra config) and the tensor list of model.safetensors.

Model (UCEForExpressionPrediction, model_type: uce)

input_ids [B, 2048]
  -> Embedding(266403, 5120)            frozen ESM2 gene-token table (embedding_requires_grad: false)
  -> LayerNorm(5120)                    use_embedding_layer_norm: true
  -> Linear(5120 -> 512) -> LayerNorm   input_projection_type: linear
  -> + sinusoidal positional encoding   max_sequence_length 2048
  -> TransformerEncoder x 8             d_model 512, nhead 4, dim_feedforward 2048 (expansion_factor 4),
                                        GELU, post-norm (norm_first: false), dropout 0.1
  -> position 0 (CLS slot)              embedding_reduction: cls
  -> Linear(512 -> 512) -> LayerNorm -> F.normalize
  => cell_embedding [B, 512], unit L2 norm   (output_embedding_dim 512)

Training head (not needed for embeddings): an MLP decoder with layer dims [1024, 512, 512, 1] (decoder_dropout 0.1) that takes the concatenation of the cell embedding and a target gene's projected token embedding and emits one logit.

Parameter count from the safetensors header (118 tensors, all F32, 1,393,937,921 elements):

block	elements
uce.embedding_layer.embedding.weight (frozen ESM2 table, 266,403 x 5,120)	1,363,983,360
uce.pos_encoder.pe (fixed sinusoidal buffer, 2048 x 512)	1,048,576
trained weights (transformer 25,219,072; input projector 2,622,976; output projector 263,680; decoder 790,017; embedding LayerNorm 10,240)	28,905,985

Weights are stored in fp32 (config.json: dtype float32); training ran under bf16 autocast.

Token vocabulary (vocab_size 266,403)

Read from vocab/all_species_gene_dict_v2026-09.json and training/config.yaml:

• ids 0-3: special tokens (pad_token_idx 0, cls_token_idx 1, chrom_token_right_idx 2 = end of a chromosome block; no gene uses ids 0-3).
• ids 4-237,410: protein_embedding_id of the 11 non-zebrafish species' genes, contiguous per species in the order homo_sapiens, mus_musculus, macaca_mulatta, callithrix_jacchus, pan_troglodytes, rattus_norvegicus, monodelphis_domestica, sus_scrofa, aotus_nancymaae, tupaia_chinensis, microcebus_murinus.
• ids 237,411-239,337: chromosome-start tokens = chrom_token_offset (237,411) + chromosome_id for those 11 species (chromosome ids 0-1,926; 1,927 tokens).
• ids 239,338-266,376: danio_rerio genes (27,039), appended in v2026-09.
• ids 266,377-266,402: danio_rerio chromosome-start tokens (chromosome ids 28,966-28,991, chosen so that the same offset 237,411 still applies).

264,446 gene tokens + 1,953 chromosome tokens in total; there are no collisions between gene and chromosome token ids.

Objective: masked expressed-gene prediction

Per cell, in the DataLoader (uce_suite/data/cell_sentence_sampler.py, parameters from training/config.yaml: dataset):

1. Align the cell's raw counts to the species' vocabulary; expressed = count > 0.
2. Hold out mask_prop = 15 % of the expressed genes (their sampling weight is set to 0).
3. Draw sample_size = 1,024 gene tokens with replacement, with probability proportional to log1p(count).
4. Group the drawn tokens by chromosome, shuffle the chromosome blocks, sort by genomic location inside each block, wrap each block in [CHROM_x] ... [CHROM_END], prepend [CLS], pad to pad_length = 2,048.
5. Targets: positive_sample_num = 512 expressed genes (the held-out genes first, topped up from the other expressed genes) and negative_sample_num = 512 non-expressed genes. The decoder predicts expressed / not expressed for each from the cell embedding; the loss is binary cross-entropy with logits, averaged over the 1,024 targets. Classes are balanced by construction, so chance level is ln 2 = 0.6931.

Optimisation (training/config.yaml: training; train.py for the schedule)

setting	value
steps	262,144 (= 2^18); max_steps, no epochs
batch	per_device_train_batch_size 64 x 4 GPUs (A100, per the experiment description) x gradient_accumulation_steps 1 = global batch 256
cells seen	262,144 x 256 = 67,108,864 sentences, i.e. about 0.46 passes over the 146.9 M-cell train split
optimiser	AdamW (Hugging Face Trainer default), peak learning_rate 1e-4, warmup_steps 500, cosine decay (lr_scheduler_type="cosine" in train.py), weight_decay 0.01
precision	bf16: true (autocast; fp32 master weights), torch_compile: true (inductor), DDP
data	26 species specs, split 0.995 / 0.005 / 0 with seed 42; 10 dataloader workers, persistent
evaluation	none during training (eval_strategy: 'no'); checkpoints every 5,000 steps
initialisation	from scratch (experiment tag from_scratch); the gene-token table was loaded from all_tokens_v2026-09.torch and never updated

---

3. Benchmark results (uce_brain_v1)

Rows for multi_v09_262k copied from the benchmark summary. All metrics are computed on a fixed stratified subsample of 60,000 cells per dataset, with donor-grouped cross-validation (GroupKFold on the donor column), so a classifier cannot score by memorising a donor's batch effect. Embeddings are extracted with deterministic cell sentences (seeded by row index) and mask_prop = 0.

Metric definitions (from the benchmark):

• loss: mean per-cell training objective (BCE over 512 expressed + 512 non-expressed target genes, training mask); lower is better.
• kNN F1 / acc: cosine k-nearest-neighbour label transfer (k = 15, distance-weighted) to held-out donors; macro-F1 and accuracy, mean over folds.
• probe F1 / acc: multinomial logistic regression on standardised embeddings, same folds.
• donor mix: within each cell type, normalised entropy of donor ids among the 30 nearest neighbours (1 = donors fully mixed), mean over labels.
• label ASW / batch ASW: scIB-style silhouettes on a 10k subsample, (s(label)+1)/2 and mean(1 - |s(donor)|).
• PR / PCs95: participation ratio and number of PCs for 95 % variance of the 512-d embedding (a diagnostic of effective dimensionality, not a score).
• genes mapped: fraction of the dataset's genes present in the model's vocabulary.

Evaluation datasets

dataset	species	what it is	primary label / groups	cells in cache
down_syndrome	homo_sapiens	fetal cortex, Down syndrome vs control (10x multiome nuclei; CxG 0df5d172): 8 CL cell types, 15 author types, 30 donors	cell_type / donor_id	248,998
pediatric_glioma	homo_sapiens	pediatric high-grade glioma (HTAN CHOP, 10x 3' v3 nuclei; CxG 843c1459): 7 CL cell types, 16 patients, tumor/normal	cell_type / donor_id	401,253
hmba_v05_macaque	macaca_mulatta	HMBA whole-brain v0.5 macaque, 10x multiome nuclei not in pretraining: 348 subclasses / 43 classes / 9 neighborhoods, 85 ROIs, 4 donors (one donor = 79 % of cells)	subclass / donor	342,406
hmba_v05_marmoset	callithrix_jacchus	HMBA whole-brain v0.5 marmoset, 10x multiome nuclei not in pretraining: 365 subclasses / 42 classes / 9 neighborhoods, 106 ROIs, 4 donors	subclass / donor	729,273

The two HMBA v0.5 sets are non-human-primate whole-brain atlases whose cells were not seen in training; they are the cross-species transfer test.

multi_v09_262k, primary label

dataset	loss (lower better)	kNN F1	kNN acc	probe F1	probe acc	donor mix	label ASW	batch ASW	PR	PCs95	genes mapped
down_syndrome (cell_type)	0.4024	0.7632	0.8900	0.7918	0.9089	0.7047	0.5427	0.8564	26.1453	76	1.0000
pediatric_glioma (cell_type)	0.4355	0.8129	0.9679	0.8352	0.9713	0.2884	0.5838	0.8543	31.1589	76	1.0000
hmba_v05_macaque (subclass)	0.3840	0.5058	0.7013	0.5871	0.7657	0.3156	0.5358	0.8674	16.9694	64	1.0000
hmba_v05_marmoset (subclass)	0.3934	0.1779	0.5885	0.2304	0.6405	0.6962	0.5508	0.8603	12.9521	53	1.0000

multi_v09_262k, additional labels (donor-grouped kNN and probe)

dataset	label	kNN acc	kNN macro-F1	probe acc	probe macro-F1
down_syndrome	author_cell_type	0.7877	0.6692	0.8467	0.7438
down_syndrome	disease	0.4596	0.4542	0.4846	0.4756
down_syndrome	development_stage	0.1849	0.0906	0.2335	0.1212
pediatric_glioma	merged_cellType	0.9679	0.8129	0.9713	0.8349
pediatric_glioma	TumorNormal	0.8190	0.7802	0.8330	0.8053
pediatric_glioma	molecularClass	0.4691	0.2342	0.4832	0.2383
hmba_v05_macaque	class	0.8865	0.5120	0.9090	0.5398
hmba_v05_macaque	neighborhood	0.9491	0.4976	0.9545	0.5222
hmba_v05_macaque	roi	0.1631	0.0177	0.2193	0.0198
hmba_v05_marmoset	class	0.7312	0.3866	0.7764	0.4245
hmba_v05_marmoset	neighborhood	0.8226	0.5393	0.8521	0.5869
hmba_v05_marmoset	roi	0.0451	0.0093	0.0518	0.0111

disease, TumorNormal and molecularClass are donor-level labels; with 16-30 donors each grouped fold holds out only a few donors, so treat those rows as coarse.

Held-out pretraining loss

Mean per-cell training objective on the first 20,000 cells of this model's own validation split, per pretraining cache (identical cells and targets for every model that shares the cache):

bg_callithrix_jacchus	bg_homo_sapiens	census_homo_sapiens	census_macaca_mulatta	census_mus_musculus	discover_danio_rerio	discover_homo_sapiens	nemo_homo_sapiens	nemo_mus_musculus
0.3739	0.3814	0.3786	0.3457	0.3675	0.3356	0.3625	0.3940	0.3607

Headline comparison: macaque subclass transfer

hmba_v05_macaque, subclass label (348 subclasses), donor-grouped kNN, same 60,000 cells:

model	trained on	kNN macro-F1	kNN acc	loss	genes mapped
multi_v09_262k (this model)	13 species / 26 caches (147.7 M), v2026-09 vocab, 262k steps	0.5058	0.7013	0.3840	1.0000
human_v09_131k (species ablation)	human only, 4 caches (109.7 M), same vocab and recipe, 131k steps	0.3902	0.6235	0.4739	1.0000
human_census_legacy_150k	human CELLxGENE Census only (96.6 M), legacy 145k vocab, 150k steps at global batch 512	0.3531	0.5744	0.5209	0.5343

The multi-species model transfers markedly better to an unseen non-human-primate whole-brain atlas than either human-only model (0.506 vs 0.390 vs 0.353 macro-F1). The legacy model maps only 53 % of the macaque genes (and none of the marmoset genes, so it has no marmoset row); the ablation shares this model's vocabulary and recipe and differs only in the training data and step count.

One training seed per model; the benchmark's own guidance is that differences of a few 1e-3 in loss or about 0.01 in F1 are within run-to-run noise.

---

4. How to load

The repository is laid out so that its root is a checkpoint directory for uce_suite.inference (config.json + model.safetensors next to each other, and the run's config.yaml beside them, which is the first place load_cell_sentence_params looks):

brain-uce-pilot-mix-v2/
  README.md                          this card
  MANIFEST.md                        sizes + sha256 of every file
  config.json                        UCEConfig (HF PretrainedConfig)
  model.safetensors                  weights, incl. the frozen 266,403 x 5,120 gene-token table
  config.yaml                        resolved training config; supplies the cell-sentence parameters
  training/
    config.yaml                      same file, kept with the provenance set
    config_unresolved.yaml           the config with Hydra interpolations unresolved
    git_info.json                    commit / branch / dirty flag of the code that trained the model
    dataset_indices_summary.json     train / val cell counts per cache
  vocab/
    all_species_gene_dict_v2026-09.json        species -> gene symbol -> {protein_embedding_id, chromosome_id, location}
    gene_names_v2026-09/gene_names_<species>.txt   each species' vocabulary, one symbol per line, in token-id order

You need the uce_suite package from the uce-training-suite repository (commit in training/git_info.json; Python >= 3.12, torch >= 2.6, transformers >= 4.57.1, safetensors, datasets). load_uce_checkpoint contains a workaround for transformers >= 5, which otherwise re-initialises the trained linear layers after from_pretrained; the checkpoint itself was saved with transformers 5.15.0.

Untested snippet (the model was not run while preparing this card):

from huggingface_hub import snapshot_download
from uce_suite.inference import (
    load_uce_checkpoint, load_cell_sentence_params, load_gene_artifacts, embed_dataset,
)

repo = snapshot_download("KuanP/brain-uce-pilot-mix-v2")   # or a local clone
species = "macaca_mulatta"    # must be a top-level key of vocab/all_species_gene_dict_v2026-09.json

# 1. cell-sentence parameters, read from <repo>/config.yaml (pad_length 2048, chrom_token_offset 237411, ...)
params = load_cell_sentence_params(repo)

# 2. gene artifacts: the DATASET's gene axis (one symbol per line, line i == column i of your matrix)
#    and the species' entry of the shared gene mapping
gene_names, gene_mapping = load_gene_artifacts(
    gene_names_path=f"{repo}/vocab/gene_names_v2026-09/gene_names_{species}.txt",  # if your dataset uses the vocab gene axis
    gene_mapping_path=f"{repo}/vocab/all_species_gene_dict_v2026-09.json",
    species=species,
)

# 3. weights: config.json + model.safetensors at the repo root
model = load_uce_checkpoint(repo, device="cuda")   # fp32 weights; bf16 autocast is applied at embed time

# 4. embed an Arrow dataset (datasets.save_to_disk layout, see section 5)
emb = embed_dataset(
    model, "/path/to/my_cells.dataset", gene_names, gene_mapping, params,
    species=species, batch_size=64, num_workers=8, device="cuda",
    mask_prop=0.0,                     # default for embeddings; do not mask genes
    output_path="my_cells_uce.npz",    # optional: cell_embeddings + cell_indices
)
# emb: float32 [n_cells, 512], unit L2 norm, row i == row i of the dataset

CLI equivalent (uce-infer = python -m uce_suite.inference.cli). The saved config.yaml is a multi-species config without a top-level dataset.gene_names_path / dataset.species, so --gene-names, --gene-mapping and --species must be passed explicitly:

uce-infer --checkpoint /path/to/brain-uce-pilot-mix-v2 \
          --dataset /path/to/my_cells.dataset \
          --gene-names /path/to/brain-uce-pilot-mix-v2/vocab/gene_names_v2026-09/gene_names_macaca_mulatta.txt \
          --gene-mapping /path/to/brain-uce-pilot-mix-v2/vocab/all_species_gene_dict_v2026-09.json \
          --species macaca_mulatta --mode embed --output my_cells_uce.npz

Mapping to a training-run layout: the loader normally expects <run>/<timestamp>/checkpoint-<step>/{config.json,model.safetensors} with config.yaml one level up; here all three sit in the same directory, which the loader's first lookup (<checkpoint>/config.yaml) covers. config.yaml still contains the absolute cluster paths of the training caches; only its dataset.{pad_length, positive_sample_num, negative_sample_num, mask_prop, sample_size, cls_token_idx, chrom_token_offset, chrom_token_right_idx, pad_token_idx} keys are read at inference. Sanity checks after loading: embedding norms ~1.0, pairwise cosine similarities well below 1.0, and (with compute_per_cell_loss) a mean loss well below ln 2 = 0.6931; a loss pinned at 0.6931 means the weights did not survive loading.

---

5. Input expectations

• Raw UMI counts, not normalised or log-transformed. The sentence sampler treats count > 0 as "expressed" and samples tokens with probability proportional to log1p(count); a normalised matrix changes both.
• Dataset format: a Hugging Face datasets directory (save_to_disk layout) with, per cell, expr_indices (ascending column indices of the nonzero genes) and expr_values (their counts) — the sparse layout used in training — or a dense expr vector. Gene alignment is positional: line i of the gene_names.txt you pass is column i of the matrix. The training caches were built with the species' vocabulary list as their gene axis, so the simplest route is to reorder / subset your matrix to vocab/gene_names_v2026-09/gene_names_<species>.txt and pass that file as gene_names_path. The scripts/build_cache/03_build_arrow_cache.py script in the code repository converts .h5ad files into this layout and writes the matching gene_names.txt.
• Gene symbols are matched by exact string against the species' entry in the gene mapping; every vocabulary symbol is upper-case (e.g. mouse A730008H23RIK, LOC100041708). Genes absent from the mapping are dropped for that cell; the dataset view logs the retained fraction. Symbols, not Ensembl ids.
• Species must be one of the 12 top-level keys of vocab/all_species_gene_dict_v2026-09.json (number of vocabulary genes in parentheses): homo_sapiens (20,076), mus_musculus (22,198), macaca_mulatta (21,513), callithrix_jacchus (21,743), pan_troglodytes (22,543), rattus_norvegicus (23,154), monodelphis_domestica (21,248), sus_scrofa (22,080), aotus_nancymaae (21,779), tupaia_chinensis (21,079), microcebus_murinus (19,994), danio_rerio (27,039). The 13th training species, macaca_nemestrina (pig-tailed macaque), has no entry of its own: it was trained through the macaca_mulatta vocabulary (vocab_key: macaca_mulatta in training/config.yaml), so use that key for it.
• Cell-sentence settings, read from config.yaml: dataset and required to match training exactly: pad_length 2048, sample_size 1024, positive_sample_num 512, negative_sample_num 512, mask_prop 0.15 (training; use 0.0 for embeddings), cls_token_idx 1, pad_token_idx 0, chrom_token_right_idx 2, chrom_token_offset 237411. A wrong offset does not raise; it silently produces well-formed, meaningless embeddings.
• Per-cell RNG: sentences are sampled, so embeddings are stochastic in principle; embed_dataset seeds each cell's sampler with its row index (deterministic=True), which makes results reproducible across runs and comparable across models.
• Not shipped: all_tokens_v2026-09.torch (the 266,403 x 5,120 ESM2 gene-token table, 5.4 GB) is not included because the identical table is stored inside model.safetensors as uce.embedding_layer.embedding.weight; it is only needed to start a new training run.

---

6. Caveats

• Pilot research model. One seed, one recipe, no hyper-parameter search; benchmark differences of about 0.01 F1 are within noise. The 512-d embedding is low-rank in practice (participation ratio 13-31 on the evaluation sets above).
• License not yet chosen. The front matter says license: other because the authors have not selected a license; until they do, please contact them before any use beyond evaluation.
• Pretraining sources: CELLxGENE Census build 2025-11-08 (human, mouse, rhesus macaque, marmoset, chimpanzee); BICAN NeMO and BrainGenome (multi-species brain); CELLxGENE Discover (human, mouse, rat, mouse lemur, pig, zebrafish). The attribution and terms of use of those original sources apply to the data this model was trained on.
• Coverage is uneven: 74 % of the cells are human and 19 % mouse; eight of the thirteen species contribute under 0.2 % each, and zebrafish was added only in this corpus version. Expect weaker embeddings for the small species.
• Species outside the vocabulary (any key not listed above) cannot be embedded without building a new gene mapping and token table.
• No evaluation ran during training; the numbers in section 3 come from the separate uce_brain_v1 benchmark run after training finished.