Source-linked AI summary
ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction
Seyone Chithrananda, Gabriel Grand, Bharath Ramsundar
TL;DR
Molecular property prediction is constrained by scarce labeled data, while the use of large unlabeled molecular structure collections remains underexplored. This paper introduces ChemBERTa, evaluates transformer pretraining and representation choices on MoleculeNet, and finds that performance scales with more pretraining data while remaining generally below state-of-the-art baselines. It also releases a curated 77M-SMILES PubChem dataset for pretraining.
Problem
Scarce labeled molecular data and underexplored use of abundant unlabeled structures motivate studying transformers for molecular representation learning.
Method
ChemBERTa applies masked-language-model pretraining to curated PubChem SMILES, varying pretraining scale, tokenization, and representation before MoleculeNet fine-tuning.
Results
ChemBERTa approaches but generally does not beat Chemprop baselines, while scaling from 100K to 10M pretraining examples yields average ΔROC-AUC = +0.110 and ΔPRC-AUC = +0.059.
Takeaways & Limitations
Transformer pretraining is a promising avenue for molecular property prediction, and the released 77M-SMILES dataset supports further large-scale self-supervised study.
Takeaways & Limitations
Large-scale pretraining is resource-intensive and generated an estimated 17.1 kg CO2eq of emissions.
Abstract
from arXiv · showhide
GNNs and chemical fingerprints are the predominant approaches to representing molecules for property prediction. However, in NLP, transformers have become the de-facto standard for representation learning thanks to their strong downstream task transfer. In parallel, the software ecosystem around transformers is maturing rapidly, with libraries like HuggingFace and BertViz enabling streamlined training and introspection. In this work, we make one of the first attempts to systematically evaluate transformers on molecular property prediction tasks via our ChemBERTa model. ChemBERTa scales well with pretraining dataset size, offering competitive downstream performance on MoleculeNet and useful attention-based visualization modalities. Our results suggest that transformers offer a promising avenue of future work for molecular representation learning and property prediction. To facilitate these efforts, we release a curated dataset of 77M SMILES from PubChem suitable for large-scale self-supervised pretraining.
1 Motivation
Molecular property prediction benefits from GNNs, but limited labeled data makes learning from abundant unlabeled molecular structures an important open challenge. Transformers provide a potential alternative representation-learning approach, motivating systematic study of their use with SMILES.
- Limited labeled data remains a fundamental challenge because each new molecular label requires costly, time-consuming laboratory testing.
- Transformers have become robust self-supervised representation learners in NLP, supported by strong downstream transfer and maturing software infrastructure.
- Millions of available SMILES strings make transformers an alternative to expert-crafted and GNN-learned molecular fingerprints.
- Key open questions concern how pretraining size, SMILES tokenization, and string representation affect downstream molecular prediction.
- The work addresses these questions through one of the first systematic evaluations of transformers on molecular property prediction tasks.
2 Related Work
Earlier cheminformatics language models learned molecular representations from SMILES, while transformer applications were comparatively limited and recent. Existing studies left the effects of pretraining scale, tokenization, and string representation insufficiently explored.
- Cheminformatics has long used SMILES language models, typically RNN sequence-to-sequence systems supporting auxiliary tasks such as focused library generation.
- Transformer work in chemistry had largely focused on reaction prediction before broader molecular property-prediction studies emerged.
- Prior molecular-property transformer studies used relatively small pretraining datasets of 861K ChEMBL compounds and 2M ZINC compounds.
- Studies using 18.7M ZINC compounds did not explore pretraining dataset size, tokenizer, or string representation effects.
- This work parallels systematic GNN pretraining research by investigating transformer self-supervision for molecular property prediction.
3 Methods
ChemBERTa combines transformer pretraining on curated PubChem SMILES with supervised fine-tuning and evaluation on MoleculeNet tasks. The methods vary pretraining scale, tokenization, and molecular representation while comparing against established baselines.
- ChemBERTa uses a RoBERTa implementation with 12 attention heads and 6 layers, yielding 72 distinct attention mechanisms.
- 3.2 Finetuning on MoleculeNet: Baselines include Chemprop’s directed message-passing neural network, Random Forest, and Support Vector Machine using 2048-bit Morgan fingerprints.
- 3.1 PreTraining on PubChem 77M: Masked-language pretraining masks 15% of input tokens, uses a 52K-token vocabulary and 512-token maximum sequence length, and trains for 10 epochs except on the 10M subset.
- 3.1 PreTraining on PubChem 77M: The pretraining corpus contains 77M unique, canonicalized, globally shuffled PubChem SMILES divided into 100K, 250K, 1M, and 10M subsets.
- 3.2 Finetuning on MoleculeNet: Models are evaluated on BBBP, ClinTox, HIV, and Tox21 classification tasks spanning 1.5K–41.1K examples and medicinal chemistry applications.
- 4.1 Tokenizers: The study compares HuggingFace BPE tokenization with a custom regex-based SmilesTokenizer and examines SMILES and SELFIES representations.
- Table 1 reports both ROC-AUC and PRC-AUC for ChemBERTa pretrained on 10M PubChem compounds and Chemprop baselines.
- 3.2 Finetuning on MoleculeNet: Fine-tuning appends a linear classification layer, backpropagates through the base model, and uses up to 25 epochs with early stopping on ROC-AUC.
4 Results
ChemBERTa approaches strong Chemprop baselines while benefiting from larger pretraining datasets. Results also compare tokenization and molecular string representations, and use attention visualization to inspect chemically relevant model behavior.
- Scaling pretraining data: ∆ROC-AUC = +0.110 and ∆PRC-AUC = +0.059 when scaling pretraining from 100K to 10M.These averages cover the evaluated BBBP, ClinTox, and Tox21 analyses; HIV was omitted from this scaling analysis.
- Downstream performance: ChemBERTa approaches, but does not beat, Chemprop baselines on the evaluated MoleculeNet tasks.The reported comparison identifies Chemprop as the strong baseline.
- Tokenizers: SmilesTokenizer narrowly outperformed BPE on Tox21 SR-p53 by ∆PRC-AUC = +0.015.The two identical models were pretrained on PubChem-1M, and the authors call for further benchmarking on additional datasets.
- SMILES vs. SELFIES: SELFIES produced no significant difference in downstream performance from SMILES on Tox21 SR-p53.The authors state that further benchmarking is needed to validate this finding.
- Attention visualization: BertViz inspection identified neurons selective for chemically relevant functional groups and aromatic rings, alongside neurons tracking bracket closures.The analysis contrasted ChemBERTa attention with molecular graph visualization from an attention-based GNN.
5 Discussion
ChemBERTa gains predictive power from MLM pretraining and scales with more pretraining data, but remains below state-of-the-art on evaluated tasks. The authors identify broader evaluation, tuning, scaling, and sample-efficiency work as future priorities.
- MLM pretraining boosts predictive power on selected MoleculeNet tasks, although ChemBERTa generally remains below state-of-the-art.The possible exception is Tox21.
- ChemBERTa’s current analysis covers only a small portion of the hypothesis space under consideration.
- Future evaluations include all MoleculeNet tasks, more systematic hyperparameter tuning, larger masking rates, and multitask finetuning.
- Future pretraining will scale from the current subsets toward the full PubChem 77M dataset and larger collections such as ZINC-15.The authors note that this expansion will require substantially improved engineering infrastructure.
- Hybrid graph-transformer models may improve sample efficiency while retaining the scalability of attention-based architectures.
Broader Impact
The paper connects molecular property prediction to drug-discovery goals while acknowledging that large-scale pretraining has environmental costs. It reports estimated emissions and recommends sustainable cloud-computing practices.
- Improving molecular property prediction accuracy and generalizability contributes to accelerating development of potentially life-saving medicines.
- Large-scale pretraining is resource-intensive and has a growing carbon footprint.
- 17.1 kg CO2eq of emissions were estimated for the paper’s pretraining.The authors state that Google Cloud Platform offset 100% of these emissions through its carbon-neutral certification.
- The authors encourage researchers to consider environmental impacts and use sustainable cloud-compute solutions where possible.