FSW-GNN: A Bi-Lipschitz WL-Equivalent Graph Neural Network

Response to Summary

We thank the reviewer for the detailed and accurate summary of our work. We would like to clarify one detail: the Fourier-Sliced Wasserstein (FSW) embedding for multisets is not merely almost injective, but rather fully injective and, moreover, bi-Lipschitz. These properties play a crucial role in ensuring the theoretical robustness and practical effectiveness of the proposed FSW-GNN architecture.

Response to Weaknesses

While the extension of the DS metric to node-featured graphs is an important step, to me this seems somewhat disconnected from the rest of the paper, considering that the proof of the lower Lipschitzness of FSW-GNN solely relies on the vector norms on the features being l1, giving rise to piecewise linear functions.

We only use the $l\_1$ norm in an intermediate step in the proof. Our results hold for any $\ell_p$-norm on the output Euclidean space, and any $p$-Wasserstein distance for $p\in[1,\infty]$. Our proof presents a novel technique to show that functions that are injective and piecewise linear are bi-Lipschitz, which we believe has applications outside the scope of this paper.

In its current stage, to me the contribution seems rather incremental on top of [2], which proposed the FSW multiset embeddings. The fact that using these injective multiset embeddings in a MPNN also yields WL-equivalent MPNNs, which is emphasized as one of the main results of the work, does not come surprising to me, and the brevity of the proof of Theorem 3.2 underlines this.

We emphasize that the primary contribution of our work lies not in achieving WL-equivalence, which has already been achieved in previous works, but in establishing the stronger property of bi-Lipschitzness. Our model is the first to provide such a guarantee, marking a significant advancement in the theoretical understanding of graph neural networks.

While it may seem a-priori intuitive that employing a bi-Lipschitz multiset aggregation should yield a bi-Lipschitz MPNN, proving this is a substantial theoretical challenge. Our proof relies on heavy machinery from the recent theory of $\sigma$-subanalytic functions [FSW], as well as the piecewise-linearity technique presented in Lemma 3.4.

In the analysis of the bi-Lipschitzness (Theorem 3.3), the obtained Lipschitz constants are not made explicit. The proof also seems rather generic to me, as it relies only on the set of all graphs up to a certain number of nodes being finite, and the underlying distance functions to be piecewise linear w.r.t. the features. With a similar reasoning, one could also argue that all ReLU MLPs are Lipschitz continuous -- however, the constant could be arbitrarily bad. As such, I can only see limited practical utility in such a statement.

We agree that explicitly bounding the Lipschitz constants would add value to our analysis. However, it is important to note that until the recent work of [FSW], the existence of such constants had not been established even for multisets. Explicitly bounding these constants would require entirely different theoretical tools, and is beyond the scope of this work. Nevertheless, we believe that this does not detract from the theoretical significance of our contribution, as it establishes a foundational result that opens the door for future explorations into bounding these constants.

To provide additional insight into the practical implications of our bi-Lipschitzness guarantee, we include an empirical evaluation of the bi-Lipschitz distortion incurred by our method and its competitors. Please refer to Figure 1 in the revised manuscript.

With a similar reasoning, one could also argue that all ReLU MLPs are Lipschitz continuous -- however, the constant could be arbitrarily bad. As such, I can only see limited practical utility in such a statement.

We agree that the argument in our bi-Lipschitzness proof could be used to show that ReLU MLPs that are injective (as a function defined on vectors) are bi-Lipschitz. However, this argument cannot be used to show that MPNNs with ReLU MLPs are bi-Lipschitz if they are injective (as functions defined on graphs).

A direct implication of the results in [FWT] is that MPNNs with piecewise-linear aggregation functions are never injective and, consequently, do not satisfy the assumptions of Lemma 3.4, nor do they have the stronger property of bi-Lipschitzness. Furthermore, the results in [FWT] imply that MPNNs with smooth aggregation functions are also never bi-Lipschitz. This distinction is strongly supported by our distortion evaluation experiment (Figure 1).

I thank the authors for their rebuttal and their thoughtful responses to my review. I sincerely apologize for the inaccurate wording regarding the injectivity of FSW-GNN, and I have updated the review accordingly.

Re technical contributions:

We only use the $\ell_1$ norm in an intermediate step in the proof

I thank the authors for the explanation, and I understand that the Lipschitzness also holds for any other norm beyond $\ell_1$, relying on the basic fact that all norms on finite dimensional vector spaces are equivalent.

