Graph Positional and Structural Encoder

Several approaches have been proposed to overcome the aforementioned limitations of standard MPNNs: Morris et al. (2019) propose higher-order MPNNs, Gutteridge et al. (2023) and Barbero et al. (2024) optimize information flow on graphs through graph rewiring, while Ding et al. (2024) draw inspiration from deep state-space models and RNNs.

Positional encodings were originally implemented as a series of sinusoidal functions in the Transformer model to capture the ordinal position of words in a sentence (Vaswani et al., 2017). However, capturing the positions of nodes in a graph is harder, as nodes of a graph lack such a canonical ordering. Many recent works on graph Transformers (GT) thus use the graph Laplacian eigenvectors as the positional encodings (Rampášek et al., 2022; Kreuzer et al., 2021), which are direct analogues to the sinusoids in Euclidean space (Spielman, 2012). Other methods for encoding positional information include electrostatic potential encodings (Kreuzer et al., 2021), shortest-path distances (Ying et al., 2021), and tree-based encodings (Shiv & Quirk, 2019). Structural encodings, on the other hand, have been developed to encode rich local and global connectivity patterns on graph-structured data. The random walk encoding, for example, has shown great success when used with GTs, particularly on small molecular graph benchmarks (Rampášek et al., 2022; Dwivedi & Bresson, 2021; Dwivedi et al., 2022b). Other notable structural encodings include the heat kernel (Kreuzer et al., 2021; Mialon et al., 2021), subgraphs (Bouritsas et al., 2022; Zhao et al., 2022; Chen et al., 2022), and node degree centralities (Ying et al., 2021).

PSEs are also useful for MPNNs, besides GTs. They can be directly used as additional node features for an MPNN (Dwivedi et al., 2022a; Lim et al., 2022; Wang et al., 2022b). Other works designed approaches to process the PSEs separately from the original node features. For example, LSPE (Dwivedi et al., 2022b) processes the PSEs with a separate channel and applies an auxiliary loss. SignNet and BasisNet explicitly design components in the prediction model to handle the graph Laplacian eigenvectors, aiming to resolve the sign- and basis-ambiguity issues (Lim et al., 2022). Despite this great amount of work in developing methods to better utilize hand-crafted PSEs, it is still unclear how to systematically encode information from multiple types of PSEs to effectively augment GNNs for diverse applications. GPSE thus represents the first of its kind in an attempt to tackle this fundamental problem.

We design a diverse collection of six PSEs for GPSE to learn against, including the Laplacian eigenvectors (4) and eigenvalues (4), the electrostatic positional encodings (7), the random walk structural encodings (20), the heat kernel structural encodings (20), and the cycle counting graph encodings (7). In short, positional encodings inform the relative position of each node in the graph, while structural encodings describe the local connectivity patterns around a node (Figure 1B). See Appendix A for precise mathematical definitions of all PSEs used in GPSE training.

At a high level, our GPSE model is an MPNN consisting of stacked graph convolution blocks with residual gating, and skip-connections in-between. The mathematical formulation of the architecture can be found in Appendix B.1. We illustrate below that this careful design enables GPSE to excel at learning powerful PSE representations. Particularly, we present different architectural choices and how they remedy the two key challenges of over-smoothing and over-squashing (see Appendix C for technical details).

Evidence suggests that over-smoothing and over-squashing in graph networks relate to the graph’s curvature (Topping et al., 2022; Nguyen et al., 2022; Giraldo et al., 2022). Formulations of graph curvature (Forman, 2003; Ollivier, 2009; Sreejith et al., 2016; Topping et al., 2022) aim to encode how a node’s neighborhood looks like a clique (positive curvature), a grid (zero curvature), or a tree (negative curvature). A clique’s high connectivity leads to rapid smoothing, while a tree’s exponentially increasing size of $k$-hop neighborhood causes over-squashing. These phenomena are in competition, and negating both is impossible in an MPNN using graph rewiring (architectures using modified adjacency for node aggregation). However, there seems to be a sweet spot where the two effects are not minimized by themselves, but their sum is minimized. This minimum is sought in Giraldo et al. (2022) and Nguyen et al. (2022). Some of our following choices of architecture are justified by this search of a sweet spot in the smoothing-squashing trade-off.

Deep GNN As several of the target PSEs, such as the Laplacian eigenvectors, require having a global view of the query graph, it is crucial for the encoder to capture long-range dependencies accurately. To accomplish this, we need to use an unconventionally deep MPNN with 20 layers. However, if a graph network suffers from over-smoothing, having this many layers will result in approximately uniform node features  (Oono & Suzuki, 2020; Li et al., 2019).

