Towards Principled Graph Transformers

Authors: luis  müller
daniel  kusuma
blai  bonet
christopher  morris
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:07:50

Abstract: The expressive power of graph learning architectures based on the $k$-dimensional Weisfeiler-Leman ($k$-WL) hierarchy is well understood. However, such architectures often fail to deliver solid predictive performance on real-world tasks, limiting their practical impact. In contrast, global attention-based models such as graph transformers demonstrate strong performance in practice, but comparing their expressive power with the $k$-WL hierarchy remains challenging, particularly since these architectures rely on positional or structural encodings for their expressivity and predictive performance. To address this, we show that the recently proposed Edge Transformer, a global attention model operating on node pairs instead of nodes, has 3-WL expressive power when provided with the right tokenization. Empirically, we demonstrate that the Edge Transformer surpasses other theoretically aligned architectures regarding predictive performance while not relying on positional or structural encodings.

Main paper: 
 limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We provide a discussion of the limitations of the ET, specifically its high runtime and memory complexity, in Section 7.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
edge transformers: 
 The ET was originally designed to improve the systematic generalization abilities of machine learning models. To borrow the example from Bergen et al. [6] , a model that is presented with relations such as MOTHER(x, y), indicating that y is the mother of x, could generalize to a more complex relation GRANDMOTHER(x, z), indicating that z is the grandmother of x if MOTHER(x, y) ∧ MOTHER(y, z) holds. The particular form of attention used by the ET, which we will formally introduce hereafter, is designed to explicitly model such more complex relations. Indeed, leveraging our theoretical results of Section 4, in Section 5, we formally justify the ET for performing systematic generalization. We will now formally define the ET; see Appendix D for a complete description of our notation.In general, the ET operates on a graph G with nodes V (G) and consecutively updates a 3D tensor state X ∈ R n×n×d , where d is the embedding dimension and X ij or X(u) denotes the representation of the node pair u := (i, j) ∈ V (G) 2 ; see Figure 1 for a visualization of this construction. Concretely, the t-th ET layer computesX (t) ij := FFN X (t-1) ij + TriAttention LN X (t-1) ij, for each node pair (i, j), where FFN is a feed-forward neural network, LN denotes layer normalization [2] and TriAttention is defined asEQUATIONwhich computes a tensor product between a three-dimensional attention tensor α and a threedimensional value tensor V, by multiplying and summing over the second dimension. Here,EQUATIONis the attention score between the features of tuples (i, l) and (l, j), andEQUATIONwe call value fusion of the tuples (i, l) and (l, j) with ⊙ denoting element-wise multiplication. Moreover, W Q , W K , W V1 , W V2 ∈ R d×d are learnable projection matrices; see Figure 2 for an overview of the tensor operations in triangular attention and see Algorithm 1 for a comparison to standard attention [46] in pseudo-code. Note that similar to standard attention, triangular attention can be straightforwardly extended to multiple heads.As we will show in Section 4, the ET owes its expressive power to the special form of triangular attention. In our implementation of the ET, we use the following tokenization, which is sufficient to obtain our theoretical results.Tokenization We consider graphs G := (V (G), E(G), ℓ) with n nodes and without self-loops, where V (G) is the set of nodes, E(G) is the set of edges, and ℓ : V (G) → N assigns an initial color to each node. We construct a feature matrix F ∈ R n×p that is consistent with ℓ, i.e., for nodes i and j in V (G), F i = F j if and only if, ℓ(i) = ℓ(j). Note that, for a finite subset of N, we can always construct F , e.g., using a one-hot encoding of the initial colors. Additionally, we consider an edge feature tensor E ∈ R n×n×q , where E ij denotes the edge feature of the edge (i, j) ∈ E(G). If no edge features are available, we randomly initialize learnable vectors x 1 , x 2 ∈ R q and assign x 1 to Algorithm 1 Comparison between standard attention and triangular attention in PYTORCH-like pseudo-code. function ATTENTION(X : n × d) Q, K, V ← linear(X).chunk(3) # no op Ã ← einsum(id, jd → ij, Q, K)A ← softmax( Ã/ √ d, -1) O ← einsum(ij, jd → id, A, V) return linear(O) end function function TRI_ATTENTION(X : n × n × d) Q, K, V 1 , V 2 ← linear(X).chunk(4) V ← einsum(ild, ljd → iljd, V 1 , V 2 ) Ã ← einsum(ild, ljd → ilj, Q, K) A ← softmax( Ã/ √ d, -1) O ← einsum(ilj, iljd → ijd, A, V) return linear(O) end function E ij if (i, j) ∈ E(G).Further, for all i ∈ V (G), we assign x 2 to E ii . Lastly, if (i, j) ̸ ∈ E(G) and i ̸ = j, we set E ij = 0. We then construct a 3D tensor of input tokens X ∈ R n×n×d , such that for node pair (i, j) ∈ V (G) 2 ,EQUATIONwhere ϕ : R 2p+q → R d is a neural network. Extending Bergen et al. [6] , our tokenization additionally considers node features, making it more appropriate for the graph learning setting.Efficiency The triangular attention above imposes a O(n 3 ) runtime and memory complexity, which is significantly more efficient than other transformers with 3-WL expressive power, such as the higherorder transformers in Kim et al. [27] and Kim et al. [28] with a runtime of O(n 6 ). Nonetheless, the ET is still significantly less efficient than most graph transformers, with a runtime of O(n 2 ) [32, 42, 53] . Thus, the ET is currently only applicable to mid-sized graphs; see Section 7 for an extended discussion of this limitation.Positional/structural encodings Additionally, GNNs and graph transformers often benefit empirically from added positional/structural encodings [13, 32, 42] . We can easily add PEs to the above tokens with the ET. Specifically, we can encode any PEs for node i ∈ V (G) as an edge feature in E ii and any PEs between a node pair (i, j) ∈ V (G) 2 as an edge feature in E ij . Note that typically, PEs between pairs of nodes are incorporated during the attention computation of graph transformers [32, 53] . However, in Section 6, we demonstrate that simply adding these PEs to our tokens is also viable for improving the empirical results of the ET.Readout Since the Edge Transformer already builds representations on node pairs, making predictions for node pair-or edge-level tasks is straightforward. Specifically, let L denote the number of Edge Transformer layers. Then, for a node pair (i, j) ∈ V (G) 2 , we simply readout X (L) ij , where on the edge-level we restrict ourselves to the case where (i, j) ∈ E(G). In the case of nodes, we can for example read out the diagonal of X (L) , that is, the representation for nodei ∈ V (G) is X (L)ii . In addition to the diagonal readout, we also design a more sophisticated read out strategy for nodes which we describe in Appendix A.1.With tokenization and readout as defined above, the ET can now be used on many graph learning problems, encoding both node and edge features and making predictions for node pair-, edge-, node-, and graph-level tasks. We refer to a concrete set of parameters of the ET, including tokenization and positional/structural encodings, as a parameterization. We now move on to our theoretical result, showing that the ET has the same expressive power as the 3-WL.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: We provide all proofs in Appendix E, where we state all assumptions in the respective theorem statements. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
