ChemBERTa-2: Towards Chemical Foundation Models

In plain words

Chemists write a molecule as a short text string of atoms and bonds (a SMILES string). ChemBERTa-2 treats those strings as sentences in a chemical language and trains a BERT-style language model on very many of them, without labels. The idea is the same as for ordinary text: learn general patterns from a huge unlabelled collection first, then reuse the model on small labelled datasets about a molecule's properties.

The authors compare two ways of pretraining. One hides parts of the string and asks the model to fill them in. The other asks the model to predict a set of about 200 computed chemical properties. They also check how much a bigger pretraining set and a lower pretraining loss help the downstream tasks. Their main finding is that the property-prediction pretraining usually gives better fine-tuned models, and that the results are competitive with a graph-based baseline on several standard benchmarks. The paper reports these results itself, and the code is public but partial.

TL;DR

• A RoBERTa-style transformer, pretrained on SMILES strings from up to 77M PubChem compounds.
• Two objectives: masked language modelling (MLM) and multi-task regression (MTR) on 200 RDKit descriptors.
• Author-reported: MTR-pretrained models tend to fine-tune better; ChemBERTa-2 beats D-MPNN on 6 of 8 MoleculeNet tasks in Table 1.
• Reuse only with care: the architecture search results are not tabulated, the 77M dataset is not in the repository, and several code defaults differ from the paper.

Problem

Labelled molecular property data is scarce, while SMILES strings are easy to collect §1. Earlier SMILES transformers were pretrained on 1.6M to 18.7M compounds (MolBERT, ChemBERTa-1 at 10M, SMILES-BERT at 18.7M) §1. The paper asks whether scaling the pretraining data, and choosing the pretraining task, improves fine-tuning on MoleculeNet tasks §1 §4.

Method

ChemBERTa-2 is a RoBERTa encoder from the HuggingFace implementation §3. The pretraining corpus is the 77M unique SMILES set from ChemBERTa-1; the strings are canonicalized and globally shuffled §3. A fixed set of 100k compounds is held out for validation, and training subsets of 5M, 10M and 77M are drawn from the rest §3. Each subset is used for both objectives §3.

1. Masked language modelling. Follow the RoBERTa procedure: mask 15% of the tokens of each input and train the model to recover them. The vocabulary has 591 tokens and the maximum sequence length is 512 §3.1.
2. Multi-task regression. For each compound, compute 200 molecular properties from the SMILES with RDKit. None needs experimental data §3.1. Train the model to predict all 200 at once, with each label mean-normalized per task §3.1. The regression head reads the first (start) token roberta_regression.py. The objectives are sketched in fig. 1.
3. Pretraining setup. Models train on AWS EC2 instances with Nvidia T4 GPUs. Early stopping patience equals one full pass over the dataset, so that every compound is seen at least once §3.1.
4. Architecture and hyperparameter search. The search draws 50 random configurations, varying hidden size, number of attention heads, dropout, intermediate size, number of layers and learning rate. Models have between 5M and 46M parameters §3.2. Each configuration is trained on the 5M subset for both MLM and MTR §3.2. Five configurations per objective, chosen independently and spanning a range of validation losses, are then trained on the 10M and 77M subsets §3.2.
5. Fine-tuning on MoleculeNet. The datasets have 1.5K to 8.0K examples. Each gets a scaffold split of 80/10/10 from the DeepChem splitter §3.3. Models fine-tune for up to 100 epochs with early stopping on validation loss. The learning rate, seed and batch size are tuned with the Optuna backend of the HuggingFace trainer §3.3. Regression labels are normalized to zero mean and unit standard deviation; classification uses balanced class weights §3.3.
6. Analysis. Compare the objectives and dataset sizes by pretraining loss and downstream score §4.1 fig. 2 fig. 3 fig. 4. Visualize BACE and BBBP embeddings with UMAP and contrast them with ECFP fingerprints §5 fig. 5.

Pretraining in the code uses the HuggingFace Trainer. Masked language modelling uses the HuggingFace collator with masking rate 0.15; regression uses a dense layer with ReLU, then a linear output over all labels, with an MSE loss on normalized labels utils.py roberta_regression.py roberta_regression.py.

# Pretrain one ChemBERTa-2 model  [§3] [§3.1]
# Input:  smiles (canonical, shuffled, list of str)
# Input:  objective in {mlm, mtr}
# Output: encoder with a masked-LM head or a regression head

