Exploring the limits of pre-trained embeddings in machine-guided protein design: a case study on predicting AAV vector viability

Effective representations of protein sequences are widely recognized as a cornerstone of machine learning-based protein design. Yet, protein bioengineering poses unique challenges for sequence representation, as experimental datasets typically feature few mutations, which are either sparsely distributed across the entire sequence or densely concentrated within localized regions. This limits the ability of sequence-level representations to extract functionally meaningful signals. In addition, comprehensive comparative studies remain scarce, despite their crucial role in clarifying which representations best encode relevant information and ultimately support superior predictive performance. In this study, we systematically evaluate multiple ProtBERT and ESM2 embedding variants as sequence representations, using the adeno-associated virus capsid as a case study and prototypical example of bioengineering, where functional optimization is targeted through highly localized sequence variation within an otherwise large protein. Our results reveal that, prior to fine-tuning, amino acid-level embeddings outperform sequence-level representations in supervised predictive tasks, whereas the latter tend to be more effective in unsupervised settings. However, optimal performance is only achieved when embeddings are fine-tuned with task-specific labels, with sequence-level representations providing the best performance. Moreover, our findings indicate that the extent of sequence variation required to produce notable shifts in sequence representations exceeds what is typically explored in bioengineering studies, showing the need for fine-tuning in datasets characterized by sparse or highly localized mutations.

💡 Research Summary

This paper conducts a systematic evaluation of pre‑trained protein language model embeddings—specifically ProtBERT and ESM2—in the context of machine‑guided protein design where experimental datasets contain only a few, highly localized mutations. Using the adeno‑associated virus (AAV) capsid as a case study, the authors compare several embedding variants: (i) a global sequence embedding derived from the CLS token, (ii) an amino‑acid‑level embedding obtained by mean‑pooling all residue tokens, and (iii) a projected sequence embedding (CLS token passed through a dense layer with tanh activation) available only for ProtBERT. One‑hot encoding (OHE) serves as a traditional baseline.

The dataset comprises over 29 000 AAV2 capsid variants in which mutations are confined to residues 561–588 (a 27‑residue region of the 735‑residue protein). Each variant is annotated with two binary labels: viability (whether the variant produces functional viral particles) and design strategy (ML‑based versus non‑ML‑based design). This highly localized mutational landscape mirrors real‑world bioengineering scenarios where functional changes are driven by a small number of residues.

Unsupervised analyses (hierarchical agglomerative clustering and t‑SNE visualizations) reveal that OHE produces larger inter‑sequence distances and more clusters, while all embedding types generate compact representations with limited separation. None of the embeddings, including the projected variant, clearly separate viable from non‑viable sequences; however, they do show modest grouping by design strategy, indicating that embeddings capture some high‑level design information even without supervision.

Supervised experiments employ logistic regression, random forests, and deep neural networks to predict viability. Without any fine‑tuning, the amino‑acid‑level embeddings consistently outperform the global CLS embeddings and OHE, achieving the highest area‑under‑curve (AUC) and accuracy across models. This suggests that, for datasets dominated by sparse, local mutations, residue‑level information is more predictive than a compressed whole‑protein representation.

Crucially, the authors explore task‑specific fine‑tuning of the pre‑trained models using the viability labels. After fine‑tuning, the global sequence embeddings surpass all other representations, delivering the best predictive performance. Fine‑tuning is performed either by updating the entire network with a low learning rate or by retraining only the final classification head; both strategies improve performance, but the full‑network update yields the greatest gain for the CLS‑based embeddings. The projected embedding remains the weakest performer even after fine‑tuning.

The study concludes that pre‑trained embeddings are not universally optimal for protein engineering tasks involving limited, localized sequence variation. While amino‑acid‑level embeddings are advantageous for direct supervised learning without additional adaptation, global sequence embeddings become superior once they are fine‑tuned on task‑specific data. Consequently, researchers should consider both the nature of their mutational dataset and the availability of labeled examples when selecting an embedding strategy. The findings provide practical guidance for future work in viral vector optimization, therapeutic enzyme engineering, and antibody or peptide design, emphasizing the importance of fine‑tuning to unlock the full potential of transformer‑based protein representations.