theoschiff-biie/capo-binding-esm2-20260705-2355-7396-best

CAPO · ESM-2 multi-species ACE2–RBD binding predictor

Predicts, from a single SARS-CoV-2 RBD amino-acid sequence, whether that RBD binds each of 10 ACE2 orthologs — a 10-way multi-label classifier fine-tuned on top of facebook/esm2_t6_8M_UR50D with its last 4 transformer layers unfrozen plus an MLP head.

Produced end-to-end by CAPO (Compute-Aware Automated Protein-Model Optimization), run binding-esm2-20260705-2355-7396. This repo contains the best checkpoint only.

Model description

• Base encoder: ESM-2 (8M params, 6 layers, hidden size 320), facebook/esm2_t6_8M_UR50D.
• Head: masked mean-pooling over residue embeddings (CLS/EOS/PAD excluded) → Linear(320 → 256) → GELU → Dropout(0.1) → Linear(256 → 10).
• Architecture: custom ProteinBindingModel (see How to use — this is not an AutoModelForSequenceClassification).
• Outputs: 10 logits, one per species. Output-head order is exactly: human_new, mouse, rat, dog, cat, cattle, horse, ihbat, monkey, mink.
• Strategy: last-4-layers-unfrozen fine-tune. The frozen-encoder + head model is the baseline it was compared against.

Intended use & limitations

• Use: rank / classify ACE2-ortholog binding of RBD variants (e.g. cross-species spillover screening) for the SARS-CoV-2 Omicron BA.1 RBD mutational landscape.
• Input: an RBD amino-acid sequence (training used a fixed 201-aa RBD window).
• Limitations: trained on one BA.1 RBD mutagenesis library against 10 fixed orthologs — out-of-distribution RBD backbones, other species, or non-RBD proteins are unsupported. Labels are sparse and class-imbalanced (~80% negative), so use the per-species tuned thresholds below rather than a flat 0.5 where recall matters.

Training data

theoschiff-biie/ace2_binding — a validated mirror of the task's referenced (non-existent) BIIE-AI/ace2_binding. Yeast-display sequencing of a SARS-CoV-2 Omicron BA.1 RBD mutagenesis library screened against 10 ACE2 orthologs. The dataset's own train/test split was used; sparse multi-label rows are handled with a masked BCE loss (only observed (sequence, species) pairs contribute).

Training procedure

Setting	Value
Epochs	5
Optimizer	AdamW (weight_decay 0.01, max_grad_norm 1.0)
Learning rate	head 3e-4 · encoder 5e-5
Effective batch size	128 (per-device 128, grad-accum 1)
Precision	bf16
Sequence length	201 aa (fixed)
Loss	masked multi-label BCE
Hardware	1× A100-SXM4-40GB (Lambda)

A 150-step canary passed with 0 non-finite loss/grad before the full loop.

Evaluation

Held-out test set, macro-averaged over the 10 species.

Model	Macro-MCC (tuned)	Macro-MCC (@0.5)
This model (last-4 unfrozen)	0.8229	0.8261
Frozen-encoder baseline	0.4557	—
Δ vs baseline	+0.3672	—

Best validation: MCC 0.8958, val loss 0.1122 (rose monotonically over 5 epochs).

Per-species (fine-tuned model, tuned threshold)

Species	Tuned thr.	MCC	AUROC	Support
human_new	0.84	0.807	0.950	26,079
mouse	0.70	0.762	0.909	45,932
rat	0.54	0.967	0.997	15,370
dog	0.77	0.755	0.919	31,088
cat	0.65	0.920	0.972	75,548
cattle	0.76	0.759	0.932	53,255
horse	0.72	0.858	0.927	28,083
ihbat	0.68	0.845	0.927	36,475
monkey	0.85	0.624	0.941	16,637
mink	0.49	0.931	0.985	113,398

(AUROC shown is threshold-independent; monkey is the hardest species, consistent with its low prevalence of 8.9%.)

Files

• model.safetensors — encoder + head weights (~31 MB).
• config.json — {model_id, num_labels: 10, hidden: 320, proj: 256, head_dropout: 0.1, architecture: "ProteinBindingModel"}.

How to use

The checkpoint is a custom ProteinBindingModel (ESM-2 encoder + pooled MLP head), not a stock transformers class. Rebuild the module, then load the weights:

import json, torch, torch.nn as nn
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer, EsmModel

REPO = "theoschiff-biie/capo-binding-esm2-20260705-2355-7396-best"
SPECIES = ["human_new","mouse","rat","dog","cat","cattle","horse","ihbat","monkey","mink"]
cfg = json.load(open(hf_hub_download(REPO, "config.json")))

tok = AutoTokenizer.from_pretrained(cfg["model_id"])
SPECIAL = {tok.cls_token_id, tok.eos_token_id, tok.pad_token_id}  # excluded from the mean

class ProteinBindingModel(nn.Module):
    def __init__(self, model_id, num_labels=10, head_dropout=0.1, proj=256):
        super().__init__()
        self.encoder = EsmModel.from_pretrained(model_id, add_pooling_layer=False)
        h = self.encoder.config.hidden_size
        self.head = nn.Sequential(
            nn.Linear(h, proj), nn.GELU(), nn.Dropout(head_dropout),
            nn.Linear(proj, num_labels),
        )

    def _pool(self, hs, ids, mask):
        sp = torch.ones_like(ids, dtype=torch.bool)
        for t in SPECIAL:
            sp &= ids != t
        m = (mask.bool() & sp).unsqueeze(-1).to(hs.dtype)
        return (hs * m).sum(1) / m.sum(1).clamp(min=1.0)

    def forward(self, input_ids, attention_mask):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        return self.head(self._pool(out.last_hidden_state, input_ids, attention_mask))

model = ProteinBindingModel(cfg["model_id"], cfg["num_labels"], cfg["head_dropout"], cfg["proj"])
model.load_state_dict(load_file(hf_hub_download(REPO, "model.safetensors")), strict=False)
model.eval()

enc = tok("SEQ...RBD...SEQ", return_tensors="pt", truncation=True, max_length=203)
with torch.no_grad():
    probs = torch.sigmoid(model(enc["input_ids"], enc["attention_mask"]))[0]
print({s: float(p) for s, p in zip(SPECIES, probs)})  # apply per-species tuned thresholds

_pool masks the ESM special tokens (CLS/EOS/PAD) exactly as training did. For bit-identical pooling, use src/models/model.py from the CAPO run directory.

Provenance

• CAPO run: binding-esm2-20260705-2355-7396 (RUN_REPORT.md in the run dir).
• Model selection: bypassed — user-specified ESM2-8M + two-regime comparison.
• Cost: infra $12.26 (6h 9m on 1× A100-SXM4 @ $1.99/h); total run $38.43.
• Push: checkpoints/best/ only (private). checkpoints/last/ is retained locally and intentionally not pushed.

License follows the ESM-2 base model (MIT, Meta AI).