Predicting molecular properties with covariant compositional networks

Density functional theory¹ (DFT) is the workhorse of modern quantum chemistry due to its ability to calculate many properties of molecular systems with high accuracy. However, this accuracy comes at a significant computational cost, generally making DFT too costly for applications such as chemical search and drug screening, which may involve hundreds of thousands of compounds. Methods that help overcome this limitation could lead to rapid developments in biology, medicine, and materials engineering.

Recent advances in machine learning, specifically, deep learning,² combined with the appearance of large datasets of molecular properties obtained both experimentally and theoretically,^3–5 present an opportunity to learn to predict the properties of compounds from their chemical structure rather than computing them explicitly with DFT. A machine learning model could allow for rapid and accurate exploration of huge molecular spaces to find suitable candidates for a desired molecule.

Central to any machine learning technique is the choice of a suitable set of descriptors or features used to parametrize and describe the input data. A poor choice of features will limit the expressiveness of the learning architecture and make accurate predictions impossible. On the other hand, providing too many features may make training difficult, especially when training data is limited. Hence, there has been a significant amount of work on designing good features for molecular systems. Predictions of energetics based on molecular geometry have been explored extensively, using a variety of parametrizations.^6,7 This includes bond types and/or angles,⁸ radial distance distributions,⁹ the Coulomb matrix and related structures,^10–12 the Smooth Overlap of Atomic Positions (SOAP),^13,14 permutation-invariant distances,¹⁵ symmetry functions for atomic positions,^16,17 moment tensor potentials,¹⁸ and scattering networks.¹⁹ 

Recently, the problem of learning from the structure of chemical bonds alone, i.e., the molecular graph, has attracted a lot of interest, especially in the light of the appearance of a series of neural network architectures designed specifically for learning from graphs.^20–26 Much of the success of these architectures stems from their ability to pick up on structures in the graph at multiple different scales, while satisfying the crucial requirement that the output be invariant to permutations of the vertices (which, in molecular learning, correspond to atoms). However, as we will show, the specific way that most of these architectures ensure permutation invariance still limits their representational power.

In this paper, we propose a new type of neural network for learning from molecular graphs, based on the the idea of covariant compositional networks (CCNs), recently introduced in Ref. 27. Importantly, CCNs are based on explicitly decomposing compound objects (in our case, molecular graphs) into a hierarchy of subparts (subgraphs), offering a versatile framework that is ideally suited to capture the multiscale nature of molecular structures from functional groups through local structures to the overall shape. In a certain sense, the resulting models are the neural networks analog of coarse graining. In addition, CCNs offer a more nuanced take on permutation invariance than other graph learning algorithms; while the overall output of a CCN is still invariant to the permutation of identical atoms, internally, the activations of the network are covariant rather than invariant, allowing us to better preserve information about the relationships between atoms. We demonstrate the success of this architecture through experiments on benchmark datasets, including QM9³ and the Harvard Clean Energy Project.⁴ 

This paper addresses the problem of predicting chemical properties directly from each compound’s molecular graph, which we will denote as $G$ (Fig. 1, left). As such, it is related to the sizable literature on learning from graphs. In the kernel machines domain, this includes algorithms based on random walks,^28–30 counting subgraphs,³¹ spectral ideas,³² label propagation schemes with hashing,^33,34 and even algebraic ideas.³⁵ 

Recently, a sequence of neural network based approaches has also appeared, starting with Ref. 36. Some of the proposed graph learning architectures^21,22,37 directly seek inspiration from the type of classical convolutional neural networks (CNNs) that are used for image recognition.^38,39 These methods involve moving a filter across vertices to create feature representations of each vertex based on its local neighborhoods. Other notable studies on graph neural networks include Refs. 24 and 40–42. Very recently, Ref. 25 showed that many of these approaches can be seen as message passing schemes and coined the term message passing neural networks (MPNNs) to refer to them collectively.

