RBD-AIM — SARS-CoV-2 RBD antibody epitope classifier (weights)
Trained weights for the RBD-AIM antibody epitope classifier from Terekhov SS, Ivanisenko NV et al., "Mining antibody functionality via AI-guided structural landscape profiling", Nature Communications 2026. 📄 Paper
Given a paired antibody Fab (heavy + light chain sequences), the model predicts
which of 12 SARS-CoV-2 RBD epitope classes it binds:
A, B, C, D1, D2, E1, E21, E22, E3, F1, F2, F3.
Architecture
Two AntiBERTy encoders (one per chain) → per-chain pooling (CLS or CDR)
→ concatenation → linear head → multi-label logits.
Files
| Pattern | What it is |
|---|---|
NN_cv_0_{CLS,CDR}_{0..4}.pth |
Fold-0 ensemble members (trained holding out cross_val_test_0) |
NN_cv_{0..4}_{CLS,CDR}_full.pth |
"Full" models (trained on all non-test data) |
CLS = CLS-token pooling head; CDR = mean-over-CDR-residues pooling head.
How to load
The checkpoints use the original module names (model_l.*, model_h.*,
epitopes.*). The reference inference/serving code (with the loader that remaps
these to a clean model and runs the ensemble) is here:
https://github.com/<your-org>/rbdaim-mlops (rbdaim.predictor.EpitopePredictor).
import torch
sd = torch.load("NN_cv_0_CLS_full.pth", map_location="cpu") # OrderedDict
# keys: model_l.* (light encoder), model_h.* (heavy encoder), epitopes.* (head)
Evaluation
Cross-validated performance of the NN-model (AntiBERTy CLS+CDR ensemble), from out-of-fold predictions over the 10-fold CV (n=1,311 labeled; n=392 PDB-validated; 12-class, random ≈ 0.08). These reproduce the paper's Fig. 2h NN-model row to within ±0.001 (precision/F1 are class-size-weighted):
| accuracy | precision (wtd) | F1 (wtd) | MCC | MCC (PDB) |
|---|---|---|---|---|
| 0.670 | 0.686 | 0.668 | 0.636 | 0.559 |
The full RBD-AIM pipeline adds a structure/iPAE-based class adjustment on top of this NN-model (paper Fig. 2h: accuracy 0.673, F1 0.671, MCC 0.639).
Released checkpoints are the fold-0 ensemble members (NN_cv_0_*_{0..4}) and
the five "full" models (NN_cv_{0..4}_*_full). Reproducing the cross-validated
table above uses the released out-of-fold predictions; a single held-out fold
(fold 0, n=118) scores ≈ 0.62 accuracy — a high-variance subsample, not the
model's characterization.
Intended use & limitations
- Intended: research triage / prioritization of antibodies by likely RBD epitope.
- Out of scope: clinical/diagnostic decisions. Predictions are computational hypotheses only.
- Trained on ~1,300 labeled SARS-CoV-2 RBD antibodies; generalization to distant scaffolds or other antigens is unvalidated.
License & attribution
Released under the MIT License. This is an AntiBERTy-derived model; please preserve both notices:
- RBD-AIM code/method — MIT, © 2025 tanishAI — https://github.com/tanishAI/rbd-aim
- AntiBERTy — MIT, © 2023 Jeffrey Ruffolo — https://github.com/jeffreyruffolo/AntiBERTy
Citation
@article{terekhov2026rbdaim,
title = {Mining antibody functionality via AI-guided structural landscape profiling},
author = {Terekhov, S. S. and Ivanisenko, N. V. and others},
journal = {Nature Communications},
year = {2026},
url = {https://www.nature.com/articles/s41467-026-70553-6}
}