brain-uce-pilot-mix-v3-d768l14v2

A pilot Universal Cell Embedding (UCE)-style model for brain single-cell RNA-seq, trained from scratch on "data mix v3" (mixture C), the same designed mixture over the v2026-09 corpus (26 caches, 147,675,214 cells, 13 species) as KuanP/brain-uce-pilot-mix-v3, for 262,144 steps at global batch 256 on 8 A100 GPUs. It keeps the UCE input (a chromosome-ordered "cell sentence" of frozen ESM2 gene tokens) and changes the network and objective to follow the public 33-layer UCE model more closely: a separate target-gene layer and an MLP cell head, a pre-norm RMSNorm / QK-norm / SwiGLU encoder (14 layers, d 768), L2-normalised gene tokens, mask 0.20 with dataset-expressed negatives and dropout 0.05. It has 114.0M trained parameters and outputs a 768-d, L2-normalised cell embedding; all earlier pilots in this series output 512-d embeddings.

Benchmark name	mixv3_d768l14v2_262k
Training run	brain_mixC_v2026_09_d768l14v2_noeval_262144steps/2026-10-18_03-33-38
Code	training: uce-training-suite (remote in training/git_info.json; a private repository, so this commit is not publicly reachable), branch ucsc-brain, commit 3c5ffe9c7893783a345fc510fb7914728e122cc9 plus uncommitted working-tree changes that include the v2 model options (listed in training/git_info.json; not published). Inference: UCE-brain, branch d768l14-v2 (section 4)
Weights	model.safetensors, 5,918,299,836 bytes, 151 fp32 tensors, final model at step 262,144
Parameters	114,014,209 trained + the frozen 266,403 x 5,120 ESM2 gene-token table (1,363,983,360)
Embedding	768-d, unit L2 norm (dot product = cosine similarity)
Status	pilot research model; license not yet chosen (see Caveats)

---

1. Training data: "data mix v3" (mixture C), the same draw as brain-uce-pilot-mix-v3

The corpus (CELLxGENE Census 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 0.995 / 0.005 train / held-out split (seed 42) and the mixture parameters (training/config.yaml: dataset.mixture) are those of brain-uce-pilot-mix-v3. The trainer materialised the same 67,108,864 samples: training/mixture_summary.json equals the v3 release's file except for cluster paths and an added all-zero heldout_sample counter, and training/dataset_indices_summary.json is byte-identical, so the held-out cells are the same. The difference is how the samples were consumed: 262,144 steps x global batch 256 instead of 131,072 x 512. The d768/12L baseline mixv3_d768l12_262k in section 3 was trained on the same draw at the same batch size and step count.

Mixture rules (see the v3 card for details): species mass proportional to (train cells)^0.7; brain cells re-weighted to at least 67 % within a species; dataset weight proportional to n^0.5 within each (species, brain / non-brain) stratum; sci-RNA-seq3 at half weight; no cache repeated more than 4x in expectation; the two CELLxGENE Discover datasets used for evaluation (0df5d172-5888-4cec-9a96-e637636bbfef, 843c1459-e852-429c-9339-17b5760e7716) and NeMO libraries duplicated in the BrainGenome caches excluded.

species	samples	share	of which brain	brain share within species
homo_sapiens	40,033,420	59.65 %	26,079,672	65.1 %
mus_musculus	15,516,985	23.12 %	9,254,226	59.6 %
macaca_mulatta	4,493,327	6.70 %	4,379,539	97.5 %
callithrix_jacchus	2,588,586	3.86 %	2,588,586	100.0 %
danio_rerio	1,734,868	2.59 %	863,391	49.8 %
rattus_norvegicus	568,217	0.85 %	568,217	100.0 %
pan_troglodytes	551,043	0.82 %	551,043	100.0 %
microcebus_murinus	416,833	0.62 %	220,333	52.9 %
sus_scrofa	375,664	0.56 %	251,695	67.0 %
monodelphis_domestica	265,449	0.40 %	265,449	100.0 %
macaca_nemestrina	238,504	0.36 %	238,504	100.0 %
aotus_nancymaae	190,748	0.28 %	190,748	100.0 %
tupaia_chinensis	135,220	0.20 %	135,220	100.0 %
total (13 species, 26 caches)	67,108,864	100 %	45,586,623	67.9 %

"Brain" is the tissue_general == "brain" label only (table identical to the v3 card, because the draw is identical). training/mixture_summary.json reports brain_share 0.712 instead: that is the mixer's expected share under its broader brain definition (tissue_general brain, central nervous system, nervous system or spinal cord).

---

2. Architecture and training recipe

Verified against config.json (the saved UCEConfig), training/config.yaml (the resolved Hydra config), training/training_args.json 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)
  -> L2-normalise each token (5120 dims)     token_norm: l2 (no LayerNorm(5120): use_embedding_layer_norm false)
  -> Linear(5120 -> 768) -> GELU -> LayerNorm   input_projection_type: linear_gelu_ln
  -> + sinusoidal PE x 0.0360844 (= 1/sqrt(768))   pe_scale; max_sequence_length 2048
  -> 14 pre-norm blocks                      encoder_impl: uce, norm_first: true; d_model 768, 6 heads x 128;
                                             RMSNorm (norm_type: rmsnorm); QK-norm (RMSNorm over head_dim on q, k);
                                             SwiGLU FFN, hidden 2048 (ffn_type: swiglu); no bias terms
                                             (encoder_bias: false); dropout 0.05 on PE, attention, FFN and residual branches
  -> final RMSNorm                           final_norm: true
  -> position 0 (CLS slot)                   embedding_reduction: cls
  -> cell head (cell_head_type: mlp, cell_head_dims [1024, 768, 768]):
       [Linear 768->1024, LayerNorm, GELU, Dropout 0.05]
       [Linear 1024->768, LayerNorm, GELU, Dropout 0.05]
       [Linear 768->768,  LayerNorm, GELU, Dropout 0.05]
       Linear 768->768
  -> F.normalize
  => cell_embedding [B, 768], unit L2 norm   (output_embedding_dim 768)

Training head (not needed for embeddings): each target gene's frozen token, L2-normalised, goes through a separate target-gene layer Linear(5120 -> 768) -> GELU -> LayerNorm (separate_gene_layer: true; earlier pilots reused the encoder's input projector here). The decoder takes the concatenation [cell embedding 768 | gene 768] and applies decoder_layer_dims [1536, 2048, 512, 128, 1] (Linear -> LayerNorm -> GELU -> Dropout 0.05 per hidden layer, no dropout on its input: decoder_input_dropout: false) to emit one logit.

What changed relative to the earlier pilots

	brain-uce-pilot-mix-v3	mixv3_d768l12_262k (d768/12L baseline)	this model
trained parameters	28,905,985	90,315,009	114,014,209
encoder	8 x 512, 4 heads, post-norm LayerNorm, GELU FFN 2048	12 x 768, 6 heads, post-norm LayerNorm, GELU FFN 3072	14 x 768, 6 heads, pre-norm RMSNorm + final RMSNorm, QK-norm, SwiGLU FFN 2048, no biases
token input	LayerNorm(5120) -> Linear -> LayerNorm -> GELU, PE x 1	same	L2 norm -> Linear -> GELU -> LayerNorm, PE x 1/sqrt(768)
cell head, output	Linear + LayerNorm, 512-d	Linear + LayerNorm, 512-d	3 MLP blocks + Linear on the CLS row, 768-d
decoder gene input	shared input projector	shared input projector	separate target-gene layer
decoder	Dropout -> 1024-512-512-1	Dropout -> 1280-512-512-1	1536-2048-512-128-1, no input dropout
dropout	0.1	0.1	0.05
masked genes / negatives	15 % / any vocabulary gene with zero count	same	20 % / dataset-expressed (below)
optimiser	AdamW (0.9, 0.999, 1e-8), wd 0.01	same	AdamW (0.9, 0.98, 1e-9), wd 0
learning rate	1e-4, 500 warm-up steps, cosine to 0	5e-5, 1,000 warm-up, cosine to 0	5e-5, 2,000 warm-up, cosine to a 1.5e-5 floor
steps x global batch, GPUs	131,072 x 512, 8 H100	262,144 x 256, 8 H100	262,144 x 256, 8 A100

