MolE-RTD-25pct-S2

MolE-RTD-25% encoder + ChEMBL supervised pretraining (Step 2). Step 1: RTD at 25% masking on 415M ZINC molecules. Step 2: frozen encoder + multi-task BCE on 641 ChEMBL assays. Best variant for hard pharmacokinetic tasks (half-life, volume of distribution, CYP metabolism).

Based on MolE by Recursion Pharmaceuticals and DeBERTa-v3.

---

Two-stage pre-training

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

Higher mask probability (25% vs 15%) motivated by ELECTRA: small-vocab domains (211 tokens here) make RTD task trivially easy → mask more tokens to maintain difficulty → better representations on hard tasks.

---

Results

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

Best on 5/22 tasks. Wins on hard PK: half_life (+0.183 over 15%), vdss (+0.065), CYP substrates (+0.05–0.08). Loses marginally on saturated tasks (herg −0.061, bioavailability −0.036).

Task	RTD-25%+S2	RTD-15%+S2	Δ
half_life_obach (Spearman)	0.473	0.290	+0.183
vdss_lombardo (Spearman)	0.623	0.558	+0.065
cyp2d6_substrate (AUPRC)	0.617	0.538	+0.079
cyp2c9_substrate (AUPRC)	0.372	0.321	+0.051
dili (AUROC)	0.903	0.868	+0.035

ASAP-Polaris Potency challenge: PARITY with #1 leaderboard. Aggregated MAE 0.500±0.022, Pearson 0.787±0.027.

TYK2 ΔΔG: RMSE 1.498, Pearson R 0.688 — beats Morgan ECFP4 (1.608/0.634).

---

Architecture

Component	Config
Layers	12
Hidden size	768
Intermediate size	3072
Attention heads	12 × 64
Vocabulary	211 atom environments (radius-0 Morgan)
Mask probability (Step 1)	25%
Relative attention	Yes (p2c + c2p)

---

How to use

Clean encoder state dict — 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-25pct-S2", "encoder_weights_rtd25_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()

Complete fine-tuning guide: FINETUNE_AND_REPRESENT.md.

---

Citation

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