def pretrain(smiles, objective):
    # Tokenize each string to at most 512 tokens  [§3.1]
    tokens = [tokenize(s, max_length=512) for s in smiles]
    if objective == "mlm":
        # Mask 15% of the tokens; predict the originals  [§3.1]
        inputs, targets = mask_tokens(tokens, rate=0.15)
    else:
        # 200 RDKit descriptors computed from each SMILES  [§3.1]
        raw = [rdkit_descriptors(s) for s in smiles]
        # Mean-normalize each property  [§3.1]
        # (the code also divides by the std: see Paper vs code)
        targets = raw - mean_per_property(raw)
        inputs = tokens
    model = encoder_with_head(objective)
    # Stop after the validation loss fails to improve for one
    # full pass over the data  [§3.1]
    return train_until_early_stop(model, inputs, targets)

# Fine-tune on one MoleculeNet task  [§3.3]
# Input:  dataset with a scaffold split 80/10/10
# Input:  task in {regression, classification}
# Output: model with the lowest validation loss

def finetune(dataset, task, n_trials):
    # Regression: z-score labels with training-set mean and std  [§3.3]
    # Classification: balanced class weights  [§3.3]
    #   (no class weighting found in the code)
    best = None
    for trial in range(n_trials):
        # Learning rate: log-uniform from 1e-6 to 1e-4  [§3.3]
        # [code: chemberta/finetune/finetune.py#L277]
        learning_rate = sample_log_uniform(1e-6, 1e-4)
        # Seed: integer from 1 to 40 in the code  [§3.3]
        # [code: chemberta/finetune/finetune.py#L281]
        seed = sample_int(1, 40)
        # Epochs: paper allows up to 100; code caps the search at 10  [§3.3]
        # [code: chemberta/finetune/finetune.py#L92]
        epochs = sample_int(1, 10)
        # Batch size: paper varies it; code fixes it to 64  [§3.3]
        # [code: chemberta/finetune/finetune.py#L282-L284]
        model = train(dataset, learning_rate, seed, epochs, batch_size=64)
        if best is None or model.valid_loss < best.valid_loss:
            best = model
    return best

Hyperparameters

Values in the first block come from the paper; the second block lists values the code sets when the paper is silent or differs. "Not stated" means neither the paper nor the code gives it.

Results

Author-reported. Setup: eight MoleculeNet tasks §3.3 with a scaffold split 80/10/10, and ChemBERTa-2 variants pretrained at 5M, 10M and 77M. Regression columns are RMSE (lower is better) and classification columns are ROC-AUC (higher is better) tab. 1. The paper does not state the units of the regression columns, nor whether each cell is a single run, a mean over seeds or the best seed. D-MPNN is from the chemprop library tab. 1. The RF and GCN baselines are not described beyond the table. Grover was not benchmarked because its benchmark procedure differs tab. 1. Hardware: Nvidia T4 GPUs on AWS EC2 §3.1.

Bold marks the best value per column, as in the paper's table. The ChemBERTa-1 source and setup are not stated.

Headline claims §4:

• The authors describe ChemBERTa-2 as competitive on almost every task and beating D-MPNN on 6 of 8 tasks. Counting from Table 1, that holds for BACE (regression), Clearance, Delaney, Lipo, BBBP and SR-p53. It does not hold for BACE (classification) or ClinTox §4 tab. 1.
• MTR-pretrained models tend to beat MLM-pretrained ones on downstream tasks §4.1.
• Pretraining loss for MLM tracks MTR loss for the same configuration at 5M, so MLM can be used for the architecture search §4.1 fig. 2.
• Training to convergence on 77M rather than 5M lowers pretraining loss by 25 to 35%, consistently across configurations §4.1 fig. 3.
• Lower pretraining loss transfers almost linearly to Lipophilicity RMSE but not clearly to BACE classification §4.1 fig. 4.
• From UMAP plots of BACE and BBBP, the authors conclude that ChemBERTa embeddings from pretrained MLM and MTR models are, on average, a stronger prior representation for a variety of downstream tasks §5 fig. 5. No clustering metric is reported.

Paper vs code

The official code is the MIT-licensed bert-loves-chemistry repository at commit 695bc28, owned by co-author Seyone Chithrananda. The training code is in chemberta/train, chemberta/finetune and chemberta/utils. The README describes a ChemBERTa notebook collection, links the ChemBERTa-2 arXiv paper, and says the DeepChem implementation was still in progress. The repository does not contain the 77M dataset, the canonicalization step or the paper's exact scripts for the reported runs, so the paper's numbers cannot be reproduced from this commit alone.