Regardless of the specifics, the two major issues that graph learning methods need to grapple with are invariance to permutations and capturing structures at multiple different scales. Let A denote the adjacency matrix of $G$⁠, and suppose that we change the numbering of the vertices by applying a permutation σ. The adjacency matrix will then change to A′, with

However, topologically, A and A′ still represent the same graph. Permutation invariance means that the representation $ϕ(G)$ learned or implied by our graph learning algorithm must be invariant to these transformations. Naturally, in the case of molecules, invariance is restricted to permutations that map each atom to another atom of the same type.

The multiscale property is equally crucial for learning molecular properties. For example, in the case of a protein, an ideal graph learning algorithm would represent $G$ in a manner that simultaneously captures structures at the level of individual atoms, functional groups, interactions between functional groups, subunits of the protein, and the protein’s overall shape.

The idea of representing complex objects in terms of a hierarchy of parts has a long history in machine learning.^43–48 We have recently introduced a general framework for encoding such part-based models in a special type of neural network, called covariant compositional networks.²⁷ In this paper, we show how the CCN formalism can be specialized to learning from graphs, specifically, the graphs of molecules. Our starting point is the following definition.

One of the messages of the present paper is that invariant aggregation functions, of the type popularized by message passing neural networks, are not the most general way to build compositional models for compound objects, such as graphs. To understand why this is the case, once again, an analogy with image recognition is helpful. Classical CNNs face two types of basic image transformations: translations and rotations. With respect to translations (barring pooling, edge effects, and other complications), CNNs behave in a quasi-invariant way, in the sense that if the input image is translated by an integer amount (t_x, t_y), the activations in each layer ℓ = 1, 2, …, L translate the same way: the activation of neuron $ni,jℓ$ is simply transferred to neuron $ni+t1,j+t2ℓ$⁠, i.e., $f′i+t1,j+t2ℓ=fi,jℓ$⁠. This is the simplest manifestation of a well studied property of CNNs called equivariance.^52,53

With respect to rotations, however, the situation is more complicated: if we rotate the input image by, e.g., 90°, not only the part of the image that fell in the receptive field of neuron $ni,jℓ$ will move to the receptive field of a different neuron $n−j,iℓ$ but also the orientation of the receptive field will change. For example, features which were previously picked up by horizontal filters will now be picked up by vertical filters. Therefore, in general, $f′−j,iℓ≠fi,jℓ$ (Fig. 2).

It can be shown that one cannot construct a CNN for images that behaves in a quasi-invariant way with respect to both translations and rotations, unless every filter is directionless. It is, however, possible to construct a CNN in which the activations transform in a predictable and reversible way, $f′−j,iℓ=R(fi,jℓ)$⁠, for some fixed function R. This phenomenon is called steerability and has a significant literature in both classical signal processing^54–58 and the neural networks field.⁵⁹ 

The situation in compositional networks is similar. The comp-net and message passing architectures that we have examined so far, by the virtue of the aggregation function being symmetric in its arguments, are all quasi-invariant (with respect to permutations) in the sense that if $N$ and $N′$ are two comp-nets for the same graph differing only in a reordering σ of the vertices of the underlying graph $G$⁠, and $ni$ is a neuron in $N$ while $nj′$ is the corresponding neuron in $N′$⁠, then $fi=fj′$ for any permutation $σ∈Sn$⁠.

Quasi-invariance amounts to asserting that the activation f_i at any given node must only depend on $Pi=ej1,…,ejk$ as a set and not on the internal ordering of the atoms $ej1,…,ejk$ making up the receptive field. At first sight, this seems desirable since it is exactly what we expect from the overall representation ϕ(G). On closer examination, however, we realize that this property is potentially problematic since it means that $ni$ loses all information about which vertex in its receptive field has contributed what to the aggregate information f_i. In the CNN analogy, we can say that we have lost information about the orientation of the receptive field. In particular, if higher up in the network f_i is combined with some other feature vector f_j from a node with an overlapping receptive field, the aggregation process has no way of taking into account which parts of the information in f_i and f_j come from shared vertices and which parts do not (Fig. 3).

