Neural message passing on high order paths

There are many factors pertaining to molecular graphs that motivate the development of our model. In this section we discuss, in more depth, the limitations of GNNs with respect to specific aspects of molecules that motivate our model. These include molecular substructures like rings and functional groups, molecular geometry as characterized by internal coordinates as well as stereochemistry.

Molecular substructures play an important role in determining molecular properties for example functional groups influence the chemical reactions a molecule undergoes. By only aggregating over neighbours, GNNs cannot learn about these larger substructures in a single propagation layer. On the other hand, by passing messages over larger neighbourhoods, in every layer we could directly learn about these structures. Furthermore, We could directly indicate if the path that a message is travelling on contains a simple functional group like alcohol (ROH) or passes through a larger functional group. For example, atoms in the neighbourhood of a functional group could receive a message along a path indicating a functional group is in the neighbourhood.

Molecular geometry is the three dimensional arrangement of atoms in a molecule and influences several properties, including the reactivity, polarity and biological activity of the molecule. An important application of GNNs is predicting quantum mechanical properties of molecules, which are heavily dependent on the geometry of the molecule. The 3D configuration of a molecule can be fully specified by (1) bond lengths—the distance between two bonded atoms, (2) bond angles—the angle formed between three neighbouring atoms, and (3) dihedral angles between four consecutive atoms. In fact the potential energy is typically modelled as a sum of terms involving each of these three. Current GNN approaches to quantum chemistry incorporate neighbouring geometry by using bond distances as edge features [3], but do not directly account for the relative orientation of neighbouring atoms and bonds—a framework that could do so would be advantageous.

Stereochemistry involves the relative spatial arrangement of atoms in molecules, specifically, stereoisomers—which are molecules with the same discrete graph but different three-dimensional orientation of atoms. For example, enantiomers—non-superimposable mirror images of molecules and cis-trans isomers, that only differ through the rotation of a functional group. Even if they use interatomic distances as edge features, GNNs will have limited ability to distinguish stereoisomers, since these molecules only differ through the relative orientation of atoms. In general, at every propagation step, GNNs should learn representations over each node's extended neighbourhood to encode the relationships between nodes in that neighbourhood.

We generalize message passing neural networks (MPNNs) to aggregate across larger neighbourhoods by passing messages along simple paths of higher order neighbours. We describe the general framework in section 3. We experiment with various molecular property prediction task and a node classification task in citation networks. Our specific contributions are two-fold:

• We devise a simple extension to any message passing neural network to learn representations over larger node neighbourhoods within each propagation layer by simply augmenting the message function to aggregate over additional neighbours.
• By summing over additional neighbours we enable the use of path features such as bond angles for paths of length two and dihedral angles for paths of length three and thus encoding the full molecular geometry and orientation, so that MPNNs can distinguish isomers.

Message passing neural networks operate on graphs G with n nodes each with feature vector $\mathbf{x}_v \in \mathbb{R} ^f$ that specify what kind of atom the node is, among other possible features. There are n × n edge feature vectors $\boldsymbol{e}_{vw}\in \mathbb{R} ^e$ that specify what kind of bond type atoms v, w have. The forward pass has two phases, a message passing phase and a readout phase.

The message passing phase runs for T propagation steps and is defined in terms of message functions M_t and node update functions U_t . During the message passing phase, hidden states $\boldsymbol{h}^t_v$ at each node in the graph are updated based on messages $\boldsymbol{m}^{t+1}_v$ according to:

$$\begin{align*} \boldsymbol{m} _v ^{t+1} = \sum _{w\in \mathcal{N} _v } M_t(\boldsymbol{h}_v ^t, \boldsymbol{h}_w^t, \boldsymbol{e}_{vw}) , \ \ \ \boldsymbol{h}^{t+1}_v = U_t(\boldsymbol{h}^{t}_v, \boldsymbol{m}^{t+1}_v), \ \ \ \mathbf{y} = \texttt{Readout}(\{\boldsymbol{h}^T _v, \boldsymbol{x} _v \}_{v\in G}). \end{align*}$$

We have $\boldsymbol{x}_v = \boldsymbol{h}_v^{0}$. The message node v receives aggregates over its neighbours $\mathcal{N}_v$, in this case, by simple summation. We then readout predictions $\mathbf{y}$ based on final node embeddings.

The first graph neural network model was proposed by [6] and many variants have been recently proposed [5, 7, 8]. Our focus is on the general framework of neural message passing from [3]. We review relevant GNN models and their use in Molecular Deep learning in this section.

2.2.1. Molecular deep learning

2.2.2. Higher order GNNs

2.2.3. Path augmented transformer

2.2.4. Structured transformer

2.2.5. 3D GCN

2.2.6. Directional message passing

