Polynormer: Polynomial-Expressive Graph Transformer in Linear Time
Authors: Chenhui Deng
Zichao Yue
Zhiru Zhang
Venue: ICLR 2024 poster (2024)
Keywords: Graph Transformers
Polynomial Networks
Large-Scale Graph Learning
Node Classification
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: A linear graph transformer that performs well on homo/heterophilic graphs by learning high-degree equivariant polynomials.

Abstract: Graph transformers (GTs) have emerged as a promising architecture that is theoretically more expressive than message-passing graph neural networks (GNNs). However, typical GT models have at least quadratic complexity and thus cannot scale to large graphs. While there are several linear GTs recently proposed, they still lag behind GNN counterparts on several popular graph datasets, which poses a critical concern on their practical expressivity. To balance the trade-off between expressivity and scalability of GTs, we propose Polynormer, a polynomial-expressive GT model with linear complexity. Polynormer is built upon a novel base model that learns a high-degree polynomial on input features. To enable the base model permutation equivariant, we integrate it with graph topology and node features separately, resulting in local and global equivariant attention models. Consequently, Polynormer adopts a linear local-to-global attention scheme to learn high-degree equivariant polynomials whose coefficients are controlled by attention scores. Polynormer has been evaluated on $13$ homophilic and heterophilic datasets, including large graphs with millions of nodes. Our extensive experiment results show that Polynormer outperforms state-of-the-art GNN and GT baselines on most datasets, even without the use of nonlinear activation functions. Source code of Polynormer is freely available at: [github.com/cornell-zhang/Polynormer](https://github.com/cornell-zhang/Polynormer).

Published as a conference paper at ICLR 2024
POLYNORMER: POLYNOMIAL-EXPRESSWV GRAPH TRANSFORMER IN LINEAR TIME
Chenhui Deng, Zichao Yue, Zhiru Zhang Cornell University, Ithaca, USA (cd574, zy383, thiruz]@cornnn (coone edu
ABSTRACT
Graph transformers (GTs) have emerged as promising architecture that is the- oretically more expressive than message-passinng graph neural networks (GNNs). However, typical GT models have at least quadratic complexity and thus can- not scale to large graphs. While there are several linear GT's recently proposed, they still lag behind GNN counterparts on several popular graph datasets, which poses a critical concern on their practical expressivity. To balance the trade-off be- tween expressivity and scalability of GTs, propose Polynormer, polynomial- expressive GT model with linear complexity. Polynormer built upon a novel base model that learns high-degree polynomial on input features. To enable the base model permutation equivariant, we integrate it with graph topology and node features separately, resulting in local and global equivariant attention mod- els. Consequently, Polynormer adopts linear local-to- -loobba attention scheme to learn high-degree equivariant polynomials whose coefficients are controlled by at- tention scores. Polynormer has been evaluated on 13 homophilic and heterophilic datasets, including large graphs with millions of nodes. Our extensive experiment results show that Polynormer outperforms state-of-the-re thetttt GNN and GT baselines on most datasets, even without the use of nonlinear activation functions. Source code of Polynormer is freely available at: huthbb/com//oomelllzzaa///////oo
INTRODUCTION
As conventional graph neural networks (GNNs) are built upon the message passing scheme by ex- changing information between adjacent nodes, they are known suffer from wer-smoothing and
over-squashing issues (Oono & Suzuki 2020 Alon & Yahav 2021 Di Giovanni et al. 2023) re-
sulting in their limited expressive power to (approximatelyy represent complex functions (Xu al.]
2018 Oono & Suzuki, 2020). Inspired by the advancements of Transformer-based models in lan- guage and vision domains (Vaswani et al. 2017 Dosovitskiy et al. 202T), graph transformers (GTs) have become increasingly popular in recent years, which allow nodes to attend to all other nodes in graph and inherently overcome the aforementioned limitations of GNNs. In particular, Kreuzer et al. (202T) have theoretically shown that GT's with unbounded layers are universal equivariant function approximators on graphs. However, it still unclear how to unlock the expressivity potential of GT's practice since the number of GT layers typically restricted
small constant.
In literature, several prior studies have attempted to enhance GT expressivity by properly involv- ing inductive bias through positional encoding (PE) and structural encoding (SE). Specifically, YYing let al. (202T): Chen et al. (2022a); Zhao et al. 2023); Ma et al. (2023) integrate several SE methods with GT to incorporate critical structural information such as node centrality, shortest path dis- tance, and graph substructures. Moreover, Kreuzer et al. (202T); Dwivedi et al.) (2022); |Rampasek]
al. p022p; p: [B al. (2023) introduce various PE approaches based upon Laplacian eigenpairs. Nonetheless, these methods generally involve nontrivial overheads to compute PE/SE, and mostly adopt the self-attention module the vanilla Transformer model that has quadratic complexity with respect to the number of nodes in graph, prohibiting their applications in large-scale node classi- fication tasks. To address the scalability challenge, many linear GTs have been recently proposed. Concretely, Choromanski et al. [2021> Zhang al. (2022); Shirzad et al. (2023); Kong et al. (2023) aim to sparsify the attf-nttention matri via leveraging node sampling or expander graphs, while Wu et al. (2022 2023) focus kernel-based approximations the self-attention matrix. Unfor-

---

Published as a conference paper at ICLR 2024
tunately, both prior work (Platonov et al. 2023) and our empirical results indicate that those linear GT models still underperform sfate-of-the- art GNN counterparts on several popular datasets, which poses a serious concern regarding the practical advantages of linear GTs over GNNs.
In this work, we provide an orthogonal way to case the tension between expressivity and scalabil- ity of GTs. Specifically, propose Polynormer, linear GT model that is polynomial-expressiee an L-layer Polynormer can expressively represent a polynomial of degree 2", which maps input node features to output node representations and is equívariant node permutations. Note that the polynomial expressivity is well motivated by the Weierstrass theorem which guarantees any smooth function can be approximated by polynomial (Stone) 1948). To this end, we first introduce a base attention model that explicitly learns polynomial function whose coefficients are determined by the attention scores among nodes. By imposing the permutation equivariance constraint on polyno- mial coefficients based on graph topology and node features separately, we derive local and global attention models respectively from the base model. Subsequently, Polynormer adopts linear local- to-global attention paradigm for learning node representations, which is common practice for efficient transformers in language and vision domains yet less explored on graphs. To demonstrate the efficacy of Polynormer, 22 competitive GNNs and GTs on 13 node classification datasets that include homophilic and het- erophilic graphs with up to millions of nodes. We believe this is possibly one of the most extensive comparisons in literature. Our main technical contributions are summarized as follows:
conduct extensive experiments by comparing Polynormer against
To the best of our knowledge, we are the first propose polynomial-expressive graph trans- former, which is achieved by introducing a novel attention model that explicitly learns a high-degree polynomial function with its coefficients controlled by attention scores.
By integrating graph topology and node features into polynomial coefficients separately, we derive local and global equivariant attention modules. As a result, Polynormer harnesses the local-to-global attention mechanism to learn polynomials that are equivariant to node permutations.
Owing to the high polynomial expressivity, Polynormer without any activation function is able to surpass state-of-the-art GNN and GT baselines on multiple datasets. When further combined with ReLU activation, Polynormer improves accuracy over those baselines by margin of up to 4. 4.06% across 11 out of 13 node classification datasets, including both homophilic and heterophilic graphs.
Through circumventing the computation of dense attention matrices and the expensive PE/SE methods used by prior arts, our local-to-global attention scheme has linear complexity in regard to the graph size. This renders Polynormer scalable to large graphs with millions of nodes.
BACKGROUND
There is an active body of research on GT's and polynomial networks, from which we draw inspi- GT with linear complexity. In the following, we present preliminaries for both areas and provide an overview of their related work.
ration to build a polynomial-exprstte
Graph transformers exploit the Transformer architecture (Vaswani et al. 2017) graphs. Specif- ically, given an n-node graph G and its node feature matrix X E Raxa where d represents the node feature dimension, a GT layer first projects X into query, key, and value matrices, i.c., Q ==XWq K = XWK, V = XWy, where Wo. WK, Wv E dxx are three trainable weight matrices. Subsequently, the output X' with self-attention is calculated
S QKT
'X' = softmaz(S)
Vd
where S E Raxx is the self- self-attention matrix. For simplicity of illustration, we assume query, key, and value have the same dimension and only consider the single-head self-attention without bias terms. The extension to the multi-head attention is standard and straightforward.
As Equation] I completely ignores the graph topology, various PE/SE methods have been proposed to incorporate the critical graph structural information into GTs. Concretely, Kreuzer et al. (202l); Dwivedi et al. (2022): y Rampasek et al. (2022); Bo et al. (2023) use the top- Laplacian eigenpairs as node PEs, while requiring nontrivial computational costs to learn the sign ambiguity of Laplacian
eigenvectors. Similarly, SE methods also suffer from high complexity for computing the distance
of all node pairs sampling graph substructures (Ying et al. 2027 Chen al. 2022a Zhha al.)
2

---

Published as a conference paper at ICLR 2024
2023
Ma al., 2023). Apart from the expensive PE/SE computation, most of these approaches
follow Equation D to compute the dense attention matrix S, resulting their quadratic complexity with regard to the number of nodes. While there are scalable GTs recently proposed by linearizing the attention matrix and not involving PE/SE, they lack a thorough analysis of their expressivity in
practice and may perform worse than state-of-the-ar GNNs (Choromanski et al. 2021: /Zhag al.)
2022 (Shrzzd et al., [023]] Kong et al.) 2023]] [W et al. [022] (023). In this work, we provide novel way balance the expressivity and scalability of GTs, via introducing a linear GT model that
can expressively represent high-degree polynomial.
Polynomial networks aim to learn a function approximator where each element of the output is expressed as a polynomial of the input features. Formally, we denote the mode-m vector product of with a vector R. by T xm u. Given the input feature x E R and the output y E RR, polynomial networks learn polynomial P R" _ IR with degree R E N:
a tensor T E 1211xxxx....>
yj = P(a)) by + t [1)T
x // w!
W R
II x, (2)
X1 x X2 Xg xx
T2 mxnx.-. X
where by R and (Wir) E of output Note that Equation input and output of the polynomial P are matrices or higher-order tensors. Moreover, suppose æ and
1-=1 are learnable parameters for the j-th element
2 can be naturally extended to more general case where both the
have the same dimension (e.g., y E IR",, we say polynomial 'P is permutation equivariant if for any permutation a of the indices {1, 2, ... n), the following property holds:
P(a x) = a- P(x) = a y
(3)
Note that when E R" represents node features for an n-node graph (i.c., each node feature has a scalar value), the aforementioned equivariance property essentially indicates the polynomial P equivariant to node permutations. It is worth mentioning that the polynomial networks are fun- damentally distinct from the polynomial graph filtering methods widely explored in spectral-based GNNs, which purely focus on the polynomials of graph shift operators rather than node features.
In the literature, the notion of learnable polynomial functions can be traced back to the group method of data handling (GMDH), which learns partial descriptors that capture quadratic correlations be- tween specific pairs of input features (Ivakhnenko, 1971). Later, the pi-sigma network (Shin &
Ghosh 199T) and its extensions (Voutriaridis al., 2003 Li 2003) have been proposed cap- ture higher-order feature interactions, while they struggle with scaling to nigh-dimensional input features. Recently, Chrysos et al.] (2020) introduce P-nets that utilize a special kind of skip con- nections to efficiently implement the polynomial expansion with high-dimensional features. In addition, Chrysos et al. (2022) express as polynomials a collection of popular neural networks, such AlexNet (Krizhevsky et al. 2012)), ResNet (He al., 2016), and SENet (Hu et al., 2018), and improve their polynomial expressivity accordingly. While these approaches have demonstrated promising results on image and audio classification tasks, they not directly applicable to graphs.
To learn polynomials on graph-structured data, Maron et al. (2018) first introduce the basis of con- stant and linear functions on graphs that are equivariant to node permutations. Subsequently, series of expressive graph models have been developed by leveraging polynomial functions (Maron et al.) 2019 Chen et al. 2019 Azizian & Lelarge 2020: Hua al. 2022). More recently, Puny et al. (2023) demonstrate that the polynomial expressivity is a finer-grained measure than the traditional Weisteiler-Lehman (WL) hierarchy for assessing the expressive power of graph learning models. Besides, they devise a graph polynomial model that achieves strictly better than 3- BL expressive power with quadratic complexity. However, these prior studies either only consider polynomials on local structures or cannot scale to large graphs due to their high complexity. In contrast, our work learns a high-degree polynomial on node features while integrating graph topology into polynomial coefficients. As a result, the learned polynomial function captures both local and global structural information with linear complexity, rendering it applicable to large-scale graphs
METHODOLOGY
In this work, we follow the common setting that there are a graph g = (V.E) and its node fea- ture matrix X xad where n and d denote the number of nodes and node feature dimension, respectively. Our goal is to design a selyyomial-exppesiiiii GT model F that produces node repre- sentations Y F(g, X) PX))), where Pg is high-degree polynomial on X whose learnable
3

