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 proposes a framework that utilizes Random Walk Positional Embeddings (RWPE) as extra features to boost the performance of GNNs. In particular, positional embeddings are updated as a separate forward network in each layer. The framework has demonstrated improved quality by injecting its positional embeddings and feed-forward structures in several GNNs.