\section{Introduction}

The modern understanding of strong interactions in the Standard Model of particle physics is based on the theory of Quantum Chromodynamics (QCD), a non-abelian gauge theory where the fundamental degrees of freedom are carried by the quark and gluon fields. Despite its great success at very high energies, we are unable to achieve the same precision in the low-energy regime using full QCD, since the theory becomes nonperturbative.
In particular, theoretical computations of the properties of baryons and nuclei from QCD are extremely difficult even for the smallest nuclei, and  phenomenological models are usually employed, instead.

The Skyrme model \cite{Skyrme61} offers an alternative approach to this problem. It constitutes
 a nonlinear field theory of mesons which corresponds to an effective field theory for low-energy QCD in the large $N_c$ expansion.
In the Skyrme model, baryons and nuclei emerge as topological solitons, i.e., classical solutions with localized energy density which are stabilized due to the nontrivial topology of the vacuum manifold.
As a consequence, many non-perturbative features of low-energy QCD, like the conservation of baryon number, the extended character of nucleons, or the global symmetries of QCD and the corresponding symmetry breaking patterns, follow from built-in properties of the Skyrme model.
This field of research has experienced significant progress in recent years, as different generalizations of the model, like the addition of higher derivative terms \cite{JACKSON1985101} or additional degrees of freedom (DoF)---e.g., vector mesons \cite{ADKINS1984251,MeissnerZahed,Meissner:1987ge,Vecmes2,Sutcliffe:2010et,Ma:2019ery}---, or more general potential terms \cite{Marleau:1990nh}, have been proposed  to better reproduce the observed nuclear properties \cite{Adam_2010,Adam:2013wya,Gillard:2015eia,Gudnason:2016mms,Naya_2018}. Further, improved quantization methods which go beyond the moduli space quantization of spin and isospin have significantly contributed to this recent progress, see, e.g.,  \cite{Halcrow:2016spb,Gudnason:2018ysx,Halcrow:2020gbm}.

On the other hand, the first observations of gravitational waves by LIGO opened a new window for the exploration of matter at ultra high densities, like at the cores of Neutron Stars (NS), which are thought to be the most dense objects allowed by General Relativity (GR) before collapsing to a black hole. Indeed, recent \cite{Abbott_2017}---and prospect---observations of mergers of NS binaries will allow us to constrain the equation of state (EoS) of nuclear matter at such high densities. In particular, since the Skyrme model and its generalizations allow to find star-like solutions when coupled to GR, these observations may serve us to determine whether the (generalized) Skyrme model is a consistent way to describe the properties of nuclei and nuclear matter at a large range of scales in a unified manner.

Different models for NS as Skyrme solitons have been previously proposed, for example, in \cite{NelmesPiette,Adam_2015a}. These models are interesting from a theoretical point of view, because they allow to obtain the EoS of NS cores from a relatively simple field theoretic description. However, none of the Skyrmion star models present in the literature have achieved a good agreement with current observational data of NS \cite{Naya_2019}.
In this paper, we present an EoS for NS based on a generalized Skyrme model
which satisfies {\em all} recent observational constraints of NS, such as the maximum mass limit or the deformability as measured in coalescent binary systems.

In this article we will use units in which $c =1$. For masses and lengths we use either nuclear physics units (MeV and fm) or astrophysical units (solar masses $M_\odot$ and km).

\section{Skyrme crystals}

The Skyrme model is an effective field theory of strong interactions at low energies which emerges in the large $N_c$ limit of QCD. It is defined via the Lagrangian
\begin{equation}
    \mathscr{L}_{SK}
    \frac{-f_\pi^2}{4}\Tr{L_\mu L^\mu}+\frac{1}{32e^2}\Tr{[L_\mu,L_\nu][L^\mu,L^\nu]}-\mu^2\mathcal{U},
    \label{Lsk}
\end{equation}
with $f_\pi$ the pion decay constant and $e$ the Skyrme coupling constant. Also, $L_\mu=U^\dagger\partial_\mu U$ is the left invariant Maurer-Cartan form associated to the SU(2)-valued Skyrme field $U(x)$, and $\mathcal{U}=\mathcal{U}(U)$ is a potential.  For the pion mass potential $\mathcal{U}_\pi = (1/2)\, {\rm tr}\, (1-U)$, the parameter $\mu$ is related to the pion mass $m_\pi$ via $\mu = (1/2) f_\pi m_\pi$.

In order to obtain finite energy configurations, one imposes constant boundary values of $U$ at $\abs{x}\rightarrow \infty$, so that the physically relevant Skyrme field configurations define maps
$
    U:S^3\rightarrow SU(2)\simeq S^3,
$
and thus the Skyrme model presents topological solitons (Skyrmions), whose topological charge equals the topological degree of these maps,
\begin{equation}
    \mathcal{B}=\int B^0 d^3 x, \,
    \text{with}\;
