[Paper Review] Graph Neural Networks with Learnable Structural and Positional Representations

Vijay Prakash Dwivedi, Anh Tuan Luu+3 more|arXiv (Cornell University)|Oct 15, 2021

Machine Learning in Materials Science62 citations

TL;DR

The paper introduces MP-GNNs-LSPE, which decouples learnable structural and positional encodings to improve expressive power, using RWPE or LapPE initializations and a specialized loss, achieving notable gains on molecular and other graph benchmarks.

ABSTRACT

Graph neural networks (GNNs) have become the standard learning architectures for graphs. GNNs have been applied to numerous domains ranging from quantum chemistry, recommender systems to knowledge graphs and natural language processing. A major issue with arbitrary graphs is the absence of canonical positional information of nodes, which decreases the representation power of GNNs to distinguish e.g. isomorphic nodes and other graph symmetries. An approach to tackle this issue is to introduce Positional Encoding (PE) of nodes, and inject it into the input layer, like in Transformers. Possible graph PE are Laplacian eigenvectors. In this work, we propose to decouple structural and positional representations to make easy for the network to learn these two essential properties. We introduce a novel generic architecture which we call LSPE (Learnable Structural and Positional Encodings). We investigate several sparse and fully-connected (Transformer-like) GNNs, and observe a performance increase for molecular datasets, from 1.79% up to 64.14% when considering learnable PE for both GNN classes.

Motivation & Objective

• Address the canonical-positioning limitation in generic GNNs by introducing learnable positional encodings that are decoupled from structural representations.
• Propose LSPE as a generic framework to augment MP-GNNs and Transformer GNNs with learnable positional information.
• Evaluate LSPE on molecular benchmarks (ZINC, MOLTOX21, MOLPCBA) and non-molecular graphs to demonstrate broad improvements.
• Show that learnable PE (especially RWPE) improves performance across sparse and fully-connected GNN variants while maintaining scalable complexity.

Proposed method

• Propose MP-GNNs-LSPE architecture that decouples structural features h and learnable positional features p, updated per layer.
• Introduce two initial positional encodings: Laplacian PE (LapPE) and Random Walk PE (RWPE); RWPE uses only RWii entries to keep complexity low.
• Define LSPE layer updates with h^{l+1} = f_h([h^l; p^l], { [h_j^l; p_j^l] }_{j in N_i}, e_{ij}^l), e^{l+1}_{ij} = f_e(h^l_i, h^l_j, e^l_{ij}), p^{l+1}_i = f_p(p^l_i, { p^l_j }_{j in N_i}, e^l_{ij}).
• Initial node/edge feature embeddings: h_i^{l=0} = LL_h(h_i^{in}) and e_{ij}^{l=0} = LL_e(e_{ij}^{in}).
• Positional encoding losses: optional Loss_LapEig(p^{l=L}) to encourage a Laplacian-eigenvector-like coordinate system, with a final total loss = TaskLoss([h^{L}; p^{L}]) + alpha Loss_LapEig(p^{L}).
• Instantiation: apply LSPE to both sparse MP-GNNs (GatedGCN, PNA, SAN, GraphiT) and Transformer GNNs, forming GatedGCN-LSPE, PNA-LSPE, SAN-LSPE, GraphiT-LSPE.

Experimental results

Research questions

• RQ1 Does decoupling structural and positional encodings via LSPE improve expressivity of MP-GNNs and Transformer GNNs on diverse graph tasks?
• RQ2 How do RWPE and LapPE compare as initial PEs within LSPE across different architectures and datasets?
• RQ3 What is the impact of learning PE at every layer versus injecting a fixed pre-computed PE only at input or final layers?
• RQ4 What are the practical performance gains of LSPE on molecular benchmarks (ZINC, MOLTOX21, MOLPCBA) and non-molecular graphs?
• RQ5 Does LSPE maintain scalability and training efficiency relative to standard MP-GNNs and Transformer GNNs?

Key findings

• Without positional encodings, GNN performance degrades across datasets; incorporating PE significantly improves results.
• RWPE as initial PE with LSPE yields substantial improvements, notably a 64.14% relative gain on ZINC for some models.
• GatedGCN-LSPE achieves the best ZINC test MAE of 0.090, outperforming several baselines and reaching state-of-the-art among domain-agnostic GNNs.
• LSPE provides comparable or better performance on MOLTOX21 and MOLPCBA, sometimes matching SOTA on these datasets.
• Sparse GNNs augmented with LSPE outperform Transformer GNNs with LSPE on molecular graphs, while also offering lower training times.
• Learning positional representations at every layer yields the best performance, and RWPE helps mitigate sign ambiguities associated with LapPE.