MolE-RTD-ChEMBL

MolE molecular encoder upgraded from MLM pre-training to DeBERTa-v3-style Replaced Token Detection (RTD), pre-trained on 1.36 million ChEMBL molecules.

Based on MolE by Recursion Pharmaceuticals and DeBERTa-v3. For the larger ZINC-415M version 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 the input by replacing tokens with plausible alternatives, and a discriminator (full-size, 12-layer) to detect which tokens were replaced. The discriminator sees a training signal on every token at every step (vs. only masked tokens in MLM), yielding richer representations per compute.

The discriminator is used as the encoder for downstream tasks.

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)
RTD λ	50
Relative attention	Yes (p2c + c2p)
Absolute position embeddings	No

Pre-training details

Setting	Value
Dataset	antoinebcx/smiles-molecules-chembl (~1.36M molecules)
Max atoms	96 heavy atoms
Steps	1,000,000 (best val at step 470,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	~2.5 days
Best val_mean_loss	1.2013 (step 470k)

---

Results

Fine-tuned on TDC BBBP (BBB_Martins, scaffold split, 3 seeds, 40 epochs):

Condition	Test AUROC	Std
MolE-RTD-ChEMBL (pretrained)	0.8782	±0.019
Random-init (same architecture)	0.8276	±0.039

Δ = +5.07 AUROC over random init, with ~2× lower variance.

---

How to use

Installation

git clone https://github.com/caithmac/mole-rtd
cd mole-rtd
pip install -e DeBERTa/
pip install -e mole_public/
pip install huggingface_hub

Load the encoder

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=64, 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,
)

def load_encoder(ckpt_path: str) -> AtomEnvEmbeddings:
    raw = torch.load(ckpt_path, 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
    cfg = ModelConfig.from_dict(DISC_CFG)
    encoder = AtomEnvEmbeddings(cfg)
    encoder.load_state_dict(enc_sd, strict=False)
    return encoder

ckpt_path = hf_hub_download("caithmac/MolE-RTD-ChEMBL", "mole_rtd_chembl_best.ckpt")
encoder = load_encoder(ckpt_path)
encoder.eval()

---

Citation

@misc{mole-rtd-chembl,
  author = {caithmac},
  title  = {MolE-RTD-ChEMBL: Replaced Token Detection pre-training of the MolE molecular encoder on 1.36M ChEMBL molecules},
  year   = {2026},
  url    = {https://huggingface.co/caithmac/MolE-RTD-ChEMBL}
}