UPM

Unified Perturbation Modeling

🤗HuggingFace | GitHub | bioRxiv (Coming soon)

Overview  ·  Quick start  ·  H1 example  ·  Python API  ·  Model details

How Much Can AI Virtual Cells Gain from Unified Perturbational Modeling?

Overview

UPM predicts single-cell gene expression after an intervention, starting from a cell's control state. Genetic, chemical and cytokine interventions share a model that maps the starting expression profile and a perturbation representation to the predicted cellular response. Multiple interventions can be supplied together as one condition.

3 intervention modalities Genetic · chemical · cytokine

22,768 gene channels Native expression panel

1.09B parameters Standalone UPM checkpoint

Intervention	What you provide	What UPM predicts
🧬 Genetic	Gene name with crispri or crispra	Expression after repression or activation
💊 Chemical	Molecule name and dose in nM	A dose-conditioned expression response
🔬 Cytokine	Cytokine name	Expression after cytokine stimulation
Combination	Multiple components in one condition	A joint response using averaged component representations

Names are selected from the checkpoint's bundled vocabulary. Every condition uses the same supplied pool of control cells.

Quick start

1. Install

Use Python 3.10–3.12 and PyTorch ≥ 2.1. For GPU inference, install a CUDA-compatible PyTorch build first.

git clone https://github.com/youlab2025/UPM.git
cd UPM
python -m pip install -e .

The GitHub checkout includes the H1 example files used below. Both repositories currently require authorized access.

2. Download the complete checkpoint

hf auth login
upm-download --repo-id youlab2025/UPM \
  --revision 689cbbc8500e3b4f5ed7f7fad3ce1f2d299bd514 \
  --output checkpoints/upm

This fetches all six required files, checks their SHA-256 values and creates a local checkpoint directory. Allow approximately 10 GB of free disk space for the Hub cache and independent local copy. Running the same command again resumes downloads through the Hub cache. Older Hub installations may use huggingface-cli login.

Checkpoint contents · approximately 4.8 GB

File	Purpose
model.safetensors	FP32 inference weights, approximately 4.35 GB
embeddings.npz	Fixed perturbation embeddings, approximately 0.42 GB
config.json	Model architecture and dimensions
genes.json	Expression gene identifiers and aliases
perturbations.json	Gene, molecule and cytokine vocabularies
manifest.json	Bundle metadata and SHA-256 checksums

Load the directory as a complete bundle. The weights and fixed embeddings are both required for inference.

3. Run the H1 example

upm-infer \
  --model checkpoints/upm \
  --control examples/vcc/controls.h5ad \
  --conditions examples/vcc/conditions.json \
  --input-space log1p-cp10k \
  --device cpu \
  --batch-size 32 \
  --output predictions.h5ad

The output contains 96 predicted cells × 17,956 genes: 32 cells for each of the three conditions. Use --device cuda for GPU inference, or --device auto to select CUDA when available. Choose a new output filename for each run.

For automatic checkpoint download, use --model youlab2025/UPM --revision 689cbbc8500e3b4f5ed7f7fad3ce1f2d299bd514 in place of the local model path.

H1 example

A compact dataset from the Virtual Cell Challenge 2025 H1 human embryonic stem cell training split is included in examples/vcc. It contains 32 non-targeting controls and 32 independently sampled observed cells each for STAT1, MED12 and TADA1 CRISPRi.

File	Cells	Role
controls.h5ad	32	Starting control states used as inference input
observed.h5ad	96	Experimentally observed perturbation cells for inspection
conditions.json	3 conditions	STAT1, MED12 and TADA1 CRISPRi requests
provenance.json	—	Source rows, seed 0 sampling, normalization and checksums

Both H5AD files retain all 17,956 measured genes. Expression in .X is already log1p(CP10K), .var_names contains Ensembl gene IDs and .var["gene_name"] contains gene symbols. Control and observed cells are disjoint.

Expected example output: 32 controls × 3 conditions → 96 predictions. The observed perturbation cells are supplied for inspection; inference starts from controls.h5ad.

