Graph Neural Networks with Learnable Structural and Positional Representations

1 INTRODUCTION . GNNs have recently emerged as a powerful class of deep learning architectures to analyze datasets where information is present in the form of heteregeneous graphs that encode complex data connectivity . Experimentally , these architectures have shown great promises to be impactful in diverse domains such as drug design ( Stokes et al. , 2020 ; Gaudelet et al. , 2020 ) , social networks ( Monti et al. , 2019 ; Pal et al. , 2020 ) , traffic networks ( Derrow-Pinion et al. , 2021 ) , physics ( Cranmer et al. , 2019 ; Bapst et al. , 2020 ) , combinatorial optimization ( Bengio et al. , 2021 ; Cappart et al. , 2021 ) and medical diagnosis ( Li et al. , 2020c ) . Most GNNs ( such as Defferrard et al . ( 2016 ) ; Sukhbaatar et al . ( 2016 ) ; Kipf & Welling ( 2017 ) ; Hamilton et al . ( 2017 ) ; Monti et al . ( 2017 ) ; Bresson & Laurent ( 2017 ) ; Veličković et al . ( 2018 ) ; Xu et al . ( 2019 ) ) are designed with a message-passing mechanism ( Gilmer et al. , 2017 ) that builds node representation by aggregating local neighborhood information . It means that this class of GNNs is fundamentally structural , i.e . the node representation only depends on the local structure of the graph . As such , two atoms in a molecule with the same neighborhood are expected to have similar representation . However , it can be limiting to have the same representation for these two atoms as their positions in the molecule are distinct , and their role may be specifically separate ( Murphy et al. , 2019 ) . As a consequence , the popular message-passing GNNs ( MP-GNNs ) fail to differentiate two nodes with the same 1-hop local structure . This restriction is now properly understood in the context of the equivalence of MP-GNNs with Weisfeiler-Leman ( WL ) test ( Weisfeiler & Leman , 1968 ) for graph isomorphism ( Xu et al. , 2019 ; Morris et al. , 2019 ) . 1Code : https : //github.com/vijaydwivedi75/gnn-lspe The said limitation can be alleviated , to certain extents , by ( i ) stacking multiple layers , ( ii ) applying higher-order GNNs , or ( iii ) considering positional encoding ( PE ) of nodes ( and edges ) . Let us assume two structurally identical nodes in a graph with the same 1-hop neighborhood , but different with respect to 2-hop or higher-order neighborhoods . Then , stacking several layers ( Bresson & Laurent , 2017 ; Li et al. , 2019 ) can propagate the information from a node to multiple hops , and thus differentiate the representation of two far-away nodes . However , this solution can be deficient for long-distance nodes because of the over-squashing phenomenon ( Alon & Yahav , 2020 ) . Another approach is to compute higher-order node-tuple aggregations such as in WL-based GNNs ( Maron et al. , 2019 ; Chen et al. , 2019 ) ; though these models are computationally more expensive to scale than MP-GNNs , even for medium-sized graphs ( Dwivedi et al. , 2020 ) . An alternative technique is to consider a global positioning of the nodes in the graph that can encode a graph-based distance between the nodes ( You et al. , 2019 ; Dwivedi et al. , 2020 ; Li et al. , 2020b ; Dwivedi & Bresson , 2021 ) , or can inform about specific sub-structures ( Bouritsas et al. , 2020 ; Bodnar et al. , 2021 ) . Contribution . In this work , we turn to the idea of learning positional representation that can be combined with structural GNNs to generate more expressive node embedding . Our main intent is to alleviate the lack of canonical positioning of nodes in arbitrary graphs to improve the representation power of MP-GNNs , while keeping their linear complexity for large-scale applications . For this objective , we propose a novel framework , illustrated with Figure 1 , that enables GNNs to learn both structural and positional representations at the same time ( thus named MPGNNs-LSPE ) . Alongside , we present a random-walk diffusion based positional encoding scheme to initialize the positional representations of the nodes . We show that the proposed architecture with learnable PE can be used with any graph network that fits to the MP-GNNs framework , and improves its performance ( 1.79 % to 64.14 % ) . In our demonstrations , we formulate LSPE instances of both sparse GNNs , such as GatedGCNs ( Bresson & Laurent , 2017 ) and PNA ( Corso et al. , 2020 ) and fully-connected Transformers-based GNNs ( Kreuzer et al. , 2021 ; Mialon et al. , 2021 ) . Our numerical experiments on three standard molecular benchmarks show that different instantiations of MP-GNNs with LSPE surpass the previous state-of-the-art ( SOTA ) on one dataset by a considerable margin ( 26.23 % ) , while achieving SOTA-comparable score on the other two datasets . The architecture also shows consistent improvements on three non-molecular benchmarks . In addition , our evaluations find the sparse MP-GNNs to be outperforming fully-connected GNNs , hence suggesting greater potential towards the development of highly efficient , yet powerful architectures for graphs . 2 RELATED WORK . In this section , we review briefly the three research directions theoretical expressivity of GNNs , graph positional encoding , and Transformer-based GNNs . Theoretical expressivity and Weisfeiler-Leman GNNs . As the theoretical expressiveness of MPGNNs is bounded by the 1-WL test ( Xu et al. , 2019 ; Morris et al. , 2019 ) , they may perform poorly on graphs that exhibit several symmetries ( Murphy et al. , 2019 ) , and additionally some message-passing functions may not be discriminative enough ( Corso et al. , 2020 ) . To this end , k-order EquivariantGNNs were introduced in Maron et al . ( 2018 ) requiring O ( nk ) memory and speed complexities . Although the complexity was improved to O ( n2 ) memory and O ( n3 ) respectively ( Maron et al. , 2019 ; Chen et al. , 2019 ; Azizian & Lelarge , 2020 ) , it is still inefficient compared with the linear complexity of MP-GNNs . Graph Positional Encoding . The idea of positional encoding , i.e . the notion of global position of pixels in images , words in texts and nodes in graphs , plays a central role in the effectiveness of the most prominent neural networks with ConvNets ( LeCun et al. , 1998 ) , RNNs ( Hochreiter & Schmidhuber , 1997 ) , and Transformers ( Vaswani et al. , 2017 ) . For GNNs , the position of nodes is more challenging due to the fact that there does not exist a canonical positioning of nodes in arbitrary graphs . Despite these issues , graph positional encoding are as much critical for GNNs as they are for ConvNets , RNNs and Transformers , as demonstrated for prediction tasks on graphs ( Srinivasan & Ribeiro , 2019 ; Cui et al. , 2021 ) . Nodes in a graph can be assigned index positional encoding ( PE ) . However , such a model must be trained with the n ! possible index permutations or else sampling needs to be done ( Murphy et al. , 2019 ) . Another PE candidate for graphs can be Laplacian Eigenvectors ( Dwivedi et al. , 2020 ; Dwivedi & Bresson , 2021 ) as they form a meaningful local coordinate system , while preserving the global graph structure . However , there exists sign ambiguity in such PE as eigenvectors are defined up to ±1 , leading to 2k number of possible sign values when selecting k eigenvectors which a network needs to learn . Similarly , the eigenvectors may be unstable due to eigenvalue multiplicities . You et al . ( 2019 ) proposed learnable position-aware embeddings based on random anchor sets of nodes , where the random selection of anchors has its limitations , which makes their approach less generalizable on inductive tasks . There also exists methods that encode prior information about a class of graphs of interest such as rings for molecules ( Bouritsas et al. , 2020 ; Bodnar et al. , 2021 ) which make MP-GNNs more expressive . But the prior information regarding graph sub-structures needs to be pre-computed , and sub-graph matching and counting require O ( nk ) for k-tuple sub-structure . Transformer-based GNNs . Although sparse MP-GNNs are very efficient , they are susceptible to the information bottleneck limitation ( Alon & Yahav , 2020 ) in addition to vanishing gradient ( similar to RNNs ) on tasks when long-range interactions between far away nodes are critical . To overcome these limitations , there have been recent works that generalize Transformers to graphs ( Dwivedi & Bresson , 2021 ; Kreuzer et al. , 2021 ; Ying et al. , 2021 ; Mialon et al. , 2021 ) which alleviates the long-range issue as ‘ everything is connected to everything ’ . However , these methods either use non-learnable PEs to encode graph structure information ( Dwivedi & Bresson , 2021 ; Ying et al. , 2021 ; Mialon et al. , 2021 ) , or inject learned PEs to the Transformer network that relies on Laplacian eigenvectors ( Kreuzer et al. , 2021 ) , thus inheriting the sign ambiguity limitation . A detailed review of the above research directions is available in the supplementary Section B . We attempt to address some of the major limitations of GNNs by proposing a novel architecture with consistent performance gains . 3 PROPOSED ARCHITECTURE . In this work , we decouple structural and positional representations to make it easy for the network to learn these two critical characteristics . This is in contrast with most existing architectures s.a. Dwivedi & Bresson ( 2021 ) ; Beani et al . ( 2021 ) ; Kreuzer et al . ( 2021 ) that inject the positional information into the input layer of the GNNs , and You et al . ( 2019 ) that rely on distance-measured anchor sets of nodes limiting general , inductive usage . Given the recent theoretical results on the importance of informative graph PE for expressive GNNs ( Murphy et al. , 2019 ; Srinivasan & Ribeiro , 2019 ; Loukas , 2020 ) , we are interested in a generic framework that can enable GNNs to separate positional and structural representations to increase their expressivity . Section 3.1 will introduce our approach to augment GNNs with learnable graph PE . Our framework can be used with different GNN architectures . We illustrate this flexibility in Sections C.1 and C.2 where the decoupling of structural and positional information is applied to both sparse MP-GNNs and fully-connected GNNs . 3.1 GENERIC FORMULATION : MP-GNNS-LSPE Notation . Let G = ( V , E ) be a graph with V being the set of nodes and E the set of edges . The graph has n = |V| nodes and E = |E| edges . The connectivity of the graph is represented by the adjacency matrix A ∈ Rn×n where Aij = 1 if there exists an edge between the nodes i and j ; otherwise Aij = 0 . The degree matrix is denoted D ∈ Rn×n . The node features and positional features for node i is denoted by hi and pi respectively , while the features for an edge between nodes i and j is indicated by eij . A GNN model is composed of three main components ; an embedding layer for the input features , a stack of convolutional layers , and a final task-based layer , as in Figure 1 . The layers are indexed by ` and ` = 0 denotes the input layer . Standard MP-GNNs . Considering a graph which has available node and edge features , and these are transformed at each layer , the update equations for a conventional MP-GNN layer are defined as : MP-GNNs : h ` +1i = fh ( h ` i , { h ` j } j∈Ni , e ` ij ) , h ` +1i , h ` i ∈ Rd , ( 1 ) e ` +1ij = fe ( h ` i , h ` j , e ` ij ) , e ` +1ij , e ` ij ∈ Rd , ( 2 ) where fh and fe are functions with learnable parameters , and Ni is the neighborhood of the node i . The design of functions fh and fe depends on the GNN architecture used , see Zhou et al . ( 2020 ) for a review . As Transformer neural networks ( Vaswani et al. , 2017 ) are a special case of MP-GNNs ( Joshi , 2020 ) , Eq . ( 1 ) can be simplified to encompass the original Transformers by dropping the edge features and making the graph fully connected . Input features and initialization . The node and edge features at layer ` = 0 are produced by a linear embedding of available input node and edge features denoted respectively by hini ∈ Rdv , einij ∈ Rde : h ` =0i = LLh ( h in i ) = A 0hini + a 0 ∈ Rd , e ` =0ij = LLe ( einij ) = B0einij + b0 ∈ Rd , where A0 ∈ Rd×dv , B0 ∈ Rd×de and a0 , b0 ∈ Rd are the learnable parameters of the linear layers . Positional Encoding . Existing MP-GNNs that integrate positional information usually propose to concatenate the PE with the input node features , similarly to Transformers ( Vaswani et al. , 2017 ) : MP-GNNs-PE : h ` +1i = fh ( h ` i , { h ` j } j∈Ni , e ` ij ) , h ` +1i , h ` i ∈ Rd , ( 3 ) e ` +1ij = fe ( h ` i , h ` j , e ` ij ) , e ` +1ij , e ` ij ∈ Rd , ( 4 ) with initial h ` =0i = LLh ( [ hini pini ] ) = D0 [ hini pini ] + d0 ∈ Rd , ( 5 ) and e ` =0ij = LLe ( e in ij ) = B 0einij + b 0 ∈ Rd , ( 6 ) where pini ∈ Rk is the input PE of node i , D0 ∈ Rd× ( dv+k ) , d0 ∈ Rd are parameters for the linear transformation . Such architecture merges the positional and structural representations together . It has the advantage to keep the same linear complexity for learning , but it does not allow the positional representation to be changed and better adjusted to the task at hand . Decoupling position and structure in MP-GNNs . We decouple the positional information from the structural information such that both representations are learned separately resulting in an architecture with Learnable Structural and Positional Encodings , which we call MP-GNNs-LSPE . The layer update equations are defined as : MP-GNNs-LSPE : h ` +1i = fh ( [ h ` i p ` i ] , { [ h ` j p ` j ] } j∈Ni , e ` ij ) , h ` +1i , h ` i ∈ Rd , ( 7 ) e ` +1ij = fe ( h ` i , h ` j , e ` ij ) , e ` +1ij , e ` ij ∈ Rd , ( 8 ) p ` +1i = fp ( p ` i , { p ` j } j∈Ni , e ` ij ) , p ` +1i , p ` i ∈ Rd , ( 9 ) The difference of this architecture with the standard MP-GNNs is the addition of the positional representation update Eq . ( 9 ) , along with the concatenation of these learnable PEs with the node structural features , Eq . ( 7 ) . As we will see in the next section , the design of the message-passing function fp follows the same analytical form of fh but with the use of the tanh activation function to allow positive and negative values for the positional coordinates . It should be noted that the inclusion of the edge features , e ` ij in the h or p update is optional as several MP-GNNs do not include edge features in their h updates . Nevertheless , the architecture we present is made as generic so as to be used for future extensions in a convenient way . Definition of initial PE . The choice of the initial PE is critical . In this work , we consider two PEs : Laplacian PE ( LapPE ) and Random Walk PE ( RWPE ) . LapPE are defined in Section B.2 as pLapPEi = [ Ui1 , Ui2 , · · · , Uik ] ∈ Rk . LapPE provide a unique node representation and are distancesensitive w.r.t . the Euclidean norm . However , they are limited by the sign ambiguity , which requires random sign flipping during training for the network to learn this invariance ( Dwivedi et al. , 2020 ) . Inspired by Li et al . ( 2020b ) , we propose RWPE , a PE based on the random walk ( RW ) diffusion process ( although other graph diffusions can be considered s.a. PageRank ( Mialon et al. , 2021 ) ) . Formally , RWPE are defined with k-steps of random walk as : pRWPEi = [ RWii , RW2ii , · · · , RWkii ] ∈ Rk , ( 10 ) where RW = AD−1 is the random walk operator . In contrast of Li et al . ( 2020b ) which uses the full matrix RWij for all pairwise nodes , we adopt a low-complexity usage of the random walk matrix by considering only the landing probability of a node i to itself , i.e . RWii . Note that these PE do not suffer from the sign ambiguity of LapPE , so the network is not required to learn additional invariance . RWPE provide a unique node representation under the condition that each node has a unique k-hop topological neighborhood for a sufficient large k. This assumption can be discussed . If we consider synthetic strongly regular graphs like the CSL graphs ( Murphy et al. , 2019 ) , then all nodes in a graph have the same RWPE for any k value , since they are isomorphic by construction . However , despite RWPE being the same for all nodes in a graph , these PE are unique for each class of isomorphic graphs , resulting in a perfect classification of the CSL dataset , see Section A.1 . For graphs such as Decalin and Bicyclopentyl ( Sato , 2020 ) , nodes which are not isomorphic receive different RWPE for k ≥ 5 , also in Section A.1 . Finally , for real-world graphs like ZINC molecules , most nodes receive a unique node representation for k ≥ 24 , see Figure 2 for an illustration , where the two molecules have 100 % and 71.43 % unique RWPEs respectively . Section A.3 presents a detailed study . Experimentally , we will show that RWPE outperform LapPE , suggesting that learning the sign invariance is more difficult ( as there exist 2k possible sign flips for each graph ) than not exactly having unique node representation for each node . As mentioned above for CSL , RWPE are related to the problem of graph isomorphism and higher-order node interactions . Precisely , iterating the random walk operator for a suitable number of steps allows coloring non-isomorphic nodes , thus distinguishing several cases of non-isomorphic graphs on which the 1-WL test , and equivalently MP-GNNs , fail s.a. the CSL , Decalin and Bicyclopentyl graphs . We refer to Section A.2 for a formal presentation of the iterative algorithm . Finally , the initial PE of the network is obtained by embedding the LapPE or RWPE into a d-dimensional feature vector : p ` =0i = LLp ( p PE i ) = C 0pPEi + c 0 ∈ Rd , where C0 ∈ Rd×k , c0 ∈ Rd . ( 11 ) Positional loss . As we separate the learning of the structual and positional representations , it is possible to consider a specific positional encoding loss along with the task loss . A natural candidate is the Laplacian eigenvector loss ( Belkin & Niyogi , 2003 ; Lai & Osher , 2014 ) that enforces the PE to form a coordinate system constrained by the graph topology . As such , the final loss function of MP-GNNs-LSPE is composed of two terms : Loss = LossTask ( [ h ` =L p ` =L ] ) + α LossLapEig ( p ` =L ) , ( 12 ) where h ` =L ∈ Rn×d , p ` =L ∈ Rn×k , k is the dimension of learned PE , ` = L is the final GNN layer , and α > 0 an hyper-parameter . Observe also that we enforce the final positional vectors p ` =L to have centered and unit norm with mean ( p ` =L· , k ) = 0 , ‖p ` =L· , k ‖ = 1 , ∀k to better approximate the Laplacian eigenvector loss defined by LossLapEig ( p ) = 1k trace ( pT∆p ) + λk ∥∥pT p − Ik∥∥2F with λ > 0 and ‖ · ‖2F being the Frobenius norm . 3.2 INSTANCES OF LSPE WITH MP-GNNS AND TRANSFORMER GNNS We instantiate two classes of GNN architectures , both sparse MP-GNNs and fully-connected Transformer GNNs using our proposed LSPE framework . For sparse MP-GNNs , we consider GatedGCN ( Bresson & Laurent , 2017 ) and PNA ( Corso et al. , 2020 ) , while we extend the recently developed SAN ( Kreuzer et al. , 2021 ) and GraphiT ( Mialon et al. , 2021 ) with LSPE to develop TransformerLSPE architectures . We briefly demonstrate here how a GNN can be instantiated using LSPE ( Eqs . ( 7-9 ) ) by developing GatedGCN-LSPE ( Eqs . ( 14-16 ) ) , while the complete equations for the four models are defined in Section C of the supplementary material , given the space constraint . GatedGCN-LSPE : Originally , GatedGCNs are sparse MP-GNNs equipped with a soft-attention mechanism that is able to learn adaptive edge gates to improve the message aggregation step of GCN networks ( Kipf & Welling , 2017 ) . Our proposed extension of this model with LSPE is defined as : h ` +1 , e ` +1 , p ` +1 = GatedGCN-LSPE ( h ` , e ` , p ` ) , h ∈ Rn×d , e ∈ RE×d , p ∈ Rn×d , ( 13 ) with h ` +1i = h ` i + ReLU ( BN ( A ` 1 [ h ` i p ` i ] + ∑ j∈N ( i ) η ` ij A ` 2 [ h ` j p ` j ] ) ) , ( 14 ) e ` +1ij = e ` ij + ReLU ( BN ( η̂ ` ij ) ) , ( 15 ) p ` +1i = p ` i + tanh ( C ` 1p ` i + ∑ j∈N ( i ) η ` ij C ` 2p ` j ) , ( 16 ) where η ` ij = σ ( η̂ ` ij ) / ( ∑ j′∈N ( i ) σ ( η̂ ` ij′ ) + ) , η̂ ` ij = B ` 1h ` i + B ` 2h ` j + B ` 3e ` ij , h ` i , e ` ij , p ` i , η ` ij , η̂ ` ij ∈ Rd , A ` 1 , A ` 2 ∈ Rd×2d and B ` 1 , B ` 2 , B ` 3 , C ` 1 , C ` 2 ∈ Rd×d . Notice the p-update in Eq . ( 16 ) follows the same analytical form as the h-update in Eq . ( 14 ) except for the difference in activation function , and omission of BN , which was not needed in our experiments .

Summary: This paper is concerning the Positional Encoding (PE) for GNNs.  PE augments the typical GNNs to distinguish isomorphic nodes. However, existing PE models such as Laplacian eigenvectors require huge computational resources. This manuscript proposes LSPE that augments the input to the nodes AND the embedding vectors with PE elements.  The LSPE models iteratively updates the embedding for PEs, in addition to the node feature embeddings.  The update formula of the PE embedding is similar to those of the node feature embedding, thus the computational requirements of the LSPE does not .  The manuscript tests two ways of PEs, one is based on the Laplacian, and the other is based on the random-walk. The manuscript also proposes the PE=only loss to foster the training.  Experimental results shows that the propose LSPE can improve the graph regression of the ZINC dataset greatly, and also can achieve some improvements in graph classification tasks on the MOLTOX21 and the MOLPCBA datasets.