B^\mu=\frac{1}{24\pi^2}\varepsilon^{\mu\nu\rho\sigma}\Tr{L_\nu L_\rho L_\sigma}
\end{equation}
the baryon density current.
The Skyrme model \eqref{Lsk} describes an interacting theory for the Goldstone bosons associated to the (broken) chiral symmetry, but baryons emerge as topological solitons, whose topological charge corresponds to the baryon number \cite{Witten:1983tx}. Furthermore, the Skyrme model has been applied to the study of matter at extremely high densities, required to describe the EoS of NS. To do so, one needs to find the lowest energy solutions of the Skyrme model for the very large baryon number of NS, typically  $N\sim N_{\odot}\sim10^{57}$.

It is well known \cite{Kugler:1989uc,Naya_2019} that the lowest energy solutions of the standard Skyrme model (described by the Lagrangian density \eqref{Lsk}) for very large baryon number consist of crystalline cubic  lattices of $B=4$ Skyrmions---which can be thought of as $\alpha$ particles.
The energy per baryon of such solutions as a function of the lattice parameter of the unit cell, $l$, is given by :
\begin{equation}
    E(l)=E_0\qty[a\qty(\frac{l}{l_0}+\frac{l_0}{l})+b]
    \label{energy/baryon}
\end{equation}
being $a=0.474$ and $b=0.0515$ adimensional quantities whose numerical values were obtained in \cite{CASTILLEJO1989801}.
We fit the values of energy (per baryon) and lattice length corresponding to the minimum energy configuration, $E_0 = 923.32$ MeV and $l_0^{-3} = n_0 = 0.16$ $\text{fm}^{-3}$, to reproduce the energy per baryon of infinite nuclear matter at nuclear saturation density $n_0$, \cite{Vinhas_2015}. Note that our values  slightly differ from those originally proposed in \cite{CASTILLEJO1989801} due to the different fit \footnote{In \cite{CASTILLEJO1989801} $E_0$ and $l_0$ are fitted to the nucleon in the standard Skyrme model parametrization which, on its part, uses the fit to the nucleon and Delta resonance masses. For our purposes, a fit to infinite nuclear matter is much more natural. In addition, using the (nonrelativistic) rigid rotor quantization to calculate the (highly relativistic) Delta mass is intrinsically problematic \cite{Adam:2016drk}.} \footnote{In principle, the pion mass term $(1/4) m_\pi^2 f_\pi^2 l^3$ should be added, but it turns out that its contribution to the Skyrme crystal is negligible for $l\le l_0$ \cite{NelmesPiette, CASTILLEJO1989801}}. From this expression one may obtain the energy per baryon as a function of the pressure \cite{NelmesPiette}, i.e., the EoS of the Skyrme crystal (at zero temperature).

Indeed, by the thermodynamical definition of pressure at zero temperature,
\begin{equation}
    p=-\pdv{E}{V}\equiv-\pdv{E(l)}{l^3}=-\frac{1}{3l^2}\dv{E(l)}{l},
\end{equation}
we have
\begin{equation}
    p(l)=a\frac{E_0}{3l^2}\qty(\frac{l_0}{l^2}-\frac{1}{l_0}) .
    \label{p(l)}
\end{equation}
This expression for the pressure vanishes at the finite length $l=l_0$, which is a well-known property of infinite nuclear matter at saturation density $n_0$. Further we shall argue below that the standard Skyrme crystal should provide the leading contribution to the nuclear EoS close to nuclear saturation. This explains the fit of the Skyrme crystal parameters $l_0$ and $E_0$ to the infinite nuclear matter values.

The above expression $p(l)$ can be inverted (solved for $l$),
\begin{equation}
    \frac{l_0^2}{l^2}=\frac{1}{2}\left( 1+\sqrt{1+\frac{12}{E_0a}pl_0^3} \right) ,
    \label{y(p)}
\end{equation}
and we may substitute the resulting $l(p)$ into \eqref{energy/baryon} to obtain the energy per baryon of the crystal as a function of the pressure, i.e. the equation of state for the Skyrme crystal (at zero temperature).

\section{The Generalized Skyrme model}
\label{sec:2}
Since it is an effective theory, the Skyrme model can be extended by adding higher order terms to the Lagrangian. The  only possible Lorentz-invariant extra term with at most second order time derivatives of the Skyrme field is \cite{Adam_2010}
\begin{equation} \label{sextic}
    \mathscr{L}_6=-\lambda^2\pi^4B_\mu B^\mu,
\end{equation}
with $\lambda$ a coupling parameter. Thus, the generalized Skyrme Lagrangian reads ${\mathscr{L}_{SK}^{gen}=\mathscr{L}_{SK}+\mathscr{L}_6}$.

