ESM-2 GFP Fitness Prediction Benchmark

A high-performance downstream regression and ranking head for predicting Green Fluorescent Protein (avGFP) log-fluorescence fitness from ESM-2 embeddings.

This repository contains the architecture, feature caching scripts, training logic, and baseline benchmarks for predicting protein fitness using the Meta AI ESM-2 (650M) protein language model.

---

Project Executive Summary

Traditional single-token representations ([CLS] or mean pooling) often dilute important local mutation signals in Deep Mutational Scanning (DMS) datasets.

This project implements BioYeti Multi-Vector Hybrid Pooling (CLS + Mean + Max), concatenating representation vectors into a unified 3840-dimensional embedding space per protein sequence. A non-linear regression head trained on top of these cached features achieves state-of-the-art ranking (Spearman ρ = 0.4787) and linear correlation (Pearson r = 0.6523), outperforming standard tree-based and linear baselines.

---

Final Benchmark Results

All evaluations were conducted on the validation split (N = 5,366) of the InstaDeepAI/true-cds-protein-tasks (avGFP task) dataset.

Model / Strategy	Pearson r ↑	Spearman ρ ↑	RMSE ↓
ESM2-650M + BioYeti Multi-Pool Head (Ours)	0.6523	0.4787	0.6432
ESM2-650M + LightGBM	0.6134	0.4055	0.6834
ESM2-650M + Ridge Regression (α = 10.0)	0.5387	0.4008	0.7986
ESM2-650M + SGD Linear	0.4886	0.3761	1.0104
Mean-Pooling Only (Baseline)	0.3068	0.2214	0.9102

---

Model Weights & Download

A prebuilt BioYeti head is attached to every release:

• bioyeti_head.safetensors — the trained head weights in HuggingFace safetensors format.
• config.json — the model card (backbone, embedding dim, training setup, validation metrics).

Latest release: https://github.com/Monster-ML/BioYeti/releases/latest

from safetensors.torch import load_file

state = load_file("models/bioyeti_head.safetensors")  # dict {tensor_name: tensor}
print(list(state.keys()))

The release workflow (.github/workflows/release.yml) builds the weights on a free CPU runner using a small ESM-2 backbone (8M), so the artifact is reproducible without a GPU. For the full 650M-model weights that match the table above, train locally on a GPU and publish them:

python src/extract.py
python src/train.py                       # saves models/bioyeti_head.pt
python src/export_model.py --checkpoint models/bioyeti_head.pt

---

Dataset Details

• Source: Hugging Face Datasets Hub (InstaDeepAI/true-cds-protein-tasks, subset fluorescence).
• Origin: Experimental Deep Mutational Scanning (DMS) of Aequorea victoria GFP (avGFP) fitness landscape (Sarkisyan et al.).
• Task: Sequence-level log-fluorescence intensity regression.
• Train split: 21,464 sequences.
• Validation split: 5,366 sequences.
• Test split: 27,217 sequences.

---

Architecture & Technical Specification

1. Feature Extraction Engine

• Backbone: facebook/esm2_t33_650M_UR50D — frozen transformer weights, 33 layers, hidden dim d = 1280.
• Multi-Vector Pooling: the three representations are concatenated along the feature axis into one 3840-dim vector:

h_hybrid = [ h_CLS ∥ h_Mean ∥ h_Max ]   ∈   ℝ³⁸⁴⁰     (d_ESM2 = 1280 → 1280 × 3 = 3840)

• h_CLS — sequence context token representation (learned sequence-level summary).
• h_Mean — mask-weighted global baseline sequence composition.
• h_Max — isolates non-zero localized mutation spikes across amino-acid substitutions (padding flushed to −∞ so it never wins the max).

2. BioYeti Multi-Pool Head Architecture

Input (3840-dim concatenated vector)
 │
 ├── Linear Layer (3840 → 512)
 ├── BatchNorm1d
 ├── GELU Activation
 ├── Dropout (p = 0.3)
 │
 ├── Linear Layer (512 → 128)
 ├── BatchNorm1d
 ├── GELU Activation
 ├── Dropout (p = 0.2)
 │
 └── Linear Layer (128 → 1) → Scalar Log-Fluorescence Output

3. Composite Fitness Loss Function

The model optimizes both absolute-value calibration and relative rank ordering simultaneously:

L_composite = L_MSE + λ · L_rank        (λ = rank_weight)

where L_MSE fits prediction magnitudes and L_rank is a pairwise margin / hinge term over all pairs in the batch that drives Spearman ρ.

---

Repository Layout

