You can deploy this model and start earning money today!

caithmac/
MolE-RTD-25pct

Name: MolE-RTD-25pct
Availability: InStock

2111.09543

Views

No views yet

MolE-RTD-25pct

MolE molecular encoder trained with DeBERTa-v3-style Replaced Token Detection at 25% masking rate on 415M ZINC-Curated molecules. Higher mask probability than standard RTD-15% — motivated by ELECTRA paper findings that small-vocab domains benefit from more aggressive masking.

Based on MolE by Recursion Pharmaceuticals and DeBERTa-v3. For the 15% variant see caithmac/MolE-RTD-ZINC415M.

Model description

MolE represents molecules as sequences of Morgan fingerprint atom environments (radius-0, vocab ~211 tokens) and encodes them with a DeBERTa disentangled-attention transformer. Bond distances from the molecular graph are passed as relative position biases — no absolute positional embeddings.

RTD pre-training trains a generator (small, 3-layer) to corrupt 25% of input tokens, and a discriminator (full-size, 12-layer) to detect which tokens were replaced. GDES (Gradient-Disentangled Embedding Sharing) prevents RTD loss from corrupting shared embeddings.

Architecture

Component	Config
Discriminator layers	12
Discriminator hidden size	768
Discriminator intermediate size	3072
Discriminator attention heads	12 × 64
Generator layers	3
Generator hidden size	256
Shared embedding size	768
Vocabulary	211 atom environments (radius-0 Morgan)
Mask probability	25%
RTD λ	50
Relative attention	Yes (p2c + c2p)
Absolute position embeddings	No

Pre-training details

Setting	Value
Dataset	SZU-ADDG/ZINC-Curated (~415M molecules)
Max atoms	96 heavy atoms
Steps	1,000,000
Effective batch size	512 (128 × 4 GPUs)
Optimizer	AdamW (lr=1e-4, weight_decay=0.01)
LR schedule	Cosine with 10k warmup steps
Hardware	4 × NVIDIA A100-SXM4-40GB
Wall time	~29 hours

Results (with Step 2 ChEMBL supervised pretraining)

Full TDC 22-task sweep. RTD-25% beats RTD-15% on hard pharmacokinetic tasks (vdss_lombardo +0.065, half_life_obach +0.183, cyp2d6_substrate +0.079). On easier tasks (herg, bioavailability), 15% is slightly better. Summary: 25% masking improves generalization on low-signal tasks at cost of slightly worse performance on saturated tasks.

Task	RTD-25%+S2	RTD-15%+S2	Best
half_life_obach (Spearman)	0.473	0.290	25%
vdss_lombardo (Spearman)	0.623	0.558	25%
cyp2d6_substrate (AUPRC)	0.617	0.538	25%
cyp2c9_substrate (AUPRC)	0.372	0.321	25%
dili (AUROC)	0.903	0.868	25%
herg (AUROC)	0.757	0.818	15%
bioavailability_ma (AUROC)	0.644	0.680	15%

ASAP-Polaris Potency (pIC50): PARITY with #1 leaderboard (MolE author's entry). Aggregated MAE 0.500±0.022.

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

How to use

python

from collections import OrderedDict
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", "mole_rtd_25pct_final.ckpt")
raw = torch.load(ckpt, map_location="cpu", weights_only=False)
sd = raw.get("state_dict", raw)

gen_w = sd["model.generator.embeddings.word_embeddings.weight"]
bias  = sd["model.disc_word_bias"]
enc_sd = OrderedDict()
enc_sd["embeddings.word_embeddings.weight"] = gen_w + bias
for k, v in sd.items():
    if not k.startswith("model.discriminator."): continue
    if ".embeddings.word_embeddings." in k: continue
    enc_sd[k[len("model.discriminator."):]] = v

encoder = AtomEnvEmbeddings(ModelConfig.from_dict(DISC_CFG))
encoder.load_state_dict(enc_sd, strict=False)
encoder.eval()

For Step 2 version (recommended for downstream use), see caithmac/MolE-RTD-25pct-S2.

Citation

bibtex

@misc{mole-rtd-25pct,
  author = {caithmac},
  title  = {MolE-RTD-25pct: RTD pre-training of MolE at 25% masking rate on 415M ZINC molecules},
  year   = {2026},
  url    = {https://huggingface.co/caithmac/MolE-RTD-25pct}
}