Unfortunately, neither large $\mathcal{B}$ solutions for $\mathscr{L}_{SK}^{gen}$ nor the corresponding EoS have been found, to our knowledge. However, at sufficiently high densities---for instance, those which occur at the core of a neutron star, which can reach several times the nuclear saturation density $n_0$---, the sextic term (\ref{sextic}) provides the most important contribution to the EoS, related to the $\omega$ meson repulsion of nuclear matter \cite{Adam:2015lra}. The sextic term alone defines a barotropic perfect fluid with energy density $\rho_6 = \lambda^2 \pi^4 n^2 =p$ (see below), where $p$ is the pressure and $n$ the baryon number density. The EoS $\rho_6 =p$ is maximally stiff with a speed of sound equal to 1, which explains its dominance at high density.

$\mathscr{L}_{SK}^{gen}$ has another interesting submodel which will be relevant for us, the so-called BPS (=Bogomolnyi-Prasad-Sommerfield) Skyrme model
$
    \mathscr{L}_{BPS}=\mathscr{L}_6-\mu^2\mathcal{U}(U).
$
This model supports topological soliton configurations saturating a BPS energy bound \cite{Adam_2010}, hence the name of the model.
Minimally coupling this submodel to gravity \cite{Adam_2015b}, we obtain its stress-energy tensor
which still is of the perfect fluid form, $T_{BPS}^{\mu\nu}=(p+\rho)u^\mu u^\nu-p g^{\mu\nu}$, with the following definitions (here $g := \abs{\det{g_{\rho\sigma}}}$),
\begin{equation}
    u^\mu=\frac{B^\mu}{\sqrt{g_{\rho\sigma}B^\rho B^\sigma}},\qquad p=\lambda^2\pi^4g^{-1}g_{\rho\sigma}B^\rho B^\sigma-\mu^2\mathcal{U}.
    \label{pressuredensitydef}
\end{equation}
and $\rho=p+2\mu^2\mathcal{U}$. Further, the proper baryon number density is $n=u^\mu (g^{-\frac{1}{2}}B_\mu ) = \sqrt{g^{-1}g_{\mu\nu}B^\mu B^\nu}$.
Note that this perfect fluid is, in general, non-barotropic, since the potential term $\mathcal{U}$ introduces a dependence on the Skyrme field in $p$ and $\rho$, such that no simple algebraic relation can be found between them.
 Nevertheless, one may still perform a mean-field approximation and obtain an effective, barotropic EoS for the BPS Skyrme fluid, which offers the interesting possibility to compare the results obtained within the exact and the mean-field approaches \cite{Adam_2015b}.

In the case of interest here, however, we will introduce a constant effective potential $\mu^2 \mathcal{U} = \rho_0$, which is supposed to take into account the effects of the subleading contributions above a certain threshold value $p_{\rm PT}$ for the pressure, see below. This is equivalent to choosing the theta-term potential of ref. \cite{Adam_2015b} and implies the barotropic EoS
\begin{equation}
\rho = \rho_6 + \rho_0 = \lambda^2\pi^4 n^2 +    \rho_0 = p + 2\, \rho_0
\label{BPSEoS}
\end{equation}
already at the full field-theory level.

\section{A generalized equation of state}
\label{subsec:hybrid model}

Both the standard Skyrme model and the BPS submodel have been previously used to describe nuclear matter inside NS \cite{Naya_2019}. However, it is clear from these attempts that the true equation of state for Skyrme matter should take into account both models in a unified fashion, because the results from approximating the full model with either of the two submodels deviate from the most recent observational data of NS, and do so in opposite directions. For example, the maximum mass of NS are either too small (for pure skyrmion crystals) or too large (for BPS Skyrmion stars) as compared with the current constraints \cite{Naya_2019}.
As explained, the generalized Skyrme model has not been solved yet for large baryon number. Nevertheless, we may still obtain some information of these high baryon number solutions by scaling arguments of the energy terms for the different submodels of the complete Lagrangian.