The design follows the public UCE 33-layer model (Rosen et al.: separate gene_embedding_layer, full_block cell head and decoder, L2-normalised tokens, dropout 0.05, mask 0.20, negatives from the source dataset's genes), with a modern pre-norm encoder in place of its post-norm ReLU stack; QK-norm and the final norm are there for training stability.

Parameter count

From the safetensors header (151 tensors, all F32, 1,479,570,433 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 1 x 768)	1,572,864
trained: encoder (14 blocks + final RMSNorm)	99,116,288
trained: input projector uce.input_gene_embedding_projector	3,934,464
trained: cell head uce.cell_head	2,760,960
trained: target-gene layer gene_embedding_layer (training only)	3,934,464
trained: decoder expression_decoder (training only)	4,268,033
trained total	114,014,209

105,811,712 of the trained parameters (encoder, input projector, cell head) are used to compute an embedding; the training log of the run reports the same total (Trainable: 114,014,209). Weights are stored in fp32 (config.json: dtype float32); training ran under bf16 autocast.

Token vocabulary (vocab_size 266,403)

Identical to brain-uce-pilot-mix-v3 (vocab/ here has the same sha256 for every file): ids 0-3 special tokens (pad_token_idx 0, cls_token_idx 1, chrom_token_right_idx 2); ids 4-237,410 genes of the 11 non-zebrafish species; ids 237,411-239,337 their chromosome-start tokens (chrom_token_offset 237,411 + chromosome_id); ids 239,338-266,376 danio_rerio genes; ids 266,377-266,402 danio_rerio chromosome-start tokens. training/config.yaml: dataset.gene_mapping_path points at all_species_gene_dict_v2026-09.json, the file shipped in vocab/.

Objective: masked expressed-gene prediction with dataset-expressed negatives

Per cell (training/config.yaml: dataset):

1. Align the cell's raw counts to the species' vocabulary; expressed = count > 0.
2. Hold out mask_prop = 20 % of the expressed genes from the sentence.
3. Draw sample_size = 1,024 gene tokens with replacement, probability proportional to log1p(count); group 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.
4. Targets: 512 expressed genes (held-out genes first, topped up from the other expressed genes) and 512 non-expressed genes; loss = binary cross-entropy with logits over the 1,024 targets (chance ln 2 = 0.6931).
5. Dataset-expressed negatives (negatives_universe: dataset_expressed): the 512 negatives are drawn only from genes that have zero counts in the cell and are detected (count >= 1) in at least one of up to 10,000 sampled cells of the same source dataset in the same cache (685 cache x source-dataset groups). Earlier pilots drew negatives from every vocabulary gene with zero count, which includes genes the source assay never measures, a dataset-identity signal the embedding had to carry. Averaged over cells, a dataset's detected set holds 82.7 % of its cache's genes (median over the 685 groups 84.0 %).

Steps 1-4 are the same as in brain-uce-pilot-mix-v3 apart from the mask fraction. Dataset-expressed negatives only change training targets; embeddings depend on steps 1-3 only.

Optimisation (training/config.yaml: training, training/training_args.json)

setting	value
steps	262,144 (= 2^18); max_steps, no epochs
batch	per_device_train_batch_size 32 x 8 A100 (one node, ilc cluster) x gradient_accumulation_steps 1 = global batch 256
samples seen	262,144 x 256 = 67,108,864 sentences from the materialised mixture (section 1)
optimiser	adamw_torch_fused, betas (0.9, 0.98), eps 1e-9, weight_decay 0 (i.e. Adam), max_grad_norm 1.0
learning rate	peak 5e-5, warmup_steps 2,000, cosine_with_min_lr with min_lr_rate 0.3 (floor 1.5e-5, reached at step 262,144)
precision	bf16: true (autocast; fp32 master weights), torch_compile: true (inductor), DDP; the collator trims padding to a multiple of 256 tokens (exact: padded positions are masked out of attention)
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; residual output projections at std 0.02 / sqrt(28) (scaled_residual_init); the gene-token table was loaded from all_tokens_v2026-09.torch and never updated
wall clock	42 h 56 min (2026-10-17 to 2026-10-19), no restarts

Training loss. The logged training loss falls from 0.696 (step 20) to 0.406 (mean of the last 1,000 steps 0.4062; 0.4105 around step 131,072); the full curve, every 20 steps, is in training/trainer_state.json. This number is not comparable with the earlier pilots' training losses (brain-uce-pilot-mix-v3 0.371 and mixv3_d768l12_262k 0.369 at their last logged step): mask 0.20 and dataset-expressed negatives make the task harder (the removed negatives are genes never detected in the cell's source dataset, which are easy to call non-expressed).