This example demonstrates the interface. It is not a model accuracy benchmark or a cell-by-cell pairing of predicted and observed cells. Data source: Virtual Cell Challenge 2025.

Python API

Load the checkpoint, align the input gene panel once and predict a response:

import anndata as ad

from upm import UPM, Perturbation
from upm.preprocessing import prepare_controls

adata = ad.read_h5ad("examples/vcc/controls.h5ad")
controls = prepare_controls(adata.X, input_space="log1p-cp10k")

model = UPM.from_pretrained(
    "youlab2025/UPM",
    revision="689cbbc8500e3b4f5ed7f7fad3ce1f2d299bd514",
    device="cpu",
)
model.align_genes(adata.var_names.tolist())

predicted = model.predict(
    controls,
    Perturbation("crispri", "STAT1"),
    batch_size=32,
)
print(predicted.shape)  # (32, 17956)

Try combinations and other intervention modalities

The examples below reuse model and controls from above. Inspect the bundled vocabularies to choose supported names.

# Two gene interventions in one condition.
gene_pair = model.predict(controls, [
    Perturbation("crispri", "STAT1"),
    Perturbation("crispri", "MED12"),
])

# Gene × drug: 1,000 nM = 1 µM.
drug_name = model.perturbations["molecules"][0]
gene_drug = model.predict(controls, [
    Perturbation("crispri", "STAT1"),
    Perturbation("molecule", drug_name, dose=1000),
])

# Cytokine stimulation.
cytokine_name = model.perturbations["cytokines"][0]
cytokine = model.predict(
    controls, Perturbation("cytokine", cytokine_name),
)

Model details

Property	This checkpoint
Model	UPM · standalone inference checkpoint used in the paper
Parameters	1,088,438,272
Inference precision	FP32
Native gene channels	22,768
Configured MLP layers	10
Hidden dimension	4,096
Gene embedding dimension	5,120
Molecule embedding dimension	2,048
Bundled vocabulary	20,115 genes · 571 molecules · 90 cytokines
Prediction mode	Deterministic VAE mean
Input	Control-cell expression and perturbation conditions
Output	Predicted log-expression in the input gene order
Hardware	CPU or CUDA
Checkpoint format	Safetensors + NumPy + JSON

The parameter count and configuration above are read from this release's files. Vocabulary sizes describe available bundled representations, rather than counts of experimentally validated conditions.

Using your own cells

Supply control cells from the cell context you intend to perturb. The CLI accepts a control-only AnnData file and a conditions JSON file.

• Raw counts: use --input-space counts. Keep all measured gene columns for normalization; add --layer counts if counts are stored in .layers["counts"].
• Already normalized: use --input-space log1p-cp10k.
• Gene IDs in a column: use --gene-id-column ensembl_id when Ensembl IDs are in .var["ensembl_id"].
• A new gene panel in Python: load a fresh model instance, then call align_genes once for that panel.

Input and output contracts

The Python predict method expects dense or prepared float32 log1p(CP10K) expression. Use prepare_controls for raw counts or sparse input. Molecule doses must be finite, nonnegative numbers in nM.

Unsupported perturbation names are rejected by default. Unmatched expression genes retain their control expression after native normalization and are marked by .var["upm_covered"]; use --unknown-gene error to require complete coverage. A panel with no matching genes is rejected.

The CLI writes predicted expression in .X, condition and control identities in .obs, gene coverage in .var, and inference provenance in .uns["upm"]. Output files preserve input gene order and are not overwritten.

Resources

• Inference code and full usage guide: youlab2025/UPM on GitHub
• Checkpoint and example files: youlab2025/UPM on Hugging Face
• Open-X-Cell: Training and inference code · Model checkpoints and examples
• Questions and reproducible issues: GitHub Issues
• License: CC BY-NC 4.0 (Creative Commons Attribution-NonCommercial 4.0 International).

Citation

If you use this code, please cite our paper.

@misc{upm2026,
  title={How Much Can AI Virtual Cells Gain from Unified Perturbational Modeling?},
  year={2026},
  note={Manuscript in preparation},
  url={https://github.com/youlab2025/UPM}
}