Indeed, consider for example the case of the Skyrmion crystal, whose energy per baryon is given by \eqref{energy/baryon}, and let $\sigma\in (0,1]$. A scale transformation of the space coordinates of the form $x\mapsto x/\sigma$ can be understood as a mapping between crystalline solutions, respectively, with lattice size $l$ and $\sigma l$. On the other hand, since the lattice length is a function of the pressure, we conclude that two solutions at different pressures $p$ and $p'$ which have a lattice length of $l(p)$ and $l'(p')$ respectively, are related through a scale transformation $\sigma(p,p')$ such that $l'(p')=\sigma(p,p')l(p)$. In particular, any configuration with lattice length $l(p)$ will be related to the zero pressure crystal (minimum energy configuration) via $l(p)=\sigma(p)l_0$, where $\sigma(p)=\sigma(0,p)$ can be seen as a function relating the pressure of the crystal and the scaling parameter. Furthermore, taking into account \eqref{y(p)}, we find
\begin{equation}
    \sigma (p)=\sqrt{\frac{2}{1+\sqrt{1+\frac{12l_0^3}{E_0a}p}}}.
\end{equation}
This expression has indeed the correct limits of $\sigma(p\rightarrow\infty)\rightarrow 0$ and $\sigma(p=0)=1$ .

This equivalence between pressure and scaling allows us to write the energy per baryon of the Skyrmion crystal at any pressure (i.e. $\sigma\neq 1$) as a simple function of $\sigma =l/l_0$,
$    E(\sigma)=a \, E_0\left( \sigma+\sigma^{-1}\right) + b \, E_0 $.
Obviously, the contribution from the term proportional to $\sigma$ becomes negligible for large pressure, whereas the term proportional to $\sigma^{-1}$ dominates in this regime (${\sigma\ll1}$).

Next, consider the sextic term contribution to the energy (and energy per baryon) of a fluid element $\Omega$
\begin{equation}
  \frac{E_6}{\mathcal{B}}=\frac{(\int_\Omega d^3 x \sqrt{g}\, \rho_6) }{(\int_\Omega d^3 x \sqrt{g}\,  n)},
\end{equation}
which transforms as $E_6\mapsto \sigma^{-3}E_6$ under a scaling of spacetime coordinates. This implies that the sextic contribution will dominate the energy per baryon at sufficiently high pressure. Therefore, we may assume that a solution of the complete model will provide an EoS which tends to the EoS of the submodel $\rho_6$ at high pressure, with an asymptotic energy per baryon of $E_6/\mathcal{B}=\rho_6/n = \lambda\pi^2\sqrt{p}$. This is, therefore, the asymptotic behavior of the energy per baryon at high pressure also for the full model.

On the other hand, as the pressure decreases to a certain value (which depends on $\lambda$), $E_6/\mathcal{B}$ becomes of the order of the energy per baryon of the Skyrme crystal, and the BPS approximation to the complete solution will start to fail.
For even lower $p$, the contribution of $E_6/\mathcal{B}$ will be subleading in comparison to the Skyrme crystal.

This supports the idea that a transition of some kind must take place within this generalized model, between the crystalline phase of the standard Skyrme model and the perfect fluid phase of the BPS model. A quantitative prediction of the pressure value $p_{PT}$ where this transition occurs, as well as the determination of its character---a smooth crossover or a phase transition---would require the knowledge of the full solution or, at least, the value of the parameter $\lambda$, because the contribution to the energy per baryon of the sextic term strongly depends on $\lambda$.

In \cite{Adam_2015a}, the BPS submodel was used to model the full neutron star core and, therefore, the model parameters $\lambda$ and $\mu$ were fitted to match with the infinite nuclear matter approximation at zero pressure. In the present case, however, the Skyrme crystal describes the low-pressure region and, therefore, should be fitted to nuclear matter.
In this section, we will propose an EoS for the generalized model corresponding to $\mathscr{L}_{SK}^{gen}$. The value of $\lambda$ will be determined, instead, by the behavior of the EoS in the limit of very high pressure, in which, as argued, it can be approximated by only the sextic term, see below.

\begin{figure}[b!]
\hspace*{-0.4cm}
    \centering
    \includegraphics[scale=0.4]{EOS.pdf}
    \caption{Comparison of the Skyrme crystal and the generalized model EoS to other neutron star  EoS usually considered in the literature, namely, APR4 \cite{Akmal:1998cf}, WFF1 \cite{Wiringa:1988tp} and BCPM \cite{Vinhas_2015} (the BCPM EoS turns out to be numerically very similar to SLy4 \cite{Douchin:2001sv}). In this figure, we show the curve for  the generalized model with $p_{PT}=50$ MeV/fm$^3$, but the range of possible values of $p_{PT}$ ($p_{*}$) that yield results consistent with observations corresponds to the blue (yellow) stripe: $p_{PT} \in [25, 50]$ MeV/fm$^3$, $p_* \in [0.5, 2.0]$ MeV/fm$^3$.}
    \label{fig:EoS}
\end{figure}

From the previous considerations, we can construct a generalized EoS which takes into account both the standard Skyrme and BPS submodels at different regimes, based on simple assumptions on the behavior of the full solutions  in the low and high pressure regimes, without knowing these solutions explicitly. Indeed, we will assume that the low pressure solutions of the complete model are still Skyrme crystals whose energy is approximately described by \eqref{energy/baryon}.
In the fluid high-pressure phase, we will assume that the sextic term provides the most important contribution, and the complete solutions can be well described by a BPS Skyrme model. We can model this behavior by introducing a certain value of the pressure, $p_{PT}$, above which the solutions are described by a BPS fluid. Therefore, the generalized EoS $\rho_{\text{Gen}}(p)$ must satisfy
\begin{equation}
    \rho_{\text{Gen}}(p)\simeq\left\{\mqty{\rho_{SK}(p),&\quad p<< p_{PT}\\\text{const.}+p,&\quad p>> p_{PT}.} \right.
    \label{limitsGenEoS}
\end{equation}
A simple way of parametrizing this behavior that yields a smooth transition between these two regimes is to consider an EoS of the form
\begin{equation}
\rho_{\text{Gen}}(p)=(1-\alpha(p))\rho_{SK}+\alpha(p)(p+\rho_{SK}(p_{PT})),
\label{GenEoS}
\end{equation}
where $\alpha (p)$ is a function which interpolates between the two regimes, i.e.,
$\alpha \to 0$ for $p/p_{\rm PT} \to 0$ and $\alpha \to 1$ for $p/p_{\rm PT} \to \infty$. Concretely, we consider the interpolating functions
\begin{equation}
    \alpha(p,p_{\rm PT},\beta)=\frac{\left( \frac{p}{p_{PT}}\right)^\beta}{1+\left( \frac{p}{p_{PT}}\right)^\beta}
    \label{interpol}
\end{equation}
as in \cite{Adam:2020djl}. Here, smaller values of $\beta$ produce a more gradual transition, whereas larger values correspond to a faster transition between the two regimes. For the transition between the Skyrme crystal and the BPS fluid at $p_{\rm PT}$, we have to choose the rather gradual transition $\beta = 0.9$, because otherwise the resulting energy density \eqref{GenEoS} would lead to acausal propagation (a speed of sound larger than one) in some regions inside the star for some values of $p_{PT}$.

As a result of this interpolation, the energy density contribution from the crystal becomes less and less important as $p$ grows, freezing at its value at $p_{PT}$ for sufficiently high pressures, playing the role of an effective potential energy for the BPS Skyrme model. The $p$ dependence for $p>p_{PT}$ is taken into account by $\rho_6$, which is known to provide the leading contribution for large $p$. Therefore, the generalized EoS \eqref{GenEoS} is effectively equivalent to that of a BPS Skyrme model with a theta potential \cite{Adam:2013wya} for $p>>p_{PT}$. In the following section, we will see that the value of $p_{PT}$ determines the maximum mass of a NS, so we may adjust the value of $p_{PT}$ to agree with the current maximum mass limit for NS.

To obtain the baryon density $n$ in the generalized model, we use the well-known Euler relation
 $   {\rho = -p + \pdv{\rho}{n}\: n},$
which yields a differential equation for $n$, that we integrate using $n(p=0) \equiv n_0 = 0.16$ $\text{fm}^{-3}$ as initial condition to obtain the curve $n(p)$. The result, and the corresponding EoS $\rho (p)$, are shown in Fig. \ref{fig:EoS}, where other EoS have been included for comparison.

\subsection{Addition of crust: the hybrid EoS}
The generalized Skyrme EoS (\ref{GenEoS}), by construction, only describes nuclear matter above nuclear saturation \footnote{we use the recent value  $n_0 = 0.160\:{\rm fm}^{-3}$ for the nuclear saturation density, see \cite{Vinhas_2015}}.
Below saturation density, nuclear matter in a NS is known to be in a rather inhomogeneous state, resulting from a competition between nuclear and electromagnetic forces (e.g., "nuclear pasta" phases \cite{Chamel_2008}).
In principle, the (generalized) Skyrme model can be coupled to the electromagnetic interaction, so these low-density phases are fundamentally within its scope. Full field-theoretical calculations for this coupled system and for large $\mathcal{B}$ are, however, not feasible, and a macroscopic (hydrodynamical) treatment is currently unknown.
On the other hand, the standard methods of nuclear physics, such as many-body techniques, can be used to describe these low-density NS crust regions and are completely reliable there.  This motivates us to consider a hybrid version of \eqref{GenEoS} in which, at a sufficiently low density $n_{*}$ (or, equivalently, $p_*$), a neutron star crust EoS  $\rho_{\text{BCPM}}(p)$ is glued,
\begin{equation}
    \rho_{\text{Hyb}}(p)\simeq\left\{\mqty{\rho_{\text{BCPM}}(p),&\quad p\leq p_{*}\\\rho_{\text{Gen}}(p),&\quad p\geq p_{*}.} \right.
    \label{HybridEoS}
\end{equation}
Concretely, we choose the BCPM EoS of \cite{Vinhas_2015}, based on the Brueckner-Hartree-Fock (BHF) approach (plus the
BCPM density functional for the crust). For the crust and the outer core $n\lesssim n_0$, nuclear matter is well understood, and standard nuclear physics EoS like \cite{Vinhas_2015} should provide a precise description of NS matter. Again, we choose a smooth transition between the two regimes, using the interpolating function (\ref{interpol}).
Now we choose the faster transition $\beta=2$, exactly as was done in \cite{Adam:2020djl} (replacing $p_{\rm PT}$ by $p_*$).

\section{Observational constraints}
To determine the static properties of the resulting NS, we simply insert the hybrid EoS \eqref{HybridEoS} into the relativistic equations of hydrodynamical equilibrium, the so-called Tolman-Oppenheimer-Volkoff (TOV) equations  \cite{Tolman:1939jz,Oppenheimer:1939ne}.
In this hybrid EoS, there are only two free parameters, namely the values of $p_{*}$ and $p_{PT}$ corresponding to the low and high density parts of the hybrid EoS. Here we show that recent astrophysical and gravitational wave observations actually tightly constrain the value ranges for both parameters.
 For example, from the mass-radius curves for different values of these parameters, we find that only the value of $p_{PT}$ affects the maximum NS mass in the model. Thus, we could for example constrain the value of $p_{PT}$ using the maximum mass limit for nonrotating NS of $M/M_{\odot}=2.16^{+0.17}_{-0.15}$ proposed in \cite{Rezzolla_2018}. However, given the recent GW observations of GW190425, with a total mass of $3.4^{+0.3}_{-0.1}M_{\odot}$ and mass ranges of components varying from 1.12 to 2.52 $M_\odot$ \cite{Abbott:2020uma} and GW190814, a compact binary merger between a $22.2-24.3 M_{\odot}$ black hole and a secondary object which falls in the mass gap  ${(2.50-2.67 M_{\odot})}$ \cite{Abbott:2020khf}, we have allowed the range of values of $p_{PT}$ to yield stars of maximum mass up to $\sim 2.7\, M_{\odot}$
 \begin{figure}

     \includegraphics[scale=0.55]{MvsR_v5.pdf}
     \caption{Mass-Radius relation for the hybrid model (red curves) for different combinations of values of ${p_{*}=\{0.5,1,2\}\text{MeV/fm}^{3}}$ and ${p_{\text{PT}}=\{25,40,50\}\text{MeV/fm}^{3}}$. The red shaded region corresponds to the accessible region of the hybrid model with $p_{*}$ and $p_{PT}$ within the given ranges (see Fig. 1).}
     \label{fig:MvR}
 \end{figure}

 In Fig. \ref{fig:MvR} we show different mass-radius curves of the hybrid model corresponding to different values of $p_{PT}$. We can see a good agreement, for any pair $(p_{*},p_{\text{PT}})$ within the ranges ${p_*\in[0.5,2]\,\text{MeV/fm}^{3}}$ and ${p_{PT}\in[25,50]\,\text{MeV/fm}^{3}}$, with the most likely mass-radius relation for the NS corresponding to the GW170817 event \cite{Abbott_2017}. We haven't included the corresponding data of the second BNS event, namely, GW190425, since it was less informative on matter effects than GW170817 , although our data is still compatible with this event as well, specially for lower values of $p_{PT}$  \cite{Abbott:2020uma}. In the same figure, we represent the masses of some of the heavier pulsars measured by the NICER collaboration, PSR $\mathrm{J} 1614-2230\left(1.928 \pm 0.017 \mathrm{M}_{\odot}\right)$\cite{Fonseca:2016tux}, PSR {$\mathrm{J} 0348+0432$}$\left(2.01 \pm 0.04 M_{\odot}\right)$ \cite{Antoniadis:2013pzd} and {$\mathrm{PSR}\,\mathrm{J}0740+6620$} $\left(2.14_{-0.09}^{+0.10} M_{\odot}\right)$ \cite{Cromartie:2019kug}, as well as the most probable M-R region from combined observations of GW and these heavy pulsars \cite{Landry:2020vaw}. Also, other constraints from NICER, chiral EFT and multimessenger observations are represented, adapted from \cite{Constraints2020neutron} and \cite{Greif:2020pju}.

The observed gravitational waveform can also be used to
place direct constraints on the tidal deformability of NS. Indeed, the  waveform produced by the coalescence of two NS at the early phase of the inspiral depends on the underlying EoS mostly
through the tidal Love number \cite{Hinderer:2009ca}. However, the individual Love numbers for the two stars
cannot be disentangled in the observed gravitational waveform.
Instead, what is measured is the so-called effective tidal deformability $\tilde{\Lambda}$, a mass weighted average of the deformabilities of the individual stars in the merger \cite{Flanagan_2008}. Similarly, the two component masses are not measured directly, but the chirp mass, $M_c=m_{1} \,{q^{3 / 5}}/{(1+q)^{1 / 5}}$
where $q=m_1/m_2$ is the mass ratio,
can actually be tightly constrained. In the case of the GW170817 event, the chirp mass was constrained to $1.188_{-0.002}^{+0.004}$
at the $90\%$ confidence level, and  the mass ratio was constrained to be in the range $0.7-1$ within the same confidence level, whereas the effective tidal deformability was inferred to be smaller than $800$ \cite{Abbott:2018wiz}.

\begin{figure}[bht]
\hspace*{-0.5cm}
    \includegraphics[scale=0.4]{Lambdatilde.pdf}
    \caption{$\tilde \Lambda$ as a function of the mass ratio. The orange shaded regions correspond to the $50\%$ (dark) and $90 \%$ (light) credible regions for the joint posterior of $\tilde\Lambda$ and $q$ PDFs as obtained in \cite{Abbott:2018wiz} assuming a low spin prior. Notation for curves from the EoS (\ref{HybridEoS}): Hyb\_\_$p_{PT}$\_\_$p_*$.
    }
    \label{fig:LAMBDATIDAL}
\end{figure}

Such measurements allow to reduce the set of Skyrme models able to reproduce the NS properties. Following \cite{Yagi_2013}, we have solved the Einstein equations for slowly rotating Skyrmion stars with the hybrid EoS using the Hartle-Thorne formalism \cite{HartleI,HartleThorne85} and obtained the dimensionless tidal deformability of stars described by this model as a function of their TOV mass. On the other hand, since the chirp mass
of the binary progenitor of GW170817 is well measured,
for any given EoS the effective deformability reduces to a simple EoS-dependent function of the mass ratio.
These curves, together with the constraints commented above, are represented in Fig. \ref{fig:LAMBDATIDAL}, from where it follows that our new EoS is compatible with the data from \cite{Abbott:2018wiz} for the ranges of $p_{*}$ and $p_{\text{PT}}$ considered. Future measurements of the tidal deformability of NS will allow us to further constrain these ranges, since we find that the curves $\Tilde{\Lambda}(q)$ depend on the particular values of both parameters.

\subsection{The sextic term and the $\omega$ meson}
As explained, the sextic term \eqref{sextic} automatically provides the leading contribution at high densities if it is included in the effective action. On the other hand, this term is physically justified because it provides the leading contribution of the $\omega$ meson repulsion in a derivative expansion which results from integrating out the $\omega$ field from an extended Lagrangian which includes both pions and vector mesons \cite{ADKINS1984251,Meissner:1987ge,Vecmes2}.  This relation to the $\omega$ meson not only motivates the sextic term, but it also leads to an expression of its coupling constant $\lambda$ in terms of the physical parameters of the $\omega$ meson. Indeed, it can be shown \cite{InmediumBPS} that
$ \lambda^2=\frac{g_{\omega}^2}{2\pi^4m_\omega^2}$,
 where $m_\omega$ and $g_\omega$ are, respectively, the mass and coupling constant of the $\omega$ vector meson. For the empirical values $m_\omega = 783 \,\text{MeV}$ and
 $g^2_\omega/(4\pi)\sim 12$ \cite{Meissner:1987ge}, we get $\lambda^2 \sim 10\, \text{MeV\,fm}^3$.

The $\lambda$ coupling constant appearing in front of the sextic term in the generalized Lagrangian does not directly show up in the generalized EoS proposed in this work, whose parameters are constrained by the observations of maximum mass and deformability. However, the generalized Skyrme model EoS approaches the EoS of the BPS submodel for sufficiently large pressure, by assumption. We can, therefore, extract an effective value of $\lambda$ by taking the limit of infinite pressure and using the BPS EoS \eqref{BPSEoS}, which immediately implies
$
    n^2=(\rho+p)/{(2\lambda^2\pi^4)}.
$
The effective value of $\lambda$ is then given by
\begin{equation}
    \lambda_{eff}=\lim\limits_{p\rightarrow\infty}\frac{1}{n(p)\pi^2}\sqrt{\frac{\rho_{Gen}(p)+p}{2}}.
\end{equation}
For the range of values $p_{PT}\in[25,50]\,\text{MeV/fm}^{3}$, we find that $\lambda_{eff}^2\in[10 ,13.5]\,\text{MeV\,fm}^{3}$. The values so obtained for $\lambda_{eff}^2$ are, therefore, perfectly compatible with the values obtained by assuming that the sextic term in the generalized Lagrangian results from integrating out the $\omega$ vector meson.

\section{Conclusions}
In this letter, we propose a completion of standard nuclear physics EoS at low densities---known to be reliable there---by an EoS based on the generalized Skyrme model in the uncharted territory above nuclear saturation density $n_0$. In the simplest version of Skyrme models, where electromagnetic effects, quantum corrections or the proton-neutron mass difference are not taken into account, they can describe nuclear matter only for $n\ge n_0$, by construction. The use of the generalized Skyrme model at densities $n_0<n\leq n_{\text{max}}$ is based on the assumptions that $i)$ strong-interaction effects (nuclear repulsion) are more important than degeneracy pressures in that region, $ii)$ the extended character of nucleons---which is automatic in the Skyrme model---is relevant at high pressure and $iii)$ nucleons are the only relevant DoF inside NS cores (no exotic contributions). This last assumption is shared by many NS models. Here, $n_{\text{max}}$ corresponds to the central density of the maximum mass NS, which is $n_{\text{max}}=7.1 n_0$ for $p_{PT}=50{\rm MeV}/{\rm fm}^3$ and $n_{\text{max}}=5.3 n_0$ for $p_{PT}=25{\rm MeV}/{\rm fm}^3$, safely below the deconfinement phase transition density $\sim 40n_0$.\footnote{In the recent paper \cite{cite-NaturepQCD}, some arguments in favor of quark matter cores of heavy NS were given. Their analysis, however, is based on the assumption that the speed of sound does not exceed the so called ``conformal bound'' $c^2_s=1/3$, which obviously does not hold in our model.  }