For graphs with a large number of nodes and edges, passing messages along paths becomes very expensive and, as in GraphSage [2], sampling a subset of paths of higher order neighbours is necessary. However, for molecules, where the number of neighbours is usually ${\leqslant}4$ this is not necessary. Furthermore, one can include domain specific path features in the message function. We describe two examples of these path features below.

3.1.1. Molecular substructures

We can incorporate whether the path travels through a molecular substructure by considering paths of at least length 2, where we have a message function that sums over two neighbouring atoms v → w → y. Along with their node and edge features, the possible path features include ring features—ie one hot indication if any atoms are in (specific) rings as well as if the path is a functional group (ROH) or within a larger functional group:

$$\begin{align} \boldsymbol{m} _v ^{t+1} = \sum _{w\in \mathcal{N} _v } \sum _{\substack{y \in \mathcal{N}_w \\ y\neq v}} M_t( \boldsymbol{h}_{v}^t, \boldsymbol{p}_{v\to y}) \ \ \ \ \ \boldsymbol{p}_{v\to y} = \begin{bmatrix} \boldsymbol{h}_w^t & \boldsymbol{h}_y ^t \\ \boldsymbol{e}_{vw} & \boldsymbol{e}_{wy} \end{bmatrix}. \end{align} \tag{ 2 }$$

3.1.2. Molecular geometry

Considering paths of length 3, where we have a message function that sums over three neighbouring atoms v → w → y → x. Along paths of length three additional features include two bond angles $\boldsymbol{\alpha}_{vwy} \& \ \boldsymbol{\alpha}_{wyx}$ and the dihedral angle $\boldsymbol{\varphi}_{vwyx}$ between the planes defined by the pairs of atoms (v, w) and (y, x). Effectively, messages passed over three consecutive neighbours contain information about the entire molecular geometry (see figure 2):

$$\begin{align} \boldsymbol{m}_v ^{t+1} = \sum _{w\in \mathcal{N} _v } \sum _{\substack{y \in \mathcal{N}_w \\ y\neq v}} \sum _{\substack{x \in \mathcal{N}_y \\ x \neq w,v }} M_t( \boldsymbol{h}_{v}^t, \boldsymbol{p}_{v\to x}) \ \ \ \ \ \boldsymbol{p}_{v\to x} = \begin{bmatrix} \boldsymbol{h}_w^t & \boldsymbol{h}_y ^t& \boldsymbol{h}_x ^t \\ \boldsymbol{e}_{vw} & \boldsymbol{e}_{wy}& \boldsymbol{e}_{yx} \\ \boldsymbol{\alpha}_{vwy}&\boldsymbol{\alpha}_{wyx}& \boldsymbol{\varphi}_{vwyx} \end{bmatrix}. \end{align} \tag{ 3 }$$

Zoom In Zoom Out Reset image size

We compare the performance of our model against a few baselines on a variety of molecular property prediction tasks involving different datasets of undirected molecular graphs with different sizes and distributions. One with moderately large graphs with 6–8 rings and two with much smaller molecules and 1–2 rings. Figure 3 has examples of molecules from the datasets in these tasks, they include:

• ESOL: Delaney [19] predicting the aqueous solubility of 1144 molecules.
• QM8: Ruddigkeit et al [20] predicting 16 electronic spectra values calculated using density functional theory for 21 786 organic molecules.
• CEP: the photovoltaic efficiency of 20 000 organic molecules from The Harvard Clean Energy Project [21].

Zoom In Zoom Out Reset image size

For QM8 we use the Path MPNN with message given by equation (3) since molecular geometry is very important for the targets in QM8, the path features such as bond lengths, bond angles and dihedrals are computed on the fly using rdkit [22] given supplied atomic coordinates in the data. We use the message function from the interaction networks model in [23] $\boldsymbol{m}_{v\to x} = \tau ( \texttt{Concat}[\boldsymbol{p}_{v\to x}] )$ that passes the concatenated path features through a single layer neural net with relu activation and then uses graph attention [7] to aggregate incoming messages passed over paths:

$$\begin{align} M_t(\boldsymbol{h}_v ^t, \boldsymbol{p}_{v\to x}) = a_{v\to x} \boldsymbol{m}_{v\to x} \ \ \ \textrm{where } \ \ \ a_{v\to x} = \frac{e^{\boldsymbol{a}^\top \boldsymbol{m}_{v\to x}}}{\sum_{v\to x } e^{\boldsymbol{a}^\top \boldsymbol{m}_{v\to x}}}, \end{align} \tag{ 4 }$$

where $\mathbf{a}$ is the attention weight vector and the summation is over the simple path v → x which is the same triple sum in (3). For ESOL and CEP we pass messages over paths of length two as specified by equation (2) in order to use path features for molecular substructures (which are important for the targets in these datasets as well), these features also computed using rdkit. For both models, the node update function concatenates incoming messages with the current node state and feeds it through a dense layer $U_t = \sigma ( [\boldsymbol{h}_v^t, \boldsymbol{m}_v^{t+1} ])$. After propagation through message passing layers, we use the set2set model [24] in the same way as [3] as the readout function to combine the node hidden features into a fix-sized hidden vector.

