dn-gh/erba-kinetics-reproduction

ERBA reproduction — Enzyme-Reaction Bridging Adapter

From-scratch PyTorch reproduction of ERBA from Multimodal Protein Language Models for Enzyme Kinetic Parameters: From Substrate Recognition to Conformational Adaptation (Wang et al., CVPR 2026), https://arxiv.org/abs/2603.12845 — no official code was released, so this implements the paper text. Where the paper is ambiguous, the choice made here is documented below.

What is implemented

Component	Paper	File
MRCA — Molecular Recognition Cross-Attention (enzyme queries × substrate keys/values, residual + LN)	Eq. 5	model.py
G-MoE — Geometry-aware MoE, n=4 low-rank experts, top-2 routing, router on pooled seq–substrate + geometry context	Sec. 3.3, Fig. 3	model.py
ESDA — diagonal-free RBF MMD² alignment of Z¹, Z² to Z⁰ (median heuristic)	Eqs. 12–13	losses.py
Heteroscedastic Gaussian head on log₁₀ values (NLL)	Eq. 14	losses.py, model.py
Total loss L = L_task + 0.01·L_G-MoE + 0.1·L_ESDA	Eq. 15	model.py
Substrate MPNN (RDKit graph, Morgan-fingerprint prior as in CatPred) → per-atom tokens	Sec. 3.2, Fig. 2	encoders.py
Pocket E-GNN over active-site Cα coords → per-residue geometry tokens	Sec. 3.2	encoders.py
LoRA r=8, α=16, dropout 0.1 on the top third of PLM layers; backbone otherwise frozen; AdamW lr 1e-4	Sec. 4.1	train.py
CatPred split protocol: 80/10/10, no (sequence, SMILES) pair shared across splits	Sec. 4.1	data.py

Usage

pip install -r requirements.txt
cd <this repo's folder>
python train.py --csv kcat.csv --backbone esm2-3b --pocket-dir pockets/ \
    --epochs 30 --batch-size 8 --repo-id <you>/erba-kcat --space-id <trackio space>

Run from the folder containing the scripts (flat imports). Quick smoke test on a laptop/CPU: --backbone esm2-8m --no-geometry --epochs 1.

Input CSV columns: uniprot_id (optional, needed for pockets), sequence, smiles (canonical), value (log₁₀ of k_cat / K_m / K_i). Datasets per the paper are BRENDA+SABIO-RK curated as in CatPred/EITLEM (k_cat 23,197; K_m 41,174; K_i 11,929 records) — see https://github.com/maranasgroup/CatPred and https://github.com/XvesS/EITLEM-Kinetics for the underlying data.

Pocket files (--pocket-dir): <uniprot_id>.npz with ca_coords [N,3] (from ESMFold/OpenFold), pocket_idx [n] active-site residue indices, and optionally resnames [N] (3-letter). Without it the model runs the no-geometry ablation.

Documented ambiguities (paper does not specify)

1. Which PLM layer is tapped as H⁰ / how many "top layers" get LoRA → last hidden state used; LoRA on the top third of layers.
2. MRCA head count / d_k → 8 heads, d_k = D = 1024.
3. Geometry–sequence fusion op in G-MoE → cross-attention (H¹ queries, H_g keys/values) + residual + LN, by symmetry with MRCA.
4. Expert structure → per-token low-rank (rank 16) updates, per the paper's "pocket-local low-rank adaptation" wording.
5. L_G-MoE formula → Switch-style load-balancing loss (Eq. 15 lists the term but the HTML/PDF text does not give its form).
6. D=1024 vs ESM2-3B's 2560-dim hidden → linear projection to D.
7. Batch size, epochs, warmup, seeds, SMILES length cap → defaults here: batch 8, ≤30 epochs, cosine schedule with 10% warmup, early stopping (patience 5) on validation PCC, seed 42.

Target numbers (paper, in-distribution test)

Endpoint	R²	PCC	RMSE
k_cat (ESM2-3B + ERBA)	0.54	0.74	0.79
K_m	0.61	0.69	0.70
K_i	0.61	0.78	1.26