Technical Summary: ProtSent (Protein Sentence Transformers)

Problem Statement

Protein Language Models (pLMs), such as ESM-2, generate per-residue representations that effectively capture evolutionary and structural information. However, their standard mean-pooled sequence embeddings are not explicitly trained to reflect functional, evolutionary, or structural similarity between entire proteins. Consequently, the embedding space geometry does not guarantee that biologically related proteins are spatially proximate. This limitation hinders the use of simple nearest-neighbor lookups for downstream tasks, which typically require training task-specific classifiers (probes) on top of the frozen embeddings. The authors draw an analogy to Natural Language Processing (NLP), where BERT-style models faced similar limitations until the introduction of Sentence-BERT, which used contrastive learning to restructure the embedding space so that semantically similar sentences become neighbors.

Methodology

The authors propose ProtSent (Protein Sentence Transformers), a contrastive fine-tuning framework designed to adapt pLM backbones into general-purpose embedding models.

Architecture and Training

• Backbone: The framework utilizes ESM-2 (both 35M and 150M parameter variants) wrapped in the SentenceTransformers framework. All Transformer layers are fine-tuned end-to-end.
• Embedding Generation: Sequence-level embeddings are obtained by mean-pooling over non-padding residue tokens.
• Loss Function: The primary objective is MultipleNegativesRankingLoss (MNRL), which optimizes the cosine similarity between positive pairs while treating other positives in the batch as negatives. For Deep Mutational Scanning (DMS) data, which provides continuous fitness scores rather than discrete pairs, the authors employ CoSENTLoss, a ranking-based objective that preserves the ordering of continuous scores.
• Data Sources: ProtSent is trained on five heterogeneous protein-pair datasets, sampled in a round-robin fashion to ensure equal representation of different biological signals:
  1. Pfam Families: 32.9M domains representing evolutionary and functional homology.
  2. Pfam Hard Negatives: 1.8M anchors with structurally derived "hard" negatives (mutants with high sequence similarity but deleterious effects) to force fine-grained discrimination.
  3. AlphaFold DB (AFDB) Structural Pairs: 133.9M sequences grouped by Foldseek structural clusters, providing large-scale structural supervision.
  4. STRING-DB: 36.5M protein-protein interaction pairs (medium confidence) capturing functional associations.
  5. Deep Mutational Scanning (DMS): 2.2M pairs from ProteinGym and clinical data, used as an auxiliary regression signal via CoSENTLoss.
• Evaluation Protocol: The authors evaluate performance using a frozen embedding protocol. The fine-tuned model encodes sequences, and a simple k-nearest-neighbor (KNN) probe (k=3, Euclidean distance) is trained on the resulting embeddings. This approach measures the intrinsic quality of the embedding space geometry rather than the capacity of a learned classifier.

Key Contributions

1. Framework Introduction: The introduction of ProtSent, a contrastive fine-tuning framework that combines five distinct protein-pair datasets with round-robin sampling to create general-purpose embeddings.
2. Performance Gains: Demonstration that contrastive fine-tuning yields substantial improvements on tasks dependent on embedding neighborhood quality, specifically remote homology detection and structural retrieval, without task-specific supervision.
3. Scalability: Validation that improvements are consistent across two model scales (35M and 150M parameters) and span diverse biological tasks (function, structure, engineering, and mutation).
4. Ablation Analysis: Systematic isolation of the contribution of each training data source and sampling strategy, revealing that Pfam families provide the dominant signal while hard negatives may introduce overly aggressive gradients for some tasks.
5. Resource Release: Public release of the models, code, training data, and training recipes.

Results

The authors evaluated ProtSent on 23 downstream tasks (binary classification, multiclass classification, and regression) and a structural retrieval benchmark (SCOPe-40).

• KNN Probe Performance:
  • ESM-2 150M: Improved 15 of 23 tasks. Notable gains include +105% on remote homology detection, +17% on variant effect prediction (GB1), and +19.9% Recall@1 on SCOPe-40 structural retrieval.
  • ESM-2 35M: Improved 16 of 23 tasks. Notable gains include +40.5% on remote homology and +15.5% Recall@1 on SCOPe-40.
  • Task Specifics: Binary classification tasks (e.g., PPI prediction, solubility) showed broad improvements. Regression tasks showed mixed results; while fluorescence and beta-lactamase fitness improved, some stability tasks degraded, suggesting the contrastive objective does not always preserve ordinal relationships within a single protein's mutational neighborhood.
• Structural Retrieval (SCOPe-40): ProtSent achieved consistent double-digit improvements in Recall@K across all thresholds, indicating that the embedding space successfully reorganizes to place structurally similar proteins closer together.
• Few-Shot Learning: In low-data regimes (N=50 to N=1000), ProtSent showed significant relative improvements over the baseline, particularly on tasks requiring structural or functional grouping (e.g., +244.5% on remote homology at N=100).
• Ablation Insights:
  • Removing Pfam families caused the largest performance drop, confirming evolutionary homology as the primary signal.
  • Removing Hard Negatives surprisingly improved the mean delta across tasks, suggesting that in-batch negatives alone may be sufficient and that hard negatives can perturb the embedding space for certain tasks.
  • Removing AlphaFold DB pairs significantly reduced performance on structural and enzyme classification tasks, highlighting the value of structural supervision.
  • Round-robin sampling proved robust, with proportional sampling yielding comparable results.

Significance and Claims

The paper claims that ProtSent successfully restructures the protein embedding space to better capture protein function and structure through contrastive fine-tuning, without requiring task-specific supervision. The authors position ProtSent as a "foundation model for universal use," arguing that if contrastive training successfully restructures the neighborhood structure, a simple neighbor-identity probe should be the most sensitive detector of improvement.

The work distinguishes itself from prior efforts (e.g., ProtTucker, PLMSearch) by its scale (training on over 70 million protein pairs from five sources) and breadth (evaluating on 23 diverse tasks rather than a single similarity axis). The authors acknowledge limitations, noting that the training data is restricted to the relationship types captured by the five sources and that they do not compare directly against specialized retrieval systems on matched benchmarks. They conclude that contrastive fine-tuning is a broadly applicable strategy for adapting pretrained pLMs for retrieval, clustering, and few-shot transfer.