If $W$ is piecewise Lipschitz on $S_{1}$ and $S_{2}$ where $S_{1}$ is the
diagonal and $S_{2}$ is off-diagonal, then $\|S_{U}W-S_{n}W\|$ converges to 0
almost surely.

###### Proof.

Since the case of $X$ is essentially the same with that of $W$, we only prove
the case of $W$.

  1. 1.

As $n\|S_{U}W-S_{n}W\|_{\infty}\geq\|S_{U}W-S_{n}W\|$, it suffices to prove
that
$n\|S_{U}W-S_{n}W\|_{\infty}=\text{max}_{i,j}|W(u_{(i)},u_{(j)})-W(\frac{i}{n},\frac{j}{n})|$
converges to 0 almost surely. Similar to Lemma 2.9.5, using Lipschitz
condition of $W$ and Glivenko-Cantelli theorem concludes the proof.

  2. 2.

This statement is stronger than the one above. The proof of the last item can
be adapted here. As $W$ is $A_{1}$ Lipschitz on off-diagonal region and
$A_{2}$ Lipschitz on diagonal,

$\displaystyle n\|S_{U}W-S_{n}W\|_{\infty}$
$\displaystyle=\text{max}_{i,j}\left|W(u_{(i)},u_{(j)})-W(\frac{i}{n},\frac{j}{n})\right|$
$\displaystyle=$ $\displaystyle\max\left(\text{max}_{i\neq
j}\left|W(u_{(i)},u_{(j)})-W(\frac{i}{n},\frac{j}{n})\right|,\text{max}_{i=j}\left|W(u_{(i)},u_{(j)})-W(\frac{i}{n},\frac{j}{n})\right|\right).$

Using Lipschitz condition on diagonal and off-diagonal part of $W$ and
Glivenko-Cantelli theorem concludes the proof.

∎

With all lemmas stated, we are ready to prove the main theorem.

See 2.5.2

###### Proof.

Using the triangle inequality

$\displaystyle\operatorname{RMSE}_{U}(\Phi_{c}\left([W,\text{Diag}(X)]\right),\Phi_{d}\left([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})]\right))$
$\displaystyle=\left\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-\frac{1}{\sqrt{n}}\Phi_{d}\left([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})]\right)\right\|$
$\displaystyle=\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)$
$\displaystyle+S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-\Phi_{d}S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{x_{n}})])$
$\displaystyle+\Phi_{d}S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\frac{1}{\sqrt{n}}\Phi_{d}([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{X_{n}})])\|$
$\displaystyle\leqslant\underbrace{\left\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)\right\|}_{\text{First
term: discretization error}}$
$\displaystyle+\underbrace{\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-\Phi_{d}S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])\right\|}_{\text{Second
term: sampling error}}$
$\displaystyle+\underbrace{\left\|\Phi_{d}S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\frac{1}{\sqrt{n}}\Phi_{d}\left([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})]\right)\right\|}_{\text{Third term:
estimation error}}$ (2.13)

The three terms measure the different sources of error. The first term is
concerned with the discretization error. The second term concerns the sampling
error from the randomness of $U$. This term will vanish if we consider only
$S_{n}$ instead of $S_{U}$ for IGN-small. The third term concerns the edge
probability estimation error.

For the first term, it is similar to the sketch in Section 2.5.3.
$\|S_{U}\Phi_{c}([W,\text{Diag}(X)])-S_{U}\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])\|=\|S_{U}(\Phi_{c}([W,\text{Diag}(X)])-\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]))\|$,
if suffices to upper bound
$\|\Phi_{c}([W,\text{Diag}(X)])-\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])\|_{L_{\infty}}$
according to property of $S_{U}$ in Lemma 2.9.3. Since
$\|\Phi_{c}([W,\text{Diag}(X)])-\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])\|_{L_{\infty}}\leqslant
C(\|W-\widetilde{W_{n}}\|_{L_{\infty}}+\|\text{Diag}(X)-\text{Diag}(\widetilde{X_{n}})\|_{L_{\infty}})$
by Proposition 2.4.1, and $\|W-\widetilde{W_{n}}\|_{L_{\infty}}$ converges to
0 in probability according to Lemma 2.9.4, we conclude that the first term
will converges to 0 in probability.

For the third term
$\|\Phi_{d}S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\frac{1}{\sqrt{n}}\Phi_{d}([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})])\|$$\\\
=\|\frac{1}{\sqrt{n}}(\Phi_{d}([W_{n\times
n},\text{Diag}(\widetilde{x_{n}})])-\Phi_{d}([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})]))\|$$=\|\Phi_{d}([W_{n\times
n},\text{Diag}(\widetilde{x_{n}})])-\\\ \Phi_{d}([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})])\|_{\text{pn}}$,

it suffices to control the $\|[W_{n\times
n},\text{Diag}(\widetilde{x_{n}})]-[\widehat{W}_{n\times
n},\text{Diag}(\widetilde{x_{n}})]\|_{\text{pn}}=\frac{1}{n}\|W_{n\times
n}-\widehat{W}_{n\times n}\|_{2}\leqslant\|W_{n\times n}-\widehat{W}_{n\times
n}\|_{2,\infty}$, which will also goes to 0 in probability as $n$ increases
according to the statistical guarantee of edge probability estimation of
neighborhood smoothing algorithm [165], stated in Theorem 2.2.1. Therefore by
Proposition 2.4.1, the third term also goes to 0 in probability.

Therefore the rest work is to control the second term
$\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-\\\
\Phi_{d}S_{U}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)\|$.
Again, we use the triangle inequality

Second term
$\displaystyle=\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-\Phi_{d}S_{U}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)\right\|$
$\displaystyle\leqslant\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|+\left\|S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-\Phi_{d}S_{U}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)\right\|$
$\displaystyle=\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|+\left\|S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-\Phi_{d}S_{n}([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})])\right\|$
$\displaystyle=\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|$
$\displaystyle\leqslant\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-S_{U}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|+\left\|S_{U}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|$
$\displaystyle=\underbrace{\left\|S_{U}\left(\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})]\right)-\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|}_{\text{term
$a$}}+\underbrace{\left\|(S_{U}-S_{n})\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|}_{\text{term
$b$}}$

The second equality holds because
$S_{U}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])=S_{n}([\widetilde{W_{n,E}},\widetilde{X_{n,E}}])$
by definition of $\widetilde{W_{n,E}}$ and IGN-small (See Remark 6 for more
discussion). The third equality holds by the definition of IGN-small. We will
bound the term a)
$\|S_{U}(\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\Phi_{c}([\widetilde{W_{n,E}},\widetilde{X_{n,E}}]))\|$
and b)
$\|(S_{U}-S_{n})\Phi_{c}([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})])\|$
next.

For term a)
$\|S_{U}(\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\Phi_{c}([\widetilde{W_{n,E}},\widetilde{X_{n,E}}]))\|$,
if suffices to prove that
$\|\Phi_{c}([\widetilde{W_{n}},\text{Diag}(\widetilde{X_{n}})])-\Phi_{c}([\widetilde{W_{n,E}},\widetilde{X_{n,E}}]))\|_{L_{\infty}}$
converges to 0 in probability. According to Proposition 2.4.1, it suffices to
bound the
$\|[\widetilde{W_{n}},\widetilde{X_{n}}]-[\widetilde{W_{n,E}},\widetilde{X_{n,E}}]\|_{L_{\infty}}$.
Because
$[\widetilde{W_{n}},\widetilde{X_{n}}]-[\widetilde{W_{n,E}},\widetilde{X_{n,E}}]\|_{L_{\infty}}$
$=\|\widetilde{W_{n}}-\widetilde{W_{n,E}}\|_{L_{\infty}}+\|\text{Diag}(\widetilde{X_{n}})-\text{Diag}(\widetilde{X_{n,E}})\|_{L_{\infty}})\leqslant\|\widetilde{W_{n}}-W\|_{L_{\infty}}+\|\widetilde{W_{n,E}}-W\|_{L_{\infty}}+\|\text{Diag}(\widetilde{X_{n}})-\text{Diag}(X)\|_{L_{\infty}}+\|\text{Diag}(\widetilde{X_{n,E}})-\text{Diag}(X)\|_{L_{\infty}}$,
we only need to upper bound $\|\widetilde{W_{n}}-W\|_{L_{\infty}}$,
$\|\widetilde{W_{n,E}}-W\|_{L_{\infty}}$,
$\|\text{Diag}(\widetilde{X_{n}})-\text{Diag}(X)\|_{L_{\infty}})$ and
$\|\text{Diag}(\widetilde{X_{n,E}})-\text{Diag}(X)\|_{L_{\infty}})$, which are
proved by Lemma 2.9.4 and Lemma 2.9.5 respectively.

For term b)
$\|(S_{U}-S_{n})\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\|$

$\displaystyle\left\|(S_{U}-S_{n})\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|$
$\displaystyle=\left\|(S_{U}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|$
$\displaystyle\leqslant\left\|(S_{U}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)\right\|+\left\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{n}\Phi_{c}\left([W,\text{Diag}(X)]\right)\right\|$
$\displaystyle+\left\|S_{n}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{n}\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)\right\|$
$\displaystyle=\left\|(S_{U}(\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-\Phi_{c}([W,\text{Diag}(X)]))\right\|+\left\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{n}\Phi_{c}\left([W,\text{Diag}(X)]\right)\right\|$
$\displaystyle+\left\|S_{n}(\Phi_{c}\left([\widetilde{W_{n,E}},\text{Diag}(\widetilde{X_{n,E}})]\right)-\Phi_{c}([W,\text{Diag}(X)]))\right\|$

For the first and last term, by the property of $S_{U},S_{n}$ and $\Phi_{c}$,
it suffices to bound $\|W-\widetilde{W_{n,E}}\|_{L_{\infty}}$ and
$\|\text{Diag}(X)-\text{Diag}(\widetilde{X_{n,E}})\|_{L_{\infty}}$. Without
loss of generality, We only prove the case for $W$. As
$\|W-\widetilde{W_{n,E}}\|_{L_{\infty}}$ converges to 0 almost surely by Lemma
2.9.5, we conclude that the first and last term converges to 0 almost surely
(therefore in probability). For the second term
$\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-S_{n}\Phi_{c}\left([W,\text{Diag}(X)]\right)\|$,
$\Phi_{c}\left([W,\text{Diag}(X)]\right)$ is piecewise Lipschitz on diagonal
and off-diagonal according to Lemma 2.9.6 , and it converges to 0 almost
surely according to the second part of Lemma 2.9.7.

As all terms converge to 0 in the probability or almost surely, we conclude
that
$\|S_{U}\Phi_{c}\left([W,\text{Diag}(X)]\right)-\Phi_{d}([\widehat{W}_{n\times
n},\text{Diag}(\widetilde{X_{n}})])\|$ converges to 0 in probability. ∎

######  Remark 6.

Note that we can not prove
$S_{n}\cdot\Phi_{c}(\widetilde{W_{n,E}})=\Phi_{d}\cdot
S_{n}(\widetilde{W_{n,E}})$ in general. The difficulty is that starting with
$\widetilde{W_{n,E}}$ of chessboard pattern, after the first layer, pattern
like Figure 2.7(e) may appear in $\sigma\circ T_{1}(\widetilde{W_{n}})$. If
$T_{2}$ is just a average/integral to map $\mathbb{R}^{n^{2}\times{1}}$ to
$\mathbb{R}$, then $S_{n}\circ T_{2}\circ\sigma\circ
T_{1}(\widetilde{W_{n}})=T_{2}\circ\sigma\circ T_{1}(\widetilde{W_{n}})$ will
not be equal to $T_{2}\circ\sigma\circ T_{1}(S_{n}\widetilde{W_{n}})$. The
reason is that both $\sigma\circ T_{1}(\widetilde{W_{n}})$ and $\sigma\circ
T_{1}(S_{n}\widetilde{W_{n}})$ will no longer be of chessboard pattern (Figure
2.7(e) may occur). The diagonal in the $\sigma\circ T_{1}(\widetilde{W_{n}})$
has no effect after taking integral in $T_{2}$ as it is of measure 0. On the
other hand, the diagonal in the matrix $\sigma\circ
T_{1}(S_{n}\widetilde{W_{n}})$ will affect the average. Therefore in general,
$S_{n}\Phi_{c}(\widetilde{W_{n,E}})=\Phi_{d}S_{n}(\widetilde{W_{n,E}})$ does
not hold.

####  2.9.4 IGN-small can Approximate Spectral GNN

Definition of Spectral GNN. The spectral GNN (SGNN) here stands for GNN with
multiple layers of the following form $\forall j=1,\ldots d_{\ell+1}$,

$\displaystyle{\quad}z_{j}^{(\ell+1)}=\sigma\left(\sum_{i=1}^{d_{\ell}}h_{ij}^{(\ell)}(L)z_{i}^{(\ell)}+b_{j}^{(\ell)}1_{n}\right)\in\mathbb{R}^{n}$
(2.14)

where $L=D(A)^{-\frac{1}{2}}AD(A)^{-\frac{1}{2}}$ stands for normalized
adjacency,444We follow the same notation as [78], which is different from the
conventional notation. $z_{j}^{\ell},b_{j}^{\ell}\in\mathbb{R}$ denotes the
embedding and bias at layer $\ell$. $d_{\ell}$ stands for the number of output
channels in $\ell$-th layer.
$h:\mathbb{R}\rightarrow\mathbb{R},h(\lambda)=\sum_{k\geq
0}\beta_{k}\lambda^{k},h(L)=\sum_{k}\beta_{k}L^{k}$, i.e., we apply $h$ to the
eigenvalues of $L$ when it is diagonalizable. Extending $h$ to multiple input
output channels which are indexed in $i$ and $j$, we have
$h_{ij}^{(\ell)}(\lambda)=\sum_{k}\beta_{ijk}^{(\ell)}\lambda^{k}$. By
defining all components of spectral GNN for graphon, the continuous version of
spectral GNN can also be defined. See [78] for details.

We first prove IGN can approximate spectral GNN arbitrarily well, both for
discrete SGNN and continuous SGNN. Next, we show that such IGN belongs to IGN-
small. We need the following simple assumption to ensure the input lies in a
compact domain.

######  AS 5.

There exists an upper bound on $\|x\|_{L_{\infty}}$ for the discrete case and
$\|X\|_{L_{\infty}}$ in the continuous case.

######  AS 6.

$\text{min}(D(A)_{\textnormal{mean}})\geq c_{\textnormal{min}}$ where
$D(A)_{\textnormal{mean}}$ is defined to be
$\frac{1}{n}\text{Diag}(A\mathbf{1})$. The same lower bound holds for graphon
case.

######  Lemma 2.9.8.

Assume AS1-AS6 and $DMD$ arbitrarily well in $L_{\infty}$ sense on a compact
domain

###### Proof.

Given diagonal matrix $D$ and matrix $M$, to implement $DMD$ with linear
equivariant layers of $2$-IGN, we first use operation 14-15 in Table 2.1 to
copy diagonal elements in $D$ to rows and columns of two matrix
$D_{\textnormal{row}}$ and $D_{\textnormal{col}}$. Then calculating $DMD$
becomes entry-wise multiplication of three matrix
$D_{\textnormal{row}},M,D_{\textnormal{col}}$. Assuming all entries of $D$ and
$M$ lies in a compact domain, we can use MLP (which is part of IGN according
to Remark 1) to approximate multiplication arbitrarily well [31, 67]. for
illustration.

To implement $\frac{1}{n}Mx$ with linear equivariant layers of $2$-IGN, first
map $x$ into a diagonal matrix $\text{Diag}(x)$ and concatenate it with $M$ as
the input $[\text{Diag}(x),M]\in\mathbb{R}^{n\times n\times 2}$ to $2$-IGN.
Apply “copy diagonal to all columns” to the first channel and use MLP to
uniformly approximates up to arbitrary precision $\epsilon$ the multiplication
of first channel with the second channel. Then use operation “copy row mean”
to map $\mathbb{R}^{n\times n}\rightarrow\mathbb{R}^{n}$ to get the
$\frac{1}{n}Mx$ within $\epsilon$ precision. See Figure 2.8. ∎

######  Remark 7.

Linear layers in $2$-IGN can not implement matrix-matrix multiplication in
general. When we introduce the matrix multiplication component, the expressive
power of GNN in terms of WL test provably increases from 2-WL to 3-WL [109]).

######  Theorem 2.9.9.

Given $n$, $\epsilon$, and $\textnormal{SGNN}_{\theta_{1}}(n)$, there exists a
$2$-IGN $\textnormal{IGN}_{\theta_{2}}(n)$ such that it approximates
$\textnormal{SGNN}_{\theta_{1}}(n)$ on a compact set (support of input feature
$x_{n}$) arbitrarily well in $L_{\infty}$ sense.

Figure 2.8: An illustration of how we approximate the major building blocks of
SGNN: $\frac{1}{n}Ax$.

###### Proof.

Since IGN and SGNN has the same non-linearity. To show that IGN can
approximate SGNN, it suffices to show that IGN can approximate linear layer of
SGNN, which further boils down to prove that IGN can approximate $Lx$.

Here we assume the input of $2$-IGN is $A\in\mathbb{R}^{n\times n}$ and
$x\in\mathbb{R}^{n\times d}$. We need to first show how
$L=D(A)^{-\frac{1}{2}}AD(A)^{-\frac{1}{2}}$ can be implemented by linear
layers of IGN. This is achieved by noting that
$L=\frac{1}{n}D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}$
where $D(A)_{\textnormal{mean}}$ is normalized degree matrix
$\frac{1}{n}\text{Diag}(A\mathbf{1})$. Representing $L$ as
$\frac{1}{n}D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}$
ensures that all entries in $A$ and $D(A)_{\textnormal{mean}}$ lies in a
compact domain, which is crucial when we extending the approximation proof to
the graphon case.

Now we show how
$Lx=\frac{1}{n}D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}x$
is implemented. First, it is easy to see that $2$-IGN can calculate exactly
$D(A)_{\textnormal{mean}}$ using equivariant layers. Second, as approximating
a) $f(a,b)=ab$ and b) $f(a)=\frac{1}{\sqrt{a}}$ can achieved by MLP on compact
domain, approximating
$D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}$
can also achieved by $2$-IGN layers according to Lemma 2.9.8. Third, we need
to show
$\frac{1}{n}D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}x$
can also be implemented. This is proved in Lemma 2.9.8.

There are two main functions we need to approximate with MLP: a)
$f(x)=1/\sqrt{a}$ and b) $f(a,b)=ab$.

For a) the input is entries of $D(A)_{\textnormal{mean}}$ which lie in
$[0,1]$. By classical universal approximation theorem [31, 67], we know MLP
can approximate a) arbitrarily well.