The solution is to regard the $Pi$ receptive fields as ordered sets and explicitly establish how f_i co-varies with the internal ordering of the receptive fields. To emphasize that henceforth the $Pi$ sets are ordered, we will use parentheses rather than braces to denote them.

It remains to explain how to define the Φ aggregation functions so as to guarantee covariance. Specifically, we show how to construct Φ such that if the $Fc1,…,Fck$ inputs of a given node $na$ at level ℓ are covariant k’th order P–tensors, then $Fa=Φ(Fc1,…,Fck)$ will also be a k’th order P–tensor. The aggregation function that we define consists of five operations executed in sequence: promotion, stacking, reduction, mixing, and an elementwise nonlinear transform. (See Fig. 5.) Practically relevant CCNs tend to have multiple channels, so each $Fci$ is actually a sequence of d separate P–tensors $Fci(1),…,Fci(d)$⁠. However, except for the mixing step, each channel behaves independently, so for simplicity, for now we drop the channel index.

Each child tensor $Fci$ captures information about a different receptive field $Pci$⁠, so before combining them we must “promote” each $Fci$ to a $Pa×⋯×Pa$ tensor $F̃ci$⁠, whose dimensions are indexed by the vertices of $Pa$ rather than the vertices in each $Pci$⁠. Assuming that $Pci=eq1,…,eqPci$ and $Pa=ep1,…,epPa$⁠, this is done by defining a $|Pa|×|Pci|$ indicator matrix

and setting

where, once again, Einstein summation is in use, so summation over $j1′,…,jk′$ is implied. Effectively, the promotion step aligns all the child tensors by permuting the indices of $Fci$ to conform to the ordering of the atoms in $Pa$ and padding with zeros where necessary.

Now that the promoted tensors $F̃c1,…,F̃cs$ all have the same shape, they can be stacked to form a $Pa×⋯×Pa$ dimensional k + 1’th order tensor T, with

It is easy to see that T itself transforms as a P-tensor of order (k + 1).

We may also add additional information to T by taking its tensor product with another P-tensor. In particular, to explicitly add information about the local topology, we may tensor multiply T by $[A↓Pa]i,j=Aepi,epj$⁠, the restriction of the adjacency matrix to $Pa$⁠. This will give an order (k + 3) tensor $S=T⊗A↓Pa$⁠. Otherwise, we just set S = T. Note that in most other graph neural networks, the topology of the underlying graph is only accounted for implicitly, by the pattern in which different activations are combined. Being able to add the local adjacency matrix explicitly greatly extends the representational power of CCNs.

Stacking and the optional tensor product increase the order of our tensors from k to k + 1 or k + 3. We reduce this to k again by a combination of generalized tensor contractions, more accurately called tensor reductions. For example, one way to drop the rank of S from k + 3 to k is to sum out three of its indices,

Note that while this notation does not make it explicit, $iα1,iα2$ and $iα3$ must be removed from amongst the indices of Q. Another way to drop the rank is to contract over three indices,

where $δiα1,iα2,iα3$ is the generalized Kronecker symbol. A third way is to contract over two indices and sum over one index and so on. The crucial fact is that the result of each of these tensor operations is still a covariant P–tensor.

In general, the number of different ways that an order k + q tensor S can be reduced to order k depends on both q and k. For example, when k = 2 and q = 3, there are 50 different possible tensor reductions (excluding diagonals). By contrast, when q = 1 (i.e., we are not multiplying by the adjacency matrix), we only have k + 1 possibilities, corresponding to summing S over each of its dimensions. No matter which case we are in, however, and how many contractions Q₁, …, Qj_d′ our network actually computes (in our experiments using second order nodes, we compute 18 different ones), what is important is that the resulting order k tensors satify the P–tensor property.