Our proof presents a novel technique to show that functions that are injective and piecewise linear are bi-Lipschitz, which we believe has applications outside the scope of this paper.

I agree that this is a nice result, and acknowledge the effort it takes to rigorously prove this fact -- however intuitively, this seems immediately obvious to me, and without a quantitative lower bound that is somewhat related to the model, I maintain my opinion that the bi-Lipschitz bound is of limited relevance.

While it may seem a-priori intuitive that employing a bi-Lipschitz multiset aggregation should yield a bi-Lipschitz MPNN, proving this is a substantial theoretical challenge. Our proof relies on heavy machinery from the recent theory of $\sigma$-subanalytic functions [FSW], as well as the piecewise-linearity technique presented in Lemma 3.4.

I am not sure I agree that this presents a "substantial theoretical challenge". From my understanding, the "heavy machinery" (i.e., $\sigma$-subanalytic functions, finite witness theorem) is only used indirectly through Theorem 4.1 of [1], and the proof of Theorem 3.2 (L1011-1017) is a rather straightforward application of said theorem. I kindly ask the authors to further elaborate on their specific technical contribution in this regard, should I have overlooked some particular difficulty.

Re oversquashing:

Our argument at this point is just that our method dramatically outperforms competing MPNNs on these graph transfer tasks (Figure 3), which are widely recognized benchmarks for the oversquashing problem. Regardless of whether this conception of the community is completely correct, this experiment shows the potential in using FSW aggregations to alleviate oversquasing. It is possible that additional tools will be required for more challenging, yet undeveloped, adversarial oversquasing examples.

I thank the authors for their response. Taking the definition of oversquashing from [2] (and as defined in the review), from my understanding the said graph transfer tasks are necessary, but not sufficient for a graph ML model to alleviate oversquashing (as there is only one message that has to be propagated through the graph). Nevertheless, I fully agree that conceptions in the community about this vary slightly, and I want to emphasize that this was a very minor comment.

Re experiments:

We added Figure 1, which shows that our method yields considerably lower bi-Lipschitz distortion in practice in comparison with other methods.

I thank the authors for adding this interesting experiment, and the results for FSW-GNN's empirical distortion rates indeed look promising.

Conclusion:

I thank the authors for their responses and particularly for adding the empirical distortion rate experiment in Figure 1, which adds to the validity of the proposed method. To acknowledge the effort put into this, I am raising my score. Yet, I maintain my perception that the theoretical contributions are limited and that the work is too incremental on top of [3] and too close to SortMPNN from [4] (as also brought up by reviewer GwhL) to justify a higher score. As said in the review, I encourage the authors to keep working on a more quantitative version of Theorem 3.3.

[1] Tal Amir, Steven J. Gortler, Ilai Avni, Ravina Ravina, and Nadav Dym. Neural injective functions for multisets, measures and graphs via a finite witness theorem. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

[2] Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint arXiv:2006.05205, 2020.

[3] Tal Amir and Nadav Dym. Fourier sliced-Wasserstein Embedding for Multisets and Measures, 2024. URL https://arxiv.org/abs/2405.16519.

[4] Yair Davidson and Nadav Dym. On the hölder stability of multiset and graph neural networks. arXiv preprint arXiv:2406.06984, 2024.

I thank the authors for their explanations regarding the technical difficulty of their proofs.

I do agree that upon first glance, Theorem 3.10 from the referenced paper [1] appears to follow from Lemma 3.4, though I did not rigorously check the entire proof and I cannot rule out that there are no subtleties I might have missed.

I further want to mention that in the theory of subanalytic sets, the inequality of Lemma 3.4 is a special case of closely related to the Łojasiewicz inequality (cf. Theorem 6.4 of [2]). In particular, the lemma states that the Łojasiewicz exponent is 1 in the special case of piecewise linear functions. Following the proof from [2], it seems to me that this would be fairly straightforward to deduce. While I did not find an exact statement of Lemma 3.4 in the textbooks that I looked at, the lemma still strikes me as too elementary and general for not being standard knowledge in fields like real analytic geometry or nonsmooth optimization. That said, I recognize the importance of including and rigorously proving this result for the paper's completeness and self-containedness. Yet, I feel that Lemma 3.4 alone does not qualify as a strong technical contribution; this is, of course, quite subjective.

I will maintain my score.

[2] Edward Bierstone and Pierre D. Milman. Semianalytic and subanalytic sets. Publications mathématiques de l’I.H.É.S., tome 67 (1988), p. 5-42.