Residual connections & gating mechanism A first attempt at reducing the smoothing is to exploit the proven ability of residual connections in reducing over-smoothing (Li et al., 2019). Using a gating mechanism in aggregation helps reduce the over-smoothing even further. Indeed, gating allows the network to reduce the weight of some edges and in the limit effectively re-wire the graph by completely or partially ignoring some edges. We argue that it is possible for gating to act as a graph sparsification device, which decreases the graph curvature and have been shown to alleviate over-smoothing (Giraldo et al., 2022; Rong et al., 2020).

Virtual node In addition, we use a virtual node (VN) (Gilmer et al., 2017) to enable global message passing; as the virtual node has access to the states of all nodes, it allows for (a) better representation of graph-level information and (b) faster propagation of information between nodes that are further apart, and thus faster convergence of states. In technical terms, adding the virtual node drastically increases the connectivity of the graph and in turn its curvature (Appendix C, Prop. C.4), and consequently decreases over-squashing. Alternatively, one can see that the Cheeger constant (another measure of bottleneckness (Topping et al., 2022)) of the graph increases after adding the virtual node.

Random node features One critical question is whether an MPNN is expressive enough to learn all the target PSEs. In particular, some PSEs, such as the Laplacian eigenvalues, may require distinguishability beyond 1-WL (Fürer, 2010). Despite the known 1-WL expressiveness limitation of a standard MPNN when using constant node features (Xu et al., 2019; Morris et al., 2019), random node features can help MPNNs surpass 1-WL expressiveness (Sato et al., 2021; Abboud et al., 2021; Kanatsoulis & Ribeiro, 2022). Thus, we base our encoder architecture on an MPNN coupled with random input node features, as shown in Figure 1A.

We empirically validate that together, the above architectural design choices lead to an effective graph encoder that finds the balance between smoothing and squashing (§3.5), and even has an elevated expressiveness due to random features (§3.3). A detailed ablation study highlighting the importance of our architectural choices is available in Table H.1.

Given a query graph structure $G=(V,E)$, we first generate a $k$-dimensional feature from a standard normal distribution for each node, $\mathbf{X}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, which is then passed through the GPSE model to extract the final representations (Figure 1A). The representations are then decoded into the target PSEs using multiple independent MLP heads, one per PSE (Figure 1B). We compute the reconstruction loss based on the sum of $\ell_{1}$ and cosine similarity losses (Appendix B.2), and optimize GPSE by minimizing this reconstruction loss. Details about hyperparameters can be found in Table B.4.

Loss function The aforementioned combination of $\ell_{1}$ and cosine similarity losses ensures that the model captures both (1) the direction of a particular PSE as a signal on the graph (via cosine similarity loss), and (2) the magnitude of the PSE (via $\ell_{1}$ loss). From a graph signal processing perspective, both types of information are crucial to describe the characteristics of PSEs in the form of a graph signal.

Training dataset PCQM4Mv2 (Hu et al., 2021) is a typical choice of pre-training dataset for molecular tasks. However, since GPSE only extracts features from graph structures (e.g., methane, CH4, would be treated as the same graph as silane, SiH4), the amount of training samples reduces to 273,920 after extracting unique graphs. Instead, we train GPSE with MolPCBA (Hu et al., 2020a) with 323,555 unique molecular graphs and an average number of 25 nodes. We randomly select 5% validation and 5% testing data fixed across runs, and use the remaining data for training GPSE. An ablation study on training datasets considered is available in Table H.5.

In these experiments, we demonstrate that GPSE provides more performance improvements over traditional PSEs for a wide range of GNN models. Additionally, we show competitive performance achieved by GPSE against the complementary self-supervised learning (SSL) pre-training approaches.

GPSE-augmented GPS is highly competitive on molecular graph benchmarks We compare the performance of the GPS model augmented with our GPSE encodings versus the same model using (a) no PSE, (b) random features as PSE, (c) LapPE and RWSE, and (d) concatenation of PSEs from §2.1 on four common molecular property prediction benchmarks (Dwivedi et al., 2022a; Hu et al., 2020a, 2021). For ZINC (Gómez-Bombarelli et al., 2018) and PCQM4Mv2 (Hu et al., 2020a), we use their subset versions following Dwivedi et al. (2022a) and Rampášek et al. (2022), respectively.

