MolE-RTD-25pct

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

Pre-training details

---

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.

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

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

@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}
}