license: apache-2.0
library_name: safetensors
pretty_name: 'TaxEmbed: a hyperbolic embedding of the NCBI tree of cellular life'
tags:
- biology
- bioinformatics
- taxonomy
- phylogenetics
- tree-of-life
- ncbi-taxonomy
- hyperbolic-embeddings
- poincare-ball
- representation-learning
TaxEmbed — a hyperbolic embedding of the complete NCBI tree of cellular life
TaxEmbed places all 1,102,163 named cellular taxa of the NCBI Taxonomy
(new_taxdump, release 2026-06-09) into a single 100-dimensional Poincaré
ball. It is a standalone, reusable representation of organismal relatedness:
a differentiable, fixed-dimensional metric of how organisms relate, learned from
one identical training recipe that holds from a few thousand taxa up to the full
1.1 million.
In the geometry:
- radius (embedding norm
‖x‖) tracks taxonomic depth — the root (cellular organisms) sits near the centre, tips near the rim (Pearson r ≈ 0.95 vs. root distance); - direction encodes learned lineage — taxa in the same clade point the same way (family-level kNN purity 0.907, angular separation up to 7.6× at family rank).
Files
| File | Size | Contents |
|---|---|---|
cellular_embedding.safetensors |
441 MB | The embedding matrix. Tensor key embedding, shape (1102163, 100), float32. |
taxid_to_index.tsv |
16 MB | Two columns taxid → idx. Row i of the matrix is the taxon whose idx == i. |
LICENSE |
— | Apache-2.0. |
All release metadata (dimension, taxon count, taxdump date, geometry) is embedded
in the safetensors header (__metadata__).
Quick start
import numpy as np, pandas as pd
from safetensors.numpy import load_file
emb = load_file("cellular_embedding.safetensors")["embedding"] # (1102163, 100) float32
idx = pd.read_csv("taxid_to_index.tsv", sep="\t").set_index("taxid")["idx"]
def poincare_distance(u, v):
"""Geodesic distance in the Poincaré ball."""
sq = np.sum((u - v) ** 2)
return np.arccosh(1 + 2 * sq / ((1 - u @ u) * (1 - v @ v)))
# Homo sapiens (9606) vs Mus musculus (10090)
d = poincare_distance(emb[idx[9606]], emb[idx[10090]])
print(f"Poincaré distance human–mouse: {d:.3f}")
# radius encodes depth: deeper taxa sit closer to the rim
print("‖x‖ human:", np.linalg.norm(emb[idx[9606]])) # ~ deep taxon, near the rim
To look up a taxon by name, join taxid_to_index.tsv against the NCBI
names.dmp from the same new_taxdump release (2026-06-09).
Model details
- Geometry: Poincaré ball, curvature −1,
d = 100. - Coverage: 1,102,163 named cellular taxa — Eukaryota 878,524; Bacteria 216,507; Archaea 7,131 (+ the root node). This mirrors NCBI's own sampling (~80% Eukaryota, ~0.6% Archaea); it is not a balanced census of the domains.
- Training signal: the depth-weighted transitive closure of the taxonomy (every ancestor–descendant pair), topology only — no branch lengths and no molecular data.
- Recipe (five load-bearing components): Euclidean tangent parametrization, a softmax/NLL loss over Poincaré distances, a depth curriculum, a radial nudge toward depth-based target radii, and a scale-aware effective batch. This single recipe is what carries the embedding to 1.1M taxa, the scale at which the prior (Nickel–Kiela) approach collapses.
Intended use
A general substrate for organismal relatedness that machine-learning models can consume directly: a differentiable phylogenetic prior, a relatedness metric for retrieval and taxonomy quality-control, or a factor to fuse with protein/other embeddings in one vector space. As one demonstrated application, the paper shows that taxonomy and protein function occupy linearly separable subspaces of a protein language model, so protein embeddings can be linearly debiased against phylogeny using this geometry.
Limitations
- Transductive. The embedding is fixed to the 2026-06-09 taxdump. A taxon added later has no coordinate; out-of-sample projection is future work.
- Topology only. It captures internal taxonomy consistency, not branch-length or molecular truth.
- Sampling skew. ~80% Eukaryota / ~0.6% Archaea, inherited from NCBI; this bounds any per-domain claim (the Archaeal arm in particular is thin).
- "Complete" means the named cellular taxa retained by the reproducible clean pass, not every NCBI node. Viruses are excluded by design.
License
Apache-2.0. Copyright 2025–2026 the TaxPointCare authors (@jcoludar and contributors). The implementation began as a fork of the Nickel & Kiela (2017) Poincaré-embeddings reference code and has since been fully reimplemented.
Citation
If you use this embedding, please cite:
Koludarov, I. & Rost, B. A hyperbolic embedding of the complete NCBI tree of cellular life (1.1 million named taxa) from a single recipe (2026).