Large-Scale Chemical Language Representations Capture Molecular Structure and Properties

Jerret Ross; Brian Belgodere; Vijil Chenthamarakshan; Inkit Padhi; Youssef Mroueh; Payel Das

Source-linked AI summary

Large-Scale Chemical Language Representations Capture Molecular Structure and Properties

Jerret Ross, Brian Belgodere, Vijil Chenthamarakshan, Inkit Padhi, Youssef Mroueh, Payel Das

arXiv:2106.09553v3Jun 17, 2021cs.LGcs.CLq-bio.BM

TL;DR

Limited molecular property labels make supervised learning difficult across the vast chemical space. The paper trains MoLFormer on 1.1 billion unlabeled molecular SMILES and finds that its representations outperform existing baselines on several downstream tasks while remaining competitive on others.

Problem
Scarce molecular property labels and an astronomically large chemical space limit supervised learning for molecular property prediction.
Method
MoLFormer learns molecular representations through masked-language-model pretraining with efficient linear attention on 1.1 billion chemical SMILES.
Results
MoLFormer-XL outperforms existing supervised and unsupervised language-model and GNN baselines on several molecular property tasks and performs competitively on others.
Takeaways & Limitations
The results provide evidence that SMILES-trained molecular language models capture chemical structure sufficient to predict diverse properties, including quantum-chemical properties.

Abstract

from arXiv · show

Models based on machine learning can enable accurate and fast molecular property predictions, which is of interest in drug discovery and material design. Various supervised machine learning models have demonstrated promising performance, but the vast chemical space and the limited availability of property labels make supervised learning challenging. Recently, unsupervised transformer-based language models pretrained on a large unlabelled corpus have produced state-of-the-art results in many downstream natural language processing tasks. Inspired by this development, we present molecular embeddings obtained by training an efficient transformer encoder model, MoLFormer, which uses rotary positional embeddings. This model employs a linear attention mechanism, coupled with highly distributed training, on SMILES sequences of 1.1 billion unlabelled molecules from the PubChem and ZINC datasets. We show that the learned molecular representation outperforms existing baselines, including supervised and self-supervised graph neural networks and language models, on several downstream tasks from ten benchmark datasets. They perform competitively on two others. Further analyses, specifically through the lens of attention, demonstrate that MoLFormer trained on chemical SMILES indeed learns the spatial relationships between atoms within a molecule. These results provide encouraging evidence that large-scale molecular language models can capture sufficient chemical and structural information to predict various distinct molecular properties, including quantum-chemical properties.

Main

MOLFORMER-XL is a large-scale molecular SMILES transformer trained on 1.1 billion molecules with efficient linear attention. Its representations capture chemical, structural, and spatial information and support diverse molecular property prediction, including quantum-chemical tasks.

Motivation: Scarce labels and an astronomically large plausible chemical space motivate generalizable molecular representation learning.The plausible chemical space is estimated at 10^60 to 10^100 molecules requiring annotation.
Model and training: MOLFORMER-XL uses efficient linear attention and is trained on a corpus of 1.1 billion molecules.The model is presented as a molecular SMILES transformer, with MOLFORMER-XL identified as the best-performing variant.
Evaluation: The study evaluates MOLFORMER through classification, regression, and ablation experiments spanning 10 benchmark datasets and multiple property domains.The datasets cover quantum mechanical, physical, biophysical, and physiological properties of small molecules.
Results: MOLFORMER representations perform better than or on par with state-of-the-art GNNs that learn precise graph topology and related information.The reported comparisons include GNNs using information such as bond distances.
Representation analysis: Analyses show that MOLFORMER captures molecular substructures and spatial interatomic distances from SMILES annotations alone.The study emphasizes that quantum-chemical prediction from SMILES alone is non-trivial because such properties depend strongly on accurate 3D geometry.

Results and Discussion

MoLFormer learns a universal molecular representation from large-scale SMILES data and is evaluated across classification and regression benchmarks using frozen or fine-tuned embeddings. Results and ablations show strong downstream performance, while privileged 3D information remains important for quantum-chemical energy prediction.

