Two axes in the gene-embedding space of single-cell foundation models

What a protein language model already knows, and what requires single-cell pretraining

Code and results for the paper of the same name by Ihor Kendiukhov (University of Tübingen). It reimplements the gene-embedding probe battery of Taggart & Green (2025) once, runs it on twelve gene-embedding sources over a frozen 17,874-gene universe shared by all of them, attaches a resampling interval to every value and a paired test to every comparison, and adds four established gene-embedding benchmarks so that no conclusion rests on one probe collection.

The compiled paper is paper/main.pdf; the supplement is paper/supplementary.pdf.

What the study finds

• Uncertainty removes most of the comparisons. Of fifteen probes, only five resolve between a frozen ESM2 protein embedding and the best trained single-cell model after multiple-testing correction — one in favour of ESM2, four in favour of trained models. Ten of the fifteen probes rest on fewer than 60 genes.
• The protein baseline wins subcellular-compartment classification decisively: balanced accuracy 0.178 [0.150, 0.208] against 0.096 [0.080, 0.116] for the best trained model. The margin survives matched representation width, three of four classifiers, both compartment granularities and both gene universes.
• Single-cell pretraining adds information over the protein prior only on probes that depend on transcriptional covariation. Measured against a noise-matched control, the gain is resolvable in 21 of 120 model-by-probe cells and every one of them is such a probe.
• Grouping probes by biological topic does not work. On transcription-factor DNA-binding-domain family — a regulatory question with a protein-domain answer — protein-language-model sources reach 0.91–0.93 balanced accuracy against at most 0.33 for any trained model, while losing transcription-factor–target ranking to them. Probes have to be grouped by what determines the label, not by subject matter.
• On four established benchmarks the best trained checkpoint is beaten on all six scores by a protein-sequence representation, and on five of the six by an embedding of 32 hand-computed sequence descriptors. Removing homologous gene pairs does not change the ordering.
• A co-expression profile with no pretraining at all is the strongest source on repressor directionality and beats every trained checkpoint on three of the six established benchmarks.

Repository layout

two-axes/
├── paper/
│   ├── main.tex, main.pdf          manuscript
│   ├── supplementary.pdf           supplement
│   ├── references.bib
│   └── figures/                    all figures, PDF and PNG
│
├── analysis/                       the single implementation every source runs through
│   ├── harness/
│   │   ├── embeddings.py           registry, frozen universe, width and scaling control
│   │   ├── probes.py               the fifteen probes and the full-pipeline bootstrap
│   │   ├── refdata.py              HPA, ENCODE, UniProt, DGIdb loaders and curated gene sets
│   │   ├── genesets.py             GO, Reactome, hu.MAP 2.0, STRING, protein families
│   │   └── seqfeatures.py          32 interpretable protein-sequence descriptors
│   ├── experiments/                one script per experiment, each writing a JSON
│   ├── results/                    every result, including bootstrap replicate series
│   ├── tables/                     LaTeX fragments generated from results/
│   └── data/common_universe.json   the frozen 17,874-gene universe
│
├── models/                         per-model embedding extraction (phase 0) and per-model notes
├── docs/probe_battery.md           what each probe measures
├── DATA.md                         external data and model weights: how to obtain them
├── requirements.txt
└── CITATION.cff

Reproducing

1. Obtain the external data and model checkpoints (DATA.md) and extract each model's gene embedding with the corresponding phase0_extract_embeddings.py under models/.
2. Fetch the benchmark resources:

   python analysis/experiments/e9a_fetch_reference_data.py

3. Build the derived columns (co-expression reference, sequence descriptors, UCE input projection, independently computed protein-language-model column):

   python analysis/experiments/e13_coexpression_column.py
   python analysis/experiments/e7b_uce_projection.py
   python analysis/experiments/e3b_esm2_650m_own.py

4. Run the battery with uncertainty, once per source:

   python analysis/experiments/e4_battery_one.py geneformer_v2_104m --B 400

5. Run the remaining experiments (e5, e5b, e7, e9b, e12, e1e2e8, …), then rebuild the tables and figures:

   python analysis/experiments/build_outputs.py
   python analysis/experiments/build_figures.py

Every number in the paper is produced by build_outputs.py or build_figures.py from a file in analysis/results/; nothing is transcribed by hand. e0_replication_check.py verifies the harness against the per-model pipelines under models/.

What can be rebuilt without the embeddings

The extracted embeddings are too large to distribute, so steps 1 and 3 above are the expensive part. Two things do not need them, because they read only the committed result files:

python analysis/experiments/build_figures.py      # all seven figures
python analysis/experiments/build_supplement.py   # all seven supplementary tables

Both reproduce the published output byte for byte from a fresh clone. build_outputs.py is the one exception: the gene-embedding sources table reports each representation's width and vocabulary size, which are read from the embeddings themselves, so it needs step 1. If an embedding is missing it says so and names the path DATA.md tells you to populate.

Sources evaluated

Source	Gene representation
Geneformer V1-10M / V2-104M / V2-316M	static, learned from expression
scGPT (whole human)	static, learned from expression
AIDO.Cell-100M	contextual, forward pass on an input cell
scFoundation	contextual, forward pass on an input cell
tGPT	static word-embedding matrix
ESM2-15B (the tensor distributed with UCE)	frozen protein language model
ESM2-650M (computed here)	frozen protein language model, symbol-to-sequence mapping documented in full
UCE-100M input projection	trained projection of the protein embedding
HPA co-expression	154 cell-type expression profile, no pretraining
Sequence descriptors	32 hand-computed protein properties, no learning

Citation

See CITATION.cff.

License

MIT, see LICENSE.