counts_flash_hog_4gpu_a6_long_u2_lowilr_sample_tissue_cellcount_matched_random

Random 75% matched to the a6 DataRater selection's tissue x cell-count profile (volume- and composition-matched control).

UCE (Universal Cell Embedding) model trained on CELLxGENE 2025-01-30. This is the final model at 131,072 steps. One of a set of runs that share an identical training recipe and differ only in which training cells were selected, so their representations can be compared against each other.

Training recipe

Steps	131,072
Global batch size	512
Learning rate	5e-05 (cosine, 500 warmup)
Weight decay	0.0001
Precision	bf16
Architecture	8L / d512 / 4-head, frozen 5120-d gene embeddings
Parameters	773.7M total, 28.9M trainable
Sequence length	2048 tokens
Data selection	ExcludeDatasetSelector(dataset_id=53d208b0-2cfd-4366-9866-c3c6114081bc) + SampleCellCountMatchedRandomSelector(scores_path=/scratch/m000137-pm06/kuan/data_rater_infer_output/counts_datarater_flash_hog_4gpu_a6_long_u2_lowilr/final/all_scores.npy, dataset_path=/scratch/m000137-pm06/kuan/cellxgene_cache_dataset_wo_filter_2025-01-30_sparse/cellxgene_all_partitions_2025-01-30_wo_filter_sparse_sparse_raw.dataset, top_percentile=75, sample_keys=['dataset_id', 'donor_id', 'tissue', 'assay', 'disease'], tissue_key=tissue, aggregation=median, seed=42)
Source dataset	cellxgene_2025_top75_counts_flash_hog_4gpu_a6_long_u2_lowilr_sample_tissue_cellcount_matched_random_exclude_ts_sparse_marlowe (62,634,100 cells / 61,888 genes)

The gene embedding table is a frozen nn.Parameter, so it lives in the state dict -- hence the ~3 GB model.safetensors. You do not need all_tokens.torch at inference.

Usage

Requires uce-training-suite (uce_suite).

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

local = snapshot_download("KuanP/counts_flash_hog_4gpu_a6_long_u2_lowilr_sample_tissue_cellcount_matched_random")           # private repo: needs auth

model  = load_uce_checkpoint(local, device="cuda")
params = load_cell_sentence_params(local)             # reads config.yaml in this repo
gene_names, gene_mapping = load_gene_artifacts(
    f"{local}/gene_names.txt", f"{local}/all_species_gene_dict.json", species="human",
)

emb = embed_dataset(
    model, dataset_path="/path/to/cells.dataset",
    gene_names=gene_names, gene_mapping=gene_mapping, **params.as_kwargs(),
)

Use load_uce_checkpoint rather than a bare from_pretrained: on transformers >=5 from_pretrained re-runs _init_weights after loading and overwrites the trained nn.Linear weights. The loader force-reloads the state dict to undo that. The symptom if it is skipped is a per-cell loss pinned at log(2) = 0.693 and noise-like embeddings.

Tokenisation

config.yaml ships in this repo and is what load_cell_sentence_params reads. These values must match at inference or the embeddings are silently wrong:

pad_length	2048
cls_token_idx	1
chrom_token_offset	143574
chrom_token_right_idx	2
pad_token_idx	0
vocab_size	145469