We first highlight that GPSE-augmented GPS achieves a remarkable 0.0648 MAE on ZINC (Table 2), not only significantly outperforming other PSEs but even challenging SOTA results. Similarly, GPS+GPSE also achieves the best result on PCQM4Mv2 amongst models that (a) are not ensemble methods and (b) do not have access to 3D information. These two strategies are highly engineering- and domain-oriented, thus do not reflect their utilities on general graph learning tasks, which we aim to demonstrate instead.

Moreover, we note that GPSE always performs better than, or at least on par with, standard PSEs, while concatenating multiple PSEs (AllPSE/LapPE+RWSE) always leads to worse performance than using individual PSEs. On ZINC & PCQM4Mv2, GPSE improves results comfortably beyond standard deviation, while in the worst case (e.g., MolHIV) it recovers the best PSE result, with differences well within a standard deviation. We discuss why GPSE works better on some datasets than others in Appendix I.

GPSE as a universal PSE augmentation The utility of GPSE encodings is not specific to GPS. In Table 2, we show that GPSE also significantly enhances two variants of SAT (Chen et al., 2022) on ZINC. Additionally, we show that augmenting different MPNN methods and the graph Transformer with GPSE universally results in remarkable improvements on ZINC: 56.24% reduction in test MAE on average compared to baselines that do not make use of any PSE, outperforming all other PSEs and their combinations (Table 3). Notably, concatenating LapPE and RWSE does not yield any benefit beyond using either one of them, demonstrating the intricacy of effectively leveraging information from multiple PSEs. We additionally perform a similar set of experiments on PCQM4Mv2 (Table H.2), and obtain the same improvements over explicitly computed PSEs, validating the success of GPSE further.

Feature augmentation using GPSE vs. SSL pre-training Our GPSE feature augmentation is related to self-supervised learning (SSL) pre-training approaches (Hu et al., 2020b; You et al., 2020b; Xie et al., 2022; Xu et al., 2021) in that both transfer knowledge from a large pre-training dataset to another for downstream evaluation. However, our approach is a substantial departure from previous SSL approaches in two distinct aspects:

1. 1.

   The trained GPSE model only serves as a feature extractor that can be coupled with any type of downstream prediction model, which will be trained from scratch.

2. 2.

   GPSE extracts representations solely from the graph structure and does not make use of the domain-specific features such as atom and bond types (Hu et al., 2020b), allowing GPSE to be utilized on any graph dataset.

To compare the performance of SSL pre-training and GPSE feature augmentation, we use the MoleculeNet (Wu et al., 2018; Hu et al., 2020b) datasets. For the downstream model, we use the identical GINE architecture (Hu et al., 2020b) from Sun et al. (2022). Finally, the extracted representations from GPSE are concatenated with the atom embeddings and are then fed into the GINE model.

We note that GPSE-augmented GINE achieves the best performance on three out of the eight MoleculeNet datasets against previously reported performances (Table 4). Moreover, GPSE augmentation improves performance over the baseline across all eight datasets, unlike some of the previously reported results that show negative transfer. Together, these results corroborate with the findings from Sun et al. (2022) that rich features can make up for the benefits of SSL pre-training. In our case, the GPSE encodings act as the rich features that contain positional and structural information from the graphs.

We also highlight that the Table 4 results are achieved in a setup where GPSE is at a comparative disadvantage: As a general-purpose feature extractor trained on a separate dataset, GPSE cannot leverage atom and bond features of the downstream graphs unlike typical molecular graph SSL methods. When GraphLoG (Xu et al., 2021) is similarly used as a feature extractor only, for a fair comparison, it is well outperformed by GPSE and even suffers from negative transfer, highlighting the power of GPSE as a feature extractor. With this in mind, GPSE can potentially be combined with other SSL methods to enhance them in future work.

GPSE can be used on arbitrary types of graphs as it is trained using the graph structures alone, in contrast to the SSL pre-training methods. Here, we show that GPSE is transferable to general graph datasets beyond molecular data, even under extreme out-of-distribution (OOD) cases.

Transferability to molecular graph sizes We use Peptides-struct and Peptides-func from the Long Range Graph Benchmark (Dwivedi et al., 2022c) to test whether GPSE can still work when the downstream (macro-)molecular graphs are substantially larger than those used for training. Despite this difference in graph sizes, GPSE outperforms explicitly computed PSEs when used with GPS as well as a GCN architecture optimized for long-range benchmarks by Tönshoff et al. (2023) (Table 5). More strikingly, GPS+GPSE challenges SOTA results for Peptides-struct, surpassing Graph MLP-Mixer (He et al., 2022). The improved performance emphasizes the ability of GPSE to better extract global information from query graphs by providing a more informative initial encoding for the global attention mechanism in GPS.