the expressivity of edge transformers: 
 Here, we relate the ET to the folklore Weisfeiler-Leman (k-FWL) hierarchy, a variant of the k-WL hierarchy for which, for k > 2, (k -1)-FWL is as expressive as k-WL [19] . Specifically, we show that the ET can simulate the 2-FWL, resulting in 3-WL expressive power. To this end, we briefly introduce the 2-FWL and then show our result. For detailed background on the k-WL and k-FWL hierarchy, see Appendix D. Folklore Weisfeiler-Leman Let G := (V (G), E(G), ℓ) be a node-labeled graph. The 2-FWL colors the tuples from V (G) 2 , similar to the way the 1-WL colors nodes [36] . In each iteration, t ≥ 0, the algorithm computes a coloring C 2,F t : V (G) 2 → N and we write C 2,F t (i, j) or C 2,F t (u) to denote the color of tuple u := (i, j) ∈ V (G) 2 at iteration t. For t = 0, we assign colors to distinguish pairs (i, j) in V (G) 2 based on the initial colors ℓ(i), ℓ(j) of their nodes and whether (i, j) ∈ E(G) or i = j. For a formal definition of the initial node pair colors, see Appendix D. Then, for each iteration, t > 0, the coloring C 2,F t is defined asC 2,F t (i, j) := recolor (C 2,F t-1 (i, j), M t-1 (i, j)) ,where recolor injectively maps the above pair to a unique natural number that has not been used in previous iterations andM t-1 (i, j) := { {(C 2,F t-1 (i, l), C 2,F t-1 (l, j)) | l ∈ V (G)} }.We show that the ET can simulate the 2-FWL, resulting in at least 3-WL expressive power. Further, we show that the ET is also, at most, as expressive as the 3-WL. As a result, we obtain the following theorem; see Appendix E for a formal statement and proof details. Theorem 1 (Informal). The ET has exactly 3-WL expressive power.Note that following previous works [33, 37, 39] , our expressivity result is non-uniform in that our result only holds for an arbitrary but fixed graph size n; see Proposition 7 and Proposition 8 for the complete formal statements and proof of Theorem 1.In the following, we provide some intuition of how the ET can simulate the 2-FWL. Given a tuple (i, j) ∈ V (G) 2 , we encode its color at iteration t with X (t) ij . Further, to represent the multiset{ {(C 2,F t-1 (i, l), C 2,F t-1 (l, j)) | l ∈ V (G)} },we show that it is possible to encode the pair of colors(C 2,F t-1 (i, l), C 2,F t-1 (l, j)) via X (t-1) il W V1 ⊙ X (t-1) lj W V2 ,for node l ∈ V (G). Finally, triangular attention in Equation ( 1), performs weighted sum aggregation over the 2-tuple of colors (C 2,F t-1 (i, l), C 2,F t-1 (l, j)) for each l, which we show is sufficient to represent the multiset; see Appendix E. For the other direction, namely that the ET has at most 3-WL expressive power, we simply show that the recolor function can simulate the value fusion in Equation (3), as well as the triangular attention in Equation (1) . Interestingly, our proofs do not resort to positional/structural encodings. The ET draws its 3-WL expressive power from its aggregation scheme, the triangular attention. In Section 6, we demonstrate that this also holds in practice, where the ET performs strongly without additional encodings. In what follows, we use the above results to derive a more principled understanding of the ET in terms of systematic generalization, for which it was originally designed. Thereby, we demonstrate that graph theoretic results can also be leveraged in other areas of machine learning, further highlighting the benefits of theoretically grounded models.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We provide the model definition in Section 3, as well as PyTorch-like pseudocode for the triplet attention in Algorithm 1. In addition, we provide all experimental details in Section 6 as well as the chosen hyper-parameters and optimizers in Table 6 .Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