For b) the input is $(D(A)_{\textnormal{mean}}^{-1/2},A)$ for normalized
adjacency matrix calculation, and $(L,x)$ for graph signal convolution.

To ensure the uniform approximation, we need to ensure all of them lie in a
compact domain. This is indeed the case as all entries in
$D(A)_{\textnormal{mean}},A,x$ are all upper bounded

  1. 1.

every entry in $A$ is either 0 or 1 therefore lies in a compact domain.

  2. 2.

similarly, all entries $D(A)_{\textnormal{mean}}$ lies in $[c_{\text{min}},1]$
by AS6, and therefore $D(A)_{\textnormal{mean}}^{-\frac{1}{2}}$ also lies in a
compact domain. As $L(A)$ is the multiplication of
$D(A)_{\textnormal{mean}}^{-1/2},A,D(A)_{\textnormal{mean}}^{-1/2}$, every
entry of $L(A)$ also lies in compact domain.

  3. 3.

input signal $x$ has bounded $l_{\infty}$-norm by assumption AS5.

  4. 4.

all coefficient for operators is upper bounded and independent from $n$ by
AS2.

Since we showed the $L(A)x$ can be approximated arbitrarily well by IGN,
repeating such processes and leveraging the fact that $L$ has bounded spectral
norm, we can then approximate $L^{k}(A)x$ up to $\epsilon$ precision. The
errors $\epsilon$ depend on the approximation error of the MLP to the relevant
function, the previous errors, and uniform bounds as well as uniform
continuity of the approximated functions. ∎

######  Theorem 2.9.10.

Given $\epsilon$, and a spectral GNN c$\textnormal{SGNN}_{\theta_{1}}$, there
exists a continuous $2$-IGN c$\textnormal{IGN}_{\theta_{2}}$such that it
approximates c$\textnormal{SGNN}_{\theta_{1}}$ on a compact set (input feature
$X$) arbitrarily well.

###### Proof.

In the continuous case,
$Lx=\frac{1}{n}D(A)_{\textnormal{mean}}^{-\frac{1}{2}}AD(A)_{\textnormal{mean}}^{-\frac{1}{2}}x$
in the discrete case will be replaced with
$D(W)^{-\frac{1}{2}}WD(W)^{-\frac{1}{2}}X$ where $D(W)$ is a diagonal graphon
defined to be $D(W)(i,i)=\int_{0}^{1}W(i,j)dj$.

We show that all items listed in proof of Theorem 2.9.9 still holds in the
continuous case

  * •

we consider the $W$ instead in the continuous case, where all entries still
lies in a compact domain $[0,1]$.

  * •

similarly all entries of the continuous analog of
$D(A)_{\textnormal{mean}},D(A)_{\textnormal{mean}}^{-\frac{1}{2}}$, and $T(W)$
also lies in a compact domain according to AS6.

  * •

the statements about input signal $X$ and the coefficient for linear
equivariant operators also holds in the continuous setting.

Therefore we conclude the proof. Now we are ready to prove that those IGN that
can approximate SGNN well is a subset of IGN-small. ∎

######  Lemma 2.9.11.

With slight abuse of notation, let $\widetilde{W_{n,E}}$ be graphon of
chessboard pattern. Let $\widetilde{X_{n,E}}$ be a graphon signal with 1D
chessboard pattern.
$S_{n}\circ\widetilde{W_{n,E}}\widetilde{X_{n,E}}=(S_{n}\widetilde{W_{n,E}})(S_{n}\widetilde{X_{n,E}})$.

###### Proof.

Since
$S_{n}\circ\widetilde{W_{n,E}}\widetilde{X_{n,E}}=S_{n}\circ\int_{j\in[0,1]}\widetilde{W_{n,E}}(i,j)\widetilde{X_{n,E}}(j)dj=\left(...,\frac{1}{\sqrt{n}}\int_{j\in[0,1]}\widetilde{W_{n,E}}(\frac{i}{n},j)\widetilde{X_{n,E}}(j),...\right)$,
it suffices to analyze $i$-th component
$\frac{1}{\sqrt{n}}\int_{j\in[0,1]}\widetilde{W_{n,E}}(\frac{i}{n},j)\widetilde{X_{n,E}}(j)$.

Since $\widetilde{W_{n,E}},\widetilde{X_{n,E}}$ are of chessboard pattern, we
can replace integral with summation.

$\displaystyle S_{n}\circ\widetilde{W_{n,E}}\widetilde{X_{n,E}}(i)$
$\displaystyle=\frac{1}{\sqrt{n}}\int_{j\in[0,1]}\widetilde{W_{n,E}}(\frac{i}{n},j)\widetilde{X_{n,E}}(j)$
$\displaystyle=\frac{1}{\sqrt{n}}\frac{1}{n}\sum_{j\in[n]}\widetilde{W_{n,E}}(\frac{i}{n},\frac{j}{n})\widetilde{X_{n,E}}(\frac{j}{n})$
$\displaystyle=\sum_{j\in[n]}\frac{1}{n}\widetilde{W_{n,E}}(\frac{i}{n},\frac{j}{n})(S_{n}\widetilde{X_{n,E}})(j)$
$\displaystyle=\sum(S_{n}\widetilde{W_{n,E}})(i,j)(S_{n}\widetilde{X_{n,E}})(j)$
$\displaystyle=\left((S_{n}\widetilde{W_{n,E}})(S_{n}\widetilde{X_{n,E}})\right)(i)$

Which concludes the proof. Note that our proof does make use of the property
of multiplication between two numbers. ∎

######  Remark 8.

The whole proof only relies on that $\widetilde{W_{n,E}}$ and
$\widetilde{X_{n,E}}$ have checkerboard patterns. Therefore replacing the
multiplication with other operations (such as a MLP) will still hold.

See 2.5.3

###### Proof.

To prove this, we only need to show that
$S_{n}\Phi_{c,\textnormal{approx}}([\widetilde{W_{n,E}},\widetilde{f_{n,E}}])=\\\
\Phi_{d,\textnormal{approx}}S_{n}([\widetilde{W_{n,E}},\widetilde{f_{n,E}}])$.
Here $\Phi_{c,\textnormal{approx}}$ and $\Phi_{d,\textnormal{approx}}$ denotes
those specific IGN in Theorems 2.9.9 and 2.9.10 constructed to approximate
SGNN.

To build up some intuition, let $\Phi_{\textnormal{SGNN}}$ denotes the
spectral GNN that $\Phi_{\textnormal{approx}}$ approximates. it is easy to see
that
$S_{n}\Phi_{c,\textnormal{SGNN}}([\widetilde{W_{n,E}},\widetilde{f_{n,E}}])=\Phi_{d,\textnormal{SGNN}}S_{n}([\widetilde{W_{n,E}},\widetilde{f_{n,E}}])$
due to Lemma 2.9.11 and Lemma 2.9.6.2. To show the same holds for
$\Phi_{\textnormal{approx}}$, note that the only difference between
$\widetilde{W_{n,E}}\widetilde{f_{n,E}}$ implemented by SGNN and approximated
by $\Phi_{\textnormal{approx}}$ is that $\Phi_{\textnormal{approx}}$ use MLP
to simulate multiplication between numbers. According to Remark 8, the
approximated version of $\widetilde{W_{n,E}}\widetilde{f_{n,E}}$ still
commutes with $S_{n}$.

Since nonlinear layer $\sigma$ in $\Phi_{\textnormal{approx}}$ also commutes
with $S_{n}$ according to Lemma 2.9.6.2, we can combine the result above and
conclude that $\Phi_{\textnormal{approx}}$ commutes with $S_{n}$. Therefore
$\Phi_{\textnormal{approx}}$ belongs to IGN-small, which finishes the proof. ∎

##  Chapter 3 On the Connection Between MPNN and Graph Transformer

###  3.1 Introduction

In this chapter, we study the connection between MPNN and Graph Transformer.
MPNN (Message Passing Neural Network) [54] has been the leading architecture
for processing graph-structured data. Recently, transformers in natural
language processing [142, 75] and vision [45, 59] have extended their success
to the domain of graphs. There have been several pieces of work [159, 153, 86,
120, 82] showing that with careful position embedding [100], graph
transformers (GT) can achieve compelling empirical performances on large-scale
datasets and start to challenge the dominance of MPNN.

Figure 3.1: MPNN + VN and Graph Transformers. Table 3.1: Summary of
approximation result of MPNN + VN on self-attention layer. $n$ is the number
of nodes and $d$ is the feature dimension of node features. The dependency on
$d$ is hidden.

| Depth | Width | Self-Attention | Note
---|---|---|---|---
Theorem 3.3.1 | $O(1)$ | $O(1)$ | Approximate | Approximate self attention in Performer [29]
Theorem 3.4.2 | $O(1)$ | $O(n^{d})$ | Full | Leverage the universality of equivariant DeepSets
Theorem 3.5.2 | $O(n)$ | $O(1)$ | Full | Explicit construction, strong assumption on $\mathcal{X}$
Proposition 3.10.6 | $O(n)$ | $O(1)$ | Full | Explicit construction, more relaxed (but still strong) assumption on $\mathcal{X}$

MPNN imposes a sparsity pattern on the computation graph and therefore enjoys
linear complexity. It however suffers from well-known over-smoothing [96, 114,
21] and over-squashing [3, 140] issues, limiting its usage on long-range
modeling tasks where the label of one node depends on features of nodes far
away. GT relies purely on position embedding to encode the graph structure and
uses vanilla transformers on top. 111GT in this chapter refers to the practice
of tokenizing graph nodes and applying standard transformers on top [159, 82].
There exists a more sophisticated GT [86] that further conditions attention on
edge types but it is not considered in this chapter.  It models all pairwise
interactions directly in one layer, making it computationally more expensive.
Compared to MPNN, GT shows promising results on tasks where modeling long-
range interaction is the key, but the quadratic complexity of self-attention
in GT limits its usage to graphs of medium size. Scaling up GT to large graphs
remains an active research area [152].

Theoretically, it has been shown that graph transformers can be powerful graph
learners [82], i.e., graph transformers with appropriate choice of token
embeddings have the capacity of approximating linear permutation equivariant
basis, and therefore can approximate 2-IGN (Invariant Graph Network), a
powerful architecture that is at least as expressive as MPNN [110]. This
raises an important question that whether GT is strictly more powerful than
MPNN. Can we approximate GT with MPNN?

One common intuition of the advantage of GT over MPNN is its ability to model
long-range interaction more effectively. However, from the MPNN side, one can
resort to a simple trick to escape locality constraints for effective long-
range modeling: the use of an additional _virtual node (VN)_ that connects to
all input graph nodes. On a high level, MPNN + VN augments the existing graph
with one virtual node, which acts like global memory for every node exchanging
messages with other nodes. Empirically this simple trick has been observed to
improve the MPNN and has been widely adopted [54, 69, 68] since the early
beginning of MPNN [54, 8]. However, there is very little theoretical study of
MPNN + VN [72].

In this work, we study the theoretical property of MPNN + VN, and its
connection to GT. We systematically study the representation power of MPNN +
VN, both for certain approximate self-attention and for the full self-
attention layer, and provide a depth-width trade-off, summarized in Table 3.1.
In particular,

  * •

With $O(1)$ depth and $O(1)$ width, MPNN + VN can approximate one self-
attention layer of Performer [29] and Linear Transformer [77], a type of
linear transformers [139].

  * •

Via a link between MPNN + VN with DeepSets [162], we prove MPNN + VN with
$O(1)$ depth and $O(n^{d})$ width ($d$ is the input feature dimension) is
permutation equivariant universal, implying it can approximate self-attention
layer and even full-transformers.

  * •

Under certain assumptions on node features, we prove an explicit construction
of $O(n)$ depth $O(1)$ width MPNN + VN approximating 1 self-attention layer
arbitrarily well on graphs of size $n$. Unfortunately, the assumptions on node
features are rather strong, and whether we can alleviate them will be an
interesting future direction to explore.

  * •

Empirically, we show 1) that MPNN + VN works surprisingly well on the recently
proposed LRGB (long-range graph benchmarks) datasets [44], which arguably
require long-range interaction reasoning to achieve strong performance 2) our
implementation of MPNN + VN is able to further improve the early
implementation of MPNN + VN on OGB datasets and 3) MPNN + VN outperforms
Linear Transformer [77] and MPNN on the climate modeling task.

###  3.2 Preliminaries

We denote ${\bm{X}}\in\mathbb{R}^{n\times d}$ the concatenation of graph node
features and positional encodings, where node $i$ has feature
${\bm{x}}_{i}\in\mathbb{R}^{d}$. When necessary, we use ${\bm{x}}^{(l)}_{j}$
to denote the node $j$’s feature at depth $l$. Let $\mathcal{M}$ be the space
of multisets of vectors in $\mathbb{R}^{d}$. We use
$\mathcal{X}\subseteq\mathbb{R}^{n\times d}$ to denote the space of node
features and the $\mathcal{X}_{i}$ be the projection of $\mathcal{X}$ on
$i$-th coordinate. $\|\cdot\|$ denotes the 2-norm.
$[{\bm{x}},{\bm{y}},{\bm{z}}]$ denotes the concatenation of
${\bm{x}},{\bm{y}},{\bm{z}}$. $[n]$ stands for the set $\\{1,2,...,n\\}$.

######  Definition 13 (attention).

We denote key and query matrix as
${\bm{W}}_{K},{\bm{W}}_{Q}\in\mathbb{R}^{d\times d^{\prime}}$, and value
matrix as ${\bm{W}}_{V}\in\mathbb{R}^{d\times d}$ 222For simplicity, we assume
the output dimension of self-attention is the same as the input dimension. All
theoretical results can be extended to the case where the output dimension is
different from $d$.. Attention score between two vectors
${\bm{u}},{\bm{v}}\in\mathbb{R}^{d\times 1}$ is defined as
$\alpha({\bm{u}},{\bm{v}})=\text{softmax}({\bm{u}}^{T}{\bm{W}}_{Q}({\bm{W}}_{K})^{T}{\bm{v}})$.
We denote $\mathcal{A}$ as the space of attention $\alpha$ for different
${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}$. We also define unnormalized
attention score $\alpha^{\prime}(\cdot,\cdot)$ to be
$\alpha^{\prime}({\bm{u}},{\bm{v}})={\bm{u}}^{T}{\bm{W}}_{Q}({\bm{W}}_{K})^{T}{\bm{v}}$.
Self attention layer is a matrix function ${\bm{L}}:\mathbb{R}^{n\times
d}\rightarrow\mathbb{R}^{n\times d}$ of the following form:
${\bm{L}}({\bm{X}})=\text{softmax}({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{T}){\bm{X}}{\bm{W}}_{V}$.

####  3.2.1 MPNN Layer

######  Definition 14 (MPNN layer [54]).

An MPNN layer on a graph $G$ with node features ${\bm{x}}^{(k)}$ at $k$-th
layer and edge features ${\bm{e}}$ is of the following form

${\bm{x}}_{i}^{(k)}=\gamma^{(k)}\left({\bm{x}}_{i}^{(k-1)},\tau_{j\in\mathcal{N}(i)}\phi^{(k)}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i}\right)\right)$

Here
$\gamma:\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}\rightarrow\mathbb{R}^{d}$
is update function,
$\phi:\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d_{e}}\rightarrow\mathbb{R}^{d^{\prime}}$
is message function where $d_{e}$ is the edge feature dimension,
$\tau:\mathcal{M}\rightarrow\mathbb{R}^{d}$ is permutation invariant
aggregation function and $\mathcal{N}(i)$ is the neighbors of node $i$ in $G$.
Update/message/aggregation functions are usually parametrized by neural
networks. For graphs of different types of edges and nodes, one can further
extend MPNN to the heterogeneous setting. We use $1,...,n$ to index graph
nodes and ${\bm{n}}$ to denote the virtual node.

######  Definition 15 (heterogeneous MPNN + VN layer).

The heterogeneous MPNN + VN layer operates on two types of nodes: 1) virtual
node and 2) graph nodes, denoted as vn and gn, and three types of edges: 1)
vn-gn edge and 2) gn-gn edges and 3) gn-vn edges. It has the following form

${\bm{x}}_{{\bm{n}}}^{(k)}=\gamma_{{\bm{n}}}^{(k)}\left({\bm{x}}_{i}^{(k-1)},\tau_{j\in[n]}\phi^{(k)}_{\text{vn-
gn}}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i}\right)\right)$
(3.1)

for the virtual node, and

$\begin{split}{\bm{x}}_{i}^{(k)}&=\gamma_{\text{gn}}^{(k)}({\bm{x}}_{i}^{(k-1)},\tau_{j\in\mathcal{N}_{1}(i)}\phi^{(k)}_{\text{gn-
vn}}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i}\right)\\\
&+\tau_{j\in\mathcal{N}_{2}(i)}\phi^{(k)}_{\text{gn-
gn}}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i})\right)\end{split}$
(3.2)

for graph node. Here $\mathcal{N}_{1}(i)$ for graph node $i$ is the virtual
node and $\mathcal{N}_{2}(i)$ is the set of neighboring graph nodes.

Our proof of approximating self-attention layer ${\bm{L}}$ with MPNN layers
does not use the graph topology. Next, we introduce a simplified heterogeneous
MPNN + VN layer, which will be used in the proof. It is easy to see that
setting $\phi^{(k)}_{\text{gn-gn}}$ to be 0 in Definition 15 recovers the
simplified heterogeneous MPNN + VN layer.

######  Definition 16 (simplified heterogeneous MPNN + VN layer).

A simplified heterogeneous MPNN + VN layer is the same as a heterogeneous MPNN
+ VN layer in Definition 15 except we set $\theta_{\text{gn-gn}}$ to be 0.
I.e., we have

${\bm{x}}_{{\bm{n}}}^{(k)}=\gamma_{{\bm{n}}}^{(k)}\left({\bm{x}}_{i}^{(k-1)},\tau_{j\in[n]}\phi^{(k)}_{\text{vn-
gn}}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i}\right)\right)$

for the virtual node, and

${\bm{x}}_{i}^{(k)}=\gamma_{\text{gn}}^{(k)}\left({\bm{x}}_{i}^{(k-1)},\tau_{j\in\mathcal{N}_{1}(i)}\phi^{(k)}_{\text{gn-
vn}}\left({\bm{x}}_{i}^{(k-1)},{\bm{x}}_{j}^{(k-1)},{\bm{e}}_{j,i}\right)\right)$

for graph nodes.