---

3. Benchmark results (uce_brain_v1)

All metrics are computed on fixed per-dataset cell subsets 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), mask_prop = 0, bf16 autocast. The primary label per dataset is the benchmark config's label_col. Numbers below are read from the benchmark's results.parquet.

Compared models:

• mixv3_d768l14v2_262k: this model.
• mixv3_131k_gb512: KuanP/brain-uce-pilot-mix-v3 (512-d, 28.9M trained, same samples).
• mixv3_d768l12_262k: the d768/12-layer post-norm baseline on the same draw, steps and batch (512-d, 90.3M trained, mask 0.15, vocabulary negatives).
• public_uce_33l: the public UCE paper model (Rosen et al., 33l_8ep_1024t_1280.torch: 33 layers, d 1280, 1280-d output, legacy 145,469-token table), embedded with the public repository's own code (fp32, pad 1536) on the same cells. Its table covers fewer of the benchmark genes (0.95 of the human sets, 0.76 macaque, 0.78 marmoset; marmoset cells are embedded through its human table by identical gene symbol), and its code's fixed filter_genes(min_cells=10) step leaves 0.82-0.87 (human), 0.71 (macaque) and 0.74 (marmoset) of the benchmark genes actually fed to the model.

Ranks are among the 16 models scored in the benchmark (15 on the marmoset sets, where the human-only legacy model has no gene overlap), excluding continuations of this checkpoint; continuations of other checkpoints (the d768/12L BICAN mid-training arms) are included. This checkpoint's own BICAN mid-trained continuation (mixv3_d768l14v2_262k_bicanA, not released) scores higher than this model on 8 of 9 sets in both kNN and probe macro-F1; counting it, this model ranks 1st-3rd on kNN and 2nd on probe macro-F1 (15th on pediatric_glioma). All numbers are from the benchmark results as of 2026-10-20 (results.parquet sha256 e2a240ad...829f).

Metric definitions:

• 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 embedding (a diagnostic of effective dimensionality, not a score).
• loss: mean per-cell training objective with the model's own mask fraction and vocabulary-wide negatives (the benchmark has no dataset-expressed path); for this model that is mask 0.20 against 0.15 for the other suite models, so its loss column is not comparable.

Evaluation datasets

