Chemistry Guided Molecular Graph Transformer
Anonymous Author(s)
AffiliationAddress email
Abstract
Classic methods to calculate molecular properties are insufficient for large amountsof data.
The Transformer architecture has achieved competitive performance ongraph-level prediction by introducing general graphic embedding.
However, thedirect spatial encoding strategy ignores important inductive bias for moleculargraphs, such as aromaticity and interatomic forces.
In this paper, inspired by the intrinsic properties of chemical molecules, we propose a chemistry-guided moleculargraph Transformer.
Specifically, motif-based spatial embedding and distanceguided multi-scale self-attention for graph Transformer are proposed to predictmolecular 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 performancecompared to various state-of-the-art methods.
Codes are available at https://anonymous.4open.science/r/chemistry-graph-transformer-451E/ .
Introduction
Molecular properties prediction is one of the most critical tasks and plays a vital role in manydownstream 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 forrepresentation learning in deep neural networks.
In particular, graph neural networks (GNNs)
[1]have achieved considerable progress in chemistry calculation [2].
However, standard GNNs inthe neighborhood aggregation paradigm suffer from limited discriminative power in distinguishinghigh-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 anew powerful workhorse for graph representation learning.
Recently, Graphormer [8] added a set ofgraph spatial-related encodings into the vanilla Transformer and achieved impressive performanceon graph-level prediction tasks.
Although it has a clear advantage of generality, spatial encodingstrategy ignores some most important molecular properties that commonly persist in nature, likearomaticity [9] and interatomic forces [10].
Aromaticity [11] is a property of cyclic molecular structure that gives highly increased stability andcompletely different chemical properties.
Aromatic rings exhibit robust stabilization and connectivitywith their delocalized pi-electron around the rings.
As shown in Fig.
1, fluorine atoms (F) on C6F5OHare 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 dissociationconstant.
That lower value of the acid dissociation constant represents strong acid.
However, for
[Figure] [Figure] C4F9OH with more fluorine atoms, the pK a is only 7 . 05 through regular bonds, and the connectivitystrength is simply determined by spatial distance.
Inspired by this phenomenon, we propose amotif-based shortest path distance (SPD) embedding following the inductive bias as the interactionbetween atoms in a single aromatic ring is vital.
Interatomic forces are composed of attraction and repulsion.
Interatomic forces can also be describedas potential energy.
Lennard-Jones potential [12] is the simplest interatomic interaction model widelyused:
[Equation]
where r is the distance between atoms.
As shown in Fig.
2, different distance scales lead to differentkinds of forces.
Two atoms repel each other at a close distance and attract each other at a furtherdistance.
Thus, motivated by multi-scale self-attention [13], we propose a distance-guided multi-scaleself-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 helprepresent molecule graphs more naturally with chemical intuition.
Firstly, we redefine distanceembedding related to connectivity and propose motif-based spatial embedding to explicit modelmolecular aromatic.
Secondly, we propose a distance-guided multi-scale self-attention.
Only nodepairs 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[14] and PCQM4M-LSC [15].
The results show that our methods achieve superior performancecompared 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 thenode set, n = |
V | is the number of nodes and E is the edge set.
Let X
= { x 1 , x 2 , ..., x n } be thenode 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 × ddenote the input, where d is the hidden dimension.
The self-attention module projects the input H intothree matrices Q, K, V by three parameter matrices W Q
[Equation]
× d Q , W K ∈ R d × d K , W V ∈ R d ×
[Equation]
V . These three matrices are projected back to each position by W O ∈ R hd V ×
d where h is the number ofheads.
The multi-head self-attention can be formalized as follows:
[Equation]
[Equation]
[Figure]
[Equation]
Apart from the self-attention module, a position-wise feed-forward network (FFN) is also one of thekey components in Transformer.
Let H l be the hidden state of the l th layer, H l +1 can be calculatedas follow:
[Equation]
where LayerNorm( · ) represents layer normalization in [16].
Methods
In this section, we redefine distance according to connectivity and propose motif-based spatialembedding as the bias of self-attention in each layer.
Then, we propose distance-guided multi-scaleself-attention in each token that pays attention to different distance scales in different layers.
Inaddition to these methods, we utilize centrality encoding, and edge encoding in Graphormer [8] foraugmenting structural information.
3.1 Motif-Based Spacial Embedding
We suppose the function distance between node pairs within a strong connectivity component issmaller than its counterpart between node pairs among unconnected components.
It is guided bythe vital concept of aromaticity in chemistry.
In chemistry, aromaticity is a property of cyclic andplanar molecular structures with pi bonds in resonance that gives stronger stability than saturatedcompounds and other geometric non-cyclic arrangments with the same set of atoms.
The delocalizedelectrons will significantly impact the molecular property as a whole other than on the individualatom that consists of the molecule.
Because most aromatic components in molecular graphs are rings,we consider using a bi-connected component [17] 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 whetherit is connected.
We use a shared shortened distance embedding to model the effect of delocalizedelectrons on the individual atoms of aromatic rings.
As in Fig. 3, the benzene ring is regarded as abi-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 [17].
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]
[Figure] [Figure] Finally, for self-attention of h th head A h , motif-based spacial embedding utilizes motif-based SPDto calculate embedding and serves as a bias term:
[Equation]
motif ( i, j )) .
Structural Debilitating Problem For data with solid connectivity, structural information debilitatingproblem is introduced.
To analyze the problem, we split ZINC (the complete details are providedin the section of experiments) by the diameter of the motif-based graph.
The diameter of the graphmeans the maximum SPD of all node pairs.
As results shown in Fig.
4, the performance in smalldiameter is relatively poor.
3.2 Distance Guided Multi-Scale Self-Attention
[Figure]
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 tokensin a sequential language signal.
Thus, multi-scale self-attention [13] is not directly applicable tograph-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 distancebetween nodes exceeds ω multi :
[Equation]
Since scale selection is essential to capture hierarchical information, we design a Fibonacci scaledistribution from the intuitive graph perspective.
Let L ( L ≥ 2) denote the number of total layers andthe distance scale of l th layer ω multil is computed as follow:
[Equation]
[-z l = (cid:40) α · z l +
 + (2 − α ) · z l +2 , l/k < L/k −2 , l/k = L/k −1 , l/k-][Equation]
