Xu-AI4Science/occupancyflow-affinity-esmc-300m

OccupancyFlow — affinity predictor (ESM C 300M)

Fine-tuned ESM C 300M (Synthyra/ESMplusplus_small) for protein–protein binding free energy (ΔG) prediction from sequence alone, used to build the occupancy-weighted human interactome in GainGod-Xu/OccupancyFlow.

This is the current, primary checkpoint, selected from a benchmark of 8 fine-tuned encoders (ESM-2 8M/35M/150M/650M, ESM C 300M/600M, Ankh base/large) trained and evaluated identically on the chain-group build of PPB-Affinity (Liu et al., Scientific Data 2024) — one sequence pair per complex on each side, matching how every other published affinity predictor defines the task. It supersedes an earlier ESM-2 650M checkpoint (Xu-AI4Science/occupancyflow-affinity-esm2-650m) selected on a chain-pair build that split multi-chain complexes (e.g. antibody heavy+light chains) into one-to-one chain pairs — mis-scoring 51% of the S79 benchmark's multi-chain complexes against a whole-complex label. See affinity_prediction/dataset/CHAIN_GROUPS.md for the full comparison.

Benchmark (chain-group build, one prediction per complex)

test set	n	MAE (kcal/mol)	RMSE	Pearson r	Spearman ρ
S90	90	1.535	2.085	0.601	0.551
S79	79	1.532	2.039	0.699	0.705

Best of 8 encoders on combined S90+S79 MAE (1.534, vs. ESM-2 650M's 1.589 under the same chain-group evaluation). Full 8-encoder table: benchmark_all_encoders.csv in this repo, or experiments/benchmark/benchmark_summary.csv in the GitHub repo.

Architecture

Each side (ligand, receptor) is encoded independently by the base model and mean-pooled; the two embeddings [h_A, h_B, |h_A - h_B|, h_A ⊙ h_B] are concatenated and passed to a 3-layer MLP head (512, 256 hidden units, GELU, dropout 0.1) predicting ΔG. Encoder and head are fine-tuned jointly with MSE loss, AdamW (lr 1e-5 encoder / 5e-4 head), early stopping on validation RMSE. max_len 2048. See affinity_prediction/model/train_finetune_benchmark.py in the GitHub repo for the exact training script.

Files

• best_ft_model.pt — fine-tuned checkpoint (encoder + MLP head state dict)
• log.json — training config, per-epoch history, final S79/S90 metrics
• predictions_s79.csv, predictions_s90.csv — per-complex predictions on both held-out sets
• benchmark_all_encoders.csv — all 8 encoders' results side by side, same evaluation protocol

Usage

Not a drop-in AutoModel — the checkpoint holds the combined encoder+head state dict from the custom PairRegressor module in affinity_prediction/model/train_finetune_benchmark.py. Load via that script's own load_encoder/model-construction functions rather than AutoModel.from_pretrained.

Citation / provenance

Part of the OccupancyFlow project (protein–protein binding affinity → competition-adjusted interactome weighting → disease-module network analysis). See the GitHub repo's top-level README and docs/data_provenance.md for full methodology and known limitations.