Source-linked AI summary
ChemBERTa-2: Towards Chemical Foundation Models
Walid Ahmad, Elana Simon, Seyone Chithrananda, Gabriel Grand, Bharath Ramsundar
TL;DR
ChemBERTa-2 investigates whether large-scale SMILES pretraining can transfer foundation-model advances to molecular machine learning, where labeled prediction data is scarce. By optimizing pretraining across datasets, hyperparameters, and MLM/MTR objectives, the study achieves competitive MoleculeNet performance while showing that transfer varies across downstream datasets.
Problem
Labeled data for molecular prediction is typically scarce, motivating investigation of whether readily available SMILES libraries can support transferable chemical foundation models.
Method
ChemBERTa-2 optimizes molecular pretraining by comparing MLM and MTR objectives across hyperparameter configurations and datasets of 5M, 10M, and 77M SMILES.
Results
ChemBERTa-2 achieves competitive results on nearly all MoleculeNet tasks and outperforms D-MPNN on 6 out of 8 tasks.
Takeaways & Limitations
MLM loss can guide MTR configuration selection, while improvements in pretraining transfer unevenly across downstream datasets.
Takeaways & Limitations
The study does not investigate why some downstream datasets benefit more from pretraining improvements than others.
Abstract
from arXiv · showhide
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.
1 Motivation
Transformer-based molecular models use SMILES to learn pretrained molecular representations. ChemBERTa-2 extends this approach with MLM and MTR pretraining over 77 million SMILES strings.
- ChemBERTa introduced a BERT-like transformer pretrained on 10M compounds to learn molecular fingerprints.
- ChemBERTa-2 uses MLM and MTR over 77 million SMILES strings to explore scaling molecular pretraining.
- SMILES represents molecules as character sequences with atom and bond symbols and few grammar rules.
- Figure 1 depicts the MLM and MTR pretraining tasks and the training pipeline used in the paper.
2 Related Work
Molecular pretraining also includes graph-based, contrastive, dual-representation, and translation-based approaches alongside transformer models.
- SNAP uses graph pretraining based on node-attribute masking and structural similarity.
- Grover scales graph-transformer pretraining to 100 million parameters on 10M compounds, while MolGNet pretrains a 53M-parameter message-passing model on 11M compounds.
- Other methods use contrastive learning, combine graph and transformer pretraining, or translate between SMILES and InChi representations.
3 Methods
ChemBERTa-2 is trained on shuffled subsets of a 77M-SMILES corpus using MLM or MTR, with hyperparameter selection performed before larger-scale pretraining and downstream evaluation.
- The corpus contains 77M unique SMILES, globally shuffled after canonicalization, with a fixed 100k-compound validation set and 5M, 10M, and 77M training subsets.
- MLM masks 15% of input tokens and trains the model to identify them.
- MTR predicts 200 RDKit-computed molecular properties simultaneously after mean-normalizing each task’s labels.
- 3.2 Hyperparameter Search: The study searches 50 configurations spanning architecture and optimization hyperparameters, with models ranging from 5M to 46M parameters.
- 3.2 Hyperparameter Search: Five MLM and five MTR configurations selected on 5M compounds are subsequently trained on the 10M and 77M datasets.
- Evaluation covers MoleculeNet regression and classification tasks spanning 1.5K–8.0K examples and medicinal chemistry applications.
4 Results
ChemBERTa-2 achieves competitive MoleculeNet performance, with MTR generally transferring better than MLM and larger pretraining datasets lowering pretraining loss. Transfer from pretraining to downstream tasks varies by dataset.
- ChemBERTa-2 outperforms D-MPNN on 6 out of 8 MoleculeNet tasks and is competitive on nearly all evaluated tasks.
- Table 1 compares MLM and MTR models trained on 5M, 10M, and 77M compounds against existing architectures using ROC-AUC and RMSE.
- 4.1 Selection of Pretraining Method: MTR-pretrained models tend to outperform MLM-pretrained models on every downstream finetuning task.
- 4.1 Selection of Pretraining Method: MLM loss corresponds well with MTR loss for a given architecture, supporting MLM-first architecture search before MTR training.
- 4.1 Selection of Pretraining Method: 25-35%: training to convergence on 77M rather than 5M unique SMILES improves pretraining loss across MLM and MTR configurations.
- 4.1 Selection of Pretraining Method: Pretraining improvements transfer differently across datasets: Lipophilicity RMSE improves almost linearly, whereas the pattern is less clear for BACE Classification.
5 Dimension Reduction of ChemBERTa Embeddings
The study uses UMAP to inspect pretrained ChemBERTa representations on BACE and BBBP, comparing them with ECFP embeddings and examining their suitability as prior representations for downstream tasks.
- Representation analysis: UMAP was used to inspect pretrained ChemBERTa representations on BACE and BBBP classification tasks, alongside ECFP embeddings.The analysis asks whether pretrained representations cluster molecules according to task labels without further MoleculeNet fine-tuning.
- Transfer analysis: Figure 4 compares finetuning performance with pretraining loss for MLM and MTR, using Lipophilicity RMSE and BACE ROC-AUC.Dotted lines are linear models fitted to the datapoints.
- Representation analysis: Large extra SMILES fragments were removed with RDKit’s LargeFragmentChooser before generating transformer and ECFP embeddings.This preprocessing avoided salts that are irrelevant to blood-brain barrier permeability.
- Findings: ChemBERTa embeddings from both masked-language and multi-task regression pretraining were found to be stronger prior representations on average for downstream fine-tuning.The UMAP model used Jaccard distance with 25 neighbors, two components, and min_dist = 0.001.
6 Discussion
The discussion presents ChemBERTa-2 as an updated transformer for molecular property prediction, while noting uneven transfer across tasks and unresolved concerns about scope and terminology.
- Discussion: ChemBERTa-2 introduces an updated transformer architecture for molecular property prediction and achieves more competitive downstream baseline results.The authors also use the pretraining pipeline to study strategies for language-model pretraining.
- Future work: Future work will compare against Grover and other graph-based architectures and extend pretraining to larger datasets.
- Transferability: Transfer from improved pretraining varies by dataset: some finetuning tasks benefit greatly, whereas others do not.Possible influences include task type, molecular structural features, dataset size, and other dataset characteristics not explored here.
- Responsible release: The models are open sourced despite dual-use risks, with the authors currently viewing novel-molecule synthesis as limiting potential harms.They state that this risk balance may shift as the research continues.
- Terminology: The authors acknowledge criticism of the term foundation model because training data for large language models is often heavily biased.They argue that chemical foundation model is appropriate because their pretraining data derives from fundamental chemical calculations.
7 Appendix
The appendix documents practical techniques for running large-scale pretraining efficiently across heterogeneous and interruption-prone computing environments.
- Training efficiency: Training on smaller machines required reducing batch size and correspondingly decreasing the learning rate.This adjustment was used to mitigate the effects of smaller batches.
- Data loading: A custom wrapper around the text-based HuggingFace loader improved efficiency compared with the default CSV loader.The wrapper parses text into tabular format for large-scale pretraining.
- Fault tolerance: AWS spot-instance interruptions made model-restart support critical for long-running training.HuggingFace’s restart system allowed training to resume part-way through after interruptions.