dataset	species	what it is	primary label / groups	cells scored
down_syndrome	homo_sapiens	fetal cortex, Down syndrome vs control (10x multiome nuclei; CxG 0df5d172): 8 CL cell types, 30 donors	cell_type / donor_id	60,000
pediatric_glioma	homo_sapiens	pediatric high-grade glioma (HTAN CHOP, 10x 3' v3 nuclei; CxG 843c1459): 7 CL cell types, 16 patients	cell_type / donor_id	60,000
hmba_v05_macaque	macaca_mulatta	HMBA whole-brain v0.5 macaque, 10x multiome nuclei not in pretraining: 4 donors	subclass / donor	57,661
hmba_v05_marmoset	callithrix_jacchus	HMBA whole-brain v0.5 marmoset, 10x multiome nuclei not in pretraining: 4 donors	subclass / donor	54,423
human_striatum	homo_sapiens	BICAN HMBA basal-ganglia atlas, caudate + putamen + nucleus accumbens ROIs, 9 donors	cell_type / donor_id	59,766
macaque_striatum	macaca_mulatta	BICAN HMBA basal-ganglia atlas, caudate + putamen + nucleus accumbens ROIs, 7 donors	cell_type / donor_id	59,822
hmba_v05_macaque_subclass_full	macaca_mulatta	same cells as hmba_v05_macaque, every subclass with >= 20 cells (127), capped at 300 cells per subclass	subclass / donor	22,798
hmba_v05_marmoset_subclass_full	callithrix_jacchus	same for marmoset (355 subclasses)	subclass / donor	86,368
human_striatum_full	homo_sapiens	human_striatum cells, every cell type with >= 20 cells (29), capped at 3,000 per type	cell_type / donor_id	53,766

The two Discover datasets were excluded from the training mixture. The striatum sets are cells of the BICAN basal-ganglia caches, which are part of the pretraining corpus, so they measure cell typing on in-distribution data; the HMBA v0.5 sets are the cross-species transfer test.

Donor-grouped kNN macro-F1, primary label (headline)

dataset (label)	this model	mixv3_131k_gb512 (brain-uce-pilot-mix-v3)	mixv3_d768l12_262k	public_uce_33l	rank of this model
down_syndrome (cell_type)	0.7971	0.7606	0.7612	0.7866	1 of 16
pediatric_glioma (cell_type)	0.8425	0.8226	0.8181	0.8255	1 of 16
hmba_v05_macaque (subclass)	0.7818	0.5534	0.6192	0.6556	1 of 16
hmba_v05_marmoset (subclass)	0.3166	0.1978	0.2325	0.2746	1 of 15
human_striatum (cell_type)	0.9166	0.7751	0.8071	0.9268	2 of 16
macaque_striatum (cell_type)	0.7956	0.6348	0.6711	0.7139	1 of 16
hmba_v05_macaque_subclass_full (subclass)	0.5161	0.3749	0.4206	0.3963	1 of 16
hmba_v05_marmoset_subclass_full (subclass)	0.1641	0.0990	0.1167	0.1381	1 of 15
human_striatum_full (cell_type)	0.8444	0.7170	0.7542	0.8517	2 of 16

First of 16 on 7 of 9 sets (on down_syndrome only by 0.0105 over the public 33L, which is within the likely noise, so read it as a tie); second to the public 33L on the two human striatum sets (0.9166 vs 0.9268 and 0.8444 vs 0.8517). Against brain-uce-pilot-mix-v3 the gain is +0.020 to +0.228 (largest on HMBA macaque, +0.228, and the striatum sets, +0.127 to +0.161); against the d768/12L baseline trained on the same samples it is +0.024 to +0.163.

Donor-grouped linear-probe macro-F1, primary label

dataset (label)	this model	mixv3_131k_gb512	mixv3_d768l12_262k	public_uce_33l	rank of this model
down_syndrome (cell_type)	0.8119	0.7881	0.7977	0.7992	1 of 16
pediatric_glioma (cell_type)	0.8106	0.8208	0.8185	0.7903	15 of 16 (first: matryoshka_v09_131k 0.8453)
hmba_v05_macaque (subclass)	0.7557	0.6093	0.6736	0.6774	1 of 16
hmba_v05_marmoset (subclass)	0.3580	0.2370	0.2788	0.3127	1 of 15
human_striatum (cell_type)	0.9396	0.8359	0.8831	0.9387	1 of 16
macaque_striatum (cell_type)	0.8445	0.7474	0.7958	0.7773	1 of 16
hmba_v05_macaque_subclass_full (subclass)	0.4965	0.4213	0.4233	0.4227	1 of 16
hmba_v05_marmoset_subclass_full (subclass)	0.2158	0.1419	0.1764	0.1817	1 of 15
human_striatum_full (cell_type)	0.8815	0.7998	0.8375	0.8685	1 of 16

Trade-off: more donor structure is kept

dataset	batch ASW: this model	mixv3_131k_gb512	mixv3_d768l12_262k	public_uce_33l	rank (lowest = 1)	donor mix: this model	mixv3_131k_gb512	mixv3_d768l12_262k	public_uce_33l	rank (lowest = 1)
down_syndrome	0.7274	0.8382	0.8570	0.7897	1 of 16	0.6543	0.7028	0.7014	0.6941	1 of 16
pediatric_glioma	0.7724	0.8465	0.8649	0.7974	1 of 16	0.2565	0.2955	0.2986	0.2624	1 of 16
hmba_v05_macaque	0.6798	0.8426	0.8754	0.7468	1 of 16	0.2561	0.3059	0.3074	0.3202	1 of 16
hmba_v05_marmoset	0.7447	0.8542	0.8858	0.7822	1 of 15	0.6539	0.6942	0.7014	0.7071	1 of 15
human_striatum	0.7378	0.8516	0.8691	0.7362	2 of 16	0.7261	0.7858	0.7926	0.7486	1 of 16
macaque_striatum	0.7484	0.8689	0.9069	0.7817	1 of 16	0.6302	0.6900	0.6993	0.7158	1 of 16
hmba_v05_macaque_subclass_full	0.6602	0.8218	0.8527	0.7397	1 of 16	0.2496	0.2843	0.2915	0.3113	1 of 16
hmba_v05_marmoset_subclass_full	0.7246	0.8454	0.8775	0.7600	1 of 15	0.6042	0.6216	0.6266	0.6429	1 of 15
human_striatum_full	0.7690	0.8777	0.8921	0.7900	1 of 16	0.6869	0.7550	0.7593	0.7017	1 of 16

The model has the lowest batch ASW of all 16 models on 8 of 9 sets (on human striatum the public 33L is 0.0016 lower) and the lowest donor-mixing entropy on all 9: a cell's nearest neighbours of the same type are concentrated on fewer donors than in any other model. Better cell-type separation (label ASW highest of 16 on 7 of 9 sets and second on human_striatum, 0.665 against the public 33L's 0.671, but 12th of 15 on hmba_v05_marmoset_subclass_full, 0.389, below brain-uce-pilot-mix-v3 at 0.416 and the d768/12L baseline at 0.446) is accompanied by more retained donor / batch structure; this is a cross-model observation with one seed per model, not a measured trade-off. Since all classification scores above are donor-grouped, this structure does not inflate them, but for integration-style uses (clustering across donors, donor-mixed neighbourhoods) expect to need batch correction more than with the earlier pilots. The embedding is also lower-rank: participation ratio 4.7-16.4 on these sets, below the d768/12L baseline's on all 9 (8.6-35.3).

