StaMP v2 release: 5-seed ensemble weights, ProteinMPNN backbone, inference, honest model card
2f40071 verified | license: mit | |
| library_name: pytorch | |
| tags: | |
| - biology | |
| - protein | |
| - protein-stability | |
| - ddg | |
| - mutation-effect | |
| - proteinmpnn | |
| pipeline_tag: other | |
| # StaMP — Stability from Message-Passing | |
| **StaMP** is a small, fast, structure-conditioned predictor of **single-point mutation stability** | |
| (ΔΔG, kcal/mol) for proteins. Give it a backbone structure and it scores how much any single | |
| amino-acid substitution stabilises or destabilises the fold, or scans **every** possible | |
| substitution at every residue (a full in-silico deep-mutational scan) in a few seconds on CPU. | |
| > **Convention:** positive ΔΔG = **destabilizing**, negative = stabilizing. ΔΔG(x→x) = 0 by construction. | |
| It is built for the first filtering step of a protein-engineering campaign (enzymes, antibodies, | |
| therapeutics, binders): rank thousands of candidate point mutations cheaply, then take only the | |
| promising ones to the bench. | |
| - **Trainable parameters:** 591,893 (~592k). Released as a **5-seed ensemble**. | |
| - **Backbone:** frozen ProteinMPNN `v_48_020` (MIT, Dauparas et al. 2022), used only as a feature extractor — nothing in it is trained. | |
| - **Speed:** a full single-mutation scan of a ~60-residue domain is ~3 s on 8 CPU threads (features dominate; the head is negligible). | |
| ## How it works | |
| All heavy components are **frozen**; only a tiny head is trained. | |
| 1. **Frozen ProteinMPNN features**, computed once per chain → `R[L, 533]`: | |
| - the 3 encoder-layer outputs (3×128, structure-only), | |
| - the `W_s(wt)` sequence embedding (128), | |
| - conditional log-probs averaged over 8 random decoding orders (21) — a zero-shot stability signal. | |
| 2. **Structural neighbour context:** for each residue, the **48 nearest residues** on ProteinMPNN's own Cα k-NN graph are mean-pooled and concatenated with the centre features plus two burial counts (#neighbours within 8 Å and 12 Å). | |
| 3. **Gated head (592k params):** a gated linear unit → LayerNorm → MLP emits one score per amino acid; the prediction is `o[mut] − o[wt]`, which forces ΔΔG(x→x)=0 exactly. | |
| The neighbour context is what separates StaMP from a per-residue baseline: it adds **+0.026 Spearman** | |
| on the Megascale test split and carries over to the external benchmarks. | |
| ## Benchmarks | |
| Trained and evaluated on the **Megascale** dataset (Tsuboyama et al., 2023), using the homology-clustered | |
| train/val/test split released by the ThermoMPNN authors (mmseqs2 at 25% identity; 239/31/28 domains). | |
| 5-seed ensemble, selected on validation Spearman (never on test). | |
| ### Megascale held-out test (in-domain, n = 28,172) | |
| | Metric | StaMP (v2, 5-seed) | ThermoMPNN (ref) | | |
| |---|---|---| | |
| | Spearman | **0.734** | ~0.73 | | |
| | Pearson | **0.767** | ~0.75 | | |
| | RMSE (kcal/mol) | 0.705 | ~0.71 | | |
| Per-domain mean Spearman 0.786 / Pearson 0.823. StaMP slightly edges ThermoMPNN's headline | |
| Megascale Spearman with **~5× fewer trainable parameters** (592k vs ~2.7M). | |
| ### External generalisation (not seen in training) | |
| Structures from RCSB; mutations mapped by PDB residue number. We also apply our **own homology filter**: | |
| any benchmark protein with ≥25% identity to a Megascale *training* domain is removed (the unfiltered | |
| numbers that are usually quoted are inflated by homology). | |
| | Benchmark | n | Spearman | Pearson | | |
| |---|---|---|---| | |
| | S669 (all) | 667 | 0.453 | 0.415 | | |
| | **S669 (homology-filtered)** | 488 | **0.384** | **0.368** | | |
| | Ssym-direct | 342 | 0.655 | 0.666 | | |
| | *ThermoMPNN, Ssym-direct (published)* | 342 | — | *0.72* | | |
| ## ⚠️ Honest caveats — read before using | |
| - **In-domain vs zero-shot.** The headline Megascale numbers are **in-domain**: train and test are different | |
| proteins but the same assay and the same kind of small (≤72-residue) AlphaFold-modelled domains. This is | |
| **not** a zero-shot number. | |
| - **Performance drops out of domain.** On homology-filtered S669 (real crystal structures of larger, often | |
| multi-domain proteins) Spearman falls to ~0.38. On Ssym StaMP (0.67 Pearson) is **below** ThermoMPNN's | |
| published 0.72. StaMP trained only on AlphaFold models of tiny domains, so it is weaker on crystal | |
| structures of full proteins. We tested coordinate-noise augmentation to close this gap and it did **not** | |
| help — the gap is a training-data limitation, not a robustness trick. | |
| - **Direction.** Only direct (forward) mutations are validated; inverse-mutation (Ssym-inverse) needs mutant | |
| structures and was not evaluated. | |
| - **Not a replacement for ΔΔG from physics or experiment.** Use StaMP as a cheap ranking / prescreen, not as a | |
| calibrated free energy. | |
| ## Usage | |
| ```bash | |
| pip install torch numpy | |
| # predict one mutation (position is 1-based along the chain) | |
| python stamp/predict.py --pdb your.pdb --chain A --mut A42G | |
| # ...or use the PDB residue number | |
| python stamp/predict.py --pdb your.pdb --chain A --mut A42G --pdb_numbering | |
| # full single-mutation saturation scan -> TSV, ranked | |
| python stamp/predict.py --pdb your.pdb --chain A --scan --out your_scan.tsv | |
| ``` | |
| `predict.py` expects a **single-chain** backbone (N, CA, C, O); extract your chain of interest first. | |
| The ProteinMPNN backbone weights (`weights/v_48_020.pt`) and the five head checkpoints | |
| (`weights/stamp_s*.pt`) ship in this repo. | |
| ## Files | |
| ``` | |
| stamp/predict.py end-to-end inference (single mutation or full scan) | |
| stamp/protein_mpnn_utils.py ProteinMPNN model + PDB featurizer (MIT, Dauparas et al.) | |
| stamp/weights/stamp_s{0..4}.pt 5 trained head checkpoints (each stores its own cfg + feature norms) | |
| stamp/weights/v_48_020.pt frozen ProteinMPNN backbone (MIT) | |
| ``` | |
| ## License | |
| MIT. StaMP's head and code are released under MIT. The bundled ProteinMPNN code and `v_48_020.pt` | |
| weights are from the ProteinMPNN project (Dauparas et al., *Science* 2022), also MIT. | |
| ## Citation | |
| If you use StaMP, please also cite the work it builds on: | |
| - Dauparas et al., "Robust deep learning–based protein sequence design using ProteinMPNN," *Science*, 2022. | |
| - Tsuboyama et al., "Mega-scale experimental analysis of protein folding stability in biology and design," *Nature*, 2023. | |
| - Dieckhaus et al. (ThermoMPNN), whose Megascale split and benchmarking protocol StaMP follows. | |
| Built by [Modular](https://huggingface.co/Modularcomputing). | |