Intuitively, adding the virtual node (VN) to MPNN makes it easy to compute
certain quantities, for example, the mean of node features (which is hard for
standard MPNN unless the depth is proportional to the diameter of the graph).
Using VN thus makes it easy to implement for example the mean subtraction,
which helps reduce over-smoothing and improves the performance of GNN [158,
166]. See more connection between MPNN + VN and over-smoothing in Section
3.6.4.

####  3.2.2 Assumptions

We have two mild assumptions on feature space
$\mathcal{X}\subset\mathbb{R}^{n\times d}$ and the regularity of target
function ${\bm{L}}$.

######  AS 1.

$\forall i\in[n],{\bm{x}}_{i}\in\mathcal{X}_{i},\|{\bm{x}}_{i}\|<C_{1}$. This
implies $\mathcal{X}$ is compact.

######  AS 2.

$\|{\bm{W}}_{Q}\|<C_{2},\|{\bm{W}}_{K}\|<C_{2},\|{\bm{W}}_{V}\|<C_{2}$ for
target layer ${\bm{L}}$. Combined with AS1 on $\mathcal{X}$, this means
$\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})$ is both upper and lower bounded,
which further implies $\sum_{j}e^{\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})}$
be both upper bounded and lower bounded.

####  3.2.3 Notations

We provide a notation table for references.

Table 3.2: Summary of important notations.

Symbol | Meaning
---|---
${\bm{X}}\in\mathcal{X}\subset\mathbb{R}^{n\times d}$ | graph node features
${\bm{x}}_{i}\in\mathbb{R}^{1\times d}$ | graph node $i$’s feature
$\tilde{{\bm{x}}}_{i}\in\mathbb{R}^{1\times d}$ | approximated graph node $i$’s feature via attention selection
$\mathcal{M}$ | A multiset of vectors in $\mathbb{R}^{d}$
${\bm{W}}_{Q}^{(l)},{\bm{W}}_{K}^{(l)},{\bm{W}}_{V}^{(l)}\in\mathbb{R}^{d\times d^{\prime}}$ | attention matrix of $l$-th self-attention layer in graph transformer
$\mathcal{X}$ | feature space
$\mathcal{X}_{i}$ | projection of feature space onto $i$-th coordinate
${\bm{L}}^{\text{ds}}_{i}$ |  $i$-th linear permutation equivariant layer in DeepSets
${\bm{L}},{\bm{L}}^{\prime}$ | full self attention layer; approximate self attention layer in Performer
${\bm{z}}_{{\bm{n}}}^{(l)},{\bm{z}}_{i}^{(l)}$ | virtual/graph node feature at layer $l$ of heterogeneous MPNN + VN
$\alpha_{{\bm{n}}}$ | attention score in MPNN + VN
$\alpha(\cdot,\cdot)$ | normalized attention score
$\alpha_{\text{GATv2}}(\cdot,\cdot)$ | normalized attention score with GATv2
$\alpha^{\prime}(\cdot,\cdot)$ | unnormalized attention score. $\alpha^{\prime}({\bm{u}},{\bm{v}})={\bm{u}}{\bm{W}}_{Q}({\bm{W}}_{K})^{T}{\bm{v}}^{T}$
$\alpha^{\prime}_{\text{GATv2}}(\cdot,\cdot)$ | unnormalized attention score with GATv2. $\alpha^{\prime}_{\text{GATv2}}({\bm{u}},{\bm{v}}):={\bm{a}}^{T}\operatorname{LeakyReLU}\left({\bm{W}}\cdot\left[{\bm{u}}\|{\bm{v}}\right]+{\bm{b}}\right)$
$\mathcal{A}$ | space of attentions, where each element $\alpha\in\mathcal{A}$ is of form $\alpha({\bm{u}},{\bm{v}})=\text{softmax}({\bm{u}}{\bm{W}}_{Q}({\bm{W}}_{K})^{T}{\bm{v}}^{T})$
$C_{1}$ | upper bound on norm of all node features $\|{\bm{x}}_{i}\|$
$C_{2}$ | upper bound on the norm of ${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}$ in target ${\bm{L}}$
$C_{3}$ | upper bound on the norm of attention weights of $\alpha_{{\bm{n}}}$ when selecting ${\bm{x}}_{i}$
$\gamma^{(k)}(\cdot,\cdot)$ | update function
$\theta^{(k)}(\cdot,\cdot)$ | message function
$\tau(\cdot)$ | aggregation function

###  3.3 $O(1)$-Depth $O(1)$-Width MPNN + VN for Unbiased Approximation of
Attention

The standard self-attention takes $O(n^{2})$ computational time, therefore not
scalable for large graphs. Reducing the computational complexity of self-
attention in Transformer is active research [139]. In this section, we
consider self-attention in a specific type of efficient transformers,
Performer [29] and Linear Transformer [77].

One full self-attention layer ${\bm{L}}$ is of the following form

${\bm{x}}_{i}^{(l+1)}=\sum_{j=1}^{n}\frac{\kappa\left({\bm{W}}_{Q}^{(l)}{\bm{x}}_{i}^{(l)},{\bm{W}}_{K}^{(l)}{\bm{x}}_{j}^{(l)}\right)}{\sum_{k=1}^{n}\kappa\left({\bm{W}}_{Q}^{(l)}{\bm{x}}_{i}^{(l)},{\bm{W}}_{K}^{(l)}{\bm{x}}_{k}^{(l)}\right)}\cdot\left({\bm{W}}_{V}^{(l)}{\bm{x}}_{j}^{(l)}\right)$
(3.3)

where $\kappa:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$ is the
softmax kernel $\kappa({\bm{x}},{\bm{y}}):=\exp({\bm{x}}^{T}{\bm{y}})$. The
kernel function can be approximated via
$\kappa({\bm{x}},{\bm{y}})=\langle{\Phi({\bm{x}}),\Phi({\bm{y}})}\rangle_{\mathcal{V}}\approx\phi({\bm{x}})^{T}\phi({\bm{y}})$
where the first equation is by Mercer’s theorem and
$\phi(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}^{m}$ is a low-dimensional
feature map with random transformation. For Performer [29], the choice of
$\phi$ is taken as
$\phi({\bm{x}})=\frac{\exp\left(\frac{-\|{\bm{x}}\|_{2}^{2}}{2}\right)}{\sqrt{m}}\left[\exp\left({\bm{w}}_{1}^{T}{\bm{x}}\right),\cdots,\exp\left({\bm{w}}_{m}^{T}{\bm{x}}\right)\right]$
where ${\bm{w}}_{k}\sim\mathcal{N}\left(0,I_{d}\right)$ is i.i.d sampled
random variable. For Linear Transformer [77],
$\phi({\bm{x}})=\operatorname{elu}({\bm{x}})+1$.

By switching $\kappa({\bm{x}},{\bm{y}})$ to be
$\phi({\bm{x}})^{T}\phi({\bm{y}})$, and denote
${\bm{q}}_{i}={\bm{W}}_{Q}^{(l)}{\bm{x}}_{i}^{(l)},{\bm{k}}_{i}={\bm{W}}_{K}^{(l)}{\bm{x}}_{i}^{(l)}\text{
and }{\bm{v}}_{i}={\bm{W}}_{V}^{(l)}{\bm{x}}_{i}^{(l)}$, the approximated
version of Equation 3.3 by Performer and Linear Transformer becomes

$\begin{split}{\bm{x}}_{i}^{(l+1)}&=\sum_{j=1}^{n}\frac{\phi\left({\bm{q}}_{i}\right)^{T}\phi\left({\bm{k}}_{j}\right)}{\sum_{k=1}^{n}\phi\left({\bm{q}}_{i}\right)^{T}\phi\left({\bm{k}}_{k}\right)}\cdot{\bm{v}}_{j}\\\
&=\frac{\left(\phi\left({\bm{q}}_{i}\right)^{T}\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j}\right)^{T}}{\phi\left({\bm{q}}_{i}\right)^{T}\sum_{k=1}^{n}\phi\left({\bm{k}}_{k}\right)}.\\\
\end{split}$ (3.4)

where we use the matrix multiplication association rule to derive the second
equality.

The key advantage of Equation 3.4 is that
$\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)$ and
$\sum_{j=1}^{n}\phi({\bm{k}}_{j})\otimes{\bm{v}}_{j}$ can be approximated by
the virtual node, and shared for all graph nodes, using only $O(1)$ layers of
MPNNs. We denote the self-attention layer of this form in Equation 3.4 as
${\bm{L}}_{\text{Performer}}$. Linear Transformer differs from Performer by
choosing a different form of $\phi({\bm{x}})=\operatorname{Relu}({\bm{x}})+1$
in its self-attention layer ${\bm{L}}_{\text{Linear-Transformer}}$.

In particular, the VN will approximate
$\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)$ and
$\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j}$, and
represent it as its feature. Both $\phi\left({\bm{k}}_{j}\right)$ and
$\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j}$ can be approximated
arbitrarily well by an MLP with constant width (constant in $n$ but can be
exponential in $d$) and depth. Note that
$\phi({\bm{k}}_{j})\otimes{\bm{v}}_{j}\in\mathbb{R}^{dm}$ but can be reshaped
to 1 dimensional feature vector.

More specifically, the initial feature for the virtual node is
$\bm{1}_{(d+1)m}$, where $d$ is the dimension of node features and $m$ is the
number of random projections $\omega_{i}$. Message function + aggregation
function for virtual node $\tau\phi_{\text{vn-
gn}}:\mathbb{R}^{(d+1)m}\times\mathcal{M}\rightarrow\mathbb{R}^{(d+1)m}$ is

$\begin{split}&\tau_{j\in[n]}\phi_{\text{vn-
gn}}^{(k)}(\cdot,\\{{\bm{x}}_{i}\\}_{i})=[\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right),\\\
&\texttt{ReshapeTo1D}{(\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j})}]\end{split}$
(3.5)

where $\texttt{ReshapeTo1D}{(\cdot)}$ flattens a 2D matrix to a 1D vector in
raster order. This function can be arbitrarily approximated by MLP. Note that
the virtual node’s feature dimension is $(d+1)m$ (where recall $m$ is the
dimension of the feature map $\phi$ used in the linear transformer/Performer),
which is larger than the dimension of the graph node $d$. This is consistent
with the early intuition that the virtual node might be overloaded when
passing information among nodes. The update function for virtual node
$\gamma_{{\bm{n}}}:$
$\mathbb{R}^{(d+1)m}\times\mathbb{R}^{(d+1)m}\rightarrow\mathbb{R}^{(d+1)m}$
is just coping the second argument, which can be exactly implemented by MLP.

VN then sends its message back to all other nodes, where each graph node $i$
applies the update function
$\gamma_{\text{gn}}:\mathbb{R}^{(d+1)m}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$
of the form

$\begin{split}&\gamma_{\text{gn}}({\bm{x}}_{i},[\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right),\texttt{ReshapeTo1D}{(\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j})}])\\\
&=\frac{\left(\phi\left({\bm{q}}_{i}\right)\sum_{j=1}^{n}\phi\left({\bm{k}}_{j}\right)\otimes{\bm{v}}_{j}\right)^{T}}{\phi\left({\bm{q}}_{i}\right)^{T}\sum_{k=1}^{n}\phi\left({\bm{k}}_{k}\right)}\end{split}$
(3.6)

to update the graph node feature.

As the update function $\gamma_{\text{gn}}$ can not be computed exactly in
MLP, what is left is to show that error induced by using MLP to approximate
$\tau\phi_{\text{vn-gn}}$ and $\gamma_{\text{gn}}$ in Equation 3.5 and
Equation 3.6 can be made arbitrarily small.

######  Theorem 3.3.1.

Under the AS1 and AS2, MPNN + VN of $O(1)$ width and $O(1)$ depth can
approximate ${\bm{L}}_{\text{Performer}}$ and ${\bm{L}}_{\text{Linear-
Transformer}}$ arbitrarily well.

###### Proof.

We first prove the case of ${\bm{L}}_{\text{Performer}}$. We can decompose our
target function as the composition of $\tau_{j\in[n]}\phi_{\text{vn-
gn}}^{(k)}$, $\gamma_{\text{gn}}$ and $\phi$. By the uniform continuity of the
functions, it suffices to show that 1) we can approximate $\phi$, 2) we can
approximate operations in $\gamma_{\text{gn}}$ and $\tau\phi_{\text{vn-gn}}$
arbitrarily well on the compact domain, and 3) the denominator
$\phi\left({\bm{q}}_{i}\right)^{T}\sum_{k=1}^{n}\phi\left({\bm{k}}_{k}\right)$
is uniformly lower bounded by a positive number for any node features in
$\mathcal{X}$.

For 1), each component of $\phi$ is continuous and all inputs
${\bm{k}}_{j},\mathbf{q}_{j}$ lie in the compact domain so $\phi$ can be
approximated arbitrarily well by MLP with $O(1)$ width and $O(1)$ depth [31].

For 2), we need to approximate the operations in $\gamma_{\text{gn}}$ and
$\tau\phi_{\text{vn-gn}}$, i.e., approximate multiplication, and vector-scalar
division arbitrarily well. As all those operations are continuous, it boils
down to showing that all operands lie in a compact domain. By assumption AS1
and AS2 on ${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}$ and input feature
$\mathcal{X}$, we know that ${\bm{q}}_{i},{\bm{k}}_{i},{\bm{v}}_{i}$ lies in a
compact domain for all graph nodes $i$. As $\phi$ is continuous, this implies
that $\phi({\bm{q}}_{i}),\sum_{j=1}^{n}\phi({\bm{k}}_{j})\otimes{\bm{v}}_{j}$
lies in a compact domain ($n$ is fixed), therefore the numerator lies in a
compact domain. Lastly, since all operations do not involve $n$, the depth and
width are constant in $n$.

For 3), it is easy to see that
$\phi\left({\bm{q}}_{i}\right)^{T}\sum_{k=1}^{n}\phi\left({\bm{k}}_{k}\right)$
is always positive. We just need to show that the denominator is bound from
below by a positive constant. For Performer,
$\phi({\bm{x}})=\frac{\exp\left(\frac{-\|{\bm{x}}\|_{2}^{2}}{2}\right)}{\sqrt{m}}\left[\exp\left({\bm{w}}_{1}^{T}{\bm{x}}\right),\cdots,\exp\left({\bm{w}}_{m}^{T}{\bm{x}}\right)\right]$
where ${\bm{w}}_{k}\sim\mathcal{N}\left(0,I_{d}\right)$. As all norm of input
${\bm{x}}$ to $\phi$ is upper bounded by AS1,
$\exp(\frac{-\|{\bm{x}}\|_{2}^{2}}{2})$ is lower bounded. As $m$ is fixed, we
know that $\|{\bm{w}}^{T}_{i}{\bm{x}}\|\leq\|{\bm{w}}_{i}\|\|{\bm{x}}\|$,
which implies that ${\bm{w}}_{i}^{T}{\bm{x}}$ is lower bounded by
$-\|{\bm{w}}_{i}\|\|{\bm{x}}\|$ which further implies that
$\exp({\bm{w}}_{i}^{T}{\bm{x}})$ is lower bounded. This means that
$\phi\left({\bm{q}}_{i}\right)^{T}\sum_{k=1}^{n}\phi\left({\bm{k}}_{k}\right)$
is lower bounded.

For Linear Transformer, the proof is essentially the same as above. We only
need to show that $\phi({\bm{x}})=\operatorname{elu}({\bm{x}})+1$ is
continuous and positive, which is indeed the case. ∎

Besides Performers, there are many other different ways of obtaining linear
complexity. In Section 3.7.2, we discuss the limitation of MPNN + VN on
approximating other types of efficient transformers such as Linformer [147]
and Sparse Transformer [28].

###  3.4 $O(1)$ Depth $O(n^{d})$ Width MPNN + VN

Figure 3.2: The link between MPNN and GT is drawn via DeepSets in Section 3.4
of this chapter and Invariant Graph Network (IGN) in [82]. Interestingly, IGN
is a generalization of DeepSets [110].

We have shown that the MPNN + VN can approximate self-attention in Performer
and Linear Transformer using only $O(1)$ depth and $O(1)$ width. One may
naturally wonder whether MPNN + VN can approximate the self-attention layer in
the full transformer. In this section, we show that MPNN + VN with $O(1)$
depth (number of layers), but with $O(n^{d})$ width, can approximate 1 self-
attention layer (and full transformer) arbitrarily well.

The main observation is that MPNN + VN is able to exactly simulate (not just
approximate) equivariant DeepSets [162], which is proved to be universal in
approximating any permutation invariant/equivariant maps [162, 131]. Since the
self-attention layer is permutation equivariant, this implies that MPNN + VN
can approximate the self-attention layer (and full transformer) with $O(1)$
depth and $O(n^{d})$ width following a result on DeepSets from [131].

We first introduce the permutation equivariant map, equivariant DeepSets, and
permutation equivariant universality.

######  Definition 17 (permutation equivariant map).

A map ${\bm{F}}:\mathbb{R}^{n\times k}\rightarrow\mathbb{R}^{n\times l}$
satisfying ${\bm{F}}(\sigma\cdot{\bm{X}})=\sigma\cdot{\bm{F}}({\bm{X}})$ for
all $\sigma\in S_{n}$ and ${\bm{X}}\in\mathbb{R}^{n\times d}$ is called
permutation equivariant.

######  Definition 18 (equivariant DeepSets of [162]).

Equivariant DeepSets has the following form
${\bm{F}}({\bm{X}})={\bm{L}}_{m}^{\text{ds}}\circ\nu\circ\cdots\circ\nu\circ{\bm{L}}^{\text{ds}}_{1}({\bm{X}})$,
where ${\bm{L}}^{\text{ds}}_{i}$ is a linear permutation equivariant layer and
$\nu$ is a nonlinear layer such as ReLU. The linear permutation equivariant
layer in DeepSets has the following form
${\bm{L}}^{\text{ds}}_{i}({\bm{X}})={\bm{X}}\bm{A}+\frac{1}{n}\mathbf{11}^{T}{\bm{X}}{\bm{B}}+\mathbf{1}\bm{c}^{T}$,
where ${\bm{A}},{\bm{B}}\in\mathbb{R}^{d_{i}\times d_{i+1}}$,
${\bm{c}}\in\mathbb{R}^{d_{i+1}}$ is the weights and bias in layer $i$, and
$\nu$ is ReLU.

######  Definition 19 (permutation equivariant universality).

