\section{Introduction}

In many domains, data can be represented as a graph, where entities interact, have meaningful relations and a global structure. The need to be able to infer and gain a better understanding of such data rises in many instances such as social networks, citations and collaborations, chemoinformatics, epidemiology etc. In recent years, along with the major evolution of artificial neural networks, graph learning has also gained a new powerful tool - graph neural networks (GNNs). Since first originated \citep{Gori2005, Scarselli2009} as recurrent algorithms, GNNs have become a central interest and the main tool in graph learning. 

Perhaps the most commonly used family of GNNs are message-passing neural networks \citep{Gilmer2017}, built by aggregating messages from local neighborhoods at each layer. Since information is only kept at the vertices and propagated via the edges, these models' complexity scales linearly with $|V|+|E|$, where $|V|$ and $|E|$ are the number of vertices and edges in the graph, respectively. In a recent analysis of the expressive power of such models, \citep{xu2018how,morris2018weisfeiler} have shown that message-passing neural networks are at most as powerful as the first Weisfeiler-Lehman (WL) test, also known as vertex coloring. The $k$-WL tests, are a hierarchy of increasing power and complexity algorithms aimed at solving graph isomorphism. This bound on the expressive power of GNNs led to the design of new architectures \citep{morris2018weisfeiler,maron2019provably} mimicking higher orders of the $k$-WL family, resulting in more powerful, yet complex, models that scale super-linearly in $|V|+|E|$, hindering their usage for larger graphs.  


Although expressive power bounds on GNNs exist, empirically in many datasets, GNNs are able to fit the train data well. This indicates that the expressive power of these models might not be the main roadblock to a successful generalization. Therefore, we focus our efforts in this paper on strengthening GNNs from a \emph{generalization} point of view. 
Towards improving the generalization of GNNs we propose the Low-Rank Global Attention (LRGA) module which can be augmented to any GNN. Standard dot-product global attention modules \citep{vaswani2017attention} apply $|V|\times|V|$ attention matrix to node data with $O(|V|^3)$ computational complexity making them impractical for large graphs. To overcome this barrier, we define a $\kappa$-rank attention matrix, where $\kappa$ is a parameter, that requires $O(\kappa|V|)$ memory and can be applied in $O(\kappa^2|V|)$ computational complexity.

To theoretically justify LRGA we focus on a GNN model family possessing maximal expressiveness ({i.e.}, universal) but vary in the generalization properties of the family members. 
\citep{murphy2019relational,loukas2019graph, dasoulas2019coloring, loukas2020hard} showed that adding node identifiers to GNNs improves their expressiveness, often making them universal. In this work, we prove that even adding \emph{random} features to the network's input, as suggested in \citep{sato2020random}, a framework we call Random Graph Neural Network (RGNN), GNN models are universal in probability. 

The improved generalization properties of LRGA-augmented GNN models is then showcased for the RGNN framework, where we show that augmenting it with LRGA algorithmically aligns with the $2$-folklore WL (FWL) algorithm; $2$-FWL is a strictly more powerful graph isomorphism algorithm than vertex coloring (which bounds message passing GNNs). To do so, we adopt the notion of algorithmic alignment introduced in \citep{Xu2019algoalign}, stating that a neural network aligns with some algorithm if it can simulate it with simple modules, resulting in provable improved generalization. We opt to use monimials in the role of simple modules and prove the alignment using polynomial kernels. Lastly, we bound the sample complexity of the model when learning the $2$-FWL update rule. Although our bound is exponential in the graph size, it nevertheless implies that RGNN augmented with LRGA can provably learn the $2$-FWL step, when training each module independently with two-layer MLP.






We evaluate our model on a set of benchmark datasets including tasks of graph classification and regression, node labeling and link prediction from \citep{dwivedi2020benchmarking, Hu2020}. LRGA improves state of the art performance in most datasets, often with a significant margin. We further perform ablation study in the random features framework to support our theoretical propositions.
\section{Related Work}

\textbf{Attention mechanisms.} The first work to use an attention mechanism in deep learning was \citep{Bahdanau2015a} in the context of natural language processing. Ever since, attention has proven to be a powerful module, even becoming the only component in the transformer architecture \citep{vaswani2017attention}. Intuitively, attention provides an adaptive importance metric for interactions between pairs of elements, {e.g.}, words in a sentence, pixels in an image or nodes in a graph.
A natural drawback of classical attention models is the quadratic complexity generated by computing scores among pairs. Methods to reduce the computation complexity were introduced by \citep{lee2018set} which introduced the set-transformer and addressed the problem by inducing point methods used in sparse Gaussian processes. Linearized versions of attention were suggested by \citep{shen2018} factorizing the attention matrix and normalizing separate components and \citep{katharopoulos2020transformers} suggesting  using fixed feature maps for linearizing the attention module. Our suggested linearized attention differs from those as we do not use fixed feature functions and use different normalization.

\textbf{Attention in graph neural networks. } In the field of graph learning, most attention works \citep{Li2015,Velickovic2018,Abu-El-Haija2018,Bresson2017Gated, Lee2018} restrict learning the attention scores to the local neighborhoods of the nodes in the graph. Motivated by the fact that local aggregations cannot capture long range relations which may be important when node homophily does not hold, global aggregation in graphs using node embeddings have been suggested by \citep{You2019position,Pei2020}. In a way, LRGA combines both approaches, allowing global weighted aggregations via the simple structure of global attention.


\textbf{Generalization in graph neural networks.}
Although being a pillar stone of modern machine learning, the generalization capabilities of NN are still not very well understood, {e.g.}, see \citep{bartlett2017spectrallynormalized,golowich2019sizeindependent}. Due to the irregular structure of graph data and the weight sharing nature of GNN, investigating their generalizing capabilities poses an even greater challenge. Despite the nonstandard setting, few works were able to construct generalization bounds for GNN via \emph{VC dimension} \citep{article}, \emph{uniform stability} \citep{verma2019stability}, \emph{Rademacher Complexity} \citep{garg2020generalization} and \emph{Neural Tangent Kernel} \citep{du2019graph}.





\section{Preliminaries and notations} \label{sec:prelim}
 We denote a graph by $G=(V,E, {\bm{X}})$ where $V$ is the vertex set of size $|V|=n$, $E$ is the edge set, and adjacency ${\bm{A}}$. ${\bm{X}}=({\bm{x}}_1,\dots,{\bm{x}}_n)^T$ represents the input vertex features.  A vertex $v_i\in V$ carries an input feature vector ${\bm{x}}_i \in \mathbb R^{d_0}$; in turn, ${\bm{X}}^{l}\in \mathbb R^{n\times d_l}$ represents the output of the $l^{th}$ layer of a neural network. 
 
 A common form of evaluating GNNs is by their ability to distinguish different graphs, described by \emph{graph isomorphism} which is an equivalence relation between graphs. The isomorphism type tensor of a graph $G$ is a tensor ${\tens{Y}}\in\mathbb R^{n^2 \times d_\text{iso}}$ which holds the isomorphism types of all pairs $(i,j)\in [n] \times [n]$. Given a pair $(i,j)$, which represents either an edge or a node of graph $G$, ${\tens{Y}}_{i,j}$ summarizes all the information this pair carries in graph $G$. More precisely put, \emph{isomorphism type} is an equivalence relation defined by: $(i,j)$ and $(i',j')$ have the same isomorphism type \emph{iff} the following conditions hold: (i) $i=j$ $\iff$ $i'=j'$; (ii) ${\bm{x}}_i={\bm{x}}_{i'}$ and ${\bm{x}}_j={\bm{x}}_{j'}$; and (iii) $(i,j)\in E$ $\iff$ $(i',j')\in E$. One way to build an isomorphism type tensor for graph $G$ is ${\tens{Y}} = [{\bm{I}}, \mathbf{1}\otimes {\bm{X}}, {\bm{X}}\otimes \mathbf{1}, {\bm{A}}]$
, where brackets denote concatentation in the feature dimension, ${\bm{I}}$ is the identity matrix, $(\mathbf{1}\otimes {\bm{X}})_{i,j,:}={\bm{x}}_j$, and similarly (with a slight abuse of notation) $({\bm{X}}\otimes \mathbf{1})_{i,j,:}={\bm{x}}_i$. 


\section{Low-rank global attention (LRGA)}

\label{sec:model}
We propose the Low-Rank Global Attention (LRGA) module that can augment any graph neural network layer, denoted here generically as $\mathrm{GNN}$, in the following way:
\begin{align} \label{e:model}
    {\bm{X}}^{l+1} \leftarrow \left[{\bm{X}}^{l}, \mathrm{LRGA}({\bm{X}}^{l}), \mathrm{GNN}({\bm{X}}^{l}) \right]
