SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ameya Velingker*
ameyav@google.com
Ali Kemal Sinop*
asinop@google.com
Ira Ktena
iraktena@google.com
Petar Velicković
petarv@google.com
Sreenivas Gollapudi
sgollapu@google.com
Graph Neural Networks (GNNs) have emerged as a powerful technique for learning on relational data. Owing to the relatively limited number of message passing steps they perform—and hence a smaller receptive field—there has been significant interest in improving their expressivity by incorporating structural aspects of the underlying graph. In this paper, we explore the use of affinity measures as features in graph neural networks, in particular measures arising from random walks, including effective resistance, hitting and commute times. We propose message passing networks based on these features and evaluate their performance on a variety of node and graph property prediction tasks. Our architecture has low computational complexity, while our features are invariant to the permutations of the underlying graph. The measures we compute allow the network to exploit the connectivity properties of the graph, thereby allowing us to outperform relevant benchmarks for a wide variety of tasks, often with significantly fewer message passing steps. On one of the largest publicly available graph regression datasets, OGB-LSC-PCQM4Mv1, we obtain the best known single-model validation MAE at the time of writing.
Graph Neural Networks (GNNs) constitute a powerful tool for learning meaningful representations in non-Euclidean domains. GNN models have achieved significant successes in a wide variety of node prediction [18, 31], link prediction [55, 52], and graph prediction [11, 50] tasks. These tasks naturally emerge in a wide range of applications, including autonomous driving [8], neuroimaging [36], combinatorial optimization [14, 34], and recommender systems [51], while they have enabled significant scientific advances in the fields of biomedicine [45], structural biology [23], molecular chemistry [42] and physics [3].
Despite the predictive power of GNNs, it is known that the expressive power of standard GNNs is limited by the 1-Weisfeiler-Lehman (1-WL) test [47]. Intuitively, GNNs possess at most the same power in terms of distinguishing between non-isomorphic (sub-)graphs, while having the added benefit of adapting to the given data distribution. For some architectures, two nodes with different local structures have the same computational graph, thus thwarting distinguishability in a standard GNN. Even though some attempts have been made to address this limitation with higher-order GNNs [32], most traditional GNN architectures fail to distinguish between such nodes.
A common approach to improving the expressive power of GNNs involves encoding richer structural/positional properties. For example, distance-based approaches form the basis for works such as
Position-aware Graph Neural Networks [52], which capture positions/locations of nodes with respect to a set of anchor nodes, as well as Distance Encoding Networks [29], which use the first few powers of the normalized adjacency matrix as node features associated with a set of target nodes.
Here, we take an approach that is inspired by this line of work but departs from it in some crucial ways: we seek to capture both distance and connectivity information using general-purpose node and edge features without the need for specifying any anchor or target nodes.
Contributions: We propose the use of affinity metrics as features in a GNN. Specifically, we consider statistics that arise from random walks in graphs, such as hitting time and commute time between pairs of vertices (see Sections 3.1 and 3.2). We present a means of incorporating these statistics as scalar edge features in a message passing neural network (MPNN) [15] (see Section 3.4). In addition to these scalar features, we present richer vector-valued resistive embeddings (see Section 3.3), which can be incorporated as node or edge feature vectors in the network. Resistive embeddings are a natural way of embedding each node into Euclidean space such that the squared  $L^2$ -distance between nodes recovers the commute time. We show that such embeddings can be incorporated into MPNNs, even for larger graphs, by efficiently approximating them using sketching and dimensionality reduction techniques and prove a novel additive approximation for hitting time (see Section 4).
Moreover, we evaluate our networks on a number of benchmark datasets of diverse scales (see Section 5). First, we show that our networks outperform other baselines on the PNA dataset [9], which includes 6 node and graph algorithmic tasks, showing the ability of affinity measures to exploit structural properties of graphs. We also evaluate the performance on a number of graph and node tasks for datasets in the Open Graph Benchmark (OGB) collection [21], including molecular and citation graphs. In particular, our networks with scalar effective resistance edge features achieve the state of the art on the OGB-LSC PCQM4Mv1 dataset, which was featured in a KDD Cup 2021 competition for large scale graph representation learning.
Finally, we provide intuition for why affinity-based measures are fundamentally different from aforementioned distance-based approaches (see Section 3.5) and bolster it with detailed theoretical and empirical results (see Appendix D) showing favorable results for affinity-based measures.
Our work builds upon a wealth of graph theoretic and graph representation learning works, while we focus on a supervised, inductive setting.
Even though GNN architectures were originally classified as spectral or spatial, we abstain from this division as recent research has demonstrated some equivalence of the graph convolution process regardless of the choice of convolution kernels [2, 7]. Spectrally-motivated methods are often theoretically founded on the eigendecomposition of the graph Laplacian matrix (or an approximation thereof) and, hence, corresponding convolutions capture different frequencies of the graph signal. Early works in this space include ChebNet [10] and its more efficient 1-hop version by Kipf et al. [24], which offers a linear function on the graph Laplacian spectrum. Levie et al. [28] proposed CayleyNets, an alternative rational filter.
Message passing neural networks (MPNNs) [15] perform a transformation of node and edge representations before and after an arbitrary aggregator (e.g. sum). Graph attention networks (GATs) [44] aimed to augment the computations of GNNs by allowing graph nodes to "attend" differently to different edges, inspired by the success of transformers in NLP tasks. One of the most relevant works was proposed by Beaini et al. [4], i.e. directional graph networks (DGN). DGN uses the gradients of the low-frequency eigenvectors of the graph Laplacian, which are known to capture key information about the global structure of the graph and prove that the aggregators they construct using these gradients lead to more discriminative models than standard GNNs according to the 1-WL test. Prior work [32] used higher-order ( $k$ -dimensional) GNNs, based on  $k$ -WL, and a hierarchical variant and proved theoretically and experimentally the improved expressivity in comparison to other models.
Other notable works include Graph Isomorphism Networks (GINs) [47], which represent a simple, maximally-powerful GNN over discrete-featured inputs. Hamilton et al. [18] proposed a method to construct node representations by sampling a fixed-size neighborhood of each node and then performing aggregation over it, which led to impressive performance on large-scale inductive benchmarks. Bouritsas et al. [6] use topologically-aware message passing to detect and count graph substructures,
while Bodnar et al. [5] propose a message-passing procedure on cell complexes motivated by a novel color refinement algorithm that proves to be powerful for molecular benchmarks. Meanwhile, Horn et al. [19] propose a layer based on persistent homology to incorporate global topological information.
**Expressivity.** Techniques that improve a GNN's expressive power largely fall under three broad directions. While we focus on the feature-based direction in this paper, we also acknowledge that it in no way compels the GNN to use the additional provided features. Hence, we briefly survey the other two, as an indication of the future research in affinity-based GNNs we hope this work will inspire.
Another avenue involves modulating the message passing rule to make advantage of the desired computations. Popular recent examples of this include DGN [4] and LSPE [12]. DGNs leverage the graph's Laplacian eigenvectors, but they do not merely use them as input features; instead, they define a directional vector field based on the eigenvectors, and use it explicitly to anisotropically aggregate neighbourhoods. LSPE features a "bespoke pipeline" for processing positional inputs.
The final direction is to modulate the graph over which messages are passed, usually by adding new nodes that correspond to desired substructures. An early proponent of this is the work of [32], which explicitly performs message passing over  $k$ -tuples of nodes at once. Recently, scalable efforts in this direction focus on carefully chosen substructures, e.g., junction trees [13], cellular complexes [5].
3 Affinity Measures and GNNs
3.1 Random Walks, Hitting and Commute Times
Let  $G = (V, E)$  be a graph of vertices  $V$  and edges  $E \subseteq V \times V$  between them. We define several natural properties of a graph that arise from a random walk. A random walk on  $G$  starting from a node  $u$  is a Markov chain on the vertex set  $V$  such that the initial vertex is  $u$ , and at each time step, one moves from the current vertex to a neighbor, chosen with probability proportional to the weight of outgoing edges. We will use  $\pi$  to denote the stationary distribution of this Markov Chain. For random walks on weighted, undirected graphs, we know that  $\pi_u = \frac{d_u}{2M}$ , where  $d_u$  is the weighted degree of node  $u$ , and  $M$  is the sum of edge weights.
The hitting time  $H_{uv}$  from  $u$  to  $v$  is defined as the expected number of steps for a random walk starting at  $u$  to hit  $v$ . We can also define the commute time between  $u$  and  $v$  as  $K_{uv} = H_{uv} + H_{vu}$ , the expected round-trip time for a random walk starting at  $u$  to reach  $v$  and then return to  $u$ .
3.2 Effective Resistance
A closely related quantity is the measure of effective resistances in undirected graphs. This quantity corresponds to the effective resistance if the whole graph was replaced with a circuit where each edge becomes a resistor with resistance equal to the reciprocal of its weight. We will use  $\mathsf{Res}(u,v)$  to denote the effective resistance between nodes  $u$  and  $v$ . For undirected graphs, it is known that [30] the effective resistance is proportional to the commute time,  $\mathsf{Res}(u,v) = \frac{1}{2M} K_{uv}$ .
Our broad goal is to incorporate effective resistances and hitting times as edge features in an MPNN. In Section 3.4, we will show they can provably improve MPNNs' expressivity.
3.3 Resistive Embeddings
Effective resistances allow us to define the resistive embedding, a mapping that associates each node  $v$  of a graph  $G = (V, E, W)$ , where  $W$  are the non-negative edge weights, with an embedding vector. Before we specify the resistive embedding, we define a few terms. Let  $L = D - A$  be the graph Laplacian of  $G$ , where  $D \in \mathbb{R}^{n \times n}$  is the diagonal matrix containing the weighted degree of each node and  $A \in \mathbb{R}^{n \times n}$  is the adjacency matrix, whose  $(i, j)^{th}$  entry is equal to the edge weight between  $i$  and  $j$ , if it exists; and 0 otherwise. Let  $B$  be the  $m \times n$  edge-node incidence matrix, where  $|V| = n$  and  $|E| = m$ , defined as follows: The  $i$ -th row of  $B$  corresponds to the  $i$ -th edge  $e_i = (u_i, v_i)$  of  $G$  and has a  $+1$  in the  $u_i$ -th column and a  $-1$  in the  $v_i$ -th column, while all other entries are zero. Finally we will use  $C \in \mathbb{R}^{m \times m}$  to denote the conductance matrix, which is a diagonal matrix with  $C_{ii}$  being the weight of  $i^{th}$  edge. It is easy to verify that  $B^T C B = L$ . Even though  $L$  is not invertible, its null-space consists of the indicator vectors for every connected component of  $G$ . For example, if  $G$  is
connected, then  $L$ 's nullspace is spanned by the all-1's vector  $^2$ . Hence, for any vector  $x$  orthogonal to all-1's,  $L \cdot L^{\dagger}x = x$ , where  $L^{\dagger}$  is the pseudo-inverse.
We can express effective resistance between any pair of nodes using Laplacian matrices [30] as  $\operatorname{Res}(u, v) = (\mathbf{1}_u - \mathbf{1}_v)^T L^\dagger (\mathbf{1}_u - \mathbf{1}_v)$ , where  $\mathbf{1}_v$  is an  $n$ -dimensional vector specifying the indicator for node  $v$ . We are now ready to define the resistive embedding.
Definition 3.1. (Effective Resistance Embedding)  $\mathbf{r}_v = C^{1/2}BL^\dagger\mathbf{1}_v$
A key property is that the effective resistance between two nodes in the graph can be obtained easily from the distance between their corresponding embeddings (see proof in Appendix C):
Lemma 3.2. For any pair of nodes  $u, v$ , we have  $\| \mathbf{r}_u - \mathbf{r}_v \|_2^2 = \operatorname{Res}(u, v)$ .
One can easily check that any rotation of  $\mathbf{r}$  also satisfies Lemma 3.2, since rotations preserve Euclidean distances; more generally, if  $U$  is an orthonormal matrix, then  $U\mathbf{r}$  is also a valid resistive embedding. This poses a challenge if we want to use the resistive embeddings as node or edge features: we want a way to enforce that a (G)NN using them will do so in a way that is invariant or equivariant to any rotations of the embeddings. In our current work, we rely on data augmentation: at every training iteration, we apply random rotations to the input ER embeddings.
Remark. While data augmentation is a popular approach for promoting invariant and equivariant predictions, it is only hinting to the network that such predictions are favourable. It is also possible, in the spirit of the geometric deep learning blueprint [7], to combine ER embeddings with an  $O(n)$ -equivariant GNN, which rigorously enforces rotational equivariance. A popular approach to building equivariant GNNs has been proposed by [39], though it focuses on the full Euclidean group  $E(n)$  rather than  $O(n)$ . We leave this exploration to future work.
Definition 3.3. Let  $\mathbf{p} \coloneqq \sum_{u} \pi_{u} \mathbf{r}_{u}$  be the mean of effective resistance embedding.
We might view  $\mathbf{p}$  as a "weighted mean" of  $\mathbf{r}$ . We will define the hitting time radius,  $H_{\mathrm{max}}$ , of a given graph as the maximum hitting time between any two nodes:
Definition 3.4 (Hitting Time Radius).  $H_{\mathrm{max}} \coloneqq \max_{u,v} H_{u,v}$
We will need the following to bound the hitting times we computed:
Lemma 3.5. For any node  $u$ ,  $\| \mathbf{r}_u - \mathbf{p}\|^2 \leq \frac{H_{\max}}{M}$ .
The proof follows easily from the fact that  $\mathbf{p}$  is a convex combination of all  $\mathbf{r}$ 's and Jensen's inequality.
3.4 Incorporating Features into MPNNs
We reiterate that our main aim is to demonstrate (theoretically and empirically) that there are good reasons to incorporate affinity-based measures into GNN computations.
In the simplest instance, a method that improves a GNN's expressive power may compute additional features (positional or structural) which would assist the GNN in discriminating between examples it otherwise wouldn't (easily) be able to. These features are then appended to the GNN's inputs for further processing. For example, it has been shown that endowing nodes with a one-hot based identity is already sufficient for improving expressive power [33]; this was then relaxed to any randomly-sampled scalar feature by [38]. It is, of course, possible to create dedicated features that even count substructures of interest [6]. Further, the adjacency information can be factorised [37] or eigendecomposed [12] to provide useful structural embeddings for the GNN.
We will focus our attention on exactly this class of methods, as it is a lightweight and direct way of demonstrating improvements from these computations. Hence, our baselines will all be instances of the MPNN framework [15], which we will attempt to improve by endowing them with affinity-based features. We start by theoretically proving that these features indeed improve expressive power:
Theorem 3.6. MPNNs that make use of any one of (a.) effective resistances, (b.) hitting times, (c.) resistive embeddings are strictly more powerful than the WL-1 test.
2Throughout this section, we will assume that our graph is connected. However everything applies to disconnected graphs, too.
Note that the average of all  $\mathbf{r}_u$ 's will be 0. If the graph is regular, then  $\mathbf{p}$  will also be 0.
Proof. Since the networks in question arise from augmenting standard MPNNs with additional node/edge features, we have that these networks are at least as powerful as the 1-WL test.
In order to show that these networks are strictly more powerful than the 1-WL test, it suffices to show the existence of a graph for which our affinity measure based networks can distinguish between certain nodes that a standard GNN (limited by the 1-WL test) cannot.
We present an example of a 3-regular graph on 8 nodes in Figure 1. It is well-known that a standard GNN that is limited by the 1-WL test cannot distinguish any pair of nodes in a regular graph, as the computation tree rooted at any node in the graph looks identical. However, there are three isomorphism classes of nodes in the above graph (denoted by different colors), namely,  $V_{1} = \{1,2\}$ ,  $V_{2} = \{3,4,7,8\}$ , and  $V_{3} = \{5,6\}$ .
We now show that GNNs with affinity based measures can distinguish between a node in  $V_{i}$  and a node in  $V_{j}$ , for  $i \neq j$ . We note that the hitting time from  $a$  to  $b$  depends only on the isomorphism classes of  $a$  and  $b$ . Thus, we write  $r_{i,j}$  as the effective resistance between a node in  $V_{i}$  and a node in  $V_{j}$ . Note that  $r_{i,j} = r_{j,i}$ , and it is easy to verify that:
$$
\begin{array}{l} r _ {1, 1} = 2 / 3, r _ {2, 2} = 1 5 / 2 8, r _ {3, 3} = 4 / 7 \\ r _ {1, 2} = r _ {2, 1} = 1 8 5 / 3 3 6 \\ r _ {2, 3} = r _ {3, 2} = 2 0 9 / 3 3 6. \\ \end{array}
$$
Hence, it follows that in a message passing step of an MPNN that uses effective resistances, vertices in  $V_{1}$ ,  $V_{2}$ , and  $V_{3}$  will aggregate feature multisets  $\{r_{1,1}, r_{1,2}, r_{1,2}\} = \{2/3, 185/336, 185/336\}$ ,  $\{r_{2,1}, r_{2,2}, r_{2,3}\} = \{185/336, 15/28, 209/336\}$ , and  $\{r_{2,3}, r_{2,3}, r_{3,3}\} = \{209/336, 209/336, 4/7\}$ , respectively, all of which are all distinct multisets. Hence, such an MPNN can distinguish nodes in  $V_{i}$  and  $V_{j}$ ,  $i \neq j$  for a suitable aggregation function.
If, instead of effective resistance, we use hitting time features or resistive embeddings, our results still hold. This is because, as we showed previously, the effective resistance between nodes is a function of the two hitting times in either direction, as well as of the resistive embeddings of the nodes. In other words, if either hitting time features or resistive embeddings are used as input features for an MPNN, this MPNN would be able to compute the effective resistance features by applying an appropriate function (e.g., Lemma 3.2 for the case of resistive embeddings). Having computed these features, the MPNN can distinguish any two graphs that the MPNN with effective resistance features can.  $\square$
3.5 Effective Resistance vs. Shortest Path Distance
It is interesting to ask how effective resistances compare with shortest path distances (SPDs) in GNNs, given the plethora of recent works that make use of SPDs (e.g., [49, 52, 29]). The most direct comparison of our effective resistance-based MPNNs would be to use SPDs as edge features in the MPNNs. However, note that SPDs along graph edges are trivial (unlike effective resistances, which incorporate useful information about the global graph structure).
An alternative to edge features would be to use (a) SPDs to a small set of anchor nodes as features in an MPNN (e.g., P-GNN [52]) or (b) a dense featurization incorporating shortest paths between all pairs of nodes (e.g., the dense attention mechanism in Graphormer [49]). We remark that the latter approach typically incurs an  $O(n^{2})$  overhead, which our MPNN-based approach avoids.
We empirically compare our models to MPNNs that use approaches (a) and (b). Results on the PNA dataset show that our effective resistance-based GNNs outperform these approaches. Furthermore, we complement the empirical results with a theoret-
ical result (see Appendix D) showing that under a limited number of message-passing steps, effective resistance features can allow one to distinguish structures that cannot be done using shortest path features. This provides some insight into why effective resistances can capture structure in GNNs
Figure 1: Degree 3 graph on 8 nodes, with isomorphism classes indicated by colors. While nodes of the same color are structurally identical, nodes of different colors are not. A standard GNN limited by the 1-WL cannot distinguish between nodes of different colors. However, affinity based networks that use effective resistances, hitting times, or resistive embeddings can distinguish every pair of such nodes.
that SPDs are unable to. We describe the result below. We point the reader to Appendix D for a proof of Theorem D.1.
4 Efficient Computation of Affinity Measures
In order to use our features, it is important that they be computable efficiently. In this section, we show how to compute or approximate the various random walk-based affinity measures. We share the results of our method on large-scale graphs in Section 5.4.
4.1 Reducing Dimensionality of Resistive Embeddings
Given their higher dimensionality, we might find it beneficial to use resistive embeddings as GNN features instead of the effective resistance. Now, the difficulty with using resistive embeddings directly as GNN features is the fact that the embeddings have dimension  $m$ , which can be quite large, e.g., up to  $n^2$  for dense graphs. It was shown by [41] that one can reduce the dimensionality of the embed
Table 1:  $\log \left( {MSE}\right)$  on the PNA test dataset.
Model	Node tasks	Graph tasks
Avg score	SSSP	Ecc	Lap feat	Conn	Diam	Spec rad
GAT	-1.730	-2.213	-1.935	-2.644	-0.618	-1.430	-1.538
GCN	-1.592	-2.283	-1.978	-1.698	-0.618	-1.432	-1.541
MPNN	-2.665	-2.235	-2.419	-3.116	-1.887	-2.681	-3.652
MPNN (rand. feat.)	-2.490	-2.136	-1.808	-3.873	-1.696	-2.614	-2.813
DGN (features)	-2.743	-2.165	-1.911	-4.184	-1.858	-2.814	-3.528
DGN (features) + ER embed + HT	-2.938	-2.360	-2.949	-3.689	-1.744	-3.060	-3.823
ER GNN	-2.779	-2.146	-1.869	-3.945	-1.962	-2.940	-3.811
ER (node) embed.	-2.658	-2.245	-2.493	-3.533	-1.649	-2.886	-3.144
ER (edge) embed.	-2.789	-2.266	-2.125	-4.253	-1.664	-2.807	-3.617
Hitting Times	-2.816	-2.189	-1.904	-4.397	-1.888	-2.796	-3.720
All ER features	-3.106	-2.789	-3.082	-4.047	-1.858	-2.894	-3.962
ding while approximately preserving Euclidean distances. The idea is to use a random projection via a constructive version of the Johnson-Lindenstrauss Lemma:
Lemma 4.1 (Constructive Johnson-Lindenstrauss). Let  $x_{1}, x_{2}, \ldots, x_{n} \in \mathbb{R}^{d}$  be a set of  $n$  points; and let  $\alpha_{1}, \alpha_{2}, \ldots, \alpha_{m} \in \mathbb{R}^{n}$  be fixed linear combinations. Suppose  $\Pi$  is a  $k \times d$  matrix whose entries are chosen i.i.d. from a Gaussian  $N(0,1)$  and consider  $\hat{x}_i := \frac{1}{\sqrt{k}} \Pi x_i$ .
Then, it follows that for  $k \geq C \log (mn) / \epsilon^2$ , with probability  $1 - o(1)$ , we have  $\| \sum_{j} \alpha_{i,j} \widehat{x}_j \|^2 = (1 \pm \epsilon) \| \sum_{j} \alpha_{i,j} x_j \|_2^2$  for every  $1 \leq i \leq m$ .
Using triangle inequality, we can see that the inner products between fixed linear combinations are preserved up to some additive error: