MolE-RTD-15pct-S2

MolE-RTD-15% encoder + ChEMBL supervised pretraining (Step 2) — the full MolE two-stage protocol. Step 1: RTD-15% on 415M ZINC molecules. Step 2: frozen encoder + multi-task binary classification on 641 ChEMBL assays (512k molecules). Best all-round encoder for TDC ADMET benchmarks.

Based on MolE by Recursion Pharmaceuticals and DeBERTa-v3.

---

Two-stage pre-training

Stage	Objective	Data	Steps	Batch
Step 1	RTD (replaced token detection)	ZINC-Curated 415M molecules	1,000,000	512
Step 2	Sparse multi-task BCE	ChEMBL 36: 512k mols × 641 assays	80,000	512

Step 2 teaches biological context: the encoder learns to predict 641 bioactivity readouts from [CLS] embeddings. This injects pharmacological signal that pure SSL (Step 1 only) lacks.

Step 2 details

Setting	Value
Assays	641 (ChEMBL 36, ≥500 molecules/assay)
Molecules	511,898
TDC test molecules removed	12,834
Label sparsity	94.0%
Head	Linear(768, 641)
Loss	BCE with logits
Optimizer	AdamW (lr=5e-6, weight_decay=0.01)
LR schedule	Cosine with warmup
Hardware	1 × NVIDIA A100-SXM4-40GB
Wall time	~6-8 hours

---

Architecture

Component	Config
Layers	12
Hidden size	768
Intermediate size	3072
Attention heads	12 × 64
Vocabulary	211 atom environments (radius-0 Morgan)
Relative attention	Yes (p2c + c2p)
Position embeddings	No (bond-graph relative positions)

---

Results

Full TDC 22-task benchmark (frozen encoder + 2-layer MLP head, 3 seeds, scaffold split):

5 tasks beat MolE paper, 1 matches, 16 trail. Best on: bioavailability_ma (+0.040 AUROC), cyp2d6_veith (+0.012 AUPRC), dili (+0.016 AUROC), ld50_zhu (−0.058 MAE), ppbr_az (−0.057 MAE).

Task	Metric	Ours	MolE Paper
bbb_martins	AUROC	0.878	0.903
bioavailability_ma	AUROC	0.680	0.640
dili	AUROC	0.868	0.852
herg	AUROC	0.818	0.835
hia_hou	AUROC	0.973	0.984
ames	AUROC	0.799	0.834
lipophilicity	MAE↓	0.459	0.406

Recommendation: use caithmac/MolE-RTD-25pct-S2 for hard PK tasks (half_life, vdss, CYP substrates); use this model for saturated tasks (herg, bioavailability, hia_hou).

---

How to use

This is a clean encoder state dict (no generator/discriminator wrappers). Load directly:

import torch
from huggingface_hub import hf_hub_download
from DeBERTa.deberta.config import ModelConfig
from mole.training.models.mole import AtomEnvEmbeddings

DISC_CFG = dict(
    embedding_size=768, hidden_size=768, intermediate_size=3072,
    num_hidden_layers=12, num_attention_heads=12, attention_head_size=64,
    attention_probs_dropout_prob=0.1, hidden_dropout_prob=0.1,
    hidden_act="gelu", layer_norm_eps=1e-7, max_position_embeddings=0,
    max_relative_positions=512, position_buckets=0, norm_rel_ebd="layer_norm",
    pos_att_type="p2c|c2p", position_biased_input=False, relative_attention=True,
    share_att_key=True, type_vocab_size=0, vocab_size=211,
)

ckpt = hf_hub_download("caithmac/MolE-RTD-15pct-S2", "encoder_weights_rtd15_s2.pt")
encoder = AtomEnvEmbeddings(ModelConfig.from_dict(DISC_CFG))
encoder.load_state_dict(torch.load(ckpt, map_location="cpu", weights_only=False), strict=False)
encoder.eval()

See FINETUNE_AND_REPRESENT.md for complete fine-tuning + embedding extraction guide.

---

Citation

@misc{mole-rtd-15pct-s2,
  author = {caithmac},
  title  = {MolE-RTD-15pct-S2: Two-stage molecular encoder (RTD + ChEMBL supervised pretraining)},
  year   = {2026},
  url    = {https://huggingface.co/caithmac/MolE-RTD-15pct-S2}
}