GRPE: R ELATIVE P OSITIONAL E NCODING FOR G RAPH T RANSFORMER
A BSTRACT
Designing an efficient model to encode graphs is a key challenge of molecular representation learning.
Transformer built upon efficient self-attention is a natural choice for graph processing, but it requires explicit incorporation of positional information.
Existing approaches either linearize a graph to encode absolution position in the sequence of nodes, or encode relative position with another node using bias terms.
The former loses preciseness of relative position from linearization, while the latter loses a tight integration of node-edge and node-spatial information.
In this work, we propose relative positional encoding for a graph to overcome the weakness of the previous approaches.
Our method encodes a graph without linearization and considers both node-spatial relation and node-edge relation.
We name our method Graph Relative Positional Encoding dedicated to graph representation learning.
Experiments conducted on various molecular property prediction datasets show that the proposed method outperforms previous approaches significantly.
Our code is publicly available at https://github.com/lenscloth/GRPE .
I NTRODUCTION
The average cost of drug discovery has increased drastically in the last decade with a declining success rate of developing new therapeutics.
Large-scale screening with deep neural networks draws considerable attention to lower the cost of drug discovery, especially in the lead finding and lead optimization phase.
By representing a molecule in a graph, a graph neural network can be utilized on the screening by predicting important properties, e g. , drug-likeness, solubility or synthesizability.
Therefore, graph representation learning has become a key technique for drug discovery.
Transformer introduced by Vaswani et al.
(2017) has been effective in graph representation learning, overcoming inductive biases of graph convolutional networks by using self-attention.
On the other hand, explicit representations of position in graph convolutional networks are lost, thus incorporating graph structure on the hidden representations of self-attention is a key challenge.
We categorize existing works into two: (a) linearizing graph with graph Laplacian to encode the absolute position of each node (Dwivedi & Bresson, 2020; Kreuzer et al., 2021) (b) encoding position relative to another node with bias terms (Ying et al., 2021).
The former loses preciseness of position due to linearization, while the latter loses a tight integration of node-edge and node-spatial information.
Ours can be interpreted as overcoming the weakness of the two, which considers both relative position and the interaction with node.
To avoid losing relative position due to the linearization, we propose to adopt the relative positional encoding by Shaw et al.
Its efficacy has been verified in natural language processing, however, the efficacy is not yet studied in a graph.
Since the original work is designed for 1D word sequences, we reformulate to incorporate graph-specific properties.
To this end, we introduce two sets of learnable positional encoding vectors to represent spatial relation or edge between two nodes.
Our method considers the interaction between node features and the two encoding vectors to integrate both node-spatial relation and node-edge information.
We name our method for graph representation learning as G raph R elative P ositional E ncoding (GRPE).
[Figure]
We extensively conducted experiments to validate the efficacy of our proposed method on various tasks from graph classification to graph regression.
Models built with our relative positional encoding achieve state-of-the-art performance on various molecule property prediction datasets, showing the efficacy of our proposed method.
R ELATED W ORK
Existing works leverage Transformer architecture to learn graph representation.
We categorize those methods as follows.
Earlier models adopt Transformer without explicit encoding of positional information on a graph.
Veli ˇ ckovi ´ c et al.
(2017) replace graph convolution operation with self-attention module where attention is only performed within neighboring nodes.
Rong et al.
(2020) stack self-attention module next to the graph convolutional networks iteratively to consider long-range interaction between nodes.
In their method, affinity is considered only on the graph convolutional networks, and positional information is not given on self-attention.
Later works employ absolute positional encoding to explicitly encode positional information of graph on Transformer.
Their main idea is to linearize a graph into a sequence of nodes, and an absolute positional encoding is added to the input feature.
Dwivedi & Bresson (2020) adopted graph
Laplacian as a positional encoding, where each cell of encoding represents partitions after graph min-cut.
Nodes sharing many partitions after graph min-cut would have similar graph Laplacian vectors.
Kreuzer et al.
(2021) employ a learnable positional encoding with a Transformer where its input is the Laplacian spectrum of a graph.
Due to the linearization of a graph, those approaches lose the preciseness of position on the graph.
Meanwhile, encoding relative positional information has been studied to avoid losing the preciseness of position.
Graphormer introduced by Ying et al.
(2021) encodes relative position on scaled dot product attention map by adding bias terms.
However, the bias terms are parameterized only relative position such as shortest path distance or edge type, and interaction with node features is lost.
On the other hand, Shaw et al.
(2018) introduce relative positional encoding, which considers tight integration of node and relative position, however, their method is designed to process 1D word sequences.
We extend the systematic design of Shaw et al.
(2018) for a graph.
Our method considers the interaction between nodes and graph-specific properties such as edges, and encodes precise positional information of a graph.
B ACKGROUND
Notation.
We denote a set of nodes on the graph { n i } i =1: N and a set of edges on the graph { e ij | j ∈
N i } i =1: N , where N is the number of nodes and N i is a set neighbors of a node n i .
Both n i and e ij are positive integer numbers to index the type of nodes or edges, e.g. , atom numbers or bond types of a molecule.
ψ ( i, j ) denotes a function encodes structural relationship between the node n i and n j .
3. T RANSFORMER
Transformer is built by stacking multiple self-attention layers.
Self-attention maps a query and a set of key pairs to compute an attention map.
Values are weighted summed with the weight on the attention map to output the hidden feature for the following layer.
Specifically, x i ∈ R
d x denotes the input feature of the node n i , and z i ∈
R d z denotes the output feature of the self-attention module.
The self-attention module computes query q , key k , and value v with independent linear transformations: W query ∈ R d x × d z , W key ∈ R d x × d z and W value ∈ R d x × d z .
[Equation] x , k = [Equation]
x and v = W value x.
[Equation]
The attention map is computed by applying a scaled dot product between the queries and the keys.
[Equation]
and
[Equation]
The self-attention module outputs the next hidden feature by applying weighted summation on the values.
[Equation]
z is later fed into a feed forward neural network with a residual connection (He et al., 2016).
However, we defer detailed explanations since it is out of the scope of our paper.
In practice, self-attention module with multi-head is adopted.
3. G RAPH WITH T RANSFORMER
To encode graph structure in Transformer, previous methods focus on encoding graph information into either the attention map or input features fed to Transformer.
Graphormer (Ying et al., 2021) adopted two additional terms on the self-attention module to encode graph information on the attention map.
[Equation]
ψ represents spatial relation between n i and n j , which outputs the shortest path distance between the two nodes.
Learnable scalar bias b ψ ( i,j ) encodes spatial relation between two nodes, e.g. , b l is a bias representing two nodes that are l -hop apart.
An embedding vector E e ij is a feature representing edge between the node n i and the node n j , and w is a learnable vector.
E e ij · w encodes edge between the two nodes.
Moreover, Graphormer adds centrality encoding into the input x which represents the number of edges of a node.
Graphormer encodes graph on the attention map with the bias terms parameterized by either spatial relation or edge without considering node feature.
However, we additionally consider node-spatial relation and node-edge relation.
We will explain the details in
Section 4.2.
Dwivedi & Bresson (2020) and Kreuzer et al.
(2021) utilize graph Laplacian (Belkin & Niyogi, 2003)
λ ∈ R | N |× d x as positional encodings on the input feature x ; λ are the top- d x smallest eigenvectors of I − D − 12 AD − 12 where I is an identity matrix,
A is an adjacency matrix and D is a degree matrix.
Each cell of a graph Laplacian vector represents partitions after graph min-cut, and neighbouring nodes sharing the many partitions would have similar graph Laplacian.
The graph Laplacian λ i represents the structure of a graph with respect to node n i .
[Equation]
Kreuzer et al.
(2021) adopt an additional Transformer model f to produce learnable positional encoding: ˆ x i
= x
i + ˆ
λ i where ˆ λ = f ( λ ) .
By adding the graph Laplacian into input x , graph information can be encoded in both the attention map and the hidden representations.
Their methods lose relative positional information during the linearization of a graph to absolute positional encodings.
However, our method encodes a graph directly on the attention map without linearization, thus relative positional information is encoded without loss.
Spatial Relation 𝝍(𝒊,𝒋)
𝑛 !
𝑛 " 𝑛 # 𝑛 $ 𝑛 %
Edge 0 E d g e
E d g e𝑛 !
𝑛 " 𝑛 # 𝑛 $ 𝑛 % 𝑛 !
𝑛 " 𝑛 # 𝑛 $ 𝑛 %0 0 1 1 11 1 2 2 2 far far far far
Edge 𝒆 𝒊𝒋 𝑛 !
𝑛 " 𝑛 # 𝑛 $ 𝑛 % 𝑛 !
𝑛 " 𝑛 # 𝑛 $ 𝑛 % [Figure] [Figure]
no [Figure]
𝑞
[Figure]
𝑘
[Figure]
𝑣
[Figure]
z: 5×𝑑 ’
[Figure]
𝑏 # [Figure] [Figure]
[Figure]
Relative Relation between Nodes
Graph Relative Positional Encoding
[Figure]
O UR A PPROACH
We reformulate relative positional encoding for 1D sequence data (Shaw et al., 2018) to incorporate graph-specific properties, such as spatial relation or edges between nodes.
Our distinction is twofold.
First, we integrate node-spatial and node-edge information on the attention map that was not leveraged in all existing works of graph Transformer.
For that, we propose node-aware attention which considers the interactions existing in two pairs: node-spatial relation and node-edge relation.
Second, we also encode a graph to the hidden representation of self-attention.
For that, we propose graph-encoded values that directly encode relative positional information on the features of value by addition.
Our node-aware attention applies the attention mechanism in a node-wise manner, while our graph-encoded value applies the attention mechanism in a channel-wise manner.
4. R ELATIVE P OSITIONAL E NCODING FOR G RAPH
We define two encodings to represent relative positional relation between two nodes in a graph.
The first is spatial encoding P , and we define the encodings for query, key, and value respectively:
P query , P key , P value ∈ R L × d z .
Each vector of P represents the spatial relation between two nodes, e.g. , P l represents the spatial relation of two nodes where their shortest path distance is l .
L is the maximum shortest path distance that our method considers.
The second is edge encoding E , and we define the encodings for query, key, and value respectively:
E query , E key , E value ∈ R E
×
d z .
E e ij is a vector representing edge between two nodes n i and n j .
E is the number of types of edge.
The spatial encodings and the edge encodings are shared throughout all layers.
4. N ODE -A WARE A TTENTION
We propose two terms to encode graph on the attention map with two newly proposed encodings.
The first term is b spatial .
It encodes graph by considering interaction between node feature and spatial relation in graph.
[Equation]
The second term is b edge .
It encodes graph by considering interaction between node feature and edge in graph.
[Equation]
Finally, the two terms are added to scaled dot product attention map to encode graph information.
[Equation]
Our two terms consider node-spatial relation and node-edge relation, but Graphormer did not consider the interaction with node feature.
For instance, any two nodes with the same distance apart have the same bias b ψ ( i,j ) on Eq.
(4), however our b spatial deploys different values according to the node features of query and key.
4. G RAPH -E NCODED VALUE
We propose to encode a graph to the hidden features of self-attention, when values are weighted summed with the attention map.
We encode both spatial encoding and edge encoding into value via summation:
[Equation]
Our method directly encodes graph information into the hidden features of value.
The attention weight ˆ a is applied equally for all channels, while our graph-encoded value enriches the feature of each channel.
Therefore, node-aware attention is a node-wise attention mechanism while graphencoded value is a channel-wise attention mechanism.
Our method encodes a graph into the hidden features without losing preciseness of positional information thanks to our relative positional encoding.
However, previous approaches encode a graph on the initial input x by linearizing graph e.g. , centrality encoding (Ying et al., 2021) or graph Laplacian (Kreuzer et al., 2021; Dwivedi & Bresson, 2020) which loss preciseness of position.
4. C OMPLEXITY A NALYSIS
A naive implementation of computing all pairs of b spatial requires time complexity of O ( N 2 d z ) , since it requires performing inner product between all node pairs.
Instead, we pre-compute an inner product of all possible pairs of node features and spatial encoding vectors P which requires time complexity of O ( NLd z ) .
Then we assign the pre-computed value according to the indices of node pairs.
Likewise, for the b edge , we pre-compute an inner product of all possible pairs of node features and edge encoding vectors E query which requires time complexity of O ( NEd z ) .
The time complexity is reduced significantly with our implementation since L and E are much smaller than the number of nodes N .
E XPERIMENT 5. I MPLEMENTATION DETAILS 5.1. V IRTUAL NODE
Following Gilmer et al.
(2017) and Ying et al.
(2021), we adopt a special node called virtual node which is connected to all other nodes.
The role of a virtual node is similar to special tokens such as a classification token (Devlin et al., 2018), where its output feature z is used as the input for the branch that predicts downstream tasks.
We additionally define two encoding vectors to define both spatial relation P virtual and edge E virtual for query, key and value respectively.
Note that, the virtual node does not involve to find the shortest path between two nodes.
Throughout all experiments, we add the virtual node to a graph to perform the downstream tasks.
5.1. S PATIAL R ELATION
We utilize shortest path distance ψ ( i, j ) to describe the spatial relation between two nodes n i and n j .
L is the maximum distance of the shortest path that we consider, and we utilize a special encoding vector P far for the node pairs with a distance more than L .
For the nodes pairs that are unreachable, we utilize another special encoding vector P unreachable .
Finally, for the pairs that are connected with the virtual node, we utilize another special encoding vector P virtual .
[Table]
[Table]
[Table]
E DGE
Some pair of nodes are not connected with edges.
Therefore, we utilize a special encoding vector
E no for the pairs of node that are not connected with any edges; { ( n i , n j ) | i ̸ =
j and j / ∈ N i } i =1: | N | .
For the pair of two identical nodes where i =
j , we use a special embedding vector E self .
Finally, for the pairs that are connected with the virtual node, we utilize another special encoding vector E virtual .
5. M OLECULE PROPERTY PREDICTION
† indicates the models adopted Transformer for learning graph representation.
Bold faced text indicates the best result.
We summarize the model configurations of our experiments in Table 1.
We validate our method on the tasks of the molecule property prediction such as OGBG-MolPCBA (MolPCBA) (Hu et al., 2020), OGBG-MolHIV (MolHIV) (Hu et al., 2020) and ZINC (Dwivedi et al., 2020).
MolPCBA consists of 437,929 graphs and the task is to predict multiple binary labels indicating various molecule properties.
The evaluation metric is average precision (AP).
MolHIV is a small dataset that consists of 41,127 graphs.
The task is to predict a binary label indicating whether a molecule inhibits HIV virus replication or not.
The evaluation metric is area under the curve (AUC).
ZINC is also a small dataset that consists of 12,000 graphs, and the task is to regress a molecule property.
The evaluation metric is mean absolute error (MAE).
All experiments are conducted for 5 times, and we report the mean and the standard deviation of the experiments.
We adopt the linear learning rate decay, and the learning rate starts from 2 ∗ 10 − 4 and ends at 1 ∗ 10 − 9 .
We set L to 5.
For ZINC dataset, we adopt GRPE-Small configuration with less than 500k parameters for a fair comparison.
For MolHIV and MolPCBA datasets, we initialize the parameter of the models with the weight of a pretrained model trained on PCQM4M (Hu et al., 2020) dataset.
Table 2 shows the results on ZINC dataset, our model achieve state-of-the-art MAE score.
Tableshows the results on MolPCBA dataset, our model achieves state-of-the-art AP score.
Table 3 shows the results on MolHIV dataset, our model achieves the state-of-the-art AUC with less parameters than Graphormer (Ying et al., 2021).
5. OGB L ARGE S CALE C HALLENGE
We validate our method on two datasets of OGB large scale challenge (Hu et al., 2020).
The two datasets aim to predict the DFT-calculated HOMO-LUMO energy gap of molecules given their molecular graphs.
We conduct experiments on both the PCQM4M and PCQM4Mv2 datasets, which are currently the biggest molecule property prediction datasets containing about 4 million graphs
[Table] [Table]
in total.
PCQM4Mv2 contains the DFT-calcuated 3D strcuture of molecules.
For our experiments, we only utilize 2D molecular graphs not 3D structures.
Throughout experiments, we set L to 5.
We adopt a GRPE-Standard for fair comparisons with Graphormer (Ying et al., 2021).
We linearly increase learning rate up to 2 ∗ 10 − 4 for 3 epochs and linearly decay learning rate upto 1 ∗ 10 − 9 for 400 epochs.
We are unable to measure the test MAE of PCQM4M, because the test dataset is deprecated as PCQM4Mv2 is newly released.
Table 5 shows the results on PCQM4M dataset.
Our model achieves the state-of-the-art validation
MAE score.
Table.
6 shows the results on PCQM4Mv2 dataset.
Our model achieves the second best result on both the validation dataset and the test dataset.
However, our model use only half of the parameters compared to the best model
. 5. A BLATION STUDY 5.4. C OMPONENTS OF GRPE
We validate the effect of components of GRPE on ZINC dataset.
We adopt GRPE-Small.
Tableshows the ablation study results.
The first row is identical to the plain Transformer without any positional encodings, and it obviously shows the highest error.
Adding either b spatial or b edge lowers the error, and using them together further lowers the error.
Finally, adding our Graph-Encoded Value does help to improve the performance.
5.4. M AXIMUM SHORTEST PATH DISTANCE L
We validate the effects of the maximum shortest path distance L .
We adopt GRPE-Standard, and the models are trained from scratch.
Figure 3 shows the ablation study result.
Increasing L means that a model can identify the position of nodes that are further away.
The AUC consistently improves by increasing L from one to four, but L more than four does not further improve performance.
5.4. S HARING SPATIAL ENCODING AND EDGE ENCODING
We conduct ablation studies about the effects of sharing spatial encoding P and edge encoding E for all layers.
We adopt GRPE-Small.
We conduct five independent runs.
The results in Tableshow that sharing two encodings does improve MAE but not significantly.
Empirically, sharing the encodings does not significantly change the performance of a model especially on large datasets.
[Table]
[Table]
Kipf & Welling,
[Table]
et al.,
[Table]
Kipf & Welling,
[Table]
al.
[Table]
[Table] [Table]
[Figure]
Is
[Table]
579k
[Table]
[Table]