I thank the authors for their comment.

We believe having read the entirety of your answer that this is what you meant as well, but these particular words could be misleading.

I apologize if my wording was unclear, and I have updated the text.

In addition, the proof includes an important novel technical part, which was not covered by previous work. Lemma 3.4, which shows that piecewise-linear semi-metrics with the same zero set, defined on compact sets, are equivalent. This result may seem trivial, but it is not. For example, when the domain is not compact, there exist counterexamples for which the Lemma statement does not hold.

We stress that we are not claiming that our work would be accepted to a journal on real analytic geometry. This paper focuses on making graph neural networks provably bi-Lipschitz and demonstrating the benefits of this property experimentally. This was not done before, and is unlikely to be achieved by researchers focused solely on real analytic geometry.

I appreciate this explanation. I also think that significant theoretical novelty is not necessary in works with a strong empirical part, as long as the theoretical aspects support and enhance the empirical findings. However, I believe that in this particular case, the proof of Bi-Lipschitzness is a core part of the contribution. My intention was not to diminish its value but to express some skepticism about the claimed novelty of Lemma 3.4.

Response to Weaknesses

On node-level vs. graph-level tasks: We thank the reviewer for this question, and would like to clarify the connection between the graph-level and node-level output. The bi-Lipschitzness of our graph-level embedings necessarily implies bi-Lipschitzness for the individual output vertex features. Lack of bi-Lipschitzness in the latter would preclude bi-Lipachitzness for the former, both theoretically and practically, as the latter is derived from the former.

The purpose of the long-range experiments was to empirically demonstrate that the output node features computed by our GNN exhibit better information retention compared to those produced by other architectures. This aligns well with our theory and strengthens the practical relevance of our method.

Response to Weaknesses

While the paper provides proof for the bi-Lipschitz properties, it relies on empirical evidence and some conjecture to suggest that the model’s stability holds as depth increases.

Thank you for this insightful question. First, we’d like to clarify that our bi-Lipschitz proof hold for abitrary depth. It is true that distortion increases with depth in our method. However, the key advantage of our approach is that the distortion grows much more slowly compared to other methods. This is due to the higher distortion incurred per iteration in methods that lack bi-Lipschitzness. In fact, methods that are not bi-Lipschitz are theoretically proven to have infinite distortion on some pairs of graphs, even after just one iteration.

This distinction is clearly demonstrated in our empirical distortion evaluation (Figure 1). Even after 10 iterations, the distortion incurred by our method remains significantly lower than that of just one iteration of most competing methods. This highlights the stability of our approach and its robustness to increasing depth.

Response to Questions

TMD requires proper features for every vertex, which can not be available in practical cases, where it fails. Can this method be extended in those cases where partial node features are missing?

Yes, this method can be extended to cases with partial node features. One approach is to augment the vertex feature vectors with a 0-1 indicator that specifies whether the corresponding vertex has a feature. In cases where a feature is missing, a placeholder vector (e.g., a zero vector) can be used as the dummy feature.

What are the average node numbers and edge numbers of the dataset that are used in the experiments?

We added these statistics in Appendix A, thanks for the suggestion.

How does the FSW-GNN handle graphs with different structures, such as highly dense versus sparse graphs, in terms of both performance and computational efficiency?

Computation wise: FSW-GNN is designed to handle both dense and sparse graphs efficiently. In particular, our implementation leverages PyTorch’s support for sparse tensors, allowing us to process sparse graphs with significant computational savings. In terms of runtime, FSW-GNN is computationally more intensive than simple GNNs that utilize basic message aggregation functions such as sum pooling. Yet, our method takes only four times longer to run compared to these simpler methods. This additional computational cost is expected due to the more sophisticated message aggregation used in FSW-GNN.

Performance wise: It can be that the two graphs in Table 4 where SortMPNN outperforms our method are exactly the two graphs with relatively high average degree. We added a comment discussing this to the main text.

How critical is the bi-Lipschitz property for real-world applications beyond synthetic tasks?

This is an important question. Indeed the purpose of all our experiment is to address it. Our experiments suggest that our bi-Lipschitz architecture is very helpful for the transductive task, and for long range tasks.

The paper focuses on long-range tasks that benefit from deep message passing. How does FSW-GNN perform on tasks involving smaller graphs or shallow networks, like the MUTAG dataset?