the logic of edge transformers: 
 After borrowing the ET from systematic generalization in the previous section, we now return the favor. Specifically, we use a well-known connection between graph isomorphism and first-order logic to obtain a theoretical justification for systematic generalization reasoning using the ET. Recalling the example around the GRANDMOTHER relation composed from the more primitive MOTHER relation in Section 3, Bergen et al. [6] go ahead and argue that since self-attention of standard transformers is defined between pairs of nodes, learning explicit representations of GRANDMOTHER is impossible and that learning such representations implicitly incurs a high burden on the learner. Conversely, the authors argue that since the ET computes triangular attention over triplets of nodes and computes explicit representations between node pairs, the Edge Transformer can systematically generalize to relations such as GRANDMOTHER. While Bergen et al. [6] argue the above intuitively, we will now utilize the connection between first-order logic (FO-logic) and graph isomorphism established in Cai et al. [10] to develop a theoretical understanding of systematic generalization; see Appendix D for an introduction to first-order logic over graphs. We will now briefly introduce the most important concepts in Cai et al. [10] and then relate them to systematic generalization of the ET and similar models.Language and configurations Here, we consider FO-logic statements with counting quantifiers and denote with C k,m the language of all such statements with at most k variables and quantifier depth m. A configuration is a map between first-order variables and nodes in a graph. Concretely, configurations let us define a statement φ in first-order logic, such as three nodes forming a triangle, without speaking about concrete nodes in a graph G = (V (G), E(G)). Instead, we can use a configuration to map the three variables in φ to nodes v, w, u ∈ V (G) and evaluate φ to determine whether v, w and u form a triangle in G. Of particular importance to us are k-configurations f where we map k variables x 1 , . . . , x k in a FO-logic statement to a k-tuple u ∈ V (G) k such that u = (f (x 1 ), . . . , f (x k )). This lets us now state the following result in Cai et al. [10] , relating FO-logic satisfiability to the k-FWL hierarchy. Here, C k,F t denotes the coloring of the k-FWL after t iterations; see Appendix D for a precise definition. Theorem 2 (Theorem 5.2 [10] , informally). Let G := (V (G), E(G)) and H := (V (H), E(H)) be two graphs with n nodes and let k ≥ 1. Let f be a k-configuration mapping to tuple u ∈ V (G) k and let g be a k-configuration mapping to tuple v ∈ V (H) k . Then, for every t ≥ 0,C k,F t (u) = C k,F t (v), if and only if u and v satisfy the same sentences in C k+1,t whose free variables are in {x 1 , x 2 , . . . , x k }.Together with Theorem 1, we obtain the above results also for the embeddings of the ET for k = 2. Corollary 3. Let G := (V (G), E(G)) and H := (V (H), E(H)) be two graphs with n nodes and let k = 2. Let f be a 2-configuration mapping to node pair u ∈ V (G) 2 and let g be a 2-configuration mapping to node pair v ∈ V (H) k . Then, for every t ≥ 0,X (t) (u) = X (t) (v),if and only if u and v satisfy the same sentences in C 3,t whose free variables are in {x 1 , x 2 }.Systematic generalization Returning to the example in Bergen et al. [6] , the above result tells us that a model with 2-FWL expressive power and at least t layers is equivalently able to evaluate sentences in C 3,t , including GRANDMOTHER(x, z) = ∃y MOTHER(x, y) ∧ MOTHER(y, z) , i.e., the grandmother relation, and store this information encoded in some 2-tuple representation X (t) (u), where u = (u, v) and X (t) (u) encodes whether u is a grandmother of v. As a result, we have theoretical justification for the intuitive argument made by Bergen et al. [6] , namely that the ET can learn an explicit representation of a novel concept, in our example the GRANDMOTHER relation.However, when closely examining the language C 3,t , we find that the above result allows for an even wider theoretical justification of the systematic generalization ability of the ET. Concretely, we will show that once the ET obtains a representation for a novel concept such as the GRANDMOTHER relation, at some layer t, the ET can re-combine said concept to generalize to even more complex concepts. For example, consider the following relation, which we naively write as GREATGRANDMOTHER(x, a) = ∃z∃y MOTHER(x, y) ∧ MOTHER(y, z) ∧ MOTHER(z, a) .At first glance, it seems as though GREATGRANDMOTHER ∈ C 4,2 but GREATGRANDMOTHER ̸ ∈ C 3,t for any t ≥ 1. However, notice that the variable y serves merely as an intermediary to establish the GRANDMOTHER relation. Hence, we can, without loss of generality, write the above as GREATGRANDMOTHER(x, a) = ∃y ∃a MOTHER(x, a) ∧ MOTHER(a, y) a is re-quantified and temporarily bound ∧MOTHER(y, a) , i.e., we re-quantify a to temporarily serve as the mother of x and the daughter of y. Afterwards, a is released and again refers to the great grandmother of x. As a result, GREATGRANDMOTHER ∈ C 3,2 and hence the ET, as well as any other model with at least 2-FWL expressive power, is able to generalize to the GREATGRANDMOTHER relation within two layers, by iteratively re-combining existing concepts, in our example the GRANDMOTHER and the MOTHER relation. This becomes even more clear, by writing GREATGRANDMOTHER(x, a) = ∃y GRANDMOTHER(x, y) ∧ MOTHER(y, a) , where GRANDMOTHER is a generalized concept obtained from the primitive concept MOTHER. To summarize, knowing the expressive power of a model such as the ET in terms of the Weisfeiler-Leman hierarchy allows us to draw direct connections to the logical reasoning abilities of the model. Further, this theoretical connection allows an interpretation of systematic generalization as the ability of a model with the expressive power of at least the k-FWL to iteratively re-combine concepts from first principles (such as the MOTHER relation) as a hierarchy of statements in C k+1,t , containing all FO-logic statements with counting quantifiers, at most k + 1 variables, and quantifier depth t.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We provide full access to the code needed to reproduce our experiments. All datasets can be downloaded freely and are automatically downloaded and processed within our code. We provide detailed instructions on installation (including package versions) and execution of our code in the README file of our codebase. Guidelines:• The answer NA means that paper does not include experiments requiring code. We provide information about the data splits and how they were selected in Section 6. Further, for every dataset and benchmark, we detail the hyper-parameters and the optimizer used in Table 6 . Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. Justification: This paper conducts foundational research in the area of graph learning. While certainly our work could be used both for positive and negative societal impact, we do not foresee any immediate positive or negative societal impact. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experimental evaluation: 
 We now investigate how well the ET performs on various graph-learning tasks. We include tasks on graph-, node-, and edge-level. Specifically, we answer the following questions.Q1 How does the ET fare against other theoretically aligned architectures regarding predictive performance? Q2 How does the ET compare to state-of-the-art models? Q3 How effectively can the ET benefit from additional positional/structural encodings?The source code for our experiments is available at https://github.com/luis-mueller/ towards-principled-gts. To foster research in principled graph transformers such as the ET, we provide accessible implementations of ET, both in PyTorch and Jax.Datasets We evaluate the ET on graph-, node-, and edge-level tasks from various domains to demonstrate its versatility.On the graph level, we evaluate the ET on the molecular datasets ZINC (12K), ZINC-FULL [14] , ALCHEMY (12K), and PCQM4MV2 [21] . Here, nodes represent atoms and edges bonds between atoms, and the task is always to predict one or more molecular properties of a given molecule. Due to their relatively small graphs, the above datasets are ideal for evaluating higher-order and other resource-hungry models.On the node and edge level, we evaluate the ET on the CLRS benchmark for neural algorithmic reasoning [47] . Here, the input, output, and intermediate steps of 30 classical algorithms are translated into graph data, where nodes represent the algorithm input and edges are used to encode a partial ordering of the input. The algorithms of CLRS are typically grouped into eight algorithm classes: Sorting, Searching, Divide and Conquer, Greedy, Dynamic Programming, Graphs, Strings, and Geometry. The task is then to predict the output of an algorithm given its input. This prediction is made based on an encoder-processor-decoder framework introduced by Velickovic et al. [47] , which is recursively applied to execute the algorithmic steps iteratively. We will use the ET as the processor in this framework, receiving as input the current algorithmic state in the form of node and edge features and outputting the updated node and edge features, according to the latest version of CLRS, available at https://github.com/google-deepmind/clrs. As such, the CLRS requires the ET to make both node-and edge-level predictions.Finally, we conduct empirical expressivity tests on the BREC benchmark [49] . BREC contains 400 pairs of non-isomorphic graphs with up to 198 nodes, ranging from basic, 1-WL distinguishable CIN [8] 0.079 ±0.006 -0.022 ±0.002 Graphormer-GD [54] 0.081 ±0.009 -0.025 ±0.004 SignNet [30] 0.084 ±0.006 0.113 ±0.002 0.024 ±0.003 BasisNet [22] 0.155 ±0.007 0.110 ±0.001 -PPGN++ [41] 0.071 ±0.001 0.109 ±0.001 0.020 ±0.001 SPE [22] 0 Baselines On the molecular regression datasets, we compare the ET to an 1-WL expressive GNN baseline such as GIN(E) [52] .On ZINC (12K), ZINC-FULL and ALCHEMY, we compare the ET to other theoretically-aligned models, most notably higher-order GNNs [8, 37, 39] , Graphormer-GD, with strictly less expressive power than the 3-WL [54] , and PPGN++, with strictly more expressive power than the 3-WL [41] to study Q1. On PCQM4MV2, we compare the ET to state-of-the-art graph transformers to study Q2. To study the impact of positional/structural encodings in Q3, we evaluate the ET both with and without relative random walk probabilities (RRWP) positional encodings, recently proposed in Ma et al. [32] . RRWP encodings only apply to models with explicit representations over node pairs and are well-suited for the ET.On the CLRS benchmark, we mostly compare to the Relational Transformer (RT) [12] as a strong graph transformer baseline. Comparing the ET to the RT allows us to study Q2 in a different domain than molecular regression and on node-and edge-level tasks. Further, since the RT is similarly motivated as the ET in learning explicit representations of relations, we can study the potential benefits of the ET provable expressive power on the CLRS tasks. In addition, we compare the ET to DeepSet and GNN baselines in Diao and Loynd [12] and the single-task Triplet-GMPNN in Ibarz et al. [24] .On the BREC benchmark, we study questions Q1 and Q2 by comparing the ET to selected models presented in Wang and Zhang [49] . First, we compare to the δ-2-LGNN [37], a higher-order GNN with strictly more expressive power than the 1-WL. Second, we compare to Graphormer [53] , an empirically strong graph transformer. Third, we compare to PPGN [33] with the same expressive power as the ET. We additionally include the 3-WL results on the graphs in BREC to investigate how many 3-WL distinguishable graphs the ET can distinguish in BREC.Experimental setup See Table 6 for an overview of the used hyperparameters.For ZINC (12K), ZINC-FULL, and PCQM4MV2, we follow the hyperparameters in Ma et al. [32] .For ALCHEMY, we follow standard protocol and split the data according to Morris et al. [39] . Here, we simply adopt the hyper-parameters of ZINC (12K) from Ma et al. [32] but set the batch size to 64.We choose the same hyper-parameters as the RT for the CLRS benchmark. Also, following the RT, we train for 10K steps and report results over 20 random seeds. To stay as close as possible to the experimental setup of our baselines, we integrate our Jax implementation of the ET as a processor into the latest version of the CLRS code base. In addition, we explore the OOD validation technique presented in Jung and Ahn [25] , where we use larger graphs for the validation set to encourage size generalization. This technique can be used within the CLRS code base through the experiment parameters.Finally, for BREC, we keep the default hyper-parameters and follow closely the setup used by Wang and Zhang [49] for PPGN. We found learning on BREC to be quite sensitive to architectural choices, possibly due to the small dataset sizes. As a result, we use a linear layer for the FFN and additionally apply layer normalization onto X il W Q , X lj W K in Equation ( 2) and V ilj in Equation ( 3).For ZINC (12K), ZINC-FULL, PCQM4MV2, CLRS, and BREC, we follow the standard train/validation/test splits. For ALCHEMY, we split the data according to the splits in Morris et al. [39] , the same as our baselines.All experiments were performed on a mix of A10, L40, and A100 NVIDIA GPUs. For each run, we used at most 8 CPU cores and 64 GB of RAM, with the exception of PCQM4MV2 and ZINC-FULL, which were trained on 4 L40 GPUs with 16 CPU cores and 256 GB RAM.Table 3 : Number of distinguished pairs of non-isomorphic graphs on the BREC benchmark over 10 random seeds with standard deviation. Baseline results (over 1 random seed) are taken from Wang and Zhang [49] . For reference, we also report the number of graphs distinguishable by 3-WL. Results and discussion In the following, we answer questions Q1 to Q3. We highlight first , second, and third best results in each table.We compare results on the molecular regression datasets in Table 1 . On ZINC (12K) and ALCHEMY, the ET outperforms all baselines, even without using positional/structural encodings, positively answering Q1. Interestingly, on ZINC-FULL, the ET, while still among the best models, does not show superior performance. Further, the RRWP encodings we employ on the graph-level datasets improve the performance of the ET on all three datasets, positively answering Q3. Moreover, in Table 5 , we compare the ET with a variety of graph learning models on ZINC (12K), demonstrating that the ET is highly competitive with state-of-the-art models. We observe similarly positive results in Table 4 where the ET outperforms strong graph transformer baselines such as GRIT [32] , GraphGPS [42] and Graphormer [53] on PCQM4MV2. As a result, we can positively answer Q2. Val. MAE (↓) # Params EGT [23] 0.0869 89.3M GraphGPSSmall [42] 0.0938 6.2M GraphGPSMedium [42] 0.0858 19.4M TokenGTORF [28] 0.0962 48.6M TokenGTLap [28] 0.0910 48.5M Graphormer [53] 0.0864 48.3M GRIT [32] 0 [16] 0.083 ±0.003 Graphormer-GD [54] 0.081 ±0.009 CIN [8] 0.079 ±0.006 Graph-MLP-Mixer [20] 0.073 ±0.001 PPGN++ [41] 0.071 ±0.001 GraphGPS [42] 0.070 ±0.004 SPE [22] 0.069 ±0.004 Graph Diffuser [18] 0.068 ±0.002 Specformer [7] 0.066 ±0.003 GRIT [32] 0.059 ±0.002 ET 0.062 ±0.004 ET+RRWP 0.059 ±0.004In Table 2 , we compare results on CLRS where the ET performs best when averaging all tasks or when averaging all algorithm classes, improving over RT and Triplet-GMPNN. Additionally, the ET performs best on 4 algorithm classes and is among the top 3 in 7/8 algorithm classes. Interestingly, only some models are best on a majority of algorithm classes. These results indicate a benefit of the ETs' expressive power on this benchmark, adding to the answer of Q2. Further, see Table 7 in Appendix B.2 for additional results using the OOD validation technique.Finally, on the BREC benchmark, we observe that the ET cannot distinguish all graphs distinguishable by 3-WL. At the same time, the ET distinguishes more graphs than PPGN, the other 3-WL expressive model, providing an additional positive answer to Q1; see Table 3 . Moreover, the ET distinguishes more graphs than δ-2-LGNN and outperforms Graphormer by a large margin, again positively answering Q2. Overall, the positive results of the ET on BREC indicate that the ET is well able to leverage its expressive power empirically. 
limitations: 
 While proving to be a strong and versatile graph model, the ET has an asymptotic runtime and memory complexity of O(n 3 ), which is more expensive than most state-of-the-art models with linear or quadratic runtime and memory complexity. We emphasize that due to the runtime and memory complexity of the k-WL, a trade-off between expressivity and efficiency is likely unavoidable. At the same time, the ET is highly parallelizable and runs efficiently on modern GPUs. We hope that innovations for parallelizable neural networks can compensate for the asymptotic runtime and memory complexity of the ET. In Figure 4 in the appendix, we find that we can use low-level GPU optimizations, available for parallelizable neural networks out-of-the-box, to dampen the cubic runtime and memory scaling of the ET; see Appendix C for runtime and memory experiments and an extended discussion. 
conclusion: 
 We established a previously unknown connection between the Edge Transformer and 3-WL, and enabled the Edge Transformer for various graph learning tasks, including graph-, node-, and edge-level tasks. We also utilized a well-known connection between graph isomorphism testing and first-order logic to derive a theoretical interpretation of systematic generalization. We demonstrated empirically that the Edge Transformer is a promising architecture for graph learning, outperforming other theoretically aligned architectures and being among the best models on ZINC (12K), PCQM4MV2 and CLRS. Furthermore, the ET is a graph transformer that does not rely on positional/structural encodings for strong empirical performance. Future work could further explore the potential of the Edge Transformer in neural algorithmic reasoning and molecular learning by improving its scalability to larger graphs, in particular through architecture-specific low-level GPU optimizations and model parallelism.  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We release neither data nor models as part of this work. Further, our experiments are conducted on comparatively small, curated, task-specific datasets used for benchmarking graph learning models. Hence, our work does not pose immediate risks for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We provide data source and license information in Appendix B.1. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators.  
: 
 Given a graph G, we construct a 3D tensor where we embed information from each node pair into a d dimensional vector. developed for systematic generalization problems [6] , to the field of graph learning. Systematic (or compositional) generalization refers to the ability of a model to generalize to complex novel concepts by combining primitive concepts observed during training, posing a challenge to even the most advanced models such as GPT-4 [15] .x iℓ x ℓj x iℓ x ℓj Project Project Project Project Dot-Product & Scale Element-wise Product Softmax ℓ ℓ Project ∀(i, j) ∈ V (G) 2 1 × d 1 × d 1 × d 1 × d n × n × n n × n × n n × n × d n × n × n × d n × n × dSpecifically, we contribute the following:1. We propose a concrete implementation of the Edge Transformer that readily applies to various graph learning tasks. 2. We show theoretically that this Edge Transformer implementation is as expressive as the 3-WL without the need for positional/structural encodings. 3. We demonstrate the benefits of aligning models with the k-WL hierarchy by leveraging well-established results from graph theory and logic to develop a theoretical understanding of systematic generalization in terms of first-order logic statements. 4. We demonstrate the superior empirical performance of the resulting architecture compared to a variety of other theoretically motivated models, as well as competitive performance compared to state-of-the-art models in molecular regression and neural algorithmic reasoning tasks.the ET offers provable 3-WL expressivity with O(n 3 ) runtime and memory complexity, several orders of magnitude more efficient than the corresponding 3-WL expressive transformer in Kim et al. [28] .For an overview of graph transformers, see Müller et al. [40] .Finally, systematic generalization has recently been investigated both empirically and theoretically [6, 15, 26, 43] . In particular, Dziri et al. [15] demonstrate that compositional generalization is lacking in state-of-the-art transformers such as GPT-4. 
 
 Appendix: Here, we present details about implementing the ET in practice.In what follows, we propose a pooling method from node pairs to nodes, which allows us also to make predictions for node-and graph-level tasks. For each node i ∈ V (G), we computeReadOut(i) := j∈[n] ρ 1 X (L) ij + ρ 2 X (L) ji ,where ρ 1 , ρ 2 are neural networks and X (L) is the node pair tensor after L ET layers. We apply ρ 1 to node pairs where node i is at the first position and ρ 2 to node pairs where node i is at the second position. We found that making such a distinction has positive impacts on empirical performance. Then, for graph-level predictions, we first compute node-level readout as above and then use common graph-level pooling functions such as sum and mean [51] or set2seq [48] on the resulting node representations. We use this readout method in our molecular regression experiments in Section 6. [25] , for Triplet-GMPNN [24] and ET, respectively.In Table 7 , following [25] , we present additional experimental results on CLRS when using graphs of size 32 in the validation set. We compare to both the Triplet-GMPNN [24] , as well as the TEAM [25] baselines. In addition, in Figure 3 , we present a comparison of the improvements resulting from the OOD validation technique, comparing Triplet-GMPNN and the ET. Finally, in Table 8 , we compare different modifications to the CLRS training setup that are agnostic to the choice of processor.We present detailed results for the algorithms in CLRS. See Table 11 for divide and conquer algorithms, Table 12 for dynamic programming algorithms, Table 13 for geometry algorithms, Table 15 for greedy algorithms, Table 10 for search algorithms, Table 9 for sorting algorithms, and Table 16 for string algorithms. We compare the standard deviation of Deep Sets, GAT, MPNN, PGN, RT, and ET following the comparison in Diao and Loynd [12] . Table 17 compares the standard deviation over all algorithms in the CLRS benchmark. We observe that the ET has the lowest overall standard deviation. The table does not contain results for Triplet-GMPNN [24] since we do not have access to the test results for each algorithm on each seed that are necessary to compute the overall standard deviation. However, Table 18 compares the standard deviation per algorithm class between Triplet-GMPNN and the ET. We observe that Triplet-GMPNN and the ET have comparable standard deviations except for search and string algorithms, where Triplet-GMPNN has a much higher standard deviation than the ET.Here, we provide additional information on the runtime and memory requirements of the ET in practice. Specifically, in Figure 4 , we provide runtime scaling of the ET with and without low-level GPU optimizations in PyTorch on an A100 GPU with bfloat16 precision. We measure the time for the forward pass of a single layer of the ET on a single graph (batch size of 1) with n ∈ {100, 200, ..., 700} nodes and average the runtime over 100 repeats. We sample random Erdős-Renyi graphs with edge probability 0.05. We use an embedding dimension of 64 and two attention heads. We find that the automatic compilation into Triton [45] , performed automatically by torch.compile, improves the runtime and memory scaling. Specifically, with torch.compile enabled, the ET layer can process graphs with up to 700 nodes and shows much more efficient runtime scaling with the number of nodes. 4 : Runtime of the forward pass of a single ET layer in PyTorch in seconds for graphs with up to 700 nodes. We compare the runtime with and without torch.compile (automatic compilation into Triton [45] ) enabled. Without compilation, the ET goes out of memory after 600 nodes. [45] . We use torch.compile in our molecular regression experiments. In addition, we want to highlight FlashAttention [11] , available for the standard transformer, as an example of architecture-specific hardware optimizations that can reduce runtime and memory requirements.Runtime per dataset/benchmark Here, we present additional runtime results for all of our datasets. We present the runtime of a single run on a single L40 GPU of ZINC (12K), ALCHEMY (12K), and BREC. For ZINC-FULL and PCQM4MV2, we present the runtime of a single run on 4 L40 GPUs; see Table 20 .On CLRS, the experiments in our work are run on a mix of A10 and A100 GPUs. To enable a fair comparison, we rerun each algorithm in CLRS in a single run on a single A100 GPU and report the corresponding runtime in Table 19 . Finally, we note that these numbers only reflect the time to run the final experiments and significantly more time was used for preliminary experiments over the course of the research project.Here, we define our notation. Let N := {1, 2, 3, . . . }. For n ≥ 1, let [n] := {1, . . . , n} ⊂ N. We use { {. . . } } to denote multisets, i.e., the generalization of sets allowing for multiple instances for each of its elements.Graphs A (node-)labeled graph G is a triple (V (G), E(G), ℓ) with finite sets of vertices or nodes V (G), edges E(G) ⊆ {{u, v} ⊆ V (G) | u ̸ = v} and a (node-)label function ℓ : V (G) → N. Then ℓ(v) is a label of v, for v in V (G).If not otherwise stated, we set n := |V (G)|, and the graph is of order n. We also call the graph G an n-order graph. For ease of notation, we denote the edge {u, v} in E(G) by (u, v) or (v, u). We define an n-order attributed graph as a pair G = (G, F ), where G = (V (G), E(G)) and F in R n×p for p > 0 is a node feature matrix. Here, we identify V (G) with [n], then F (v) in R 1×p is the feature or attribute of the node v ∈ V (G). Given a labeled graph (V (G), E(G), ℓ), a node feature matrix F is consistent with ℓ if ℓ(v) = ℓ(w) for v, w ∈ V (G) if and only if F (v) = F (w).The neighborhood of a vertex v in V (G) is denoted by N (v) := {u ∈ V (G) | (v, u) ∈ E(G)} and the degree of a vertex v is |N (v)|. Two graphs G and H are isomorphic and we write G ≃ H if there exists a bijection φ : V (G) → V (H) preserving the adjacency relation, i.e., (u, v) is in E(G) if and only if (φ(u), φ(v)) is in E(H). Then φ is an isomorphism between G and H. In the case of labeled graphs, we additionally require that l(v) = l(φ(v)) for v in V (G), and similarly for attributed graphs. Moreover, we call the equivalence classes induced by ≃ isomorphism types and denote the isomorphism type of G by τ G . We further define the atomic type atp : V (G) k → N, for k > 0, such that atp(v) = atp(w) for v and w in V (G) k if and only if the mapping φ :V (G) k → V (G) k where v i → w i induces a partial isomorphism, i.e., we have v i = v j ⇐⇒ w i = w j and (v i , v j ) ∈ E(G) ⇐⇒ (φ(v i ), φ(v j )) ∈ E(G).Matrices Let M ∈ R n×p and N ∈ R n×q be two matrices then [M N ] ∈ R n×(p+q) denotes column-wise matrix concatenation. We also write R d for R 1×d . Further, let M ∈ R p×n and N ∈ R q×n be two matrices then M N ∈ R (p+q)×n denotes row-wise matrix concatenation.For a matrix X ∈ R n×d , we denote with X i the ith row vector. In the case where the rows of X correspond to nodes in a graph G, we use X v to denote the row vector corresponding to the node v ∈ V (G).The Weisfeiler-Leman algorithm We describe the Weisfeiler-Leman algorithm, starting with the 1-WL. The 1-WL or color refinement is a well-studied heuristic for the graph isomorphism problem, originally proposed by Weisfeiler and Leman [50] . 1 Intuitively, the algorithm determines if two graphs are non-isomorphic by iteratively coloring or labeling vertices. Formally, let G = (V, E, ℓ) be a labeled graph, in each iteration, t > 0, the 1-WL computes a vertex coloring C 1 t : V (G) → N, depending on the coloring of the neighbors. That is, in iteration t > 0, we setC 1 t (v) := recolor C 1 t-1 (v), { {C 1 t-1 (u) | u ∈ N (v)} } ,for all vertices v in V (G), where recolor injectively maps the above pair to a unique natural number, which has not been used in previous iterations. In iteration 0, the coloring C 1 0 := ℓ. To test if two graphs G and H are non-isomorphic, we run the above algorithm in "parallel" on both graphs. If the two graphs have a different number of vertices colored c in N at some iteration, the 1-WL distinguishes the graphs as non-isomorphic. It is easy to see that 1-WL cannot distinguish all non-isomorphic graphs [10] .The k-dimensional Weisfeiler-Leman algorithm Due to the shortcomings of the 1-WL or color refinement in distinguishing non-isomorphic graphs, several researchers, e.g., Babai [3] , Cai et al. [10] , devised a more powerful generalization of the former, today known as the k-dimensional Weisfeiler-Leman algorithm (k-WL), operating on k-tuples of nodes rather than single nodes.Intuitively, to surpass the limitations of the 1-WL, the k-WL colors node-ordered k-tuples instead of a single node. More precisely, given a graph G, the k-WL colors the tuples from V (G) k for k ≥ 2 instead of the nodes. By defining a neighborhood between these tuples, we can define a coloring similar to the 1-WL. Formally, let G be a graph, and let k ≥ 2. In each iteration, t ≥ 0, the algorithm, similarly to the 1-WL, computes a coloring C k t : V (G) k → N. In the first iteration, t = 0, the tuples v and w in V (G) k get the same color if they have the same atomic type, i.e., atp k (v) = atp k (u). Then, for each iteration, t > 0, C k t is defined byEQUATIONwith M t (v) the multisetEQUATIONand where ϕ j (v, w) := (v 1 , . . . , v j-1 , w, v j+1 , . . . , v k ). That is, ϕ j (v, w) replaces the j-th component of the tuple v with the node w. Hence, two tuples are adjacent or j-neighbors if they are different in the jth component (or equal, in the case of self-loops). Hence, two tuples v and w with the same color in iteration (t -1) get different colors in iteration t if there exists a j in [k] such that the number of j-neighbors of v and w, respectively, colored with a certain color is different.We run the k-WL algorithm until convergence, i.e., until for t in NC k t (v) = C k t (w) ⇐⇒ C k t+1 (v) = C k t+1 (w), for all v and w in V (G) k holds.Similarly to the 1-WL, to test whether two graphs G and H are non-isomorphic, we run the k-WL in "parallel" on both graphs. Then, if the two graphs have a different number of nodes colored c, for c in N, the k-WL distinguishes the graphs as non-isomorphic. By increasing k, the algorithm gets more powerful in distinguishing non-isomorphic graphs, i.e., for each k ≥ 2, there are non-isomorphic graphs distinguished by (k + 1)-WL but not by k-WL [10] .The folklore k-dimensional Weisfeiler-Leman algorithm A common and well-studied variant of the k-WL is the k-FWL, which differs from the k-WL only in the aggregation function. Instead of Equation ( 6), the "folklore" version of the k-WL updates k-tuples according toM F t (v) := { {(C k,F t-1 (ϕ 1 (v, w)), ..., C k,F t-1 (ϕ k (v, w))) | w ∈ V (G)} }, resulting in the coloring C k,F t : V (G) k → N,and is strictly more powerful than the k-WL. Specifically, for k ≥ 2, the k-WL is exactly as powerful as the (k -1)-FWL [19] .Computing k-WL's initial colors Let G = (V (G), E(G), ℓ) be a labeled graph, k ≥ 2, and let v := (v 1 , . . . , v k ) ∈ V (G) k be a k-tuple. Then, we can present the atomic type atp(v) by a k × k matrix K over {1, 2, 3}. That is, the entryK ij is 1 if (v i , v j ) ∈ E(G), 2 if v i = v j ,and 3 otherwise. Further, we ensure consistency with ℓ, meaning that for two k-tuples v := (v 1 , . . . , v k ) ∈ V (G) k and w := (w 1 , . . . , w k ) ∈ V (G) k , then C k 0 (v) = C k 0 (w), if and only if, atp(v) = atp(w) and ℓ(v i ) = ℓ(w i ), for all i ∈ [k]. Note that we compute the initial colors for both k-WL and the k-FWL in this way.We begin with a short review of Cai et al. [10] . We consider our usual node-labeled graph G = (V (G), E(G), ℓ) with n nodes. However, we replace ℓ with a countable set of color relations C 1 , . . . , C n , where for a node v ∈ V (G),C i (v) ⇐⇒ ℓ(v) = i.Note that Cai et al. [10] consider the more general case where nodes can be assigned to multiple colors simultaneously. However, for our work, we assume that a node is assigned to precisely one color, and hence, the set of color relations is at most of size n. We can construct first-order logic statements about G. For example, the following sentence describes the existence of a triangle formed by two nodes with color 1:∃x 1 ∃x 2 ∃x 3 E(x 1 , x 2 ) ∧ E(x 1 , x 3 ) ∧ E(x 2 , x 3 ) ∧ C 1 (x 1 ) ∧ C 1 (x 2 ) .Here, x 1 , x 2 , and x 3 are variables which can be repeated and re-quantified at will. Statements made about G and a subset of nodes in V (G) are of particular importance to us. To this end, we define a k-configuration, a function f : {x 1 , . . . , x k } → V (G) that assigns a node in V (G) to each one of the variables x 1 , . . . , x k . Let φ be a first-order formula with free variables among x 1 , . . . , x k . Then, we write G, f |= φ if φ is true when the variable x i is interpreted as the node f (x i ), for i = 1, . . . , k.Cai et al. [10] define the language C k,m of all first-order formulas with counting quantifiers, at most k variables, and quantifier depth bounded by m, and the language C k = m≥0 C k,m . For example, the sentence ∀x∃!3y E(x, y) in C 2 describes 3-regular graphs; i.e., graphs where each vertex has exactly 3 neighbors.We define the equivalence relation ≡ k,m over pairs (G, f ) made of graphs G and k-configurations f as (G, f ) ≡ k,m (H, g) if and only ifG, f |= φ ⇐⇒ H, g |= φfor all formulas φ in C k,m whose free variables are among x 1 , . . . , x k .We can now formulate a main result of Cai et al. [10] . Let G and H be two graphs, let k ≥ 1 and m ≥ 0 be non-negative integers, and let f and g be k-configurations for G and H respectively. Ifu = (f (x 1 ), . . . , f (x k )) ∈ V (G) k and v = (g(x 1 ), . . . , g(x k )) ∈ V (H) k , then C k,F m (u) = C k,F m (v) ⇐⇒ (G, f ) ≡ k,m (H, g) .Here, we first generalize the GNN from Grohe [19] to the 2-FWL. Higher-order GNNs with the same expressivity have been proposed in prior works by Azizian and Lelarge [1] . However, our GNNs have a special form that can be computed by the Edge Transformer.Formally, let S ⊆ N be a finite subset. First, we show that multisets over S can be injectively mapped to a value in the closed interval (0, 1), a variant of Lemma VIII.5 in Grohe [19] . Here, we outline a streamlined version of its proof, highlighting the key intuition behind representing multisets as m-ary numbers. Let M ⊆ S be a multiset with multiplicities a 1 , . . . , a k and distinct k values. We define the order of the multiset as k i=1 a i . We can write such a multiset as a sequence x (1) , . . . , x (l) where l is the order of the multiset. Note that the order of the sequence is arbitrary and that for i ̸ = j it is possible to have x (i) = x (j) . We call such a sequence an M -sequence of length l. We now prove a slight variation of a result of Grohe [19] . Lemma 4. For a finite m ∈ N, let M ⊆ S be a multiset of order m -1 and let x i ∈ S denote the ith number in a fixed but arbitrary ordering of S. Given a mapping g : S → (0, 1) whereg(x i ) := m -i ,and an M -sequence of length l given by x (1) , . . . , x (l) with positions i (1) , . . . , i (l) in S, the sumj∈[l] g(x (j) ) = j∈[l] m -i (j)is unique for every unique M .Proof. By assumption, let M ⊆ S denote a multiset of order m -1. Further, let x (1) , . . . , x (l) ∈ M be an M -sequence with i (1) , . . . , i (l) in S. Given our fixed ordering of the numbers in S we can equivalently write M = ((a 1 , x 1 ), . . . , (a n , x n )), where a i denotes the multiplicity of ith number in M with position i from our ordering over S. Note that for a number m -i there exists a corresponding m-ary number written as 0.0 . . .1 i . . .Then the sum,j∈[l] g(x (j) ) = j∈[l] m -i (j) = i∈S a i m -i ∈ (0, 1)and in m-ary representation 0.a 1 . . . a n .Note a i = 0 if and only if there exists no j such that i (j) = i. Since the order of M is m -1, it holds that a i < m. Hence, it follows that the above sum is unique for each unique multiset M , implying the result.Recall that S ⊆ N and that we fixed an arbitrary ordering over S. Intuitively, we use the finiteness of S to map each number therein to a fixed digit of the numbers in (0, 1). The finite m ensures that at each digit, we have sufficient "bandwidth" to encode each a i . Now that we have seen how to encode multisets over S as numbers in (0, 1), we review some fundamental operations about the m-ary numbers defined above. We will refer to decimal numbers m -i as corresponding to an m-ary number 0.0 . . .1 i . . . ,where the ith digit after the decimal point is 1 and all other digits are 0, and vice versa.To begin with, addition between decimal numbers implements counting in m-ary notation, i.e., m -i + m -j corresponds to 0.0 . . . Shifting further applies to general decimal numbers in (0, 1). Let x ∈ (0, 1) correspond to an m-ary number with l digits, 0.a 1 . . . a l .Then, m -i • x corresponds to 0.0 . . . 0 a 1 . . . a l i+1,...,i+l .Before we continue, we show a small lemma stating that two non-overlapping sets of m-ary numbers preserve their uniqueness under addition. Lemma 5. Let A and B be two sets of m-ary numbers for some m > 1. Ifmin x∈A x > max y∈B y, then for any x 1 , x 2 ∈ A, y 1 , y 2 ∈ B, x 1 + y 1 = x 2 + y 2 ⇐⇒ x 1 = x 2 and y 1 = y 2 .Proof. The statement follows from the fact that ifmin x∈A x > max y∈B y,then numbers in A and numbers in B do not overlap in terms of their digit range. Specifically, there exists some l > 0 such that we can writex := 0.x 1 . . . x l y := 0. 0 . . . 0 l y 1 . . . y k ,for some k > l and all x ∈ A, y ∈ B. As a result,x + y = 0.x 1 . . . x l y 1 . . . y k .Hence, x + y is unique for every unique pair (x, y). This completes the proof.We begin by showing the following proposition, showing that the tokenization in Equation ( 4) is sufficient to encode the initial node colors under 2-FWL. Proposition 6. Let G = (V (G), E(G), ℓ) be a node-labeled graph with n nodes. Then, there exists a parameterization of Equation ( 4)with d = 1 such that for each 2-tuples u, v ∈ V (G) 2 , C 2,F 0 (u) = C 2,F 0 (v) ⇐⇒ X(u) = X(v).Proof. The statement directly follows from the fact that the initial color of a tuple u := (i, j) depends on the atomic type and the node labeling. In Equation ( 4), we encode the atomic type with E ij and the node labels with[E ij F i F j ]The concatenation of both node labels and atomic type is clearly injective. Finally, since there are at most n 2 distinct initial colors of the 2-FWL, said colors can be well represented within R, hence there exists an injective ϕ in Equation ( 4) with d = 1. This completes the proof.We now show Theorem 1. Specifically, we show the following two propositions from which Theorem 1 follows. Proposition 7. Let G = (V (G), E(G), ℓ) be a node-labeled graph with n nodes and F ∈ R n×p be a node feature matrix consistent with ℓ. Then for all t ≥ 0, there exists a parametrization of the ET such thatC 2,F t (v) = C 2,F t (w) ⇐= X (t) (v) = X (t) (w), for all pairs of 2-tuples v and w ∈ V (G) 2 .and by the induction hypothesis above,X (t-1) jl = T (t-1) jl T (t-1) jl n 2 X (t-1) lm = T (t-1) lm T (t-1) lm n 2, where we expanded sub-matrices. Specifically,W V1 1 , W V2 1 , W V1 2 , W V2 2 ∈ R d 2 ×d. We then setW Q = W K = 0 W V1 1 = [βI 0] W V1 2 = [0 0] W V2 1 = [0 I] W V2 2 = [0 0].Here, W Q and W K are set to zero to obtain uniform attention scores. Note that then for all j, l, k, α jlm = 1 n , due to normalization over l, and we end up with Equation ( 9) ash 1 (X (t-1) ) jm = 1 n n l=1 V (t-1) jlmwhereV (t-1) jlm = T (t-1) jl • βI + T (t-1) jl n 2 • 0 0 ⊙ T (t-1) lm • 0 + T (t-1) lm n 2 • I 0 = β • T (t-1) jl • T (t-1) lm n 2 0 .We now conclude our proof as follows. Recall that the Edge Transformer layer computes the final representation X (t) asX (t) jm = FFN X (t-1) jm + h 1 (X (t-1) ) jm W O = FFN T (t-1) jm T (t-1) jm n 2 + β n n l=1 T (t-1) jl • T (t-1) lm 0 W O = W O :=I FFN T (t-1) jm T (t-1) jm n 2 + β n n l=1 T (t-1) jl • T (t-1) lm 0 = FFN T (t-1) jm + β n n l=1 T (t-1) jl • T (t-1) lm T (t-1) jm n 2 = Eq.8 FFN T (t) jm T (t-1) jm n 2for some FFN. Note that the above derivation only modifies the terms inside the parentheses and is thus independent of the choice of FFN. We have thus shown that the ET can compute Equation (8).To complete the induction, let f : R 2 → R 2 be such thatf T (t) jm T (t-1) jm n 2 = T (t) jm T (t) jm n 2.Since our domain is compact, f is continuous, and hence we can choose FFN to approximate f arbitrarily close. This completes the proof.Next, we show the other direction of Theorem 1 under mild and reasonable assumptions. First, we say that a recoloring function, that maps structures over positive integers into positive integers, is (effectively) invertible if its inverse is computable. All coloring functions used in practice (e.g., hashbased functions, those based on pairing functions, etc) are invertible. Second, the layer normalization operation is a proper function if it uses statistics collected only during training mode, and not during evaluation mode. Proposition 8. Let recolor be an invertible function, and let us consider the 2-FWL coloring algorithm using recolor. Then, for all parametrizations of the ET with proper layer normalization, for all node-labeled graphs G = (V (G), E(G), ℓ), and for all t ≥ 0:C 2,F t (v) = C 2,F t (w) =⇒ X (t) (v) = X (t) (w), for all pairs of 2-tuples v and w in V (G) 2 .Proof. We first claim that there is a computable function Z : 2 , independent of the graph G and its order. The proof of the claim is by induction on t. For t = 0, by definition, C 2,F 0 (v) identifies the atomic type atp 2 (v) which defines X (0) (v) (since the atomic type tells if v is an edge in G, and the labels of the vertices in v).N * × N → R p , where N * = {0} ∪ N, such that X (t) (v) = Z(t, C 2,F t (v)) for all v ∈ V (G)For t > 0 and v = (i, j), the function Z(t, C 2,F t (v)) proceeds as follows. First, it uses the invertibility of recolor to obtain the pairC 2,F t-1 (i, j), { { C 2,F t-1 (i, l), C 2,F t-1 (l, j) | l ∈ V (G)} } .Then, by inductive hypothesis using the function Z(t -1, •), it obtains the pairX (t-1) (i, j), { { X (t-1) (i, l), X (t-1) (l, j) | l ∈ V (G)} } .Finally, it computesX (t) (i, j) = FFN X (t-1) (i, j) + TriAttention LN X (t-1) (i, j)under the assumption that the layer normalization is a proper function. The statement of the proposition then follows directly from the claim sinceX (t) (v) = Z(t, C 2,F t (v)) = Z(t, C 2,F t (w)) = X (t) (w) .Note that unlike the result in Proposition 7, the above result is uniform, in that the concrete choice of recolor and the function Z does not depend on the graph size n. Finally, Theorem 1 follows from Proposition 7 and Proposition 8.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: Our main claims are that the ET has 3-WL expressive power, which we prove in Theorem 1 and that the ET surpasses theoretically aligned graph models, which we demonstrate in Table 1 and Table 3 . Further, we claim that the ET is competitive with state-of-the-art models on a variety of tasks, which we demonstrate in Table 5 and Table 2 .Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Strictly speaking, the 1-WL and color refinement are two different algorithms. The 1-WL considers neighbors and non-neighbors to update the coloring, resulting in a slightly higher expressive power when distinguishing vertices in a given graph; see[19] for details. For brevity, we consider both algorithms to be equivalent. 