The reduction step can potentially produce quite a large number of order k tensors, Q₁, …, Q_r. We reduce this number by linearly mixing Q₁, …, Q_r, i.e., taking d′ < r linear combinations of the form

This is again a covariant operation, and the mixing weights are the actual learnable parameters of our neural network.

It is natural to interpret $Q̃(1),…,Q̃(d′)$ as d′ separate channels in neural network terminology. Recall that we also allow the input tensors to have multiple channels, but up to now the corresponding index has been suppressed. The mixing step is the point where these channels are all allowed to interact, so (2) becomes

and the learnable weights of the network at each level form a third order tensor $Wℓ=(wc,c′,j(ℓ))c,c′,j$⁠. The channel index c is not to be confused with c₁, …, c_jk denoting the children of node a. Equation (2) is the main mechanism whereby CCNs are able to learn increasigly complex topological features as we move up the network.

Finally, to get the actual activation of our neuron $na$⁠, we add an additive bias b_ℓ,c and an elementwise nonlinearity σ (specifically, the ReLU operator $σ(x)=max0,x$⁠), as is standard in neural networks. Thus, ultimately, the output of the aggregation function is the collection of P–covariant tensors

with $c∈1,2,…,d′$⁠. As usual, in neural networks, the W_ℓ weight tensors are learned by backpropagation and some form of stochastic gradient descent.

We tested our CCN framework on three different types of datasets that involve learning the properties of molecules from their structure: (a) four relatively small datasets of small molecules that are standard benchmarks in the kernel literature (MUTAG, PTC, NCI1 and NCI109), (b) the Harvard Clean Energy Project,⁴ and (c) QM9.³ The first two types of datasets are pure graph classification/regression tasks. QM9 also has physical (spatial) features, which go somewhat beyond the scope of the present paper. Therefore, on QM9, we conducted separate experiments with and without these physical features [QM9(b) vs. QM9(a)].

In each case, we used second order CCNs (CCN2D) and included the tensor product with the restricted adjacency matrix, as described in Sec. IV B. However, for computational reasons, we only used 18 of the 50 possible contractions. The base features of each vertex were initialized with computed histogram alignment kernel features⁶³ of depth up to 10: each vertex receives a base label $li=concatj=110Hj(i)$ where $Hj(i)∈Rd$ (with d being the total number of distinct discrete node labels) is the vector of relative frequencies of each label for the set of vertices at a distance equal to j from vertex i. The network was chosen to be three levels deep, with the number of output channels at each level fixed to 10.

To run our experiments, we used our own custom-built neural network library called GraphFlow.⁶⁴ Writing our own deep learning software became necessary because at the time we started work on the experiments, none of the standard frameworks such as TensorFlow,⁶⁵ PyTorch,⁶⁶ or MXNet⁶⁷ had efficient support for the type of tensor operations required by CCN. GraphFlow is a fast, general purpose deep learning library that offers automatic differentiation, dynamic computation graphs, multithreading, and Graphics Processing Unit (GPU) support for tensor contractions. In addition to CCN, it also implements other graph neural networks, including neural graph fingerprints,²¹ PSCN,³⁷ and gated graph neural networks.⁴⁰ We also provide a reference implementation of CCN1D and CCN2D in PyTorch.⁶⁸ 

In each experiment, we used 80% of the dataset for training, 10% for validation, and 10% for testing. For the kernel datasets, we performed the experiments on 10 separate training/validation/test stratified splits and averaged the resulting classification accuracies. Our training technique used mini-batch stochastic gradient descent with the Adam optimization method⁶⁹ and a batch size of 64. The initial learning rate was set to 0.001 and decayed linearly after each step toward a minimum of 10⁻⁶.

