Towards Principled Graph Transformers

We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer states

The total number of parameters for the model is not provided in the parameter table

We provide the parameters for each model in the general response above. As can be observed, on the molecular regression tasks, the ET is well within the 500K parameter budget.

Next, the reviewers says

While ET is more efficient than some other high-order transformers, it is still significantly inferior to most graph transformers with O(n^2). This limits the applicability of ET

See our general response for a discussion of the practical usefulness of the ET. In summary, we argue that there exists still plenty of application and interest for a model with high runtime complexity, particularly, if this complexity can be translated into strong predictive performance.

Finally, the reviewer says

The author's contribution is confined to applying the ET model to the graph, without proposing their own method

See our general response for a discussion on the novelty and contribution of our work. In summary, we argue that the novelty of our work does not come from proposing a new method. Rather, the main contribution of our work is in providing a connection between the ET, systematic generalization and graph learning.

Questions

Further, we are happy to address the question of the reviewer below. First, the reviewer asks

Can you provide possible ways to reduce the complexity of ET?

As our work is partly motivated by demonstrating that high expressivity can also result in strong empirical performance, we do not aim to reduce the theoretical complexity of the ET in this work. At the same time, we show that in practice, the high runtime and memory cost can be reduced by leveraging a parallelizable implementation, together with low-level GPU optimizations; see the limitation section in our paper. More generally, while the ET may not be applicable to large graphs yet, we regard our work as a first step at a theoretical understanding of equivariant graph transformers with high expressive power that also perform strongly in practice.

Next, the reviewer asks

Why has the performance changed so much compared to the submitted version of ICML in 2024? Especially on the Searching dataset.

In a previous version, we had not included the graph features that are available in CLRS, and used by baseline models, into the ET. These features proved particularly useful on the Searching datasets, explaining the improved performance.

Finally, the reviewer asks

What do Average and All algorithms represent respectively?

In “Average”, we average performance across the six algorithm classes, whereas in “All algorithms” we average performance over all algorithms.

We want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.

We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer states

My major concern is the real usefulness of ET, since the high runtime and memory complexity may offset the importance of ET

See our general response for discussion of the practical usefulness of the ET. In summary, we argue that there exists still plenty of application and interest for a model with high runtime complexity, particularly, if this complexity can be translated into strong predictive performance.

Next, the reviewer says

Also, it is not clear what is the parameter count, which raises the question that the better performance of ET may be mainly from large number of parameters

We provide the parameter counts in the general response. We want to highlight that on ZINC and Alchemy, we adhere to the 500K parameter budget and that the ET’s strong performance cannot be explained by a comparatively large number of parameters.

Finally, the reviewer says

It seems that the ET in this paper is mainly from existing work, and the difference needs to be explained

See our general response for a discussion on the novelty and contribution of our work. In summary, we argue that the novelty of our work does not come from proposing a new method. Rather, the main contribution of our work is in providing a connection between the ET, systematic generalization and graph learning. To still answer the question about the difference between the ET in our work and the original ET, note that on the architectural side, we provide a sufficient tokenization and readout for the ET to simulate 2-FWL. In what follows, we provide a detailed explanation in response to the reviewer’s first question, which also touches on this concern.

Questions

The reviewer asks

But it is not quite clear to me what are the main differences of the tokenization compared with existing ones, and why this tokenization can help the 3-WL

