A 3D-Shape Similarity-based Contrastive Approach to Molecular Representation Learning
Abstract
Molecular shape and geometry dictate key biophysical recognition processes, yet many graph neural networks disregard 3D information for molecular property prediction.
Here, we propose a new contrastive-learning procedure for graph neural networks, Mol ecular C ontrastive L e a rning from S hape S imilarity ( MolCLaSS ), that implicitly learns a three-dimensional representation.
Rather than directly encoding or targeting three-dimensional poses, MolCLaSS matches a similarity objective based on Gaussian overlays to learn a meaningful representation of molecular shape.
We demonstrate how this framework naturally captures key aspects of three-dimensionality that two-dimensional representations cannot and provides an inductive framework for scaffold hopping.
Introduction and Background
Molecular shape is critical for biophysical processes, yet encoding relevant three-dimensional features remains challenging for many molecular property prediction tasks, especially when an understanding of three-dimensional shape is limited or unknown [1].
Numerous methods have been developed to effectively encode individual conformers that are both appropriate and highly-effective for conformerlevel prediction tasks such as predicting quantum chemical properties of single conformational poses [2; 3; 4; 5; 6].
However, these approaches are poorly suited for representing complete molecules since relying on a single low-energy conformer to represent a diverse conformational ensemble is inherently limiting.
Dietterich et al.
[7] first recognized and addressed this challenge in the development of the multiple-instance learning framework [8].
Recent studies have explored deep multiple-instance learning approaches for learning on conformational ensembles [9; 10], yet are computationallydemanding due to the need to encode each conformer independently.
Furthermore, methods that encode three-dimensional information often do not provide a strong performance benefit over 2D baselines [10].
Herein, we propose to inject graph neural networks with implicit 3D shape information through a supervised contrastive approach.
Learning implicit 3D representations have been recently explored by Stärk et al.
[11] and Liu et al.
[12], but these approaches do not consider three-dimensional relationships between molecules and hence do not learn a direct measure of molecular shape similarity.
In contrast, we propose to learn key features of molecular shape through direct comparison with the use of 3D molecular similarity kernels based on Gaussian overlays [13].
Specifically, methods such as the Rapid Overlay of Chemical Structures (ROCS)
[14] provide a fast and scalable method for matching molecular shapes based both on molecular volumes and electrostatic matching on predefined pharmacophore features.
Our framework, Mol ecular C ontrastive L e a rning on S hape S imilarity (MolCLaSS), naturally aligns a pairwise shape similarity objective with supervised contrastive learning framework to learn meaningful representations of molecular shape (Figure 1).
Our work contributes directly in two ways: 1) MolCLaSS provides fast, scalable, and inductive approximation of 3D-shape similarity scores between molecules directly from their 2D graphs, and does not require conformer generation and shape alignment for inference, and 2) we demonstrate how MolCLaSS learns meaningful molecular embeddings that naturally capture 3D-shape and features that 2D, topological methods cannot.
The resulting pre trained can used for shape
Related Work
Graph Neural Networks for Small Molecules Graph neural networks have been widely developed for predicting small molecule properties and activities [15; 16; 17] on both 2D and 3D tasks.
The properties of individual conformers have been effectively modeled by leveraging 3D spatial features for a range of quantum chemical and property prediction tasks [3; 4; 5; 6].
Adams et al.
[18] recently described a hybrid approach for conformationally-invariant 3D approaches.
Furthermore, "4D" multiple-instance learning methods have recently been developed that operate over sets of conformers that are each modeled as a graph [9; 10].
Our work builds on this prior work to learn improved and compact representations.
Pre-Training Graph Neural Networks Graph pretraining methods are an active area of research.
Hu et al.
[19] first reported a pretraining strategy based on self-supervised node pre-training, followed by supervised pre-training with masking.
You et al.
[20] recently demonstrated the effectiveness of a contrastive learning approach on graphs, which was further extended by Wang et al.
Recently, Stärk et al.
[11] and Liu et al.
[12] developed self-supervised approaches to maximize three-dimensional information for 2D-graph neural network pretraining with promising results.
These approaches aim to maximize mutual information with the goal of matching a 2D representation to a three-dimensional pose, but they do not explicitly consider three-dimensional similarity between molecules.
Our work here presents a supervised contrastive approach for pretraining neural networks that complements the self-supervised approaches above.
Kernel-Based Approximations of Molecular Shape
Our work is closely related to prior work on learning low dimensional embeddings from molecular similarity kernels.
Raghavendra and
Maggiora [23] introduced a method to learn molecular basis vectors by directly decomposing
Tanimoto similarities.
The SCISSORS method by Haque & Pande [24; 25; 26] generalizes a kernel
PCA approach to molecular similarity measures, including ROCS, that provides a fast approximation for molecular similarity.
These prior approaches are naturally transductive; to obtain new scores, molecules must be scored against the resulting basis set and predicted based on a least-squares estimate.
Our approach couples the intuition of Haque and Pande [24] with modern graph neural networks to directly learn a meaningful molecular space that is naturally inductive, i.e. can generalize to new and unseen molecules without the need for additional conformer generation and scoring.
Problem Formulation and Methods
We invoke a variation of the classic similar property principle and assume that two molecules are similar if they can adopt similar molecular shapes [27; 28].
Prior approaches for encoding threedimensional information typically operate on a molecular input x i and corresponding spatial features s i that includes explicit atomic coordinate, distance, or angle information.
These approaches rely on designing an expressive transformation f θ ( x i , s i ) that can accurately learn to map similar molecules to similar parts of chemical space, and assume that a single or several low-energy conformers encode relevant geometry.
Rather than operate over explicit spatial representations, we instead adopt an implicit strategy that leverages a predefined similarity kernel over pairs of molecular inputs k ( x i , x j ) .
Critically, by leveraging a well-defined 3D similarity function using inner products, we avoid the need to explicitly encode spatial features s i , and can focus on learning an invariant function f θ ( x i ) .
As illustrated in Figure 1, our approach naturally fits a supervised contrastive learning framework: given a set of molecules, our goal is to learn an expressive representation satisfies a pairwise similarity constraint k ( x i , x j ) .
Here, we decompose this representation into graph encoder f θ and projection heads g θ to flexibly model multiple outputs.
In this study, we use Gaussian shape and color overlays as an intuitive measure of shape similarity [13; 1], and follow the insight of Haque and Pande [24, 25] to leverage the Tanimoto kernel [29]:
[Equation]
This interpretation defines a molecular embedding space, where three-dimensional shape similarities can be conveniently modeled based on inner products.
Rather than approximate embeddings z i via linear decomposition methods, our objective is to learn an inductive model that generates z i directly from a molecular graph.
Here, we use an encoder based on the Graph Isomorphism Network with
Edge features [19] followed by two projection heads to model Shape and Color separately:
[Equation]
with f θ ( x i )
=
h i = (cid:88) p
[Equation]
[Equation]
c
[Equation]
Above, v tp corresponds to the hidden state of the node p at step t (final step T ), with h i as the final graph representation of molecule i using sum pooling.
The projection heads g θ ( h i ) are parameterized by Color- and Shape-dependent MLPs with trainable weight matrices U and V , and σ is a ReLU nonlinearity.
Finally, we optimize the network to directly predict ShapeTanimoto and ColorTanimoto scores via minimization of the following loss function:
[Equation]
Here, we define the loss as the mean squared error over all the pairs in a batch (size N ) between the predicted Color and Shape scores based on the Tanimoto kernel (Eq. 1) and the calculated Gaussian overlay k { s,c } ( x i , x j ) based on conformer generation and alignment.
We balance the individual objectives L s and L c with the tuneable hyperparameter λ to adjust the influence of pharmacophore features (but typically set to 1 ).
ShapeTanimoto
[Table]
Experiments and Results
Prediction of Shape- and Color-Tanimoto Scores from 2D Graphs We systematically investigated the ability of contrastive models to directly predict ShapeTanimoto and ColorTanimoto scores of drug-like molecules from the ChEMBL database [30; 31].
For these studies, we generated a complete all-by-all similarity matrix of Shape- and ColorTanimoto scores [13; 14] for 100k molecules from ChEMBL, with maximum overlay scores recorded from pairwise comparisons of up toconformers generated per molecule [32; 33] (see Appendix for complete details). We refer to the complete data set as ROCS100k .
We assessed the ability of both fingerprint- and graph-based models to accurately predict pairwise, 3D similarity scores directly from their 2D graph representations.
Although thissetup requires computationally-intensive conformer generation and exhaustive similarity scoring, this cost is amortized over the training data.
At inference time, predicted Shape- and ColorTanimoto similarities are directly obtained and circumvent the need for explicit three-dimensional representations.
As illustrated in Table 4, we directly compared our approach to a Morgan fingerprint [34] and dense neural network baseline, with positive performance gains for increasingly large data set sizes.
Our graph neural network-based approach based on GINEConv graph layers with independent projection heads for ShapeTanimoto and ColorTanimoto predictions exhibits a clear improvement over hashed fingerprint representations, even with significantly less data.
For example, MolCLaSS trained with only 10k examples achieves nearly identical performance to fingerprint-based models trained on 100k examples.
Critically, these dense networks learn a non-trivial transformation of the input data.
Indeed, as shown in Table 4 (first row), there is nearly no correlation between bulk Tanimoto scores on 2D representations and their 3D Tanimoto scores.
Notably, the MolClASS network can directly predict 3D similarity scores with good accuracy (ShapeTanimoto MAE = 0 .
030 , ShapeTanimoto MAE = 0 .
at only a fraction of the computational cost.
At inference, predicting 3D similarity scores on tens of thousands of molecules takes only seconds, replacing both the need for conformer generation and overlays and representing an improvement of nearly 10 4 − 10 5 times in speed.
MolCLaSS Representations Capture 3D Shape Similarity.
Given the strong performance of
MolCLaSS, we investigated the qualities of the learned molecular embeddings.
We specifically visualized our hold-out test set of ChEMBL molecules in their graph embedding ( h i ) and projection heads ( z s i , z c i ) (Figure 2) against a fingerprint baseline, and colored them based on calculated three-dimensional descriptors [28; 35] of their low-energy conformers, including radius of gyration (Figure 2A) and the first principal moment of inertia (Figure 2B).
As shown, both the graph-layers and shape-projection layer provide clear localization based on these 3D properties when compared to Morgan fingerprints.
As expected, the color projection head does not exhibit the same localization, as pharmacophore features are less dependent on overall molecular shape.
A key consequence of the Tanimoto similarity objective (Eq. 1) is that it induces a Euclidean structure over the projected vector space z [36].
Indeed, we found an excellent correlation (Pearson r = 0 . 87 ) between pairwise Euclidean distance and ShapeTanimoto scores (see Appendix).
To further probe the difference in representations we performed a nearest neighbors analysis using the shape projections in MolCLaSS (Figure 3).
In our analysis, we consistently find that MolCLaSS preserves molecular shape and size.
Scaffolds hops based on ring mutations are found nearby while maintaining excellent overall shape similarity.
In contrast, topological fingerprints largely favor substructure matching and exhibits a wider range of molecular shapes.
Together, these two studies illustrate how the MolCLaSS framework can capture relevant shape measures through a supervised approach.
found, nearest neighbors have
Conclusions and Future Directions
Our studies above outline a preliminary roadmap for learning three-dimensional representations based on 3D similarities.
In contrast to prior work, MolCLaSS learns via supervised pairwise comparisons, and hence is able to relate and differentiate molecules of varying size and shape.
We have demonstrated how the MolCLaSS network itself provides a direct and fast approximation method for approximating shape- and 3D pharmacophore-based and further illustrated how meaningful three-dimensional features are naturally learned through this inductive framework.
Our ongoing work seeks to further improve the MolCLaSS framework to improve predictive performance and to explore its application as a pretrained model for broad molecular property prediction tasks.
Acknowledgments and Disclosure of Funding
We gratefully acknowledge members of the Department of Artificial Intelligence and Machine
Learning at Genentech Research and Early Development for helpful feedback and support.
Appendix
All experiments were performed using Python and standard numerical libraries.
For cheminformatics analysis, all molecules were processed using either OpenEye Applications and Toolkits or the opensource cheminformatics library RDKit [35].
We implemented all experiments in Python using
PyTorch 1.11 [37] and PyTorch Geometric [38].
Appendix A: Molecular Dataset Properties and Creation
All data were downloaded directly from ChEMBL31 release of the ChEMBL database [30; 31].
Molecules were filtered based on OpenEye Filter using the drug-likeness that limits.
The resulting detailed statistics of the molecules and some of their properties are shown in Table 2, with a representative sample of molecules shown in Figure 4.
OMEGA Conformer Generation
For each molecule, we generate a conformational ensemble using OpenEye Applications ( 2022.1.1 ) using OMEGA ( v.4.2.0 ) [32; 33].
Conformational ensembles were generated with the optimized default fastROCS settings with multiprocessing: maxconfs=10, ewindow=15, flipper=
False, mpi=128 which been shown to accurately recapitulate binding poses.
Molecules with ambiguous or undefined stereochemistry were dropped during the conformer generation process.
OpenEye fastROCS Scoring
We use OpenEye’s
GPU-accelerated fastROCS Toolkit ( v.2.2.2.1 ) to calculate ShapeTanimoto and ColorTanimoto scores.
For each conformer database, we generate an all-by-all similarity matrix that scores every conformer of every molecule against all conformers of the rest of the database, saving the maximum score between two molecules.
In practice, generation of the ROCS100k training dataset generates 4.95 billion ShapeTanimoto and ColorTanimoto scores, each, and requires 7 hours on 3 NVIDIA 3090 RTX GPUs with 80 processes.
Computational Complexity of All-by-All Conformer ROCS Exhaustive similarity comparisons between n molecules scales at O ( n 2 ) .
As an upper limit, for k conformers are generated per molecule, generation of the all-conformer by all-conformer similarity matrix scales at O ( k 2 n 2 ) .
For symmetric similarity measures like ColorTanimoto and ShapeTanimoto, O (( n )( n − 1) / 2) comparisons are required with no self-comparisons.
Similarity, for k sampled conformers we require O (( kn )( kn − 1) / 2) pairwise comparisons.
[Table]
Appendix B: Network Architecture and Training Details
All studies were trained using the generated ROCS100k dataset.
The size dependence studies illustrated in 4 using 10k and 50k examples use subsets of the full 100k molecules.
A separate validation set of 8,548 molecules from ChEMBL31 were additionally used for model tuning and selection.
All results in Table 4 are reported on an independent, random test set of 49,621 molecules from ChEMBL31. [Table]
All neural networks were trained using the Adam optimizer (learning rate = 1 × 10 − 3 to 1 × 10 − 4 ) and a batch size of 2048 for up to 4,000 epochs, using the early stopping criterion based on the validation set described above.
The model architecture and hidden dimensions are specified in Appendix Table 4.
All networks use five graph encoding layers with two project heads (single-hidden layer MLPs).
The entire network is trained on Shape- and Color-Tanimoto targets bounded within [0,1], using mean squared error as the loss criterion and trained to early stopping.
Neural Network Performance
Our summary of overall performance results and metrics for different model architectures are shown in Table 4.
Below, we include scatter plots of our best model trained on the full ROCS100k dataset in
Figure 5.
Correlation of Euclidean Distance with ShapeTanimoto and ColorTanimoto
The project heads g θ ultimately learn a meaningful Euclidean distance.
As shown in the plots below, molecules closer in the embedded latent space also tend to have a much higher Shape and
ColorTanimoto score (Figure 6).
[Table]
Graph Encoder f θ 5 x [GINEConv + BatchNorm w/ ReLU] 5 × (512 → 1024 → 512)
Projection
[Table]
θ
MLP + ReLU 512 → 1024 →
ECFP4 + FF-NN
Encoder f θ
MLP w/ ReLU 2048 → 2048 →
Projection [Table] θ
MLP + ReLU 512 → 1024 →