where α (0 ≤ α ≤ 2) , S and k are hyper-parameters.
Here α controls the scale distribution, Scontrols 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 [14] and
OGB-LSC [15].
Then, we provide ablation studies to demonstrate the significance of each proposedcomponent.
Finally, we analyze that distance-guided multi-scale self-attention effectively eases thedebilitating structural problem.
Datasets ZINC is one of the most popular real-world molecular datasets of 250K graphs.
Similar tobenchmark [18], we choose the subset (12K) to regress molecular constrained solubility.
In KDD Cup2021, OGB-LSC was provided to the community to encourage the development of state-of-the-artgraph 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, thetask 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 [19], GAT [20], GIN [21] and
MoleculeX [22].
In addition, we compare the recently proposed random walk-based method CRAWL[23].
Especially, Transformer-based model SAN [24] , GT[25] 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-upsteps followed by a linear decay learning rate scheduler.
For the ZINC dataset, the large model is notencouraged.
Thus parameters of the model are limited to under 500 K .
And we conduct experimentswith 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 mostmodels 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 inthe proposed methods.
The results are shown in Table 4.
Spatial Embedding We compare atom-based spatial embedding to the proposed motif-based spatialembedding.
Although motif-based SPD introduces structural debilitating, it outperforms the vanillaatom-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 resultshows that graph Transformer architecture with multi-scale self-attention yields a large margin ofperformance improvement.
[Table]
4.3 Structural Debilitating Problem
We analyze the debilitating structural problem discussed above.
We conduct ablation experimentswith 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 performanceon various diameter graphs.
Especially for shorter distances, with relatively more vital interatomicforces, multi-scale self-attention achieves improvement more significantly.
Limitations
Although our methods achieve considerable progress, there are still some limitations.
Firstly, wemerely treat all kinds of rings as bi-connected components.
The chemical properties of aromatic rings [Table] [Table] differ significantly from those of regular rings.
Some sophisticated designs involving chemical bondproperties 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 oftheir 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 [29], whichwe did not introduce into our method.
Conclusions
This paper proposes a chemistry-guided transformer for the molecular graph forecasting task, whichachieves state-of-art performance.
It can be seen as an example of how simple chemical intuitions cansignificantly help us in chemical property forecasting tasks.
To add chemistry-inspired inductive biasinto 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 considerinteratomic forces and eliminate the debilitating structural problem.
Lastly, experiments showthat the proposed method considerably improves on two large benchmark datasets compared withstate-of-the-art algorithms.