• MTR label scaling. Paper: mean-normalize each property §3.1. Code: divides by the standard deviation as well, a z-score roberta_regression.py compute_norms.py.
• Pretraining early stopping. Paper: patience equals one pass over the dataset §3.1. Code: patience 3 evaluations, with an evaluation every 50 steps flags.py train_roberta.py train_roberta.py. These two rules agree only for a particular dataset size and batch size.
• Pretraining validation split. Paper: a fixed set of 100k compounds §3. Code default: random split with frac_train 0.95 unless an eval file is given flags.py utils.py. The paper's own split file is not in the repository.
• Vocabulary. Paper: 591 tokens §3.1. Code default: 600 flags.py, and the default tokenizer is a HuggingFace hub model, seyonec/SMILES_tokenized_PubChem_shard00_160k flags.py.
• Regression head. Paper: not described §3.1. Code: dense layer, ReLU, dropout, linear output on the first token roberta_regression.py.
• Model size. Paper: hidden size and FFN size are searched, 5M to 46M parameters §3.2. Code: hidden size = 64 × heads and FFN = 4 × the flag, so the default flag gives 12288 train_roberta.py flags.py. The help text calls the flag the intermediate size, so the 4× factor is easy to miss.
• Fine-tuning epochs. Paper: up to 100 epochs §3.3. Code: the search caps epochs at 10 by default finetune.py finetune.py.
• Fine-tuning batch size. Paper: varied in the search §3.3. Code: fixed to the flag default 64 finetune.py finetune.py.
• Fine-tuning class weights. Paper: balanced class weights §3.3. Code: no class weighting found in chemberta/finetune or chemberta/utils.
• Fine-tuning normalization. Matches the paper: z-score from training-set labels finetune.py §3.3.
• Scaffold split. Matches the paper: DeepChem scaffold split for the MoleculeNet tasks molnet_dataloader.py §3.3. The 80/10/10 ratio is not visible in the repository.
• MLM details. Masking rate 0.15 matches §3.1 flags.py. The replacement rule (for example the usual 80/10/10 mix of mask, random and unchanged tokens) is set by the HuggingFace collator and not stated in the paper.
• Canonicalization and shuffling. Paper: canonicalized and globally shuffled §3. Code: no such step found in the Python sources.
• Figures and UMAP. The UMAP settings match the paper in the notebook cells that use 25 neighbours ChemBERTA_dimensionaliy_reduction_BBBP.ipynb. Other cells in the same notebook use 100 neighbours, and the paper does not say which setting produced Figure 5 fig. 5.

The 200 properties are not listed in the repository. The script takes every RDKit descriptor except Ipc, so the count depends on the installed RDKit version compute_norms.py.

Pitfalls

• The chosen configurations are not listed. The paper says five configurations per objective were kept for the 10M and 77M runs, but does not name them or give their sizes §3.2. Only the 5M results are plotted fig. 3. An implementer cannot reproduce the reported 77M models from the text.
• Dataset is not reproducible here. The 77M set is ChemBERTa-1's corpus; the repository has only sample files, such as 1k PubChem and 100k ZINC strings . Canonicalization and the 100k validation set are not in the repository.
• Loss scale depends on the label normalization. Mean-only and z-score targets give different MSE scales, so pretraining losses from different code paths are not comparable §3.1 roberta_regression.py.
• "One pass" early stopping is not a fixed number of steps. Its meaning depends on dataset and batch size; the code uses an evaluation count instead §3.1 flags.py.
• Regression units are missing. Table 1 gives RMSE values without units, and the Clearance column is around 50 while the others are below 3. Do not compare RMSE across columns tab. 1.
• Single numbers without spread. Table 1 gives one value per cell and does not say whether it is a single run, a mean or the best seed. The Figure 3 caption notes large run-to-run variability in pretraining loss fig. 3 §4.1.
• Some classification scores are below chance. MLM ClinTox ROC-AUC values of 0.239 to 0.349 appear in Table 1 without comment tab. 1. Check the metric before using these rows.
• HIV and Tox21 are named but absent. The fine-tuning text lists them §3.3, yet Table 1 has eight columns with no HIV or Tox21 results.
• Duplicate names. BACE appears as both a regression and a classification task, and the reader must tell them apart by the metric in the header tab. 1.
• Cost hidden in the abstract. Pretraining on 77M compounds with a 5M-to-46M model search is a substantial compute cost. The paper gives no GPU-hour total; it notes that spot instances were used and that interrupted runs were restarted from checkpoints §7.1.
• Release and dual use. The authors released trained models and state their view on dual-use risk in the Discussion §6. Check that view before publishing derived models.
• Visual evidence for embeddings. The claim that ChemBERTa embeddings are a stronger prior rests on UMAP plots, with no metric §5 fig. 5.

Lineage

• Builds on: ChemBERTa-1, which supplies the corpus and is the starting point §1 §3. The code README gives its arXiv id as 2010.09885.
• Builds on: RoBERTa, used as the encoder through the HuggingFace implementation §3.
• Compared against: D-MPNN via chemprop, Random Forest and GCN tab. 1; Grover was not compared because of benchmark differences tab. 1. Related graph pretraining methods are discussed in §2.
• Superseded by: not stated in this card's sources.