Given a compact domain $\mathcal{X}$ of $\mathbb{R}^{n\times d_{\text{in}}}$,
permutation equivariant universality of a model ${\bm{F}}:\mathbb{R}^{n\times
d_{\text{in}}}\rightarrow\mathbb{R}^{n\times d_{\text{out}}}$ means that for
every permutation equivariant continuous function
${\bm{H}}:\mathbb{R}^{n\times d_{\text{in}}}\rightarrow\mathbb{R}^{n\times
d_{\text{out}}}$ defined over $\mathcal{X}$, and any $\epsilon>0$, there
exists a choice of $m$ (i.e., network depth), $d_{i}$ (i.e., network width at
layer $i$) and the trainable parameters of ${\bm{F}}$ so that
$\|{\bm{H}}({\bm{X}})-{\bm{F}}({\bm{X}})\|_{\infty}<\epsilon$ for all
${\bm{X}}\in\mathcal{X}$.

The universality of equivariant DeepSets is stated as follows.

######  Theorem 3.4.1 ([131]).

DeepSets with constant layer is universal. Using ReLU activation the width
$\omega:=\text{max}_{i}d_{i}$ ($d_{i}$ is the width for $i$-th layer of
DeepSets) required for universal permutation equivariant network satisfies
$\omega\leq
d_{\text{out}}+d_{\text{in}}+\left(\begin{array}[]{c}n+d_{\text{in}}\\\
d_{\text{in}}\end{array}\right)=O(n^{d_{\text{in}}})$.

We are now ready to state our main theorem.

######  Theorem 3.4.2.

MPNN + VN can simulate (not just approximate) equivariant DeepSets:
$\mathbb{R}^{n\times d}\rightarrow\mathbb{R}^{n\times d}$. The depth and width
of MPNN + VN needed to simulate DeepSets is up to a constant factor of the
depth and width of DeepSets. This implies that MPNN + VN of $O(1)$ depth and
$O(n^{d})$ width is permutation equivariant universal, and can approximate
self-attention layer and transformers arbitrarily well.

Table 3.3: Baselines for Peptides-func (graph classification) and Peptides-
struct (graph regression). The performance metric is Average Precision (AP)
for classification and MAE for regression. Bold: Best score.

Model | # Params. | Peptides-func | Peptides-struct
---|---|---|---
Test AP before VN | Test AP after VN $\uparrow$ | Test MAE before VN | Test MAE after VN $\downarrow$
GCN | 508k | 0.5930$\pm$0.0023 | 0.6623$\pm$0.0038 | 0.3496$\pm$0.0013 | 0.2488$\pm$0.0021
GINE | 476k | 0.5498$\pm$0.0079 | 0.6346$\pm$0.0071 | 0.3547$\pm$0.0045 | 0.2584$\pm$0.0011
GatedGCN | 509k | 0.5864$\pm$0.0077 | 0.6635$\pm$0.0024 | 0.3420$\pm$0.0013 | 0.2523$\pm$0.0016
GatedGCN+RWSE | 506k | 0.6069$\pm$0.0035 | 0.6685$\pm$0.0062 | 0.3357$\pm$0.0006 | 0.2529$\pm$0.0009
Transformer+LapPE | 488k | 0.6326$\pm$0.0126 | - | 0.2529$\pm$0.0016 | -
SAN+LapPE | 493k | 0.6384$\pm$0.0121 | - | 0.2683$\pm$0.0043 | -
SAN+RWSE | 500k | 0.6439$\pm$0.0075 | - | 0.2545$\pm$0.0012 | -

###### Proof.

Equivariant DeepSets has the following form
${\bm{F}}({\bm{X}})={\bm{L}}^{\text{ds}}_{m}\circ\nu\circ\cdots\circ\nu\circ{\bm{L}}^{\text{ds}}_{1}({\bm{X}})$,
where ${\bm{L}}^{\text{ds}}_{i}$ is the linear permutation equivariant layer
and $\nu$ is an entrywise nonlinear activation layer. Recall that the linear
equivariant layer has the form
${\bm{L}}^{\text{ds}}_{i}({\bm{X}})={\bm{X}}\bm{A}+\frac{1}{n}\mathbf{11}^{T}{\bm{X}}{\bm{B}}+\mathbf{1}\bm{c}^{T}$.
As one can use the same nonlinear entrywise activation layer $\nu$ in MPNN +
VN, it suffices to prove that MPNN + VN can compute linear permutation
equivariant layer ${\bm{L}}^{\text{ds}}$. Now we show that 2 layers of MPNN +
VN can exactly simulate any given linear permutation equivariant layer
${\bm{L}}^{\text{ds}}$.

Specifically, at layer 0, we initialized the node features as follows: The VN
node feature is set to 0, while the node feature for the $i$-th graph node is
set up as ${\bm{x}}_{i}\in\mathbb{R}^{d}$.

At layer 1: VN node feature is $\frac{1}{n}\mathbf{11}^{T}{\bm{X}}$, average
of node features. The collection of features over $n$ graph node feature is
${\bm{X}}{\bm{A}}$. We only need to transform graph node features by a linear
transformation, and set the VN feature as the average of graph node features
in the last iteration. Both can be exactly implemented in Definition 16 of
simplified heterogeneous MPNN + VN.

At layer 2: VN node feature is set to be 0, and the graph node feature is
${\bm{X}}\bm{A}+\frac{1}{n}\mathbf{11}^{T}{\bm{X}}{\bm{B}}+\mathbf{1}\bm{c}^{T}$.
Here we only need to perform the matrix multiplication of the VN feature with
${\bm{B}}$, as well as add a bias ${\bm{c}}$. This can be done by implementing
a linear function for $\gamma_{\text{gn}}$.

It is easy to see the width required for MPNN + VN to simulate DeepSets is
constant. Thus, one can use 2 layers of MPNN + VN to compute linear
permutation equivariant layer ${\bm{L}}^{\text{ds}}_{i}$, which implies that
MPNN + VN can simulate 1 layer of DeepSets exactly with constant depth and
constant width (independent of $n$). Then by the universality of DeepSets,
stated in Theorem 3.4.1, we conclude that MPNN + VN is also permutation
equivariant universal, which implies that the constant layer of MPNN + VN with
$O(n^{d})$ width is able to approximate any continuous equivariant maps. As
the self-attention layer ${\bm{L}}$ and full transformer are both continuous
and equivariant, they can be approximated by MPNN + VN arbitrarily well. ∎

Thanks to the connection between MPNN + VN with DeepSets, there is no extra
assumption on $\mathcal{X}$ except for being compact. The drawback on the
other hand is that the upper bound on the computational complexity needed to
approximate the self-attention with wide MPNN + VN is worse than directly
computing self-attention when $d>2$.

###  3.5 $O(n)$ Depth $O(1)$ Width MPNN + VN

The previous section shows that we can approximate a full attention layer in
Transformer using MPNN with $O(1)$ depth but $O(n^{d})$ width where $n$ is the
number of nodes and $d$ is the dimension of node features. In practice, it is
not desirable to have the width depend on the graph size.

In this section, we hope to study MPNN + VNs with $O(1)$ width and their
ability to approximate a self-attention layer in the Transformer. However,
this appears to be much more challenging. Our result in this section only
shows that for a rather restrictive family of input graphs (see Assumption 3
below), we can approximate a full self-attention layer of transformer with an
MPNN + VN of $O(1)$ width and $O(n)$ depth. We leave the question of MPNN +
VN’s ability in approximate transformers for more general families of graphs
for future investigation.

We first introduce the notion of $({\bm{V}},\delta)$ separable node features.
This is needed to ensure that VN can approximately select one node feature to
process at each iteration with attention $\alpha_{{\bm{n}}}$, the self-
attention in the virtual node.

######  Definition 20 ($({\bm{V}},\delta)$ separable by $\alpha$).

Given a graph $G$ of size $n$ and a fixed ${\bm{V}}\in\mathbb{R}^{n\times
d}=[{\bm{v}}_{1},...,{\bm{v}}_{n}]$ and $\bar{\alpha}\in\mathcal{A}$, we say
node feature ${\bm{X}}\in\mathbb{R}^{n\times d}$ of $G$ is $({\bm{V}},\delta)$
separable by some $\bar{\alpha}$ if the following holds. For any node feature
${\bm{x}}_{i}$, there exist weights
${\bm{W}}^{\bar{\alpha}}_{K},{\bm{W}}^{\bar{\alpha}}_{Q}$ in attention score
$\bar{\alpha}$ such that $\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})>\max_{j\neq
i}\bar{\alpha}({\bm{x}}_{j},{\bm{v}}_{i})+\delta$. We say set $\mathcal{X}$ is
$({\bm{V}},\delta)$ separable by $\bar{\alpha}$ if every element
${\bm{X}}\in\mathcal{X}$ is $({\bm{V}},\delta)$ separable by $\bar{\alpha}$.

Table 3.4: Test performance in graph-level OGB benchmarks [69]. Shown is the mean $\pm$ s.d. of 10 runs.  Model | ogbg-molhiv | ogbg-molpcba | ogbg-ppa | ogbg-code2 |
---|---|---|---|---|---
AUROC $\uparrow$ | Avg. Precision $\uparrow$ | Accuracy $\uparrow$ | F1 score $\uparrow$ |
GCN | 0.7606 $\pm$ 0.0097 | 0.2020 $\pm$ 0.0024 | 0.6839 $\pm$ 0.0084 | 0.1507 $\pm$ 0.0018 |
GCN+virtual node | 0.7599 $\pm$ 0.0119 | 0.2424 $\pm$ 0.0034 | 0.6857 $\pm$ 0.0061 | 0.1595 $\pm$ 0.0018 |
GIN | 0.7558 $\pm$ 0.0140 | 0.2266 $\pm$ 0.0028 | 0.6892 $\pm$ 0.0100 | 0.1495 $\pm$ 0.0023 |
GIN+virtual node | 0.7707 $\pm$ 0.0149 | 0.2703 $\pm$ 0.0023 | 0.7037 $\pm$ 0.0107 | 0.1581 $\pm$ 0.0026 |
SAN | 0.7785 $\pm$ 0.2470 | 0.2765 $\pm$ 0.0042 | – | – |
GraphTrans (GCN-Virtual) | – | 0.2761 $\pm$ 0.0029 | – | 0.1830 $\pm$ 0.0024 |
K-Subtree SAT | – | – | 0.7522 $\pm$ 0.0056 | 0.1937 $\pm$ 0.0028 |
GPS | 0.7880 $\pm$ 0.0101 | 0.2907 $\pm$ 0.0028 | 0.8015 $\pm$ 0.0033 | 0.1894 $\pm$ 0.0024 |
MPNN + VN + NoPE | 0.7676 $\pm$ 0.0172 | 0.2823 $\pm$ 0.0026 | 0.8055 $\pm$ 0.0038 | 0.1727 $\pm$ 0.0017 |
MPNN + VN + PE | 0.7687 $\pm$ 0.0136 | 0.2848 $\pm$ 0.0026 | 0.8027 $\pm$ 0.0026 | 0.1719 $\pm$ 0.0013 |

The use of $({\bm{V}},\delta)$ separability is to approximate hard selection
function arbitrarily well, which is stated below and proved in Section 3.10.1.

######  Lemma 3.5.1 (approximate hard selection).

Given $\mathcal{X}$ is $({\bm{V}},\delta)$ separable by $\bar{\alpha}$ for
some fixed ${\bm{V}}\in\mathbb{R}^{n\times d}$, $\bar{\alpha}\in\mathcal{A}$
and $\delta>0$, the following holds. For any $\epsilon>0$ and $i\in[n]$, there
exists a set of attention weights ${\bm{W}}_{i,Q},{\bm{W}}_{i,K}$ in $i$-th
layer of MPNN + VN such that
$\alpha_{{\bm{n}}}({\bm{x}}_{i},{\bm{v}}_{i})>1-\epsilon$ for any
${\bm{x}}_{i}\in\mathcal{X}_{i}$. In other words, we can approximate a hard
selection function $f_{i}({\bm{x}}_{1},...,{\bm{x}}_{n})={\bm{x}}_{i}$
arbitrarily well on $\mathcal{X}$ by setting $\alpha_{{\bm{n}}}=\bar{\alpha}$.

With the notation set up, We now state an extra assumption needed for deep
MPNN + VN case and the main theorem.

######  AS 3.

$\mathcal{X}$ is $({\bm{V}},\delta)$ separable by $\bar{\alpha}$ for some
fixed ${\bm{V}}\in\mathbb{R}^{n\times d}$, $\bar{\alpha}\in\mathcal{A}$ and
$\delta>0$.

######  Theorem 3.5.2.

Assume AS 1-3 hold for the compact set $\mathcal{X}$ and ${\bm{L}}$. Given any
graph $G$ of size $n$ with node features ${\bm{X}}\in\mathcal{X}$, and a self-
attention layer ${\bm{L}}$ on $G$ (fix
${\bm{W}}_{K},{\bm{W}}_{Q},{\bm{W}}_{V}$ in $\alpha$), there exists a $O(n)$
layer of heterogeneous MPNN + VN with the specific aggregate/update/message
function that can approximate ${\bm{L}}$ on $\mathcal{X}$ arbitrarily well.

The proof is presented in the Section 3.10. On the high level, we can design
an MPNN + VN where the $i$-th layer will select $\tilde{{\bm{x}}}_{i}$, an
approximation of ${\bm{x}}_{i}$ via attention mechanism, enabled by Lemma
3.5.1, and send $\tilde{{\bm{x}}}_{i}$ to the virtual node. Virtual node will
then pass the $\tilde{{\bm{x}}}_{i}$ to all graph nodes and computes the
approximation of $e^{\alpha({\bm{x}}_{i},{\bm{x}}_{j})},\forall j\in[n]$.
Repeat such procedures $n$ times for all graph nodes, and finally, use the
last layer for attention normalization. A slight relaxation of AS3 is also
provided in Section 3.10.

Table 3.5: Evaluation on PCQM4Mv2 [68] dataset. For GPS evaluation, we treated the _validation_ set of the dataset as a test set, since the _test-dev_ set labels are private.  Model | PCQM4Mv2 |
---|---|---
Test-dev MAE $\downarrow$ | Validation MAE $\downarrow$ | Training MAE | # Param.
GCN | 0.1398 | 0.1379 | n/a | 2.0M
GCN-virtual | 0.1152 | 0.1153 | n/a | 4.9M
GIN | 0.1218 | 0.1195 | n/a | 3.8M
GIN-virtual | 0.1084 | 0.1083 | n/a | 6.7M
GRPE [115] | 0.0898 | 0.0890 | n/a | 46.2M
EGT [71] | 0.0872 | 0.0869 | n/a | 89.3M
Graphormer [133] | n/a | 0.0864 | 0.0348 | 48.3M
GPS-small | n/a | 0.0938 | 0.0653 | 6.2M
GPS-medium | n/a | 0.0858 | 0.0726 | 19.4M
MPNN + VN + PE (small) | n/a | 0.0942 | 0.0617 | 5.2M
MPNN + VN + PE (medium) | n/a | 0.0867 | 0.0703 | 16.4M
MPNN + VN + NoPE (small) | n/a | 0.0967 | 0.0576 | 5.2M
MPNN + VN + NoPE (medium) | n/a | 0.0889 | 0.0693 | 16.4M

###  3.6 Experiments

We benchmark MPNN + VN for three tasks, long range interaction modeling in
Section 3.6.2 and OGB regression tasks in Section 3.6.3. The code is available
https://github.com/Chen-Cai-OSU/MPNN-GT-Connection.

####  3.6.1 Dataset Description

ogbg-molhiv and ogbg-molpcba [69] are molecular property prediction datasets
adopted by OGB from MoleculeNet. These datasets use a common node (atom) and
edge (bond) featurization that represent chemophysical properties. The
prediction task of ogbg-molhiv is a binary classification of molecule’s
fitness to inhibit HIV replication. The ogbg-molpcba, derived from PubChem
BioAssay, targets to predict the results of 128 bioassays in the multi-task
binary classification setting.

ogbg-ppa [153] consists of protein-protein association (PPA) networks derived
from 1581 species categorized into 37 taxonomic groups. Nodes represent
proteins and edges encode the normalized level of 7 different associations
between two proteins. The task is to classify which of the 37 groups does a
PPA network originate from.

ogbg-code2 [153] consists of abstract syntax trees (ASTs) derived from the
source code of functions written in Python. The task is to predict the first 5
subtokens of the original function’s name.

OGB-LSC PCQM4Mv2 [68] is a large-scale molecular dataset that shares the same
featurization as ogbg-mol* datasets. It consists of 529,434 molecule graphs.
The task is to predict the HOMO-LUMO gap, a quantum physical property
originally calculated using Density Functional Theory. True labels for
original test-dev and test-challange dataset splits are kept private by the
OGB-LSC challenge organizers. Therefore for the purpose of this chapter, we
used the original validation set as the test set, while we left out random
150K molecules for our validation set.

####  3.6.2 MPNN + VN for LRGB Datasets

We experiment with MPNN + VN for Long Range Graph Benchmark (LRGB) datasets.
Original paper [44] observes that GT outperforms MPNN on 4 out of 5 datasets,
among which GT shows significant improvement over MPNN on Peptides-func and
Peptides-struct for all MPNNs. To test the effectiveness of the virtual node,
we take the original code and modify the graph topology by adding a virtual
node and keeping the hyperparameters of all models unchanged.

Results are in Table 3.3. Interestingly, such a simple change can boost MPNN +
VN by a large margin on Peptides-func and Peptides-struct. Notably, with the
addition of VN, GatedGCN + RWSE (random-walk structural encoding) after
augmented by VN outperforms all transformers on Peptides-func, and GCN
outperforms transformers on Peptides-struct.

####  3.6.3 Stronger MPNN + VN Implementation

Next, by leveraging the modularized implementation from GraphGPS [120], we
implemented a version of MPNN + VN with/without extra positional embedding.
Our goal is not to achieve SOTA but instead to push the limit of MPNN + VN and
better understand the source of the performance gain for GT. In particular, we
replace the GlobalAttention Module in GraphGPS with DeepSets, which is
equivalent to one specific version of MPNN + VN. We tested this specific
version of MPNN + VN on 4 OGB datasets, both with and without the use of
positional embedding. The results are reported in Table 3.4. Interestingly,
even without the extra position embedding, our MPNN + VN is able to further
improve over the previous GCN + VN & GIN + VN implementation. The improvement
on ogbg-ppa is particularly impressive, which is from 0.7037 to 0.8055.
Furthermore, it is important to note that while MPNN + VN does not necessarily
outperform GraphGPS, which is a state-of-the-art architecture using both MPNN,
Position/structure encoding and Transformer, the difference is quite small –
this however, is achieved by a simple MPNN + VN architecture.

We also test MPNN + VN on large-scale molecule datasets PCQMv2, which has
529,434 molecule graphs. We followed [120] and used the original validation
set as the test set, while we left out random 150K molecules for our
validation set. As we can see from Table 3.5, MPNN + VN + NoPE performs
significantly better than the early MPNN + VN implementation: GIN + VN and GCN
+ VN. The performance gap between GPS on the other hand is rather small:
0.0938 (GPS) vs. 0.0942 (MPNN + VN + PE) for the small model and 0.0858 (GPS)
vs. 0.0867 (MPNN + VN + PE) for the medium model.