We are happy to provide further details on the proposed tokenization for the ET. For the tokenization to be sufficient for the 2-FWL, the initial embeddings of each node pair need to injectively encode the initial colors of each node, as well as an indicator of whether the nodes are connected by an edge. Such a tokenization is already vaguely described in the original paper (https://arxiv.org/abs/2112.00578). In this work, we merely formalize this tokenization such that it can be used in our formal proofs. Other than the tokenization and readout, we do not modify the ET at all.

Next, the reviewer asks

In Table 1, ET with RRWP has a little bit lower performance on QW9. Can the authors elaborate on this? Since generally, transformers with positional encodings will lead to improved performance

We suspect that adding positional encodings (PEs) only translates to improved empirical performance if the PEs provide additional features either useful for expressivity or directly for prediction. However, since the ET already has a high base expressivity, on this particular task, the additional RRWP encodings may not be useful for either distinguishing graphs or directly for predicting the targets and instead lead to slight overfitting. A takeaway from this finding may be that the base expressivity of the ET is often enough, as the RRWP encodings do not prove substantial for the good performance of the ET. On the other hand, models such as the SOTA graph transformer GRIT rely on PEs for good performance; see https://arxiv.org/abs/2305.17589, Table 5.

Finally, the reviewer asks

But maybe as an open question, if incorporated with PEs, will the theoretical results change?

We agree with the reviewer that this is an interesting question for future work. In particular, if one were able to show that RRWP encodings can distinguish graphs indistinguishable by the 3-WL, one would obtain a theoretical guarantee that ET+RRWP leads to increased expressivity. At the same time, we regard it as a strength of the ET that the model does not rely on PEs but can still successfully leverage them, if provided and informative on the given task; see e.g., our results on ZINC 12K.

We want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.

We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer says

Previous work [1] has demonstrated that graph transformers (GTs) with proper positional encodings (PEs) can be more powerful than any WL test

The work in [1] presents a universality result for transformers with Laplacian encodings. However, the transformer in [1] is not equivariant to node permutations, since the Laplacian encodings are not invariant to sign- and basis changes in the underlying eigenvalue decomposition. While a universal model is more expressive than any WL test, it is questionable whether the transformer in [1] is able to learn permutation equivariance from limited data. Theoretically, we expect worse generalization from non-equivariant models on graph tasks; see e.g., Petrache and Trivedi [2]. In fact, the lack of permutation equivariance could well explain the worse empirical performance of the model in [1] compared to the ET, e.g., on ZINC 12K. In summary, our theoretical results should be viewed under the assumption of permutation equivariance, where 3-WL is considered to have very strong expressive power.

Next, the reviewer says

While the authors mention how to apply ET for node-level tasks, they provide no corresponding results

We refer the reviewer to the experiments on the CLRS benchmark. There, many tasks involve node- and edge-level predictions. The reviewer also mentions that

Given the cubic complexity of ET, it is likely not scalable to most graphs for node-level tasks.

We agree that the ET currently cannot be scaled to thousands of nodes and that most large-scale transductive node tasks are out of scope for the ET. Nonetheless, our results on CLRS indicate that the ET can be effectively used for node-level tasks.

Finally, the reviewer mentions that

Some relevant GT models haven't been discussed or compared (e.g., [1-3])

We are happy to include those models into our discussion. In particular, we will add a paragraph discerning expressivity results for permutation equivariant models as compared to non-equivariant but universal models such as the one in [1].

We want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.

References

[1] Kreuzer et al., Rethinking Graph Transformers with Spectral Attention, 2021 https://arxiv.org/abs/2106.03893

[2] Petrache and Trivedi, Approximation-Generalization Trade-offs under (Approximate) Group Equivariance, 2023 https://arxiv.org/abs/2305.17592

Thanks for authors' detailed response. By node-level tasks, I’m actually referring to some larger graphs for node classification (e.g., citation/biological networks) rather than neural algorithmic reasoning tasks. I’m interested in seeing how well the proposed approach scales with larger graphs.

We want to ask the reviewer to increase their score if they are satisfied with our response.

I'm really not a big fan of explicitly asking reviewers to change their scores in this way. Authors could say "reconsidering their score" instead.

That said, I do appreciate the major contributions of this work and increase my score to 7.

We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer.

Regarding evaluation, the reviewer says that

I find it is needed to include more benchmark datasets in graph learning […] to comprehensively demonstrate its practical predictive power compared with state-of-the-art (SOTA) methods, e.g., those graph transformers Grit, Exphormer, GraphGPS

as well as

Besides datasets from algorithmic reasoning, currently, in my opinion, ET is only evaluated on two widely-used graph benchmark datasets

and

It would be interesting to see its long-range modeling capability on the LRGB benchmark as it's a transformer

While we certainly evaluate the ET only on a subset of all possible graph learning benchmarks we argue that our results already provide a comprehensive view on the predictive performance of the ET. In particular, we evaluate the ET on three different types of benchmarks:

• Molecular property prediction (ZINC, Alchemy, QM9)
• Algorithmic reasoning (CLRS)
• Graph isomorphism testing (BREC)

Most notably, the CLRS benchmark contains a total of 30 different datasets and includes a variety of graph-, node- and edge-level tasks. In addition, the CLRS benchmark inherently evaluates models on size generalization, a problem recently highlighted as an important challenge in graph learning; see https://arxiv.org/abs/2402.02287, Section 3.1, Challenge III.4. Hence, we argue that CLRS is a high-quality benchmark and note that it is also widely adopted; see e.g., https://arxiv.org/abs/2209.11142, https://arxiv.org/abs/2403.04929, https://arxiv.org/abs/2402.01107, https://arxiv.org/abs/2404.03441, https://arxiv.org/abs/2312.05611, https://arxiv.org/abs/2210.05062.

Finally, note that our concrete selection of tasks is partly also to validate our theoretical findings and not only to demonstrate SOTA performance on the most popular datasets; see also our remark in the general response on the intended scope of our work, in particular, regarding an insightful theoretical and empirical investigation of the ET.

The above being said, we worked hard during the rebuttal to provide encouraging preliminary results on two additional datasets, see below.

PCQM4Mv2

We evaluate on one random seed, as is common on this benchmark.

ZINC-full

We evaluate on 2 random seeds. The ET has 283,969 parameters.

Notably, on PCQM4Mv2 our 6M model outperforms the GraphGPS model at the same parameter budget while also being on par with transformer models such as TokenGT which has almost 5x as many parameters. On ZINC-full, we observe that the results of the ET are competitive with the best models, albeit not SOTA results.

We note that achieving SOTA results on new datasets typically involves extensive hyper-parameter tuning. Due to the limited time during the rebuttal we were only able to test a few configurations and also at only a fraction of the number of epochs (1000 instead of 2000 on ZINC-full and 40 instead of 200 on PCQM4Mv2) of the baseline models. Nonetheless, our results indicate that with more time for properly tuning hyper-parameters and running the same number of epochs, the ET is a viable contender for SOTA results on both datasets.

Finally, the reviewer says

In terms of the evaluation of expressive power, it would be better to include SOTA methods such as those subgraph-GNNs, potentially also discussing their complexity compared to ET

We refer the reviewer to Wang and Zhang 2023 (https://arxiv.org/abs/2304.07702), Table 2. Here, the authors have already evaluated a variety of subgraph GNNs on BREC. Most notably, the ET beats 7 out of 9 models on this benchmark.

In what follows, we address the question of the reviewer. First, the reviewer asks

How does the setting of QM9 used in this paper differ from the setting that is commonly used in equivariant GNNs

In this work, for QM9, we adopt the setting used for SpeqNets (https://arxiv.org/abs/2203.13913), where tasks are learnt jointly and the resulting MAE is the average across all tasks, where we note that all targets are normalized by subtracting the mean and dividing by the standard deviation. We follow this setting to enable a comparison to the higher-order WL SpeqNets. We will add this description to the paper and plan to include results on QM9 in the more widely adopted single-task setting in the future.

Next, the reviewer asks

Are there any particular benefits to having an expressive model without using PE or SE? […] because it would be more challenging to analyze expressive power with PE/SE, therefore it would be easier to understand a model's expressive power without using them?

and mentions that

I feel like depending on how PE is used, it is still possible to analyze the expressive power, e.g., via graph substructure counting

While it is true that the expressive power of PEs can be understood, in general, we argue that knowing which PEs to add is typically highly task-dependent; see e.g., the study in GraphGPS (https://arxiv.org/abs/2205.12454), as PEs are typically based on graph heuristics. This also holds for subgraph GNNs, where assumptions have to be made about the type of substructures that are relevant for downstream tasks. In contrast, the ET without any additional PEs already performs strongly, across tasks in molecular regression, algorithmic reasoning and graph isomorphism testing.

We want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.