.
├── README.md                 # System instructions and documentation
├── LICENSE                   # MIT license
├── requirements.txt          # Python runtime dependencies
├── requirements-dev.txt      # Dev tooling (ruff)
├── requirements-notebooks.txt# Optional deps for the notebooks
├── .gitignore                # Git exclusions (caches, weights, data)
├── pyproject.toml            # Ruff / formatting configuration
├── package.json              # Husky + lint-staged git hooks (npm)
├── .github/
│   ├── workflows/
│   │   ├── ci.yml            # Lint + format + notebook validation
│   │   └── release.yml       # Builds .safetensors and creates a release
│   └── dependabot.yml        # Automated dependency updates
├── src/
│   ├── __init__.py
│   ├── extract.py            # Pre-extracts 3840-dim embeddings from ESM-2
│   ├── models.py             # BioYeti Neural Head & Composite Loss definition
│   ├── train.py              # Training execution script (saves checkpoint)
│   └── export_model.py       # Converts a checkpoint to .safetensors
├── scripts/
│   └── build_release_model.py# CI: trains + exports the release model
├── docs/
│   └── RESEARCH.md           # Curated research papers & reading path
├── notebooks/
│   ├── 01_dataset_eda.ipynb          # Explore the avGFP DMS dataset
│   └── 02_pooling_showdown.ipynb     # CLS vs Mean vs Max vs Hybrid (cookbook)
├── examples/
│   └── cookbook_protein_fitness.md   # Recipes for new datasets, tuning, eval
├── benchmark.py              # Comparative evaluation script vs. ML baselines
└── results/
    └── benchmark_report.csv  # Auto-generated performance metric comparison

---

Codebase Overview

• src/extract.py — Runs facebook/esm2_t33_650M_UR50D over the dataset and caches the hybrid 3840-dim embeddings as cached_train_features.pt / cached_val_features.pt (TensorDatasets). See the three pooling operations (CLS / mean / max) implemented here.
• src/models.py — Defines BioYetiMultiPoolHead, the MLP regression head (3840 → 512 → 128 → 1), and CompositeFitnessLoss, which blends MSE with a pairwise margin-ranking term.
• src/train.py — Loads the frozen backbone, pre-extracts embeddings in fp16, trains the head with AdamW + a cosine-warmup schedule, and reports Pearson r / Spearman ρ on the validation split. Saves a checkpoint to models/bioyeti_head.pt.
• src/export_model.py — Converts a trained checkpoint into HuggingFace .safetensors for distribution.
• scripts/build_release_model.py — What the release workflow runs: trains a demo head on a small ESM-2 backbone on CPU and exports .safetensors + config.json.
• benchmark.py — Trains cheap baselines (Ridge, SGD Linear, LightGBM) on the cached features and writes the head-to-head comparison to results/benchmark_report.csv.

---

Developer Setup

CI runs ruff (lint + format) on every push and PR; a local husky pre-commit hook does the same on staged files.

# Python tooling
pip install -r requirements.txt
pip install -r requirements-dev.txt

# Git hooks (optional, needs Node.js)
npm install          # installs husky + lint-staged, activates the pre-commit hook

The pre-commit hook runs ruff check --fix and ruff format on staged .py files. To bypass it in a pinch: git commit --no-verify.

CI / CD

• .github/workflows/ci.yml — On push/PR: ruff check, ruff format --check, and notebook JSON validation.
• .github/workflows/release.yml — On a v* tag: installs CPU-only torch, trains the demo head, exports .safetensors, and uploads it to a GitHub release.
• .github/dependabot.yml — Weekly dependency bump PRs for pip, npm, and GitHub Actions.

Cut a new release by tagging:

git tag v0.1.1
git push origin v0.1.1

---

Research & Learning Path

• docs/RESEARCH.md — A curated, annotated list of the papers this project builds on: the ESM/ProtTrans language models, the Sarkisyan et al. GFP landscape, FLIP/ProteinGym benchmarks, zero-shot mutation-effect methods (EVE, DeepSequence), and ranking losses — plus a recommended reading order.

Notebooks

• notebooks/01_dataset_eda.ipynb — EDA on the avGFP DMS dataset: split sizes, label distribution, sequence lengths, amino-acid composition.
• notebooks/02_pooling_showdown.ipynb — Cookbook-style ablation of CLS vs Mean vs Max vs Hybrid pooling on a small 8M-parameter ESM-2 (~2 min on CPU). Reproduces why hybrid pooling wins before you spend GPU time.

pip install -r requirements-notebooks.txt
jupyter lab

Cookbook & Examples

• examples/cookbook_protein_fitness.md — Copy-pasteable recipes: applying the pipeline to a new DMS dataset or raw CSV, tuning the composite loss, fighting VRAM limits, ProteinGym-style evaluation, exporting embeddings, and rebuilding the benchmark table.

---

Quickstart & Pipeline Execution

Step 1: Environment Setup

git clone https://github.com/Monster-ML/BioYeti.git
cd BioYeti
pip install -r requirements.txt

Step 2: Extract Embeddings (Feature Caching)

Runs facebook/esm2_t33_650M_UR50D over the Hugging Face dataset and caches the 3840-dim tensors locally (~20 s on GPU):

python src/extract.py

Step 3: Train the BioYeti Neural Head

Trains the multi-layer perceptron head on top of the cached representations:

python src/train.py

Step 4: Run Baseline Comparison Benchmarks

Runs fast evaluation against Ridge Regression, SGD Regressor, and LightGBM:

python benchmark.py

---

Citation & Attributions

• ESM-2 Backbone: Lin et al., Language models enable zero-shot prediction of the effects of mutations on protein function. bioRxiv (2022).
• Dataset: Sarkisyan et al., Local fitness landscape of the green fluorescent protein. Nature (2016). Curated under InstaDeepAI/true-cds-protein-tasks.
• BioYeti Head Architecture: Custom multi-vector pooling strategy designed for mutational fitness prediction.