Evaluation setup: MoLFormer is pretrained with masked language modeling on large-scale chemical SMILES and evaluated as a task-agnostic representation on MoleculeNet classification and regression tasks.Rotary embeddings produced more stable and faster-converging training than absolute embeddings.
Evaluation setup: MOLFORMER-XL uses mean-pooled final hidden-state embeddings and supports frozen evaluation with a task-specific fully connected model or end-to-end fine-tuning.MOLFORMER-XL was pretrained on approximately 1.1 B molecules from PubChem and Zinc using linear attention with rotary positional embeddings.
Classification tasks: The benchmark compares MOLFORMER-XL against nine baselines across six classification tasks, including fingerprint models, graph neural networks, and self-supervised or geometry-aware methods.The baselines include RF, SVM, GIN-based methods, GraphMVP-C, DimeNet, GeomGCL, GEM, and ChemBERTa.
Regression tasks: MOLFORMER-XL is evaluated on QM9, QM8, ESOL, FreeSolv, and Lipophilicity, including quantum-chemical regression tasks that are challenging without privileged 3D geometry.QM9 and QM8 predict several quantum-chemical measures.
Quantum-chemical energy prediction: On QM9 atomization-energy and enthalpy prediction, SchNet and DimeNet outperform MOLFORMER-XL by roughly factors of 8 and 10, respectively, while directly encoding 3D information.This supports the value of SMILES-based universal representations while confirming the importance of privileged geometric information for quantum-chemical energy prediction.
Ablation results: Fine-tuning decisively outperforms frozen embeddings on QM9 across all pretraining dataset sizes, and rotary embeddings surpass absolute embeddings once dataset size exceeds 1 Billion molecules.Rotary embeddings trail absolute embeddings on smaller datasets before winning at the largest scale.

Insights into MOLFORMER

MOLFORMER-XL embeddings reflect chemical similarity and molecular substructure, while its attention patterns capture short- and medium-range spatial relationships between atoms beyond SMILES-sequence neighbors.

Chemical similarity: MOLFORMER-XL embedding Euclidean distances were compared with molecular-fingerprint Tanimoto distances and maximum-common-subgraph atom counts to assess chemical similarity.The analysis used random molecule pairs to examine relationships between embedding-space distances, fingerprint similarity, and shared molecular structure.
Spatial attention: Attention in MOLFORMER with linear or full attention and rotary positional embeddings strongly matches short- and medium-range interatomic distances.Short distances are ≤2 Å, medium distances are 2-4 Å, and short distances mostly reflect typical covalent bonds such as the 1.5 Å C-C single bond.
Spatial attention: Attention shows weak, around 0.2 similarity with long interatomic distances of ≥4 Å.The pattern indicates that MOLFORMER captures spatial relations between atomic tokens that need not be neighbors in the SMILES sequence.
Spatial attention: Figure 3 compares learned attention coefficients with covalent-bond connectivity and 3D distances between atom pairs in an intermediate MOLFORMER-XL layer.The analysis used two molecules from the QM9 test set, with complete attention matrices across layers provided in Figures 5 and 6 of the Supplementary Information.

Conclusion

The work shows that large-scale self-supervised molecular language models can support molecular property prediction across diverse tasks, despite SMILES not explicitly encoding topology. MOLFORMER also shows potential for faster in silico screening, while requiring experimental and scientific validation to mitigate misuse.

Model capabilities: Large-scale self-supervised molecular language models can predict diverse molecular properties despite SMILES not explicitly encoding molecular topology.The approach combines self-supervised training on a large corpus with an expressive transformer architecture, linear attention, and parallelized training.
Benchmark performance: MOLFORMER outperforms existing graph-based baselines across a wide variety of molecular regression and classification benchmarks.The validated prediction range extends from quantum chemical to physiological properties.
Interpretability: Attention analysis shows that MOLFORMER trained on SMILES sequences is aware of interatomic relationships.The analysis supports the model’s ability to learn structural information from molecular language representations.
Applications and limitations: MOLFORMER has immediate potential for faster in silico screening in material design and drug discovery, but misuse without wet-lab validation can have harmful implications.The conclusion emphasizes proper experimental and scientific validation before relying on such technology.

Methods

MoLFormer uses a linear-attention transformer encoder with a modified rotary positional-embedding scheme for efficient chemical-SMILES pretraining. Training combines canonicalized PubChem and ZINC molecules, masked-language-model objectives, and distributed sequence-length bucketing to scale computation.

Model architecture: MoLFormer uses a 12-layer transformer encoder with linear attention to avoid vanilla attention’s quadratic sequence-length cost.Linear attention and bucketing reduced the required GPUs from roughly 1000 to 16.
Positional embeddings: The modified rotary embedding rotates the original keys and queries before the feature map, training faster and achieving better performance than absolute embeddings.Preliminary linear RoFormer experiments performed worse than the absolute-position counterpart; the proposed modification trained more gracefully and performed better.
Pretraining data: Pretraining data combine PubChem’s 111 million molecules and ZINC’s over 1 billion molecules after RDKit canonicalization and tokenizer-based vocabulary construction.Datasets were constructed with varying proportions of PubChem and ZINC.
Pretraining objective: Masked-language-model pretraining selects 15% of tokens, masking 80% of selected tokens, replacing 10% randomly, and leaving 10% unchanged.The method follows the masked language model objective defined in the cited prior work.
Distributed training: Training ran for 4 epochs with a fixed learning rate of 1.6e−4 and 1600 molecules per GPU across 16 GPUs and 2 Infiniband-connected servers.Increasing the GPU count required increasing the learning rate by up to a factor of 8.

Code availability

The authors provide the MoLFormer training and fine-tuning code, attention-visualization notebooks, and pretrained model instances through the project’s GitHub repository.

Code availability: Python code for MoLFormer training and fine-tuning, attention-visualization notebooks, and instances of pretrained models are available on GitHub.Other enquiries can be directed to the corresponding authors.

Figures and Tables

This section contains three figures: Figures 1, 2, and 3.

Figure 1 is included in the section.
Figure 2 is included in the section.
Figure 3 is included in the section.

Figure and Table Captions · Extended data

The extended data captions document comparisons of MoLFormer-XL and its variants across classification and regression benchmarks, including effects of model size, pretraining data, fine-tuning, and rotary embeddings. They also describe analyses relating molecular embeddings to structural similarity metrics.

Extended data: Extended data Table 1 compares MoLFormer-XL with smaller or less extensively pretrained MoLFormer models on six classification benchmarks.The benchmarks are BBBP, HIV, Sider, Clintox, Tox21, and BACE.
Extended data: Extended data Table 2 compares fine-tuned MoLFormer-XL with smaller or less extensively pretrained MoLFormer models on five regression benchmarks.The benchmarks are QM9, QM8, ESOL, FreeSolv, and Lipophilicity, using the captioned metrics.
Extended data: Extended data Table 3 compares MoLFormer variants on the QM9 test set using average MAE and average standard MAE.The variants differ in pretraining data, downstream fine-tuning, and the presence or absence of rotary embeddings.
Extended data: The variants in Extended data Table 3 are pretrained on QM9 only, PubChem only, or PubChem+ZINC datasets.The caption states that the best candidate variant for Table 8 is chosen using average MAE, with lower Mean Absolute Error preferred.
Extended data: Extended data Table 4 reports correlations between embedding-based similarities, fingerprint similarities, and structural similarity measured by maximum common subgraph size.The analysis uses 10000 randomly selected molecule pairs from PubChem and includes MoLFormer-XL and ChemBERTa embeddings.
Extended data: The structural analysis in Extended data Table 4 also relates maximum common subgraph atom counts to Euclidean distances in embedding space.These correlations assess how molecular embeddings reflect structural relationships.

Related Work … A.2 Parallelization and Computing Environment

The paper situates MoLFORMER among pretrained language, molecular representation, and geometry-aware learning approaches, then details its model choices, optimizer, linear attention, positional embeddings, and distributed computing environment.

Related Work: Transformer-based masked and prefix language models improve downstream language understanding through self-supervised pretraining and increased model scale.MLMs predict randomly masked tokens, whereas PLMs add task-specific text tags to input sequences.
Related Work: Molecular representation learning has used chemical fingerprints, supervised neural networks, SMILES-based recurrent models, and graph convolutional networks.Molecular graphs can also incorporate costly 3D geometry and pairwise interatomic distances to model spatial interactions.
Related Work: Large-scale unsupervised pretraining has extended contextual language representations from proteins to SMILES-based molecular property prediction.Prior SMILES models did not attempt to predict a diverse range of molecular properties.
Related Work: Molecular self-supervised methods have used motif prediction, substructure and cheminformatics information, cross-view mutual information, and geometry-based pretraining.These views include 1D SMILES, 2D graphs, and 3D geometry, with GEM modeling atom–bond–angle relations.
A.1 MOLFORMER Model and Pre-Training Details: MOLFORMER pretraining details cover optimization, linear attention, positional embeddings, and implementation choices.The appendix introduces these details before describing individual components and experiments.
A.1.1 Optimizer: Fused Lamb was selected because of superior training performance, robustness to large batch sizes, and reduced need for learning-rate warmups or optimizer modifications.Other optimizers became unstable when training configurations changed without substantial modification.
A.1.2 Linear Attention: 32 was the selected feature-map size for generalized linear attention, balancing computation speed against minimal performance loss relative to FAVOR45.Generalized Features simplify the FAVOR45 feature map.
A.1.3 Rotary versus Absolute position embeddings: MOLFORMER with linear attention and rotary embeddings achieved better validation loss than the absolute-position counterpart, motivating its use throughout the paper.Training used distributed multi-node, multi-GPU infrastructure with PyTorch Distributed Data Parallel, NCCL, RDMA GPU-direct technology, and localized HuggingFace Datasets.

