TaxoFormer

Predict a taxonomic lineage directly from a protein sequence — from domain down to species, across 37 ranks, constrained to a known taxonomy tree.

TaxoFormer fine-tunes ESM-2 (650M) with a LoRA adapter (rank 8) and an autoregressive transformer decoder over a 15,000-token taxonomy vocabulary. Code: https://github.com/msparsa/taxoformer

Files in this repo

file	description
model.safetensors	trained delta: LoRA + attention pooling + taxonomy decoder (~300 MB, bf16). Base ESM-2 is not bundled — it is downloaded from the esm package at load time.
config.json	model hyperparameters
phylo2_mapping.json	tokenizer vocabulary (token ⇄ taxon name, 118 slots / 37 ranks)
parent_to_child_mapping.csv	legal parent→child taxonomy edges (for valid-tree decoding)

Usage

pip install git+https://github.com/msparsa/taxoformer

from taxoformer import TaxoFormer

model = TaxoFormer.from_pretrained("msparsa/taxoformer")
seq = "MGNKWSKGWPQIRERIRRTPPAAAEGVGAVSQDLDKHGAVTSSNMNN..."   # an HIV-1 protein
print(" > ".join(model.predict(seq, method="leaf_reconstruct")))
# Viruses > Riboviria > Pararnavirae > ... > Retroviridae > Orthoretrovirinae > Lentivirus > ...

Four decoding methods are available — greedy, min_edit, beam, leaf_reconstruct (default). You can also get a confidence score or the top-K valid trees:

model.predict(seq, method="leaf_reconstruct", return_confidence=True)  # {"lineage":..., "confidence":0..1}
model.predict_topk(seq, k=5)   # top-5 valid lineages, each with confidence + rel_prob

confidence is the model's confidence in the broad placement (domain → phylum/class), the geometric-mean per-rank probability over the broad ranks. It is positively (if weakly) correlated with how much of the true lineage is recovered (r ≈ +0.34 on a labeled set across all four superkingdoms), but it is not a correctness guarantee — short or novel proteins can be placed confidently but wrongly.

See the GitHub repo and the paper for benchmarks and details.

Limitations

Accuracy varies by taxon and input quality: strong on well-represented groups (e.g. HIV/retroviruses, many bacteria and eukaryotes), weaker on sparsely-seen or fragmentary entries, where it can confidently predict a wrong lineage. Within animals it resolves the broad placement (metazoan/vertebrate) reliably but typically predicts a generic vertebrate (mammal/bird/fish) rather than the exact class/order — so trust the domain→class part of the path more than the species leaf. Very short / low-complexity proteins are the most likely to be placed in the wrong superkingdom (often a confident collapse onto a short bacterial lineage).

Partial training set. This checkpoint was trained on a ~30M-sequence subset, not the full ~400M-sequence corpus, so it is not the model's ceiling — accuracy should improve with more data.

Sequence length. Training used only sequences of length ≤ 300 aa. Longer proteins (up to ESM-2's 1022-residue limit) still run but are out of distribution, so performance may drop.

Citation

@article{Parsa2026.06.06.730618,
    author = {Parsa, Mohammad and Azimian, Kooshiar and Wei, Kathy Y.},
    title = {TaxoFormer: Hierarchical Transformer for Predicting the Full Taxonomic Lineage of Protein Sequences},
    elocation-id = {2026.06.06.730618},
    year = {2026},
    doi = {10.64898/2026.06.06.730618},
    publisher = {Cold Spring Harbor Laboratory},
    URL = {https://www.biorxiv.org/content/early/2026/06/09/2026.06.06.730618},
    eprint = {https://www.biorxiv.org/content/early/2026/06/09/2026.06.06.730618.full.pdf},
    journal = {bioRxiv}
}

License

Code MIT. Base ESM-2 weights are © Meta AI under their own license.