Recurrent Distance-Encoding Neural Networks for Graph Representation Learning

We thank reviewer hzhG for taking the time to review our work, and are pleased she/he found our work to be well-written and our experiments to be good. We address the raised concerns as follows:

Q1. RNN here is strange, the assumption here is that the information flow from K-hop to K-1 hop. but information can also go "outward", from k-2 hop to k-1 hop. RNN is not good to model the graph information.

Ans: It’s true that each edge in the graph allows “bidirectional” information flow. However, if we focus on a particular node for which we want to generate a representation, it’s natural to categorize the other nodes into “shells” (Figure 1(b)). Nodes that belong to the same shell have the same shortest distance from the target node (considering running a breadth-first search from the target node). Then a recurrent neural network, which runs from the outermost shell to the innermost one, would be a reasonable choice to encode the neighborhood structure.

Besides, our model indeed allows “bidirectional” information flow. For two different target nodes (Node 1 and Node 2), their shells are different, the K-1 hop and K hop neighbors of Node 1 can become the K hop and K-1 hop respectively for Node 2, and the information propagates reversely when we generate the representation for Node 2. Therefore, considering the set of node representations as a whole, we won’t lose any structural information. This argument is supported by our rigorous theoretical analysis of expressiveness

Q2. I assume all the GNN drawback will also be here in this Model (over smoothing, etc). The model can be regarded as another interpretation of massage passing(even though the weight is different), for the target node, the model will aggregate all the K-hop node information.

Ans: GRED doesn’t have the drawbacks of MPNN because the computational graphs of GRED and MPNN are different. We have added a straightforward example to Appendix E of the updated PDF to illustrate their difference. To aggregate all nodes in the K-hop neighborhood, MPNN needs to perform message passing for K rounds. As a result, MPNN’s computational graph for the target node grows exponentially with K because the same neighbor will appear many times in the computational graph. The exponentially growing number of nodes are mixed together by MPNN to generate the representation for the target node, resulting in the known drawbacks (e.g., over-smoothing, over-squashing). On the contrary, the computational graph of GRED includes each node in the K-hop neighborhood exactly once, and GRED further utilizes LRU to control how signals of distant nodes are propagated toward the target node. Therefore, GRED can more effectively leverage information beyond the local neighborhood than MPNN.

Q3. how the performance like if we change LRU to the vanilla RNN. seems LRU for the long range information pass is the key factor to win against MPNN.

Ans: We have conducted experiments to compare LRU with vanilla RNN:

We replace the LRU component of GRED with a vanilla RNN:

$

\mathbf{s}_{v, k}^{(l)} = \text{tanh} ( \mathbf{W}\_{\text{rec}} \mathbf{s}\_{v, k - 1}^{(l)} + \mathbf{W}\_{\text{in}}\mathbf{x}\_{v, K - k}^{(l)})

$

where $\mathbf{W}\_{\text{rec}} \in \mathbb{R}^{d_s \times d_s}$ and $\mathbf{W}\_{\text{in}} \in \mathbb{R}^{d_s \times d}$ are two trainable real weight matrices. We use the same number of layers and the same $K$ for both models. We can observe that the performance of GRED with a vanilla RNN drops significantly. On CIFAR10 and ZINC where $K$ is not large, GRED$\_{\text{RNN}}$ still outperforms the best MPNN, but on Peptides-func where long-range information is needed (8 layers with $K=40$ per layer), the performance of GRED$\_{\text{RNN}}$ is worse than the best MPNN because vanilla RNNs are known to be difficult to train on long sequences, unlike the LRU.

We kindly ask the reviewer to consider raising her/his score if she/he thinks our response can address the concerns. We would be happy to answer further questions!

We thank reviewer 8Mos for taking the time to review our work and are pleased she/he found: 1) our presentation to be very good | clear; 2) our theory to be interesting | rigorous; and 3) our empirical results to be well-organized and contributing to a comprehensive understanding of GRED's performance.

We admit that we missed several references on K-hop GNNs, and we agree that a more detailed discussion of these methods would help improve the understanding of our contributions. We address these concerns as follows:

Q1. The novelty of the proposed architecture is somewhat limited | Lack of a comparative analysis of existing K-hop GNNs.

Ans: Compared with existing K-hop GNNs, the novelty of our proposed architecture is a new way to encode the K-hop neighborhood, i.e., with LRU. The use of LRU ensures our model theoretically sound expressiveness as well as ability to address the over-squashing problem in long-range reasoning tasks. To see why the novelty of our model is significant, we provide a comparative analysis of existing K-hop GNNs here:

Among K-hop GNNs, MixHop [3] uses powers of the adjacency matrix to access k-hop nodes. k-hop GNN [2] iteratively applies MLPs to combine two consecutive hops and propagates information towards the target node. While they have shown that higher hop information can improve the expressiveness of MPNNs, they suffer from over-squashing, i.e., they cannot effectively utilize information from distant nodes because the signals of distant nodes would decay rapidly due to iterative neighborhood mixing. As a result, the largest neighborhood these approaches operate on is 3-hop. The drawback of these approaches is also pointed out by the SPN paper [1]. The SPN paper [1] proposes a conceptual update formula which is similar to ours:

$$h\_u^{(t+1)} = \text{COM} ( h\_u^{(t)}, \text{AGG}\_{u, 1}, …, \text{AGG}\_{u, K})$$

However, SPN simply uses weighted summation as COM to aggregate K hops. The weighted summation cannot guarantee the expressiveness of the model since it cannot uniquely represent sequences. On the contrary, the LRU component of GRED can achieve expressiveness and long-range modeling capacity simultaneously: LRU can express an injective mapping of sequences, and at the same time prevent distant information from decaying by learning proper eigenvalues of the transition matrix. We have updated the related work to include a comparison of K-hop GNNs (marked in red in the updated PDF)

To validate that GRED can encode the K-hop neighborhood more effectively than SPN, We evaluate GRED on two additional datasets from TUDataset and compare against SPN. As shown in Part III of “Additional Experimental Results”, GRED can outperform SPN with the same number of hops.

[1] Shortest Path Networks for Graph Property Prediction. Learning on Graphs Conference. 2022

[2] k-hop graph neural networks. Neural Networks. 2020

[3] MixHop: Higher-Order Graph Convolutional Architectures via Sparsified Neighborhood Mixing. ICML. 2019

We thank reviewer e7jX for taking the time to review our work, and we are pleased she/he found our evaluation to be very thorough and our results to validate the effectiveness of our model. We address the raised concerns as follows:

Q1. A sensitivity analysis of GRED with respect to the choice of K is missing.

Ans: This is now in Appendix D of the updated PDF. Please refer to Part I of our “Additional Experimental Results” thread in our general response for a thorough analysis.

Q2. The training efficiency analysis is conducted against a transformer model only, …, each node will have its own sequence of sets of nodes at distance k, so there might not be any shared computation that can be leveraged…

Ans: It’s true that each node has its own sequence of set representations, but the computations for all nodes in the graph can be done in batches. In the implementation, we use a matrix mask to indicate each node’s $k$-hop neighbors (just like the adjacency matrix indicates each node’s 1-hop neighbors). These masks are computed only once in preprocessing using the shortest path algorithm, and they are used to do aggregation at each hop. Then the input to LRU would be a batch of sequences of length K+1 (including the target node for each sequence). We have added a code snippet of GRED in Appendix B of the updated PDF, to show how the computation is done.

To compare the training efficiency of GRED with MPNN, we use the same number of layers but the MPNN counterpart only aggregates 1-hop neighbors at each layer. We measure the average training time per epoch:

We can observe that on ZINC and CIFAR10, the additional time required to use a larger neighborhood is marginal. On Peptides-func, the difference is greater since GRED uses $K=40$. However, this extra time, which is still less than the time required by GRIT, is acceptable considering the significant performance improvement

Q3. How would this method perform on non long-range benchmarks? Would it underperform compared to MPNNs who might only need a local receptive field to solve a task?

Ans: The datasets in Table 1 and Table 2 (ZINC) are from [4], which is a widely used benchmark. These datasets are not particularly long-range. PATTERN and CLUSTER model communities in social networks and all nodes can be reached within 3 hops. MNIST and CIFAR10 are nearest-neighbor graphs of pixels, and just like their normal versions local information is important. Our model performs well on these datasets. We further evaluate our model on two additional datasets from TUDataset (please refer to Part III of our “Additional Experimental Results” thread) and our model also achieves good performance.

Our model can adapt to both long-range and non-long-range tasks by learning proper eigenvalues of the LRU transition matrix that control how fast the filters decay as the distance increases. For a long-range task, the learned eigenvalues are close to 1 which better keeps distant information, while for a non-long-range task, the magnitudes of the learned eigenvalues are smaller (please refer to Figure 3 for an illustration of the eigenvalues after training). Additionally, the choice of $K$ gives our model some flexibility to adjust the receptive field, and we found the optimal $K$ for a non-long-range task lies between 1 and the diameter (Part I of “Additional Experimental Results”)

[4] Benchmarking Graph Neural Networks. JMLR. 2022

We kindly ask the reviewer to consider raising her/his score if she/he thinks our response can address the concerns. We would be happy to answer further questions!

We thank reviewer bs5m for taking the time to review our work and are pleased she/he found our work to be well-written and easy to understand. First, we would like to ask the reviewer to take a look at our post “Additional Experimental Results” for updated results and the sensitivity analysis with respect to the value of $K$. We address the raised concerns as follows:

Q1. What is the "skip" in Figure 1(a)?