Transferability to graph connectivity patterns We further test if GPSE generalizes to graph connectivity patterns distinct from its training dataset. Particularly, we use the superpixel graph benchmarking datasets CIFAR10 and MNIST from Dwivedi et al. (2022a), which are $k$-nearest neighbor graphs with $k\!=\!8$, and employ connectivity patterns that significantly differ from those found in molecular graph datasets. Impressively, GPS+GPSE again achieves comparable results against GPS + computed PSEs (Table 5).

Transferability to extreme OOD node-classification benchmarks Taking the evaluation of GPSE one step further, we test its ability to provide useful information to transductive node classification tasks, where the graphs contain hundreds of thousands of nodes, which are completely out of distribution from the GPSE training dataset. We use both MPNN (GCN, GraphSAGE (Hamilton et al., 2017a), GATv2 (Veličković et al., 2018; Brody et al., 2022)) and graph Transformer baselines. Remarkably, on 8/10 model-dataset combinations, GPSE attains the best or equal-best result, as well as the best overall results, while LapPE performs below the no-encoding baseline for MPNNs (Table 6).

Meanwhile, the indifference in performance on the Proteins dataset for MPNNs is not unexpected, as the connectivity structures of the protein interaction network do not contribute to the proteins’ functions as meaningfully. Instead, the identity of the proteins’ interacting partners is of importance, commonly referred to as homophily in the graph representation learning community (Zhu et al., 2020) or more generally known as the Guilt-by-Association principle in network biology (Cowen et al., 2017). This result provides valuable insights into the usefulness of GPSE as an augmentation: It is more beneficial when the underlying graph structure is informative for the downstream tasks.

Given that GPSE can recover different PSEs so well (Table 1), it is natural to wonder whether it boosts standard MPNN expressiveness. We first confirm that GPSE encodings surpass 1-WL distinguishability by observing a clear visual separation of GPSE encodings on 1-WL indistinguishable graph pairs (Figure E.1). More information regarding graph isomorphism, the WL test and their connections to GNN expressivity is discussed in Appendix E.

To rigorously and systematically illustrate the expressivity of GPSE encodings further, we employ two synthetic benchmarks (Dwivedi et al., 2022a; Abboud et al., 2021) that require beyond 1-WL power, and use an MPNN model with 1-WL expressivity in GIN. Indeed, we find that GPSE provides extra power to the base MPNN to correctly distinguish graph isomorphism classes (Table 7). This expressivity boost by GPSE is particularly impressive, considering that (1) GPSE is pre-trained on MolPCBA, whose graph structures are not designed to be 1-WL indistinguishable like these synthetic graphs, and (2) naively adding random features to the input does not provide the same improvement.

We further point out that, in fact, augmenting the base GIN model using common PSEs like LapPE and RWSE readily archives nearly perfect graph isomorphism classification, corroborating with previous theoretical results on distance-encodings (Li et al., 2020) and spectral invariants (Fürer, 2010). This finding partially explains why GPSE provides additional power and also why previous methods using LapPE achieve perfect classification on these tasks (He et al., 2022). Finally, we note the performance difference between LapPE/RWSE and GPSE is not unexpected, as random input features only act as a patch to the MPNN expressivity limitation, rather than fully resolving it. Thus, developing more powerful and practically scalable GPSE models that losslessly capture the latent semantics of various PSEs is a vital avenue to explore in the future.

We demonstrate the efficiency and scalability advantages of GPSE over hand-crafted PSEs through two sets of scaling experiments. In the first, we compare the computation time for different PSEs against GPSE as we vary the number of graphs, represented as a percentage of the MolPCBA dataset. In the latter set of experiments, we instead generate four datasets of 1000 synthetic graphs, but scale up the individual graph sizes in each dataset. Our results are presented in Appendix G, where we show that GPSE is not only considerably faster to compute than explicitly computing PSEs, but also scales much better than explicit PSE computation as both the number of graphs and graph sizes increase.

The benefit of GPSE further compounds and leads to orders-of-magnitude faster computation times when we compute and combine all PSEs required for a fair comparison (AllPSE). One primary computational advantage of GPSE is that its complexity remains unchanged at inference time, regardless of the number of PSE types used to train the model. In our study, we restricted ourselves to six different PSEs, but future work could easily include more complex and specialized PSEs yet claim the same efficiency properties for PSE extraction. However, for AllPSE, the computational costs of stacking an increasing number of encodings would quickly become untenable.