Towards Principled Graph Transformers
Authors: Luis Müller
Daniel Kusuma
Blai Bonet
Christopher Morris
Venue: NeurIPS 2024 poster (2024)
Keywords: graph transformers
expressivity
Weisfeiler and Leman
Weisfeiler and Lehman
Primary area: graph_neural_networks
TL;DR: We study graph transformers that are both theoretically grounded in the Weisfeiler-Leman hierarchy as well as perform comparative with state-of-the-art on graph learning benchmarks.

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.

Towards Principled Graph Transformers
Luis Müller RWTH Aachen University luis www.mueellerrrr www.-awchee.nc.. de
Daniel Kusuma RWTH Aachen University
Blai Bonet
Christopher Morris RWTH Aachen University
Universitat Pompeu Fabra
Abstract
The expressive power of graph learning architectures based on the k-dimensional Weisfeiler-Leman (k-WL) hierarchy is well understood. However, such architec- tures often fail to deliver solid predictive performance on real-world tasks, lim- iting their practical impact. In contrast, global attention-based models such graph transformers demonstrate strong performance in practice. However, com- paring their expressivity with the k-WL hierarchy remains challenging, particu- larly since attention- based architectures rely on positional or structural encod- ings. 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 ar- chitectures regarding predictive performance and is competitive with state-of-the- art models on algorithmic reasoning and molecular regression tasks while not relying on positional or structural encodings. Our code is available at https; //github. com/ luis-muellet wtwemmmeeleetttttdedeeee pwwpnds--iininnppeeeeee ggg
Introduction
Graph Neural Networks (GNNs) are the de-facto standard in graph learning 17 44 29 51 but suffer from limited expressivity in distinguishing non-isomorphic graphs in terms of the 1-dimensional Peeiseellr-Lemade algorithm (1-WL) 51]. Hence, recent works introduced higher-order GNNs, aligned with the k-dimensional Weisfeiler-Leman (k-WL) hierarchy for graph isomorphism testing The k-WL hierarchy 10 55), offering a deep theoretical understanding of k-WL-aligned GNNs. While theoretically intriguing. higher-order GNNs often fail to deliver stte-of-the-att performance on real-world problems, making theoretically grounded models less relevant in practice I 37 39 In contrast, graph transformers [8
36. 37 39], resulting in more expressivity with an increase in
draws from a rich history in graph theory and logic [3] 4
5
recently demonstrated srate-of-the- art empirical performance. However, they draw their expressive power mostly from positionaa//srructuaal encodings (PEs), making difficult understand these models in terms of an expressivity hierarchy such as the WL. While a few works theoretically aligned graph we are not aware any works reporting empirical results for 3-WL-equivalent graph transformers on established graph learning datasets.
20 32
53
transformers with the k-WL hierarchy (2 28
In this work, aim to set the ground for graph learning architectures that are theoretically aligned with the higher-order Weisfeiler-Leman hierarchy while delivering strong empirical performance and, the same time, demonstrate that such an alignment creates powerful synergies between transformers and graph learning. Hence, we close the gap between theoretical expressivity and real-world predictive power. To this end, apply the Edge Transformer (ET) architecture, initially
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Wi.2 e V(Gy
UotsCrodoct Scale
Slement-wian Product
(a) Graph
(b) Node pair teusor X
Figure I: Tokenization the Edge Transformer. Given graph G, we construct 3D tensor where we embed information from each node pair into dimensional vector.
Figure 2: Tensor operations a single triangular attention head; see Algorithm| I comparison to standard attention in pseudo-code.
developed for systematic generalization problems [6]. to the field of graph learning. Systematic (or compositional) generalization refers to the ability of model generalize complex novel concepts by combining primitive concepts observed during training, posing challenge to even the most advanced models such GPT-4 [15].
Specifically, we contribute the following:
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 as expressive as the 3-WL without the need for positionnal/strucuuaal encodings.
3. We demonstrate the benefits of aligning models with k-WL hierarchy by leveraging well-established results from graph theory and logic to develop thcoretical understanding systematic generalization in terms first-order logic statements.
4. We demonstrate the superior empirical performance of the resulting architecture compared to a variety other theoretically motivated models, as well as competitive performance compared to ffte-of-theear models in molecular regression and neural algorithmic reasoning tasks.
2 Related work
Many graph learning models with higher-order WL expressive power exist, notably 8-k-GNNs [3] SetqNets
39 k-IGNN 35 34 PPGN [33, and the more recent PPGN+ [T] Moreover, Lipman
et al. 31 devise a low-rank attention module possessing the same power as the folklore 2-WL
and Bodnar et Weisfeiler-Leman in graph learning, see Morris et al. [38]
[ propose CIN with an expressive power of least 3-WL. For an overview of
Many graph transformers exist, notably Graphormer art graph transformers typically rely on posittionll/srrucuuaal encodings, which makes it challenging to derive a theoretical understanding their expressive power. The Relational Transformer (RT) [12] operates over both nodes and edges and, similar to the ET, builds relational representations, that is, representations on edges. Although the RT integrates edge information into self-attention and hence does not need to resort to positional/struucural encodings, the RT is theoretically poorly understood, much like other graph transformers. Graph transformers with higher-order expressive power are Graphormer-O0 GD [54] and TokenGT [28 as well as the higher-order graph transformers in Kim et al. 277. However, Graphormer-GD is strictly less expressive than the 3-WL (54] Further, Kim
53 and GraphGPS 42
However, state-of-the-
[27 and Kim et al. (28) align transformers with A-IGNs and, thus, obtain the theoretical expressive power of the corresponding k-WL but do not empirically evaluate their transformers for 2. addition, these higher-order transformers suffer from O(n runtime and memory complexity. For k 3.

---

the ET offers provable 3-WL expressivity with (n)²) runtime and memory complexity, several orders of magnitude more efficient than the corresponding 3-WL expressive transformer in Kim et al. 22 For an overview of graph transformers, see Müller et al. [40]
Finally, systematic generalization has recently been investigated both empirically and theoretically 15 26 43]. In particular, Dziri et al. [15] demonstrate that compositional generalization is lacking state-of-the-arl transformers such GPT-4.
3 Edge Transformers
The ET was originally designed to improve the systematic generalization abilities of machine learning models. To borrow the example from Bergen et [ a model that is presented with relations such as MOTHER(Z; (z, y), indicating that y is the mother of z, could generalize to more complex relation GRANDMOTHER(2, 2), indicating that is the grandmother 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 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.
if (OTHER( y) A MOTHER (y. 2)
In general, the ET operates on graph G with nodes (GG) and consecutively updates 3D tensor state X E Raxxxd where d is the embedding dimension and Xi or X(u) denotes the representation of the node pair u = (i.j) E (G)); see Figure for visualization of this construction. Concretely, the t-th ET layer computes
FFN (t-1) + TriAttenntonnLN ( ))),
for each node pair (i,j), where FFN is feed-forward neural network, LN denotes layer normalization 2 and TriAttention is defined as
n
FriAttention(X.jj Σ
Oal Vaj.
(I)
which computes tensor product between three-dimensional atlention tensor a and three- dimensional value tensor V, by multiplying and summing over the second dimension. Here,
O:ly == softmax LE|D]
X (X,, E
(2)
the attention score between the features of tuples (i, and (l,j), and
Vay = X; @ WV,
we call value fusion of the tuples (i, and (1,j) with @ denoting element-wise multiplication. Moreover, Wº, W WV, V E Rdxd are learnable projection matrices; see Figure 2 for an overview of the tensor operations in triangular attention and see Algorithm [ 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 the ET owes its expressive power 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), E) with n nodes and without self-loope where V(G) the set of nodes, E(G) the set of edges, and V(G) + assigns an initial color to each node. We construct feature matrix F E Raxp that is consistent with E, i.e. for nodes and j in VGG), F: F, if and only if, {(i) E()). Note that, for finite subset of N, we can always construct F, e.g., using one-hot encoding of the initial colors. Additionally, we consider an edge feature tensor E Raxxxy, where E, denotes the edge feature the edge (i. E F.G)). If no edge features available, randomly initialize learnable vectors 31, 32 IR9 and assign æ1 to

---

Algorithm Comparison between standard attention and triangular attention in PYTORCH1-ike pseudo-code.
function ATTENTION(X. : n x d)
function TRI TRI_ATTER IILATTENNIOONX n2 X d)
Q,K. V + (mmea((X)ccunkkk) op
Q, K, V V + linear(X) chkkkk/cckkkk (4) einsum (ild, + iljd, V V2)
A f- einsum [dd, - ij, Q, K)
A + einsum (id, ljd + ilj, Q K)
A softmax AA/vd-11 -1)
A softmax (attmaaxAA//dda -2)
O < einsum [ij.jd id, A, V)
+ einsum (ilj, iljd ijd, A, V)
return linear 0)
return linear O)
end function
end function
By if (i,j) E B(G). Further, for all i V(G), we assign 22 to Eg Lastly, if (i,j) & E(G) and i i j. we set Egj- O. We then construct 3D tensor of input tokens X E ggxxxx such that for node pair (i.j) E VGG),
X; == O
F. Fj),
where d: 2222 + Fd is neural network. Extending Bergen al. [6], our tokenization additionally considers node features, making it more appropriate for the graph learning setting.
Efficiency The triangular attention above imposes a On²") runtime and memory complexity, which significantly more efficient than other transformers with 3-WL expressive power, such as the higher- order transformers in Kim et al. 2 and Kim et al. (28) with runtime of (nn)). Nonetheless, the ET still significantly less efficient than most graph transformers, with a runtime of (n?²) 32 42
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 empir- We can easily add PEs to the above
ically from added positional//truucuuaal encodings [33 32 42 tokens with the ET. Specifically, we can encode any PEs for node (G) as an edge feature in Ea and any PEs between a node pair (i,j) E VG)² as an edge feature in Esj. Note that typically, PEs between pairs of nodes are incorporated during the attention computation of graph transformers 32 [33. 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 predic- tions for node pair- edge-level tasks is straightforward. Specifically, let L denote the number of Edge Transformer layers. Then, for node pair (i, V(G)², we simply readout X X(L) 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 node E V(G) is X 8 In addition to the diagonal readout, we also design a more sophisticated read out strategy for nodes which we describe in Appendix A.I
(L)
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//truutuuaal encodings, as parameterization. We now move on to our theoretical result, showing that the ET has the same expressive power as the 3-WL.
4 The expressivity of Edge Transformers
Here, we relate the ET to the folklore Weisfeiler-Lemaa (k-FWL) hierarchy, variant of the k-WL hierarchy for which, for k > 2, (k (65-1)- 1)-FWL 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 Weisaniler-Leman Let G = (V(G), E) be node-labeled graph. The 2-FWL colors the tuples from V(G)², similar to the way the 1-WL colors nodes [36]. In each iteration, 2
the algorithm computes a coloring C.FF: :V(G)2 + N and we write C22 (i, or C, C22P(u) (u) to denote the color of tuple = (i.5) E V(G)* at iteration t. For t = O, we assign colors to distinguish pairs
(i.j) in V (G)² based on the imitial colors ((t), {(j of their nodes and whether (i,j) E [()) or = j. For formal definition of the initial node pair colors, see Appendix D Then, for each iteration, > 0,
the coloring CZ defined as
CZ recolor recolor((C(( (1.5), M- M-1(t.,))).
where recolor injectively maps the above pair unique natural number that has not been used in previous iterations and
Mc-1ii,) - ((C2)00 C (
We show that the ET can simulate the 2-FWL, resulting in least 3-WL expressive power. Further, we show that the ET is also, at most, expressive as the 3-WL. As result, we obtain the following theorem; see Appendix
[ for a formal statement and proof details.
Theorem (Informal). The ET has exactly 3- WL expressive power.
Note that following previous works 33 result only holds for an arbitrary but fixed graph size n; see Proposition and Proposition g for the complete formal statements and proof of Theorem
37
39
our expressivity result is non-uniform in that our
In the following, we provide some intuition of how the ET can simulate the 2-FWL. Given tuple (i.j) VGG)?, we encode its color at iteration with x Further, represent the multiset
$ 10 (1.1), (2.(0) |Le V(G)).
we show that it is possible to encode the pair of colors
(C (). (2.)( via x x w
for node E V(G). Finally, triangular attention in Equation ( performs weighted sum aggregation
over the 2-tuple of colors (C2.F (i, (l,))) for each (, 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
Interestingly, proofs do not resort positionnll/strucuuaal encodings. The ET draws its 3-WL expressive power from its aggregation scheme, the triangular attention. In Section [ 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 was originally designed. Thereby, demonstrate that graph theoretic results can also be leveraged in other areas of machine learning, further highlighting the benefits of theoretically grounded models.
5
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 GR ANDMOTHER 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 defined between pairs of nodes, learning explicit representations of GR GRANDMOTHER is impossible and that learning such representations implicitly incurs a high burden on the Iearner. Conversely, the
5

---

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 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 for an introduction to first-order logic over graphs. We will now briefly introduce the most important concepts in Cai et al. [T0] 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 Cg.m the language all such statements with at most variables and quantifier depth m. A configuration is map between first-order variables and nodes in graph. Concretely, configurations let us define a statement y in first-order logic, such as three nodes forming triangle, without speaking about concrete nodes in graph G (V BG)). Instead, we can use configuration to map the three variables in yp to nodes v, w, u E (() and evaluate yp to determine whether v, w and form a triangle in G. Of particular importance to Ii.... 2g in a FO-logic statement to a k-tuple u E V(G)* such that u = (f(x1),... f(xx)) This lets now state the following result in Cai al. [10, relating FO-logic satisfiability to the FWL hierarchy. Here, C.... denotes the coloring of the k-FWL after iterations; see Appendix D for precise definition.
k-configurations f where map variables
Theorem (Theorem 5.2 [0], informally). Let G == (V(G), F(()) and H = (V ((H)) be two graphs with nodes and let k 1. Let f be k-configurraton mapping tuple VG)* and let g be k-configuration mapping tuple VHH)². Then. for every t 2 0.
ch. (t. C.. (v).
if and only if u and v satisfy the same sentences Ck+1,, whose free variables are in {21, T2, }
Together with Theorem D we obtain the above results also for the embeddings of the ET for k = 2. Corollary 3. Let G = (V (V),,
and H == (V),,, E(H)) be two graphs with n nodes and let == 2. Let f be 2-configuration mapping node pair E VG)2 and let g be 2-configuration mapping to node pair vo E VHH). Then, for every 2 0,
( = (v).
and only if u and v satisfy the same sentences C3,t whose free variables are in [21,
Systematic generalization Returning to the example in Bergen et al. 66. the above result tells us that model with 2-FWL expressive power and at least I layers is equivalently able evaluate sentences in C3,,, including
GRANDMOTHER(Z, =) By (y(OTHER((Z, v) A MOTHER(Y, ()):
i.e. the grandmother relation, and store this information encoded in some 2-tuple representation X (u),, where u (u, v) and X (1) (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. ( namely that the ET can learn an explicit representation of a novel concept, in our example the GRANDMOTHERR relation.
However, when closely examining the language C3,t, find that the above result allows for wider theoretical justification of the systematic generalization ability of the ET. Concretely, we will show that once the ET obtains representation for novel concept such the GR GRANDMOTHER relation, at some layer f, the ET can re-combine said concept to generalize to even more complex concepts. For example, consider the following relation, which we naively write
GREATGRANDMOTHERR a) = 3:3y (MOTHER(Z, B) A MOTHER(Y. 2) A MOTHER(2,a))
At first glance, it seems as though GREATGRANDMOOTE for any 2 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
E C4.2 but GREATGR
& Ca.

---

GREATGRANDMOTHER (I, By (3a(MOTHER(z, a) (oTHER(a, v))) AMOTHER a)).
is re-quantified and temporarily bound
i.., we re-quantify a to temporarily serve as the mother of x and the daughter of y. Afterwards, a is released and again refers the great grandmother of x. As result, GREATGRANDMOTIES 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
C3,2
GREATGRANDMOTHHE (z., a): = By (GRANDMOTHER v) A MOTHER a)):
where GRANDMOTHER 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 the logical reasoning abilities of the model. Further, this theoretical connection allows an interpretation of systematic generalization the ability of model with the expressive power of at least the k-FWL iteratively re-combine concepts from first principles (such as the MOTHER relation) as hierarchy of statements in CK+1. containing all FO-logic statements with counting quantifiers, at most k + variables, and quantifier depth L.
Experimental evaluation
We now investigate how well the ET performs on various graph-leaming 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 tfate-of-the-an models?
Q3 How effectively can the ET benefit from additional positional//structura encodings?
The source code for our experiments is available at hhtps:///iith/: towards- prarrdd-principledd-tta To foster research in principled graph transformers such the ET, provide accessible implementations of ET, both in PyTorch and Jax.
Datasets We evaluate the ET graph-, node-, and edge-level tasks from various domains demonstrate its versatility.
On the graph level, evaluate the ET on the molecular datasets ZINC (12K), ZINC-FULL [4] ALCHEMY (12K), and PCQM4Mv2 [21). Here, nodes represent atoms and edges bonds between atoms, and the task is always to predict one 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 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- cewterr-rooosoo-reoooooo recursively applied to execute the algorithmic steps iteratively. We will use the ET 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 /tttt::///// both node- and ldge- predictions.
framework introduced by Velickovic et al. [47]. which is
// / com/ com/google google-deepmindeee weteddeeeddddd//eeeeeee clrs As such, the CLRS requires the ET to make
Finally, we conduct empirical expressivity tests on the BREC benchmark 449 BREC contains 400 pairs of non-isomorphie graphs with up 198 nodes, ranging from basic, 1-WL distinguishable
7

---

Table Average test results and standard deviation for the molecular regression datasets, ALCHEMY (12K) and ZINC-FULL over 5 random seeds. ZINC (12K) over 10 random seeds.
ZINC (12K) ALCHEMY (12K) ZINC FULL
Model
MAE 4
MAE 4
MAE +
GIN(E) 41.41
0.163 1008
0.180 10.006
0,180 .nmss
CIN 8
0.079 a00006
0.022
Graphormer- GD 5
0,081 20.009
0.025 20.04
SignNet 30
0.084 a010
0.113
0.024
BasisNet 2
0.155 20107
0.110
PPGN++
0,071 -0000
0.109 20.001
0,020
SPE (2
0.069 20105
0.108 e0.001
ET
0.062
0.099
0.026 -000
ET -RIWP
0.059 20.094
0.098
0.024
graphs to graphs even indistinguishable by 4-WL. In addition. BREC comes with its own training and evaluation pipeline. Let f: G - E2 be the model whose expressivity we want test, where maps from a set of graphs G to _ for some d > O. Let (G, H) be a pair of non-isomorphic graphs. During training. f is trained to maximize the cosine distance between graph embeddings f(G) and f(H). During the evaluation, BREC decides whether f can distinguish G' and H by conducting a Hotelling's 's T-square test with the null hypothesis that f cannot distinguish G and H.
On the molecular regression datasets, we compare the ET to an 1-WL expressive GNN baseline such as GIN(E) (52)
Baselines
On ZINC (12K), ZINC-FULL and ALCHEMY,, we compare the ET to other theoretically-alignes models, most notably higher-order GNNs R 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 4I to study Q1. On PCQM4Mv2, we compare the ET to state-of-the-ar graph transformers study Q2. To study the impact of positionnal/sruucuuaa 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 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 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 449 First, we compare to the 8-2-LGGN 337 a higher-order GNN with strietly 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 for an overview of the used hy hypepaarameters.
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. [3)] Here, we simply adopt the hyper-parameter of ZINC (12K) from Ma et al. 332 but set the batch size to 64.
We choose the same hyper- papeee-rramtters 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 processor into the latest version of the CLRS code base. 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
8

---

Table Average test micro Fl of different algorithm classes and average test score of all algorithms in CLRS over ten random seeds; see Appendix B.3 for test scores per algorithm and Appendix B.4 for details on the standard deviation.
Triplet-
Algorithm
Deep Sets
12
GAT 12
MPNN
12
PGN 12
RT 12
ET (ours)
GMPNN 24
Sorting
68.89
21.25
27.12
28.93
50.01
60.37
82.26
Searching
50.99
38.04
43.94
60.39
65.31
58.61
63.00
DC
12.29
15.19
16.14
51.30
66.52
76.36
64.44
Grredy
77.83
75.75
89.40
76.72
85.32
91.21
81.67
DP
68.29
63.88
68.81
71.13
83,20
81.99
83,49
Graphs
42.09
55,53
63.30
64.59
65,33
81.41
86.08
Strings
2.92
1.57
2.09
1.82
32,52
49.09
54.84
Geometry
65.47
68.94
83.03
67.78
84.55
94.09
88.22
Avg. class
48.60
41.82
49.23
52.83
66.60
74.14
75.51
All algorithms
50.29
48.08
55.15
56.57
66.18
75.98
80.13
generalization.
This technique can be used within the CLRS code base through the experiment
parameters.
Finally, for BREC, we keep the default syper-parameters and follow closely the setup used by Wang and Zhang [4] 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 WQ, X1j W in Equation ② and Vaj in Equation 3
For ZINC (12K), ZINC-FULL, PCQM4Mv2, CLRS, and BREC, we follow the standard raint/vlidationntesst 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 mix of A10, L40, and A 100 NVIDIA GPUs. For each run, we used at most 8 CPU cores and 64 GB of RAM, with the exception of PCQM4Mv2 and ZINC-FULI, which were trained on 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 I random seed) are taken from Wang and Zhang 49 For reference, we also report the number of graphs distinguishable by 3-WL.
Model
Basic
Regular
Extension
CFI
All
8-2-LGNN
60
100
216
PPGN
60
50
100
23
233
Graphormer
16
12
41
10
79
ET
60 200
50 100
100
48.1 119
258.1 21.
3-WL
60
100
60
270
Results and discussion In the following, we answer questions Q1 to Q3. We highlight second, and third best results in each table.
first
We compare results on the molecular regression datasets in Table On ZINC (12K) and ALCHEMY, the ET outperforms all baselines, even without using positional//tructurall 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 [ 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 where the ET outperforms strong graph transformer baselines such as GRIT (32), GraphGPS 42 and Graphormer (53) on PCQM4Mx2. As result, positively answer Q2.

---

Table 4: Average validation MAE the PCQM4MV2 benchmark over single random seed.
Table 5: ZINC (12K) leaderboard.
7aNC (12K)
Model
MAE +
Model
Val. MAE ())
Params
SignNer (30
0.084
SUN [16
0,083 20.008
EGT
23
0.0869
89.3M
Graphormer-GD
64
0,081 =0,009
GraphGPS Small
0.0938
6.2M
CIN 8
0,079 =0,006
GraphGPS Meduu
42
0.0858
19.4M
Graph-MMPP-Mxxe 20
0.073 10.001
TokenGTowe
23
0.0962
48.6M
PPGN+ [4]
0.071 =0.001
TokenGTuu
28
0.0910
48.5M
GraphGPS 42
0.070 =0,004
Graphormer
53
0.0864
48.3M
0.069 10.004
GRIT
32
0.0859
16.6M
Graph Diffuser
0.068 +0.000
Specformer
0.066 =0.008
GPTrans-I.
0.0809
86.0M
GRIT 32
0.059
ET
0.0840
16.8M
ET
0.062 =0.004
0.0832
16.8M
ET -KWW
0.059
In 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 Triple--GMPNN Additionally, the ET performs best on 4 algorithm classes and is among the top in 7/8 algorithm classes. Interestingly, only some models are best on 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 Moreover, the ET distinguishes more graphs than 5-2-LGNN and outperforms Graphormer by large margin, again positively answer- ing Q2. Overall, the positive results of the ET on BREC indicate that the ET is well able to leverage its expressive power empirically.
7 Limitations
While proving to be a strong and versatile graph model, the ET an asymptotic runtime and memory complexity of O(r²), 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, 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 compensate for the asymptotic runtime and memory complexity of the ET. In Figure in the appendix, we find that we can use low-level GPU optimizations, available for parallelizable neural networks out-of-the-bhe to dampen the cubic runtime and memory scaling of the ET; see Appendix for runtime and memory experiments and an extended discussion.
8 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 theoretical interpretation of systematic generalization. We demonstrated empiri- cally that the Edge Transformer is 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 graph transformer that does not rely on positional/structuual encodings for strong empirical performance. Future work could further explore the potential Edge Transformer in neural algorithmic reasoning and molecular learning by improving its scalability larger graphs, particular through architecture- specific low-level GPU optimizations and model parallelism.
10

---

Acknowledgments and Disclosure of Funding
CM and LM are partially funded by a DFG Emmy Noether grant (468502433) and RWTH Junior Principal Investigator Fellowship under Germany's 's Excellence Strategy. We thank Erik Müller for crafting the figures.
References
[1] W. Azizian and M. Lelarge. Characterizing the expressive power of invariant and equivariant graph neural networks. In International Conference Learning Representations, 2021.
[2] L. J. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. Arxiv preprint, 2016.
[3] I.. Babai. Lectures on graph isomorphism. University of Toronto, Department of Computer Science. Mimeographed lecture notes, October 1979, 1979.
[4] L. Babai and L.. Kucera. Canonical labelling of graphs in linear average time. In Symposium on Foundations of Computer Science, pages 39-46, 1979.
[5] L. Babai, P. Erdōs, and S. M. Selkow. Random graph isomorphism. SIAM Journal on Computing, pages 628 635, 1980.
[6] L. Bergen, J. O'Donnell, and D. Bahdanau. Systematic generalization with edge transformers. In Advances in Neural Information Processing Systems, 2021.
[7] D. Bo, C. Shi, L. Wang, and R. Liao. Specformer: Spectral graph neural networks meet transformers, In International Conference on Learning Representations, 2023.
[8] C. Bodnar, F. Frasca, N. Otter, Y. G. Wang, P. Liò, G. Montúfar, and M. M. Bronstein. Weisfeiler and Lehman cellular: CW networks. In Advances in Neural Information Processing Systems, 2021.
[9] M. Bohde, M. Liu, A. Saxton, and S. Ji. On the markov property of neural algorithmic reasoning: Analyses and methods. In International Conference Learning Representations, 2024.
[10] Cai, M. Fürer, and N. Immerman. An optimal lower bound on the number of variables for graph identifications. Combinatorica, 12(4):389-410) 1992.
[11] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficiaee exact attention with io-awareness. Advances in Neural Information Processing Systems, 2022.
[12] C. Diao and R. Loynd. Relational attention: Generalizing transformers for graph-structured tasks. In International Conference on Learning Representationss 2023.
[13] V. P. Dwivedi, A. T. Luu, T. Laurent, Y. Bengio, and X. Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representations, 2022.
[14] V. Dwivedi, C. K. Joshi, A. T. Luu, T. Laurent, Y. Bengio, and X. Bresson. Benchmarking graph neural networks. Journal of Machine Learning Research, 24, 2023.
[15] N. Dziri, X. Lu, M. Sclar, X. L. Li, L. Jiang, B. Y. Lin, S. Welleck, P. West, C. Bhaaavvtula, R. L. Bras, J. D. Hwang. S. Sanyal, X. Ren, A. Ettinger, Z. Harchaoui, and Y. Choi. Faith and fate: Limits of transformers on compositionality. In Advances in Neural Information Processing Systems, 2023.
[16] Frasca, B. Bevilacqua, M. M. Bronstein, and H. Maron. Understanding and extending subgraph GNNs by rethinking their symmetries. ArXiv preprint, 2022.
[17] Gilmer, S. S. Schoenholz, Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, 2017.
[18] D. Glickman and E. Yahav. Diffusing graph attention. ArXiv preprint, 2023.

---

[19] M. Grohe. The logic of graph neural networks. In Symposium on Logic in Computer Science, pages 1-17. 2021.
[20] X. He, B. Hooi, T. Laurent, A. Perold, Y. LeCun, and X. Bresson. A generalization of vit/mlp- mixer to graphs. In International Conference on Machine Learning. 2023.
[21] W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong. and Leskovec. OGB-LSC: A large-scale challenge for machine learning on graphs. In NeurlPS: Datasets and Benchmarks Track, 2021.
[22] Y. Huang, W. Lu, J. Robinson, Y. Yang, M. Zhang, S. Jegelka, and P. Li. On the stability of expressive positional encodings for graph neural networks. Arxiv preprint, 2023.
[23] M. S. Hussain, M. J. Zaki, and D. Subramanian. Global self-attention as replacement for graph convolution. In A. Zhang and H. Rangwala, editors, ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 655-665, 2022.
[24] B. Ibarz, V. Kurin, G. Papamakarios, K. Nikiforou, M. Bennani, R. Csordás, A. J. Dudzik, M. Bosnjak, A. Vitvitskyi, Y. Rubanova, A. Deac, B. Bevilacqua, Y. Ganin, C. Blundell, and P. Velickovic. A generalist neural algorithmic learner. In Learning Graphs Conference, 2022.
[25] Y. Jung and Ahn. Triplet edge attention for algorithmic reasoning. Arxiv preprint. 2023.
[26] D. Keysers, N. Schärli, N. Scales, H. Buisman, D. Furrer, S. Kashubin, N. Momchev, D. Sinopal- nikov, L. Stafiniak, T. Tihon, D. Tsarkov, X. Wang, M. van Zee, and O. Bousquet. Measuring compositional generalization: A comprehensive method on realistic data. In International Conference Learning Representations, 2020.
[27] Kim, S. Oh, and S. Hong. Transformers generalize deepsets and can be extended to graphs hypergraphs. In Advances in Neural Information Processing Systems, 2021.
[28] J. Kim, T. D. Nguyen, S. Min, S. Cho, M. Lee, H. Lee, and S. Hong. Pure transformers are powerful graph learners. ArXiv preprint, 2022.
[29] T. N. Kipf and M. Welling, Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
[30] D. Lim, Robinson, L. Zhao, T. E. Smidt, S. Sra, H. Maron, and S. Jegelka, Sign and basis invariant networks for spectral graph representation learning. ArXiv preprint, 2022.
[31] Y. Lipman, O. Puny, and H. Ben-Hamu. Global attention improves graph networks generalization. ArXiv preprint, 2020.
[32] L. Ma. C. Lin, D. Lim, A. Romero-Soriano, K. Dokania, M. Coates, P. H.S. Torr, and S.-N. Lim. Graph Inductive Biases in Transformers without Message Passing. In International Conference on Machine Learning, 2023.
[33] H. Maron, H. Ben-Hamu, H. Serviansky, and Y. Lipman. Provably powerful graph networks. In Advances in Neural Information Processing Systems, 2019.
[34] H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019.
[35] H. Maron, E. Fetaya, N. Segol, and Y. Lipman. On the universality of invariant networks. In International Conference on Machine Learning, pages 4363-4371, 2019.
[36] C. Morris, M. Ritzert, M. Fey, W. Hamilton, J. F. Lenssen, G. Rattan, and M. Grohe Weisfeiler and Leman go neural: Higher-order graph neural networks. In AAAI Conference Artificial Intelligence, 2019.
[37] C. Morris, G. Rattan and P. Mutzel. Weisfeiler and Leman go sparse: Towards higher-order graph embeddings. In Advances in Neural Information Processing Systems, 2020.
[38] C. Morris, I.., H. Maron, B. Rieck, N. M. Kriege, M. Grohe, M. Fey, and K. Borgwardt. Weisfeiler and Leman go machine learning: The story so far. ArXiv preprint, 2021.
12

---

[39] C. Morris, G. Rattan, S. Kiefer, and S. Ravanbakhsh. SpeqNets: Sparsity- aware permutation equivariant graph networks. In International Conference on Machine Learning, pages 16017- 16042, 2022.
[40] L. Müller, M. Galkin, C. Morris, and L. Rampásek. Attending to graph transformers. Transactions on Machine Learning Research, 2024.
[41] O. Puny, D. Lim, B. T. Kiani, H. Maron, and Y. Lipman. Equivariant polynomials for graph neural networks. In International Conference Machine Learning. 2023.
[42] L. Rampášek, M. Galkin, V. P. Dwivedi, A. T. Luu, G. Wolf, and Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 2022.
[43] Y. Ren, S. Lavoie, M. Galkin, D. J. Sutherland, and A. Courville, Improving compositional gen- eralization using iterated learning and simplicial embeddings. In Advances Neural Information Processing Systems, 2023.
[44] F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 2009.
[45] Tillet, H. Kung, and D. D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019.
[46] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
[47] P. Velickovic, A. P. Badia, D. Budden, R. Pascanu, A. Banino, M. Dashevskiy, R. Hadsell, and C. Blundell. The CLRS algorithmic reasoning benchmark. In International Conference on Machine Learning, 2022.
[48] O. Vinyals, S. Bengio, and M. Kudlur. Order matters: Sequence to sequence for sets. In International Conference on Learning Representations, 2016.
[49] Y. Wang and M. Zhang. Towards better evaluation of GNN expressiveness with BREC dataset. Arxiv preprint, 2023.
[50] B. Weisfeiler and A. Leman. The reduction of graph canonical form and the algebra which appears therein.
Technicheskaya Informatsia, 2(9):12-16, 1968.
[51] K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
[52] K. Xu, L. Wang. M. Yu, Y. Feng, Y. Song. Z. Wang. and D. Yu. Cross-lingual knowledge graph alignment via graph matching neural network. In Annual Meeting of the Association for Computational Linguistics, 2019.
[53] C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? In Advances in Neural Information Processing System, 2021.
[54] B. Zhang, S. Luo, L. Wang, and D. He. Rethinking the expressive power of gnns via graph biconnectivity. In International Conference on Learning Representations, 2023.
13

---

the second Table 6: Hyperparameters of the Edge Transformer across all datasets.
Hyperparameter
ZINC(I2K)
ALCHEMY
ZINC-FULL
CLRS
BREC
PCQM4MV2
Learning rate
0.001
0.001
0.001
0.00025
0.0001
0.0002
Grad. clip norm
1.0
1.0
1.0
1.0
5.0
Batch size
64
256
256
Optimizer
A
Adam
AdamW
Adam
Adam
AdamW
Num. layers
10
10
10
Hidden dim.
64
64
64
192
32
384
Num. heads
8
8
12
16
Activation
GELU
GELU
GELU
RELU
GELU
Pooling
SUM
SUM
SUM
SUM
RRWP dim.
32
32
32
128
Weight decay
le-5
le-5
le-5
0.0001
0.1
Dropout
0.0
0.0
0.0
0.0
0.0
0.1
Atlention dropout
0.2
0.2
0.2
0.0
0.0
0.1
Steps
10K
2M
# Warm-up steps
0
60K
# Epochs
2K
2K
IK
20
# Warm-up epochs
50
50
# RRWP steps
21
21
21
22
A Implementation details
Here, present details about implementing the ET in practice.
A.1 Node-level readout
In what follows, we propose pooling method from node pairs to nodes, which allows also make predictions for node- and graph-level tasks. For each node i E V(G), we compute
ReadOut(i)
Σ Pi JE(N]
x
(1)
+p (x!tt)
where Pl, P2 are neural networks and ( is the node pair tensor after L ET layers. We apply Pr node pairs where node is at the first position and f node pairs where node 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 sum and mean [51] set2seq [48) on the resulting node representations. We use this readout method our molecular regression experiments in Section
B Experimental details
Table 6 gives overview of selected
for all experiments.
See Appendix B.2 through Appendix B.4 detailed results on the CLRS benchmark. Note that in the case of CLRS, we evaluate in the single-task setting where we train a new set parameters for each concrete algorithm, initially proposed in CLRS, to be able compare against graph transformers fairly. We leave the multi-task learning proposed in Iharz et al. (2 for future work.
B.1 Data source and license
ZINC (12K), ALCHEMY (12K) and ZINC-FULL are available at https::/pyg org under an MIT license. PCQM4Mv2 available at sattts:///ghh stanford. Stst://dos/sssc/ppppeee com/
under a CC BY 4.0 license. The CLRS benchmark is available at https: thttbbbb............ thub. google-deepmind/clll clrs under an Apache 2.0 license. The BREC benchmark is available httpp: thu.. com wtwp/GraapPKK/BB... under an MIT license.
14

---

Table 7: Average test scores for the different algorithm classes and average test scores of all algorithms in CLRS with the OOD validation technique over 10 seeds; see Appendix B.3 for test scores per algorithm and Appendix for details on the standard deviation. Baseline results for Triplet-GMPNN and TEAM are taken from Jung and Ahn [25] Results in %.
Algorithm
TriPlct-GMPNN
TEAM
ET (ours)
Sorting
72.08
68.75
88.35
Searching
61.89
63.00
80.00
DC
65.70
69.79
74.70
Greedy
91.21
91.80
88.29
DP
90.08
83.61
84.69
Graphs
77.89
81.86
89.89
Strings
75.33
81.25
51.22
Gcometry
88.02
94.03
89.68
Avg. algorithm class
77.48
79.23
80.91
All algorithms
78.00
79.82
85.01
26.24
Processor
25
Triplet-GMPNN
20
15
11.71
10.26
Micro F1
6.09
6.62
6.84
5
3.28
1.2
1.46
0
-3.52
-3.63
-6.07
-10
-10.66
Sorting Searching
Greedy
DP Algorithm class
Graphs Strings Geometry
Figure 3: Difference micro with and without the OOD validation technique in Jung and Ahn 25 for TriPlet-GMPNN 24 and ET, respectively.
B.2 Experimental results OOD validation in CLRS
Table 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 the TEAM (25] baselines. addition, in Figure we present a comparison the improvements resulting from OOD validation technique, comparing Triplet-GMPNN and the ET. Finally, in Table & compare different modifications to the CLRS training setup that are agnostic to the choice processor.
B.3 CLRS test scores
present detailed results for the algorithms in CLRS. See Table for divide and conquer algorithms, Table 12 for dynamic programming algorithms, Table 13 for geometry algorithms. Table [5) 15) algorithms, Table 10 for search algorithms, Table for sorting algorithms, and Table T6 for string algorithms.
greedy
15

---

Table CLRSS30 Processor-agnostic modifications.
Processor
Markov 9
OOD Validation
25
Avg. algorithm class
All algorithms
Triplet-GMPNN
79.75
82.89
Triplet-GMPNN
X
77.65
78,00
TEAM
79.23
79.82
ET
80.91
85.02
Table 9: Detailed test scores for the ET sorting algorithms.
Algorithm
Fl-score(%)
Std. dev.(%)
FO-scor(())OOOD)
Std. dev.(%) (OOD)
Bubble Sort
93.60
3.87
87.44
13.48
Heapsort
64.36
22,41
80.96
12.97
Insertion Sort
85,71
20.68
91.74
6,83
Quicksort
85.37
8.70
93.25
9.10
Average
82.26
13.92
88.35
10.58
Table 10: Detailed test scores for the ET search algorithms.
Algorithm
F1 F1-score(%)
Std. d%v.
(1=score(%)(OOD)))
Std. dev. (OOD)
Binary Search
79.96
11.66
90.84
2.71
Minimum
96.88
1.74
97.94
0.87
Quickselect
12.43
11.72
52.64
22.04
Average
63.00
8.00
80.00
8.54
Table 11: Detailed test scores for the ET divide and conquer algorithms.
Algorithm
F1-score(%6)
Std. d%v
Fl-score($))OOD)
Std. d%v. (OOD)
Find Max. Subarray Kadande
64.44
2.24
74.70
2.59
Average
64.44
2.24
74.70
2.59
Table 12: Detailed test scores for the ET on dynamic programming algorithms.
Algorithm
Fl Fl-score(%%)
Std. dev.(%)
Fl-score(% %% (O)(OOD)
Std. dev. (OOD)
LCS Length 88.67
2.05
88.97
2.06
Matrix Chain Order
90.11
3.28
90.84
2.94
Optimal BST
71.70
5.46
74.26
10.84
Average
83.49
3.60
84.68
5.28
Table 13:
Detailed test scores for the ET on geometry algorithms.
Algorithm
F1-score(%)
Std. d%v.
F1-score(%)OOOD)
Std. dev.(%) (OOD)
Graham Scan
92.23
2.26
96.09
0.96
Jarvis March
89.09
8.92
95.18
1.46
Segments Intersect
83.35
7.01
77.78
1.16
Average
88.22
6.09
89.68
1.19
16

---

Table 14: Detailed test scores for the ET on graph algorithms.
Algorithm
Fl-score(%)
Std. dev.
Fl-score(%)XOOD)
Std. dev.(%) (OOD)
Articulation Points
93.06
0.62
95.47
2.35
Bellman-Ford
89.96
3.77
95.55
1.65
BFS
99.77
0.30
99,95
0.08
Bridges
91.95
10.00
98.28
2.64
DAG Shortest Paths
97.63
0.85
98,43
0.65
DFS
65.60
17.98
57.76
14.54
Dijkstra
91.90
2.99
97.32
7.32
Floyd- Warshall
61.53
5.34
83.57
1.79
MST-Kruskal
84.06
2.14
87.21
1.45
MST-Prim
93.02
2.41
93,00
1.61
SCC
65.80
8.13
74.58
5.31
Topological Sort
98.74
2.24
97.53
2.31
Average
86.08
4.73
89,92
3.02
Table 15: Detailed test scores for the ET on greedy algorithms.
Algorithm
Fl Fl-score(%)
Std. dev.(%)
Fl-score(% (OO)))
Std. dev.(%) (OOD)
Activity Selector
80.12
12.34
91.72
2.35
Task Scheduling
83.21
0.30
84.85
2.83
Average
81.67
6.34
88.28
2.59
B.4 CLRS test standard deviation
compare the standard deviation Deep Sets, GAT, MPNN, PGN, RT, and ET following the comparison in Diao and Loynd 12]. Vable|17]ompal the CLRS benchmark. We observe that the ET has the lowest overall standard deviation. The table does not contain results for TriPNetGGMPNN 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. observe that Triplet-GMPNN and the ET have comparable standard deviations except for search and string algorithms, where Triplet- GMPNN has much higher standard deviation than the ET.
the standard deviation over all algorithms in
!4 since do not have access to the test results for
C Runtime and memory
Here, we provide additional information 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 Py Torch forward pass of a single layer of the ET on a single graph (batch size 1) with n E (100, 200, 700) 100 repeats. We sample random Erdōs-Renyi graphs with edge
A100 GPU with bfloat16 precision. We measure the time for the
nodes and average the runtime probability 0.05. We use an embedding dimension 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.
Table 16: Detailed test scores for the ET on string algorithms.
Algorithm
Fl-score(%)
Std. dev.
F1-score (OOD)
Std. (%) (OOD)
KMP Matcher
10.47
10.28
8.67
8.14
Naive String Match
99.21
1.10
93.76
6.28
Average
54.84
5.69
51.21
7.21
17

---

Table 17: Standard deviation Deep Sets, GAT, MPNN, PGN, RT, and ET (over all algorithms and all seeds).
Model
Std. Dev. (%)
Deep Sets
29.3
GAT
32.3
MPNN
34.6
PGN
33.1
RT
29.6
ET
26.6
Table 18: Standard deviation per algorithm class of Triplet-GMPNN (over 10 random seeds) as reported in Ibarz et [2 and ET (over random seeds). Results in %.
Algorithm class
Triplet-GMPNN
ET
Sorting
12.16
15.57
Searching
24.34
3,51
Divide and Conquer
1.34
2.46
Greedy
2.95
6,54
Dynamic Programming
4.98
3.60
Graphs
6.21
6.79
Strings
23.49
8.60
Geometry
2.30
3.77
Average
9.72
6.35
Model
ET
Iriton)
0.150
Time elapsed [scronds]
0.100
0.050
0.000
100
200
300
400
500
600
700
nodes
Figure 4: Runtime of the forward pass of single ET layer in PoTchc in seconds for graphs with up 700 nodes. compare runtime with and without torch. compile (automatic compilation into Triton (55) enabled. Without compilation, the ET goes out of memory after 600 nodes.
18

---

Table 19: Runtime of single run of the ET in CLRS on a single A GPU.
Algorithm
Time in hh:mm:ss
Activity Selector
00:09:38
Articulation Points
01:19:39
Bellman Ford
00:07:55
BFS
00:07:03
Binary Search
00:05:53
Bridges
01:20:44
Bubble Sort
01:05:34
DAG Shortest Paths
00:29:15
DFS
00:27:47
Dijkstra
00:09:37
Find Maximum Subarray Kadane
00:15:25
Floyd Warshall
00:12:56
Graham Scan
00:15:55
Heapsort
00:57:14
Insertion Sort
00:10:39
Jarvis March
01:34:40
Kmp Matcher
00:57:56
LCS Length
00:08:12
Matrix Chain Order
00:15:31
Minimum
00:21:25
MST Kruskal
01:15:54
MST Prim
00:09:34
Naive String Matcher
00:51:05
Optimal BST
00:12:57
Quickselect
02:25:03
Quicksort
00:59:24
Segments Intersect
00:03:38
Strongly Connected Components
00:56:58
Task Scheduling
00:08:50
Topological Sort
00:27:40
Table 20: Runtime of a single run on the molecular regression datasets, as well as BREC, on L40
GPUs in days:hourrmmiiuesssseooo...
ZinC (12K)
ALCHEMY (12K)
CINC-FUlL
PCQM4Mv2
BREC 00:00:08:37
ET
00:06:04:52
00:02:47:51
00:23:11:05
03:10:35:11
ET+RRWP
00:06: 19:52
00:02:51:22
01:01:10:55
03:10:22:06
Num. GPUs
1
1
4
4
19

---

Hardware optimizations Efficient compilation of neural networks is already available via program- ming languages such Triton 45]. We use torch. compile our molecular regression experiments. In addition, we want to highlight FlashAttention [II], available for the standard transformer, as an example of architecture-specfii 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 single run on a single 1.40 GPU of ZiNc (12K), ALCHEEYY (12K), and BREC. For ZINC-FULL and PCQM4Mv2, we present the runtime of single run on 4 L40 GPUs; see Table 20
On CLRS, the experiments in our work are run on mix of A10 and A100 GPUs. To enable fair comparison, we rerun each algorithm in CLRS in a single run on single A100 GPU and report the corresponding runtime in Table [9. 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.
D Extended preliminaries
Here, we define our notation. Let N == 21, For 2 1, let [m] == {1, [1....,
.., nu C N. We use {... It to denote multisets, i.e., the generalization of sets allowing for multiple instances for each of its elements.
Graphs A Mabeeel graph V(G), edges E(G) S Hu, ef C VVG) a and (node-)labe) function &: V(G) + N. Then ((v) is label of v, for e in (G). not otherwise stated, we set n == [V(G),, and the graph is of order 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 pair g = (G, where G (V,
triple (V(G), E(G), E(G)) E) with finite sets of vertices or nodes
and F in Rx for p 0 is node feature matrix. Here, we identify V V(G) with in]. then (v) in IR XP the feature or attribute of the node v E V(G). Given labeled graph (V) EG)), e). node feature matrix F is consistent with l if ((t) [(w) for only if F (v) = F(w).
V(G) if and
Neighborhood and Isomorphism The neighborhood of vertex v in V(G) is denoted by N(v) = {u V(G) uv, F(G)) and the degree of vertex is [N(v). Two graphs G and H are isomorphic and we write G = H if there exists a bijection 9: V(G) - ((H) preserving the adjacency relation, i.e., (u, v) is in E(G) if and only if ((()),((()) is in E(H). Then & is an isomorphism between G and H. In the case of labeled graphs, we additionally require that I(v) - (()(v)) ((p()) 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 TO. We further define the atomic type atp: V(G) *
for k > O, such that atp(v) atp(w) for e and w in V(G if and only if the mapping yp: V(G)* V(G)* where u; - wi induces partial isomorphism, i.e., wn; wy and (U,, 0j) E E(G) = ((((x)) ppr))) E F((G).
we have 0; Uj è
Matrices Let M guxp and N Raxq be two matrices then [M NJ E gux(pt)) denotes column-wise matrix concatenation. We also write Rd for Rixd Further, let M Rxxx and N E R be two matrices then
M N
E p(++))x
denotes row-wise matrix concatenation.
For matrix X E Randd we denote with X, the ith row vector. In the case where the rows of X correspond to nodes in graph G, we use X, to denote the row vector corresponding to the node " E V()).
The Weisteiler-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,
20

---

originally proposed by Weisfeiler and Leman 150
Intuitively, the algorithm determines if two
graphs are non-isomorphic by iteratively coloring or labeling vertices. Formally, let G = (V. be a labeled graph, in each iteration, i >> 0, the 1-WL computes vertex coloring Ct V(G) - N. depending on the coloring of the neighbors. That is, in iteration > O, we set
(e) = recolor (C)),()), (C- u E (())))))))
for all vertices in V(G), where recolor injectively maps the above pair to unique natural number, which has not been used in previous iterations. In iteration 0, the coloring Co == L. To test if two graphs G and H are non-isomorper -somoooo,,,,,, we run the above algorithm in "parallel" on both graphs. If the two graphs have a different number of vertices colored e N at some iteration, the 1-WL distinguishes the graphs as non-isomorphic. It is easy to see that 1 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-isomopppic graphs, several researchers, e.g., Babai [3, Cai et al. [O]. devised more powerful generalization of the former, today known as the k-dimensional Weisfeiler- Leman algorithm (k- (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 single node. More precisely, given graph G, the k-WL colors the tuples from V(G)* for k 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 2. In each iteration, & 2 0, the algorithm, similarly to the 1-WL, computes a coloring Cf:- V(G)* + N. In the first iteration, t = 0, the tuples and w in V get the same color if they have the same atomic type, i.e., atp& (v) = atp; (u)) Then, for each iteration, > 0, C is defined by
Ck (v) (ec)
(5)
with (v) the multiset
Mc(w) = (o w)) w ((V))) (ox,,)
| w E V()))),
(6)
and where
65 (v) (.i,..., Uj-1, w, Uj+
ok)
That is, (v, w) replaces the j-th component of the tuple v with the node w. Hence, two tuples are adjacent 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 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 WL algorithm until convergence, i.e., until for in N
(0) = C* (w) = C4+1 C(+)(w),
for all v and w in V(G)* holds.
Similarly to the 1-WL, to test whether two graphs G and H are non-isomorphic, run the k-WL in "parallel" on both graphs. Then, if the two graphs have a different number of nodes colored c, for 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 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 to
K.F Mf(x) == {(C (01 (v, w)) .... Ci- (#x(v, w))) w V(G)
resulting in the coloring C...F. V(G)) - N, and is strictly more powerful than the k-WL. Specifical for & > 2, the k-WW is exactly as powerful as the (k FWLL 19].
'Strictly speaking. the 1-WL and color retinement 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 given graph: see [9 for details. For brevity, consider both algorithms be equivalent.
21

---

Computing k-WL's initial colors Let G = (V(G), E(G), E) be labeled graph, k 2 2, and let v := (v1...., "*) E V(G)* be a k-tuple Then, we can present the atomic type atp(v) by a k x matrix K over (1, That is, the entry Kij is if (o% vj) E EGG), if vs vj. and 3 otherwise. Further, we ensure consistency with l, meaning that for two k-tuples v (U vn) E (G) and w = (w1. (w...... wk) V(G)*, then
(v) C. (w),
if and only if, atp(v) = atp(w) w) and E(vy) = f(wx), for all E [k]. Note that we compute the initial colors for both k-WL and the k-FWL this way.
D.I Relationship between first-order logic and Weisfeiler-Leman
We begin with a short review of Cai et al. [oI. We consider our usual node-labeled graph G (V(G), E(G), E) with nodes. However, replace l with countable set of color rela- tions Cr,... On, where for a node E V()),
C,v) (())=
Note that Cai et al. [10] consider the more general case where nodes can be assigned to multiple colors simultaneously. However, for work, we assume that a node is assigned 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 triangle formed by two nodes with color 1:
3 3ry (x(1, x2) A F(11,as) A E(22, 23) (
Here, Ii, 22, and I3 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, function ʃ {-I,-.., T J
- V (G) that assigns node in V(G) to each one of the variables æi...., Let be a first-order formula with free variables among 21...., Tk. Then, we write
F
if (p is true when the variable n, is interpreted the node fxx4), for i 1,...,
Cai et al. [10] define the language Cs.m of all first-order formulas with counting quantifiers, at most k variables, and quantifier depth bounded by m, and the language C. = Um20 Ck. For example,
the sentence Vu (F(x,)) v)) in C2 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 as (G, f) =k.m (H,g) g) if and only if
g
for all formulas in Cr.m whose free variables are among 21, .k.-
We formulate main result of Cai et al. [00. Let G and H be two graphs, let k > 1 and m 2 0 be non-negative integers, and let f and g be k-contigurations for G and H respectively. If (f(z)), ...,, ((x)) E (GG and (g(
((H), then
Ok.F (u) O&. (v) = (G. kmm (H,9) g)
E Proofs
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 special form that can be computed by the Edge Transformer.
Formally, let S N be finite subset. First, we show that multisets over S can be injectively mapped to value in the closed interval (0, 1), a variant of Lemma VIII. in Grohe [9]. Here, we outline streamlined version of its proof, highlighting the key intuition behind representing multisets as m-ary numbers. Let M C S be multiset with multiplicities a aI...., 4g and distinct & values. We define
the order of the multiset as > -iz1 04. We can write such multiset as sequence z ..... z where
22

---

1 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 ( = و(ع We call such sequence an M -esquencc of length I. We now prove slight variation of result of Grohe [9]
Lemma 4. For finite N. let M _ S be multiset of order 1 and let 24 E S denote the ith number in a fixed but arbitrary ordering of S. Given a mapping g: - (0, 1) where
g(z;) = m
and an M-sequence of length I given by x
1
with positions j(2), ,..,,, (00) in S, the sum
Eg(xb))
Σ SER
m G
JEI
is unique for every unique M.
Proof. By assumption, let C S denote multiset of order m - 1. Further, let x: be an M M-sequence with ;(1) ...., (00) in S. Given our fixed ordering of the numbers S we can equivalently write M = II).. (an: kn)), where denotes the multiplicity of ith number in M with position from our ordering over S. Note that for a number m there exists a corresponding m-ary number written as
,..., ,..., (9) E M
0.0
Then the sum,
((((1))
((5)
SEN
JE[
Σ
a.m E 10,
iE
and in m-ary representation
0.a ... an.
Note that a; = if and only if there exists no j such that (06 Since the order of M m 1, it holds that a; implying the result.
Hence, it follows that the above sum is unique for each unique multiset M,
the finiteness of S to map each number therein to a fixed digit of the numbers in (0, 1). The finite m ensures that each digit, encode multisets over S as numbers in (0, 1), we review some fundamental operations about the m-ary numbers defined above. We will refer decimal numbers -1 as corresponding an m-ary number
Recall that S C N and that fixed an arbitrary ordering over S. Intuitively,
have sufficient "bandwidth"
encode each a;- Now that have seen how to
0.0
where the äth digit after the decimal point is 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 vn corresponds to 0.0...
for digit positions + j and
m corresponds to 0.0... 2
otherwise. We used counting in the previous result's proof to represent a multiset's multiplicities. Next, multiplication between decimal numbers implements shifting in m-ary notation, i.c.,
m m corresponds to 0.0
23

---

Shifting further applies to general decimal numbers in (0, 1). Let (0, 1) correspond to an m-ary number with / digits,
0.a1 ... dq.
Then,
m x corresponds 0.0... a1 a1...
Before we continue, we show 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 >> If
min #EA
max y
yEB
then for any 02,32 11 1/2 E B.
3 + yı 12
21
22 and y1 92-
Proof. The statement follows from the fact that
min x > max
yEB
then numbers in A and numbers in B do not overlap in terms of their digit range. Specifically, there exists some / >> 0 such that we can write
0.71 ... 21 y = 0. 0. y1
yk:
for some k > and all y, B. As result,
x + y 0.21 ... T11 yk-
Hence, x + y unique for every unique pair (z, y). This completes the proof.
We begin by showing the following proposition, showing that the tokenization in Equation is sufficient encode the initial node colors under 2-FWL.
Proposition 6. Les G (V(G), E(G), E) be node-labeled graph with nodes. Then, there exists paraneterization of Equation with d = 1 such that for each 2-tuples u, v E V(G)²,
CO.F (u) C2.((w)
cf (p) + X(u)=
Proof. The statement directly follows from the fact that the initial color of a tuple u := (i,jj) depends on the atomic type and the node labeling. In Equation (4, we encode the atomic type with Eig and the node labels with
E F. FJ
The concatenation of both node labels and atomic type is clearly injective. Finally, since there are at most n² distinct initial colors the 2-FWL, said colors can well represented within IR, hence there exists an injective é in Equation 44 with d = 1. This completes the proof.
We now show Theorem 1 Specifically, we show the following two propositions from which Theorem follows.
Proposition 7. Let G (V(G), E(G), E) be a node-labeled graph with n nodes and F E guxp be node feature matrix consistent with f. Then for all t 2 O, there exists a parametrization of the ET such that
CZ- (22(ww = x (v) x (w),
for all pairs of 2-tuples v and E 3
24

---

Proof. We begin by stating that our domain is compact since the ET merely operates on at most possible node features in F and binary edge features in E, and at each iteration there exist at most n² distinct 2-FWL colors. We prove our statement by induction over iteration t. For the base case, we can simply invoke Proposition since our input tokens are constructed according to Equation Nonetheless, we show possible initialization of the tokenization that is consistent with Equation that we will use in the induction step.
From Proposition 6 we know that the color representation of a tuple can be represented in R. We denote the color representation of tuple We choose in Equation (4 such that for each u (i,j)
j) at iteration i ( and T interchangeably.
n²
T (0)
T (0)
R2
where we store the tuple features, one with exponent 1 and once with exponent n² and where T. E I
and (1) E IR. We choose color representations T follows. First, we define an injective function fe V(G) [n²] that maps each 2-tuple to number in [n²] unique for its 2-FWL color C2 at iteration f. Note that fa can be injective because there most be [n²] unique numbers under the 2-FWL. We will use fa to map each tuple color under the 2-FWL to a unique n-ary number. We then choose & Equation (4) such that for each (i,j) VG)²,
| n fo co:
for all co > 0, by the universal function approximation theorem, which we can invoke since our x²²
domain is compact. We will use (I:!0) in the induction step; see below.
For the induction, we assume that
CZ-i(v) C2. (w) = (1) (ww
and that
( n (--((()) <
for 6 > and (1,5) E V(G)². We then want to show that there exists a parameterization of the
1-th layer such that
C2 (v) C. (w) - (v) (w)
and that
17 n (o || E2,
for all > 0 and (i.j) E V(G)2. Clearly. if this holds for all t, then the proof statement follows. Thereto, we show that the ET updates the tuple representation of tuple (j. m) as
T T FFN T (t-1)
B
ET
(t- (T(h n²
(8)
for arbitrary but fixed B. We first show that then, Equation 7 holds. Afterwards we show that the ET can indeed compute Equation (8). To show the former, note that for two 2-tuples (j, 1) and (1, m).
n² Se-1 (n (n- =n 1-0
is unique for the pair of colors
(C. (F((())))
where n n-" is a constant normalization term we will later introduce with g. Note further, that we have
|T( (T()- n (n" + (iD)
f--1(mm)) (-
8
for all 6R-1 > 0. Further, n -(h-1 fr- (1,m)) still an m--ar number with m = n. As a result, we can set B = nn n and invoke Lemma 4 to obtain that
a
(fo=0 (
+ G.) (.
n
25

---

is unique for the multiset of colors
(C
and we have that
n²
(1-1),
(
n - < x-b
mm
1=1
1=1
for all 7-1 > 0. Finally, we define
A
n
fe f1-mm)m) (i.m) v(
B = 3
(-----0 (i.m) E V(G)).
1=1
Further, because multiply with we have that
min x y ÆĒA
vEB
and as result, by Lemma
5
B
n 1(.mm)
+ Σ
n
1=1
unique for the pair
0f (( ((C2-(((,m)) (),,
and consequently for color C ((j, m)) at iteration t. Further, we have that
1) (r(t-1)
+ T
3 n
lm "I!
n² -n -t-:
n -----= (1-1)))))))))))))))11=)))) < Tt-1
for all 7i-1 0. Finally, since our domain is compact, we can invoke universal function approximation with FFN in Equation (8) to obtain
Jm n -
< Es.
for all e, 0. Further, because n f+(.m) is unique for each unique color C C2 ((j, m)). Equation follows.
remains show that the ET indeed compute Equation B). To this end, will require single transformer head in each layer. Specifically, we want this head to compute
(Ti(L n²
hn (X (t- 1) jim
1t)-1)
(9)
1=1
Now, recall the definition of the Edge Transformer head tuple (j, m)
hn (x((11-1))
aiim vit-!)
where
1 x((-1) w (x/1
O'lmm softmax
LE(n]
Vd
with
((t-1)
x((-1)
x
(t-
W
x (t-1)
V27
W V27 W V/
jim
W
Vi
km
26

---

and by the induction hypothesis above,
T n²
X ((-1)
1)
[ (1-1) 1)
X ((-1)
T(t-1)
T n1 lm
where we expanded sub-matrices. Specifically, W V1
W
W W
RRxd,
We then set
Wº W
W V [8I 0
W V 10 0]
W. V- 10 1
W 1/2 00 0j.
Here, WQ and W set zero obtain uniform attention scores. Note that then for j, a3lm due to normalization over I, and we end up with Equation as
h1 (1-1)) jm
(1-1)
jim
where
v-t-1) jim
11)
3I .11 +
(T (T((-1)) n²
0
((le-1
0 +
n2 (T
im
.
(((-1) (T(h-11) x²
3
We now conclude our proof as follows. Recall that the Edge Transformer layer computes the final representation X(() as
X
FFN
X (1-1) +
hh1 (x
wo
FFN
T jm
(I n27
+
8
( 1m
0/ o/woo)
n²
Σ T I 0)
FFN
(T((-1)) jm
+
2
jm
(r) n²? :
FFN
T jm
1)
+
1 Ex= Tul Tilt-1) 1)
T(t--)
FFN
(
(T(C-1)) 13)
B S
for some FFN. Note that the above derivation only modifies the terms inside the parentheses and is thus independent the choice of FFN. We have thus shown that the ET can compute Equation (8) To complete the induction, let f: RR2 + R² be such that
(T!
( n²
T jm
(1)
n²
jm
Since our domain is compact, f continuous, and hence we can choose FFN to approximate arbitrarily close. This completes the proof.
27

---

Next, we show the other direction of Theorem 1 under mild and reasonable assumptions. First, we say that recoloring function, that maps structures over positive integers into positive integers, is (effectively) invertilee if its inverse is computable. All coloring functions used in practice (e.g., hash- based 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 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), ((G)),D E(G), and for all 2 0:
Cº (v) C (w) = X (2) (v) x (w),
for all pairs of 2-tuples v and w in V(G)².
Proof. We first claim that there is computable function Z :N x N + R2, where N {0} U N. such that X (²(()) (l, C2FF())) (v)) for all V(G)², independent of the graph G and its order. The proof the claim by induction on For = 0, by definition, C (v) identifies the atomic type (t)) which defines X (0) (v) (since the atomic type tells if u is an edge G, and the labels of the vertices in v).
For é >> 0 and (i. the function (tt, C.F((())) of recolor to obtain the pair
(v)) proceeds as follows. First, it uses the invertibility
(2 (CL | V(G)))
Then, by inductive hypothesis using the function Z(t 1, obtains the pair
x x(--1) (ix (--1)) |Le v(G)B)
Finally, it computes
X (i))
FFN
X (1-1) (i.j) + TriAttention (1-1) (i.))
under the assumption that the layer normalization is a proper function. The statement of the proposition then follows directly from the claim since
(x) = Z(t, C (v)) = Z(t, C 2. (w)) x (w).
Note that unlike the result in Proposition] 7 the above result is uniform, in that the conerete choice of recolor and the function Z does not depend on the graph size n. Finally, Theorem follows from Proposition 7 and Proposition g
28

---

NeurIPS Paper Checklist
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 and that the ET surpasses theoretically aligned graph models, which we demonstrate in Table | and Table 3 Further, we claim that the ET is competitive with state-of-the-art models variety of tasks, which we demonstrate in Table 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 other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We provide discussion of the limitations of the ET, specifically its high runtime and memory complexity, in Section
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 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-specificatiin, asymptotic approximations only holding locally). The authors should reflect 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 few datasets with 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, 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.
applicable, 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 '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.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
29

---

Answer: [Yes]
Justification: We provide all proofs in Appendix E where 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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We provide the model definition in Section Rl well as Py Porchh-like pseu- docode for the triplet attention in Algorithm [ In addition, we provide all experimental details in Section as well the chosen hyper-parameters and optimizers Table
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer 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 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 replicate the results, access to hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with 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 (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
30

---

Question: Does the paper provide open access to the data and code, with sufficient instructions 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.
Please see the NeurIPS code and data submission guidelines (https://nnis..o cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7 //mips.ccppullcc/guue
ssionsolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide information about the data splits and how they were selected in Section 6 Further, for every dataset and benchmark, we detail the pprrmpprraeeters and the optimizer used in Table
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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We present the standard deviation over multiple random seeds for all experi- mental results; see Section 6 and Appendix B for the CLRS benchmark.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: We detail compute resources in Section 6 and runtimes needed to reproduce our experiments in Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http::ttp::// http://wwrr...ccoolal.llccctttta
Answer: [Yes]
Justification: Our research conforms to the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consider- ation due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
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, do foresee any immediate positive or negative societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

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), fairesss 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 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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 (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We release neither data models as part of this work. Further, experiments are conducted on comparatively small, curated, task-specific datasets used for benchmarking graph learning models. Hence, work does not pose immediate risks for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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.I
Guidelines:
The answer NA means that the paper does
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 used and, 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.
33

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, wwwpssppppesseeeee wwwoeee 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.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification: This work does not release new assets.
Guidelines:
The answer NA means that the paper does release assets. Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
34

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "8ipCI30tB1", "reviewer": "Reviewer_YYTZ", "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.", "questions": "N/A", "limitations": "Please refer to the Weaknesses section for details.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written. All technical steps are easy to follow.\n2. The proposed triangular attention is interesting and novel.\n3. ET achieves impressive results for both molecular regression and algorithmic reasoning tasks.", "weaknesses": "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. \n2. 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.\n3. Some relevant GT models haven't been discussed or compared (e.g., [1-3]).\n\n[1] Kreuzer et al., \"Rethinking Graph Transformers with Spectral Attention\", NeurIPS'21. \\\n[2] Zhu et al., \"On Structural Expressive Power of Graph Transformers\", KDD'23. \\\n[3] Deng et al., \"Polynormer: Polynomial-Expressive Graph Transformer in Linear Time\", ICLR'24.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer says\n\n> Previous work [1] has demonstrated that graph transformers (GTs) with proper positional encodings (PEs) can be more powerful than any WL test\n> \n\nThe 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.\n\nNext, the reviewer says\n\n> While the authors mention how to apply ET for node-level tasks, they provide no corresponding results\n> \n\nWe refer the reviewer to the experiments on the CLRS benchmark. There, many tasks involve node- and edge-level predictions. The reviewer also mentions that\n\n> Given the cubic complexity of ET, it is likely not scalable to most graphs for node-level tasks.\n> \n\nWe 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.\n\nFinally, the reviewer mentions that\n\n> Some relevant GT models haven't been discussed or compared (e.g., [1-3])\n> \n\nWe 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].\n\nWe want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.\n\n### References\n\n[1] Kreuzer et al., **Rethinking Graph Transformers with Spectral Attention,** 2021 https://arxiv.org/abs/2106.03893\n\n[2] Petrache and Trivedi, **Approximation-Generalization Trade-offs under (Approximate) Group Equivariance,** 2023 https://arxiv.org/abs/2305.17592", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Oiy6qoNdIA", "reviewer": "Reviewer_szHE", "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.", "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.\n- 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?\n    - 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": "See above.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "- It proves ET has 3-WL expressive power and demonstrates excellent results on BREC benchmark. Even though it has $O(n^3)$ complexity, it is still more efficient than other 3-WL models.\n- It shows good performance on real-world graph learning tasks.\n- The paper is well-written and easy to follow.", "weaknesses": "- My major concern is about empirical evaluation.\n    - 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.\n    - 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].\n    - It would be interesting to see its long-range modeling capability on the LRGB benchmark as it's a transformer.\n    - 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.\n\n\n[1] Wang, X., Li, P., & Zhang, M. (2024). Graph as Point Set. ICML 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer.\n\nRegarding evaluation, the reviewer says that\n\n> 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\n> \n\nas well as\n\n> Besides datasets from algorithmic reasoning, currently, in my opinion, ET is only evaluated on two widely-used graph benchmark datasets\n> \n\nand\n\n> It would be interesting to see its long-range modeling capability on the LRGB benchmark as it's a transformer\n> \n\nWhile 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:\n\n- Molecular property prediction (ZINC, Alchemy, QM9)\n- Algorithmic reasoning (CLRS)\n- Graph isomorphism testing (BREC)\n\nMost 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.\n\nFinally, 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.\n\nThe above being said, we worked hard during the rebuttal to provide encouraging preliminary results on two additional datasets, see below.\n\n### PCQM4Mv2\n\nWe evaluate on one random seed, as is common on this benchmark.\n\n| Model | Num. parameters | Validation MAE |\n| --- | --- | --- |\n| GraphGPS-small | 6M | 0.0938 |\n| GRIT | 17M | 0.0859 |\n| TokenGT | 50M | 0.0910 |\n| ET | 6M | 0.0911 |\n| ET | 17M | 0.0883 |\n\n### ZINC-full\n\nWe evaluate on 2 random seeds. The ET has 283,969 parameters.\n\n| Model | Test MAE |\n| --- | --- |\n| SignNet | 0.024 $\\pm$ 0.003 |\n| Graphormer | 0.052 $\\pm$ 0.005 |\n| Graphormer-GD | 0.025 $\\pm$ 0.004 |\n| GRIT | 0.023 $\\pm$ 0.001 |\n| PST | 0.018 $\\pm$ 0.001 |\n| ET+RRWP | 0.028 $\\pm$ 0.002 |\n\nNotably, 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.\n\nWe 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.\n\nFinally, the reviewer says\n\n> 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\n> \n\nWe 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.\n\nIn what follows, we address the question of the reviewer. First, the reviewer asks\n\n> How does the setting of QM9 used in this paper differ from the setting that is commonly used in equivariant GNNs\n> \n\nIn 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.\n\nNext, the reviewer asks\n\n> 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?\n> \n\nand mentions that\n\n> I feel like depending on how PE is used, it is still possible to analyze the expressive power, e.g., via graph substructure counting\n> \n\nWhile 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.\n\nWe want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x3588jbR21", "reviewer": "Reviewer_Ff1y", "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.", "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.\n2. 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.\n3. 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? \n\n[1] Comparing Graph Transformers via Positional Encodings, ICML 2024", "limitations": "Scalability limitations are discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.  This paper proposes a concrete implementation of the Edge Transformer, and proves that it has an expressive power of 3-WL.\n2.  The proposed ET has superior empirical performance on molecular regression and neural algorithmic reasoning tasks.", "weaknesses": "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.\n2. It seems that the ET in this paper is mainly from existing work, and the difference needs to be explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer states\n\n> My major concern is the real usefulness of ET, since the high runtime and memory complexity may offset the importance of ET\n> \n\nSee 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.\n\nNext, the reviewer says\n\n> 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\n> \n\nWe 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.\n\nFinally, the reviewer says\n\n> It seems that the ET in this paper is mainly from existing work, and the difference needs to be explained\n> \n\nSee 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.\n\n### Questions\n\nThe reviewer asks\n\n> 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\n> \n\nWe 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.\n\nNext, the reviewer asks\n\n> 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\n> \n\nWe 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.\n\nFinally, the reviewer asks\n\n> But maybe as an open question, if incorporated with PEs, will the theoretical results change?\n> \n\nWe 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.\n\nWe want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WOELaZ9Plv", "reviewer": "Reviewer_ckA7", "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.", "questions": "1、Can you provide possible ways to reduce the complexity of ET?\n2、Can you provide the total number of parameters for ET on each dataset? See Weaknesses 1.\n3、Why has the performance changed so much compared to the submitted version of ICML in 2024? Especially on the Searching dataset.\n4、What do Average and All algorithms represent respectively?", "limitations": "See Weaknesses.", "rating": 1, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1、ET bridges the gap between theoretical expressivity and practical performance.\n2、The paper provides comprehensive empirical evidence demonstrating that the Edge Transformer outperforms existing theoretically motivated models and competes well with state-of-the-art models in various graph learning tasks.", "weaknesses": "1、The total number of parameters for the model is not provided in the parameter table.\n2、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.\n3、The author's contribution is confined to applying the ET model to the graph, without proposing their own method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their effort and are happy to address the concerns raised by the reviewer. First, the reviewer states\n\n> The total number of parameters for the model is not provided in the parameter table\n> \n\nWe 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.\n\nNext, the reviewers says\n\n> 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\n> \n\nSee 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.\n\nFinally, the reviewer says\n\n> The author's contribution is confined to applying the ET model to the graph, without proposing their own method\n> \n\nSee 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.\n\n### Questions\n\nFurther, we are happy to address the question of the reviewer below. First, the reviewer asks\n\n> Can you provide possible ways to reduce the complexity of ET?\n> \n\nAs 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.\n\nNext, the reviewer asks\n\n> Why has the performance changed so much compared to the submitted version of ICML in 2024? Especially on the Searching dataset.\n> \n\nIn 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.\n\nFinally, the reviewer asks\n\n> What do Average and All algorithms represent respectively?\n> \n\nIn “Average”, we average performance across the six algorithm classes, whereas in “All algorithms” we average performance over all algorithms.\n\nWe want to ask the reviewer to increase their score if they are satisfied with our response. We are happy to answer any remaining questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their time and effort and the valuable feedback they provided for our work. Here, we address two common concerns. In addition, we provide parameter counts for each dataset.

### On the novelty and contribution of our work

Here, we address the concern that our work is not novel because we do not propose a new architecture.

Indeed, we adopt the existing ET architecture, with almost no modifications. At the same time, we want to point out that novelty does not solely originate from proposing novel architectures. Concretely, we regard the main contribution of our work as establishing the connection between the ET, its systematic generalization capabilities and graph learning (Sections 3 and 4). In particular, we provide two highly non-trivial proofs that show that the ET has exactly 3-WL expressive power. Our empirical study (Section 5) is aimed at both validating the theoretical findings, as well as demonstrating that the ET can effectively translate its expressive power into strong empirical performance. As such, our work should be viewed as a theoretical and empirical investigation rather than proposing a novel method.

### On the practical usefulness of the ET

Here, we address the concern of the practical usefulness of the ET and the results in our paper, given its cubic runtime and memory complexity.

As already addressed in our limitations section, the ET has indeed a high runtime and memory complexity. At the same time, we want to highlight that the ET shows very strong performance on smaller graphs, in particular, on molecular graphs. As such, we believe the ET to be very useful in the context of molecular representation learning. Here, we want to highlight TGT (https://arxiv.org/abs/2402.04538) and AlphaFold 3 (https://www.nature.com/articles/s41586-024-07487-w), two recent works that use models with a triangular attention mechanism for learning molecular representations. In particular, both works compute attention scores over triplets of nodes. Although the attention mechanism in these works differs from the one in the ET, their strong performance demonstrates nonetheless the usefulness of studying triangular-style attention and its impact in molecular representation learning. Most notably, the authors of TGT refer to https://arxiv.org/abs/2302.05743 to motivate their work, arguing that 

> 3rd-order interactions are crucial for geometric understanding, yet ≥4th-order interactions add little/no benefit at much higher computational cost
>

In addition, the ET performs particularly well on CLRS. We refer the reviewer to the CLRS paper (https://arxiv.org/abs/2205.15659) as well to https://arxiv.org/abs/2105.02761, for an in-depth discussion of the potential use-cases of algorithmic reasoning in neural networks.

### Number of parameters

Here, we report the number of parameters of the ET without positional encodings.

| Dataset | Num. parameters |
| --- | --- |
| ZINC 12K | 280,897 |
| Alchemy 10K | 295,356 |
| QM9 | 257,916 |
| CLRS | 1,371,264 |
| BREC | 58,112 |

Here, we report the number of parameters of ET+RRWP.

| Dataset | Num. parameters |
| --- | --- |
| ZINC 12K | 283,969 |
| Alchemy 10K | 298,428 |
| QM9 | 260,988 |

Meta-review: 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).
Decision: Accept (poster)