ChemBERTa-2: Towards Chemical Foundation Models

TL;DR

ChemBERTa-2 investigates scaling SMILES-based pretraining with MLM and multi-task regression on a 77 million SMILES corpus to learn chemical fingerprints. By comparing MLM and MTR, it shows that MTR generally yields stronger downstream performance while MLM serves as a faster proxy for architecture search, with pretraining loss scaling driving improvements across tasks. The work achieves competitive results on MoleculeNet benchmarks, provides embedding analyses via UMAP, and highlights practical considerations for large-scale pretraining and open-source release in the context of chemical foundation models.

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.

Figures (5)

• Figure 1: a) An illustration of masked language modeling (MLM) and multitask regression (MTR) pretraining tasks. b) The training pipeline implemented to achieve results in this paper.
• Figure 2: Comparing MLM and MTR pretrain losses
• Figure 3: Pretrain losses for each of the 5 model configurations that were trained on all three datasets (5M, 10M, and 77M). MLM configurations are on the left, and MTR on the right. The configurations are sorted by their loss when training with 5M compounds along the x-axis. Note that there is considerable performance variability across runs.
• Figure 4: Finetuning performance versus pretraining loss. Left Column: MLM Pretraining, Right Column: MTR Pretraining. Top Row: Lipophilicity Finetune, RMSE ($\downarrow$), Bottom Row: BACE Classification Finetune, ROC-AUC ($\uparrow$). The dotted lines represent linear models fit to the datapoints.
• Figure 5: (a) MTR-77M embeddings fit using UMAP on BACE classification task. (b) MLM-77M embeddings fit using UMAP on BACE classification task. (c) ECFP embeddings fit using UMAP on BACE classification task. (d) MTR-77M embeddings fit using UMAP on BBBP classification task. (e) MLM-77M embeddings fit using UMAP on BBBP classification task (f) ECFP embeddings fit using UMAP on BBBP classification task