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

Stage	Objective	Data	Steps	Batch
Step 1	MLM (r0→r2 prediction)	ZINC-Curated 415M molecules	1,000,000	512
Step 2	Sparse multi-task BCE	ChEMBL 36: 512k mols × 641 assays	80,000	512

Step 1 details

Setting	Value
Input radius	0 (~211 tokens)
Target radius	2 (~141k unique atom environments, min_count=100)
MLM head	Linear(768→141k) — 108M params
Mask probability	15%
Hardware	4 × NVIDIA H200 (141 GB each)
Wall time	~3 days

---

Architecture

Component	Config
Layers	12
Hidden size	768
Intermediate size	3072
Attention heads	12 × 64
Input vocabulary	211 atom environments (radius-0 Morgan)
Target vocabulary	~141k atom environments (radius-2 Morgan, min_count=100)
Relative attention	Yes (p2c + c2p)

---

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