MolE-RTD-ZINC415M

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

Based on MolE by Recursion Pharmaceuticals and DeBERTa-v3.

---

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	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
Throughput	~9 it/s

---

How to use

Installation

# Clone the MolE-RTD repo (includes mole_public + DeBERTa fork)
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:
    """Load the discriminator encoder from an RTD checkpoint."""
    raw = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = raw.get("state_dict", raw)

    # Reconstruct discriminator word embeddings: gen_weight + disc_word_bias (GDES)
    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)
    missing, unexpected = encoder.load_state_dict(enc_sd, strict=False)
    print(f"Loaded encoder: missing={len(missing)}, unexpected={len(unexpected)}")
    return encoder

# Download from HuggingFace
ckpt_path = hf_hub_download("caithmac/MolE-RTD-ZINC415M", "mole_rtd_415m_final.ckpt")
encoder = load_encoder(ckpt_path)
encoder.eval()

Encode molecules

import torch
from torch_geometric.utils import to_dense_adj, to_dense_batch
from mole.training.data.datasets import MolDataset
from mole.training.data.utils import open_dictionary
from torch_geometric.loader import DataLoader
import mole
from pathlib import Path

vocab_path = Path(mole.__path__[0]) / "training/data/vocabularies/vocabulary_207atomenvs_radius0_ZINC_guacamole.pkl"
dictionary = open_dictionary(str(vocab_path))

smiles_list = ["CCO", "c1ccccc1", "CC(=O)Oc1ccccc1C(=O)O"]
import pandas as pd
ds = MolDataset(
    smiles=pd.Series(smiles_list),
    dictionary_inp=dictionary,
    radius_inp=0,
    useFeatures_inp=False,
    cls_token=True,
)
loader = DataLoader(ds, batch_size=len(smiles_list))
batch = next(iter(loader))

with torch.no_grad():
    input_ids, input_mask = to_dense_batch(batch.x, batch.batch, fill_value=0)
    relative_pos = to_dense_adj(batch.edge_index, batch.batch, batch.edge_attr)
    out = encoder(input_ids, input_mask, attention_mask=input_mask, relative_pos=relative_pos)

# CLS token hidden state — shape [n_molecules, 768]
mol_embeddings = out["hidden_states"][-1][:, 0, :]
print(mol_embeddings.shape)  # torch.Size([3, 768])

---

Citation

If you use this model, please cite:

Original MolE:

@article{Mendez-Lucio2022,
  title   = {A geometric deep learning approach to predict binding conformations of bioactive molecules},
  author  = {Mendez-Lucio, Oscar and others},
  journal = {Nature Machine Intelligence},
  year    = {2022}
}

DeBERTa-v3:

@article{he2021debertav3,
  title   = {DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing},
  author  = {He, Pengcheng and others},
  journal = {arXiv:2111.09543},
  year    = {2021}
}