We added an experiment on MUTAG. FSW-GNN outperforms all other MPNNs we compared against. Note also that several graphs in the transductive learning tasks are also rather small.

How sensitive is FSW-GNN to changes in hyperparameters compared to traditional MPNNs?

In our experiments, FSW-GNN demonstrated low sensitivity to changes in hyperparameters. For instance, doubling or halving the hidden feature dimension had very little effect on the results for the Actor and Cora datasets.

Response to Strengths

We appreciate the reviewer’s thoughtful feedback and are glad that you found our manuscript enjoyable and well-presented.

Response to Weaknesses

Firstly, as you suggested, we have added a comparison to SortMPNN on all datasets considered in the paper, and also on the MUTAG dataset and distortion experiment suggested by other reviewers. As we wrote in our comment to all reviewers, FSW-GNN outperforms SortMPNN in most tasks: 7 of 9 tasks in Table 4, peptide-struct, all of our synthetic oversquashing experiments, and the new MUTAG and distortion experiments. In contrast, SortMPNN outperforms FSW-GNN only on peptide-func and 2 out of 9 datasets in Table 4.

Regarding the similarity to Sort-MPNN

To address the similarities and distinctions between FSW-GNN and SortMPNN, we wish to highlight two critical differences:

1. Handling variable neighborhood sizes: To accommodate vertex neighborhoods of varying size, SortMPNN uses padding, which requires the specification of a parameter $d\_{\max}$ that upper-bounds the maximal neighborhood size in the input graphs. This necessitates a priori knowledge of the maximal vertex degree across all input graphs to be processed by the model. Such a requirement can pose significant limitations in real-world use cases where this information is unavailable or unpredictable. Moreover, SortMPNN’s runtime and memory complexity scale linearly with $d\_{\max}$, making the use of arbitrarily large $d\_{\max}$ impractical. Additionally, padding introduces artificial distortion in the embedding space, which can impair performance. This is evidenced in our new experiments, as seen in Figure 3. In contrast, FSW-GNN treats vertex neighborhoods as uniform distributions - a natural and principled approach that eliminates the need for padding or prior knowledge of maximum neighborhood size. This makes FSW-GNN computationally more efficient and conceptually more generalizable.

2. Applicability to weighted graphs: SortMPNN is limited to handling simple graphs or, at best, multigraphs, as its padding-based approach is not suited for weighted graphs. FSW-GNN, on the other hand, is inherently applicable to weighted graphs due to its treatment of vertex neighborhoods as distributions. This further extends the practical applicability of our method, and is directly supported in our implementation.

Theoretical guarantees: Our work is the first to provide a theoretical bi-Lipschitzness guarantee for a GNN that hold uniformly and deterministically for all inputs. In contrast, the guarantees provided in the SortMPNN paper require a probabilistic assumption on the input, and hold only in expectation. Furthermore, our novel proof technique, which combines the Finite Witness Theorem and piecewise-linearity argument of Lemma 3.4, can be extended to more general settings. We thus believe this technique is of independent interest and has the potential to support future research beyond the scope of this work.

The authors argue that FSW-GNN excels in long-range tasks, however, the reported results in peptides-struct and peptides-func, which are considered long-range datasets, are not that good. How do the authors explain this?

We thank the reviewer for this question. Indeed, the goal of designing LRGB (Long-Range Graph Benchmark) was to design a dataset of long-range tasks on which graph transformers could outperform MPNN. Note that, in the language of the original paper [LRGB], the dataset “is not directed at proposing ‘provable LRI (Long Range Interaction)' benchmarks, which would often lead to toy datasets” but rather on proposing real-world tasks which are believed to be long range. The flip side of this is that we are not completely sure that these tasks are indeed long range. Indeed, regarding the two peptide datasets considered in this paper, the reasoning given for the long-rangeness is rather circumstantial: “Thus, the proposed Peptides-func and Peptides-struct are better suited to benchmarking of graph Transformers or other expressive GNNs, as they contain larger graphs, more data points, and challenging tasks.”

Additionally, a later paper [GAP] showed that, with appropriate hyperparameter tuning, MPNNs with depth of 6-10 can achieve very competitive results, in comparison with graph transformers. These are the results we compare with in our paper. The success of standard MPNN of reasonable depth casts some doubt on the long-rangeness of Peptides-Struct and Peptides-Func.

In contrast, the long range datasets we consider in this paper, Figures 2-3, are popular synthetic datasets, but are provably long range: an MPNN with less depth than the problem radius will never be successful.