####  3.6.4 Connection to Over-Smoothing Phenomenon

Over-smoothing refers to the phenomenon that deep GNN will produce same
features at different nodes after too many convolution layers. Here we draw
some connection between VN and common ways of reducing over-smoothing. We
think that using VN can potentially help alleviate the over-smoothing problem.
In particular, we note that the use of VN can simulate some strategies people
use in practice to address over-smoothing. We give two examples below.

Example 1: In [166], the two-step method (center & scale) PairNorm is proposed
to reduce the over-smoothing issues. In particular, PairNorm consists of 1)
Center and 2) Scale

$\tilde{\mathbf{x}}_{i}^{c}=\tilde{\mathbf{x}}_{i}-\frac{1}{n}\sum_{i}\tilde{\mathbf{x}}_{i}$

$\dot{\mathbf{x}}_{i}=s\cdot\frac{\tilde{\mathbf{x}}_{i}^{c}}{\sqrt{\frac{1}{n}\sum_{i}\left||\tilde{\mathbf{x}}_{i}^{c}\right||_{2}^{2}}}$

Where $\tilde{\mathbf{x}}$ is the node features after graph convolution and
$s$ is a hyperparameter. The main component for implementing PairNorm is to
compute the mean and standard deviation of node features. For the mean of node
features, this can be exactly computed in VN. For standard deviation, VN can
arbitrarily approximate it using the standard universality result of MLP [5].
If we further assume that the standard deviation is lower bounded by a
constant, then MPNN + VN can arbitrarily approximate the PairNorm on the
compact set.

Example 2: In [158] mean subtraction (same as the first step of PairNorm) is
also introduced to reduce over-smoothing. As mean subtraction can be trivially
implemented in MPNN + VN, arguments in [158] (with mean subtraction the
revised power Iteration in GCN will lead to the Fiedler vector) can be carried
over to MPNN + VN setting.

In summary, introducing VN allows MPNN to implement key components of [158,
166], we think this is one reason why we observe encouraging empirical
performance gain of MPNN + VN.

###  3.7 On the Limitation of MPNN + VN

Although we showed that in the main part of this chapter, MPNN + VN of varying
depth/width can approximate the self-attention of full/linear transformers,
this does not imply that there is no difference in practice between MPNN + VN
and GT. Our theoretical analysis mainly focuses on approximating self-
attention without considering computational efficiency. In this section, we
mention a few limitations of MPNN + VN compared to GT.

####  3.7.1 Representation Gap

The main limitation of deep MPNN + VN approximating full self-attention is
that we require a quite strong assumption: we restrict the variability of node
features in order to select one node feature to process each iteration. Such
assumption is relaxed by employing stronger attention in MPNN + VN but is
still quite strong.

For the large width case, the main limitation is the computational complexity:
even though the self-attention layer requires $O(n^{2})$ complexity, to
approximate it in wide MPNN + VN framework, the complexity will become
$O(n^{d})$ where $d$ is the dimension of node features.

We think such limitation shares a similarity with research in universal
permutational invariant functions. Both DeepSets [162] and Relational Network
[128] are universal permutational invariant architecture but there is still a
representation gap between the two [169]. Under the restriction to analytic
activation functions, one can construct a symmetric function acting on sets of
size $n$ with elements in dimension $d$, which can be efficiently approximated
by the Relational Network, but provably requires width exponential in $n$ and
$d$ for the DeepSets. We believe a similar representation gap also exists
between GT and MPNN + VN and leave the characterization of functions lying in
such gap as the future work.

####  3.7.2 On the Difficulty of Approximating Other Linear Transformers

In Section 3.3, we showed MPNN + VN of $O(1)$ width and depth can approximate
the self-attention layer of one type of linear transformer, Performer. The
literature on efficient transformers is vast [139] and we do not expect MPNN +
VN can approximate many other efficient transformers. Here we sketch a few
other linear transformers that are hard to approximate by MPNN + VN of
constant depth and width.

Linformer [147] projects the $n\times d$ dimension keys and values to $k\times
d$ suing additional projection layers, which in graph setting is equivalent to
graph coarsening. As MPNN + VN still operates on the original graph, it
fundamentally lacks the key component to approximate Linformer.

We consider various types of efficient transformers effectively generalize the
virtual node trick. By first switching to a more expansive model and reducing
the computational complexity later on, efficient transformers effectively
explore a larger model design space than MPNN + VN, which always sticks to the
linear complexity.

####  3.7.3 Difficulty of Representing SAN Type Attention

In SAN [86], different attentions are used conditional on whether an edge is
presented in the graph or not, detailed below. One may wonder whether we can
approximate such a framework in MPNN + VN.

In our proof of using MPNN + VN to approximate regular GT, we mainly work with
Definition 16 where we do not use any gn-gn edges and therefore not leverage
the graph topology. It is straightforward to use gn-gn edges and obtain the
different message/update/aggregate functions for gn-gn edges non-gn-gn edges.
Although we still achieve the similar goal of SAN to condition on the edge
types, it turns out that we can not arbitrarily approximate SAN.

Without loss of generality, SAN uses two types of attention depending on
whether two nodes are connected by the edge. Specifically,

$\displaystyle\hat{\bm{w}}_{ij}^{k,l}=\left\\{\begin{array}[]{lr}\frac{\bm{Q}^{1,k,l}\bm{h}_{i}^{l}\circ\bm{K}^{1,k,l}\bm{h}_{j}^{l}\circ\bm{E}^{1,k,l}\bm{e}_{ij}}{\sqrt{d_{k}}}&\text{
if }i\text{ and }j\text{ are connected in sparse graph }\\\
\frac{\bm{Q}^{2,k,l}\bm{h}_{i}^{l}\circ\bm{K}^{2,k,l}\bm{h}_{j}^{l}\circ\bm{E}^{2,k,l}\bm{e}_{ij}}{\sqrt{d_{k}}}&\text{
otherwise }\end{array}\right\\}$ (3.7) $\displaystyle
w_{ij}^{k,l}=\left\\{\begin{array}[]{cc}\frac{1}{1+\gamma}\cdot\operatorname{softmax}\left(\sum_{d_{k}}\hat{\bm{w}}_{ij}^{k,l}\right)&\text{
if }i\text{ and }j\text{ are connected in sparse graph }\\\
\frac{\gamma}{1+\gamma}\cdot\operatorname{softmax}\left(\sum_{d_{k}}\hat{\bm{w}}_{ij}^{k,l}\right)&\text{
otherwise }\end{array}\right\\}$

where $\circ$ denotes element-wise multiplication and
$\bm{Q}^{1,k,l},\bm{Q}^{2,k,l},\bm{K}^{1,k,l},\bm{K}^{2,k,l},\bm{E}^{1,k,l},\bm{E}^{2,k,l}\in$
$\mathbb{R}^{d_{k}\times d}$. $\gamma\in\mathbb{R}^{+}$is a hyperparameter
that tunes the amount of bias towards full-graph attention, allowing
flexibility of the model to different datasets and tasks where the necessity
to capture long-range dependencies may vary.

To reduce the notation clutter, we remove the layer index $l$, and edge
features, and also consider only one-attention head case (remove attention
index $k$). The equation is then simplified to

$\displaystyle\hat{\bm{w}}_{ij}=\left\\{\begin{array}[]{lr}\frac{\bm{Q}^{1}\bm{h}_{i}^{l}\circ\bm{K}^{1}\bm{h}_{j}^{l}}{\sqrt{d_{k}}}&\text{
if }i\text{ and }j\text{ are connected in sparse graph }\\\
\frac{\bm{Q}^{2}\bm{h}_{i}^{l}\circ\bm{K}^{2}\bm{h}_{j}^{l}}{\sqrt{d_{k}}}&\text{
otherwise }\end{array}\right\\}$ (3.8) $\displaystyle
w_{ij}=\left\\{\begin{array}[]{cc}\frac{1}{1+\gamma}\cdot\operatorname{softmax}\left(\sum_{d}\hat{\bm{w}}_{ij}\right)&\text{
if }i\text{ and }j\text{ are connected in sparse graph }\\\
\frac{\gamma}{1+\gamma}\cdot\operatorname{softmax}\left(\sum_{d}\hat{\bm{w}}_{ij}\right)&\text{
otherwise }\end{array}\right\\}$

We will then show that Equation 3.8 can not be expressed (up to an arbitrary
approximation error) in MPNN + VN framework. To simulate SAN type attention,
our MPNN + VN framework will have to first simulate one type of attention for
all edges, and then simulate the second type of attention between gn-gn edges
by properly offset the contribution from the first attention. This seems
impossible (although we do not have rigorous proof) as we cannot express the
difference between two attention in the new attention mechanism.

###  3.8 Related Work

Virtual node in MPNN. The virtual node augments the graph with an additional
node to facilitate the information exchange among all pairs of nodes. It is a
heuristic proposed in [54] and has been observed to improve the performance in
different tasks [68, 69]. Surprisingly, its theoretical properties have
received little study. To the best of our knowledge, only a recent paper [72]
analyzed the role of the virtual node in the link prediction setting in terms
of 1) expressiveness of the learned link representation and 2) the potential
impact on under-reaching and over-smoothing.

Graph transformer. Because of the great successes of Transformers in natural
language processing (NLP) [142, 151] and recently in computer vision [39, 45,
101], there is great interest in extending transformers for graphs. One common
belief of advantage of graph transformer over MPNN is its capacity in
capturing long-range interactions while alleviating over-smoothing [96, 114,
21] and over-squashing in MPNN [3, 140].

Fully-connected Graph transformer [43] was introduced with eigenvectors of
graph Laplacian as the node positional encoding (PE). Various follow-up works
proposed different ways of PE to improve GT, ranging from an invariant
aggregation of Laplacian’s eigenvectors in SAN [86], pair-wise graph distances
in Graphormer [159], relative PE derived from diffusion kernels in GraphiT
[112], and recently Sign and Basis Net [100] with a principled way of handling
sign and basis invariance. Other lines of research in GT include combining
MPNN and GT [153, 120], encoding the substructures [25] and efficient graph
transformers for large graphs [152].

Deep Learning on Sets. Janossy pooling [113] is a framework to build
permutation invariant architecture for sets using permuting & averaging
paradigm while limiting the number of elements in permutations to be $k<n$.
Under this framework, DeepSets [162] and PointNet [119] are recovered as the
case of $k=1$. For case $k=2$, self-attention and Relation Network [128] are
recovered [145]. Although DeepSets and Relation Network [128] are both shown
to be universal permutation invariant, recent work [169] provides a finer
characterization on the representation gap between the two architectures.

###  3.9 Concluding Remarks

in this chapter, we study the expressive power of MPNN + VN under the lens of
GT. If we target the self-attention layer in Performer and Linear Transformer,
one only needs $O(1)$-depth $O(1)$ width for arbitrary approximation error.
For self-attention in full transformer, we prove that heterogeneous MPNN + VN
of either $O(1)$ depth $O(n^{d})$ width or $O(n)$ depth $O(1)$ width (under
some assumptions) can approximate 1 self-attention layer arbitrarily well.
Compared to early results [82] showing GT can approximate MPNN, our
theoretical result draws the connection from the inverse direction.

On the empirical side, we demonstrate that MPNN + VN remains a surprisingly
strong baseline. Despite recent efforts, we still lack good benchmark datasets
where GT can outperform MPNN by a large margin. Understanding the inductive
bias of MPNN and GT remains challenging. For example, can we mathematically
characterize tasks that require effective long-range interaction modeling, and
provide a theoretical justification for using GT over MPNN (or vice versa) for
certain classes of functions on the space of graphs? We believe making
processes towards answering such questions is an important future direction
for the graph learning community.

###  3.10 Missing Proofs

In this section, we show the missing proofs of $O(n)$ heterogeneous MPNN + VN
Layer with $O(1)$ width can approximate $1$ self attention layer arbitrarily
well.

####  3.10.1 Assumptions

A special case of $({\bm{V}},\delta)$ separable is when $\delta=0$, i.e.,
$\forall i,\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})>\max_{j\neq
i}\bar{\alpha}({\bm{x}}_{j},{\bm{v}}_{i})$. We provide a geometric
characterization of ${\bm{X}}$ being $({\bm{V}},0)$ separable.

######  Lemma 3.10.1.

Given $\bar{\alpha}$ and ${\bm{V}}$, ${\bm{X}}$ is $({\bm{V}},0)$ separable by
$\bar{\alpha}$ $\Longleftrightarrow$ ${\bm{x}}_{i}$ is not in the convex hull
spanned by $\\{{\bm{x}}_{j}\\}_{j\neq i}$. $\Longleftrightarrow$ there are no
points in the convex hull of $\\{{\bm{x}}_{i}\\}_{i\in[n]}$.

###### Proof.

The second equivalence is trivial so we only prove the first equivalence. By
definition, ${\bm{X}}$ is $({\bm{V}},0)$ separable by $\bar{\alpha}$
$\Longleftrightarrow$ $\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})>\max_{j\neq
i}\bar{\alpha}({\bm{x}}_{j},{\bm{v}}_{i})\forall i\in[n]$
$\Longleftrightarrow$
$\langle{{\bm{x}}_{i},{\bm{W}}_{Q}^{\bar{\alpha}}{\bm{W}}_{K}^{{\bar{\alpha},T}}{\bm{v}}_{i}}\rangle>\max_{j\neq
i}\langle{{\bm{x}}_{j},{\bm{W}}_{Q}^{\bar{\alpha}}{\bm{W}}_{K}^{{\bar{\alpha}},T}{\bm{v}}_{i}}\rangle\forall
i\in[n]$.

By denoting the
${\bm{v}}^{\prime}_{i}:={\bm{W}}_{Q}^{\bar{\alpha}}{\bm{W}}_{K}^{{\bar{\alpha}},T}{\bm{v}}_{i}\in\mathbb{R}^{d}$,
we know that $\langle{{\bm{x}}_{i},{\bm{v}}^{\prime}_{i}}\rangle>\max_{j\neq
i}\langle{{\bm{x}}_{j},{\bm{v}}^{\prime}_{i}}\rangle\forall i\in[n]$, which
implies that $\forall i\in[n],{\bm{x}}_{i}$ can be linearly seprated from
$\\{{\bm{x}}_{j}\\}_{j\neq i}$ $\Longleftrightarrow$ ${\bm{x}}_{i}$ is not in
the convex hull spanned by $\\{{\bm{x}}_{j}\\}_{j\neq i}$, which concludes the
proof. ∎

See 3.5.1

###### Proof.

Denote $\bar{\alpha}^{\prime}$ as the unnormalized $\bar{\alpha}$. As
$\mathcal{X}$ is $({\bm{V}},\delta)$ separable by $\bar{\alpha}$, by
definition we know that $\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})>\max_{j\neq
i}\bar{\alpha}({\bm{x}}_{j},{\bm{v}}_{i})+\delta$ holds for any $i\in[n]$ and
${\bm{x}}_{i}\in\mathcal{M}$. We can amplify this by multiple the weight
matrix in $\bar{\alpha}$ by a constant factor $c$ to make
$\bar{\alpha}^{\prime}({\bm{x}}_{i},{\bm{v}}_{i})>\max_{j\neq
i}\bar{\alpha}^{\prime}({\bm{x}}_{j},{\bm{v}}_{i})+c\delta$. This implies that
$e^{\bar{\alpha}^{\prime}({\bm{x}}_{i},{\bm{v}}_{i})}>e^{c\delta}\max_{j\neq
i}e^{\bar{\alpha}^{\prime}({\bm{x}}_{j},{\bm{v}}_{i})}$. This means after
softmax, the attention score $\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})$ will be
at least $\frac{e^{c\delta}}{e^{c\delta}+n-1}$. We can pick a large enough
$c(\delta,\epsilon)$ such that
$\bar{\alpha}({\bm{x}}_{i},{\bm{v}}_{i})>1-\epsilon$ for any
${\bm{x}}_{i}\in\mathcal{X}_{i}$ and $\epsilon>0$. ∎

Proof Intuition and Outline. On the high level, $i$-th MPNN + VN layer will
select $\tilde{{\bm{x}}}_{i}$, an approximation $i$-th node feature
${\bm{x}}_{i}$ via attention mechanism, enabled by Lemma 3.5.1, and send
$\tilde{{\bm{x}}}_{i}$ to the virtual node. Virtual node will then pass the
$\tilde{{\bm{x}}}_{i}$ to all graph nodes and computes the approximation of
$e^{\alpha({\bm{x}}_{i},{\bm{x}}_{j})},\forall j\in[n]$. Repeat such
procedures $n$ times for all graph nodes, and finally, use the last layer for
attention normalization.

The main challenge of the proof is to 1) come up with
message/update/aggregation functions for heterogeneous MPNN + VN layer, which
is shown in Section 3.10.2, and 2) ensure the approximation error, both from
approximating Aggregate/Message/Update function with MLP and the noisy input,
can be well controlled, which is proved in Section 3.10.4.

We will first instantiate the Aggregate/Message/Update function for
virtual/graph nodes in Section 3.10.2, and prove that each component can be
either exactly computed or approximated to an arbitrary degree by MLP. Then we
go through an example in Section 3.10.3 of approximate self-attention layer
${\bm{L}}$ with $O(n)$ MPNN + VN layers. The main proof is presented in
Section 3.10.4, where we show that the approximation error introduced during
different steps is well controlled. Lastly, in Section 3.10.5 we show
assumption on node features can be relaxed if a more powerful attention
mechanism GATv2 [12] is allowed in MPNN + VN.

####  3.10.2 Aggregate/Message/Update Functions

Let $\mathcal{M}$ be a multiset of vectors in $\mathbb{R}^{d}$. The specific
form of Aggregate/Message/Update for virtual and graph nodes are listed below.
Note that ideal forms will be implemented as MLP, which will incur an
approximation error that can be controlled to an arbitrary degree. We use
${\bm{z}}_{{\bm{n}}}^{(k)}$ denotes the virtual node’s feature at $l$-th
layer, and ${\bm{z}}_{i}^{(k)}$ denotes the graph node $i$’s node feature.
Iteration index $k$ starts with 0 and the node index starts with 1.

##### virtual node

At $k$-th iteration, virtual node $i$’s feature ${\bm{z}}_{i}^{(k)}$ is a
concatenation of three component $[\tilde{{\bm{x}}}_{i},{\bm{v}}_{k+1},0]$
where the first component is the approximately selected node features
${\bm{x}}_{i}\in\mathbb{R}^{d}$, the second component is the
${\bm{v}}_{i}\in\mathbb{R}^{d}$ that is used to select the node feature in
$i$-th iteration. The last component is just a placeholder to ensure the
dimension of the virtual node and graph node are the same. It is introduced to
simplify notation.

