Protein Set Transformer

Viruses are the most abundant biological entities on Earth and the least tractable to annotate. Standard viromics leans on homology — BLAST, HMM profiles, shared-gene networks — but viral genomes diverge so quickly that a newly assembled metagenomic virus often has no usable hit in any reference database. The result is that the majority of sequences recovered from environmental and clinical viromes are effectively dark matter: assembled, but unrelatable to anything known.

Protein Set Transformer (PST), from Karthik Anantharaman's and Anthony Gitter's groups at the University of Wisconsin-Madison, attacks this by changing the unit of representation. Rather than reading a genome as a nucleotide string or comparing gene-by-gene, PST models a genome as an unordered set of proteins, each already embedded by a protein language model, and learns to contextualize those proteins against one another at genome scale. The output is both a contextualized embedding per protein and a single learned embedding per genome, neither of which requires a functional label or a homology match to exist.

The released pretrained model — referred to as vPST, the viral PST — was trained on over 100,000 dereplicated viral genomes and outperforms homology-based and language-model-based alternatives at relating viral genomes by shared protein content. It also shows structural and functional awareness it was never supervised for: capsid-fold-containing proteins cluster with known capsid proteins, and late-gene proteins cluster within related viruses. The work first appeared on bioRxiv in July 2024 and was published in Nature Communications in 2025.

#Key Features

• Genomes as protein sets: Each genome is represented internally as a graph of fully connected, locally adjacent protein subgraphs, with multi-head attention contextualizing protein embeddings within the genome and learning per-protein weights for a weighted-average genome representation.
• Triplet-loss pretraining: Instead of masked language modeling, training uses a triplet objective that pulls each genome toward its most related neighbor and away from a less related one, mirroring how relationship-guided genomics actually reasons about relatedness.
• PointSwap data augmentation: Hybrid training genomes are synthesized by swapping similar protein vectors between a genome and its positive partner, expanding the effective training set without new sequence data.
• Position and strand encoding: Small vectors encoding genomic position and coding strand are concatenated onto the input protein embeddings, giving the set-based model access to synteny information a plain set transformer would discard.
• Two levels of output: Protein-level and genome-level embeddings drop into downstream tasks — clustering, taxonomy, host prediction, annotation transfer — without task-specific fine-tuning.

#Technical Details

PST takes ESM-2 protein embeddings as input, in either a small (6-layer, 8M-parameter ESM-2) or large (30-layer, 150M-parameter ESM-2) variant. Released checkpoints span 5.4M parameters for the small triplet-loss models up to 177.9M for the large ones, with masked-language-model variants from 23.8M to 185.8M; hyperparameters were tuned by leave-one-group-out cross validation over viral taxonomic realms and over protein-diversity groups. Pretraining used 103,589 viruses collected from 12 public sources, encoding more than 6 million proteins. Evaluation used two held-out sets chosen to be distinct at the nucleotide level from training data: 151,255 viruses from IMG/VR v4 and 12,857 viruses from soil metagenomes in MGnify. Baselines included unweighted averages of the input ESM-2 embeddings, GenSLM gene and genome embeddings, and HyenaDNA nucleotide embeddings.

#Applications

PST embeddings support the core operations of viral ecology on sequences that homology tools cannot place: clustering environmental viruses into meaningful groups, detecting remote relationships, transferring functional annotation between contextually similar proteins, and predicting bacterial hosts. The authors demonstrate host prediction by swapping PST genome embeddings into a CHERRY-style graph neural network. The framework is not virus-specific — the model architecture takes any set of proteins from any genome — so the same recipe applies to bacterial and archaeal genomics given suitable pretraining data.

#Impact

Code is MIT licensed and pretrained checkpoints plus training and evaluation data are distributed through Dryad, downloadable via a pst download command. The main caveat is scope: the released foundation model is trained on viruses only, so its demonstrated performance does not transfer automatically to cellular genomes, and the authors frame broader microbial applicability as a hypothesis rather than a result. Within viromics, the contribution is a concrete alternative to homology for the majority of sequences that homology cannot reach.