Towards characterizing the value of edge embeddings in Graph Neural Networks

We thank the reviewer for their thoughtful review and we appreciate the positive comments! Below, we address the reviewer’s concerns and questions:

Efficient approaches that close the separation between edge-based and node-based? Thanks for pointing us to the reference by Topping et al. on graph rewiring; we’ll add it to the discussion. We agree that it would be nice to design more efficient approaches to close the separation. There are a variety of interesting heuristics for improving GNNs, though the trade-offs of such approaches are still not theoretically fully-understood. In particular, it would be very exciting to understand what representational effects (i.e. on required depth/width of a GNN) such heuristics have for specific choices of graphs and tasks.

Does the depth separation translate to time complexity separation (Question 1)? Yes, there is a separation in parallel time complexity (with one processor per node/edge in the respective models), modulo an assumption that each processor reads all of its neighbors’ input at each round. In the edge model, the procedure can be implemented in time O(n) as the reviewer states. In the node model, we prove $TB \geq \sqrt{n}$. Since the hub node has $O(n)$ neighbors, each round requires parallel time at least $O(nB)$ if the hub node reads the $B$ bits for each neighbor, so the overall parallel time complexity is at least $O(nTB) = O(n^{3/2})$. We will clarify this point in Remark 4.

Remark about input features in WL test? We agree that the k-WL refinement procedure can utilize informative features, and we apologize for the imprecise wording of Remark 10.

In Remark 10 we were trying to briefly make a somewhat subtle point (which we will expand on): prior works on expressivity of GNNs (w.r.t. the WL test) measure expressivity by asking “for a particular input, what are the possible outputs” (and they show this is characterized by number of WL refinement steps). However, particularly for GNNs that take as input both a graph and informative input labels, we would argue the central representation question is what functions the GNN can represent, i.e. “what are the possible mappings from inputs to outputs”. This is what matters for downstream learning tasks, and what motivated our framework (see also the paragraph “GNNs as a computational machine” in Sec. 3). This is also analogous to the classical representational theory for standard neural networks.

Implication for virtual nodes (Question 2)? Thanks for bringing up virtual nodes and the reference to Cai et al.; this is an interesting connection and we will add a remark to the paper. Since our main construction uses a node that is connected to the rest of the graph, it can be interpreted as showing a difficulty for memory-constrained virtual nodes (or even a motivation for designing GNNs with larger memory at the virtual node, which perhaps could be implemented by variable dimensionality of the embeddings maintained at different nodes in the GNN).

Choice of real-world benchmarks (Question 3)? We tried to choose diverse tasks that enable fair comparison of edge-based vs. node-based models (thus, we ruled out edge classification and node classification tasks, because they introduce a confounder – they require substantively different collation layers between the two architectures). ZINC/MNIST/CIFAR-10 are graph classification/regression datasets from the original GNN benchmark of Dwivedi et al. [1], and Peptides-Func/Peptides-Struct are the graph classification/regression datasets from the long-range graph benchmark [2].

Most of these graphs in fact do not have as skewed degree distributions as the synthetic examples (e.g. the chemistry graphs are subject to physical/molecular constraints, and the image graphs are constructed so that all degrees are within a factor of two – see Appendix C.5 in [1]). The goal of Table 1 was primarily to perform a controlled comparison between edge and node architectures on commonly-used GNN benchmark datasets. Of course, expanding the range of datasets (or even constructing more challenging ones reflective of natural GNN tasks) and exploring the impact of e.g. degree statistics on performance is an interesting future direction.

Notation in Lem. 2 / Def. 6 / Prop. 3? The overline notation in Lem. 2 means set complement, and the Delta({0,1}^V) in Def. 6 means the set of distributions over {0,1}^V. We will clarify these (and Prop. 3).

Ablation on the number of layers? We used the same number of layers for both architectures in Table 1 so as to equalize everything except the design choice “node-based vs edge-based”. Anecdotally we did not find significant changes with deeper node architectures, but unfortunately doing a thorough sweep over depths for both architectures would have been computationally taxing.

[1] Dwivedi et al. “Benchmarking Graph Neural Networks”. JMLR, 2023.

[2] Dwivedi et al. “Long Range Graph Benchmark”. NeurIPS 2022.

We thank the reviewer for their thoughtful review and we appreciate the positive comments! Below, we address the reviewer’s concerns and questions:

Interpretation of our main conclusions: The reviewer is correct that it’s unsurprising that adding edge embeddings may provide additional (representational) power. However, the goal of our work is not just to answer this yes/no question, but rather to investigate when edge embeddings help, how much they help (in particular, what are instances in which edge embeddings are particularly helpful), and understand the mechanism by which they help. Our main theoretical conclusions are that:

(1) adding edge embeddings (substantially) improves representational power in terms of required depth when there are hub nodes, and conversely does not improve representational power when the degree is bounded;

(2) this phenomenon is a consequence of memory constraints (and in particular, their interplay with depth) and not present under the standard theoretical lens where the only constraint on the protocol is symmetry.

We note these results can be viewed as GNN parallels of well-studied depth separation results for feedforward architectures [1] and Transformers [2] in which the goal is, similarly, to understand when and how much depth helps. These results were very influential in understanding the benefits of depth for classical architectures—but such theory is much less developed for GNNs.

Our experimental results validate the “hub nodes” finding (Table 2) and demonstrate a noticeable but small gain on real-world benchmarks (Table 1). We hope that these results will enrich the conversation on when to use edge embeddings, motivate the search for real-world benchmarks with larger performance gaps (if they exist), and inspire the development of architectures that match the representational power of edge embeddings with better computational efficiency.

We hope that this addresses the reviewer’s concern about whether our conclusions are “illuminating”.

Interpretation of Table 1: We agree with the reviewer that the closest analogue of the theory would be an experiment of the form “edge-based architectures achieve the same accuracy with lower depth”. Table 1 is morally equivalent so long as one believes that higher depth improves accuracy: “edge-based architectures achieve higher accuracy with the same depth”. We chose the latter because sweeping over depths to match accuracies is computationally expensive (and, ultimately, accuracy subject to compute/size constraints is an important desideratum in its own right), but we agree that an experiment trying to determine some representational “thresholds” as a function of depth would be scientifically interesting.

Interpretation of Table 2: Table 2 is actually about the planted model (section 8.2), not the Ising model – we apologize if the location of the table caused this confusion. Since the planted model is edge-based, it is obvious that the edge-based architecture is representationally powerful enough to solve the task, but Table 2 shows that the learning procedure also works (i.e. there are no unexpected training difficulties). Table 2 also shows that the node-based architecture empirically cannot learn (even though e.g. there is a node for every edge, since the graph is a star), consistent with our theoretical finding that hub nodes cause issues for node-based architectures.

[1] Telgarsky, “Benefits of depth in neural networks”, COLT 2016.

[2] Sanford et al. “Transformers, parallel computation, and logarithmic depth” ICML 2024.

We thank the reviewer for their thoughtful review and we appreciate the positive comments!

Regarding the computational challenge posed by dense graphs: we agree that mitigating this challenge while maintaining the representational power of edge-based architectures is an interesting direction for future work. A variety of heuristics have been proposed in the GNN literature to try to address related issues (e.g. graph rewiring, as mentioned by reviewer m8Mu). Developing a fuller theoretical understanding for these practical approaches is an exciting direction, and we are happy to add discussion to this effect in the paper.

We thank the reviewer for their thoughtful review and we appreciate the positive comments, particularly that the reviewer finds our work to be “of significant interest to the graph machine learning community”!

Regarding the size of the performance gain for edge-based GNNs on the real-world experiments: we agree that the gain is marginal. However, we believe this is still a valuable experimental outcome. Note that edge-based GNNs are a pre-existing architecture with already numerous applications, and our goal in this paper was not to make an argument that they are inherently superior to node-based GNNs but rather to understand their benefits or drawbacks in controlled settings. A priori, edge-based architectures could have been much better or much worse than node-based architectures on existing benchmarks; Table 1 indicates that they are in fact slightly better.

This outcome is consistent with our theoretical results since most of the graphs in these benchmarks do not have as skewed degree distributions as our theoretical constructions and synthetic examples (e.g. the chemistry graphs from ZINC/Peptides-func/Peptides-struct are subject to physical/molecular constraints, and the computer vision graphs from MNIST/CIFAR-10 are constructed in such a way that all degrees are within a factor of two – see Appendix C.5 in [1]). We agree completely that understanding whether there are naturally-arising “harder” benchmarks is an interesting direction for future research.

[1] Dwivedi et al. “Benchmarking Graph Neural Networks”. JMLR, 2023.