Our first set of experiments involved three standard “graph kernels” datasets: (1) MUTAG, which is a dataset of 188 mutagenic aromatic and heteroaromatic compounds,⁷⁰ (2) PTC, which consists of 344 chemical compounds that have been tested for positive or negative toxicity in lab rats,⁷¹ and (3) NCI1 and NCI109, which have 4110 and 4127 compounds, respectively, each screened for activity against small cell lung cancer and ovarian cancer lines.⁷² In each of these datasets, each molecule has a discrete label (i.e., toxic/non-toxic and aromatic/heteroaromatic) and the goal is to predict this label. We compare CCN 2D against the graph kernel results reported in Ref. 62 (C-SVM algorithm with the Weisfeiler–Lehman kernel,³³ Weisfeiler–Lehman edge kernel,³³ shortest paths graph kernel,²⁹ graphlet kernel,³¹ and the multiscale Laplacian graph kernel⁶²), neural graph fingerprints²¹ (with up to 3 levels and a hidden size of 10), and the “patchy-SAN” convolutional algorithm (PSCN).³⁷ The results are presented in Table I.

The Harvard Clean Energy Project (HCEP)⁴ dataset consists of 2.3 × 10⁶ organic compounds that are candidates for use in solar cells. The inputs are molecular graphs (derived from their SMILES strings), and the regression target is the power conversion efficiency (PCE). The experiments were ran on a random sample of 50 000 molecules.

On this dataset, we compared CCN to the following algorithms: Lasso, ridge regression, random forests, Gradient Boosted Trees (GBT), the optimal assignment Wesifeiler–Lehman graph kernel,³³ neural graph fingerprints,²¹ and PSCN.³⁷ For the first four of these baseline methods, we created simple feature vectors from each molecule: the number of bonds of each type (i.e., number of H—H bonds, number of C—O bonds, etc.) and the number of atoms of each type. Molecular graph fingerprints use atom labels of each vertex as base features. For ridge regression and Lasso, we cross-validated over λ. For random forests and GBT, we used 400 trees and cross validated over maximum depth, minimum samples for a leaf, minimum samples to split a node, and learning rate ( for GBT). For neural graph fingerprints, we used up to 3 layers and a hidden layer size of 10. In PSCN, we used a patch size of 10 with two convolutional layers and a dense layer on top as described in their paper (Tables II–IV).

QM9 has recently emerged as a molecular dataset of significant interest. QM9 contains ∼134 K organic molecules with up to nine atoms (C, H, O, N, and F) out of the GDB-17 universe of molecules. Each molecule is described by a SMILES string, which is converted to the molecular graph. The molecular properties are then calculated using DFT at the level of either B3LYP or 6-31G(2df, p), returning the spatial configurations of each atom along with thirteen molecular properties:

• U₀: atomization energy at 0 K (eV),

• U: atomization at room temperature (eV),

• H: enthalpy of atomization at room temperature (eV),

• G: free energy of atomization (eV),

• ω₁: highest fundamental vibrational frequency (cm⁻¹),

• ZPVE: zero point vibrational energy (eV),

• HOMO: highest occupied molecular orbital, energy of the highest occupied electronic state (eV),

• LUMO: lowest unoccupied molecular orbital, energy of the lowest unoccupied electronic state (eV),

• GAP: difference between HOMO and LUMO (eV),

• R²: electronic spatial extent (bohr²),

• μ: norm of the dipole moment (D),

• α: norm of the static polarizability (bohr³),

• $Cv$⁠: heat capacity at room temperature (cal/mol/K).

We performed two experiments on the QM9 dataset, with the goal of providing a benchmark of CCN as a graph learning framework and demonstrating that our framework can predict molecular properties to the same level as DFT. In both cases, we trained our system on each of the thirteen target properties of QM9 independently and report the MAE for each.

