[Paper Review] ChemBERTa-2: Towards Chemical Foundation Models

Walid Ahmad, Elana P. Simon+3 more|arXiv (Cornell University)|Sep 5, 2022

Machine Learning in Materials Science133 citations

TL;DR

ChemBERTa-2 pretrains a large SMILES-based transformer (77M compounds) and compares MLM and MTR pretraining; it achieves competitive MoleculeNet results and analyzes how pretraining scales to downstream tasks.

ABSTRACT

Large pretrained models such as GPT-3 have had tremendous impact on modern natural language processing by leveraging self-supervised learning to learn salient representations that can be used to readily finetune on a wide variety of downstream tasks. We investigate the possibility of transferring such advances to molecular machine learning by building a chemical foundation model, ChemBERTa-2, using the language of SMILES. While labeled data for molecular prediction tasks is typically scarce, libraries of SMILES strings are readily available. In this work, we build upon ChemBERTa by optimizing the pretraining process. We compare multi-task and self-supervised pretraining by varying hyperparameters and pretraining dataset size, up to 77M compounds from PubChem. To our knowledge, the 77M set constitutes one of the largest datasets used for molecular pretraining to date. We find that with these pretraining improvements, we are competitive with existing state-of-the-art architectures on the MoleculeNet benchmark suite. We analyze the degree to which improvements in pretraining translate to improvement on downstream tasks.

Motivation & Objective

• Motivate chemical foundation models by transferring large-language-model pretraining concepts to SMILES-based representations.
• Explore the impact of pretraining dataset size and objectives (MLM vs MTR) on downstream molecular prediction tasks.
• Identify how improvements in pretraining translate to performance gains across MoleculeNet benchmarks.
• Provide practical guidance on pretraining strategies and hyperparameters for large-scale molecular models.

Proposed method

• Use RoBERTa-based transformer (ChemBERTa-2) with MLM pretraining on 77M SMILES from PubChem.
• Compare MLM and multi-task regression (MTR) pretraining objectives.
• Perform hyperparameter search over architecture and training settings across 5M, 10M, and 77M datasets.
• Fine-tune pretrained models on MoleculeNet tasks with scaffold-based splits and per-task normalization or class weighting as appropriate.
• Analyze transfer between pretraining losses and downstream performance.
• Open-source trained models for reproducibility.

Experimental results

Research questions

• RQ1 Does larger-scale SMILES pretraining improve molecular property prediction performance on MoleculeNet tasks?
• RQ2 How do MLM and multi-task regression pretraining compare in terms of downstream performance and training efficiency?
• RQ3 To what extent does pretraining loss translate into improvements on specific downstream datasets (e.g., solubility, toxicity)?

Key findings

• Pretraining with larger datasets (up to 77M SMILES) yields lower pretraining loss and improved downstream performance across several tasks.
• MTR pretraining generally outperforms MLM on downstream tasks, but MLM-based architecture searches can efficiently guide MTR pretraining choices due to correlated losses.
• ChemBERTa-2 configurations pretrained with MLM or MTR on 77M data achieve competitive results and outperform D-MPNN on 6 of 8 tasks in Table 1.
• Improvements in pretraining loss correlate with gains in some tasks (e.g., Lipophilicity) but transfer efficiency varies by dataset and task.
• Using MLM to select hyperparameters before committing to slower MTR training can be a practical workflow.
• Dimensionality reduction analyses (UMAP) show that ChemBERTa embeddings can serve as strong priors for downstream classification tasks.