Chemistry Guided Molecular Graph Transformer
Abstract
Classic methods to calculate molecular properties are insufficient for large amounts of data.
The Transformer architecture has achieved competitive performance on graph-level prediction by introducing general graphic embedding.
However, the direct spatial encoding strategy ignores important inductive bias for molecular graphs, such as aromaticity and interatomic forces.
In this paper, inspired by the intrinsic properties of chemical molecules, we propose a chemistry-guided molecular graph Transformer.
Specifically, motif-based spatial embedding and distance-guided multi-scale self-attention for graph Transformer are proposed to predict molecular property effectively.
To evaluate the proposed methods, we have conducted experiments on two large molecular property prediction datasets,
ZINC, and PCQM4M-LSC.
The results show that our methods achieve superior performance compared to various state-of-the-art methods.
Codes are available at https://github.com/PSacfc/chemistry-graph-transformer .
Introduction
Molecular properties prediction is one of the most critical tasks and plays a vital role in many downstream applications, from material discovery to drug design.
Classic techniques, such as Density
Functional Theory (DFT), are computationally too expensive to deal with large amounts of data.
As molecules and graphs share similar structures, it is natural to model molecules as graphs for representation learning in deep neural networks.
In particular, graph neural networks (GNNs) [1] have achieved considerable progress in chemistry calculation [2].
However, standard GNNs in the neighborhood aggregation paradigm suffer from limited discriminative power in distinguishing high-order graph structures as opposed to some low-order motif.
Moreover, when GNNs go deeper, the over-smoothing [3] problem gets worse for the message passing paradigm.
Transformer [4] has achieved remarkable progress in multiple domains [5, 6, 7] and could become a new powerful workhorse for graph representation learning.
Recently, Graphormer [8] added a set of graph spatial-related encodings into the vanilla Transformer and achieved impressive performance on graph-level prediction tasks.
EGT [9] added a dedicated pathway for pairwise structural information and proposed Edge-augmented Graph Transformer.
Although Graphormer has a clear advantage of generality, spatial encoding strategy ignores some most important molecular properties that commonly persist in nature, like aromaticity [10] and interatomic forces [11].
Aromaticity [12] is a property of cyclic molecular structure that gives highly increased stability and completely different chemical properties.
Aromatic rings exhibit robust stabilization and connectivity with their delocalized pi-electron around the rings.
As shown in Fig.
1, fluorine atoms (F) on C6F5OH are strongly influencing the hydroxy group (OH) through the aromatic ring’s pi-electrons, causing
C6F5OH to be one of the most acidic phenols with pK a = 5 . 5 , where pK a is the acid dissociation [Figure] [Figure]
constant.
That lower value of the acid dissociation constant represents strong acid.
However, for
C4F9OH with more fluorine atoms, the pK a is only 7 . 05 through regular bonds, and the connectivity strength is simply determined by spatial distance.
Inspired by this phenomenon, we propose a motif-based shortest path distance (SPD) embedding following the inductive bias as the interaction between atoms in a single aromatic ring is vital.
Interatomic forces are composed of attraction and repulsion.
Interatomic forces can also be described as potential energy.
Lennard-Jones potential [13] is the simplest interatomic interaction model widely used:
[Equation]
where r is the distance between atoms.
As shown in Fig.
2, different distance scales lead to different kinds of forces.
Two atoms repel each other at a close distance and attract each other at a further distance.
Thus, motivated by multi-scale self-attention [14], we propose a distance-guided multi-scale self-attention for molecular graphs.
In this paper, we propose a chemistry-guided molecular graph Transformer.
Inspired by Graphormer, we take nodes in a graph as tokens fed into Transformer and propose two components to help represent molecule graphs more naturally with chemical intuition.
Firstly, we redefine distance embedding related to connectivity and propose motif-based spatial embedding to explicit model molecular aromatic.
Secondly, we propose a distance-guided multi-scale self-attention.
Only node pairs under specific distance scales participate in attention bias calculation for each Transformer layer.
We have conducted extensive experiments on molecular property prediction datasets, including ZINC [15] and PCQM4M-LSC [16].
The results show that our methods achieve superior performance compared to state-of-the-art methods.
Preliminaries
Notations An undirected graph can be denoted as G =
( V, E ) where V = { v 1 , v 2 , ..., v n } is the node set, n = | V | is the number of nodes and E is the edge set.
Let X
= { x 1 , x 2 , ..., x n } be the node feature set where x i is the feature associated with node v i .
Transformers Multi-head self-attention is the main component of Transformer.
Let H ∈ R n × d denote the input, where d is the hidden dimension.
The self-attention module projects the input H into three matrices Q, K, V by three parameter matrices W
[Equation]
Q
[Equation] ∈ [Equation] × d K , W
[Equation]
∈ R
[Equation] × [Equation] V .
These three matrices are projected back to each position by W O ∈ R hd V ×
d where h is the number of heads.
The multi-head self-attention can be formalized as follows:
[Equation]
[Figure]
Apart from the self-attention module, a position-wise feed-forward network (FFN) is also one of the key components in Transformer.
Let H l be the hidden state of the l th layer, H l +1 can be calculated as follow:
[Equation]
where LayerNorm( · ) represents layer normalization in [17].
Methods
In this section, we redefine distance according to connectivity and propose motif-based spatial embedding as the bias of self-attention in each layer.
Then, we propose distance-guided multi-scale self-attention in each token that pays attention to different distance scales in different layers.
In addition to these methods, we utilize centrality encoding, and edge encoding in Graphormer [8] for augmenting structural information.
3.1 Motif-Based Spatial Embedding
We suppose the function distance between node pairs within a strong connectivity component is smaller than its counterpart between node pairs among unconnected components.
It is guided by the vital concept of aromaticity in chemistry.
In chemistry, aromaticity is a property of cyclic and planar molecular structures with pi bonds in resonance that gives stronger stability than saturated compounds and other geometric non-cyclic arrangments with the same set of atoms.
The delocalized electrons will significantly impact the molecular property as a whole other than on the individual atom that consists of the molecule.
Because most aromatic components in molecular graphs are rings, we consider using a bi-connected component (BCC) [18] of an undirected graph to label such inductive bias.
The bi-connected component is a sub-graph in which any cut of edges has no influence on whether it is connected.
We use a shared shortened distance embedding to model the effect of delocalized electrons on the individual atoms of aromatic rings.
As in Fig. 3, the benzene ring is regarded as a bi-connected component.
The distance between carbons of methyl groups reduces from 4 to 2.
Firstly, for graph G ( V, E ) , bi-connected components are found by [18].
Then, all nodes in a single biconnected component are treated as a single node to construct a new acyclic graph G bcc
( V bcc , E bcc ) and dis bcc ( s, t ) represents the SPD between v bccs and v bcct .
Thus, for nodes v i ∈ v bccs and v j ∈ v bcct , the motif-based SPD dis motif ( i, j ) can be calculated in Eq.
8 as follows:
[Equation]
Finally, for self-attention of h th head A h , motif-based spaial embedding utilizes motif-based SPD to calculate embedding and serves as a bias term:
[Equation]
[Equation] [Figure] [Figure]
Structural Debilitating Problem For data with solid connectivity, structural information debilitating problem is introduced.
To analyze the problem, we split ZINC (the complete details are provided in the section of experiments) by the diameter of the motif-based graph.
The diameter of the graph means the maximum SPD of all node pairs.
As results shown in Fig.
4, the performance in small diameter is relatively poor.
3.2 Distance Guided Multi-Scale Self-Attention
To eliminate structural debilitating and utilize interatomic forces, we introduce multi-scale selfattention.
However, unlike NLP tasks, the distance between graph nodes is not regular as tokens in a sequential language signal.
Thus, multi-scale self-attention [14] is not directly applicable to graph-related tasks.
Distance-guided multi-scale attention is then proposed to tackle this problem, and heads in different layers work on various scales.
As shown in Fig.
5, the interaction between weakly connected pairs with farther distances is ignored.
Thus, for a single head, we remove distance-related terms in the attention matrix when the distance between nodes exceeds ω multi :
[Equation]
Since scale selection is essential to capture hierarchical information, we design a Fibonacci scale distribution from the intuitive graph perspective.
Let L ( L ≥ 2) denote the number of total layers and the distance scale of l th layer ω multil is computed as follow:
[Equation]
where α (0 ≤ α ≤ 2) , S and k are hyper-parameters.
Here α controls the scale distribution, S controls the maximum scale size, and k represents that two adjacent k layers share the scale size.
Experiments
In this section, we conduct experiments on two molecular property prediction datasets, ZINC [15] and
OGB-LSC [16].
Then, we provide ablation studies to demonstrate the significance of each proposed component.
Finally, we analyze that distance-guided multi-scale self-attention effectively eases the debilitating structural problem.
Datasets ZINC is one of the most popular real-world molecular datasets of 250K graphs.
Similar to benchmark [19], we choose the subset (12K) to regress molecular constrained solubility.
In KDD Cup 2021, OGB-LSC was provided to the community to encourage the development of state-of-the-art graph ML models for sizeable molecular graph datasets.
The graph regression challenge of OGB-LSC,
PCQM4M-LSC, is a quantum chemistry dataset containing 3 .
8 M molecular graphs.
Specifically, the task is to predict the HOMO-LUMO energy gap of molecules given their 2D molecular graphs.
4.1 Graph Representation
Baseline We benchmark our model with various GNNs such as GCN [20], GAT [21], GIN [22] and
MoleculeX [23].
In addition, we compare the recently proposed random walk-based method CRAWL [24].
Especially, Transformer-based model SAN [25] , GT[26] and the challenge winning solution
Graphormer [8], is also compared.
Settings For both training procedures, we use Adam as the optimizer and adopt several warm-up steps followed by a linear decay learning rate scheduler.
For the ZINC dataset, the large model is not encouraged.
Thus parameters of the model are limited to under 500 K .
And we conduct experiments with training procedures under 4 different seeds.
The detailed experimental settings are shown in
Table 1.
[Table] [Table]
[Table]
Results Table 2 shows the performance on ZINC.
It can be observed that our model surpasses most models and approximates CRAWL with lower variance.
Table 3 also summarizes performance on
PCQM4M-LSC and our model performs better than Graphormer [8] with similar parameter quantity.
4.2 Ablation Study
We perform a series of ablation studies on the ZINC dataset to illustrate each component’s effect in the proposed methods.
The results are shown in Table 4.
Spatial Embedding We compare atom-based spatial embedding to the proposed motif-based spatial embedding.
Although motif-based SPD introduces structural debilitating, it outperforms the vanilla atom-based SPD.
This phenomenon shows that motif-based SPD is suitable for molecular graphs, and aromatic connectivity is vital for distance definition.
Multi-Scale Self-Attention
To illustrate the importance of distance-guided multi-scale information, we compare distance-guided multi-scale self-attention with complete scales in all layers.
The result shows that graph Transformer architecture with multi-scale self-attention yields a large margin of performance improvement.
[Table]
4.3 Structural Debilitating Problem
We analyze the debilitating structural problem discussed above.
We conduct ablation experiments with models with and without multi-scale self-attention on different diameter graphs.
The results in
Table 5, show that multi-scale self-attention significantly moderates the inconsistent performance on various diameter graphs.
Especially for shorter distances, with relatively more vital interatomic forces, multi-scale self-attention achieves improvement more significantly.
Limitations
Although our methods achieve considerable progress, there are still some limitations.
Firstly, we merely treat all kinds of rings as bi-connected components.
The chemical properties of aromatic rings differ significantly from those of regular rings. Some sophisticated designs involving chemical bond
[Table] [Table]
properties as filtering signals could help us separate those two cases and further improve our results.
Secondly, in our proposed methods, both edges represent the distance of a single unit, regardless of their chemical bond type.
Specifically, multi-scale self-attention is motivated by interatomic forces.
However, interatomic distance is related to bond length, bond angle, and dihedral angle [30], which we did not introduce into our method.
Conclusions
This paper proposes a chemistry-guided transformer for the molecular graph forecasting task, which achieves state-of-art performance.
It can be seen as an example of how simple chemical intuitions can significantly help us in chemical property forecasting tasks.
To add chemistry-inspired inductive bias into the graph Transformer, we propose a motif-based spatial embedding to represent aromaticityrelated connectivity.
Moreover, we introduce a distance-guided multi-scale self-attention to consider interatomic forces and eliminate the debilitating structural problem.
Lastly, experiments show that the proposed method considerably improves on two large benchmark datasets compared with state-of-the-art algorithms.