_Initial feature_ is ${\bm{z}}_{{\bm{n}}}^{(0)}=[\bm{0}_{d},{\bm{v}}_{1},0]$.

_Message function + Aggregation function_ $\tau_{j\in[n]}\phi_{\text{vn-
gn}}^{(k)}:\mathbb{R}^{2d+1}\times\mathcal{M}\rightarrow\mathbb{R}^{2d+1}$ has
two cases to discuss depending on value of $k$. For $k=1,2,...,n$,

$\begin{split}&\tau_{j\in[n]}\phi_{\text{vn-
gn}}^{(k)}({\bm{z}}_{{\bm{n}}}^{(k-1)},\\{{\bm{z}}_{i}^{(k-1)}\\}_{i})=\\\
&\begin{cases}\sum_{i}\alpha_{{\bm{n}}}({\bm{z}}_{{\bm{n}}}^{(k-1)},{\bm{z}}_{i}^{(k-1)}){\bm{z}}_{i}^{(k-1)}&k=1,2,...,n\\\
\bm{1}_{2d+1}&k=n+1,n+2\\\ \end{cases}\end{split}$ (3.9)

where ${\bm{z}}_{{\bm{n}}}^{(k-1)}=[\tilde{{\bm{x}}}_{k-1},{\bm{v}}_{k},0]$.
${\bm{z}}_{i}^{(k-1)}=[\overbrace{\underbrace{{\bm{x}}_{i}}_{d\text{
dim}},...,...}^{2d+1\text{ dim}}]$ is the node $i$’s feature, where the first
$d$ coordinates remain fixed for different iteration $k$.
$\tau_{j\in[n]}\phi_{\text{vn-gn}}^{(k)}$ use attention $\alpha_{{\bm{n}}}$ to
approximately select $k$-th node feature
$[\overbrace{\underbrace{{\bm{x}}_{k}}_{d\text{ dim}},...,...}^{2d+1\text{
dim}}]$. Note that the particular form of attention $\alpha_{{\bm{n}}}$ needed
for soft selection is not important as long as we can approximate hard
selection arbitrarily well. As the ${\bm{z}}^{(k-1)}_{{\bm{n}}}$ contains
${\bm{v}}_{k}$ and ${\bm{z}}_{i}^{(k-1)}$ contains ${\bm{x}}_{i}$ (see
definition of graph node feature in Section 3.10.2), this step can be made as
close to hard selection as possible, according to Lemma 3.10.4.

In the case of $k=n+1$, $\tau_{j\in[n]}\phi_{\text{vn-
gn}}^{(k)}:\underbrace{\mathbb{R}^{2d+1}}_{{\bm{n}}}\times\underbrace{\mathcal{M}}_{\text{set
of }\text{gn}}\rightarrow\mathbb{R}^{d}$ simply returns $\bm{1}_{2d+1}$. This
can be exactly implemented by an MLP.

_Update function
$\gamma_{{\bm{n}}}^{(k)}:\underbrace{\mathbb{R}^{2d+1}}_{{\bm{n}}}\times\underbrace{\mathbb{R}^{2d+1}}_{\text{gn}}\rightarrow\mathbb{R}^{2d+1}$_:
Given the virtual node’s feature in the last iteration, and the selected
feature in virtual node ${\bm{y}}=[{\bm{x}}_{k},...,...]$ with
$\alpha_{{\bm{n}}}$,

$\gamma_{{\bm{n}}}^{(k)}(\cdot,{\bm{y}})=\begin{cases}[{\bm{y}}_{0:d},{\bm{v}}_{k+1},0]&k=1,...,n-1\\\
[{\bm{y}}_{0:d},\bm{0}_{d},0]&k=n\\\ \bm{1}_{2d+1}&k=n+1,n+2\\\ \end{cases}$
(3.10)

where ${\bm{y}}_{0:d}$ denotes the first $d$ channels of
${\bm{y}}\in\mathbb{R}^{2d+1}$. ${\bm{y}}$ denotes the selected node
${\bm{z}}_{i}$’s feature in Message/Aggregation function.
$\gamma_{{\bm{n}}}^{(k)}$ can be exactly implemented by an MLP for any
$k=1,...,n+2$.

##### Graph node

Graph node $i$’s feature ${\bm{v}}_{i}\in\mathbb{R}^{2d+1}$ can be thought of
as a concatenation of three components $[\underbrace{{\bm{x}}_{i}}_{d\text{
dim}},\underbrace{\textsf{tmp}{}}_{d\text{
dim}},\underbrace{\textsf{partialsum}{}}_{1\text{ dim}}]$, where
${\bm{x}}_{i},\in\mathbb{R}^{d},\textsf{tmp}{}\in\mathbb{R}^{d}$ 333tmp
technicially denotes the dimension of projected feature by $W_{V}$ and does
not has to be in $\mathbb{R}^{d}$. We use $\mathbb{R}^{d}$ here to reduce the
notation clutter., and $\textsf{partialsum}{}\in\mathbb{R}$.

In particular, ${\bm{x}}_{i}$ is the initial node feature. The first $d$
channel will stay the same until the layer $n+2$.
$\textsf{tmp}{}=\sum_{j\in\text{subset
of}[n]}e^{\alpha^{\prime}_{ij}}{\bm{x}}_{j}$ stands for the unnormalized
attention contribution up to the current iteration.
$\textsf{partialsum}{}\in\mathbb{R}$ is a partial sum of the unnormalized
attention score, which will be used for normalization in the $n+2$-th
iteration.

_Initial feature_ ${\bm{z}}_{\text{gn}}^{(0)}=[{\bm{x}}_{i},\bm{0}_{d},0]$.

_Message function + Aggregate function: $\tau_{j\in[n]}\phi_{\text{gn-
vn}}^{(k)}:\mathbb{R}^{2d+1}\times\mathbb{R}^{2d+1}\rightarrow\mathbb{R}^{2d+1}$_
is just “copying the second argument” since there is just one incoming message
from the virtual node, i.e., $\tau_{j\in[n]}\phi_{\text{gn-
vn}}^{(k)}({\bm{x}},\\{{\bm{y}}\\})={\bm{y}}$. This function can be exactly
implemented by an MLP.

_Update function_
$\gamma_{\text{gn}}^{(k)}:\underbrace{\mathbb{R}^{2d+1}}_{\text{gn}}\times\underbrace{\mathbb{R}^{2d+1}}_{{\bm{n}}}\rightarrow\mathbb{R}^{2d+1}$
is of the following form.

$\begin{split}&\gamma_{\text{gn}}^{(k)}([{\bm{x}},\textsf{tmp}{},\textsf{partialsum}{}],{\bm{y}})=\\\
&\begin{cases}[{\bm{x}},\textsf{tmp}{},\textsf{partialsum}{}]&k=1\\\
[{\bm{x}},\textsf{tmp}{}+e^{\alpha^{\prime}({\bm{x}},{\bm{y}}_{0:d})}{\bm{W}}_{V}{\bm{y}}_{0:d},\\\
\textsf{partialsum}{}+e^{\alpha^{\prime}({\bm{x}},{\bm{y}}_{0:d})}]&k=2,...,n+1\\\
[\frac{\textsf{tmp}{}}{\textsf{partialsum}{}},\bm{0}_{d},0]&k=n+2\\\
\end{cases}\end{split}$ (3.11)

where $\alpha^{\prime}({\bm{x}},{\bm{y}}_{0:d})$ is the usual unnormalized
attention score. Update function $\gamma_{\text{gn}}^{(k)}$ can be arbitrarily
approximated by an MLP, which is proved below.

######  Lemma 3.10.2.

Update function $\gamma_{\text{gn}}^{(k)}$ can be arbitrarily approximated by
an MLP from $\mathbb{R}^{2d+1}\times\mathbb{R}^{2d+1}$ to $\mathbb{R}^{2d+1}$
for all $k=1,...,n+2$.

###### Proof.

We will show that for any $k=1,...,n+2$, the target function
$\gamma_{\text{gn}}^{(k)}:\mathbb{R}^{2d+1}\times\mathbb{R}^{2d+1}\rightarrow\mathbb{R}^{2d+1}$
is continuous and the domain is compact. By the universality of MLP in
approximating continuous function on the compact domain, we know
$\gamma_{\text{gn}}^{(k)}$ can be approximated to arbitrary precision by an
MLP.

Recall that

$\begin{split}&\gamma_{\text{gn}}^{(k)}([{\bm{x}},\textsf{tmp}{},\textsf{partialsum}{}],{\bm{y}})=\\\
&\begin{cases}[{\bm{x}},\textsf{tmp}{},\textsf{partialsum}{}]&k=1\\\
[{\bm{x}},\textsf{tmp}{}+e^{\alpha^{\prime}({\bm{x}},{\bm{y}}_{0:d})}{\bm{W}}_{V}{\bm{y}}_{0:d},\\\
\textsf{partialsum}{}+e^{\alpha^{\prime}({\bm{x}},{\bm{y}}_{0:d})}]&k=2,...,n+1\\\
[\frac{\textsf{tmp}{}}{\textsf{partialsum}{}},\bm{0}_{d},0]&k=n+2\\\
\end{cases}\end{split}$

it is easy to see that $k=1$, $\gamma_{\text{gn}}^{(1)}$ is continuous. We
next show for $k=2,...,n+2$, $\gamma_{\text{gn}}^{(1)}$ is also continuous and
all arguments lie in a compact domain.

$\gamma_{\text{gn}}^{(k)}$ is continuous because to a)
$\alpha^{\prime}({\bm{x}},{\bm{y}})$ is continuous b) scalar-vector
multiplication, sum, and exponential are all continuous. Next, we show that
four component ${\bm{x}},\textsf{tmp}{},\textsf{partialsum}{},\\\
{\bm{y}}_{0:d}$ all lies in a compact domain.

${\bm{x}}$ is the initial node features, and by AS1 their norm is bounded so
${\bm{x}}$ is in a compact domain.

tmp is an approximation of
$e^{\alpha^{\prime}_{i,1}}{\bm{W}}_{V}{\bm{x}}_{1}+e^{\alpha^{\prime}_{i,2}}{\bm{W}}_{V}{\bm{x}}_{2}+...$.
As $\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})$ is both upper and lower
bounded by AS2 for all $i,j\in[n]$ and ${\bm{x}}_{i}$ is bounded by AS1,
$e^{\alpha^{\prime}_{i,1}}{\bm{W}}_{V}{\bm{x}}_{1}+e^{\alpha^{\prime}_{i,2}}{\bm{W}}_{V}{\bm{x}}_{2}+...$
is also bounded from below and above. tmp will also be bounded as we can
control the error to any precision.

partialsum is an approximation of
$e^{\alpha^{\prime}_{i,1}}+e^{\alpha^{\prime}_{i,2}}+...$. For the same reason
as the case above, partialsum is also bounded both below and above.

${\bm{y}}_{0:d}$ will be $\tilde{{\bm{x}}}_{i}$ at $i$-th iteration so it will
also be bounded by AS1.

Therefore we conclude the proof. ∎

####  3.10.3 A Running Example

We provide an example to illustrate how node features are updated in each
iteration.

Time $0$: All nodes are initialized as indicated in Section 3.10.2. Virtual
node feature ${\bm{z}}_{{\bm{n}}}^{(0)}=[\bm{0}_{d},{\bm{v}}_{1},0]$. Graph
node feature ${\bm{z}}_{i}^{(0)}=[{\bm{x}}_{i},\bm{0}_{d},0]$ for all
$i\in[n]$.

Time $1$:

For virtual node, according to the definition of
$\tau_{j\in[n]}\phi_{\text{vn-gn}}^{(1)}$ in Equation 3.9, it will pick an
approximation of ${\bm{x}}_{1}$, i.e. $\tilde{{\bm{x}}}_{1}$. Note that the
approximation error can be made arbitrarily small. VN’s node feature
${\bm{z}}_{{\bm{n}}}^{(1)}=[\tilde{{\bm{x}}}_{1},{\bm{v}}_{2},0]$.

For $i$-th graph node feature, ${\bm{z}}_{{\bm{n}}}^{(0)}=\bm{1}_{d}$, and
${\bm{z}}_{i}^{(0)}=[{\bm{x}}_{i},\bm{0}_{d},0]$. According to
$\gamma_{\text{gn}}^{(k)}$ in Equation 3.11,
${\bm{z}}_{i}^{(1)}=[{\bm{x}}_{i},\bm{0}_{d},0]$.

Time $2$:

For the virtual node feature: similar to the analysis in time 1, VN’s feature
${\bm{z}}_{{\bm{n}}}^{(2)}=[\tilde{{\bm{x}}}_{2},{\bm{v}}_{3},0]$ now. Note
that the weights and bias in $\tau_{j\in[n]}\phi_{\text{vn-gn}}^{(2)}$ will be
different from those in $\tau_{j\in[n]}\phi_{\text{vn-gn}}^{(1)}$.

For $i$-th graph node feature, as
${\bm{z}}_{{\bm{n}}}^{(1)}=[\tilde{{\bm{x}}}_{1},{\bm{v}}_{2},0]$ and
${\bm{z}}^{(1)}_{i}=[{\bm{x}}_{i},\bm{0}_{d},0]$, according to
$\gamma_{\text{gn}}^{(k)}$ in Equation 3.11,
${\bm{z}}^{(2)}_{i}=[{\bm{x}}_{i},e^{\widetilde{\alpha^{\prime}_{i,1}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{1},e^{\widetilde{\alpha^{\prime}_{i,1}}}]$.
Here
$\widetilde{\alpha^{\prime}_{i,1}}:=\alpha^{\prime}({\bm{x}}_{i},\tilde{{\bm{x}}}_{1})$.
We will use similar notations in later iterations. 444To reduce the notation
clutter and provide an intuition of the proof, we omit the approximation error
introduced by using MLP to approximate aggregation/message/update function,
and assume the aggregation/message/update can be exactly implemented by neural
networks. In the proofs, approximation error by MLP is handled rigorously.

Time $3$:

Similar to the analysis above,
${\bm{z}}_{{\bm{n}}}^{(3)}=[\widetilde{{\bm{x}}_{3}},{\bm{v}}_{4},0]$.

${\bm{z}}_{i}^{(3)}=[{\bm{x}}_{i},e^{\widetilde{\alpha^{\prime}_{i,1}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{1}+e^{\widetilde{\alpha^{\prime}_{i,2}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{2},e^{\widetilde{\alpha^{\prime}_{i,1}}}+e^{\widetilde{\alpha^{\prime}_{i,2}}}]$.

Time $n$:

${\bm{z}}_{{\bm{n}}}^{(n)}=[\tilde{{\bm{x}}}_{n},\bm{0}_{d},0]$.

${\bm{z}}_{i}^{(n)}={\bm{x}}_{i},\underbrace{e^{\widetilde{\alpha^{\prime}_{i,1}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{1}+...+e^{\widetilde{\alpha^{\prime}_{i,n-1}}}{\bm{W}}_{V}\widetilde{{\bm{x}}_{n-1}}}_{n-1\text{
terms}},\\\
\underbrace{e^{\widetilde{\alpha^{\prime}_{i,1}}}+e^{\widetilde{\alpha^{\prime}_{i,2}}}+...+e^{\widetilde{\alpha^{\prime}_{i,n-1}}}]}_{n-1\text{
terms}}$.

Time $n+1$:

According to Section 3.10.2, in $n+1$ iteration, the virtual node’s feature
will be $\bm{1}_{d}$.

${\bm{z}}_{i}^{(n+1)}=[{\bm{x}}_{i},\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{k},\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}]$

Time $n+2$ (final layer):

For the virtual node, its node feature will stay the same.

For the graph node feature, the last layer will serve as a normalization of
the attention score (use MLP to approximate vector-scalar multiplication), and
set the last channel to be 0 (projection), resulting in an approximation of
$[{\bm{x}}_{i},\frac{\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{k}}{\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}},0]$.
Finally, we need one more linear transformation to make the node feature
become
$[\frac{\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}{\bm{W}}_{V}\tilde{{\bm{x}}}_{k}}{\sum_{k\in[n]}e^{\widetilde{\alpha^{\prime}_{ik}}}},\bm{0}_{d},0]$.
The first $d$ channel is an approximation of the output of the self-attention
layer for node $i$ where the approximation error can be made as small as
possible. This is proved in Section 3.10, and we conclude that heterogeneous
MPNN + VN can approximate the self-attention layer ${\bm{L}}$ to arbitrary
precision with $O(n)$ MPNN layers.

####  3.10.4 Controlling Error

On the high level, there are three major sources of approximation error: 1)
approximate hard selection with self-attention and 2) approximate equation
$\gamma_{\text{gn}}^{(k)}$ with MLPs, and 3) attention normalization in the
last layer. In all cases, we aim to approximate the output of a continuous map
${\bm{L}}_{c}({\bm{x}})$. However, our input is usually not exact ${\bm{x}}$
but an approximation of $\tilde{{\bm{x}}}$. We also cannot access the original
map ${\bm{L}}_{c}$ but instead, an MLP approximation of ${\bm{L}}_{c}$,
denoted as ${\bm{L}}_{{\mathrm{MLP}}}$. The following lemma allows to control
the difference between ${\bm{L}}_{c}({\bm{x}})$ and
${\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})$.

######  Lemma 3.10.3.

Let ${\bm{L}}_{c}$ be a continuous map from compact set to compact set in
Euclidean space. Let ${\bm{L}}_{{\mathrm{MLP}}}$ be the approximation of
${\bm{L}}_{c}$ by MLP. If we can control $\|{\bm{x}}-\tilde{{\bm{x}}}\|$ to an
arbitrarily small degree, we can then control the error
$\|{\bm{L}}_{c}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$
arbitrarily small.

###### Proof.

By triangle inequality
$\|{\bm{L}}_{c}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|\leq\|{\bm{L}}_{c}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}({\bm{x}}))\|+\|{\bm{L}}_{{\mathrm{MLP}}}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$.

