MolE-MLM-r2

MolE-MLM-r2-S2

MolE-MLM (radius-2 target prediction) + ChEMBL supervised pretraining — the MLM baseline for RTD ablation. Step 1: BERT-style MLM predicting radius-2 atom environments (~141k classes) from radius-0 inputs, on 415M ZINC molecules. Step 2: ChEMBL supervised pretraining (641 assays).

This is the control condition: same data, same architecture, same two-stage protocol — different pre-training objective. Comparing RTD vs MLM isolates the effect of replaced token detection.

For the RTD variants (main models), see caithmac/MolE-RTD-25pct-S2.

---

Why this exists

MolE originally used MLM pre-training (predicting radius-2 atom environments at masked positions). Our RTD models replace MLM with replaced token detection. To prove RTD is the improvement (not data scale, not architecture), we need an MLM baseline trained identically.

Key difference from RTD:

• MLM: predict what was at masked positions (141k-way classification, trains on 15% of tokens)
• RTD: detect which tokens were replaced (binary classification, trains on 100% of tokens)
• RTD gets 6.7× more signal per step. GDES stabilizes training. No adversarial instability.

---

Two-stage pre-training

Step 1 details

---

Architecture

---

How to use

import torch
from collections import OrderedDict
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-MLM-r2-S2", "encoder_weights_mlm_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()

---

Citation

@misc{mole-mlm-r2-s2,
  author = {caithmac},
  title  = {MolE-MLM-r2-S2: BERT-style MLM molecular encoder (MLM baseline for RTD ablation)},
  year   = {2026},
  url    = {https://huggingface.co/caithmac/MolE-MLM-r2-S2}
}