\end{align}
where the brackets denote concatenation along the feature dimension. The LRGA module is defined for an input feature matrix ${\bm{X}}\in\mathbb R^{n\times d_{\text{in}}}$ via 
\begin{equation} \label{e:lowRankAtt}
        \mathrm{LRGA}({\bm{X}})=\brac{\frac{1}{\eta({\bm{X}})} m_1({\bm{X}}) \parr{m_2({\bm{X}})^T m_3({\bm{X}})}, \; m_4({\bm{X}}) }
\end{equation}
where $m_1,m_2,m_3,m_4: \mathbb R^{n\times d_{\text{in}}} \rightarrow \mathbb R^{n\times \kappa}$ are MLPs operating on the feature dimension, that is $m({\bm{X}})=\brac{m({\bm{x}}_1),\ldots,m({\bm{x}}_n)}^T$, and $\kappa\in \mathbb N_0$ is a parameter representing the \emph{rank} of the attention module. Lastly, $\eta$ is a normalization factor:
\begin{align}\label{e:etaNorm}
    \eta({\bm{X}}) =\frac{1}{n}\left(\mathbf{1}^T m_1({\bm{X}})\right)\left(m_2({\bm{X}})^T\mathbf{1}\right),
\end{align}
where $\mathbf{1}=(1,1,\ldots,1)^T\in\mathbb R^n$. 
The matrix $\eta({\bm{X}})^{-1}m_1({\bm{X}})m_2({\bm{X}})^T$ can be thought of as a $\kappa$-rank attention matrix that acts globally on the graph's node features.


\textbf{Computational complexity.}\label{s:prelim}
Standard attention models \citep{vaswani2017attention,Luong2015} require explicitly computing the attention score between all possible pairs in the set, meaning that its memory requirement and computational cost scales as $O(n^2)$. This makes global-attention seem impractical for large sets, or large graphs in our case.
We address the global attention computational challenge by working with bounded rank ({i.e.}, $\kappa$) attention matrices, and avoid the need to construct the attention matrix in memory by replacing the standard entry-wise normalization ($\mathrm{softmax}$ or $\tanh$) with a the global normalization $\eta$. In turn, the memory requirement of LRGA is $O(n\kappa)$, and using low rank matrix-vector multiplications LRGA allows applying global attention in $O(n\kappa^2)$ computation cost.  

\textbf{Permutation Equivariance.} A common demand from GNN architectures is to respect the graph representation symmetries, namely the ordering of nodes \citep{Maron2019}. As shown in \citep{lee2018set} the set attention module is permutation equivariant. The same matrix product structure of the LRGA makes this module also permutation equivariant.





\section{Theoretical Analysis} \label{sec: theo_analysis}
In this section we establish the theoretical underpinning for LRGA. Since we want to analyse the generalization power added by LRGA, we focus on a family of GNNs with unbounded expressive power \emph{in probability} (RGNN). Under this model we show the benefit of augmenting GNNs with LRGA in terms of improved generalization via the notion of \textit{algorithmic alignment} with a powerful graph isomorphism testing algorithm (2-FWL). 



\subsection{Random graph neural networks}

We analyse LRGA under the framework of Random Graph Neural Networks (RGNNs):
\begin{definition}[Random Graph Neural Network] Let ${\mathcal{D}}$ be a probability distribution of zero mean and variance $c$, and $G=(V,E,{\bm{X}})$ a graph. RGNN is a GNN variant with random input features sampled at every forward pass {i.e.}, the input to the network is $[ {\bm{X}}, {\bm{R}}]$ where ${\bm{R}}$ are i.i.d.~samples ${\bm{R}}\in\mathbb{R}^{n\times d} \sim {\mathcal{D}}$.
\end{definition}


RGNN, suggested by \cite{sato2020random}, has related variants \citep{loukas2020hard, loukas2019graph, murphy2019relational} that use node identifiers or distinctive features, which can be viewed as \emph{constant} random features, in order to break symmetry between isomorphic nodes. Such models are proven to be universal but lose their inherent equivariance due to arbitrary prescription of node identifiers. We choose to work in the seemingly more limited setting of RGNN, which allows the network to distinguish between different nodes but does not overfit specific identifiers. Our main claims regarding this framework is that RGNN is both universal in probability and equivariant in expectation. 




\begin{proposition}[Universal]\label{prop:universalGNNS}
RGNN can approximate an arbitrary continuous graph function given random features sampled from a bounded distribution ${\mathcal{D}}$.
\end{proposition}
Here approximation is in a probabilitic sense: Let $\Omega\subset \mathbb R^{n\times d_0}\times \mathbb R^{n^2}$ be a compact set of graphs, $[{\bm{X}},{\bm{A}}]\in \Omega$, where ${\bm{A}}\in\mathbb R^{n^2}$ is the adjacency matrix. Then, given a continuous graph function $f$ defined over $\Omega$ and arbitrary $\varepsilon,\delta>0$, there exist network parameters and $d$ so that $P(|\mathrm{GNN}([{\bm{X}}, {\bm{R}}])-f([{\bm{X}},{\bm{A}}])|<\varepsilon)>1-\delta$, for all graphs $[{\bm{X}},{\bm{A}}]\in\Omega$. 
Proposition \ref{prop:universalGNNS} holds for GNN variants with a global attribute block such as \citep{battaglia2018relational}. The proof is based on the idea that random features  allow the GNN to transfer the graph's connectivity information to the node features. Once all graph information is encapsulated at the nodes, we exploit the universality of set functions \citep{zaheer2017deep} to get universality. The full proof is in Appendix \ref{app:proof_uni}. To the best of our knowledge this is the first result proving universality under the random feature assumption. 
 
\begin{proposition}[Equivariant in expectation]
RGNN is permutation equivariant in expectation.
\end{proposition}
Changing the random features at each forward pass allows RGNN to preserve equivariance in expectation. Indeed, equivariance of GNN implies that $\mathrm{GNN}({\bm{P}}\cdot[{\bm{X}},{\bm{R}}])={\bm{P}}\cdot \mathrm{GNN}([{\bm{X}},{\bm{R}}])$, for any permutation matrix ${\bm{P}}$ and input $[{\bm{X}},{\bm{R}}]$. Taking the expectation of both sides w.r.t.~${\bm{R}}\sim {\mathcal{D}}$, noting that ${\bm{P}}{\bm{R}} \sim {\bm{R}}$ and using linearity of expectation we get equivariance in expectation.  









 


\subsection{RGNN augmented with LRGA aligns with 2-FWL}

In this section we will formulate our main theoretical result, Theorem \ref{thm:algoalign}, stating that augmenting RGNN with LRGA algorithmically aligns with a powerful graph isomorphism testing algorithm called $2$-Folklore Weisfeiler-Lehman (2-FWL) \citep{grohe2015pebble,grohe2017descriptive}. We will first introduce the notion of algorithmic alignment and the 2-FWL algorithm, then formulate our main theorem, and continue in the next section with a proof. 

\textbf{Algorithmic alignment.} The notion of \emph{algorithmic alignment} was introduced in \cite{Xu2019algoalign} and shown (both theoretically and practically) to lead to improved generalization. Intuitively, a neural network ${\mathcal{N}}$ is said to be aligned with an algorithm ${\mathcal{A}}$ if ${\mathcal{N}}$ can simulate ${\mathcal{A}}$ by a composition of modules, and each module is "simple", or learnable, {i.e.}, have bounded (hopefully low) sample complexity. Our definition of algorithmic alignment is a slightly stricter version: 
\begin{definition}[Monomial Algorithmic Alignment]
A neural network ${\mathcal{N}}$ aligns with algorithm ${\mathcal{A}}$ if ${\mathcal{N}}$ can simulate ${\mathcal{A}}$ by learning only monomial functions, {i.e.}, $f({\bm{x}})={\bm{x}}^{\bm{\alpha}}$, where ${\bm{x}}\in\mathbb R^d$, ${\bm{\alpha}}\in \mathbb N^d$, and ${\bm{x}}^{{\bm{\alpha}}} = x_1^{\alpha_1}\cdot \dots \cdot  x_{d}^{\alpha_{d}}$.
\end{definition}
To motivate this choice of monomials as "simple" functions we note that \citep{arora2019fine,Xu2019algoalign} show a sample complexity bound for even-power polynomials learned by (two-layer) MLPs and we extend it to general monomials in the following proposition proved in Appendix \ref{app:monomial_sample_complexity}:

\begin{proposition}
Let a two layer MLP trained with gradient descent be denoted as the learning algorithm ${\mathcal{A}}'$. The monomial $g({\bm{x}}) = {\bm{x}}^{{\bm{\alpha}}}$, ${\bm{x}}\in\mathbb R^d$, of degree $n$, $|{\bm{\alpha}}|\leq n$, is PAC learnable with ${\mathcal{A}}'$ with a sample complexity bound:
$${\mathcal{C}}_{{\mathcal{A}}'}(g,\epsilon,\delta)={\mathcal{O}}\parr{\frac{C_{n,d}+\log(1/\delta)}{\epsilon^2}},$$
$C_{n,d}=\parr{n^2+1}^{(n+1)/2}c_{n,d}$, $\varepsilon>0$ is the error parameter and $\delta\in(0,1)$ the failure probability.
\end{proposition}
The asymptotic behaviour of $c_{n,d}$ is out of the scope of this paper. 
Therefore, a monomial algorithmic alignment of ${\mathcal{N}}$ to ${\mathcal{A}}$ means (under the assumptions and sequential training method of Theorem 3.6 in \cite{Xu2019algoalign}) that ${\mathcal{A}}$ is learnable by ${\mathcal{N}}$.  

\textbf{2-Folklore Weisfeiler-Lehman (2-FWL) Algorithm.}
$2$-FWL is part of the $k$-WL hierarchy of polynomial-time (approximate) graph isomorphism iterative algorithms that recolor $k$-tuples of vertices at each step according to neighborhoods aggregation. Upon reaching a stable coloring, the algorithm terminates and if the histograms of colors of two graphs are not the same then the graphs are deemed not isomorphic. The $2$-FWL algorithm is equivalent to $3$-WL, strictly stronger than vertex coloring (2-WL) which bounds the expressive power of GNNs. 

\begin{wrapfigure}[7]{R}{0.3\textwidth}
\vspace{-10pt}
\includegraphics[scale=0.8]{2fwl.pdf}
\vspace{5pt}
\end{wrapfigure}
In more detail, let ${\tens{Y}}^0\in\mathbb R^{n^2\times d_\text{iso}}$ represent the isomorphism types of a given graph $G=(V,E,{\bm{X}})$, that is ${\tens{Y}}_{i,j}^0\in\mathbb R^{d_\text{iso}}$ represents the isomorphism type of the pair $(i,j)$.  The 2-FWL algorithm is initialized with ${\tens{Y}}^0$. 
Let ${\tens{Y}}^l\in\mathbb R^{n^2\times d_l}$ denote the coloring tensor after the $l^{th}$ update step. An update step in the algorithm aggregates information from the multiset of neighborhood colors for each pair. We represent the multiset of neighborhood colors of the tuple $(i,j)$ with a matrix ${\bm{Z}}^l_{(i,j)}\in\mathbb R^{n\times 2d_l}$. That is, any permutation of the rows of ${\bm{Z}}^l_{(i,j)}$ represent the same multiset. The rows of ${\bm{Z}}^l_{(i,j)}$, which represent the elements in the multiset, are ${\bm{z}}_k=[{\tens{Y}}^l_{i,k},{\tens{Y}}^l_{k,j}]\in \mathbb R^{2d_l}$, $k\in[n]$. See the inset for an illustration. 
The $2$-FWL update step of a pair $(i,j)$ from ${\tens{Y}}^l$ to ${\tens{Y}}^{l+1}$ concatenates the previous pair's color and an encoding of the multiset of neighborhoods colors:
\begin{align}\label{e:FWL_update}
  {\tens{Y}}^{l+1}_{i,j} &= \brac{{\tens{Y}}^{l}_{i,j}, \mathrm{ENC}\parr{ {\bm{Z}}_{(i,j)}^l}}
\end{align}
where $\mathrm{ENC}:\mathbb R^{n\times 2d_l}\rightarrow \mathbb R^{d_{\text{enc}}}$ is a multiset injective  map invariant to the row-order of its input.

\textbf{Main result.} Consider the 2-FWL update rule in \eqref{e:FWL_update} and let ${\bm{Y}}^{l+1} \in \mathbb R^{n^2}$ denote (arbitrary) single feature dimension pealed off ${\tens{Y}}^{l+1}\in\mathbb R^{n^2\times d_{l+1}}$; we call ${\bm{Y}}^{l+1}$ a single-head of the update rule. Then,
\begin{theorem} \label{thm:algoalign}
LRGA augmented RGNN algorithmically aligns with a single head 2-FWL update step.
\end{theorem}
A corollary of this theorem is: 
\begin{corollary}
Multi-head LRGA augmented RGNN algorithmically aligns with 2-FWL.
\end{corollary}
Multi-head LRGA is a module of the form $[{\bm{X}}^l, \mathrm{LRGA}_1({\bm{X}}^l),\ldots,\mathrm{LRGA}_k({\bm{X}}^l),\mathrm{GNN}({\bm{X}}^l)]$, which is an equivalent to multi-head self-attention. In practice, we found single-head LRGA to be on par performance-wise with multi-head LRGA and therefore we focus on the single-head version in the experimental section. 


\subsection{Proof of Theorem \ref{thm:algoalign}}

To prove Theorem \ref{thm:algoalign} we need to show RGNN augmented with LRGA can simulate one head of the 2-FWL update step using only monomials as learnable functions. We achieve that by the following steps: (i) introduce the notion of node factorization to encode $n\times n$ tensor data as node features; (ii) show that RGNN can approximate node factorization of the graph's isomorphism type tensor with a single GNN layer using learnable monomial functions; (iii) show that 2-FWL update step can be formulated using matrix multiplication of monomial functions; and (iv) show LRGA can approximate a single head 2-FWL update step using learnable monomials. 

\textbf{Part (i).} 
We start with the definition of node feature factorization:
\begin{definition}[Node factorization] \label{def:richfeat}
Let ${\tens{Y}}\in\mathbb R^{n^2\times d}$ be a tensor. ${\bm{X}}\in\mathbb R^{n\times D}$ is called \emph{node factorization} of ${\tens{Y}}$ if there exists a block structure ${\bm{X}}=\brac{{\bm{X}}^1,\ldots,{\bm{X}}^k}$ so that ${\tens{Y}} = \brac{{\bm{X}}^{s_1}({\bm{X}}^{t_1})^T,\ldots,{\bm{X}}^{s_d}({\bm{X}}^{t_d})^T}$, where $(s_1,t_1),\ldots,(s_d,t_d) \in [k] \times [k]$ are index pairs. 
\end{definition}
Note that for all $i,j\in[n]$ we have ${\tens{Y}}_{i,j} = \brac{\ip{{\bm{x}}^{s_1}_{i}, {\bm{x}}^{t_1}_{j}},\ldots,\ip{{\bm{x}}^{s_d}_{i}, {\bm{x}}^{t_d}_{j}}}\in\mathbb R^d$.
Lets illustrate the definition with an example. Let ${\bm{A}}\in \set{0,1}^{n\times n}$ be the adjacency matrix of some graph $G$, and for simplicity assume that there are no node features. Then, the isomorphism type tensor of $G$ is ${\tens{Y}}^0=\brac{{\bm{I}},{\bm{A}}}\in \mathbb R^{n^2 \times 2}$. One possible way of node factoring ${\tens{Y}}^0$ is using the SVD decomposition of the adjacency matrix ${\bm{A}}$. Note that node factorization is not unique. 

\textbf{Part (ii).}
\begin{proposition}\label{prop: rich_to_random}
RGNN with skip connection can approximate node factorization of the  isomorphism type tensor ${\tens{Y}}^0$.
\end{proposition}
\begin{proof}
We will prove the case of graph $G=(V,E)$, {i.e.}, with no vertex features; the general case can be found in Appendix \ref{app: extension_prop_4}. Let  ${\bm{R}}\in \mathbb{R}^{n\times d}$ be a random node features matrix sampled i.i.d.~from ${\mathcal{D}}$. 
A single layer of standard message passing can represent $\mathrm{GNN}({\bm{R}})=d^{-0.5}[{\bm{A}}{\bm{R}}, {\bm{R}}]$, which requires learning only first degree (linear) monomials in the GNN's learnable parts.  Furthermore, $\mathrm{GNN}({\bm{R}})$ is an approximate node factorization of ${\tens{Y}}^0$, since $d^{-1}[{\bm{R}}\mR^T,{\bm{A}}{\bm{R}}\mR^T]\approx [{\bm{I}}, {\bm{A}}]={\tens{Y}}^0$, where the approximation error $d^{-1}{\bm{R}}\mR^T\approx {\bm{I}}$ can be bounded using the result in Appendix \ref{app:proof_uni}.
\end{proof}

\textbf{Part (iii).}
As shown in \citep{maron2019provably} the encoding function $\mathrm{ENC}$ from the 2-FWL update rule (see \eqref{e:FWL_update}) can be expressed as follows (derivation can be found in Appendix \ref{app:multiset_enc}):
\begin{equation}\label{e:FWL_update_rule_matrix}
    {\tens{Y}}^{l+1} = \brac{{\tens{Y}} , \brac{{\tens{Y}}^{{\bm{\beta}}}{\tens{Y}}^{{\bm{\gamma}}} \ \big\vert \ ({\bm{\beta}},{\bm{\gamma}})\in \mathbb N_0^{2d}, \abs{{\bm{\beta}}} + \abs{{\bm{\gamma}}} \leq n} }
\end{equation} \\
where for notational simplicity we denote ${\tens{Y}}={\tens{Y}}^l$ and $d=d_l$. By ${\tens{Y}}^{\bm{\beta}}$ we mean that we apply the multi-power ${\bm{\beta}}$ to the feature dimension, {i.e.}, $({\tens{Y}}^{\bm{\beta}})_{i,j}={\tens{Y}}_{i,j}^{\bm{\beta}}$. Therefore, computing the multisets encoding amounts to calculating monomials ${\tens{Y}}^{\bm{\beta}},{\tens{Y}}^{\bm{\gamma}}$ and their matrix multiplications ${\tens{Y}}^{\bm{\beta}} {\tens{Y}}^{\bm{\gamma}}$. 

\textbf{Part (iv).}
\begin{proposition} \label{prop:poly_kernels}
The node factorization of each head of ${\tens{Y}}^{l+1}$, the result of $2$-FWL update step, can be approximated via LRGA module applied to  node factorization of ${\tens{Y}}={\tens{Y}}^l$. The MLPs in the LRGA approximation need to learn only monomial functions. 
\end{proposition}
\textit{Proof.}
Let ${\bm{X}} = [{\bm{X}}^1,\dots, {\bm{X}}^k]\in \mathbb R^{n\times D}$ be a node factorization of ${\tens{Y}}={\tens{Y}}^l$. The $2$-FWL update step requires computation of polynomials of the form ${\tens{Y}}^{\bm{\beta}}$ as shown in \eqref{e:FWL_update_rule_matrix}. Using the node factorization of  ${\tens{Y}}$, ${\tens{Y}}_{i,j} = \brac{\ip{{\bm{x}}^{s_1}_{i}, {\bm{x}}^{t_1}_{j}},\ldots,\ip{{\bm{x}}^{s_d}_{i}, {\bm{x}}^{t_d}_{j}}}\in\mathbb R^d$, we can write:
\vspace{-5pt}
\begin{align} \nonumber 
    {\tens{Y}}_{i,j}^{\bm{\beta}} &= \prod_{l=1}^d \ip{{\bm{x}}_{i}^{s_l},{\bm{x}}_{j}^{t_l}}^{\beta_l}=\prod_{l=1}^d \ip{\varphi_{\beta_l}({\bm{x}}_{i}^{s_l}),\varphi_{\beta_l}({\bm{x}}_{j}^{t_l})}=
    \prod_{l=1}^d \ip{\varphi_{\beta_l}({\bm{x}}_{i}^{s}),\varphi_{\beta_l}({\bm{x}}_{j}^{t})} \\ \label{e:phi_beta}
    &= \ip{\varphi_{{\bm{\beta}}}({\bm{x}}_{i}^{s}),\varphi_{{\bm{\beta}}}({\bm{x}}_{j}^{t})}
\end{align}
where the second equality is using the feature maps $\varphi_{\beta_l}$ of the (homogeneous) polynomial kernels \citep{Vapnik1998}, $\ip{{\bm{x}}_1,{\bm{x}}_2}^{\beta_l}$; the third equality is reformulating the feature maps $\varphi_{\beta_l}$ on the vectors ${\bm{x}}_i^{s}=\brac{{\bm{x}}_i^{s_1},\ldots,{\bm{x}}_i^{s_d}}$, and ${\bm{x}}_i^{t}=\brac{{\bm{x}}_i^{t_1},\ldots,{\bm{x}}_i^{t_d}}$; and the last equality is due to the closure of kernels to multiplication. We denote the final feature map by $\varphi_{\bm{\beta}}$. 

Now, let $\psi_{\bm{\beta}}({\bm{x}}_i)=\varphi_{\bm{\beta}}({\bm{x}}_i^{s})$ and $\phi_{\bm{\beta}}({\bm{x}}_i)=\varphi_{\bm{\beta}}({\bm{x}}_i^{t})$ then we have:
$${\tens{Y}}^{\bm{\beta}} = \psi_{\bm{\beta}}({\bm{X}})\phi_{\bm{\beta}}({\bm{X}})^T,$$ where $\psi_{\bm{\beta}}({\bm{X}})$ is applying $\psi_{\bm{\beta}}$ to every row of ${\bm{X}}$.
Therefore, arbitrary head of ${\tens{Y}}^{l+1}$, {i.e.}, of the form ${\tens{Y}}^{\bm{\beta}}{\tens{Y}}^{\bm{\gamma}}$, can be written directly as a function of ${\bm{X}}$ using the feature maps $\phi_{\bm{\beta}},\psi_{\bm{\beta}},\phi_{\bm{\gamma}},\psi_{\bm{\gamma}}$:
\begin{equation}\label{e:Y_beta_gamma}
    {\tens{Y}}^{\bm{\beta}}{\tens{Y}}^{\bm{\gamma}} = \psi_{\bm{\beta}}({\bm{X}})\phi_{\bm{\beta}}({\bm{X}})^T \psi_{\bm{\gamma}}({\bm{X}})\phi_{\bm{\gamma}}({\bm{X}})^T.
\end{equation}
A node factorization of the head ${\tens{Y}}^{\bm{\beta}}{\tens{Y}}^{\bm{\alpha}}$ is therefore $\brac{\psi_{\bm{\beta}}({\bm{X}})\phi_{\bm{\beta}}({\bm{X}})^T\psi_{\bm{\gamma}}({\bm{X}}),\phi_{\bm{\gamma}}({\bm{X}})}$. Recalling the structure of the LRGA module introduced in \eqref{e:lowRankAtt}: $\mathrm{LRGA}({\bm{X}})=\brac{\eta({\bm{X}})^{-1} m_1({\bm{X}}) \parr{m_2({\bm{X}})^T m_3({\bm{X}})}, \; m_4({\bm{X}}) }$, to implement the 2-FWL head the MLPs $m_1,m_2,m_3,m_4$ need to learn the polynomial feature maps formulated in \eqref{e:Y_beta_gamma}: $m_1\approx \psi_{\bm{\beta}}$, $m_2\approx \phi_{\bm{\beta}}$, $m_3\approx \psi_{\bm{\gamma}}$, and $m_4\approx \phi_{\bm{\gamma}}$. Every coordinate of these feature maps is a monomial (proof of this fact in Appendix \ref{app:feat_map_monomial}). Lastly, 
note that 2-FWL tensors  ${\tens{Y}}^l$ are insensitive to global scaling and therefore the normalization $\eta$ has no theoretical influence (it is assumed non-zero). \qedsymbol





\section{Experiments}
We evaluated our method on various tasks including graph regression, graph classification, node classification and link prediction. The datasets we used are from two benchmarks: (i) benchmarking GNNs \citep{dwivedi2020benchmarking}; and (ii) Open Graph Benchmark (OGB) \citep{Hu2020}.  Each benchmark has its own evaluation protocol designed for a fair comparison among different models. These protocols define consistent splits of the data to train/val/test sets, set a budget on the size of the models (OGB), define a stopping criterion for reporting test results and require training with several different initializations to measure the stability of the results. We followed these protocols.\\
\textbf{Random Features Evaluation.} In addition, we also conducted a set of experiments with the random feature framework. In this experiment we focused on the PATTERN node classification dataset from \citep{dwivedi2020benchmarking} and evaluated a variety of models under the RGNN framework.\\
\textbf{Rank Ablation Study.} In this experiment we examined the relation between the rank parameter $\kappa$, which can limit the expressiveness of the attention module, and the network performance. Results are presented in Appendix \ref{sec: Rank_ablation}.\\
\textbf{Implementation details of LRGA.} We implemented the LRGA module according to the description in Section \ref{sec:model} (equations \ref{e:lowRankAtt}, \ref{e:etaNorm}) using the pytorch framework and the DGL \citep{wang2019dgl} and Pytorch geometric \citep{fey2019fast} libraries. Each LRGA module contains $4$ MLPs $m_1,m_2,m_3,m_4$. Each $m_i:\mathbb R^{d}\rightarrow \mathbb R^\kappa$ is a single layer MLP (linear with ReLU activation). The implementation of a layer is according to \eqref{e:lowRankAtt}, where in practice we added another single layer MLP,  $m_5:\mathbb R^{d+2\kappa+d_{GNN}}\rightarrow \mathbb R^{d}$, for the purpose of reducing the feature dimension size. In the OGB benchmark dataset we did not use the skip connections (better performance), and as advised in \citep{wang2019dgl}, we used batch and graph normalization at each layer. \\
\textbf{Baselines.} We compare performance with the following state of the art baselines: \textit{GCN} \citep{Kipf2016}, \textit{GraphSAGE} \citep{Hamilton2017}, \textit{GIN} \citep{xu2018how}, \textit{GAT} \citep{Velickovic2018}, \textit{GatedGCN} \citep{Bresson2017Gated}, \textit{Node2Vec} \citep{grover2016node2vec} , DeepWalk \citep{Perozzi_2014} and \textit{MATRIX FACTORIZATION} \citep{Hu2020}.


\begin{table}[h]
\captionsetup{font=,font=footnotesize}
\caption{Performance on the benchmarking GNN datasets. In bold: better performance between LRGA augmented and vanilla models; note the parameter (\#) budget. Blue represents best performance with the $100$K budget and red with the $500$K budget.}
\label{table: benchmarking}
\vspace{-6pt}
\tiny
\centering
\begin{tabular}{l?c|c?c|c?c|l?c|c?c|c?c|c}
\Xhline{2\arrayrulewidth}
\multirow{2}{*}{\textbf{Model}} & \multicolumn{2}{c?}{\textbf{PATTERN}}   & \multicolumn{2}{c?}{\textbf{CLUSTER}}   & \multicolumn{2}{c?}{\textbf{ZINC}}                        & \multicolumn{2}{c?}{\textbf{MNIST}}     & \multicolumn{2}{c?}{\textbf{CIFAR10}}   & \multicolumn{2}{c}{\textbf{TSP}}        \\ \cline{2-13} 
                                & \textbf{\#} & \textbf{Acc $\pm$ std}    & \textbf{\#} & \textbf{Acc $\pm$ std}    & \textbf{\#} & \multicolumn{1}{c?}{\textbf{MAE $\pm$ std}} & \textbf{\#} & \textbf{Acc $\pm$ std}    & \textbf{\#} & \textbf{Acc $\pm$ std}    & \textbf{\#} & \textbf{F1 $\pm$ std}      \\ \Xhline{2\arrayrulewidth}
\textbf{GCN}                    & 100K        & 63.88 $\pm$ 0.07          & 101K        & 53.44 $\pm$ 2.02          & 103K        & 0.459 $\pm$ 0.006                           & 101K        & 90.70 $\pm$ 0.21          & 101K        & 55.71 $\pm$ 0.38          & 95K         & 0.630 $\pm$ 0.001          \\
\textbf{LRGA + GCN}             & 90K         & \textbf{83.09 $\pm$ 0.73} & 91K         & \textbf{68.44 $\pm$ 0.16} & 92K         & \textbf{0.448 $\pm$ 0.009}                  & 91K         & \textbf{97.63 $\pm$ 0.11} & 91K         & \textbf{65.80 $\pm$ 0.43} & 97K         & \textbf{0.702 $\pm$ 0.001} \\ \hline
\textbf{GAT}                    & 109K        & 75.82 $\pm$ 1.82          & 110K        & 57.73 $\pm$ 0.32          & 102K        & 0.475 $\pm$ 0.007                           & 110K        & 95.53 $\pm$ 0.20          & 110K        & 64.22 $\pm$ 0.45          & 96K         & 0.671 $\pm$ 0.002          \\
\textbf{LRGA + GAT}             & 90K         & \textbf{82.54 $\pm$ 0.71} & 91K         & \textbf{69.05 $\pm$ 0.05} & 92K         & \textbf{0.421 $\pm$ 0.020}                  & 90K         & \textbf{97.47 $\pm$ 0.16} & 90K         & \textbf{68.00 $\pm$ 0.13} & 97K         & \textbf{0.680 $\pm$ 0.003} \\ \hline
\textbf{GatedGCN}               & 104K        & 84.48 $\pm$ 0.12          & 104K        & 60.40 $\pm$ 0.41          & 105K        & 0.375 $\pm$ 0.003                           & 104K        & 97.34 $\pm$ 0.14          & 104K        & 67.31 $\pm$ 0.31          & 97K         & \textbf{\textcolor{blue}{0.808 $\pm$ 0.003}} \\
\textbf{LRGA + GatedGCN}        & 93K         & \textbf{\textcolor{blue}{85.09 $\pm$ 0.11}} & 93K         & \textbf{\textcolor{blue}{69.28 $\pm$ 0.16}} & 94K         & \textbf{\textcolor{blue}{0.355 $\pm$ 0.010}}                  & 93K         & \textbf{\textcolor{blue}{98.20 $\pm$ 0.03}} & 93K         & \textbf{\textcolor{blue}{70.65 $\pm$ 0.18}} & 97K         & 0.807 $\pm$ 0.001          \\ \Xhline{2\arrayrulewidth}
\textbf{GCN}                    & 500K        & 71.89 $\pm$ 0.33          & 501K        & 68.49 $\pm$ 0.97          & 505K        & \textbf{0.367 $\pm$ 0.011}                  & 504K        & 91.39 $\pm$ 0.25          & 504K        & 54.84 $\pm$ 0.44          & -           & -                          \\
\textbf{LRGA + GCN}             & 400K        & \textbf{84.55 $\pm$ 0.57} & 400K        & \textbf{76.01 $\pm$ 0.67} & 501K        & 0.377 $\pm$ 0.009                           & 463K        & \textbf{98.34 $\pm$ 0.06} & 463K        & \textbf{68.27 $\pm$ 0.46} & -           & -                          \\ \hline
\textbf{GAT}                    & 526K        & 78.27 $\pm$ 0.18          & 528K        & 70.58 $\pm$ 0.44          & 531K        & 0.384 $\pm$ 0.007                           & 441K        & 96.50 $\pm$ 0.18           & 442K        & 66.11 $\pm$ 0.98          & -           & \textbf{-}                 \\
\textbf{LRGA + GAT}             & 533K        & \textbf{\textcolor{red}{85.82 $\pm$ 0.42}} & 267K        & \textbf{76.16 $\pm$ 0.34} & 536K        & \textbf{0.360 $\pm$ 0.004}                  & 476K        & \textbf{98.41 $\pm$ 0.08} & 476K        & \textbf{71.57 $\pm$ 0.26} & -           & -                          \\ \hline
\textbf{GatedGCN}               & 502K        & 85.56 $\pm$ 0.01          & 502K        & 73.84 $\pm$ 0.32          & 504K        & 0.282 $\pm$ 0.015                           & 500K        & 98.24 $\pm$ 0.04          & 500K        & 71.33 $\pm$ 0.39          & -           & -                          \\
\textbf{LRGA + GatedGCN}        & 486K        & \textbf{85.81 $\pm$ 0.31} & 438K        & \textbf{\textcolor{red}{76.39 $\pm$ 0.13}} & 446K        & \textbf{\textcolor{red}{0.249 $\pm$ 0.011}}                  & 486K        & \textbf{\textcolor{red}{98.47 $\pm$ 0.16}} & 487K        & \textbf{\textcolor{red}{73.48 $\pm$ 0.29}} & -           & -                         
\end{tabular}
\end{table}


\subsection{Benchmarking Graph Neural Networks \citep{dwivedi2020benchmarking}} \label{sec: benchmarking_gnn}
\textbf{Datasets.} 
This benchmark contains $6$ main datasets (full description in appendix \ref{app:dwivedi}) : (i) ZINC, graph regression task of molecular dataset evaluated with MAE metric; (ii) MNIST and CIFAR10, the image classification problem converted to graph classification using a super-pixel representation \citep{Knyazev_superpixel}; (iii) CLUSTER and PATTERN, node classification tasks which aim to classify embedded node structures \citep{abbe2017community}; (iv) TSP, a link prediction variation of the Traveling Salesman Problem \citep{joshi2019efficient} on 2D Euclidean graph. 
\\
\textbf{Evaluation protocol.} All models were evaluated with two different sets of parameter budgets and restrictions. The first set restricted to have roughly $100K$ parameters and $4$ layers, while the second set of experiments has a budget of roughly $500K$ parameters and up to $16$ layers. The learning rate and its decay are set according to a predetermined scheduler using the validation loss. The stopping criterion is set to when the learning rate reaches a specified threshold. All results are averaged over a set of predetermined fixed seeds and standard deviation is reported as well. 
\vspace{-5pt}

\textbf{Results.} Table \ref{table: benchmarking} summarizes the results of training and evaluating our model according to the evaluation protocol; We observe that LRGA improves GNN performance, often by a large margin, across all models and datasets, besides GCN on ZINC and GatedGCN in TSP, supporting our claim for improved generalization. We further note that SOTA in all datasets except TSP is achieved with LRGA augmented GNNs. In some datasets, such as CLUSTER and PATTERN, LRGA reaches top and roughly equivalent performance for all models it augmented, which emphasizes the empirical contribution of LRGA independently of the GNN variant.




\subsection{Link prediction datasets from the OGB benchmark \citep{Hu2020}}
\begin{wraptable}[9]{r}{0.55\textwidth}
\captionsetup{justification=centering,font=,font=footnotesize}
\vspace{-13pt}

\caption{ Performance on the link prediction tasks from the OGB benchmark}
\vspace{-6pt}
\label{table: ogb}
\begin{adjustbox}{width=0.55\textwidth}
\centering

\begin{tabular}{l?c|c?c|c?c|c}
\Xhline{2\arrayrulewidth}
\multirow{2}{*}{\textbf{Model}} & \multicolumn{2}{c?}{\textbf{ogbl-ppa}}         & \multicolumn{2}{c?}{\textbf{ogbl-collab}} & \multicolumn{2}{c}{\textbf{ogbl-ddi}}   \\ \cline{2-7} 
                                & \textbf{\# Param} & \textbf{Hits@100$\pm$std}  & \textbf{\# Param}  & \textbf{Hits@50$\pm$std}   & \textbf{\# Param} & \textbf{Hits@20$\pm$std}   \\ \Xhline{2\arrayrulewidth}
\textbf{Node2vec}               & 7.3M              & 0.223 $\pm$ 0.008          & 30M          & 0.489 $\pm$ 0.005          & 645K        & 0.233 $\pm$ 0.021          \\
\textbf{DeepWalk}               & 150M              & 0.289 $\pm$ 0.015          & 61M          & 0.504 $\pm$ 0.003          & 11M         & 0.264 $\pm$ 0.061          \\
\textbf{MF}                     & 147M              & 0.323 $\pm$ 0.009          & 60M          & 0.389 $\pm$ 0.003          & 1.2M        & 0.137 $\pm$ 0.047          \\ \Xhline{2\arrayrulewidth}
\textbf{GraphSage}              & 424K              & 0.165 $\pm$ 0.024          & 460K         & 0.481 $\pm$ 0.008          & 1.4M        & 0.539 $\pm$ 0.047          \\
\textbf{GCN}                    & 278K              & 0.187 $\pm$ 0.013          & 296K         & 0.447 $\pm$ 0.011          & 1.2M        & 0.370 $\pm$ 0.050          \\
\textbf{LRGA + GCN}             & 814K              & \textbf{0.342 $\pm$ 0.016} & 1M           & \textbf{0.522 $\pm$ 0.007} & 1.5M        & \textbf{0.623 $\pm$ 0.091}
\end{tabular}
\end{adjustbox}
\end{wraptable}



\textbf{Datasets.} We further evaluate LRGA on semi-supervised learning tasks including graphs with hundreds of thousands of nodes, from the OGB benchmark: (i) ogbl-ppa, a graph of proteins and biological connections as edges ;(ii) ogbl-collab, an authors collaborations graph; (iii) ogbl-ddi drug interaction network. The evaluation metric for all of the tasks is Hits@K; more details in appendix \ref{app:ogb}.  \\
\textbf{Evaluation protocol.}  All models have a hidden layer of size $256$ and the number of layers  is $3$ in ogbl-ppa and ogbl-collab and $2$ in ogbl-ddi. Test results are reported by the best validation epoch averaged over $10$ random seeds. \\
\textbf{Results.} Table \ref{table: ogb} summarizes the results on the link prediction tasks. It should be noted that the first three rows correspond to node embedding methods where the rest are GNNs. Augmenting GCN with LRGA achieves SOTA results on those datasets, while still using order of magnitude less parameters than the node embedding runner-up method.




\subsection{Random Features Ablation}
\begin{wraptable}[14]{r}{0.25\textwidth}
\vspace{-10pt}
\captionsetup{justification=centering,font=footnotesize}
\caption{Random Features Evaluation}
\vspace{-6pt}
\label{table: random_features}
\begin{adjustbox}{width=0.25\textwidth}
\begin{tabular}{l?c}

\Xhline{2\arrayrulewidth}
\multirow{2}{*}{\textbf{Model}} & \textbf{PATTERN}            \\ \cline{2-2} 
                                & \textbf{Acc $\pm$ std}      \\ \Xhline{2\arrayrulewidth}
\textbf{GCN}                    & 74.891 $\pm$ 0.713          \\
\textbf{LRGA + GCN}             & 84.118 $\pm$ 1.216 \\ \hline
\textbf{GAT}                    & 81.796 $\pm$ 0.661          \\
\textbf{LRGA + GAT}             & 85.905 $\pm$ 0.109 \\ \hline
\textbf{GraphSage}              & 85.039 $\pm$ 0.068                           \\
\textbf{LRGA + GraphSage}        & 85.229 $\pm$ 0.331                  \\ \hline
\textbf{GatedGCN}               & 85.848 $\pm$ 0.065          \\
\textbf{LRGA + GatedGCN}        & 85.944 $\pm$ 0.664 \\ \hline
\textbf{GIN}                    & 85.760 $\pm$ 0.001          \\
\textbf{LRGA + GIN}               & \textbf{86.765 $\pm$ 0.065}
\end{tabular}
\end{adjustbox}
\end{wraptable}

In this experiment we wanted to validate the theoretical analysis presented at section \ref{sec: theo_analysis}. The dataset for this evaluation is the PATTERN dataset, which is originally equipped with random features, but in contrast to the RGNN framework those features are sampled only once at the dataset creation stage. We evaluated the different models according to the RGNN framework, {i.e.}, resample the features with every forward pass. The features were sampled from a zero mean Gaussian distribution with variance $\frac{1}{d}$, where $d$ is the input feature dimension. The evaluation protocol is the same as the one used in section \ref{sec: benchmarking_gnn} and we followed the $500K$ budget. As seen from table \ref{table: random_features}, using alternating random features improves performance for all the models. \textit{GIN} and \textit{GraphSage} do not appear in the main table but according to \citep{dwivedi2020benchmarking} achieves $85.39\%$ and $50.49\%$ respectively. The LRGA augmented RGNN models maintain their superiority (even presenting a small improvement compared to Table \ref{table: benchmarking}) and serve as an empirical validation to our main theorem.



\vspace{-5pt}



\section{Conclusions}
\vspace{-5pt}
In this work, we set ourself in a path for improving the generalization power of GNNs. To do so, we introduced the LRGA module, a global self attention module, which is a variant of the dot-product self-attention with linear complexity. In order to theoretically evaluate the contribution of LRGA we analyzed our model under the RGNN framework, which is proved to be universal \emph{in probability}. Under this framework we were able to show that RGNN augmented with LRGA can align with the powerful 2-FWL isomorphism test by learning simple monomial functions, which have a known sample complexity bound. Under certain conditions the latter provides concrete generalization guarantees for RGNN augmented with LRGA. Empirically, we demonstrated augmenting GNN models with LRGA improves their performance significantly, often achieving SOTA performance.

\newpage








\section{Proof of Proposition \ref{prop:universalGNNS}}

\begin{proof}
Every graph function can be formulated as a function of the isomorphism type tensor ${\tens{Y}}\in\mathbb R^{n^2\times d}$, and we will approximate such arbitrary continuous functions with GNN.
Let $f:\mathbb R^{n^2\times d}\rightarrow \mathbb R$ be a continuous invariant graph function ({i.e.}, agnostic to ordering the graph nodes) defined over the isomorphism type tensors $\mathbb R^{n^2\times d}$. Define $\hat{f}({\bm{X}})=f({\tens{Y}})$, where ${\tens{Y}}$ is defined as in Definition \ref{def:richfeat}. $\hat{f}:K\rightarrow\mathbb R$ is an invariant set function since it is a composition of invariant and equivariant functions (see {e.g.}, \cite{Maron2019} for definition of equivariance); it is also continuous as a composition of continuous functions. Hence $\hat{f}$ can be approximated over $K$ using DeepSets \cite{zaheer2017deep} (due to DeepSets universality). Since the GNN in \cite{battaglia2018relational} includes DeepSets as a particular case it can approximate $\hat{f}$ as-well. 
\end{proof}

\section{2-FWL via polynomial kernels}
In this section, we give a full characterization of feature maps, $\varphi_{\bm{\beta}}$, of the final polynomial kernel we use to formulate the 2-FWL algorithm. 
A key tool for the derivation of the final feature map is the multinomial theorem, which we state here in a slightly different form to fit our setting. 

\textbf{Multinomial theorem. } Let us define a set of $m$ variables $x_1y_1,\dots,x_my_m$ composed of products of corresponding $x$ and $y$'s. Then,
\begin{align*} 
    (x_1y_1+\dots+x_my_m)^n = \sum_{\abs{{\bm{\nu}}}=n} {{n}\choose{{\bm{\nu}}}}\prod_{i=1}^m (x_iy_i)^{\nu_i}
\end{align*}
where ${\bm{\nu}}\in\mathbb N_0^m$, and the notation ${{n}\choose
{{\bm{\nu}}}} = \frac{n!}{\nu_1!\cdot\dots\cdot\nu_m!}$. The sum is over all possible ${\bm{\nu}}$ which sum to $n$, in total ${{n+m-1}\choose{m-1}}$ elements.

Recall that we wish to compute ${\tens{Y}}^{\bm{\beta}}_{i,j}$ as in \eqref{e:phi_beta} in the paper:
\begin{align*}
    {\tens{Y}}_{i,j}^{\bm{\beta}} &= \prod_{k=1}^d \ip{{\bm{x}}_{i}^{2k-1},{\bm{x}}_{j}^{2k}}^{\beta_k}=\prod_{k=1}^d \ip{\varphi_{\beta_k}({\bm{x}}_{i}^{2k-1}),\varphi_{\beta_k}({\bm{x}}_{j}^{2k})}=
    \prod_{k=1}^d \ip{\varphi_{\beta_k}({\bm{x}}_{i}^{\text{odd}}),\varphi_{\beta_k}({\bm{x}}_{j}^{\text{even}})} \\ 
    &= \ip{\varphi_{{\bm{\beta}}}({\bm{x}}_{i}^{\text{odd}}),\varphi_{{\bm{\beta}}}({\bm{x}}_{j}^{\text{even}})}
\end{align*}


We will now follow the equalities in \eqref{e:phi_beta} to derive the final feature map. The second equality is using the feature maps $\varphi_{\beta_k}$ of the (homogeneous) polynomial kernels \cite{Vapnik1998}, $\ip{{\bm{x}}_1,{\bm{x}}_2}^{\beta_k}$, which can be derived from the multinomial theorem.

Suppose the dimensions of ${\bm{X}}^{2k-1},{\bm{X}}^{2k}$ are $n\times D_k$ where $\sum_{k=1}^d 2D_k=D$. Then, $\varphi_{\beta_k}$ consists of monomials of degree $\beta_k$ of the form $\varphi_{\beta_k}({\bm{x}})_{\bm{\nu}} = \sqrt{{{\beta_k}\choose{{\bm{\nu}}}}}\prod_{i=1}^{D_k}x_i^{\nu_i}=\sqrt{{{\beta_k}\choose{{\bm{\nu}}}}} {\bm{x}}^{\bm{\nu}}$,  $\abs{{\bm{\nu}}}=\beta_k$. In total the size of the feaure map $\varphi_{\beta_k}$ is ${{\beta_k+D_k-1}\choose{D_k-1}}$.

The third equality is reformulating the feature maps $\varphi_{\beta_k}$ on the vectors ${\bm{x}}_i^{\text{odd}}=\brac{{\bm{x}}_i^1,{\bm{x}}_i^3,\ldots,{\bm{x}}_i^{2d-1}}\in\mathbb R^{D/2}$, and ${\bm{x}}_i^{\text{even}}=\brac{{\bm{x}}_i^2,{\bm{x}}_i^4,\ldots,{\bm{x}}_i^{2d}}\in\mathbb R^{D/2}$.


The last equality is due to the closure of kernels to multiplication. The final feature map, which is the product kernel, is composed of all possible products of elements of the feature maps, {i.e.},  $$\varphi_{\bm{\beta}}({\bm{x}}) = \parr{\prod_{k=1}^{d}\sqrt{{{\beta_k}\choose{{\bm{\nu}}_k}}}{\bm{x}}_k^{{\bm{\nu}}_k} \;\Big| \; \abs{{\bm{\nu}}_j}=\beta_j,\; \forall j\in[d]},$$
where ${\bm{x}}=\brac{{\bm{x}}_1,{\bm{x}}_2,\ldots,{\bm{x}}_d}\in\mathbb R^{D/2}$, and ${\bm{x}}_k\in\mathbb R^{D_k}$ for all $k\in[d]$. 
The size of the final feature map is $\prod_{k=1}^d{{\beta_k+D_k-1}\choose{D_k-1}}\leq N$ where $N={{n+D}\choose{D}}$.




\section{Bound on LRGA sample complexity}
\subsection{Proof of Lemma \ref{lem:poly}}
\begin{proof}
Using the multinomial theorem we have: $(\ip{{\bm{\beta}},{\bm{x}}}+1)^{n}=\sum_{{\bm{\alpha}}\in {\mathcal{B}}} d_{\bm{\alpha}} {\bm{\beta}}^{\bm{\alpha}} {\bm{x}}^{\bm{\alpha}}$, where $d_{\bm{\alpha}}$ are positive multinomial coefficients. This equation defines a linear relation between the monomial basis ${\bm{x}}^{\bm{\delta}}$ and $(\ip{{\bm{\beta}},{\bm{x}}}+1)^{n}$, for ${\bm{\beta}}\in{\mathcal{B}}$. The matrix of this system is ${\bm{V}}$ multiplied by a positive diagonal matrix with $d_{\bm{\alpha}}$ on its diagonal. By inverting this matrix and solving this system for ${\bm{x}}^{\bm{\delta}}$ the lemma is proved.  
\end{proof}

\subsection{Derivation of sample complexity bound}

Corollary 6.2 in \cite{arora2019fine} provides a bound on the sample complexity, denoted ${\mathcal{C}}_{\mathcal{A}}(g,\epsilon,\delta)$, of a polynomial  $g:\mathbb R^D\rightarrow\mathbb R$ of the form
\begin{equation}\label{e:g}
g({\bm{x}})=\sum_j a_j \ip{{\bm{\beta}}_j,{\bm{x}}}^{p_j},    
\end{equation}
where $p_j\in \set{1,2,4,6,8,\ldots}$, $a_j\in \mathbb R$, ${\bm{\beta}}_j\in\mathbb R^D$; $\epsilon,\delta$ are the relevant PAC learning constants, and ${\mathcal{A}}$ represents an over-parameterized, randomly initialized two-layer MLP trained with gradient descent:
$${\mathcal{C}}_{\mathcal{A}}(g,\epsilon,\delta)={\mathcal{O}}\parr{\frac{ \sum_j p_j \abs{a_j}\norm{{\bm{\beta}}_j}_2^{p_j} + \log(1/\delta)}{\epsilon^2}}$$

In our case $g:\mathbb R^{D+1}\rightarrow \mathbb R$ is defined as $g({\bm{x}},x_{D+1}) = \sum_{{\bm{\beta}}\in{\mathcal{B}}}a_{\bm{\beta}}\parr{\ip{{\bm{\beta}},{\bm{x}}}+x_{D+1}}^n$ where ${\mathcal{B}}=\set{{\bm{\beta}}\in \mathbb N_0^D \ \vert \ \abs{{\bm{\beta}}}\leq n}$ and by Lemma \ref{lem:poly} there exist ${\bm{a}}$ such that $g({\bm{x}},1)={\bm{x}}^{\bm{\delta}}$.
The sample complexity bound expression by Corollary 6.2 is therefore:
$${\mathcal{C}}_{\mathcal{A}}(g,\epsilon,\delta)={\mathcal{O}}\parr{\frac{ \sum_{{\bm{\beta}}\in{\mathcal{B}}} n \abs{a_{\bm{\beta}}}\norm{\hat{{\bm{\beta}}}}_2^{n} + \log(1/\delta)}{\epsilon^2}}, \quad \hat{{\bm{\beta}}}=({\bm{\beta}},1) $$


Let us bound the first term in the numerator of the sample complexity expression:

\begin{align*}
    \sum_{{\bm{\beta}}\in{\mathcal{B}}} n \abs{a_{\bm{\beta}}}\norm{\hat{{\bm{\beta}}}}_2^{n} = n\cdot\sum_{{\bm{\beta}}\in{\mathcal{B}}} \abs{a_{\bm{\beta}}}\parr{\sum_{i=1}^D\beta_i^2+1}^{n/2} \leq n\cdot \parr{n^2+1}^{n/2} \sum_{{\bm{\beta}}\in{\mathcal{B}}} \abs{a_{\bm{\beta}}}\leq \parr{n^2+1}^{(n+1)/2}c_{n,D}
\end{align*}
The first inequality is due to $\norm{\cdot}_2\leq \norm{\cdot}_1$, the second is by Lemma \ref{lem:poly} and uniting $n$ into the main term. From the above, the bound follows.




\section{Implementation Details}
In this section we describe the datasets on which we performed our evaluation. In addition, we specify the hyperparameters for the experiments section in the paper. The rest of the model configurations are determined directly by the evaluation protocols defined by the benchmarks. It is worth noting that most of our experiments ran on a single Tesla V-$100$ GPU, if not stated otherwise. We performed our parameter search only on $\kappa$ and $d$ (except for CIFAR10 and MNIST were we searched over different dropout values), since the rest of the parameters were dictated by the evaluation protocol. The models sizes were restricted by the allowed parameter budget.
\begin{table}[h]
\small
\centering
\caption{\label{table 1}Summary of the benchmarking GNN and ogb link prediction Datasets}
\vspace*{3mm}
\begin{tabular}{l|c|c|c|c|c}
\hline
Dataset                & \multicolumn{1}{l|}{\textbf{\#Graphs}} & \multicolumn{1}{l|}{\textbf{\#Nodes}} & \multicolumn{1}{l|}{\textbf{Avg. Nodes}} & \multicolumn{1}{l|}{\textbf{Avg. Edges}} & \multicolumn{1}{l}{\textbf{\#Classes}} \\ \hline
\textbf{ZINC}          & 12K                                    & 9-37                                  & 23.16                                    & 49.83                                    & -                                      \\
\textbf{CLUSTER}       & 12K                                    & 40-190                                & 117.20                                   & 4301.72                                  & 6                                      \\
\textbf{PATTERN}       & 14K                                    & 50-180                                & 117.47                                   & 4749.15                                  & 2                                      \\
\textbf{MNIST}         & 70K                                    & 40-75                                 & 70.57                                    & 564.53                                   & 10                                     \\
\textbf{CIFAR10}       & 60K                                    & 85-150                                & 117.63                                   & 941.07                                   & 10                                     \\
\textbf{TSP}           & 12K                                    & 50-500                                & 275.76                                   & 6894.04                                  & 2                                      \\
\textbf{obgl-ppa}      & 1                                      & 576,289                               & -                                        & 30,326,273                               & -                                      \\
\textbf{obgl-collab}   & 1                                      & 235,868                               & -                                        & 1,285,465                                & -                                      \\
 \hline
\end{tabular}
\end{table}



\subsection{Benchmarking Graph Neural Networks \cite{dwivedi2020benchmarking}}

\paragraph{Datasets.} 
This benchmark contains $6$ main datasets :
\begin{enumerate}[(i)]
\item \textbf{ZINC}, a molecular graphs dataset with a graph regression task where each node represents an atom and each edge represents a bond. The regression target is a property known as the constrained solubility (with mean absolute error as evaluation metric). Additionally, the node features represent the atom's type ($28$ types) and the edge features represents the type of connection ($4$ types).
The hyperparameters range which we used in our search was $\kappa\in\set{20,25,30,40}$ and $d\in\set{35,40,45,50,55}$.
For the reported results we used $\kappa=30, d=45$ and the averaged time for a single epoch (whole training) was $15.5$ seconds ($27.5$ minutes).

\item  \textbf{MNIST} and \textbf{CIFAR10}, the known image classification problem is converted to a graph classification task using Super-pixel representation \cite{Knyazev_superpixel}, which represents small regions of homogeneous intensity as nodes. The edges in the graph are obtained by applying k-nearest neighbor algorithm on the nodes coordinates. Node features are a concatenation of the Super-pixel intensity (RGB for CIFAR10 and greyscale for MNIST) and its image coordinate. Edges features are the k-nearest distances.
For the CIFAR10 and MNIST datasets our search range was $\kappa\in\set{20,25,30}$,\; $d=\set{45,50}$ and $p\in\set{0,0.1,0.2,0.3,0.5}$. The chosen hyperparameters
for the CIFAR10 dataset were $\kappa=30, d=45$ with additional dropout of $p=0.1$. The averaged time for a single epoch (whole training) is $238.5$ seconds ($4.77$ hours). We used the same hyperparameters for the MNIST dataset, besides the dropout which was changed to $p=0.2$. Average time per epoch (whole training) is $197.69$ seconds ($3.84$ hours).

\item \textbf{CLUSTER} and \textbf{PATTERN}, node classification tasks which aim to identify embedded node structures in stochastic block model graphs \cite{abbe2017community}. The goal of the task is to assign each node to the stochastic block it was originated from, while the structure of the graph is governed by two probabilities that define the inner-structure and cross-structure edges. A single representative from each block is assigned with an initial feature that indicates its block while the rest of the nodes have no features.
We searched hyperparameters over the range $\kappa\in\set{20,25,30,40}$ and $d\in\set{35,40,45,50,55}$.
The hyperparameters for the CLUSTER dataset were $\kappa=30, d=45$. Average time per epoch (whole training) is $80.34$ seconds ($1.92$ hours). For the PATTERN dataset we used $\kappa=25, d=50$. Averaged running time per epoch (whole training) is $153.83$ seconds ($3.476$ hours), on a single Tesla P-$100$.


\item \textbf{TSP}, a link prediction task that tries to tackle the NP-hard classical Traveling Salesman Problem \cite{joshi2019efficient}. Given a 2D Euclidean graph the goal is to choose the edges that participate in the minimal edge weight tour of the graph. The evaluation metric for the task is F$1$ score for the positive class. Our
hyperparameters search was in the range $\kappa\in\set{20,25,30}$ and $d\in\set{45,50,55}$, the results shown in the paper uses $\kappa = 20, d=50$ and the averaged running time per epoch (whole training) is $166.02$ seconds ($20.5$ hours), on a single Tesla P-$100$.

\end{enumerate}


\subsection{Link prediction datasets from the OGB benchmark \cite{Hu2020}}
\paragraph{Datasets.} In order to provide a more complete evaluation of our model we also evaluate it on semi-supervised learning tasks of link prediction. We searched over the same hyperparameter range $\kappa\in\set{25,50,100},\;d\in\set{150,256}$ and used $\kappa=50, d=150$ in both tasks. The two datasets were:
\begin{enumerate}[(i)]
    \item \textbf{ogbl-ppa}, an undirected unweighted graph. Nodes represent types of proteins and the edges signify biological connections between proteins. The initial node feature is a 58-dimensional one-hot-vector that indicates the origin specie of the protein. The learning task is to predict new connections between nodes. The train/validation/test split sizes are $21$M/$6$M/$3$M
    . The evaluation metric is called Hits@K \cite{Hu2020}. Averaged running time was 4.5 minutes per epoch and 1.5 hours for the whole training.
    
    \item \textbf{ogbl-collab}, is a graph that represents a network of collaborations between authors. Every author in the network is represented by a node and each collaboration is assigned with an edge. Initial node features are obtained by combining word embeddings of papers by that author (128-dimensional vector). Additionally, each collaboration is described by the year of collaboration and the number of collaborations in that year as a weight. The train/validation/test split sizes are $1.1$M/$60$K/$46$K. Similarly to the previous dataset, the evaluation metric is Hits@K. Averaged running time was 5.22 seconds per epoch and 17.4 minutes for the whole training.
\end{enumerate}