For the first term
$\|{\bm{L}}_{c}(\tilde{{\bm{x}}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$,
by the universality of MLP, we can control the error
$\|{\bm{L}}_{c}(\tilde{{\bm{x}}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$
in arbitrary degree.

For the second term
$\|{\bm{L}}_{{\mathrm{MLP}}}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$,
as ${\bm{L}}_{{\mathrm{MLP}}}$ is continuous on a compact domain, it is
uniformly continuous by Heine-Cantor theorem. This means that we can control
the
$\|{\bm{L}}_{{\mathrm{MLP}}}({\bm{x}})-{\bm{L}}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}})\|$
as long as we can control $\|{\bm{x}}-\tilde{{\bm{x}}}\|$, independent from
different ${\bm{x}}$. By assumption, this is indeed the case so we conclude
the proof. ∎

######  Remark 9.

The implication is that when we are trying to approximate the output of a
continuous map ${\bm{L}}_{c}$ on the compact domain by an MLP
${\bm{L}}_{{\mathrm{MLP}}}$, it suffices to show the input is 1)
$\|{\bm{L}}_{c}-{\bm{L}}_{{\mathrm{MLP}}}\|_{\infty}$ and 2)
$\|\tilde{{\bm{x}}}-{\bm{x}}\|$ can be made arbitrarily small. The first point
is usually done by the universality of MLP on the compact domain [31]. The
second point needs to be shown case by case.

In the Section 3.10.3, to simplify the notations we omit the error introduced
by using MLP to approximate aggregation/message/update functions (continuous
functions on the compact domain of $\mathbb{R}^{d}$.) in MPNN + VN. Lemma
3.10.3 justify such reasoning.

######  Lemma 3.10.4 ($\tilde{{\bm{x}}}_{i}$ approximates ${\bm{x}}_{i}$.
$\widetilde{\alpha^{\prime}_{i,j}}$ approximates $\alpha^{\prime}_{i,j}$.).

For any $\epsilon>0$ and $x\in\mathcal{X}$, there exist a set of weights for
message/aggregate functions of the virtual node such that
$||{\bm{x}}_{i}-\tilde{{\bm{x}}}_{i}||<\epsilon$ and
$|\alpha^{\prime}_{i,j}-\widetilde{\alpha^{\prime}_{i,j}}|<\epsilon$.

###### Proof.

By Lemma 3.5.1 We know that
$\widetilde{\alpha_{i,j}}:=\widetilde{\alpha}({\bm{x}}_{i},{\bm{x}}_{j})\rightarrow\delta(i-j)$
as $C_{3}(\epsilon)$ goes to infinity. Therefore we have

$||\tilde{{\bm{x}}}_{i}-{\bm{x}}_{i}||=||\sum_{j}\widetilde{\alpha_{i,j}}{\bm{x}}_{j}-{\bm{x}}_{i}||=||\sum(\widetilde{\alpha}_{i,j}-\delta(i-j)){\bm{x}}_{j}||<\epsilon\sum||{\bm{x}}_{j}||<nC_{1}\epsilon$
(3.12)

As $n$ and $C_{1}$ are fixed, we can make the upper bound as small as we want
by increasing $C_{3}$.

$|\alpha^{\prime}_{i,j}-\widetilde{\alpha^{\prime}_{i,j}}|=|\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})-\alpha^{\prime}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})|=|\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})-\alpha^{\prime}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})|+|\alpha^{\prime}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})-\alpha^{\prime}_{{\mathrm{MLP}}}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})|=|\alpha^{\prime}({\bm{x}}_{i}-\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})|=({\bm{x}}_{i}-\tilde{{\bm{x}}}_{i})^{T}{\bm{x}}_{j}C_{2}^{2}+\epsilon<nC_{1}\epsilon
C_{1}C_{2}^{2}+\epsilon=(nC_{1}^{2}C_{2}^{2}+1)\epsilon$. As
$\alpha^{\prime}_{i,j},\widetilde{\alpha^{\prime}_{i,j}}$ is bounded from
above and below, it’s easy to see that
$|e^{\alpha^{\prime}_{i,j}}-e^{\widetilde{\alpha^{\prime}_{i,j}}}|=|e^{\alpha^{\prime}_{i,j}}(1-e^{\alpha^{\prime}_{i,j}-\widetilde{\alpha^{\prime}_{i,j}}})|<C(1-e^{\alpha^{\prime}_{i,j}-\widetilde{\alpha^{\prime}_{i,j}}})$
can be controlled to arbitrarily degree. ∎

See 3.5.2

###### Proof.

$i$-th MPNN + VN layer will select $\tilde{{\bm{x}}}_{i}$, an arbitrary
approximation $i$-th node feature ${\bm{x}}_{i}$ via attention mechanism. This
is detailed in the message/aggregation function of the virtual node in Section
3.10.2. Assuming the regularity condition on feature space $\mathcal{X}$,
detailed in AS3, the approximation error can be made as small as needed, as
shown in Lemmas 3.5.1 and 3.10.4.

Virtual node will then pass the $\tilde{{\bm{x}}}_{i}$ to all graph nodes,
which computes an approximation of
$e^{\alpha^{\prime}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})},\forall j\in[n]$. This
step is detailed in the update function $\gamma_{\text{gn}}^{(k)}$ of graph
nodes, which can also be approximated arbitrarily well by MLP, proved in Lemma
3.10.2. By Lemma 3.10.3, we have an arbitrary approximation of
$e^{\alpha^{\prime}(\tilde{{\bm{x}}}_{i},{\bm{x}}_{j})},\forall j\in[n]$,
which itself is an arbitrary approximation of
$e^{\alpha^{\prime}({\bm{x}}_{i},{\bm{x}}_{j})},\forall j\in[n]$.

Repeat such procedures $n$ times for all graph nodes, we have an arbitrary
approximation of
$\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}{\bm{W}}_{V}{\bm{x}}_{k}\in\mathbb{R}^{d}$
and $\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}\in\mathbb{R}$. Finally, we use the
last layer to approximate attention normalization
$L_{c}({\bm{x}},y)=\frac{{\bm{x}}}{y}$, where
${\bm{x}}\in\mathbb{R}^{d},y\in\mathbb{R}$. As inputs for attention
normalization are arbitrary approximation of
$\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}{\bm{W}}_{V}{\bm{x}}_{k}$ and
${\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}}$, both of them are lower/upper
bounded according to AS1 and AS2. Since the denominator is upper bounded by a
positive number, this implies that the target function $L_{c}$ is continuous
in both arguments. By evoking Lemma 3.10.3 again, we conclude that we can
approximate its output
$\frac{\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}{\bm{W}}_{V}{\bm{x}}_{k}}{{\sum_{k\in[n]}e^{\alpha^{\prime}_{ik}}}}$
arbitrarily well. This concludes the proof.

∎

####  3.10.5 Relaxing Assumptions with More Powerful Attention

One limitation of Theorem 3.5.2 are assumptions on node features space
$\mathcal{X}$: we need to 1) restrict the variability of node feature so that
we can select one node feature to process each iteration. 2) The space of the
node feature also need to satisfy certain configuration in order for VN to
select it. For 2), we now consider a different attention function for
$\alpha_{{\bm{n}}}$ in MPNN + VN that can relax the assumptions AS3 on
$\mathcal{X}$.

More powerful attention mechanism. From proof of Theorem 3.5.2, we just need
$\alpha(\cdot,\cdot)$ uniformly select every node in ${\bm{X}}\in\mathcal{X}$.
The unnormalized bilinear attention $\alpha^{\prime}$ is weak in the sense
that $f(\cdot)=\langle{{\bm{x}}_{i}{\bm{W}}_{Q}{\bm{W}}_{K}^{T},\cdot}\rangle$
has a linear level set. Such a constraint can be relaxed via an improved
attention module GATv2. Observing the ranking of the attention scores given by
GAT [143] is unconditioned on the query node, [12] proposed GATv2, a more
expressive attention mechanism. In particular, the unnormalized attention
score
$\alpha^{\prime}_{\text{GATv2}}({\bm{u}},{\bm{v}}):={\bm{a}}^{T}\operatorname{LeakyReLU}\left({\bm{W}}\cdot\left[{\bm{u}}\|{\bm{v}}\right]+{\bm{b}}\right)$,
where $[\cdot||\cdot]$ is concatenation. We will let
$\alpha_{{\bm{n}}}=\alpha_{\text{GATv2}}$ to select features in
$\tau_{j\in[n]}\phi_{\text{vn-gn}}^{(k)}$.

Figure 3.3: In the left figure, we have one example of $\mathcal{X}$ being
$({\bm{V}},\delta)$ separable, for which $\alpha$ can uniformly select any
point (marked as red) ${\bm{x}}_{i}\in\mathcal{X}_{i}$. In the right figure,
we change $\alpha_{{\bm{n}}}$ in MPNN + VN to $\alpha_{\text{GATv2}}$, which
allows us to select more diverse feature configurations.

######  Lemma 3.10.5.

$\alpha^{\prime}_{\text{GATv2}}(\cdot,\cdot)$ can approximate any continuous
function from $\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$. For
any ${\bm{v}}\in\mathbb{R}^{d}$, a restriction of
$\alpha^{\prime}_{\text{GATv2}}(\cdot,{\bm{v}})$ can approximate any
continuous function from $\mathbb{R}^{d}\rightarrow\mathbb{R}$.

###### Proof.

Any function continuous in both arguments of $\alpha^{\prime}_{\text{GATv2}}$
is also continuous in the concatenation of both arguments. As any continuous
functions in $\mathbb{R}^{2d}$ can be approximated by
$\alpha^{\prime}_{\text{GATv2}}$ on a compact domain according to the
universality of MLP [31], we finish the proof for the first statement.

For the second statement, we can write $W$ as $2\times 2$ block matrix and
restrict it to cases where only ${\bm{W}}_{11}$ is non-zero. Then we have

$\alpha^{\prime}_{\text{GATv2}}({\bm{u}},{\bm{v}})=a^{T}\operatorname{LeakyReLU}\left(\left[{\begin{array}[]{cc}{\bm{W}}_{11}&{\bm{W}}_{12}\\\
{\bm{W}}_{21}&{\bm{W}}_{22}\\\
\end{array}}\right]\cdot\left[\begin{array}[]{c}{\bm{u}}\\\ {\bm{v}}\\\
\end{array}\right]+{\bm{b}}\right)={\bm{a}}^{T}\operatorname{LeakyReLU}\left({\bm{W}}_{11}{\bm{u}}+{\bm{b}}\right)$
(3.13)

which gives us an MLP on the first argument ${\bm{u}}$. By the universality of
MLP, we conclude the proof for the second statement.

∎

######  Definition 21.

Given $\delta>0$, We call $\mathcal{X}$ is $\delta$ nonlinearly separable if
and only if $\min_{i\neq j}d(\mathcal{X}_{i},\mathcal{X}_{j})>\delta$.

######  AS 4.

$\mathcal{X}$ is $\delta$ nonlinearly separable for some $\delta>0$. TODO: set
counter

######  Proposition 3.10.6.

If $\mathcal{X}\subset\mathbb{R}^{n\times d}$ satisfies that $\mathcal{X}_{i}$
is $\delta$-separated from $\mathcal{X}_{j}$ for any $i,j\in[n]$, the
following holds. For any ${\bm{X}}\in\mathcal{X}$ and $i\in[n]$, there exist a
$\alpha_{\text{GATv2}}$ to select any ${\bm{x}}_{i}\in\mathcal{X}_{i}$. This
implies that we can arbitrarily approximate the self-attention layer
${\bm{L}}$ after relaxing AS3 to AS3’.

###### Proof.

For any $i\in[n]$, as $\mathcal{X}_{i}$ is $\delta$-separated from other
$\mathcal{X}_{j},\forall j\neq i$, we can draw a region
$\Omega_{i}\subset\mathbb{R}^{d}$ that contains $\mathcal{X}_{i}$ and separate
$\mathcal{X}_{i}$ from other $\mathcal{X}_{j}(j\neq i)$, where the distance
from $\mathcal{X}_{i}$ from other $\mathcal{X}_{j}$ is at least $\delta$
according to the definition of Definition 21. Next, we show how to construct a
continuous function $f$ whose value in $\mathcal{X}_{i}$ is at least 1 larger
than its values in any other $\mathcal{X}_{j}$ $\forall j\neq i$.