We find that the resulting EoS provides an excellent description of NS properties, compatible with all constraints, among them the latest ones from LIGO. Our EoS contains two parameters which have a clear physical interpretation as transitions between standard nuclear matter and the Skyrme crystal ($p_*$) and between this crystal and a Skyrme fluid ($p_{PT}$). In particular, we propose a rather smooth transition between a crystalline and a fluid regime for $20\le p_{PT}\cdot{\rm fm}^3/{\rm MeV}\le 50$, whose precise position may be determined by more precise NS binary observations. Let us remark that the very recent observation of the GW190814 event \cite{Abbott:2020khf}, with a certain indication of an NS with a mass of about 2.6 $M_\odot$ \cite{fattoyev2020gw190814,Huang:2020cab}, can be easily accommodated by our generalized Skyrme model EoS, by simply choosing a slightly lower value of $p_{\rm PT}\sim 25$ MeV fm$^{-3}$ for the transition between Skyrme crystal and BPS fluid, see Fig. \ref{fig:MvR}.

We also find that the range of values for the effective coupling constant $\lambda$ of the sextic term which results from our generalized EoS and the fit to realistic NS is perfectly compatible with the range of values resulting from its relation to the $\omega$ meson.

Motivated by the results obtained with the hybrid EoS proposed above, it
would certainly be interesting to try to derive a similar EoS from an
exact solution of the generalized Lagrangian, for example, using a
crystalline ansatz for the Skyrme field, as in \cite{gencryst}. One then
could study whether this exact EoS presents a phase transition of some
kind, in the same fashion as the proposed hybrid EoS. We leave this study
for a future work.