Ans: The “skip” in Figure 1(a) denotes an identity branch. We have added a code snippet of GRED to Appendix B of the updated PDF, to show how the computation is done exactly.

Q2. For the training efficiency evaluation, it would be nice to include the memory consumption for different methods. Also, the performance gap on the task between GRIT and GRED is still significant, which makes the comparison a bit unfair. It would be better to compare with the architecture that has similar task performance.

Ans: We have added the GPU memory consumption to Table 4 of the updated PDF. We also post it here:

The gap between the performance of GRIT and the updated performance of GRED is much less significant now. Especially on Peptides-func with $K=40$ GRED can outperform GRIT. GRED still maintains higher training efficiency than GRIT.

Q3. The method introduces a new hyperparameter K to be tuned. It would be nice to show how K is selected, and how the different selections of K can affect the task performance.

Ans: We have conducted experiments on how different $K$ values affect the performance of GRED. Please refer to Part I of our “Additional Experimental Results” thread in our general response for a thorough analysis.

Q4. Although the technique helps increase the range of nodes the model can process, unlike the graph transformer-based method, the distance is still limited by the selection of K and the number of layers of the model. It would be good to show some insights into how that super long-distance information affects the model performance.

Ans: It’s true that the range of nodes each layer can access is bounded by the value of $K$. However, for real-world datasets, the graph diameter won’t increase linearly with the number of nodes. Therefore, even if the value of $K$ doesn’t seem “super long”, each layer of our model can access all nodes in the graph.

For all datasets we used in experiments, except for Peptides-func and Peptides-struct, our model fits into a single 24GB GPU even when $K$ is set to $K_{\text{max}}$ (i.e., the maximum diameter of all graphs in the dataset). For Peptides-func and Peptides-struct, we use a $K$ value which is smaller than $K_{\text{max}}$ in order to increase the depth of the model. In this sense, the choice of $K$ allows the model some flexibility when the hardware resource is constrained. In all cases, we ensure that model depth * $K$ is strictly larger than $K_{\text{max}}$, so the receptive field of the entire GRED model is the same as graph transformers and GRED is able to extract global information.

Q5. Two related papers

Ans: We have added these two references to the related work (marked in red) of the updated PDF.

We kindly ask the reviewer to consider raising her/his score if she/he thinks our response can address the concerns. We would be happy to answer further questions!

In this thread, we present additional experimental results requested by the reviewers. In part I, we show how different $K$ values (number of hops in the recurrence) affect the performance of our model. In part II, we replace the LRU component with a vanilla RNN and show the difference in performance (confirming the effectiveness of our LRU choice). In part III, we evaluate our model on two additional datasets from TUDataset. Our model generalizes well to TUDataset and outperforms another related approach SPN [1]. All the above results have been added to Appendix D of the updated PDF.

Moreover, with further hyper-parameter tuning, we have improved the performance of our model on CIFAR10 and Peptides-func: for CIFAR10 we increase the model depth with a smaller hidden dimension, and for Peptides-func we increase $K$ to 40. The hyper-parameters for our model are summarized in Appendix C. Especially, when $K=40$ the performance of GRED on Peptides-func (0.7041±0.0049) is better than the best graph transformer GRIT (0.6988±0.0082), and GRED still maintains higher efficiency than GRIT (158.9s/epoch vs 225.6s/epoch). This new result is impressive given that GRED doesn’t use any positional encoding and further validates that the architecture of GRED alone can encode the structural information.

Part I: Effect of $K$ on performance

Effect of $K$ on the performance on CIFAR10:

Effect of $K$ on the performance on ZINC:

Effect of $K$ on the performance on Peptides-func:

We use $K_{\text{max}}$ to denote the maximum diameter of all graphs in the dataset. For Peptides-func, the maximum $K$ we tried was slightly smaller than $K_{\text{max}}$ in order to fit the model into a single RTX A5000 GPU with 24GB memory. From the three tables, we can observe that larger $K$ values generally yield better performance. On CIFAR10 and ZINC, while directly using $K_{\text{max}}$ already outperforms MPNNs, the optimal value of $K$ yielding best performance lies strictly between $1$ and $K_{\text{max}}$. We think this is because information that is too far away is less important for these two tasks: indeed, interestingly, the best $K$ value for CIFAR10 is similar to the width of a convolutional kernel on a normal image. On Peptides-func, the change of performance is more monotonic in $K$. When $K=40$, GRED outperforms the best graph transformer GRIT. We observe no further performance gain when we increase $K$ to 60.

\mathbf{s}_{v, k}^{(l)} = \text{tanh} ( \mathbf{W}\_{\text{rec}} \mathbf{s}\_{v, k - 1}^{(l)} + \mathbf{W}\_{\text{in}}\mathbf{x}\_{v, K - k}^{(l)})