mixv3_d768l14v2_262k, primary label, all metrics

dataset (label)	labels	folds	loss (not comparable)	kNN F1	kNN acc	probe F1	probe acc	donor mix	label ASW	batch ASW	PR	PCs95	genes mapped
down_syndrome (cell_type)	8	5	0.4009	0.7971	0.9217	0.8119	0.9233	0.6543	0.6055	0.7274	6.5900	43	1.0000
pediatric_glioma (cell_type)	7	5	0.4359	0.8425	0.9684	0.8106	0.9644	0.2565	0.6263	0.7724	16.3831	63	1.0000
hmba_v05_macaque (subclass)	52	4	0.3720	0.7818	0.8864	0.7557	0.8843	0.2561	0.6803	0.6798	12.5072	31	1.0000
hmba_v05_marmoset (subclass)	145	4	0.3830	0.3166	0.7262	0.3580	0.7610	0.6539	0.6411	0.7447	9.5059	29	1.0000
human_striatum (cell_type)	21	5	0.3616	0.9166	0.7885	0.9396	0.8608	0.7261	0.6654	0.7378	5.3535	23	1.0000
macaque_striatum (cell_type)	25	5	0.3602	0.7956	0.7529	0.8445	0.8640	0.6302	0.6220	0.7484	4.7059	16	1.0000
hmba_v05_macaque_subclass_full (subclass)	127	4	0.3777	0.5161	0.6744	0.4965	0.6858	0.2496	0.5832	0.6602	15.9904	37	1.0000
hmba_v05_marmoset_subclass_full (subclass)	355	4	0.3831	0.1641	0.2279	0.2158	0.3106	0.6042	0.3894	0.7246	11.8193	40	1.0000
human_striatum_full (cell_type)	29	5	0.3769	0.8444	0.9328	0.8815	0.9517	0.6869	0.7617	0.7690	13.1548	29	1.0000

"labels" = labels with >= 50 cells (>= 20 on the _full sets) that are scored; "folds" = number of GroupKFold folds, min(5, number of donors).

mixv3_d768l14v2_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.8732	0.7702	0.8772	0.7827
down_syndrome	disease	0.4908	0.4802	0.4884	0.4721
down_syndrome	development_stage	0.2479	0.1304	0.2794	0.1463
pediatric_glioma	merged_cellType	0.9684	0.8425	0.9671	0.8123
pediatric_glioma	TumorNormal	0.8087	0.7723	0.8380	0.8097
pediatric_glioma	molecularClass	0.4081	0.2120	0.4705	0.2362
hmba_v05_macaque	class	0.9539	0.7033	0.9469	0.6548
hmba_v05_macaque	neighborhood	0.9702	0.6819	0.9665	0.7288
hmba_v05_macaque	roi	0.2468	0.0225	0.2568	0.0224
hmba_v05_marmoset	class	0.8183	0.5134	0.8432	0.5562
hmba_v05_marmoset	neighborhood	0.8689	0.6182	0.8868	0.6509
hmba_v05_marmoset	roi	0.0499	0.0110	0.0595	0.0130
human_striatum	tissue	0.2594	0.1583	0.2774	0.1690
macaque_striatum	tissue	0.3036	0.2232	0.3094	0.2250

disease, TumorNormal and molecularClass are donor-level labels and roi / tissue are brain-region labels; under donor-grouped folds these sit near chance for every model.

Held-out pretraining loss (not comparable across these models)

Mean per-cell objective on up to the first 20,000 cells (1,512-20,000 per cache) of the held-out split of nine pretraining caches (the same cells for the three models below, since they share the split). This model is scored with mask 0.20, the other two with mask 0.15, all with vocabulary-wide negatives:

model	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
mixv3_d768l14v2_262k	0.3600	0.3644	0.3763	0.3632	0.3625	0.3331	0.3593	0.3734	0.3348
mixv3_131k_gb512	0.3756	0.3828	0.3851	0.3503	0.3726	0.3325	0.3662	0.3959	0.3631
mixv3_d768l12_262k	0.3745	0.3865	0.3846	0.3461	0.3717	0.3269	0.3658	0.3980	0.3646

One training seed per model, so run-to-run noise was not measured; differences of about 0.01 in F1 are likely within it. Two first places in the tables above fall inside that band and are best read as ties: down_syndrome kNN (+0.0105 over the public 33L) and human_striatum probe (+0.0009).

---

4. How to use (UCE-brain, branch d768l14-v2)

Load this model with the UCE-brain package at branch d768l14-v2 (to be merged into main). That branch adds the v2 architecture options (L2 token norm, linear_gelu_ln projector, PE scaling, the pre-norm RMSNorm / QK-norm / SwiGLU encoder, the MLP cell head and the separate target-gene layer) to UCE-brain's model, and keeps every earlier released checkpoint loading as before (bitwise-identical embeddings). Its load_model is now strict: it raises on config.json keys or tensors it does not know, which no released checkpoint has (strict=False downgrades this to warnings). Older UCE-brain versions refuse to load this checkpoint: their config check rejects input_projection_type: linear_gelu_ln. The uce_suite package at the training commit 3c5ffe9 does the same; the run used uncommitted v2 changes on top of it.

git clone https://github.com/snap-stanford/UCE-brain.git && cd UCE-brain
git checkout d768l14-v2
uv sync

Requires torch>=2.4 (nn.RMSNorm); uv sync installs the locked torch 2.10. Tested with torch 2.10 / transformers 5.2 and torch 2.13 / transformers 5.15.