Decision: This paper analyzes the expressive power of Edge Transformer (graph transformer operating on node pairs) and shows that it has equivalent expressive power as 3-WL without using PEs. The paper receives 1, 7, 5, 6 scores after rebuttal, with reviewer ckA7 decreasing their score from 3 to 1. I carefully checked the reviews and rebuttals and deemed that the review content does not justify the score, and the authors have successfully responded to ckA7's concerns. Further, reviewer ckA7 did not participate into the discussion. Therefore, their score is downweighted and a final acceptance recommendation is made.

Overall, I encourage the authors to include their new results in the rebuttal, propertly address all the reviewers' concerns, and work further on addressing the computational complexity of ET to make it practically more useful (on large graphs).

Reviews: review: paper_summary: paper_summary: This work proposes an Edge Transformer (ET), a global attention model operating on node pairs instead of nodes. Authors theoretically demonstrate that ET has 3-WL expressive power with the proper tokenization. Experimental results demonstrate that the proposed model outperforms other theoretically aligned models in terms of predictive performance. Additionally, ET competes with state-of-the-art models in algorithmic reasoning and molecular regression tasks without relying on positional or structural encodings.
strength_weakness: strength_weakness: 1. Previous work [1] has demonstrated that graph transformers (GTs) with proper positional encodings (PEs) can be more powerful than any WL test. Therefore, it is unclear why the authors aim to develop a model with 3-WL expressive power, which cannot be more expressive than prior GTs (e.g., [1]). Notably, given the cubic complexity of the proposed ET, computing PEs is no longer a scalability bottleneck. Consequently, if ET empirically outperforms GT+PE, it essentially implies that the WL hierarchy is not a meaningful metric in practice, undermining the purpose of building a 3-WL expressive model. 
2. While the authors mention how to apply ET for node-level tasks, they provide no corresponding results. Given the cubic complexity of ET, it is likely not scalable to most graphs for node-level tasks.
3. Some relevant GT models haven't been discussed or compared (e.g., [1-3]).