For all models, we only use atom and bond type/distance as one-of-k hot encodings for for initial node and edge features $\boldsymbol{x}_v, \boldsymbol{e}_{uv}$, in addition to previously described path features. The models are trained using root mean squared error (RMSE) for loss. Model evaluation is done using mean absolute error (MAE) of the molecular properties in the QM8 dataset, RMSE for ESOL and percent for CEP. We use a 80%–10%–10% train, validation, test set split. We perform three runs with different randomizations and report the mean performance and standard deviation. We do not perform cross-validation.

The model was coded in pytorch [25] from scratch and trained on a GeForce RTX 2070 GPU which has 8 GB of ram. The Path MPNN model is memory intensive and the largest batch size we could use was 32 for the smaller dataset (QM8, ESOL) with 16 for the larger molecular graphs (CEP). Minimal hyperparameter optimization was needed for the path model to achieve superior performance. We used learning rate of 10⁻⁴ with the adam optimizer [26] training for 300 epochs for both the Path MPNN and baseline MPNN. We experimented with 1–3 path MPNN layers and 3–6 MPNN layers in the baseline but we found no improvement for increasing number of layers in both. For both the baseline MPNN and Path MPNN we used 128 hidden units for the node representation as well as in the message function's one layer NN with a ReLu non-linearity. In addition we use 1024 hidden units for the set2vec model with five unroll steps for both the Path MPNN and baseline MPNN. We only optimized hidden unit number and number of message passing layers along with learning rate.

We use the top performing model from Molecule Net [18] (Molnet) for each dataset. We also benchmark with the differentiable version of circular fingerprints from [1] (neural fingerprints). To highlight the importance of path features, we also compared the performance of a standard Message Passing neural net that uses three message passing layers. This uses the same message function, node update function and readout function as the Path model (except path length is one). The last benchmark is the path-augmented graph transformer network since this model is similarly built to model longer-range dependencies in molecular graphs. As can be seen in table 1, for QM8, ESOL and CEP, passing messages over paths leads to a substantial improvement in predictive accuracy.

We use the experimental setup of [8]. We sum over paths of length 3 while uniformly sampling a single second order and third order neighbour. Our base MPNN is a GCN [8] that has message function:

$$\begin{align} \boldsymbol{m} _v ^{t+1} = \sum _{w\in \mathcal{N} _v } \hat{\mathbf{A}}_{vw} \boldsymbol{h} ^t_v , \ \ U_t = \sigma ( \boldsymbol{m}_v^{t+1} ), \end{align} \tag{ 5 }$$

where $\hat{\mathbf{A}}$ is the adjacency matrix and σ is a dense layer with sigmoid activation. There is no readout function necessary for pooling and a softmax layer maps the node representations to the prediction. For a citation network the path features are just the node features and edge features connecting v to nodes that are $\ell$ nodes away, i.e.

$$\begin{align*} \boldsymbol{p}_{v_1 \to v_\ell} = \{\boldsymbol{h}_{v_1}^t, \boldsymbol{e}_{vv_1}, \dots , \boldsymbol{h}_{v_\ell}^t, \boldsymbol{e}_{v_{\ell}v_{\ell -1}} \}. \end{align*}$$

We compare with two other higher order GCN variants: Mixhop [9] and PAN [10]: Path integral graph convolution—both use powers of the adjacency to aggregate GCN layers of higher order neighbours. The results are displayed in table 2 and our model achieve similar accuracy to our baselines.

In this work we only considered very simple message functions, in general, it is not straight forward to construct message function over paths. For example, the message function from [3], maps edge features to a square matrix using a neural net—incorporating more neighbours and their edge and path features into this kind of message function introduces many design challenges. There are many graph types where the path MPNN would be roughly as useful as a standard MPNN, including most non chemical graph data. In general, we found no improvement in accuracy in passing messages over longer paths than relevant to the desired path features for example. The path MPNN has a large capacity and is very susceptible to overfitting, by using a fairly low learning rate we were able to overcome this and see comparable performance on training, validation and test sets. In further we hope to find methods to reduce the memory requirement of the Path model.

For standard MPNNs, a single step of the message passing phase for a dense graph requires $\mathcal{O}(n^2d^2)$ floating point multiplications. Naturally for the Path models this is multiplied by the path length $\ell$ but is still roughly equivalent to $\ell$ propagation steps for the standard MPNN while having increased performance benefits.

We introduce a general GNN framework based on message passing over simple paths of higher order neighbours. This allows us to use path features in addition to node and edge features, which is very useful in molecular graphs, as many informative features are characterized by the paths between atoms. We benchmarked our framework on molecular property prediction tasks and a node classification task in citation networks.