We use the QM9 dataset to benchmark the CCN architecture against the Weisfeiler–Lehman graph kernel, neural graph fingerprints, and PSCN. For this test, we consider only heavy atoms and exclude hydrogen. The CCN architecture is as described above, and settings for NGF and PSCN are as described for HCEP.

To compare the DFT error, we performed a test of the QM9 dataset with each molecule including hydrogen atoms. We used both physical atomic information (vertex features) and bond information (edge features) including the atom type, atomic number, acceptor, donor, aromatic, hybridization, number of hydrogens, Euclidean distance, and Coulomb distance between pairs of atoms. All the information is encoded in a vectorized format. Our physical features were taken directly from the dataset used in Ref. 25 without any special feature engineering.

To include the edge features into our model along with the vertex features, we used the concept of line graphs from graph theory. We constructed the line graph for each molecular graph in such a way that an edge of the molecular graph corresponds to a vertex in its line graph, and if two edges in the molecular graph share a common vertex, then there is an edge between the two corresponding vertices in the line graph. (See Fig. 6.) The edge features become vertex features in the line graph. The inputs of our model contain both the molecular graph and its line graph. The feature vectors F_ℓ between the two graphs are merged at each level ℓ.

We present results for both the CCN 1D and CCN 2D architectures. For CCN 1D, our network is seven layers with 64 input channels; at the first and second layer, the number of channels is halved, and beyond that each layer has 16 channels. For the CCN 2D architecture, we used three layers, with 32 channels at the input and 16 at the remaining layers. We report the mean average error for each learning target in its corresponding physical unit and compare it against the DFT error given in Ref. 25.

Overall, our CCN outperforms the other algorithms on a significant fraction of the experiments we implemented. On the subsampled HCEP dataset, CCN outperforms all other methods by a very large margin. For the graph kernel datasets, the SVM with the Weisfeiler–Lehman kernels achieve the highest accuracy on NCI1 and NCI109, while CCN wins on MUTAG and PTC. Perhaps this poor performance is to be expected since the datasets are small and neural networks usually require tens of thousands of training examples to be effective. Indeed, neural graph fingerprints and PSCN also perform poorly compared to the Weisfeiler–Lehman kernels. In the QM9(a) experiment, CCN obtains better results than the three other graph learning algorithms on all 13 learning targets.

In the QM9(b) experiment, the error of CCN is smaller than that of DFT itself on 11 of 12 learning targets (Ref. 25 does not have DFT errors for R2). However, other recent studies^14,25 have obtained even stronger results. Looking at our results, we find that values depending strongly on position information, such as the dipole moment and average electronic spatial extent, are predicted poorly when we include physical features. By contrast, properties that are not expected to strongly depend on spatial extent are predicted significantly better. This suggests that our spatial input features were not fully exploited and that feature engineering position information could significantly enhance the power of our CCN.

Our custom deep learning library⁶⁴ enabled all the above results to be obtained reasonably efficiently. The prediction time for CCN 1D and CCN 2D on QM9(b) comes out to 6.0 ms/molecule and 7.2 ms/molecule, respectively, making it possible to search through a million candidate molecules in less than 2 h.

In this paper, we presented a general framework called covariant compositional networks (CCNs) for learning the properties of molecules from their graphs. Central to this framework are two key ideas: (1) a compositional structure that generalizes message passing neural networks (MPNNs) and (2) the concept of covariant aggregation functions based on tensor algebra.

We argue that CCNs can extract multiscale structures from molecular graphs and keep track of the local topology in a manner the MPNNs are not able to. We also introduced the GraphFlow software library that provides an efficient implementation of CCNs. Using GraphFlow, we were able to show that CCNs often outperform existing state-of-the-art algorithms in learning molecular properties.

The authors would like to thank the Institute for Pure and Applied Mathematics and the participants of its “Understanding Many Particle Systems with Machine Learning” program for inspiring this research. Our work was supported by DARPA Young Faculty Award No. D16AP00112 and used computing resources provided by the University of Chicago Research Computing Center.