ViralEmbed

Protein language models are trained one protein at a time. That framing is convenient and it throws away real information: proteins encoded by the same genome co-evolve, interact, and constrain one another, and in a viral genome — densely packed, with minimal intergenic space — those relationships are the dominant signal. A model whose context window ends at the stop codon cannot see any of it. ESM-2, for instance, is capped at 1,024 tokens, less than a single large viral protein pair.

ViralEmbed, from the University of Hawaiʻi at Mānoa and the University of Delaware, extends the context window of a protein language model across an entire viral genome. It supports sequences up to 61,000 amino acids on a single commodity GPU, enough to cover 98% of the sequenced viral genomes in the NCBI Virus database. The mechanism that makes this affordable is a biologically induced sparse attention pattern: instead of the fixed strided or windowed patterns used by text-oriented long-context transformers, ViralEmbed restricts attention to protein pairs that a first pass has identified as probable interaction partners, so the compute budget is spent where biology says the dependencies are.

The approach sits alongside genome-scale nucleotide models such as Evo and GenSLM, but keeps amino acids as the token and the protein as the unit of representation, producing embeddings that are directly comparable to those of single-protein models while carrying genome-wide context.

#Key Features

• Interaction-inferred sparsity: Candidate protein-protein interactions are inferred from attention patterns, and the top 50 ranked pairs per genome define the sparsity prior — a content-aware pattern rather than a static one.
• Block-sparse implementation: Attention is computed in the Block Sparse Row tensor format with a block size of 32 amino acids, executed through custom Triton kernels, which is what brings 61,000-token contexts onto a single 48 GB GPU.
• ALiBi positional encoding: Replacing rotary embeddings with attention-with-linear-biases flattens the length-dependent degradation that causes rotary-encoded models to inflate perplexity well before their nominal context limit.
• Two released variants: LV-3C is trained from scratch at up to 61,000 amino acids; LV-5B initializes from ESM-2 650M and trains at up to 48,000, retaining traces of its parent's attention structure.
• Attention that recovers known interactions: Attention heads enriched for long-range signal recapitulate documented couplings such as DNA polymerase I, ribonucleotide reductase, and helicase in replication.

#Technical Details

Training used the NCBI Virus database as retrieved in June 2024: 14,436 complete viral genomes totaling 683 million amino acids, segmented at protein boundaries and concatenated. Both variants trained for 12 epochs on NVIDIA A6000 GPUs, LV-3C from random initialization at an initial learning rate of 3 × 10⁻⁴ and LV-5B by transfer from ESM-2 650M at 3 × 10⁻⁵.

Both models show substantially lower perplexity than ESM-2 on fragments longer than 3,000 amino acids and hold that advantage out to 20,000, whereas pretrained ESM-2 and a rotary-fine-tuned variant degrade sharply past their training context. For protein embedding quality measured by silhouette score against structural clusters, LV-3C scores 0.83 and LV-5B 0.67 at long sequence lengths. A single-layer taxonomic classifier trained on the embeddings shows LV-3C achieving the highest F1 across every class-count setting from 5 to 100 species, ahead of fine-tuned ESM-2 and the Protein Set Transformer. Comparisons against Longformer, Shifted Sparse Attention, and BigBird patterns at a 12,000-amino-acid context establish that the biologically induced pattern outperforms static alternatives. Validating attention-derived interaction predictions against 146 STRING-annotated interactions across five viral genomes at a 0.5 confidence threshold gives an F1 of 0.74 and accuracy of 0.88.

#Applications

The immediate payoff is better representations for viral proteins, which are chronically underrepresented in the databases behind general protein language models and are often unannotated. Genome-context embeddings improve species-level taxonomic classification, which matters for viral metagenomics and phylogenetics, and the attention maps themselves function as a screen for candidate protein-protein interactions within a viral proteome — a computational complement to interaction assays that are difficult to run at viral scale.

#Impact

ViralEmbed reframes the long-context problem for biological sequences: rather than importing a sparsity pattern designed for prose, it derives one from the interaction structure of the genome, and shows that the biologically motivated pattern beats the generic ones on the same compute budget. The code is released on GitHub under an MIT license, and the work has been peer reviewed and published in GigaScience. Access is the main friction — model checkpoints are not distributed in the repository and must be requested from the authors, so the released artifact cannot be run zero-shot as downloaded. The evaluation is also intrinsic rather than functional: it measures perplexity, embedding clustering, taxonomy classification, and agreement with STRING, not downstream function or structure prediction. The authors point to bacteriophage and eukaryotic genomes as the next extension.