scOntoFM

scOntoFM scOntoFM is a novel biologically informed foundation model that explicitly integrates hierarchical cell ontology into representation learning. Our approach consistently improves performance across diverse downstream tasks, but also significantly improves the preservation of biological structure, offering a more robust tool for single-cell analysis.

This work has been accepted to NeurIPS 2026.

Model

Architecture	HuggingFace BertForMaskedLM: 6 layers, hidden size 256, 4 heads, intermediate 512, 8.9M parameters
Input	<cls> + top-2,000 genes by median-normalized rank + <eos> (max length 2,002), vocabulary 19,941
Cell embedding	last-layer <cls> state (256-d)
Training data	27,525,118 human cells from xiangyguo/scOntoFM-pretrain-corpus (234 CELLxGENE datasets, 210 tissues, 566 cell types, 58 diseases)
Training	3 epochs, 6 × H200, DeepSpeed ZeRO-2, fp16, AdamW lr 1e-3, global batch 1,536 triplets

The corpus repository has 29,160,112 cells. Unannotated, malignant and obsolete-term cells (1,634,994) are not used in training; see the dataset card for the breakdown.

Files

Path	Description
config.json, model.safetensors	pretrained weights
training_info.pkl	training metadata (total steps, epochs, batch size, date)
vocab/token_dictionary.pkl	gene → token id (19,941 tokens)
vocab/gene_medians.pkl	per-gene non-zero medians used for tokenization
assets/celltype_similarity_LCA_deepest_allnodes.pkl	LCA distances between all Cell Ontology terms (CL 2025-07-30)
assets/CL_id_map.pkl	{cell type: {dataset id: [row indices]}} of the training cells in the corpus
assets/triplet_ind_hard_top50.pkl	27,525,118 (anchor, positive, negative) row-index triplets used for pretraining

Usage

Tokenization (median normalization + ranking) is implemented in the scOntoFM code repository. The simplest route is its inference script, which takes a raw-count .h5ad with Ensembl gene IDs:

python scripts/run_inference.py --input_h5ad my_data.h5ad --output_dir ./output   # CLS embeddings

Loading the weights directly:

from transformers import BertForMaskedLM

model = BertForMaskedLM.from_pretrained("xiangyguo/scOntoFM").eval()
out = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
cell_embedding = out.hidden_states[-1][:, 0, :]      # <cls> token, 256-d

License

Model weights and code: MIT. The Cell Ontology is CC BY 4.0. The pretraining corpus is derived from CZ CELLxGENE Discover (CC BY 4.0).