[1] Kreuzer et al., "Rethinking Graph Transformers with Spectral Attention", NeurIPS'21. \
[2] Zhu et al., "On Structural Expressive Power of Graph Transformers", KDD'23. \
[3] Deng et al., "Polynormer: Polynomial-Expressive Graph Transformer in Linear Time", ICLR'24.
questions: questions: N/A
limitations: limitations: Please refer to the Weaknesses section for details.
paper_summary: paper_summary: This paper applies Edge Transformer (ET) to the field of graph learning. It is proved that ET has 3-WL expressivity with cubic complexity, which is more efficient than existing 3-WL models. Experiments on BREC benchmark clearly demonstrate its expressive power. Results on the other three molecular datasets and CLRS benchmark are also provided.
strength_weakness: strength_weakness: - My major concern is about empirical evaluation.
    - While I appreciate ET's excellent theoretical expressive power, I find it is needed to include more benchmark datasets in graph learning (e.g., OGB, moleculeNet, GNN benchmark, Long Rang Graph Benchmark, Zinc-full, etc) to comprehensively demonstrate its practical predictive power compared with state-of-the-art (SOTA) methods, e.g., those graph transformers Grit, Exphormer, GraphGPS, etc.
    - Besides datasets from algorithmic reasoning, currently, in my opinion, ET is only evaluated on two widely-used graph benchmark datasets with performance reported from very recent SOTA methods, i.e., Zinc-12k and Alchemy-12k, making it hard to assess ET's practical performance. Although QM9 is also a popular dataset, the adopted setting is less common. Meanwhile, Zinc-12k is just a small subset of Zinc-full, and it has recently been shown that being good at Zinc-12k does not necessarily lead to the best performance on Zinc-full [1].
    - It would be interesting to see its long-range modeling capability on the LRGB benchmark as it's a transformer.
    - 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.


