The process of pre-training large transformer models (such as BERT) on huge datasets has dominated Natural Language Processing research. More recently, this technique has been used to predict chemical properties, reaction products and analyse molecular information after fine-tuning on molecular data. Pre-training is a very computationally expensive process, as such, the scalability of such models has not yet been explored in the chemical domain. In this work, we train the largest transformer in chemical literature on an aggregation of publicly available SMILES and reaction datasets (Pubchem, ZINC, ChEMBL, GDB-13, USPTO) of over 1.2B+ unique molecules. We achieve this by splitting the load over multiple machines communicating over the network in a distributed training procedure; this allows training to be performed in a reasonable time frame and in a more robust fashion. Our approach is based on the RoBERTa architecture and is entirely data-driven. A unigram SentencePiece model is trained on the dataset and used to automatically tokenize the SMILES based on the log-likelihood during training. Molecular fingerprints are often used by predictive machine learning models as a tool for feature extraction and dimensionality reduction – they are often fixed-length vectors which describe the most important features of a molecule. We explore the resulting contextual embeddings from our model’s latent space as general molecular fingerprints for enabling better predictability with pre-existing models. The RDKit Benchmarking platform was used to compare different molecular fingerprints, with our best model achieving an ROC AUC score of 0.846 and a BED ROC 20 score of 0.584 on the filtered subset I with cosine similarity compared to the more classic ECFC4 fingerprint which scored 0.603 and 0.170 respectively. This shows potential great benefits to scaling when combining with fine-tuning on downstream tasks.

Large-scale distributed pre-training of transformers for data-driven molecular fingerprinting

ACS Spring Meeting 2021

The official archive and repository for all ACS Meetings & Expositions abstracts and select posters or presentations.