We set the values of $f$ in $\mathcal{X}_{i}$ to be 1.5 and values of $f$ in
$\mathcal{X}_{j},\forall j\neq i$ to be 0. We can then interpolate $f$ in
areas outside of $\cup\mathcal{X}_{i}$ (one way is to set the values of $f(x)$
based on $d(x,\mathcal{X}_{i}$), which results in a continuous function that
satisfies our requirement. By the universality of $\alpha_{\text{GATv2}}$, we
can approximate $f$ to arbitrary precision, and this will let us select any
$\mathcal{X}_{i}$. ∎

##  Chapter 4 Graph Coarsening with Neural Networks

###  4.1 Introduction

In this chapter, we look into another common way of global modeling for large
graphs. As large scale-graphs become increasingly ubiquitous in various
applications, they pose significant computational challenges to process,
extract and analyze information. It is therefore natural to look for ways to
simplify the graph while preserving the properties of interest.

There are two major ways to simplify graphs. First, one may reduce the number
of edges, known as graph edge sparsification. It is known that pairwise
distance (spanner), graph cut (cut sparsifier), eigenvalues (spectral
sparsifier) can be approximately maintained via removing edges. A key result
[137] in the spectral sparsification is that any dense graph of size $N$ can
be sparsified to $O(Nlog^{c}N/\epsilon^{2})$ edges in nearly linear time using
a simple randomized algorithm based on the effective resistance.

Alternatively, one could also reduce the number of nodes to a subset of the
original node set. The first challenge here is how to choose the topology
(edge set) of the smaller graph spanned by the sparsified node set. On the
extreme, one can take the complete graph spanned by the sampled nodes.
However, its dense structure prohibits easy interpretation and poses
computational overhead for setting the $\Theta(n^{2})$ weights of edges. this
chapter focuses on _graph coarsening_ , which reduces the number of nodes by
contracting disjoint sets of connected vertices.

The original idea dates back to the algebraic multigrid literature [123] and
has found various applications in graph partitioning [61, 76, 89],
visualization [60, 70, 146] and machine learning [90, 52, 134].

However, most existing graph coarsening algorithms come with two restrictions.
First, they are _prespecified_ and not adapted to specific data nor different
goals.

Second, most coarsening algorithms set the edge weights of the coarse graph
equal to the sum of weights of crossing edges in the original graph. This
means the weights of the coarse graph is determined by the coarsening
algorithm (of the vertex set), leaving no room for adjustment.

With the two observations above, we aim to develop a data-driven approach to
better assigning weights for the coarse graph depending on specific goals at
hand.

We will leverage the recent progress of deep learning on graphs to develop a
framework to learn to assign edge weights _in an unsupervised manner_ from a
collection of input (small) graphs. This learned weight-assignment map can
then be applied to new graphs (of potentially much larger sizes).

In particular, our contributions are threefold.

  * •

First, depending on the quantity of interest $\mathcal{F}$ (such as the
quadratic form w.r.t. Laplace operator), one has to carefully choose
projection/lift operator to relate quantities defined on graphs of different
sizes. We formulate this as the invariance of $\mathcal{F}$ under lift map,
and provide three cases of projection/lift map as well as the corresponding
operators on the coarse graph. Interestingly, those operators all can be seen
as the special cases of doubly-weighted Laplace operators on coarse graphs
[66].

  * •

Second, we are the first to propose and develop a framework to learn the edge
weights of the coarse graphs via graph neural networks (GNN) in an
unsupervised manner. We show convincing results both theoretically and
empirically that changing the weights is crucial to improve the quality of
coarse graphs. many existing graph coarsening algorithms.

  * •

Third, through extensive experiments on both synthetic graphs and real
networks, we demonstrate that our method GOREN significantly improves common
graph coarsening methods under different evaluation metrics, reduction ratios,
graph sizes, and graph types. It generalizes to graphs of larger size (than
the training graphs), adapts to different losses (so as to preserve different
properties of original graphs), and scales to much larger graphs than what
previous work can handle. Even for losses that are not differentiable w.r.t
the weights of the coarse graph, we show training networks with a
differentiable auxiliary loss still improves the result.

###  4.2 Proposed Approach: Learning Edge Weight with GNN

####  4.2.1 High-level overview

Our input is a non-attributed (weighted or unweighted) graph $G=(V,E)$. Our
goal is to construct an appropriate “coarser" graph
${\widehat{G}}=({\widehat{V}},{\widehat{E}})$ that preserves certain
properties of $G$. Here, by a “coarser" graph, we assume that
$|{\widehat{V}}|<<|V|$ and there is a surjective map ${\pi}:V\to{\widehat{V}}$
that we call the _vertex map_. Intuitively, (see figure on the right), for any
node ${\hat{v}}\in{\widehat{V}}$, all nodes ${\pi}^{-1}({\hat{v}})\subset V$
are mapped to this _super-node_ ${\hat{v}}$ in the coarser graph
${\widehat{G}}$. We will later propose a GNN based framework that can be
trained using a collection of existing graphs _in an unsupervised manner_ , so
as to construct such a coarse graph ${\widehat{G}}$ for a future input graph
$G$ (presumably coming from the same family as training graphs) that can
preserve properties of $G$ effectively.

We will in particular focus on preserving properties of the _Laplace operator
${\mathcal{O}_{G}}$_ of $G$, which is by far the most common operator
associated to graphs, and forms the foundation for spectral methods.
Specifically, given $G=(V=\\{v_{1},\ldots,v_{N}\\},E)$ with
$w:E\to{\mathbb{R}}$ being the weight function for $G$ (all edges have weight
1 if $G$ is unweighted), let $W$ the corresponding $N\times N$ edge-weight
matrix where $W[i][j]=w(v_{i},v_{j})$ if edge $(v_{i},v_{j})\in E$ and $0$
otherwise. Set $D$ to be the $N\times N$ diagonal matrix with $D[i][i]$ equal
to the sum of weights of all edges incident to $v_{i}$. The standard _(un-
normalized) combinatorial Laplace operator_ of $G$ is then defined as
${L}=D-W$. The _normalized Laplacian_ is defined as
$\mathcal{L}=D^{-1/2}{L}D^{-1/2}=I-D^{-1/2}WD^{-1/2}$.

However, to make this problem as well as our proposed approach concrete,
various components need to be built appropriately. We provide an overview
here, and they will be detailed in the remainder of this section.

  * •

Assuming that the set of super-nodes ${\widehat{V}}$ as well as the map
${\pi}:V\to{\widehat{V}}$ are given, one still need to decide how to set up
the connectivity (i.e, edge set ${\widehat{E}}$) for the coarse graph
${\widehat{G}}=({\widehat{V}},{\widehat{E}})$. We introduce a natural choice
in Section 4.2.2, and provide some justification for this choice.

  * •

As the graph $G$ and the coarse graph ${\widehat{G}}$ have the different
number of nodes, their Laplace operators ${\mathcal{O}_{G}}$ and
${\mathcal{O}_{{\widehat{G}}}}$ of two graphs are not directly comparable.
Instead, we will compare ${\mathcal{F}}({\mathcal{O}_{G}},f)$ and
${\mathcal{F}}({\mathcal{O}_{{\widehat{G}}}},{\hat{f}})$, where
${\mathcal{F}}$ is a functional intrinsic to the graph at hand (invariant to
the permutation of vertices), such as the quadratic form or Rayleigh quotient.
However, it turns out that depending on the choice of ${\mathcal{F}}$, we need
to choose the precise form of the Laplacian ${\mathcal{O}_{{\widehat{G}}}}$,
as well as the (so-called lifting and projection) maps relating these two
objects, carefully, so as they are comparable. We describe these in detail in
Section 4.2.3.

  * •

In Section 4.2.4 we show that adjusting the weights of the coarse graph
${\widehat{G}}$ can significantly improve the quality of ${\widehat{G}}$. This
motivates a learning approach to learn a strategy (a map) to assign these
weights from a collection of given graphs. We then propose a GNN-based
framework to do so in an unsupervised manner. Extensive experimental studies
will be presented in Section 4.3.

####  4.2.2 Construction of Coarse graph

Assume that we are already given the set of super-nodes
${\widehat{V}}=\\{{\hat{v}}_{1},\ldots,{\hat{v}}_{n}\\}$ for the coarse graph
${\widehat{G}}$ together with the vertex map $\pi:V\to{\widehat{V}}$ – There
has been much prior work on computing the sparsified set ${\widehat{V}}\subset
V$ and $\pi$ [105, 104]; and if the vertex map $\pi$ is not given, then we can
simply define it by setting $\pi(v)$ for each $v\in V$ to be the nearest
neighbor of $v$ in ${\widehat{V}}$ in terms of graph shortest path distance in
$G$ [36].

To construct edges for the coarse graph
${\widehat{G}}=({\widehat{V}},{\widehat{E}})$ together with the edge weight
function ${\hat{w}}:{\widehat{E}}\to{\mathbb{R}}$, instead of using a complete
weighted graph over ${\widehat{V}}$, which is too dense and expensive, we set
${\widehat{E}}$ to be those edges “induced" from $G$ when collapsing each
_cluster ${\pi}^{-1}({\hat{v}})$_ to its corresponding super-node
${\hat{v}}\in{\widehat{V}}$: Specifically,
$({\hat{v}},{\hat{v}}^{\prime})\in{\widehat{E}}$ if and only if there is an
edge $(v,v^{\prime})\in E$ such that ${\pi}(v)={\hat{v}}$ and
${\pi}(v^{\prime})={\hat{v}}^{\prime}$.

The weight of this edge is
${\hat{w}}({\hat{v}},{\hat{v}}^{\prime}):=\sum_{(v,v^{\prime})\in
E\big{(}{\pi}^{-1}({\hat{v}}),{\pi}^{-1}({\hat{v}}^{\prime})\big{)}}w(v,v^{\prime})$

where $E(A,B)\subseteq E$ stands for the set of edges crossing sets
$A,B\subseteq V$;

i.e., ${\hat{w}}({\hat{v}},{\hat{v}}^{\prime})$ is the total weights of all
crossing edges in $G$ between clusters ${\pi}^{-1}({\hat{v}})$ and
${\pi}^{-1}({\hat{v}}^{\prime})$ in $V$. We refer to ${\widehat{G}}$
constructed this way the _${\widehat{V}}$ -induced coarse graph_.

As shown in [36], if the original graph $G$ is the $1$-skeleton of a hidden
space $X$, then this induced graph captures the topological of $X$ at a
coarser level if ${\widehat{V}}$ is a so-called $\delta$-net of the original
vertex set $V$ w.r.t. the graph shortest path metric.

Let ${\widehat{W}}$ be the edge weight matrix, and ${\widehat{D}}$ be the
diagonal matrix encoding the sum of edge weights incident to each vertex as
before. Then the standard combinatorial Laplace operator w.r.t.
${\widehat{G}}$ is simply ${{\widehat{L}}}={\widehat{D}}-{\widehat{W}}$.

Relation to the operator of [104].  Interestingly, this construction of the
coarse graph ${\widehat{G}}$ coincides with the coarse Laplace operator for a
sparsified vertex set ${\widehat{V}}$ constructed by [104]. We will use this
view of the Laplace operator later; hence we briefly introduce the
construction of [104] (adapted to our setting):

Given the vertex map $\pi:V\to{\widehat{V}}$, we set a $n\times N$ matrix $P$
by
$P[r,i]=\left\\{\begin{array}[]{ll}\frac{1}{\left|{\pi}^{-1}({\hat{v}}_{r})\right|}&\text{
if }v_{i}\in{\pi}^{-1}({\hat{v}}_{r})\\\ 0&\text{ otherwise
}\end{array}\right.$. In what follows, we denote
$\gamma_{r}:=\left|{\pi}^{-1}({\hat{v}}_{r})\right|$ for any $r\in[1,n]$,
which is the size of the cluster of ${\hat{v}}_{r}$ in $V$. $P$ can be
considered as the weighted projection matrix of the vertex set from $V$ to
${\widehat{V}}$. Let $P^{+}$ denote the Moore-Penrose pseudoinverse of $P$,
which can be intuitively viewed as a way to lift a function on ${\widehat{V}}$
(a vector in ${\mathbb{R}}^{n}$) to a function over $V$ (a vector in
${\mathbb{R}}^{N}$). As shown in [104], $P^{+}$ is the $N\times n$ matrix
where $P^{+}[i,r]=1$ if and only if ${\pi}(v_{i})={\hat{v}}_{r}$. See Section
4.10.1 for a toy example.

Finally, [104] defines an operator for the coarsened vertex set
${\widehat{V}}$ to be $\tilde{L}_{\widehat{V}}=(P^{+})^{T}{L}P^{+}$.
Intuitively, ${{\widehat{L}}}$ operators on $n$-vectors. For any $n$-vector
${\hat{f}}\in{\mathbb{R}}^{n}$, $\tilde{L}_{\widehat{V}}{\hat{f}}$ first lifts
${\hat{f}}$ to a $N$-vector $f=P^{+}{\hat{f}}$, and then perform ${L}$ on $f$,
and then project it down to $n$-dimensional via $(P^{+})^{T}$.

######  Proposition 4.2.1.

[104] The combinatorial graph Laplace operator
${\widehat{L}}={\widehat{D}}-{\widehat{W}}$ for the ${\widehat{V}}$-induced
coarse graph ${\widehat{G}}$ constructed above equals to the operator
$\tilde{L}_{\widehat{V}}=(P^{+})^{T}{L}P^{+}$.

####  4.2.3 Laplace Operator for the Coarse Graph

We now have an input graph $G=(V,E)$ and a coarse graph ${\widehat{G}}$
induced from the sparsified node set ${\widehat{V}}$, and we wish to compare
their corresponding Laplace operators. However, as ${\mathcal{O}_{G}}$
operates on ${\mathbb{R}}^{N}$ (i.e, functions on the vertex set $V$ of $G$)
and ${\mathcal{O}_{{\widehat{G}}}}$ operates on ${\mathbb{R}}^{n}$, we will
compare them by their effects on “corresponding" objects.

[105, 104] proposed to use the quadratic form to measure the similarity
between the two linear operators. In particular, given a linear operator $A$
on ${\mathbb{R}}^{N}$ and any $x\in{\mathbb{R}}^{N}$,
${\mathsf{Q}}_{A}(x)=x^{T}Ax$. The quadratic form has also been used for
measuring spectral approximation under edge sparsification. The proof of the
following result is in Section 4.10.1.

######  Proposition 4.2.2.

For any vector $\hat{x}\in{\mathbb{R}}^{n}$, we have that
${\mathsf{Q}}_{{{\widehat{L}}}}(\hat{x})={\mathsf{Q}}_{{L}}(P^{+}\hat{x})$,
where ${\widehat{L}}$ is the combinatorial Laplace operator for the
${\widehat{V}}$-induced coarse graph ${\widehat{G}}$ constructed above. That
is, set $x:=P^{+}\hat{x}$ as the lift of $\hat{x}$ in ${\mathbb{R}}^{N}$, then
$\hat{x}^{T}{{\widehat{L}}}\hat{x}=x^{T}{L}x$.

Intuitively, this suggests that if later, we measure the similarity between
${L}$ and some Laplace operator for the coarse graph ${\widehat{G}}$ based on
a loss from quadratic form difference, then we should choose the Laplace
operator ${\mathcal{O}_{{\widehat{G}}}}$ to be ${{\widehat{L}}}$ and compare
${\mathsf{Q}}_{{{\widehat{L}}}}(Px)$ with ${\mathsf{Q}}_{{L}}(x)$. We further
formalize this by considering the _lifting map_
${\mathcal{U}}:{\mathbb{R}}^{n}\to{\mathbb{R}}^{N}$ as well as a projection
map ${\mathcal{P}}:{\mathbb{R}}^{N}\to{\mathbb{R}}^{n}$, where
${\mathcal{P}}\cdot{\mathcal{U}}=Id_{n}$. Proposition 4.2.2 suggests that for
quadratic form-based similarity, the choices are
${\mathcal{U}}=P^{+},{\mathcal{P}}=P$, and
${\mathcal{O}_{{\widehat{G}}}}={{\widehat{L}}}$. See the first row in Table
4.1.

Table 4.1: Depending on the choice of $\mathcal{F}$ (quantity that we want to
preserve) and ${\mathcal{O}_{G}}$, we have different projection/lift operators
and resulting ${\mathcal{O}_{{\widehat{G}}}}$ on the coarse graph.

Quantity $\mathcal{F}$ of interest | ${\mathcal{O}_{G}}$ | Projection ${\mathcal{P}}$ | Lift ${\mathcal{U}}$ | ${\mathcal{O}_{{\widehat{G}}}}$ | Invariant under ${\mathcal{U}}$
---|---|---|---|---|---
Quadratic form ${\mathsf{Q}}$ | $L$ | $P$ | ${P^{+}}$ | Combinatorial Laplace $\widehat{L}$ | ${\mathsf{Q}}_{L}({\mathcal{U}}\hat{x})={\mathsf{Q}}_{\widehat{L}}(\hat{x})$
Rayleigh quotient ${\mathsf{R}}$ | $L$ | $\Gamma^{-1/2}{{(P^{+})}^{T}}$ | ${P^{+}}\Gamma^{-1/2}$ | Doubly-weighted Laplace $\widehat{{\mathsf{L}}}$ | ${\mathsf{R}}_{L}({\mathcal{U}}\hat{x})={\mathsf{R}}_{\widehat{{\mathsf{L}}}}(\hat{x})$
Quadratic form ${\mathsf{Q}}$ | $\mathcal{L}$ | $\widehat{D}^{1/2}PD^{-1/2}$ | $D^{1/2}{(P^{+})}\widehat{D}^{-1/2}$ | Normalized Laplace $\widehat{\mathcal{L}}$ | ${\mathsf{Q}}_{\mathcal{L}}({\mathcal{U}}\hat{x})={\mathsf{Q}}_{\widehat{\mathcal{L}}}(\hat{x})$

On the other hand, eigenvectors and eigenvalues of a linear operator $A$ are
more directly related, via Courant-Fischer Min-Max Theorem, to its Rayleigh
quotient ${\mathsf{R}}_{A}(x)=\frac{x^{T}Ax}{x^{T}x}$. Interestingly, in this
case, to preserve the Rayleigh quotient, we should change the choice of
${\mathcal{O}_{{\widehat{G}}}}$ to be the following _doubly-weighted Laplace
operator_ for a graph that is both edge and vertex weighted.

Specifically, for the coarse graph ${\widehat{G}}$, we assume that each vertex
${\hat{v}}\in{\widehat{V}}$ is weighted by
$\gamma_{\hat{v}}:=|{\pi}^{-1}({\hat{v}})|$, the size of the cluster from $G$
that got collapsed into ${\hat{v}}$. Let $\Gamma$ be the vertex matrix, which
is the $n\times n$ diagonal matrix with $\Gamma[r][r]=\gamma_{{\hat{v}}_{r}}$.
The _doubly-weighted Laplace operator_ for a vertex- and edge-weighted graph
${\widehat{G}}$ is then defined as:

$\displaystyle\widehat{{\mathsf{L}}}=\Gamma^{-1/2}({\widehat{D}}-{\widehat{W}})\Gamma^{-1/2}=\Gamma^{-1/2}{{\widehat{L}}}\Gamma^{-1/2}=(P^{+}\Gamma^{-1/2})^{T}{L}(P^{+}\Gamma^{-1/2}).$

The concept of doubly-weighted Laplace for a vertex- and edge-weighted graph
is not new, see e.g [30, 66, 157]. In particular, [66] proposes a general form
of combinatorial Laplace operator for a simplicial complex where all simplices
are weighted, and our doubly-weighted Laplace has the same eigenstructure as
their Laplacian when restricted to graphs. See Section 4.10.1 for details.
Using the doubly-weighted Laplacian for Rayleigh quotient based similarity
measurement between the original graph and the coarse graph is justified by
the following result (proof in Section 4.10.1).

######  Proposition 4.2.3.

For any vector $x\in{\mathbb{R}}^{n}$, we have that
${\mathsf{R}}_{\widehat{{\mathsf{L}}}}(\hat{x})={\mathsf{R}}_{{L}}(P^{+}\Gamma^{-1/2}\hat{x})$.
That is, set the lift of $\hat{x}$ in ${\mathbb{R}}^{N}$ to be
$x=P^{+}\Gamma^{-1/2}\hat{x}$, then we have that
$\frac{\hat{x}^{T}\widehat{{\mathsf{L}}}\hat{x}}{\hat{x}^{T}\hat{x}}=\frac{x^{T}{L}x}{x^{T}x}$.

Finally, if using the normalized Laplace ${\mathcal{L}}$ for the original
graph $G$, then the appropriate Laplace operator for the coarse graph and
corresponding projection/lift maps are listed in the last row of Table 4.1,
with proofs in Section 4.10.1.

####  4.2.4 A GNN-based Framework for Constructing the Coarse Graph

Figure 4.1: An illustration of learnable coarsening framework. Existing
coarsening algorithm determines the topology of coarse graph ${\widehat{G}}$,
while GOREN resets the edge weights of the coarse graph.

In the previous section, we argued that depending on what similarity measures
we use, appropriate Laplace operator ${\mathcal{O}_{{\widehat{G}}}}$ for the
coarse graph ${\widehat{G}}$ should be used. Now consider the specific case of
Rayleigh quotient, which can be thought of as a proxy to measure similarities
between the low-frequency eigenvalues of the original graph Laplacian and the
one for the coarse graph. As described above, here we set
${\mathcal{O}_{{\widehat{G}}}}$ as the doubly-weighted Laplacian
$\widehat{{\mathsf{L}}}=\Gamma^{-1/2}({\widehat{D}}-{\widehat{W}})\Gamma^{-1/2}$.

###### The effect of weight adjustments.

We develop an iterative algorithm with convergence guarantee (to KKT point in
4.10.2) for optimizing over edge weights of ${\widehat{G}}$ for better
spectrum alignment. As shown in the figure on the right, after changing the
edge weight of the coarse graph, the resulting graph Laplacian has eigenvalues
much closer (almost identical) to the first $n$ eigenvalues of the original
graph Laplacian. More specifically, in this figure, $G.e$ and $Gc.e$ stand for
the eigenvalues of the original graph $G$ and coarse graph ${\widehat{G}}$
constructed by the so-called Variation-Edge coarsening algorithm [104].
“After-Opt" stands for the eigenvalues of coarse graphs when weights are
optimized by our iterative algorithm. See Section 4.10.2 for the description
of our iterative algorithm, its convergence results, and full experiment
results.

A GNN-based framework for learning weight assignment map.  The discussions
above indicate that we can obtain better Laplace operators for the coarse
graph by using better-informed weights than simply summing up the weights of
crossing edges from the two clusters.

More specifically, suppose we have a fixed strategy to generate
${\widehat{V}}$ from an input graph $G=(V,E)$. Now given an edge
$({\hat{v}},{\hat{v}}^{\prime})\in{\widehat{E}}$ in the induced coarse graph
${\widehat{G}}=({\widehat{V}},{\widehat{E}})$, we model its weight
${\hat{w}}({\hat{v}},{\hat{v}}^{\prime})$ by a _weight-assignment function_
$\mu(G|_{{\pi}^{-1}({\hat{v}})\cup{\pi}^{-1}({\hat{v}}^{\prime})})$, where
$G|_{A}$ is the subgraph of $G$ induced by a subset of vertices $A$. However,
it is not clear how to setup this function $\mu$. Instead, we will learn it
from a collection of input graphs in an _unsupervised_ manner. Specifically,
we will parametrize the weight-assignment map $\mu$ by a learnable neural
network ${\mathcal{M}_{\theta}}$.

See Figure 4.1 for an illustration.

In particular, we use Graph Isomorphism Network (GIN) [156] to represent
${\mathcal{M}_{\theta}}$. We initialize the model by setting the edge
attribute of the coarse graph to be 1. Our node feature is set to be a
5-dimensional vector based on LDP (Local Degree Profile) [20]. We enforce the
learned weight of the coarse graph to be positive by applying one extra ReLU
layer to the final output. All models are trained with Adam optimizer with a
learning rate of 0.001. See Section 4.3.3 for more details. We name our model
as Graph cOarsening RefinemEnt Network (GOREN).

Given a graph $G$ and a coarsening algorithm $\mathcal{A}$, the general form
of loss is

$\displaystyle
Loss({\mathcal{O}_{G}},\mathcal{O}_{\widehat{G_{t}}})=\frac{1}{k}\sum_{i=1}^{k}|\mathcal{F}(\mathcal{O}_{G},f_{i})-\mathcal{F}(\mathcal{O}_{\widehat{G_{t}}},{\mathcal{P}}f_{i})|,$
(4.1)

where $f_{i}$ is signal on the original graph (such as eigenvectors) and
${\mathcal{P}}f_{i}$ is its projection. We use $\mathcal{O}_{\widehat{G_{t}}}$
to denote the operator of the coarse graph during training, while
$\mathcal{O}_{\widehat{G}}$ standing for the operator defined w.r.t. the
coarse graph output by coarsening algorithm $\mathcal{A}$. That is, we will
start with $\mathcal{O}_{{\widehat{G}}}$ and modify it to
$\mathcal{O}_{\widehat{G_{t}}}$ during the training.

The loss can be instantiated for different cases in Table 4.1. For example, a
loss based on quadratic form means that we choose
$\mathcal{O}_{G},\mathcal{O}_{\widehat{G_{t}}}$ to be the combinatorial