ChemBERTa-ZINC100M-SmilesTok

A RoBERTa-style chemical language model pre-trained with masked-language modeling (MLM) on ~100M SMILES from ZINC20, using an atom-level SMILES tokenizer (Schwaller regex) rather than byte-pair encoding. It is a faithful reproduction of the ChemBERTa-100M-MLM recipe whose single deliberate change is the tokenizer. This card documents the full training setup, dynamics, and the observations and decisions made during pre-training.

Quick start

from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
import torch

repo = "andrewboldi/chemberta-zinc100m-smilestok"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModel.from_pretrained(repo)               # encoder, for embeddings
# or AutoModelForMaskedLM.from_pretrained(repo)       # for fill-mask

inputs = tok("CC(=O)Oc1ccccc1C(=O)O", return_tensors="pt")  # aspirin
with torch.no_grad():
    h = model(**inputs).last_hidden_state[:, 0]       # [CLS] representation

The tokenizer ships as a portable tokenizer.json (no custom code, no deepchem dependency).

Architecture

Tokenization

Atom-level tokenization with the Schwaller SMILES regex (the scheme used by DeepChem's SmilesTokenizer): multi-character atoms and bracketed groups such as Cl, Br, [N+], [O-], [nH] are single tokens. The 581-token vocabulary covers drug/lead-like chemistry with ≈0% unknown tokens on ZINC20 and on the MoleculeNet evaluation sets. The released tokenizer.json is verified to reproduce the training tokenization exactly on 15k+ held-out SMILES. This is the one intentional deviation from the original ChemBERTa-100M-MLM, which used a byte-pair-encoding tokenizer.

Pre-training data

~100M SMILES from ZINC20 (xchen16/zinc20). Because ZINC20 is ordered by molecule size, the corpus was assembled by sampling shards at a fixed stride across the full size range (every 10th tranche) so the training distribution spans small-to-large molecules rather than a size-biased prefix. SMILES are canonicalized with RDKit and packed (int16 token ids) for training.

Pre-training procedure

Optimization & stability

Higher learning rates were unstable at this scale: 5e-4 and 2e-4 both trained well early but diverged late in the schedule. Training was stabilized with lr 1e-4, β₂ = 0.95, and gradient clipping at 0.5, which trained cleanly for the full epoch with no divergence.

Validation protocol

Validation MLM loss is measured on a representative random hold-out of 5,000 SMILES sampled across the whole corpus. (An early contiguous "first-N" split was rejected because, on the size-ordered packed data, it is a biased small-molecule slice and gives misleadingly high, erratic loss; a random hold-out is the faithful measure.) Held-out MLM loss decreased monotonically through training and continued to improve into the second half of the epoch (the language model did not overfit).

Training dynamics & scaling observation

Downstream task performance saturates early. Intermediate checkpoints (20%, 45%, and 55% of one epoch) were fine-tuned on the full MoleculeNet classification suite (6 tasks, Bemis–Murcko scaffold split). The 6-task mean ROC-AUC was essentially flat:

The change from 20%→55% is +0.004 — within run-to-run noise (individual tasks moved ±0.05 in both directions). Over the same range, held-out MLM loss kept falling: the language-modeling objective continued to improve while the downstream property-prediction signal had already plateaued. In other words, most of the transferable representation is learned in the first ~20% of one epoch on this corpus.

Decision: despite this early saturation, pre-training was nonetheless carried to a full epoch — to remove any ambiguity about convergence and to provide a clean, fully-trained single-epoch checkpoint for reuse and reproduction.

Evaluation

Evaluated on the MoleculeNet benchmark following the ChemBERTa protocol: 80/10/10 Bemis–Murcko scaffold split, fine-tune up to 25 epochs with early stopping on validation ROC-AUC. Classification tasks report ROC-AUC (and PRC-AUC for the class-imbalanced sets); multi-task datasets average over tasks. Numbers below are for the final (end-of-epoch) checkpoint.

Comparison to a byte-pair-encoding baseline

To situate the atom-level tokenizer, this model is compared against DeepChem/ChemBERTa-77M-MLM — the canonical BPE-tokenized ChemBERTa — under the identical downstream protocol (same Bemis–Murcko scaffold splits, same 25-epoch fine-tuning and hyperparameters), each model using its own tokenizer. Test ROC-AUC:

Under matched downstream evaluation the atom-level model is competitive with, and on average modestly ahead of, the BPE baseline (+0.026 mean ROC-AUC; higher on 4 of 6 datasets). This is not a controlled tokenizer ablation: the two checkpoints also differ in pre-training corpus (ZINC20 vs. PubChem) and corpus size (~100M vs. 77M molecules), so the difference reflects the combined effect of tokenizer, data, and scale — not tokenization in isolation. It is reported to show that the atom-level tokenizer is not a downgrade under standard evaluation, not to claim tokenizer superiority on its own.

Intended uses & limitations

A general-purpose chemical language model: molecular embeddings, property-prediction fine-tuning, and fill-mask over SMILES.

Distribution & limitations. The model is trained on ZINC20, a database of purchasable, drug/lead-like compounds (largely make-on-demand synthetic libraries). It is therefore well-matched to medicinal-chemistry property prediction but out-of-distribution on chemistry under-represented in ZINC20 — notably natural products and other high-Fsp³, stereochemically dense, macrocyclic or polycyclic scaffolds. Training used sequences of ≤126 SMILES tokens; much longer SMILES are out of the trained range. Not intended for clinical or safety-critical decisions.

Training compute & reproducibility

One epoch over ~100M molecules (1,558,643 optimizer steps) on a single RTX 4060 Laptop (8 GB), bf16, batch size 64. Training, tokenizer construction, and evaluation are fully scripted and reproducible.

License & citations

Released under the Apache License 2.0.

• ZINC20 — Irwin et al., J. Chem. Inf. Model. 60(12), 2020.
• Atom-level SMILES tokenization — Schwaller et al., ACS Cent. Sci. 5(9), 2019; DeepChem SmilesTokenizer.
• ChemBERTa — Chithrananda, Grand, Ramsundar, 2020.
• MoleculeNet — Wu et al., Chem. Sci. 9, 2018.