A.3 Memory Efficient Training with Adaptive Bucketing By Sequence Length … D Additional results on comparing MOLFORMER-XL with geometry-aware GNNs on regression benchmarks

The paper uses adaptive sequence-length bucketing to reduce molecular pretraining waste, evaluates scaling and fine-tuning choices, specifies dataset conventions, and compares MOLFORMER with geometry-aware GNNs.

A.3 Memory Efficient Training with Adaptive Bucketing By Sequence Length: Molecules were typically shorter than 45 tokens, so minibatches were split into multiple buckets on the fly to reduce padding-related computational waste.Adaptive bucketing avoided full-dataset preprocessing while grouping sequences by length batch by batch.
A.3 Memory Efficient Training with Adaptive Bucketing By Sequence Length: 20× speedup reduced a single PubChem epoch on one V100 from approximately 1200 hours without bucketing to around 60 hours with bucketing.The authors had not benchmarked their molecular bucketing implementation against generic micro-batching.
A.4 Pre-training Scaleout: Training time decreased slightly sub-linearly as GPUs were added and scaled approximately linearly with additional data.These estimates used MOLFORMER with rotary embeddings after 250 iterations on PubChem and PubChem+ZINC.
C Dataset, Vocabulary, and Property Units: Downstream evaluations used MoleculeNet tasks, with random splits for Table 2 tasks and scaffold splits for Table 1 tasks.The paper provides dataset descriptions in Tables 2 and 3 and directs readers to the cited source for task details.
D Additional results on comparing MOLFORMER-XL with geometry-aware GNNs on regression benchmarks: Additional results compare MOLFORMER with supervised DimeNet and self-supervised GeomGCL and GEM on MoleculeNet physical-chemistry regression benchmarks.These comparisons are reported in Table 7 and focus on geometry-aware graph neural networks.

E Additional Results and Ablations on QM9 Benchmark … H Assets and License

Additional QM9 ablations show that larger pre-training corpora, rotary embeddings, and fine-tuning improve MOLFORMER performance, while cross-validation indicates robust representations. Visualization analyses further suggest that frozen embeddings capture molecular properties and attention maps encode structure; the experiments use libraries with documented licenses.

E Additional Results and Ablations on QM9 Benchmark: MOLFORMER-XL outperforms all baselines in average MAE and average standard MAE on QM9, while linear attention retains linear time complexity.The more expressive 123-GNN performs better on most individual measures but is difficult to scale.
E Additional Results and Ablations on QM9 Benchmark: Larger pre-training datasets yield lower MAE on QM9, with MOLFORMER-XL denoting pre-training on the 1.1-billion-molecule PubChem+ZINC corpus.Rotary embeddings trail absolute embeddings on smaller datasets but win once the dataset exceeds 1 billion molecules.
E Additional Results and Ablations on QM9 Benchmark: Fine-tuning and rotary embeddings achieve the best QM9 performance compared with frozen embeddings and absolute embeddings.The ablation directly compares fine-tuned versus frozen MOLFORMER embeddings under both positional-embedding schemes.
E Additional Results and Ablations on QM9 Benchmark: Across five 80%/10%/10% cross-validation folds, low standard deviations and mean errors aligned with the main paper, suggesting robust MOLFORMER representations.Most related work reports only a single split rather than cross-validation.
F Insights into MOLFORMER- tSNE Visualization: Without task-specific fine-tuning, frozen MOLFORMER-XL embeddings discriminate BBBP penetrating from non-penetrating molecules and visibly cluster toxic versus non-toxic ClinTox molecules.The ClinTox clusters show some overlap, while toxic molecules near non-toxic clusters tend to lie at cluster borders.
F Insights into MOLFORMER- tSNE Visualization: Frozen MOLFORMER-XL embeddings also show a separation tendency between high- and low-LUMO classes, although substantial overlap remains.This visualization evaluates discretized LUMO classes without fine-tuning.
G MOLFORMER Attention Visualization and Structure Discovery: Linear- and full-attention MOLFORMER-XL variants have similar downstream performance, but their attention weights offer insight into structural information learned from SMILES.The visualized models use rotary positional embeddings and are not fine-tuned on QM9, yet their attentions represent molecular structure.
H Assets and License: The experiments use libraries including Fast-Transformers, PyTorch, RDKit, PyTorch Lightning, and HuggingFace Datasets under documented MIT, BSD, and Apache 2.0 licenses.The assets section summarizes the libraries and accompanying license terms.