Finally, we would like to comment on the similarities and differences of
our proposal to the scenarios considered in Ref. [7]. The calculations and
discussions in [7] are based on the (standard) Skyrme crystal and are, in
this sense, similar in spirit to ours. There are, nevertheless, some
important differences. First of all, in [7] more degrees of freedom are
considered, among them the dilaton to recover the scale symmetry of QCD at
large densities. At low density, this symmetry is broken spontaneously—the
dilaton freezes—and the Skyrme model is recovered. Further, higher mass
mesons are taken into account implicitly in \cite{Ma:2019ery}. The main difference for
the present purpose, however, is related to another phase transition which
is known to occur in the standard Skyrme crystal \cite{Kugler:1989uc}, namely the
transition from a skyrmion phase to a half-skyrmion phase as density
increases. In \cite{Ma:2019ery}, this transition leads to a significant stiffening of
the EoS which is mainly related to an enhanced contribution of the
symmetry energy in the half-skyrmion phase. In our case, we assume that in
the region where the Skyrme model is effective (i.e., for $p>p_*$) we are
always in the half-skyrmion phase. In addition, the Skyrme crystal
influences the EoS only via its scaling properties in our case, see Eq.
(1), and the effects of the symmetry energy are taken into account only
implicitly, by an appropriate choice of our physical parameters. The
stiffening of the EoS is caused, instead, by the sextic term, i.e., by the
$\omega$ repulsion, as explained in the main text.

\begin{acknowledgements}
The authors acknowledge financial support from the Ministry of Education, Culture, and Sports, Spain (Grant No. FPA2017-83814-P), the Xunta de Galicia (Grant No. INCITE09.296.035PR and Conselleria de Educacion), the Spanish Consolider-Ingenio 2010 Programme CPAN (CSD2007-00042), Maria de Maetzu Unit of Excellence MDM-2016-0692, and FEDER.
\end{acknowledgements}


\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}