Source-linked AI summary
Molecular representation learning with language models and domain-relevant auxiliary tasks
Benedek Fabian, Thomas Edlich, Héléna Gaspar, Marwin Segler, Joshua Meyers, Marco Fiscato, Mohamed Ahmed
TL;DR
Drug discovery needs molecular representations that support Virtual Screening and QSAR despite ambiguity in SMILES linearizations. The paper applies BERT with domain-relevant self-supervised tasks, finding that task choice strongly affects downstream performance and that MOLBERT reaches state-of-the-art results on established benchmarks.
Problem
SMILES can represent the same molecular graph through multiple valid sequences, creating ambiguity for language-model-based molecular representations used in Virtual Screening and QSAR.
Method
MOLBERT applies bidirectional BERT to molecular strings and compares self-supervised pre-training tasks, including SMILES equivalence and calculated molecular descriptor prediction.
Results
MOLBERT achieves state-of-the-art performance on established Virtual Screening and QSAR benchmarks, with fine-tuning best on all six QSAR datasets.
Takeaways & Limitations
The choice of self-supervision significantly affects downstream performance, and chemistry-relevant descriptor prediction improves the learned molecular representations.
Takeaways & Limitations
The authors leave applying MOLBERT to representations of other entities such as proteins for future work.
Abstract
from arXiv · showhide
We apply a Transformer architecture, specifically BERT, to learn flexible and high quality molecular representations for drug discovery problems. We study the impact of using different combinations of self-supervised tasks for pre-training, and present our results for the established Virtual Screening and QSAR benchmarks. We show that: i) The selection of appropriate self-supervised task(s) for pre-training has a significant impact on performance in subsequent downstream tasks such as Virtual Screening. ii) Using auxiliary tasks with more domain relevance for Chemistry, such as learning to predict calculated molecular properties, increases the fidelity of our learnt representations. iii) Finally, we show that molecular representations learnt by our model `MolBert' improve upon the current state of the art on the benchmark datasets.
1 Introduction
Molecular representations are central to drug-discovery prediction and search, but SMILES-based language models face ambiguity from multiple valid graph traversals. This work evaluates BERT with domain-relevant auxiliary pre-training tasks for molecular representations.
- Effective molecular representations support Virtual Screening and QSAR modeling and can strongly affect novel drug-candidate discovery efficiency.
- Language models have been applied to SMILES for molecular property, reaction prediction, and generative tasks.
- SMILES linearization is ambiguous because multiple valid sequences can represent the same molecular graph.
- Canonicalization reduces representational ambiguity but can introduce artifacts that distract language models.
- The study evaluates BERT molecular representations using domain-relevant auxiliary pre-training tasks on Virtual Screening and QSAR benchmarks.
2 MOLBERT
MOLBERT is a bidirectional BERT language model that combines self-supervised tasks to learn molecular representations, then applies them directly or through downstream models. Its evaluation targets Virtual Screening and QSAR applications.
- MOLBERT uses the BERT architecture as a bidirectional language model for molecular representation learning.
- MASKEDLM trains the model to predict masked token identities using cross-entropy loss.
- SMILES-EQ predicts whether two SMILES represent the same molecule when the second is either a random molecule or a synonymous permutation of the first.
- PHYSCHEMPRED predicts normalized molecular descriptors computed with RDKit using mean squared error.
- MOLBERT representations are used for similarity search, SVM-based downstream modeling, or a newly initialized downstream task head.
- Evaluation covers Virtual Screening retrieval and QSAR regression and classification benchmarks.
3 Experimental Evaluation
The experiments ablate pre-training tasks and evaluate MOLBERT against molecular-representation baselines on Virtual Screening and QSAR benchmarks. Domain-relevant descriptor prediction is especially influential, and fine-tuned MOLBERT achieves the best QSAR performance across all six datasets.
- The evaluation combines task ablations with Virtual Screening and QSAR benchmark comparisons.
- Pre-training uses the GuacaMol dataset of approximately 1.6M compounds, while the BERT-Base models contain approximately 85M parameters.
- MOLBERT embeddings generally use the pooled output, while MASKEDLM-only models use averaged sequence outputs.
- Ablation study: PHYSCHEMPRED alone achieves average BEDROC20 0.292 versus 0.266 for MASKEDLM alone, while adding MASKEDLM contributes +0.031 average BEDROC20 and SMILES-EQ slightly decreases performance.
- Descriptor grouping: Using all descriptors gives the best overall result, while SURFACE properties provide competitive supervision using only 25% of descriptors.
- Virtual Screening: MOLBERT outperforms CDDD, RDKit descriptors, and ECFC4 on Virtual Screening for both overall classification and early enrichment.
- QSAR: Fine-tuned MOLBERT achieves the best performance on all six QSAR benchmark datasets, while MOLBERT with an SVM leads the other descriptors on three.
4 Conclusions
MOLBERT uses BERT to learn molecular embeddings, and its downstream performance depends significantly on the selected self-supervised pre-training tasks. With an appropriate task set, it achieves state-of-the-art results on established Virtual Screening and QSAR benchmarks, while extension to other entities remains future work.
- MOLBERT uses BERT to learn molecular embeddings for drug discovery.
- The choice of self-supervision task significantly impacts performance on downstream tasks.
- With the right set of pre-training tasks, MOLBERT achieves state-of-the-art performance on established Virtual Screening and QSAR benchmarks.
- Future work includes applying MOLBERT to representations of other entities, such as proteins, and developing further learning strategies.
A Virtual Screening benchmark: performance per target
Figure A1 compares BEDROC20 and AUROC across Virtual Screening target datasets, sorted by BEDROC20. MOLBERT is superior for 45 of 69 targets and competitive on the remaining targets.
- 45 of 69 targets show superior MOLBERT performance, with competitiveness in all other cases.The datasets are sorted using the BEDROC20 enrichment metric.
- Figure A1 reports BEDROC20 and AUROC for each Virtual Screening target dataset.Results are shown for the best-performing MOLBERT model from the ablation study, trained for 100 epochs.
B Auxiliary task induced biases in MOLBERT models
The analysis examines how auxiliary pre-training tasks shape MOLBERT’s embedding space and relate to Virtual Screening performance. PHYSCHEMPRED encourages high similarity for SMILES permutations and low similarity for unrelated molecules, whereas several other task combinations fail to structure the space effectively.
- The analysis hypothesizes that high within-molecule permutation similarity and low unrelated-molecule similarity improve nearest-neighbor retrieval.
- PHYSCHEMPRED-inclusive models assign high similarity to permutations of the same molecule and low similarity to random ChEMBL compounds.The analysis compares pairwise cosine similarities across molecule groups and relates these behaviors to representation structure.
- MASKEDLM and MASKEDLM + PERMUTE produce similarly low similarities for same-molecule permutations and unrelated molecules, indicating poor semantic structuring.These combinations may map inputs into a large representation space without organizing it semantically.
- SMILES-EQ + PERMUTE alone does not structure the representation space, but adding MASKEDLM groups permutations and improves BEDROC.
- The t-SNE analyses visualize molecules and their SMILES permutations under task combinations that either encourage or fail to encourage disambiguation.Tables A2–A4 use the compounds listed in Table A1 and include BEDROC values in their headings.
C PhysChemPred subset ablation
The PHYSCHEMPRED ablation tests subsets of the 200 RDKit descriptors on the Virtual Screening benchmark. Although the full descriptor set performs best overall, small subsets—especially SURFACE and CHARGE—achieve nearly comparable results.
- The ablation groups the 200 PHYSCHEMPRED descriptors into related subsets and repeats the Virtual Screening benchmark.
- All physicochemical descriptor groups perform well on the benchmark.
- SURFACE with 49 descriptors and CHARGE with 18 descriptors achieve almost the same results as all physicochemical descriptors.
- The full set of 200 descriptors achieves the best overall result.