---

Published as a conference paper at ICLR 2024
X2
y1 W1,b W12byzz Wiaby Wis + Wi W
y2 Wanbox + Waabaxa + Waadbar + Waytaary + Wapa + Wagataas
X3
y (Wx)O(x
Wazdor + Wazbgka + Wag + Way + Wapa + Was
Figure 1: A toy example 3-node graph with scalar node features.
coefficients encode the information of g. To this end, we first introduce base attention model in Section 3.1 that explicitly learns high-degree polynomials. To enable the learned polynomial extend the base model to equivariant local and global (linear) attention models in Section 3.2 via incorporating graph topology and node features respectively. Finally, Section 3.3 presents the proposed Polynormer architecture that employs local-to-global attention scheme based on the equívariant attention models. As result, Polynormer preserves high polynomial expressivity while simultaneously benefiting from linear complexity.
equivariant to node permutations,
3.1 POLYNOMIAL-EXPRESSII BASE MODEL WITH ATTENTION
By following the concept of polynomial networks in Hua et al. (022); Chrysos et al. (020 2022), we provide Definition 3.T of the polynomial expressivity, which measures the capability of learning high-degree polynomial functions that map input node features into output node representations. Appendix
M provides a detailed discussion on our definition and comparison prior work.
Definition 3.1. Given an n-node graph with node features X RRnx a model P RRaxd Raad is r-polyyomial-exppexse P can be parameterized such that P(X); ; = X; 3 M where @ denotes the Hadamard product.
if for any node i and (rgree-(r-1) monomial M formed by rows in X,
Polynomial expressivity of prior graph models. For typical message-passing GNN models, we can unify their convolution layer as: P(X); E, C.j where Cr,j is the edge weight between nodes i and j. As each output is a linear combination of input node features, these models are at most wwwprmommmalxxppssiieeee
Hence, they mainly rely on the activation function to capture nonlinearity, instead of explicitly learning high-degree polynomials. In regard to GTs and high-orde GNNs (c.g., gating-basee GNNs), they only capture a subset of all possible monomials with a certain degree, resulting in their limited polynomial expressivity, as discussed in Appendix C
motivating example. simplified scenario where node features æ E R i.e., each node feature is a scalar. In this context, we consider a model PP that outputs = ((x) = [Wan ax) @ (xc + b), where W E RR and E Ik" able represent degree-2 polynomials, which consist of all possible monomials of degree up to 2 (except the constant term). Besides, as W controls the coefficients of quadratic monomials (i.e., 22 23 and Z2 we can interpret W as general attention matrix, where W.j represents the importance of node to node i.
Before constructing a polynomial-expressiv model, let us first examine a
are weight matrices. Figure [ shows that this model
Motivated by the above example, we proceed establish base model using the following defini- tion; we then theoretically analyze its polynomial expressivity.
Definition 3.2. Given the input node features X (0) gaxd and trainable weight matrices W Raxn B E paxd, a model P is defined as the base model its (-th layer is computed as:
x (1) (W x @ (x (1-) + B())
(4)
Theorem 3.3. An L-layer base model P is 2 -olyoooiaal-eeppesssssee
The complete proof for Theorem 3.3 available in Appendix[A [A] Theorem 3.3 shows that the polyno- mial expressivity of the basc model increases exponentially with the number of layers. Additionally, Appendices A and reveal that W and B control the learned polynomial coefficients, which can be viewed as attention scores for measuring the importance of different node feature interactions.
Limitations of the base model. As the matrices W and B fail to exploit any graph inductive bias, the base model learns polynomials that not equivariant node permutations. Besides, the model has quadratic complexity due to the n² size of W and thus cannot scale to large graphs. To address both limitations, derive two equivariant models with linear attention in Section 3.2
4

---

Published as a conference paper at ICLR 2024
3.2 EQUIVARIANT ATTENTION MODELS WITH POLYNOMIAL EXPRESSIVITY
For clarity, we omit the layer index (1) in the ensuing discussion unless it is explicitly referenced. Instead of learning the non-equivariant matrix B IRa in Equation 4 we replace it with learnable weights / E R sharing across nodes, i.e., B = 1,3", where 1 E IR" denotes the all-ones vector. In addition, we apply linear projections on X allow interactions among feature channels. This leads to Equation 5 where V = Wy and H = XWH with trainable matrices Wv, WH E Raxd
X =
(H + 1,3")
Subsequently, it straightforward to show that we only need to achieve permutation equivariance on the term W in Equation to build an equivariant model. To this end, we introduce the following two equivariant attention models, by leveraging graph topology and node features respectively.
Equivariant local attention. We incorporate graph topology information by setting W A, where A is (general) sparse attention matrix such that the nonzero elements in A represent attention scores of adjacent nodes. While A can be implemented by adopting any sparse attention methods previously proposed, we choose the GAT attention scheme (Veličković et al. 2017) due to its effi cient implementation. We leave the exploration of other sparse attention approaches future work. Consequently, the term AV holds the equivariance property, i.e., PP) (PV) P(AV) for any node permutation matrix P. It is noteworthy that replacing W with essentially imposes a sparsity constraint on the learned polynomial of the base model, such that polynomial coefficients are nonzero if and only if the corresponding monomial terms are formed by features of nearby nodes.
Equivariant global attention. By setting W softmar((S), where S denotes the global self- attention matrix defined in Equation the term WV in Equation becomes so soft ttaar( ()) V that is permutation equivariant (Yun al. 2019). According our discussion in Appendix A an T--layer model based on the updated Equation learns an equivariant polynomial of degree 2 with the coef- ficients determined by attention scores in S. More importantl, the learned polynomial contains all monomial basis elements that capture global and high-order feature interactions. Nonetheless, the dense attention matrix S still has quadratic complexity in regard to the number of nodes, rendering the approach not scalable to large graphs. To tackle this issue. we linearize the global attention by introducing a simple kernel trick in Equation [6] where Q XWo. K XWx k with weight ma- trices Wo, WK, and denotes the sigmoid function to guarantee the attention scores are positive. The denominator term Equation Soniiono@@eeee [ensures that the sum of attention scores is normalized node. In this way, we preserve two key properties of the so soft ftraaz function: non-negativvty and nor- malization. As we can first compute o K V whose output is then multiplied by o (Q)), Equation 6 avoidd computing the dense attention matrix, resulting in the linear global attention. In Appendix D we further demonstrate the advantages of our approach over prior kernel- kernel-based linear attention methods in terms of hyperparameter tuning and training stability on large graphs.
per
WV (QQ) (KT) V ((Q)(((T)) =Q (K
(QQ) E.
Complexity analysis. Given a graph with n nodes and m edges, suppose the hidden dimension is d < n, then the local attention model has the complexity of O(md + ndd)). Since we exploit the kernel trick to linearize the computation of global attention in Equation [6] the complexity of global attention model is reduced from O(n2d) d) (ndd). Hence, both proposed equivariant attention models have linear complexity with respect to the number of nodes/edges.
Discussion. It is noteworthy that the introduced local and global attention models are intended to enable the base model to learn high-degree equivariant polynomials, which is fundamentally distinct from the purpose of prior attention models in literature. This allows our proposed Polynormer model to outperform those attention- baselines, even without nonlinear activation functions.
3.3 THE POLYNORMER ARCHITECTURE
Guided by the aforementioned equivariant models with linear attention, Figure 2 shows the Poly- normer architecture that adopts a local-to-globee ologabaal attention scheme based on the following modules:
Local attention module. We employ the local attention approach via replacing W in Equation S with the sparse attention matrix A as discussed in Section 3.2 which explicitly learns an equivariant
5

---

Published as a conference paper at ICLR 2024
ASD Kg
A
16"
80
2000
0000
K v
Q K V
H
K
H
Sigmaid
Linear Projections
Linear Projections
Linerr Projections
(8) Local-and-GGlber Attention
(b) Local-to-Global Attention (Oure)
Figure 2: Distinctions between the attention schemes in previous work and Polynormer- (a) Prior GTs use a local-and-glote attention scheme, involving the simultancous use of local and global attentions in each layer; (b) We adopt local-to-giobal attention scheme, where local and global attention modules are applied sequentially; Note that the softmar operator in (a) and the attention normalization in (b) are omitted for brevity.
polynomial with coefficients controlled by A and 33. Since the attention scores in A are normalized the range of (0, 1). we apply a sigmoid function a on 3 to scale it into the same range as A, which prevents [3 from dominating the polynomial coefficients. This results in the local attention layer of Polynormer shown in Equation 7 By stacking Ln layers, the local attention module outputs X Xlocal = E/a (1), which captures equivariant polynomials with different degrees based on local attention. X locol then serves the input for the following global attention module.
X = AV O (H + (113"))
Global attention module. By substituting Equation Equation we obtain the global attention layer shown in Equation 8 After stacking L2 global attention layers, the global attention module takes as input X and produces the final node representations for downstream tasks.
(QQ) (Q))
X
(H o(13T)
Note that the Polynormer architecture described above does not apply any activation function on the outputs per layer. Inspired by prior polynomial networks on images (Chrysos et al.) (022), we can optionally integrate ReLU activation into each layer, which potentially further improves the model performance by involving additional nonlinearity. We denote this version by Polynormer-r.
Discussion. Compared the popular architecture of prior GTs shown in Figure a Polynormer builds upon our proposed base model, thus preserving high polynomial expressivity. Besides, our local-to-global attention scheme outperforms the prior localaand-loobe attention scheme, as empir- ically demonstrated in Section 4.3 Another important advantage of Polynormer is its ease of imple- mentation. By avoiding the enumeration of equivariant polynomials, Polynormer offers a concise way to learn high-degree equivariant polynomials, which be easily implemented using popular graph learning frameworks such as PyG (Fey & Lenssen 2019) and DGL (Wang et al. 2019). We further provide an analysis of Polynormer expressivity under the WL hierarchy in Appendix
EXPERIMENTS
We have conducted extensive evaluation of Polynormer against state-of-the-art (SOTA) GNN and GT models on both homophilic and heterophilic graphs, where nearby nodes tend to have the same or different labels. Besides, we demonstrate the scalability of Polynormer on large-scale graphs that contain millions of nodes. In addition, we perform additional ablation analysis to understand the effectiveness of the local-to-globa. attention scheme adopted by Polynormer. Finally, we visualize Polynormer attention scores showcase its ability of capturing critical global structures.
Experimental Setup. Appendix D shows the details of all 13 datasets used in our experiments, which consist of 7 homophilic and 6 heterophilic graphs. Notably, the heterophilic datasets are from Platonov et al. (2023): ; Lim et al. (202I). which have addressed serious issues (e.g.. train-test data leakage) of conventional popular datasets such chameleon and squirrel. We compare Poly- normer against 10 competitive GNNs. Note that prior graph polynomial models except (GNN (Hua al. 2022) discussed in Section 2 out of GPU memory even the smallest dataset in

---

Published as a conference paper at ICLR 2024
experiments. Besides, we also evaluate 6 GTs that have shown promising results the node clas- sification task. We report the performance results of baselines from their original papers or official leaderboards whenever possible, as those results are obtained by well-tuned models. For baselines whose results are not publicly available on given datasets, we tune their hyperparameters to achieve the highest possible accuracy. Detailed hyperparameter settings of baselines and Polynormer available in Appendix H Our hardware information is provided in Appendix F
4.1 PERFORMANCE ON HOMOPHILIC AND HATEROPHILIC GRAPHS
Table 1: Averaged node classification accuracy (%) std over 10 runs on homophilic datasets. Polynormer-r denotes Polynormer with ReLU activation. We highlight the top first, second, and third results per dataset.
Computer
Photo
CS
Physics
WikiCS
GCN
89,65 + 0.52
92.70 0.20
92.92 +: 0.12
96.18 = 0.07
77.47 0.85
GraphSAGE
91.20 + 0.29
94.59 + 0.14
93.91 0.13
96.49 = 0.06
74.77
0.95
GAT
90.78 + 0.13
93.87 + 0.11
93.61 0.14
96.17 0.08
0.82
GCNII
91.04 1 0.41
94.30 0.20
0.12
95.97 0.11
78.68 0.55
GPRGNN
89.32 0.29
94.49 0.14
0.01
96.85 0.08
0.23
APPNP
0.11
94.32
0.049
96.54 0.07
7.18
PPRGo
0.21
98.61
0.15
95.51 0.08
77.89 0.42
GGCN
9
91.50
0.05
97.07 0.05
78.41
OrderedGNN
0.103
95.10 0.20
0.100
97.00 1.00000.8 0.08
79.01 + 0.68
IGNN
83.40 + 1.33
89.92 = 0.72
12.85
96.24 0.24
71.49 + 1.05
GraphGPS
91.19 0.54
95.06 + 0.13
0.12
97.12 0.19
7
NAGphormer
91.22 + 0.14 0.14
95.49 + 0.11
9.079
97.34 0.03
0.72
Exphormer
91.47 0.17
95.35 + 0.22
0.001
96.89 0.09
+.49
NodeFormer
86.98 + 0.62
93.46 0.35
96.45 0.22
74.73
DIFFormer
91.99 0.76
95.10 0.47
0.20
96160
73.46
GOAT
90.96 + 0.90
92.96 + 1.48
94.21 4
96.24 0.24
77.00 + 0.77
Polynormer
93.18 = 0.18
96.11 + 0.23
95.51 + 0.29
97.
0.06
79.53 + 0.83
Polynormer-r
93.68 0.21
96.46 + 0.26
95.53 + 0.16
97.27 0.08
80.10 + 0.67
Performance on homophilic graphs. We first compare Polynormer with 16 popular baselines that include SOTA GNNs or GTs on common homophilic datasets. As shown in Table I Polynormer is able to outperform all baselines on 3 out of 5 datasets, which clearly demonstrate the efficacy of model. Moreover, incorporating the ReLU function can lead to further improvements in the accuracy of Polynormer. This is because the nonlinearity imposed by ReLU introduces additional higher-order monomials, which in turn enhance the quality of node representations. It is noteworthy that the accuracy gain of Polynormer over SOTA baselines is 1.65% on Computer, which is a nontrivial improvement given that those baselines have been finely tuned on these well-established homophilic datasets.
the proposed saayymmmallexpresttieee
Performance on heterophilic graphs. Table 2 reports the average results over 10 runs on het- erophilic graphs. Notably, the GT baselines underperform SOTA GNNs all datasets, which raises concerns about whether those prior GT models properly utilize the expressivity brought by the self-attention module. Moreover, there the top 3 models across 5 datasets. In contrast, by integrating the attention mechanism into the polynomial-expressiv model, Polynormer surpasses baselines more, Polynormer-r consistently outperforms baselines across all datasets with the accuracy im- provement by a margin up to 3.55% (i.e., 97.46% vs. 93.91 % on minesweeper).
baseline model that consistently ranks among
4 out of datasets. Further-
4.2 PERFORMANCE ON LARGE GRAPHS
We conduct experiments on large graphs by comparing Polynormer against GT's as well as some representative GNN models. Since the graphs in oghn-products and pokec are too large to loaded into the GPU memory for full-batch training, we leverage the random partitioning method adopted by prior GT models Wu et al. 2022 2023) to perform mini-batch training. As shown in Table 3 Polynormer-r outperforms all baselines on these large graphs with the accuracy gain of up to 4.06%. Besides, we can observe that the accuracy of Polynormer drops 1.64% when removing ReLU acti- vation on ogbn-arxiv. Thus, we believe the nonlinearity associated with Re LU is more critical on

---

Published as a conference paper at ICLR 2024
Table 2: Averaged node classification results over 10 runs on heterophilic datasets Accuracy is reported for roman-empire and amazon-ratings, and ROC AUC is reported for minesweeper, tolok- ers, and questions. Polynormer-r denotes Polynormer with ReLU activation. We highlight the top first, second, and third results per dataset.
roman-empire
amazon-ratings
minesweeper
tolokers
questions
GCN
73.69 + 0.74
48.70 : 0.63
89.75 : 0.52
83.64 +
76.09 1.27
GraphSAGE
85.74 0.67
53.63 0.39
98.51 + 0.57
82.43 + 0.44
76.44 0.62
GAT-sep
88.75 + 0.41
52.70 + 0.62
93.91 0.35
83.78 - 0.43
76.79 ::.99::00 0.71
H2GCN
60.11 = 0.32
36.17 + 0.23
89.71 + 0.31
73.35 + 1.01
63.59
1.46
GPRGNN
64.85 0.27
44.88 : 0.34
86.24 0.61
72.94
$5.48 0.91
ESGNN
+
52.74 + 0.83
90.08 + 0.70
82.76
78.86 0.92
GloGNN
59.63
36.89 + 0.14
51.08 + 1.23
73.39 1.17
65.74 1.19
GGCN
74.46
43.00 0.32
87.54 1.22
77.31 1.14
71.10 1.57
OrderedGNN
+
47.29 0.65
80.58 +
75.60 +
75.09 1.00
G2-GNN
0.78
47.93 + 0.58
91.83 0.56
82.51 + 0.80
74.82 + 0.92
DIR-GNN
91.23 0.32
17.89 + 0.39
87.05 + 0.69
81.19 + 1.05
76.13 + 1.24
IGNN
79.95 + 0.75
48.21 + 0.53
91.93 + 0.77
70.84 + 1.75
76.38 + 1.79
GraphGPS
82.00
53.10 0.42
90.63 + 0.67
83.71
71.73 1.47
NAGphormer
74.34 0.77
51.26 + 0.72
84.19 + 0.66
78.32 +
68.17 + 1.53
Exphormer
89.03 + 0.37
53.51 0.46
90.74 0.53
83.77 0.78
73.94 1.06
NodeFormer
64.49 - 0.73
43.86 + 0.35
86.71 0.88
78.10
74.27 + 1.46
DIFFormer
79.10 0.32
47,84 + 0.65
90.89 + 0.58
83.57 0.68
72.15 + 1.31
GOAT
71.59 + 1.25
44.61 + 0.50
81.09 + 1.02
83.11 1.04
75.76 + 1.66
Polynormer
92.13 0.50
54.46 0.40
96.96 0.52
84.83 33 0.72
77.95 + 1.06
Polynormer-r
92.55 + 0.37
54.81 0.49
97.46 0.36
85.91 + 0.74
78.92 92 + 0.89
Table 3: Averaged node classification accuracy (%) + std on large-scale datasets Polynormer-r denotes Polynormer with LU activation. We highlight the top first, second, and third results per dataset. OOM means out of memory.
ogbn-arxiv
ogbn-products
pokec
GCN
71.74 0.29
75.64
75.45 0.17
GAT
72.01 0.20
79.45 4.50
72.23 0.18
GPRGNN
71.10 0.12
79.76
78.83 + 0.05
LINKX
66.18 0.33
7.759
82.04 0.07
GraphGPS
70.97 0.41
OOM
OOM
NAGphormer
70.13 0.55
73.55 0.
76.59 0.25 00M
Exphormer
72.44 0.28
00M
NodeFormer
67.19 1119-000 0.83
72.93 0.13
71.00 1.30
DIFFormer
69.86 0.25
7..1
73.89 + 0.35
GOAT
72.41 1 0.40
82.00
66.37
Polynormer
71.82 0.23
82.97 0.28 0.28
85.95 - 0.07
Polynormer-r
73.46 0.16
83.82 0.11
86.10 0.05
ogbn-arxiv, which is known to be challenging dataset (Shirzad et al. 2023). 3.. In addition, Appendix provvddes the training time and memory usage of Polynormer and GT baselines. Moreover, we fur- ther evaluate Polynormer
industrial-level graph benchmark named TpuGraphs (Phothilimthanal ct al. 2024), whose results are provided in Appendix
4.3 ABLATION ANALYSIS POLYNORMER ATTENTION SCHEMES
Figure 3 shows the comparison of the SOTA baseline (red bar), Polynormer without global attention (orange bar), variant of Polynormer where the local and global attention modules are applied simultaneously to update node features per layer (blue bar), and Polynormer (green bar). We provide the detailed architecture of the Polynormer variant in Appendix G Besides, we omit the results of Polynormer without local attention, it completely ignores the graph structural information and thus performs poorly on the graph datasets.
observe that the local attention model achieves comparable results to Polynormer on homophilic graphs (Computer and Photo), while lags behind Polynormer heterophilic graphs (roman-empire and minesweeper). This observation
By comparing the orange and green bars in Figure
8

---

Published as a conference paper at ICLR 2024
SOTA Basaine
Local Attontion
Local-and-Glccen Attention
Local-to-Ginbal Attention
Accuracy
93
ROC AUC.
92
96
85
91
84
90
94
Computer
Photo
minesweeper
Figure 3: Ablation studies attention modules of Polynormer "Local Attention" means only the local attention module used. "Local-and-Global Attention" denotes the local and global atten- tion modules employed in parallel, while LLocal-to-Glober Attention" represents proposed Polynormer model where the local attention module is followed by the global attention module.
highlights the efficacy of the global attention module in Polynormer, which captures global infor- mation that proves more beneficial for heterophilic graphs (Li et al., 2022), as further illustrated in Section 4.4 Moreover, Figure demonstrates that the local-to-glone -to-bl attention adopted by Poly- normer performs better than counterpart, i.e., local-and-glohal attention, which has been widely used by prior GT models (Rampasek et al. 2022 Wu et al. 2022 2023 Kong et al. (023). We at- tribute to that the local-and-global attention may introduce the risk mixing local and non-local interactions (Di Giovanni al. 2023), while the local-to-global attention inherently avoids this is- sue. This is further confirmed by the results on Computer, where the seaal-an-gllobr attention model performs even worse than the local attention alone model.
4.4 VISUALIZATION
1.0
1.0
1.0
0.8
0.8
0.8
0.8
0.8
0.6
0.4
0.4
0.4
(a) Locally Important Nodes
(b) Globally Important Nodes
(c) Global Atin.
Polynormer
Figure 4: Visualization on the importance of nodes (columns) to each target node (row) Higher heatmap values indicate greater importance; Both subfigures (a) and (b) consider nodes are important they share the same label as the target node, while (a) has an additional constraint that these nodes most 5-hop away from the target node; Subfigure (c) measures node importance based on the corresponding global attention scores in Polynormer.
For the sake of clear visualization. which three heatmaps are drawn Figure Notably, the coordinates in each heatmap correspond node indices in the graph. We first draw Figure ()) by assigning every heatmap value if the corresponding two nodes have the same label and their shortest path distance is up to 5, and otherwise. This essentially highlights locally important neighbors per node. Similarly, we further visualize globally important nodes by removing the local constraint, resulting in Figure Figure 4() visualizes the attention matrix in the last global attention layer of Polynormer. Note that linearly scale the attention scores such that the maximum value is for the visualization purpose. By comparing Figures (aa) and (b), we observe numerous globally important nodes that can be exploited for accurate predictions target nodes. Figure (c) clearly showcases that the attention scores Polynormer effectively differentiate those globally important nodes from unimportant ones. As result, these attention scores enable Polynormer to focus crucial monomial terms that consist of a sequence of globally important nodes, allowing it to learn critical global structural information.
randomly sample 100 nodes from minesweeper, based
5 CONCLUSIONS
This work introduces Polynormer that adopts linear local-to-globe. attention scheme with high polynomial expressivity. Our experimental results indicate that Polynormer achieves SOTA results wide range of graph datasets, even without the of nonlinear activation functions.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work is supported in part by NSF Awards #2118709 and #2212371, Qualcomm Innovation Fellowship, and ACE, one of the seven centers JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA.
REFERENCES
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical im- plications. In International Conference on Learning Reprexentations, 2021. URL https:
//openreview wttp://e///tonmo 20tt/Eorrmm1111BB000000
Waiss Azizian and Marc Lelarge. Expressive power of invariant and equivariant graph neural net- works. arXiv preprint arXiv:2006. 15646, 2020.
Jinheon Back, Minki Kang, and Sung Ju Hwang. Accurate learning of graph representations with graph multiset pooling. In International Conference on Learning Representations, 2021. URL
https:://ppenreee ew net/ htrr//ormmm//// d= Gn
Deyu Bo, Chuan Shi, Lele Wang, and Renjie Liao. Specformer: Spectral graph neural networks meet transformers. In The Eleventh International Conference on Learning Representations, 2023. URL
https::/cpenreviie net Forum F 2id==dddddpppd.t....
Alcksandar Bojchevski. Johannes Gasteiger, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberezki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2464-2473, 2020.
Dexiong Chen, Leslie O' Bray, and Karsten Borgwardt. Structure- aware transformer for graph rep- resentation learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 3469-3489. PMLR, wtt.i/pppeeessweeeheeeeeeee
17-23 Jul 2022a. URL
html
Jinsong Chen, Kaiyuan Gao, Gaichao Li, and Kun He. Nagphormer: A tokenized graph transformer for node classification in large graphs. In The Eleventh International Conference on Learning Representations, 2022b.
Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. In International conference on machine learning, pp. 1725-1735. PMLR, 2020.
Zhengdao Chen, Soledad Villar, Lei Chen, and Joan Bruna. On the equivalence between graph isomorphism testing and function approximation with gnns. Advances in neural information processing systems, 32, 2019.
Eli Chien, Jianhao Peng, Pan Li, and Olgica Milenkovic. Adaptive universal generalized pagerank graph neural network. arXiv preprint arXiv:2006.07988. 2020.
Krzysztof Choromanski, Han Lin, Haoxian Chen, Tianyi Zhang, Arijit Sehanobish, Valerii Likhosh- erstov, Jack Parker- Holder,, Tamás Sarlós, Adrian Weller, and Thomas Weingarten. From block- toeplitz matrices to differential equations on graphs: towards general theory for scalable masked transformers. In International Conference on Machine Learning, 2021.
Grigorios G Chrysos, Stylianos Moschoglou, Giorgos Bouritsas, Yannis Panagakis, Jiankang Deng, and Stefanos Zafeiriou. P-nets: Deep polynomial neural networks. In Proceedings of the JEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 7325-7335, 2020.
Grigorios G Chrysos, Markos Georgopoulos, Jiankang Deng, Jean Kossaifi, Yannis Panagakis, and Anima Anandkumar. Augmenting deep classifiers with polynomial neural networks, In European Conference Computer Vision, pp. 692-716. Springer, 2022.
10

---

Published as a conference paper at ICLR 2024
Francesco Di Giovanni, T Konstantin Rusch, Michael M Bronstein, Andreea Deac, Mare Lackenby, Siddhartha Mishra, and Petar Veličković. How does over-squashing affect the power gnns? arXiv preprint ar arXiv:2306.000 035809 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openreview net. forst//foummmtom d== wbF.dTTTT
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representations, 2022. URL httt:::///penrreiewweeeee
id=wTTinnGphy
Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv: 022428 2019.
Johannes Gasteiger, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate:
Graph neural networks meet personalized pagerank. arXiv preprint
2018.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Jie Hu, Li Shen, and Gang Sun. Squeezo-and-exxitatto networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:221 18-22133,, 2020.
Chenqing Hua, Guillaume Rabusseau, and Jian Tang. High-order pooling for graph neural networks with tensor decomposition. Advances in Neural Information Processing Systems, 35:6021-6033, 2022.
Alexey Grigorevich Ivakhnenko. Polynomial theory of complex systems. IEEE transactions Systems, Man, and Cybernetics, pp. 364-378, 1971.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: 1609.009777 2016.
Kezhi Kong, Jiuhai Chen, John Kirchenbauer, Renkun Ni, C. Bayan Bruss, and Tom Gold- stein. GOAT: A global transformer on large-scale graphs. In Andreas Krause, Emma Brun- skill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Proceed- ings of Machine Learning Research, pp. 17375-17390. PMLR, 23-29 Jul 2023. URL https: //proceedings.no .mr....mmrr......... www...ppeess//ee Titeas/v202// wwtp//2///ookaaaa.aaaa html
Devin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau, and Prudencio
Tossou. Rethinking graph transformers with spectral attention. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL URL https: /ttLattttt///peet htttps:://penreview.cc net/ www.mooirdiiioiooooo.o wtwprmmmii-hhha--i--ee
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Kuo Li. A sigma-pi- neural network (spsnn). Neural Processing Letters, 17:1-19, 2003.
11

---

Published as a conference paper at ICLR 2024
Xiang Li, Renyu Zhu, Yao Cheng, Caihua Shan, Siqiang Luo, Dongsheng Li, and Weining Qian. Finding global homophily in graph neural networks when meeting heterophily. In International Conference Machine Learning, pp. 13242-13256 PMLR, 2022.
Derek Lim, Felix Hohne, Xiuyu Li, Sijia Linda Huang, Vaishnavi Gupta, Omkar Bhalerao, and Ser Nam Lim. Large scale learning non-homophilous graphs: New benchmarks and strong simple methods. Advances in Neural Information Processing Systems, 34:20887-20902, 2021.
Liheng Ma, Chen Lin, Derek Lim, Adriana Romero-Soriano, Puneet Kumar Dokania, Mark Coates, Philip H. S. Torr, and Ser Nam Lim. Graph inductive biases in transformers without message passing. ArXiv, abs/2305. 17589, 2023.
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902, 2018.
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Advances in neural information processing systems, 32, 2019.
Sunil Kumar Maurya, Xin Liu, and Tsuyoshi Murata. Simplifying approach to node classification in graph neural networks. Journal of Computational Science, 62: 62:101695, 2022.
Péter Mernyei and Cătälina Cangea. Wiki-es: A wikipedia-based benchmark for graph neural net- works. arXiv preprint arXiv:2007.02901, 2020.
Kenta Oono and Taiji Suzuki. Graph neural networks exponentially lose expressive power for node classification. In International Conference on Learning Representations, 2020. URL https:
Mangpo Phothilimthana, Sami Abu-EElHaija, Kaidi Cao, Bahare Fatemi, Michael Burrows, Charith Mendis, and Bryan Perozzi. Tpugraphs: A performance prediction dataset on large tensor com- putational graphs. Advances in Neural Information Processing Systems, 36, 2024.
Oleg Platonov. Denis Kuznedelev. Michael Diskin, Artem Babenko, and Liudmila Prokhorenkova. A critical look at the evaluation of GNNs under heterophily: Are we really making progress? In The Eleventh International Conference on Learning Representations, 2023. URL https:
7 //peereview. / for wtt.rmmmmdddtttfeeee.oem d=w.dbofwwwwwwtt 5wv
Omri Puny, Derek Lim, Bobak Toussi Kiani, Haggai Maron, and Yaron Lipman. Equivariant poly- nomials for graph neural networks. ArXiv, abs/2302.115553 2023.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Ling- peng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint arXiv:2202.08771, 2022.
Ladislay Rampasek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Pro- cessing Systems, 2022. URL http:tt::: httt:s:://pppreeeeeeeee /ttp:////ppeeeiieeeeeeeee net www.forrmmmfoo 1 60xKM
Emanuele Rossi, Bertrand Charpentier, Francesco Di Giovanni, Fabrizio Frasca, Stephan Günnemann, and Michael Bronstein. Edge directionality improves learning on heterophilic graphs. arXiv preprint 230.: 2023.
T Konstantin Rusch, Benjamin P Chamberlain, Michael W Mahoney, Michael M Bronstein, and Siddhartha Mishra. Gradient gating for deep multi-rate learning on graphs. arXiv preprint 21100:2200.00513. 2022.
Ryoma Sato. A survey on the expressive power of graph neural networks. arXiv preprint arXhv:2003.04078, 2020.
Yoan Shin and Joydeep Ghosh. The pi-sigma network: An efficient higher-order neural network for pattern classification and function approximation. In WICNN-9I-Seatie international joint conference on neural networks, volume pp. 13-18. IEEF, 1991.
12

---

Published as a conference paper at ICLR 2024
Hamed Shirzad, Ameya Velingker, B. Venkatachalam, Danica J. Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. ArXiv, abs/2303.06117, 2023.
Yunchong Song, Chenghu Zhou, Xinbing Wang, and Zhouhan Lin. Ordered gnn: Ordering message passing deal with heterophily and verrssmothinn. arXiv preprint arXiv:2302.01524 2023.
Marshall H Stone. The generalized weierstrass approximation theorem. Mathematics Magazine, 21 (5):237-254, 1948.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Lio, and Yoshua Bengio. Graph attention networks, arXiv preprint arXiv: 2017.
Christodoulos Voutriaridis, Yiannis S Boutalis, and Basil G Mertzios. Ridge polynomial networks in pattern recognition. In Proceedings EC-VIP-MC 2003. 4th EURASIP Conference focused on Video/Image Processing and Multimedia Communications (IEEE Cat. No. 03EX667), volume 2,
pp. 124. IEEE, 2003.
Minjie Wang. Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, et al. Deep graph library: A graph-centric, highly-performe package for graph neural networks. arXiv preprint urXiv:
2019.
Ruoxi Wang. Rakesh Shivanna. Derek Cheng. Sagar Jain, Dong Lin, Lichan Hong, and Ed Chi. Den v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems. In Proceedings of the web conference 2021, pp. 1785-1797, 2021.
Qitian Wu, Wentao Zhao, Zenan Li, David Wipf. and Junchi Yan. Nodeformer. scalable graph structure learning transformer for node classification. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL htttts:://ppeeeeeiieeeeeee//eeee
net./ forum? 3==aMez7GG5550
Qitian Wu, Chenxiao Yang, Wentao Zhao, Yixuan He, David Wipf, and Junchi Yan. DIFFormer: Scalable (graph) transformers induced by energy constrained diffusion. In The Eleventh Interna- fional Conference on Learning Representations, 2023. URL ptttps://eperreviewe net/
Forum?id-j62uaae
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful graph neural networks? arXiv preprint 1rXiv:1810.00822 2018.
Yujun Yan, Milad Hashemi, Kevin Swersky, Yaoqing Yang, and Danai Koutra. Two sides of the same coin: Heterophily and oversmoothing in graph convolutional neural networks. 2022
IEEE International Conference Data Mining (ICDM), pp. 1287-1292. IEEE, 2022.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform badly for graph representation? In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
net forum? OeWooOxE
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeee functions? arXiv preprint arXiv:1912.10077, 2019. 2019.
Zaixi Zhang, Qi Liu, Qingyong Hu, and Chee-Kong Lee. Hierarchical graph transformer with adaptive node sampling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave. and Kyunghyun Cho (eds.), Advances Neural Information Processing Systems, 2022. URL https:// reeireview met//foumm? ddoouumaidx
Haiteng Zhao, Shuming Ma. Dongdong Zhang, Zhi-Hong Deng, and Furu Wei. Are more layers beneficial to graph transformers? The Eleventh International Conference on Learning Repre- sentations, 2023. URL URL http:itttps./ppeeeeeeeeeee.. hewpeemerrriwe..eee////rrroo http://ooummiiiuaa-aaa/xxe-x
13

---

Published as a conference paper at ICLR 2024
Jiong Zhu, Yujum Yan, Lingxiao Zhao, Mark Heimann, Leman Akoglu, and Danai Koutra. Beyond homophily in graph neural networks: Current limitations and effective designs. Advances in neural information processing systems, 33:7793-7804, 2020.
14

---

Published as a conference paper at ICLR 2024
A PROOF FOR THEOREM 3.3
Proof. For the convenience of our proof, we denote the initial feature matrix by X E Raxd the number of layers by L, the number of nodes by n, and [n] 21, n.. Besides, we set B vl [L]. The role of B is analyzed in Appendix B By defining / to be an ordered set that consists of node indices (e.g.. I = (i1, 12, i3, i4)). we define the following function er. on coefficients:
c W)
w(l Er
W (1-1) W (1-1).
w 1/ -1
(2) 12" 1
13.44
2 times
2 times
times
where W fa represents the (ia, is) block in W at the j-th layer. Based on Equation we further define the following monomial Sr. of degree 24, where each X; denotes d-dimensional feature vector of node i:
[I W,; X) CT; W) in C X Xiz.
(10)
Next, we prove by induction that the L-layer base model produces the following node representation x! (L) for any node i [n]:
x
Σ [u(:: X) JES
(11)
where Si2 is set that represents all the combinations of choosing 24 elements from [x]] (with replacement), with the first element fixed to be i, i.e., S2 == {v E [n]2* y1 = =i.. Base case. When L we have:
(1)
3 X
X; 3 X,
Σ
( X,, 0 X,,
Σ ((;,W,))
(12)
JESY
Induction step. When L = I, suppose we have x(1
Σ TESt
f(1; W, X). Then, for L + 1,
we have:
x (+1)
W
(1+) x (1)
@ x(")
j=1
Σ
0 Σ f(l:W,
JES 2!
IES
(1))
c);;W) X,, o X ... ..... 2
ijj
JES
a (T; W X,, ...
... xx1)))
LES
Σ
al::W (1+1)) aJ;
( W) X,, o Xi, ... Xigi
LEST 1=l JESS g'
15

---

Published as a conference paper at ICLR 2024
21 O
X2
E c41(1;W) X;, C X12- O o Xinn
X Xig41
168
Σ (I:(W:WWX
(13)
LES 24+1
The first 3 equations and the last equation above are straightforward. For the fourth equation, we essentially expand the product in LHS the sum in RHS. For the second last euqation, merge the weight matrix ((++1) into the coefficient function cj(J; W), and then replace the index set J (2,...,.. } with = { +1, 42:42: 12 J. which is combined with the set {ú1, i2,-., 12 to obtain the new index set = {úr, 22 12 "1 (2 } E St S This complements our proof for Equation| 11
As for any node i, the first index E Syg is always i, we can rewrite Equation
as:
x
E
f,(1:W,
IES 24
= Σ
c1(1; () X @ ......00 ........ .',
LESt
X;
CI; W) ... Xigt
LES
X O
(L-11)
W 24
X;
X
(14)
4.4
PE
where S2t-11 is a set that represents all the combinations of choosing 2 elements from [ri] (with replacement), i.e., S2c-1 tv [n]=*-1].
As a result, for any d2gree-(24-1) monomial M formed by rows in X, we can denote it by X X,, X ) which corresponds specific oredered index set I' $2 in Equation 14 Therefore, we can always parameterize weight matrices W E [Z] such that only the ele- ments with indices 7' that determine the given monomial are Equation 14 and all other elements in W()) are set to O. Consequently, Equation 14 with parameterized (1() becomes X, @ M2"-1 which complements our proof for Theorem 3.3
Discussion. Notably, Equation | reveals that the weight matrix W controls the coefficients of the monomial terms of degree 24 in the learned polynomial. Thus, if we replace W with any attention matrix S (e.g.. the global self-attention matrix described in Section 3.2) then the attention scores in S naturally control all the monomial coefficients, which essentially capture the importance of different (global) node feature interactions with order 2 In Appendix B are going provide the analysis to demonstrate that the matrix B in Equation controls the coefficients lower-degree monomials in practice.
B ANALYSIS ON B IN EQUATION
By ignoring the weight matrix B in Equation 4 have demonstrated in Appendix A that the base model learns a polynomial that encompasses all possible monomials of degree 2 However, our proof in Appendix [ also reveals that the learned polynomial is unable to represent any monomials with degrees smaller than 24, This raises concerns since these lower-degree monomial terms are often crucial for the model predictions (Hua al.] (2022). Fortunately, we are going show that incorporating B into Equation eenables the L-layer base model to capture all monomial terms with degrees up 2 (except the constant term).
Specifically, we can consider each monomial term of the learned polynomial as generated in the following way: for each layer of the base model, select two monomials, denoted as M0()) and
16

---

Published as a conference paper at ICLR 2024
M from the terms (W() (())) and (X(I-1)++*' in Equation 4 respectively. Consequently, the generated monomial after layer ( is obtained by the Hadamard product, i.e., M = _ OM
We continue this process recursively until we obtain (1)) at the L-th layer.
Regarding the selection of monomial M2 from (X (1-1) 1) + B two scenarios arise: (i) If M2
monomial term from the polynomial ( it increases the degree of the generated monomial
M by up to the degree of (1-1) (i.e., 2/-1, as analyzed in Appendix A). (ii) If M = B does not increase the degree.
If we never choose M B for all E [L]. then the generated monomial at layer L has degree of 24, as proven in Appendix A Hence, each time we select 2 = B()) at layer 1, we reduce the degree of monomial at the L-th layer (whose original degree is 2 by 2'-1. In other words, the more we opt for M ( the smaller the degree of the monomial the L-th layer, and vice versa. In the extreme scenario where we choose M, = B for all / E [7]] we obtain a monomial 2 2D 1. Consequently, the weight matrices _ E [4]. enable the base model to capture all possible monomials with degrees ranging from up to 24.
of degree 2 - 24-1
Discussion. The above analysis demonstrates the importance of B in Equation i.c.. controlling the coefficients of lower-degree monomials, which is another key distinction of our approach to previous gating-based GNNs that ignore B in their gating units (Yan et al. 2022 Rusch et al. 2022 Song et al. 2023
C POLYNOMIAL EXPRESSIVITY OF PRIOR GRAPH MODELS
Polynomial expressivity of GTs. For the sake of simplicity, we omit the so softrar operator in GTs and assume each node feature contains a scalar value, i.e., æ E IR" for an n-node graph. Besides, we denote the weights for query, key, and value by ug. Wk, and wp respectively. Consequently, for any node i, GT layer produces æ¹ = Wyy.W. W. E3 a,,r,,0 wy w.W.wWya. Σ; 2²., which consists of 3 monomials. However, there For instance, any monomial in the set {æ:2x, i i j} cannot be captured by x/e resulting in the limited polynomial expressivity of prior GT models.
still many degree-3 monomials that GTs fail to capture.
Polynomial expressivity of high-order GNNs. We initially focus on a closely related model known 1GNN, which is recently proposed by Hua al. (2022). In tGNN, the polynomial for each target node is computed as [x), 1] @ ... 2 ai 1JW, where [æ:;, 1] represents the concatenation of the feature vector of node with a constant
As asserted by [uua et al. (2022), (GNN learns multilinear polynomial in which no variables appear with a power of 2 or higher. Consequently,
IGNN cannot represent any monomials where specific variable has a power of at least 2, limiting
its polynomial expressivity.
Furthermore, there exist multiple GNN models based on gating mechanisms that implicitly capture high-order interactions among node features (e.g., GGCN Yan et al. (022)). G2-GNN (Rusch et al., 2022), and OrderedGNN (Song et al., (023)). However, these models primarily consider local structures and, therefore, are unable to learn high-order monomials formed by features of distant nodes. Additionally, few gating-based GNNs considers leveraging the weight matrix B in Equation 4 to learn lower-degree monomial terms, as discussed in Appendix B
Conclusion. Due to the limitations of polynomial expressivity analyzed above, previous GTs and high-order GNNs still necessitate nonlinear activation functions attain reasonable accuracy. In contrast, our method achieves higher accuracy than these models even without activation functions, thanks to its high polynomial expressivity.
D COMPARISON TO KERNEL-BAASER LINEAR ATTENTION MODELS
D.1 HYPERPARAMETER TUNING
In literature, several linear GTs employ kerne-bbased methods to approximate the dense self- attention matrix within the vanilla Transformer architecture. Concretely, both GraphGPS (with
17

---

Published as a conference paper at ICLR 2024
Performer-base attention) (Rampasek al. 2022) and NodeFormer (Wu al. 2022) utilize kernel function based on positive random features (PRFs). However, PRFs introduce critical hyperparameter. denoted as m, which determines the dimension of the transformed features. In addition to this, NodeFormer introduces another crucial hyperparameter, т, controlling the tempera- ture in Gumbel-Softmax, which is integrated into the kernel function. Furthermore, the DiFFormer model (Wu et al. )023) offers a choice between two types of kernel functions in practice.
Consequently, these kernel-based GTs require extensive tuning of these critical hyperparameters, which can be time-consuming, particularly when dealing with large graphs. In contrast, our pro- posed linear global attention, as defined Equation 6 eliminates the need for any hyperparameters in model tuning.
D.2 TRAINING STABILITY
To the best of our knowledge, the work most similar Equation i cosFormer (Qin et al. 2022). which primarily focuses text data rather than graphs. The key distinction between cosFormer and our proposed linear global attention lies in their choice kernel function. CosFormer utilizes the ReLU function instead of the Sigmoid function in Equation While ReLU performs well on text data containing to 16K tokens per sequence, we observe that it leads to training instability when applied to large graphs with millions of nodes.
Specifically, we substitute ReLU into Equation 6 the denominator term E; (K?.) becomes E, (e.. This modification results in training loss of NaN during our experiments on both the ogbn-products and pokec datasets. The issue arises because ReLU only sets negative values in K, to zero while preserving positive values. As consequence, the term Σ, ReLU (K?.) accumulates potentially large positive values across millions nodes. This leads to the denominator in Equation exceediin the representational capacity of 32-bit floating-poit format and results in a NaN loss during model training.
In contrast, approach employs the Sigmoid function, which maps all elements in K,,; the range (0, 1). As result, (KKT) does not produce excessively large values, avoiding the issue of NaN loss.
E DATASET DETAILS
Table Statisties of datasets used in our experiments.
Dataset
Type
Homophily Score
Nodes
Edges
Classes
Features
Computer
Homophily
0.700
13, 752
245, 261 861
10
767
Photo
Homophily
0.772
7, 650
119, 081
8
745
CS
Homophily
0.755
18, 333
894
15
6, 805
Physics
Homophily
0.847
34, 493
247,962
5
8, 8.415
WikiCS
Homophily
0.568
11,
216, 123
10
300
roman-empire
Heterophily
0.023
22, 662
927
18
300
amazon-ratings
Heterophily
0.127
24, 492
93,050
5
300
minesweeper
Hetcrophily
0.009
10, 000
39, 402
2
7
tolokers
Heterophily
0.187
11, 758
519,
2
questions
Heterophily
0.072
48, 921
153,540
2
301
ogbn-arxiv
Homophily
0.416
169, 343
166,6, 243
40
oghn-products
Homophily
0.459
2, 449, 029
61, 859, 140
47
pokec
Heterophily
0.000
1, 632 803
30, 622, 564
2
65
Table 4 shows the statistics of all 13 datasets used in our experiments. The homophily score per dataset computed based on the metric proposed by Lim al. (2021) (higher score means more homophilic).
Train/Valle Vasid/Test splits. For Computer, Photo, CS, and Physics datasets, adhere the widely accepted practice of randomly dividing nodes into training (60%), validation (20%), and test (20%) sets (Chen et al. 2022b (Shirzad et al. 2023). As for the remaining datasets in experiments,
18

---

Published as a conference paper at ICLR 2024
the official splits provided in their respective papers (eenyei & Cangea 2020 Hu al. 2020
Lim et al. 2021 Platonov et al. 2023)
F HARDWARE INFORMATION
We conduct all experiments on a Linux machine equipped with an Intel Xeon Gold 5218 CPU (featuring 8 cores @ 2.30 GHz) and 4 RTX A6000 GPUs (each with 48 GB of memory). It is worth noting that Polynormer only requires GPU for training, while the remaining GPUs are used to run baseline experiments in parallel.
G A VARIANT OF POLYNORMER
In the following, we provide the detailed architecture of the variant of Polynormer mentioned in Section 4.3 which built upon the local-and-global attention scheme, i.e., the local and global attention layers of Polynormer are employed in parallel to update node features.
X (AV + ( E. (K? (Q)))(())))))
2 OH + o
(15)
where matrices Q. K, V. H are obtained by linearly projecting the node feature matrix X from the previous layer. It is worth pointing out that this type of attention has been widely used by prior GT models (Rampasek al. 2022 Wu al. 2022 2023 Kong et al. 2023). However, our ablation study in Section 4.3 indicates that the local-and-globa attention performs worse than the local-to-global attention adopted by Polynormer.
H HYPERPARAMETERS SETTINGS
H.1 BASELINE MODELS
For the homophilic datasets listed in Table 1. we present the results of GCN (Kipt & Welling
2016)
GAT (Veličković et al.] (017). APPNP (asteigerr al., 2018). GPRGNN (Chien al.)
2020
PPRGo (Bojchevski et al. 2202), NAGphormer (Chen et al. (2022b), and Exphormer (Shirzad et al.
2023. as reported in Chen al.
2022b); Shirzad al. (2023).
For the heterophilic datasets in Table 2 we provide the results of GCN, GraphSAGE (Hamilton
al. 2017). GAT-sep, H2GCN (Zhu et al. 2022)) GPRGNN, FSGNN (Maurya al. GloGNN (Li et al., 2022), as reported in Platonov et al. (2023).
2022). and
In the case of large-scale datasets listed in Table we include the results of GCN, GAT, GPRGNN,
LINKX (Lim et al. 202T). and GOAT (Kong et al. 2023), as reported in Hu et al. (2020) Lim et al.
2021: Zhang et al. (2022) Kong et al. (2023).
For baseline models without publicly available results on given datasets, we obtain their highest achievable accuracy through tuning critical hyperparameters as follows:
GCNI (Chen et al., 2020). We set the hidden dimension to 512, the learning rate to 0.001, and the number of epochs to 2000. We perform hyperparameter tuning on the number of layers from (5, 10}, the dropout rate from {0.3, 0.5, 0.7), from {0.3, 0.5, 0.7), and 0 from {0.5, 1.0).
GGCN (Yan et al., 2022). We set the hidden dimension to 512, the learning rate to 0.001, and the number of epochs to 2000. We perform hyperparameter tuning on the number of layers from (5, 10], the dropout rate from (0.3, 0.5, 0.7), the decay rate 7 from (0.5, 1.0, 1.5), and the exponent from (2,3).
OrderedGNN (Song et al. 2023). We set the hidden dimension to 512, the learning rate to 0.001, and the number of epochs to 2000. We perform hyperparameter tuning the number of layers from (5, 10), the dropout rate from f0.3, 0.5, 0.7), and the chunk size from {4, 16, 64).
(GCN (Hua et al.) 2022). We set the hidden dimension 512, the learning rate to 0.001, and the number of epochs to 2000. We perform hyperparameter tuning on the number of layers from (2, (2,3), the dropout rate from (0.3, 0.5, 0.7), and the rank from (256, 512).
19

---

Published as a conference paper at ICLR 2024
G2-GNN (Rusch et al., 2022). We set the hidden dimension to 512, the learning rate to 0.001, and the number of epochs to 2000. We perform hyperparameter tuning on the number of layers from {5, 10}, the dropout rate from {0.3, 0.5, 0.7), and the exponent from 22,3.4)
DIR-GNN (Rossi et al., 2023). We set the hidden dimension to 512, the learning rate to 0.001, the number of epochs to 2000, and 0.5. Besides, we choose GATConv with the type of jumping knowledge as "max". We perform hyperparameter tuning on the number of layers from {3, 5). the dropout rate from {0.3, 0.5,0.7].
GraphGPS (Rampasek et al. 2022). We choose GAT as the MPNN layer type and Performer as the global attention layer type. We set the number of layers to the number of heads 8, the hidden dimension to 64, and the number of epochs to 2000. We perform hyperparameter tuming on 0.4, 0.5).
the learning rate from {le 1 5e 4, 1e- 3e- and the dropout rate from
NAGphormer (Chen et al. 2022b). We set the hidden dimension to 512, the learning rate to 0.001, the batch size to 2000, and the number of epochs to 500. We perform hyperparameter tuning on the number of layers from {1, 2, 3). the number of heads from {1, 8), the number of hops from {3, 7, 10), and the dropout rate from {0.0, 0.1, 0.2, 0.4,0.
Exphormer (Shirzad et al., 2023). We choose GAT the local model and Exphormer as the global model. set the number epochs to 2000 and the number of heads to We perform hyperparameter tuning on the learning rate from {le - 4, le 3). the number of layers from (2, 4), the hidden dimension form {64, 80, 96], and the dropout rate from {0.0, 0.1, 0.2, 0.4,
NodeFormer (Wu et al., 2022). We set the number of epochs to 2000. Additionally, we perform hyperparameter tuning on learning rate from {le 4, le 3, le - 2). the number of layers from {1, 2, 3). the hidden dimension from {32, 64, 128), the number of heads from {1, 4). M from {30, (3), K from (5, 10), rb order from {1. 2). ). the dropout rate from {0.0, 0.3). and the temperature from {0.10, 0.15, 0.20, 0.25, 0.30, 0.40, 0.50).
DIFFormer (Wu et al., 2023). We use the "simple" kernel. Moreover, we perform hyperparameter tuning on the learning rate from {le 4, le 3, le - 2). the number of epochs from (500, 2000), the number of layers from [2,3]. the hidden dimension form {64, 128), ). the number of heads from (1,8), or from {0.1, 0.2, 0.3). and the dropout rate from {0.0, 0.1, 2, 0.3, 0.4, 0.5).
GOAT Kong et al. 2023). We set the "conv _type" to "full", the number of layers to 1 (fixed by GOAT), the number of epochs to 200, the number of centroids 4096, the hidden dimension 256, the dropout of feed forward layers to 0.5, and the batch size to 1024. We perform hyperparameter tuning on the learning rate from {le 4, Ie le 2}. the global dimension from {128, 256), and the attention dropout rate from (0.0, 0.1, 0.2, 0.3, 0.4, 0.5).
H.2 POLYNORMER
Table 5: Hyperparameters of Polynormer per dataset.
Warm-up Epochs
Local-to-Glone Epochs
Local Layers
Global Layers
Dropout
Computer
200
1000
0.7
Photo
200
1000
7
2
0.7
CS
100
1500
5
2
0.3
Physies
100
1500
5
4
0.5
WikiCS
100
1000
7
2
0.5
roman-empire
100
2500
10
2
0.3
amazon-ratings
200
2500
10
1
0.3
minesweeper
100
2000
10
3
0.3
tolokers
100
800
2
0.5
questions
200
1500
5
3
0.2
ogbn-arxiv
2000
500
7
2
0.5
ogbn-products
1000
500
10
2
0.5
pokec
2000
500
7
0.2
20

---

Published as a conference paper at ICLR 2024
Like the baseline models, set the hidden dimension to 512 and the learning rate 0.001. Ad- ditionally, we introduce a warm-up stage dedicated to training the local module. This step ensures that the node representations generated by the local module capture meaningful graph structural information before being passed to the global module. Moreover, we leverage 8 heads for our at- tention modules. For the oghn-products and pokec datasets, size of 100, 000 and 550, 000 respectively, while for the other datasets, we employ full batch train- ing. Besides, disable the local attention module ogbn-arxiv and pokec datasets by replacing GAT (Veličković et al. 2017) with GCN (Kipf & Welling 2016), which we empirically observe perform better with lower memory usage. In Table we provide the critical hyperparameters of Polynormer used with each dataset.
mini-hatch training with hatch
Additional implementation details. We can rewrite Equation as Equation T6
x= H @ AV + t113") AV
(16)
To improve training stability, we adopt a LayerNorm on H AV. Besides, we empirically find that scaling LayerNorm( H AV) by 1 - o (1,33) typically improves model accuracy. Thus, we implement the local attention module as shown in Equation and the global attention module follows similar way:
X = (1 (1- (1-0((1 @ LayerNorm(H o AV) + t0(3) o AV
(17)
Appendix [ provides the basic Polynormer implementation. We refer readers to our source code for the detailed model implementation and hyperparameter settings.
I MODEL IMPLEMENTATION
N: the number of nodes
M: the number of edges
D: the node feature dimension
LI: the number of local attenticn layers
L2: the number of global attention layers
x: input node feature matrix with shape [N, D] edge_index: input graph structure with shape [2, M]
local_cenvs: local attention layers (e.g., GATCолv from. PyG)
local_betas: trainable weights with shape [D1, DI
global_convs: global attention layers (implemented Code 2)
global_betass trainable weights with shape [12, D)]
equivariant local attention module local
for i, local_come in enumerate :Local
h = h_lins (x) beta t.s.iooiilolllleeetteee.oo
unsqueeze (0)
leccl_cccvvlcccee (x, edge_index) (h beta)
local += x
equivariant global attention medule x_local
for i, giobal_cenv in enumerate (olesaa g = g_lins|il (x)
beta = F. (ggoo betasettlllll .ttaasii)..usqueeeee (0)
(g + beta)
output linear layer out [indlinnx))
negative log-likelihood loss calculation
y_pred F. sofftma (out, dim=1)
loss = c
y_true [train_idx]
Code 1: PyTorch-style Pseudocode for Polynormer
21

---

Published as a conference paper at ICLR 2024
N: the number of nodes
M: the number of edges
D: the node feature dimension
H: the number of attention heads
x: node feature matrix reshaped to IN, D]
linear projections to get query, key, value matrices g_lins [1] (x) # reshaped to [N, D/B, H]
k_lins [i] (x) # reshaped to [N, D/B, H]
v = v_lins [i] (x) # reshaped to [N, D/B, H]
map q and k to (0, for kernel approximation
q -..iggmei (q)
F.sigmeid(k)
numerater
kv torch. torch.einsum nmh, nch -> mdh', k,
num torch. einsum nmh, mah ndh', q kv) IN, D/H, H]
denominator
k sum = terch. einsuu (' ndh -> dh', k)
den torch. einsum ndh, dh nh', q, k_sum) unsqueeze (1) [N, 1, H]
linear global attention based on kernel trick
= num/den
reshape (N, D))
layerNorm (x)
Code 2: PyTorch-style Pseudocode for Linear Global Attention Layer
J SCALABILITY ANALYSIS OF POLYNORMER
J.1 PROFILING RESULTS ON SYNTHETIC GRAPHS
Polynomer Memory Usage Graph Size
Polynomer Training Time Graph Size
GPU Memory Usage (GB)
10
Training Time Epoch (8)
10'
CPU Trsining
GPU Training
10
10
10"
-0
10
10
10
10
10
# Nodes
# Nodes
Figure 5: GPU memory usage and training time of Polynormer on synthetic graphs
To validate the linear complexity of Polynormer with respect to the number of nodes/edges in graph, we generate synthetic Erdos-Renyi (ER) graphs with node count ranging from 1, 000 to 200, 000. We control the edge probability in the ER model to achieve average node degree of around and set the node feature dimension to 100. Figure3 clearly shows that the memory usage of Polynormer linearly increases with graph size, confirming its linear space complexity. As for the time complexity of Polynormer, we observe in Figure 5 (right) that the GPU training time remains nearly constant when increasing the graph size. We attribute it to that the GPU device (RTX A6000) used our experiments can efficiently handle parallel computations on graphs with different sizes for full batch training. This observation further supported by the CPU training
22

---

Published as a conference paper at ICLR 2024
time results, which linearly increase with graph size, the CPU is less powerful than the GPU for parallel computation.
J.2 PROFILING RESULTS ON LARGE REALISTIC GRAPHS
Table 6: Training time and GPU memory usage on large graphs We underline the results obtained
via full batch training, and highlight the top first, second, and third results per dataset.
ogbn-arxiy
ogbn-products
pokec
Method
Train/Epoch (s)
Mem. (GB)
Train/Epoch
Mem. (GB)
Train/Epoch (s)
Mem. (GB)
GAT
0.16
10.64
9.04
1.36
9.57
GraphGPS
1.32
38.91
OOM
OOM
OOM
OOM
NAGphormer
4.26
5.15
9.64
7.91
38.32
6.12
Exphormer
0.74
34.04
00M
00M
00M
00M
NodeFormer
1.20
16.30
3.37
31.55
6.15
17.21
DIFFormer
0.77
24.51
1.50
16.24
3,93
16.03
GOAT
12.32
6.98
20.11
9.64
51.93
8.73
Polynormer 0.31
16.09
3.13
12.93
2.64
21.54
Tahle/sshows the profiling results of GAT, linear GTs, and Polynormer terms of training time per epoch and memory usage on large realistic graphs. Note that we perform full batch training for all models whenever possible, since it avoids the nontrivial overhead associated with graph sampling in mini-batch training.
The results demonstrate that Polynormer consistently ranks among the top 3 fastest models, with relatively low memory usage. We attribute this efficiency advantage to the implementation of Poly- normer that only involves common and highly-optimized built-in APIs from modern graph learning frameworks (c.g., PyG). In contrast, the GT baselines incorporate less common expensive com- pute kernels (e.g., complicated kernel functions, Gumbel-Softmax, and nearest neighbor search), making them more challenging for graph learning frameworks accelerate. Thus, we believe Poly- normer is reasonably efficient and scalable in practice.
K WL EXPRESSIVITY OF POLYNORMER
While have shown that Polynormer is polynomial-exxressive, is also interesting see its expressivity under the Weisfeiler-Lehman (WL. hierarchy. To this end, we are going to firstly show that Polynormer is least as expressive as 1-WL GNNs, and then introduce simple way that renders Polynormer strictly more expressive than 1-WL GNNs.
For the sake of illustration, let revisit Equation in the following:
X (WX) (X + B)
(18)
where W E Bexn B E Raxd are trainable weight matrices, and X E Raxdd represents node features. Suppose all nodes have identical features (followed by the original formulation of WL algorithm), if we replace B = 1, the term shared across nodes. As result, the term (WX) @ (X + 113") is reduced to WX. X. By properly designing the weight matrix W (c.g. the adjacency matrix with self-loop), the term W X can be reduced to 1-WL GNNs. Hence, Polynormer based Equation [8 is at least expressive I-WL GNNs.
1,3 is essentially constant feature vector
To make Polynormer strictly more expressive than 1-WL GNNs, let focus on how properly de- sign the weight matrix B. Previously, we set = 1,3T. Note that the constant vector 1 E R" here essentially the eigenvector en that corresponds the smallest (normalized) graph Laplacian eigen- value. This motivates us to design B = v2 where vs denotes the eigenvector corresponding to the second smallest (i.e., first non-trivial) eigenvalue of the normalized Laplacian. Consequently, Polynormer can be built upon the following Equation:
X (WX) (X + vz
(19)
23

---

Published as a conference paper at ICLR 2024
Notably, the vector 22 essentially encodes the node positional information from graph spectrum into polynomial coefficients B, which makes each output node feature unique in Equation 19 and thus allows Polynormer to distinguish non-isomorphic graphs that 1-WL GNNs fail to distinguish, as empirically confirmed in Figure 6 It is noteworthy that this approach fundamentally different from prior PE methods based L.pllciann eigenvectors. Specifically, prior PE methods primarily focus on concatenating node positional encodings with node features. In contrast, we incorporate the positional information into the polynomial coefficients learned by Polynormer.
To distinguish the aforementioned two options of designing the weight matrix B, we denote the Polynormer of using the constant vector v1 and the second Laplacian vector U2 by Polynormer-vl and Polynormer-v2. respectively. In the following, we empirically show that Polynormer v2 is more powerful than 1 I-WI. GNNs on distinguishing non-isomorphic graphs.
Experimental setup. Without loss of generality, we focus on the local attention layer of Polynormer, and set as well as initial node features to be 1. As the attention scores between identical features are not meaningful, we replace the sparse attention matrix introduced in Section 3.2 with the random walk matrix Á AD where A and D denote the adjacency and degree matrices, respectively.
Experimental results. As shown in Figure 6 Polynormer-v2 able to distinguish non -somoophi graphs such circular skip link (CSL) graphs which are known to be indistinguishable by the I-WL test as well as I-WL GNNs Sato (020): Rampasek et al. 2002 Dwivedi et al. (2022).
0.1085
0.1085
0.1085
Graph Readout (sum-pooling)
0.6510
0.1085
0.1085
0.1085
Polynormer-v2
0.0480
2.6285
Graph Readout (sum-pooling)
0.3564
1,8700
2
2
7.5907
0.0489
0.000
1. 1.00
1.000
Graph Readout (sum-pooling)
1.500
1,500
(8,2)
7.000
1,000
Polynormer-v2
0.000
7.730
1,000
1,000
0.4225
0.4226
Graph Resdout (sum-pooling)
Gesr. (8,3)
8.000
1,000
1.000
0.4226
Figure 6: Two example pairs of non-isomorphic graphs from Sato 2020 that cannot distin- guished by 1-WL test. Our Polynormer-v2 can distinguish them.
While Polynormer-v2 is strictly more powerful than the 1-WL algorithm, our empirically results in Table 7 indicate that it does not clearly outperform Polynormer-vl I on realistic node classifica- tion datasets. One possible reason is that each node possesses unique features in realistic datasets, which diminishes the importance of incorporating additional node positional information into Poly- normer. Hence, we implement Polynormer based on Polynormer-vl in practice to avoid computing the Laplacian eigenvector, and leave the improvement of Polynormer-v2 our future work.
24

---

Published as a conference paper at ICLR 2024
Table 7: Comparison of Polynormer-vl and Polynormer-v2.
Computer
Photo
WikiCS
roman-empire
amazon-ratings
minesweeper
Polynormrr-vl
93.18 + 0.18
96.11 0.23
79.53 0.83
92.13 0.50
54.46 + 0.40
96.96 96 0.52
Polynormer-v2
93.25 0.22
96.09 + 0.30
80.01 = 0.55
91.33 0.39
54.41 0.35
96.48 + 0.37
L CASE STUDY: AI MODEL RUNTIME PREDICTION ON TPU
Apart from comparing Polynormer against conventional GRL models mainstream graph datasets, we further evaluate it on TpuGraphs (Phothilimthana et al. 2024), a large-scale runtime prediction dataset tensor computation graphs: github.com/googleeeesenne hea.re--essaacc-cccttttssp....e
Dataset details. Unlike prior datasets for program runtime prediction that are relatively small (up 100 nodes), TpuGraphs is a runtime prediction dataset full tensor programs, represented as computation graphs. Each graph within the dataset represents the primary computation of an ML program, typically encompassing one or more training steps or a single inference step. These graphs are obtained from open- goen-source ML programs and include well-known models such as ResNet, Ef- ficientNet, Mask R-CNN, and various Transformer models designed for diverse tasks like vision, natural language processing, speech, audio, recommender systems, and generative AI. Fach data sample in the dataset comprises a computational graph, a compilation configuration, and the corre- sponding execution time when the graph is compiled with the specified configuration an accelerator tailored for ML workloads. The compilation configuration governs how the XLA (Ac- celerated Linear Algebra) compiler transforms the graph through specific optimization passes. The TpuGraphs dataset is composed two main categories based on the compiler optimization level: (1) TpuGraphs-Laayn (layout optimization) and (2) TpuGraaphs-Tile (tiling optimization). Layout configurations dictate how tensors are organized in physical memory, determining the dimension order for each input and output of operation node. Notably, TpuGraphs-Laaye Layout has 4 collections based on ML model type and compiler configuration as follows:
TPU v3,
ML model type:
NLP: computation graphs of BERT models, XLA: computation graphs of ML models from various domains such as vision, NLP, speech, audio, and recommendation.
Compiler configuration:
Default: configurations picked by XLA compiler's heuristic.
Random: randomly picked configurations.
On the other hand, tile configurations control the tile size of each fused subgraph. The TpuGraphs- Layout comprises 31 million pairs of graphs and configurations, with an average of over 7, 700 nodes per graph. In contrast, the TpuGraphs-TTin includes 13 million pairs of kernels and configurations, averaging 40 nodes per kernel subgraph.
L.1 RESULTS ON SETGGRPHS--LYOOUE COLLECTIONS.
compare Polynormer against the GraphSAGE baseline provided by Google Hamilton HamiiooHamitonnn et al. (2017). To ensure a fair comparison. we only change the model architecture while keeping all other configurations the same (i.e., no additional feature engineering). Experimen- tal results show that wwwnoommerooolooorrrleoooo outperforms GraphSAGE on all collections. In particular, Polynormer achieves 10.3% accuracy improvement over GraphSAGE and 6.3% accuracy improvement over Polynormer-local on the XLA-Random collection, which is the most challeng- ing collection since XLA consists of diverse graph structures from different domains and Random contains various compiler configurations. This showcases that the global attention of Polynormer effectively captures critical global structures in different graphs, owing to its high polynomial ex- pressivity.
As shown in Table
L.2 RESULTS ON TPUGRAPHS (LAYOUT TILE).
We further evaluate Polynormer on the whole TpuGraphs dataset that contains Layout and Tile opti- mizations. Table shows that the local attention module alone in Polynormer is able to outperform
25

---

Published as a conference paper at ICLR 2024
Table 8: Ordered pair accuracy (%) on TpuGraphs-Layout collections notes the local attention module in Polynormer.
"Polynormer-Local" de-
GraphSAGE
Polynormer-Local
Polynormer
NLP-Default 81.2
82.0
82.1
NLP-Random 93.0
94.1
93.6
XLA-Default 77.3
73.7
79.6
XLA-Random 76.3
81.3
87.6
Table 9: Averaged runtime prediction accuracy (%) on TpuGraphs "Polynormer-Local" denotes the local attention module in Polynormer.
GraphSAGE
Polynormer-Local
Polynormer
TpuGraphs (Layout+Tile)
30.1
49.8
61.9
the baseline model by 19.6%, indicating the efficacy of polynomial-expressive architecture for capturing critical local structures. Moreover, the accuracy of Polynormer gets further improved when adding the global attention module, leading 31.8% accuracy improvement over the base- line model. This confirms the effectiveness of the proposed local-to-global attention scheme.
M DISCUSSION ON THE DEFINITION OF POLYNOMIAL EXPRESSIVITY
As graph learning models primarily learn node representations by aggregating input features from a set of nodes, they can be viewed multiset pooling problem with auxiliary information about the graph topology (Baek et al., 2021 Hua et al.) [022). In the context of polynomial networks on graphs, we follow a similar notion by focusing on polynomial functions that map input node features into output node representations, with the graph topology encoded into polynomial coefficients. Notably, such polynomials have been explored by prior studies Hua et al. (2022); Chrysos et al. (2000) y WWang et al. (2021), which inspire assessing the ability to learn high-degree polynomials is well motivated based on the Weierstrass theorem, which guarantees any smooth function be approximated by polynomial Stone (948). Thus, models with higher polynomial expressivity are more capable of leamning complex functions.
formally provide Definition 3.1 Morcover,
While Definition 3.1 does not explicitly mention graph topology, the utilization of graph topology in different graph learning models affects the inclusion of distinct monomial terms in the learned polynomial function, thereby impacting the polynomial expressivity. It is also worth noting that Definition 3.] can be viewed as a simplified version of the definition on polynomial expressivity in Puny et (2023), since we consider graph structural information polynomial coefficients rather than indeterminates. This simplification allows to decouple the analysis on node features from graph topology, which renders us develop a scalable model with high polynomial expressivity.
26


Reviews:
[{"review_id": "sGKd3olUqm", "reviewer": "Reviewer_HGJh", "summary": "Summary:  \nThis paper proposes Polynormer, a graph transformer model architecture for node classification. \nFirst, the paper introduces a base attention model in Section 3.1 that explicitly represents node features as polynomials, with coefficients determined by attention scores. This is claimed to result in high polynomial expressivity (in Sec. 3.1). To make the model equivariant, the paper integrates graph topology and node features into the polynomial coefficients to derive local and global attention models. This makes the overall Polynormer architecturem which achieves linear complexity instead of quadratic. Experiments are performed on 13 datasets including both homophilic and heterephilic tasks where Polynormer improves on most datasets.", "questions": "included with weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Strengths:\n- Provides theoretical analysis of polynomial expressivity, though restricted to scalar features. It goes beyond the WL expressivity as common in graph learning literature.\n- Demonstrates the performance of the architecture on 13 datasets where comparisons with baselines make the proposed model better on 11 datasets.\n- Ablation study on a smaller dataset group shows benefits of global attention and local-to-global scheme.", "weaknesses": "Weaknesses and Questions:  \n-The theoretical expressivity claims in Section 3.1 may be overclaiming capabilities, as proofs make simplifying assumptions about scalar features that differ from real graph data (Section 4). Can this be justified further?   \n-While complexity is analyzed, runtime and memory usage are not empirically compared to baselines in Section 4.2 to demonstrate scalability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "1WULbqKMYl", "reviewer": "Reviewer_D5Yt", "summary": "The proposed Polynormer is a new graph transformer model that balances expressivity and scalability. The paper introduces a polynomial model that achieves polynomial expressiveness with linear complexity, outperforming state-of-the-art GNN and GT baselines on most datasets. The model is based on a novel polynomial attention mechanism that can capture higher-order interactions between nodes in a graph. The attention mechanism is designed to be both local and global equivariant, allowing it to capture both local and global patterns in the graph.", "questions": "- How is the degree of polynomial defined?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea to adopt attention model in the polynomial feature mapping is novel and interesting.\n- Experiments are sufficient. Many important baselines and datasets of various sizes are covered.", "weaknesses": "- I find that the proposed approach (global) may also work in transformers in other fields, e.g. NLP. Could you provide such experiments to show its capacity in dealing with different types of data?\n- Why and how could polynomial expressivity improve model performance? The point was not clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "aHqtRgPH1t", "reviewer": "Reviewer_uwZk", "summary": "This paper proposes Polynormer, a polynomial- expressive GT model with linear complexity. Polynormer is built upon a novel base model that learns a high-degree polynomial on input features, with model permutation equivariance. Polynormer has been evaluated on 13 homophilic and heterophilic datasets, including large graphs with millions of nodes, with results showing that Polynormer outperforms state-of-the-art GNN and GT baselines on most datasets.", "questions": "- See weakness above.\n\n- Can authors discuss more about the relationship between WL-test expressivity and polynomial expressivity? For example, is the proposed Polynormer strictly more powerful than 1-WL-GNNs?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The idea of designing a polynomial-expressive graph Transformer model is novel and interesting. \n\n- The resulting Polynormer model is powerful, efficient, and theoretically expressive.\n\n- The experiments are convincing, showing that Polynormer can outperform sota GNNs and GTs on a wide range of datasets.", "weaknesses": "- It is inappropriate to claim that GTs and GNNs has limited polynomial expressivity (in section 3.1 and appendix C), since the non-linearity layers are not negligible. In [1] it is shown that without softmax GTs cannot represent GNNs. And in [2] Transformers are proved to be universal approximators on sequences with the softmax layer as key component. Can you discuss the polynomial expressivity of GTs and GNNs with non-linearity layers? And since [2] proves that Transformers are universal approximators, do GTs have $\\infty$-polynomial expressivity?\n\n- The concept of graph is defined by edge connections. And the definition of polynomial expressivity is completely ignorant of graph structure, comparing to WL-test expressivity. From my opinion, polynomial expressivity defined here should be used to model expressivity on sets, not graphs. What is the motivation of modeling the polynomial expressivity of graph models?\n\n- The O(N+E) complexity claim should be supported by more experiments, like a training time (VRAM) – graph size plot on synthetic random graphs with different sizes.\n\n[1] Ying, Chengxuan, et al. \"Do transformers really perform badly for graph representation?.\" Advances in Neural Information Processing Systems 34 (2021): 28877-28888.\n\n[2] Yun, Chulhee, et al. \"Are transformers universal approximators of sequence-to-sequence functions?.\" arXiv preprint arXiv:1912.10077 (2019).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (poster)