Repository layout (its root is a checkpoint directory: config.json and model.safetensors side by side, with the run's config.yaml beside them, the first place UCE-brain looks for the run config that supplies the cell-sentence parameters):

brain-uce-pilot-mix-v3-d768l14v2/
  README.md                          this card
  MANIFEST.md                        sizes + sha256 of every file
  config.json                        UCEConfig (HF PretrainedConfig), including the v2 keys
  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 / list of uncommitted files of the training code
    dataset_indices_summary.json     train / val cell counts per cache
    mixture_summary.json             realised mixture draw per cache
    trainer_state.json               training log every 20 steps (loss, grad norm, learning rate); no optimizer state
    training_args.json               Hugging Face TrainingArguments of the run (token fields redacted)
  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

Snippet, tested on one A100 GPU with the d768l14-v2 code against a local copy of this repository (2,000 HMBA macaque cells -> float32 [2000, 768], unit norm, 151 tensors verified; UCE-brain's own GPU tests reproduce the training suite's embeddings, which the benchmark numbers above come from):

from uce_brain.data import load_gene_mapping, read_h5ad_subsampled
from uce_brain.inference import (
    resolve_checkpoint, find_gene_mapping, load_model, load_cell_sentence_params, embed_adata,
)

ckpt = resolve_checkpoint("KuanP/brain-uce-pilot-mix-v3-d768l14v2")   # Hub repo id, or a local copy
params = load_cell_sentence_params(ckpt)        # from config.yaml: pad_length 2048, chrom_token_offset 237411, ...
model = load_model(ckpt, device="cuda")         # fp32 weights; bf16 autocast is applied at embed time
gene_mapping = load_gene_mapping(find_gene_mapping(ckpt))   # vocab/all_species_gene_dict_v2026-09.json

adata = read_h5ad_subsampled("cells.h5ad", n_cells=5000)    # raw UMI counts
adata.obsm["X_uce"] = embed_adata(model, adata, gene_mapping, params, species="macaca_mulatta")
# float32 [n_cells, 768], unit L2 norm, row i == adata row i (mask_prop defaults to 0 for embeddings)

Sanity checks after loading: load_model logs that it verified 151 tensors against model.safetensors and that the model has 114,014,209 trainable parameters; embedding norms ~1.0 and width 768; pairwise cosine similarities well below 1.0 (mean 0.58 on the HMBA macaque cells); compute_per_cell_loss well below ln 2 = 0.6931: 0.373 on 512 HMBA macaque cells and 0.355 on 512 human striatum cells. It uses the training mask 0.20 but draws negatives from the cell's zero-count genes among the AnnData's mapped genes, not from the training's dataset-expressed set, so it approximates the training objective (the benchmark's loss column, computed by the training suite, spans 0.360-0.436). A loss pinned at 0.6931 means the weights did not survive loading. Downstream code that assumes 512-d UCE embeddings must read the width from the model (config.output_embedding_dim = 768).

---

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.
• Gene symbols: every vocabulary symbol is upper-case; UCE-brain's embed_adata matches symbols case-insensitively by default and resolves Ensembl var_names through var["feature_name"] (or gene_symbol_column). Genes absent from the species' vocabulary are dropped; the fraction mapped is logged.
• Species must be one of the 12 top-level keys of vocab/all_species_gene_dict_v2026-09.json (UCE-brain also accepts aliases such as human, macaque, marmoset): homo_sapiens (20,076 genes), 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). macaca_nemestrina (pig-tailed macaque) was trained through the macaca_mulatta vocabulary; 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.2 (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 embeddings that differ from the trained tokenisation (for this model, offset 1000 gives cosine 0.85 on macaque and 0.996 on human cells to the correct embeddings, with kNN scores and, on human cells, even the per-cell loss unchanged), so downstream scores will not reveal it. Always take the offset from the checkpoint (load_cell_sentence_params). negatives_universe / gene_universe_dir are training-only and are not read at inference.
• Per-cell RNG: sentences are sampled; UCE-brain seeds each cell's sampler with the cell's row position in the AnnData, so results are reproducible for the same AnnData row order and comparable across models. The same cell embedded in a different subset or at another position gets a different, equally valid draw (per-cell cosine about 0.985 to the benchmark's cached embedding, which seeds by dataset row; kNN accuracy within 0.006).
• 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. The per-dataset gene universes used for the negatives are not shipped either; both are only needed to train.

---

6. Caveats

• Pilot research model. One seed; the v2 recipe was chosen from short pilot runs, not a hyper-parameter search. Run-to-run noise was not measured; benchmark differences of about 0.01 F1 are likely within it.
• Donor structure. See the trade-off in section 3: this model keeps more donor / batch signal than any other model in the benchmark, by donor-mixing entropy on all 9 sets and by batch ASW on 8 of 9.
• Not a drop-in replacement for the 512-d pilots: different output width (768), different architecture keys in config.json, and it needs UCE-brain with v2 support (branch d768l14-v2, or main once that is merged).
• Training code is not public: the run used commit 3c5ffe9 of the private uce-training-suite repository plus uncommitted working-tree changes (the v2 model options among them). The architecture needed for inference is in UCE-brain (section 4).
• 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 designed, not uniform (same mixture as brain-uce-pilot-mix-v3): 60 % of the training samples are human and 23 % mouse; the eight smallest species contribute 0.2-0.85 % each and brain tissue 68 % of samples.
• 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.