ProtSent-V2-35M / README.md
GrimSqueaker's picture
Model card: normal card, no DMS claim, training config and results
497c50c verified
|
Raw
History Blame Contribute Delete
5.14 kB
---
library_name: sentence-transformers
tags:
- sentence-transformers
- protein
- esm2
- contrastive-learning
- protein-embeddings
- biology
license: mit
language:
- en
pipeline_tag: sentence-similarity
---
# ProtSent-V2 ESM-2 35M
Contrastively fine-tuned [ESM-2 35M](https://huggingface.co/facebook/esm2_t12_35M_UR50D) producing fixed-length protein
embeddings where biological similarity maps to embedding proximity. Intended for
retrieval, clustering, and nearest-neighbour transfer.
Retrained on a corpus decontaminated against the benchmark test sets. Predecessor:
[oriel9p/protsent-esm2-35M](https://huggingface.co/oriel9p/protsent-esm2-35M).
Other scale: [GrimSqueaker/ProtSent-V2-150M](https://huggingface.co/GrimSqueaker/ProtSent-V2-150M).
## Usage
```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
model = SentenceTransformer("GrimSqueaker/ProtSent-V2-35M")
emb = model.encode([
"MKTLLLTLVVVTIVCLDLGYT",
"MKTLLLTLVVVTIVCLDLGYN",
"AGWYRSPQEGLKPVDTFKDIV",
])
print(cos_sim(emb[0], emb[1:]))
```
Embeddings are mean-pooled over the final layer, dimension 480. Matryoshka heads at
64/128/256 are available by truncating the embedding.
## Training data
Three sources, all decontaminated (see below). **No DMS/ProteinGym component** — unlike
the V1 release, which included ProteinGym DMS pairs under a CoSENT loss.
| source | pairs used |
|---|---:|
| Pfam families | 777,306 |
| AlphaFold DB (Foldseek clusters) | 18,987,468 |
| STRING-DB v12 PPI | 15,000,000 |
| **total** | **34,764,774** |
Pfam and AFDB pairs are sampled within clusters; STRING is a fixed 15M-pair subsample
(seed 42) of the filtered pair table.
## Decontamination
Every source was searched against the benchmark test sequences with MMseqs2 `easy-search`
(corpus as query, 40% identity, 80% coverage, `--cov-mode 1`) and matching sequences
removed before training.
| corpus | rows before | rows after | removed |
|---|---:|---:|---:|
| Pfam | 28,530,684 | 27,929,772 | 2.11% |
| AlphaFold DB | 135,404,259 | 126,301,607 | 6.72% |
| STRING | 76,070,154 | 71,891,417 | 5.49% |
Filter targets: `biomap-research/fold_prediction` (remote homology) and
`Synthyra/bernett_gold_ppi` (PPI) test splits. The result was verified by semi-joining each
training file against the removal lists: zero flagged sequences remained.
SCOPe-40 was not a filter target — it has no train/test split, so filtering against it
would remove nearly all domain sequences from the corpus.
## Training configuration
| setting | value |
|---|---|
| backbone | ESM-2 35M (480 hidden, 12 layers) |
| loss | CachedMultipleNegativesRankingLoss |
| contrastive batch | 1024 per device |
| gather across devices | off |
| synthetic hard negatives | off |
| multi-dataset sampler | proportional |
| Matryoshka dims | 64 / 128 / 256 |
| max sequence length | 512 |
| optimiser | AdamW, LR 2e-4, `cosine_with_min_lr` |
| precision / attention | bf16, flash-attention-2 |
| hardware | 7x NVIDIA B300 |
| steps | 4,850 (one epoch) |
| gradient-cache mini-batch | 256 |
| warmup | 1,000 steps |
| wall clock | 10 h 53 m |
Training code: [github.com/oriel9p/ProtSent](https://github.com/oriel9p/ProtSent),
`train_esm2_35m.sh`.
## Results
SCOPe-40 structural retrieval, test split, self excluded, no-hit queries counted as
failures. Restricted to the 1,693 of 2,207 queries that have a non-self same-family protein
in the gallery.
| method | R@1 | R@10 | MAP |
|---|---:|---:|---:|
| ESM-2 35M | 0.4991 | 0.7614 | 0.4210 |
| MMseqs2 (`-s 7.5`) | 0.6556 | 0.7348 | 0.4041 |
| ProtSent-V1 35M | 0.5854 | 0.8512 | 0.5509 |
| **ProtSent-V2 35M** | **0.6852** | **0.9220** | **0.6459** |
Paired bootstrap over queries (10,000 resamples), V2 − V1: R@1 +0.0986 [+0.0762, +0.1211],
MAP +0.0943 [+0.0814, +0.1074]. Against a maximally sensitive HMMER (phmmer) this model is
statistically tied at top-1 and ahead at R@10 and MAP.
Remote homology (the task the corpus was filtered against), test split:
| model | 3-NN accuracy | linear-probe accuracy |
|---|---:|---:|
| ESM-2 35M | 0.5835 | 0.6868 |
| ProtSent-V1 35M | 0.6587 | 0.6899 |
| **ProtSent-V2 35M** | **0.6668** | **0.7016** |
## Limitations
- Under a **trained linear probe** on the final layer, this model is roughly neutral to
slightly worse than the stock ESM-2 backbone across a 23-task suite. The advantage is in
nearest-neighbour geometry, not in information a trained readout can extract.
- The final layer is not the best pooling layer for property prediction. In a layer sweep
on remote homology, an intermediate layer (~2/3 depth) scored higher for every model
tested, including the stock backbone.
- V2 differs from V1 in more than decontamination: no hard negatives, proportional
sampling, no DMS source, larger effective batch. It is not a controlled ablation of
filtering alone.
- Only the remote-homology and PPI test sets were decontamination targets; other benchmark
test sets were not filtered against.
## Citation
Paper: [ProtSent: Protein Sentence Transformers](https://doi.org/10.48550/arXiv.2605.06830)