[1] Wang, X., Li, P., & Zhang, M. (2024). Graph as Point Set. ICML 2024.
questions: questions: - How does the setting of QM9 used in this paper differ from the setting that is commonly used in equivariant GNNs (e.g., EGNN, DimeNet, etc)? It would be better if this description could be put in the paper, making it self-contained.
- Are there any particular benefits to having an expressive model without using PE or SE? I consider the major benefit would be on the computation side as computing PE can be $O(n^2)$, but ET has already had $O(n^3)$ complexity. So, for the claim in the paper, is it just 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?
    - I feel like depending on how PE is used, it is still possible to analyze the expressive power, e.g.,  via graph substructure counting. Can the authors discuss the works such as graph as point set and those subgraph-GNNs? Because they may also use PE, while proving their expressive power via graph substructure counting.
limitations: limitations: See above.
paper_summary: paper_summary: In this paper, the authors show that Edge Transformer, a global attention model operating on node pairs, has 3-WL expressive power when provided with the right tokenization. Experiments results also show that the Edge Transformer has competitive performance on  molecular regression tasks and algorithmic reasoning.
strength_weakness: strength_weakness: 1. My major concern is the real usefulness of ET, since the high runtime and memory complexity may offset the importance of ET, though the authors discuss it in Limitations. 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.
2. It seems that the ET in this paper is mainly from existing work, and the difference needs to be explained.
questions: questions: 1. The ET in this paper is a variant of existing ones, due to the specific tokenization. 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. It would be good if the authors provide more details.
2. 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.
3. The authors emphasize that the designed ET can achieve superior performance even without positional encodings. But maybe as an open question, if incorporated with PEs, will the theoretical results change? 

[1] Comparing Graph Transformers via Positional Encodings, ICML 2024
limitations: limitations: Scalability limitations are discussed.
paper_summary: paper_summary: This paper introduces a novel architecture called Edge Transformer (ET) for graph learning tasks. The ET leverages global attention on node pairs instead of individual nodes and demonstrates strong empirical performance without relying on positional or structural encodings. The authors show that ET has the expressive power of 3-WL and achieves competitive results on algorithmic reasoning and molecular regression tasks.
strength_weakness: strength_weakness: 1、The total number of parameters for the model is not provided in the parameter table.
2、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.
3、The author's contribution is confined to applying the ET model to the graph, without proposing their own method.
questions: questions: 1、Can you provide possible ways to reduce the complexity of ET?
2、Can you provide the total number of parameters for ET on each dataset? See Weaknesses 1.
3、Why has the performance changed so much compared to the submitted version of ICML in 2024? Especially on the Searching dataset.
4、What do Average and All algorithms represent respectively?
limitations: limitations: See Weaknesses.

Comments: title: Follow-up
comment: Considering that the performance of this version achieves about 30% improvement than that of ICML version on Searching dataset, while it decreases on Strings and Geometry datasets, I believe the emprical results are not convincing. Moreover, I think the response is not reasonable.
Thank the authors for clarifications and explaining the main contributions. Though this paper does not propose a new model, it is a solid paper. I am happy to increase my rating.
I have read the authors' response and other reviews and intend to keep the original rating.
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.