pagerefparensbackrefparens

# PatchGT: Transformer over Non-trainable Clusters for Learning Graph
Representations

Han Gao
University of Notre Dame
<EMAIL_ADDRESS>Han11footnotemark: 1
Tufts University
<EMAIL_ADDRESS>Huang11footnotemark: 1
University of Pennsylvania
<EMAIL_ADDRESS>Wang
University of Notre Dame
<EMAIL_ADDRESS>Liu
Tufts University
<EMAIL_ADDRESS>Equal contribution.

###### Abstract

Recently the Transformer structure has shown good performances in graph
learning tasks. However, these Transformer models directly work on graph nodes
and may have difficulties learning high-level information. Inspired by the
vision transformer, which applies to image patches, we propose a new
Transformer-based graph neural network: Patch Graph Transformer (PatchGT).
Unlike previous transformer-based models for learning graph representations,
PatchGT learns from non-trainable graph patches, not from nodes directly. It
can help save computation and improve the model performance. The key idea is
to segment a graph into patches based on spectral clustering without any
trainable parameters, with which the model can first use GNN layers to learn
patch-level representations and then use Transformer to obtain graph-level
representations. The architecture leverages the spectral information of graphs
and combines the strengths of GNNs and Transformers. Further, we show the
limitations of previous hierarchical trainable clusters theoretically and
empirically. We also prove the proposed non-trainable spectral clustering
method is permutation invariant and can help address the information
bottlenecks in the graph. PatchGT achieves higher expressiveness than 1-WL-
type GNNs, and the empirical study shows that PatchGT achieves competitive
performances on benchmark datasets and provides interpretability to its
predictions. The implementation of our algorithm is released at our Github
repo: https://github.com/tufts-ml/PatchGT.

##  1 Introduction

Learning from graph data is ubiquitous in applications such as drug design
[15] and social network analysis [37]. The success of a graph learning task
hinges on effective extraction of information from graph structures, which
often contain combinatorial structures and are highly complex. Early works [7]
often need to manually extract features from graphs before applying learning
models. In the era of deep learning, Graph Neural Networks (GNNs) [35] are
developed to automatically extract information from graphs. Through passing
learnable messages between nodes, they are able to encode graph information
into vector representations of graph nodes. GNNs have become the standard tool
for learning tasks on graph data.

While they have achieved good performances in a wide range of tasks, GNNs
still have a few limitations. For example, GNNs [36] suffer from issues such
as inadequate expressiveness [36], over-smoothing [28], and over-squashing
[2]. These issues have been partially addressed by techniques such as
improving message-passing functions and expanding node features [21, 5].

Another important progress is to replace the message-passing network with the
Transformer architecture [38, 18, 6, 24]. These models treat graph nodes as
tokens and apply the Transformer architecture to nodes directly. The main
focus of these models is how to encode node information and how to incorporate
adjacency matrices into network calculations. Without the message-passing
structure, these models may overcome some associated issues and have shown
premium performances in various graph learning tasks. However, these models
suffer from computation complexity because of the global attention on all
nodes. It is hard to capture the topological information of graphs.

As a comparison, the Transformer for image data works on image patches instead
of pixels [9, 22]. While this model choice is justified by reduction of
computation cost, recent work [31] shows that “patch representation itself may
be a critical component to the ‘superior’ performance of newer architectures
like Vision Transformers”. One intriguing question is whether patch
representation can also improve learning models on graphs. With this question,
we consider patches on graphs. Patches over graphs are justified by a “mid-
level” understanding of graphs: for example, a molecule graph’s property is
often decided by some function groups, each of which is a subgraph formed by
locally-connected atoms. Therefore, patch representations are able to capture
such mid-level concepts and bridge the gap between low-level structures to
high-level semantics.

Motivated by our question, we propose a new framework, Patch Graph Transformer
(PatchGT). It first segments a graph into patches based on spectral
clustering, which is a non-trainable segmentation method, then applies GNN
layers to learn patch representations, and finally uses Transformer layers to
learn a graph-level representation from patch representations. This framework
combines the strengths of two types of learning architectures: GNN layers can
extract information with message passing, while Transformer layers can
aggregate information using the attention mechanism. To our best knowledge, we
firstly show several limitations of previous trainable clustering method based
on GNN. We also show that the proposed non-trainable clustering can provide
more reasonable patches and help overcoming information bottleneck in graphs.

We justify our model architecture with theoretical analysis. We show that our
patch structure derived from spectral clustering is superior to patch
structures learned by GNNs [39, 4, 13]. We also propose a new mathematical
description of the information bottleneck in vanilla GNNs and further show
that our architecture has the ability of mitigating this issue when graphs
have small graph cuts. The contributions of this paper are summarized as
follows.

  * -

We develop a general framework to overcome the information bottleneck in
traditional GNNs by applying a Transformer on graph patches in Section 3. The
graph patches are from an unlearnable spectral clustering process.

  * -

We prove several new theorems for the limitations of previous pooling methods
from the 1-WL algorithm in Theorem 1 and Theorem 2. And we theoretically prove
that PatchGT is strictly beyond 1-WL and hence has better expressiveness in
Theorem 3. Also, in Section 4.4, we show that the segmentations from
hierarchical learnable clustering methods may aggregate disconnected nodes,
which will definitely hurt the performance of the transformer model.

  * -

We demonstrate the existence of information bottleneck in GNNs in Section 4.3.
When a graph consists of loosely-connected clusters, we make the first attempt
to characterize such information bottleneck. And it indicates when there is a
small graph cut between two clusters, the GNNs need to use more layers to pass
signals from one group to another. And we further demonstrate with direct
attention between groups, PatchGT could overcome such limitations.

We run an extensive empirical study and demonstrate that the proposed model
outperforms competing methods on a list of graph learning tasks. The ablation
study shows that our PatchGT is able to combine the strengths of GNN layers
and Transformer layers. The attention weights in Transformer layers also
provide explanations for model predictions.

##  2 Related Work

Transformer models have gained remarkable successes in NLP applications [16].
Recently, they have also been introduced to vision tasks [9] and graph tasks
[38, 18, 6, 24, 11, 20, 34, 40]. These models all treat nodes as tokens.
Particularly, Memory-based graph networks[1] apply a hierarchical attention
pooling methods on the nodes. GraphTrans [34] directly applies a GNN on all
nodes, followed by a transformer. Therefore, they are hard to be applied to
large graphs because of huge computation complexity.

At the same time, image patches have been shown to be useful for Transformer
models on image data [31, 9], so it is not surprising if graph patches are
also helpful to Transformer models on graph data. Graph multiset pooling [3]
applies trainable pooling methods on the nodes based on GNN. And then adopt a
global attention layer on learned clusters. We will show that such trainable
clustering has several limitations for attention mechanism in this work.

Hierarchical pooling models [39, 12, 4, 13, 27, 19, 39, 12, 4, 13, 27, 19] are
relevant to our work in that they also aggregate information from node
representations in middle layers of networks. However, these methods all form
their pooling structures based on representations learned from GNNs. As a
result, these pooling structures inherit drawbacks from GNNs [36]. They may
also aggregate nodes that are far apart on the graph and thus cannot preserve
the global structure of the input graph. Also such trainable clustering
methods need much computation for training. Furthermore, our main purpose is
to use non-trainable patches on graphs as tokens for a Transformer model,
which is different from these models.

Figure 1: Model review. We segment a graph into several patch subgraphs by
non-trainable clustering. We first extract local information through a GNN,
and the initial patch representations are summarized by the aggregation of
nodes within the corresponding patches. To further encode structure
information, we apply another patch-level GNN to update the representations of
patches. Finally, we use Transformer to extract the representation of the
entire graph based on patch representations.

##  3 Patch Graph Transformer

###  3.1 Background

In this work, we consider graph-level learning problems. Let $G=(V,E)$ denote
a graph with node set $V$ and edge set $E$. Let $\mathbf{A}$ denote its
adjacency matrix. The graph has both node features
$\mathbf{X}=(\mathbf{x}_{i}\in\mathbb{R}^{d}:i\in V)$ and edge features
$\mathbf{E}=(\mathbf{e}_{i,j}\in\mathbb{R}^{d^{\prime}}:(i,j)\in E)$. Let $y$
denote the label of graph. This work aims to learn a model that maps
$(\mathbf{A},\mathbf{X},\mathbf{E})$ to a vector representation $\mathbf{g}$,
which is then used to predict the graph label $y$.

GNN layers. A GNN uses node vectors to represent structural information of the
graph. It consists of multiple GNN layers. Each GNN layer passes learnable
messages and updates node vectors. Suppose
$\mathbf{H}=(\mathbf{h}_{i}\in\mathbb{R}^{d^{\prime\prime}}:i\in V)$ are node
vectors, a typical GNN layer updates $\mathbf{H}$ as follows.

$\displaystyle\mathbf{h}^{\prime}_{i}=\sigma(\mathbf{W}_{1}\mathbf{h}_{i}+\sum_{j:(i,j)\in
E}\mathbf{W}_{2}\mathbf{h}_{j}+\mathbf{W}_{3}\mathbf{e}_{i,j})$ (1)

Here matrices $(\mathbf{W}_{1},\mathbf{W}_{2},\mathbf{W}_{3})$ are all
learnable parameters; and $\sigma$ is the activation function. We denote the
layer function by
$\mathbf{H}^{\prime}=\mathrm{GNN}(\mathbf{A},\mathbf{E},\mathbf{H})$. If there
are no edge features, then the calculation can be written in matrix form.

$\displaystyle\mathbf{H}^{\prime}=\sigma(\mathbf{H}\mathbf{W}_{1}^{\top}+\mathbf{A}\mathbf{H}\mathbf{W}_{2}^{\top})$
(2)

###  3.2 Model design

PatchGT has three components: segmenting the input graph into patches,
learning patch representations, and aggregating patch representations into a
single graph vector. The overall architecture is shown in Figure 1. The second
and third steps are in an end-to-end learning model. Graph segmentation is
outside of the learning model, which will be justified by our theoretical
analysis later.

Forming patches over the graph. We first discuss how to form patches on a
graph. One consideration is to include an informative subgraph (e.g., a
function group, a motif) into a single patch instead of segmenting it into
pieces. A reasonable approach is to run node clustering on the input graph and
treat each cluster as a graph patch. If a meaningful subgraph is densely
connected, it has a good chance of being contained in a single cluster.

In this work, we consider spectral clustering [30, 41] for graph segmentation.
Let $\mathbf{L}=\mathbf{I}-\mathbf{D}^{-1/2}\mathbf{A}\mathbf{D}^{-1/2}$ be
the normalized Laplacian matrix of $G$, and its eigen-decomposition is
$\mathbf{L}=\mathbf{U}\mathbf{\Lambda}\mathbf{U}^{\top}$, where the eigen-
values $\mathbf{\Lambda}=\mathrm{diag}(\lambda_{1},\ldots,\lambda_{|V|})$ is
sorted in the ascending order. By thresholding eigen-values with a small
threshold $\gamma$, we get
$k=\operatorname*{arg\,max}_{k^{\prime}}\lambda_{k^{\prime}}\leq\gamma$ eigen-
vectors $\mathbf{U}_{1:k}$, then we run $k$-means to get $k$ clusters (denoted
by $\mathcal{P}$) of graph nodes. Here $\mathcal{P}=\\{C_{k^{\prime}}\subset
V:k^{\prime}=1,\ldots,k\\}$ with each $C_{k^{\prime}}$ representing a
cluster/patch. Note that the threshold $\gamma$ is a hyper-parameter, and $k$
varies depending on the underlying graph’s topology.

Computing patch representations. When we learn representations of patches in
$\mathcal{P}$, we consider both node connections within the patch and also
connections between patches. Patches form a coarse graph, which is also
referred as a patch-level graph, by treating patches as nodes and their
connections as edges. We first learn node representations using GNN layers.
Let $\mathbf{H}_{0}=\mathbf{X}$ denote the initial representations of all
nodes. Then we apply $L_{1}$ GNN layers to get node representations
$\mathbf{H}_{L_{1}}$.

$\displaystyle\mathbf{H}_{\ell}=\mathrm{GNN}(\mathbf{A},\mathbf{E},\mathbf{H}_{\ell-1}),\ell=1,\ldots,L_{1}$
(3)

Here for easier discussion, we apply GNN layers to the entire graph. We have
also tried to apply GNN layers within each patch only and found that the
performance is similar.

Then we read out the initial patch representation by summarizing
representations of nodes within this patch. Let $\mathbf{z}^{0}_{k^{\prime}}$
denote the initial patch representation, then

$\mathbf{z}_{k^{\prime}}^{0}=\frac{|C_{k^{\prime}}|}{|V|}\cdot\mathrm{readout}(\mathbf{h}^{L_{1}}_{i}:i\in
C_{k^{\prime}}),k^{\prime}=1,\ldots,k$ (4)

Here $\mathbf{h}^{L_{1}}_{i}$ is node $i$’s representation in
$\mathbf{H}_{L_{1}}$. We collectively denote these patch representations in a
matrix $\mathbf{Z}_{0}=(\mathbf{z}_{k^{\prime}}^{0}:k^{\prime}=1,\ldots,k)$.
The readout function $\mathrm{readout}(\cdot)$ is a function aggregating
information from a set of vectors. Our implementation uses the max pooling. We
use the factor $\frac{|C_{k^{\prime}}|}{|V|}$ to assign proper weights to
patch representations.

To further refine patch representations and encode structural information of
the entire graph, we apply further GNN layers to the patch-level formed by
patches. We first compute the adjacency matrix $\tilde{\mathbf{A}}$ of the
patch-level graph. If we convert the partition $\mathcal{P}$ to an assignment
matrix $\mathbf{S}=(S_{i,k^{\prime}}:i\in V,k^{\prime}=1,\ldots k)$ such that
$S_{i,k^{\prime}}=1[i\in C_{k^{\prime}}]$, then the adjacency matrix over
patches is

$\displaystyle\tilde{\mathbf{A}}=1\big{[}(\mathbf{S}^{\top}\mathbf{A}\mathbf{S})>0\big{]}.$
(5)

Note that $\tilde{\mathbf{A}}$ only has connections between patches and does
not maintain connection strength.

We then compute use $L_{2}$ GNN layers to refine patch representations.

$\displaystyle\mathbf{Z}_{\ell}=\mathrm{GNN}(\tilde{\mathbf{A}},\mathbf{0},\mathbf{Z}_{\ell-1}),\quad\ell=1,\ldots,L_{2}$
(6)

GNN layers here do not have edge features. From the last layer, we get patch
representations in $\mathbf{Z}_{L_{2}}$

Graph representation via Transformer layers. Then we use $L_{3}$ Transformer
layers to extract the representation of the entire graph. Here we use a
learnable query vector $\mathbf{q}_{0}$ to “retrieve” the global
representation $\mathbf{g}$ of the graph from patch representations
$\mathbf{Z}_{L_{2}}$.

$\displaystyle\mathbf{q}^{\prime}_{\ell}$
$\displaystyle=\mathrm{MHA}\left(\mathbf{q}_{\ell-1},\mathbf{Z}_{L_{2}},\mathbf{Z}_{L_{2}}\right),\quad\ell=1,\ldots,L_{3}$
(7) $\displaystyle\mathbf{q}_{\ell}$
$\displaystyle=\mathrm{MLP}(\mathbf{q}^{\prime}_{\ell})+\mathbf{q}_{\ell-1},\quad\ell=1,\ldots,L_{3}$
(8) $\displaystyle\mathbf{g}$ $\displaystyle=\mathrm{LN}(\mathbf{q}_{L_{3}})$
(9)

Here $\mathrm{MHA}(\cdot,\cdot,\cdot)$ is the function of a multi-head
attention layer (please refer to Chp. 10 of [42]). Its three arguments are the
query, key, and value. The two functions $\mathrm{MLP}(\cdot)$ and
$\mathrm{LN}(\cdot)$ are respectively a multi-layer perceptron and a linear
layer. Note that patch representations $\mathbf{Z}_{L_{2}}$ are carried
through without being updated. Only the query token is updated to query
information from patch representations. The final learned graph representation
is $\mathbf{g}$, from which we can perform various graph level tasks.

##  4 Theoretical Analysis

In this section, we study the theoretical properties of the proposed model. To
save space, we put all proofs in the appendix.

###  4.1 Enhancing model expressiveness with patches

On purpose we form graph patches using a clustering method that is not part of
the neural network. An alternative consideration is to learn such cluster
assignments with GNNs (e.g. DiffPool [39] and MinCutPool[4]). However, cluster
assignment learned by GNNs inherits the limitation of GNNs and hinders the
expessiveness of the entire model.

######  Theorem 1.

Suppose two graphs receive the same coloring by 1-WL algorithm, then DiffPool
will compute the same vector representation for them.

Although DiffPool and MinCutPool claims to cluster “similar” graph nodes into
clusters during pooling, but these nodes may not be connected. Because of the
limitation of GNNs, they may aggregate nodes that are far apart in the graph.
For example, nodes in the same orbit always get the same color by the 1-WL
algorithm and also the same representations from a GNN, then these nodes
always have the same cluster assignment. Merging these nodes into the same
cluster does not seem capture the high-level structure of a graph.

Another prominent pooling method is the Graph U-Net [12], which has similar
issues. We briefly introduce its calculation here. Suppose the layer input is
$(\mathbf{A},\mathbf{H})$, the model’s pooling layer projects $\mathbf{H}$
with a unit vector $\mathbf{p}$ and gets values
$\mathbf{v}=\mathbf{H}\mathbf{p}$ for all nodes, then it chooses the top $k$
nodes that have largest values in $\mathbf{v}$ and keep their representations
only. We will show that this approach is NOT invariant to node orders.

We also consider a small variant of Graph U-Net for analysis convenience.
Instead of choosing $k$ nodes with top values in $\mathbf{v}$, the variant
uses a threshold $\beta$ (either learnable or a hyper-parameter) to choose
nodes: $\mathbf{b}=\mathbf{v}\geq\beta$. Then the output of the layer is
$(\mathbf{A}[\mathbf{b},\mathbf{b}],\mathbf{H}[\mathbf{b}])$. We call the
model with the variant with thresholding as Graph U-Net-th. We show that the
variant of Graph U-Net-th is also bounded by the 1-WL algorithm.

######  Theorem 2.

Suppose two graphs receive the same coloring by 1-WL algorithm, then Graph
U-Net-th will compute the same vector representation for them.

The two theorems strongly indicate that pooling structures learned by GNNs
have the same drawback. We provide detailed analysis for Graph U-Net in
Section A.3.

In contrast, a small variant of PatchGT is more expressive than the 1-WL
algorithm. Figure 7 in Appendix shows two graph pairs that can be
distinguished by PatchGT but not the 1-WL algorithm. In this PatchGT variant,
we only need to choose the summation operation to aggregate node
representations in the same patch and multiply a scalar to the MHA output. We
put the result in the following Theorem.

######  Theorem 3.

Suppose a PatchGT uses GIN layers, uses sum-pooling as the readout function in
Equation 4, $\mathbf{\mathbf{z}^{0}_{k^{\prime}}}=\sum_{i\in
C_{k^{\prime}}}\mathbf{h}^{L_{1}}_{i}$, and multiplies the MHA output in
Equation 7 with the number $k$ of patches,
$\mathbf{q}^{\prime}_{\ell}=k\cdot\mathrm{MHA}\left(\mathbf{q}_{\ell-1},\mathbf{Z}_{L_{2}},\mathbf{Z}_{L_{2}}\right)$.
Let $\mathbf{g}_{1}$ and $\mathbf{g}_{2}$ be outputs computed from two graphs
$G_{1}$ and $G_{2}$ by a PatchGT model. There exists a PatchGT such that
$\mathbf{g}_{1}\neq\mathbf{g}_{2}$ if $G_{1}$ and $G_{2}$ can be distinguished
by the 1-WL algorithm. Furthermore, there are graph pairs $G_{1}\neq G_{2}$
that cannot be distinguished by the 1-WL algorithm, but
$\mathbf{g}_{1}\neq\mathbf{g}_{2}$ from this PatchGT model.

The first part of the conclusion is true because the patch aggregation, patch-
level GNN, and the MHA pooling can all be bijective mapping. According to
Corollary 6 of [36], the outputs of GIN layers have the same expressive power
as the 1-WL algorithm. Such expressive power is maintained in the model
output. However, when GIN layers on patches use extra structural information
on patches, the model can distinguish graphs that cannot be distinguished by
the 1-WL algorithm. We put the formal proof in Section A.4.

###  4.2 Permutation invariance

Our model depends on the patch structure formed by the clustering algorithm,
which further depends on the spectral decomposition of the normalized
Laplacian. Note that the spectral decomposition is not unique, but we show
that the clustering result is not affected by sign variant and multiplicities
associated with decomposition, and our model is still invariant to node
permutations.

######  Theorem 4.

The network function of PatchGT is invariant to node permutations.

###  4.3 Addressing information bottleneck with patch representations

Figure 2: Pooling methods on a pair of graphs that cannot be distinguished by
the 1-WL algorithm (nodes are colored by the 1-WL algorithm).

Figure 3: It is hard for a GNN to push signal from one graph cluster to the
other, but a patch-level GNN can do so with patch representations.

Alon et al. [2] recently characterize the issue of information bottleneck in
GNNs through empirical methods. Here we consider this issue on a special case
when a graph consists of loosely-connected node clusters. Note that molecule
graphs often have this property. Here we make the first attempt to
characterize the _information bottleneck_ through theoretical analysis. We
further show that our PatchGT can partially address this issue.

For convenient analysis, we consider a regular graph with degree $\tau$.
Suppose the node set $V$ of $G$ forms two clusters $S$ and $T$: ${V}=S\cup
T,\;S\cap T=\emptyset$, and there are only $m$ edges between $S$ and $T$.

We consider the difficulty of passing signal from $S$ to $T$. Let
$f^{\mathrm{GNN}}(\cdot)$ denote the network function of a GNN of $L$ layers
with ReLU activation $\sigma$ as in (2), and input
$\mathbf{X}=(\mathbf{x}_{i}\in\mathbb{R}^{d}:i\in V)\in\mathbb{R}^{|V|\times
d}$, which contains $d$-dimensional feature inputs to nodes in $G$. Let
$f_{i}^{\mathrm{GNN}}(\cdot)$ be the output at node $i$. We can ask this
question: _if we perturb the input to nodes in $S$, how much impact we can
observe at the output at nodes in $T$._ We need to avoid the case that the
impact is amplified by scaling up network parameters. In real applications,
scaling up network parameters also amplifies signals within $T$ itself, and
the signal from $S$ still cannot be well received. Here we consider relative
impact: the ratio between the impact on $T$ from $S$ over that from $T$
itself.

Let $\mathbold{\alpha}\in\mathbb{R}^{|V|\times d}$ be some perturbation on $S$
such that $\alpha_{ij}\leq\epsilon$ if $i\in S$ and $\alpha_{ij}=0$ otherwise.
Here $\epsilon$ is the scale of the perturbation. Similarly let
$\mathbold{\beta}\in\mathbb{R}^{|V|\times d}$ be some perturbation on $T$:
$\beta_{ij}\leq\epsilon$ if $i\in T$ and $\beta_{ij}=0$ otherwise. Then the
impacts on node representations $f_{i}^{\mathrm{GNN}}$, $i\in T$ from
$\mathbold{\alpha}$ and $\mathbold{\beta}$ are respectively

$\displaystyle\delta_{S\rightarrow T}=\max_{\mathbold{\alpha}}\sum_{i\in
T}\|f_{i}^{\mathrm{GNN}}(\mathbf{X}+\mathbold{\alpha})-f_{i}^{\mathrm{GNN}}(\mathbf{X})\|_{1}$
(10) $\displaystyle\delta_{T\rightarrow T}=\max_{\mathbold{\beta}}\sum_{i\in
T}\|f_{i}^{\mathrm{GNN}}(\mathbf{X}+\mathbold{\beta})-f_{i}^{\mathrm{GNN}}(\mathbf{X})\|_{1}$
(11)

where the maximum is also over all possible learnable parameters
$\|\mathbf{W}_{1}\|_{L_{1}\rightarrow
L_{1}},\|\mathbf{W}_{2}\|_{L_{1}\rightarrow L_{1}}\leq 1$ as in (2). Then we
have the following proposition to bound the ratio $\delta_{S\rightarrow
T}/\delta_{T\rightarrow T}$.

######  Proposition 1.

Given a $\tau$-regular graph $G$, a node subset $S$ with its complement $T$
such that there are only $m$ edges between $S$ and $T$, and a $L$-layer GNN,
it holds that

$\frac{\delta_{S\rightarrow T}}{\delta_{T\rightarrow T}}\leq\frac{2mL}{|T|}$
(12)

The proposition indicates that when there is a small graph cut between two
clusters, then it forms an information bottleneck in a GNN – the network needs
to use more layers to pass signal from one group to another. The bound is
still conservative: if the signal is extracted in middle layers of the
network, then passing the signal is even harder. The proposition is
illustrated in Figure 3.

In our PatchGT model, communication can happen at the coarse graph and thus
can partially address this issue. The coarse graph $\tilde{\mathbf{A}}$
consists of two nodes (we still denote them by $S,T$), and there is an edge
between $S$ and $T$. From the output $f^{\mathrm{GNN}}$, we construct the
patch representations
$(\mathbf{z}_{S},\mathbf{z}_{T})=(\frac{1}{|V|}\sum_{i\in
S}f_{i}^{\mathrm{GNN}}(\mathbf{X}),\frac{1}{|V|}\sum_{i\in
T}f_{i}^{\mathrm{GNN}}(\mathbf{X}))\in\mathbb{R}^{2\times d}$. Then we apply a
GNN layer to get node represents on the coarse graph
$(g_{S}^{\mathrm{GNN}}(\mathbf{X}),g_{T}^{\mathrm{GNN}}(\mathbf{X}))\in\mathbb{R}^{2\times
d}$:

$\displaystyle
g_{S}^{\mathrm{GNN}}(\mathbf{X})=\sigma(\mathbf{z}_{S}\mathbf{W}_{1}^{\top}+\mathbf{z}_{T}\mathbf{W}_{2}^{\top}),\quad
g_{T}^{\mathrm{GNN}}(\mathbf{X})=\sigma(\mathbf{z}_{T}\mathbf{W}_{1}^{\top}+\mathbf{z}_{S}\mathbf{W}_{2}^{\top}),$
(13)

where $\mathbf{W}_{1},\mathbf{W}_{2}\in\mathbb{R}^{d\times d}$ are learnable
parameters. We consider the impact of $\mathbold{\alpha}$ on our patch GT, let

$\displaystyle\eta_{S\rightarrow
T}=\max_{\mathbold{\alpha}}\|g^{\mathrm{GNN}}_{T}(\mathbf{X}+\mathbold{\alpha})-g^{\mathrm{GNN}}_{T}(\mathbf{X})\|_{1}$
(14) $\displaystyle\eta_{T\rightarrow
T}=\max_{\mathbold{\beta}}\|g^{\mathrm{GNN}}_{T}(\mathbf{X}+\mathbold{\beta})-g^{\mathrm{GNN}}_{T}(\mathbf{X})\|_{1},$
(15)

Then we have the following proposition on the ratio $\eta_{S\rightarrow
T}/\eta_{T\rightarrow T}$.

######  Theorem 5.

The ratio $\frac{\eta_{S\rightarrow T}}{\eta_{T\rightarrow T}}$ can be
arbitrarily close to $1$ in a PatchGT model, under the assumption of regular
graphs.

This is because $S$ and $T$ are direct neighbors in the coarse graph, then
$\alpha_{S}$ can directly impact $\mathbf{z}_{S}$, which can impact
$g^{\mathrm{GNN}}_{T}$ through messages passed by GNN layers or the attention
mechanism of Transformer layers. The right part of fig. 3 shows that patch
representation can include signals from the other node cluster.

Molecules | Ours | Trainable | Molecules | Ours | Trainable
---|---|---|---|---|---
|  |  |  |  |
Figure 4: Segmentation results from spectral clustering and trainable
clustering.

###  4.4 Comparison for different Segmentation methods

In the previous researches, there exist many hierarchical pooling models [39,
12, 4, 13, 27, 19]. The most obvious difference from the proposed method is
that the pooling/segmentation is trainable. Particularly, the pooling is from
the node respresentations learned by GNNs. In the Theorem 1 and Theorem 2, we
prove such trainable clustering methods will compute the same representations
to the nodes if 1-WL algorithm can not differentiate them. This takes two
serious problems for the graph segmentation: First, the nodes with the same
representations will be assigned to the same cluster even if they are not
connected to each other; Second, too many nodes could be assigned to one
cluster to make sure that the nodes far away from each other are in the same
cluster.

Here we compare the two segmentation results: one is from spectral clustering
and another is from Memory-based graph networks[1] which is a typical
trainable clustering method. In the first case, we find that nodes in the blue
cluster from trainable clustering are not connected. If we adopt such patch
representations by aggregating the disconnected nodes, it will definitely hurt
the performance. This can also be applied to other hierarchical pooling
methods such as Diffpool, Eigenpool, and MinCutpool.

In the second case, the spectral clustering methods segment the graph by
minimum cuts. This is helpful to solve the information bottleneck between
patches. However, the Memory-based graph networks cluster the two benzene
rings together. It will be difficult for the model to detect the existence of
these two benzene rings.

##  5 Empirical Study

In this section, we evaluate the effectiveness of PatchGT through experiments.

Datasets. We benchmark the performances of PatchGT on several commonly studied
graph-level prediction datasets. The first four are from the Open Graph
Benchmark (OGB) datasets [14] (ogbg-molhiv, ogbg-molbace, ogbg-molclintox, and
ogbg-molsider). These tasks are predicting molecular attributes. The
evaluation metric for these four datasets is ROC-AUC ($\%$). The second group
of six datasets are from the TU datasets [25], and they are DD, MUTAG,
PROTEINS, PTC-MR, ENZYMES, and Mutagenicity. Each dataset contains one
classification task for molecules. The evaluation metric is accuracy ($\%$)
over all six datasets. The statistics for the datasets is summarized in
Section A.11.

###  5.1 Quantitative evaluation

Table 1: Results $(\%)$ on OGB datasets

| ogbg-molhiv | ogbg-molbace | ogbg-molclintox | ogbg-molsider
---|---|---|---|---
GCN +VN | 75.99 $\pm$1.19 | 71.44 $\pm$ 4.01 | 88.55$\pm$2.09 | 59.84$\pm$1.54
GIN + VN | 77.07$\pm$1.49 | 76.41$\pm$2.68 | 84.06$\pm$3.84 | 57.75 $\pm$1.14
Deep LRP | 77.19$\pm$1.40 | - | - | -
PNA | 79.05$\pm$1.32 | - | - | -
Nested GIN | 78.34$\pm$1.86 | 74.33$\pm$1.89 | 86.35$\pm$1.27 | 61.2$\pm$1.15
GRAPHSNN +VN | 79.72$\pm$1.83 | - | - | -
Graphormer (pre-trained) |  80.51$\pm$0.53 | - | - | -
PatchGT-GCN | 80.22$\pm$0.84 | 86.44$\pm$1.92 |  92.21 $\pm$1.35 | 65.21 $\pm$ 0.87
PatchGT-GIN | 79.99$\pm$1.21 | 84.08$\pm$2.03 | 86.75 $\pm$1.04 | 64.90 $\pm$0.92
PatchGT-DeeperGCN | 78.13 $\pm$ 1.89 |  88.31$\pm$1.87 | 89.02$\pm$ 1.21 |  65.46$\pm$1.03

Table 2: Results $(\%)$ on TU datasets

| DD | MUTAG | PROTEINS | PTC-MR | ENZYMES | Mutagenicity
---|---|---|---|---|---|---
GCN | 71.6$\pm$2.8 | 73.4$\pm$10.8 | 71.7$\pm$4.7 | 56.4$\pm$7.1 | 50.17 | -
GraphSAGE | 71.6$\pm$3.0 | 74.0$\pm$8.8 | 71.2$\pm$5.2 | 57.0$\pm$5.5 | 54.25 | -
GIN | 70.5$\pm$3.9 | 84.5$\pm$8.9 | 70.6$\pm$4.3 | 51.2$\pm$9.2 | 59.6 | -
GAT | 71.0$\pm$4.4 | 73.9$\pm$10.7 | 72.0$\pm$3.3 | 57.0$\pm$7.3 | 58.45 | -
DiffPool | 79.3$\pm$2.4 | - | 72.7$\pm$3.8 | - | 62.53 | 77.6$\pm$2.7
MinCutPool | 80.8$\pm$2.3 | - | 76.5$\pm$2.6 | - | - | 79.9$\pm$2.1
Nested GCN | 76.3$\pm$3.8 | 82.9$\pm$11.1 | 73.3$\pm$4.0 | 57.3$\pm$7.7 | 31.2$\pm$6.7 | -
Nested GIN | 77.8$\pm$3.9 | 87.9$\pm$8.2 | 73.9$\pm$5.1 | 54.1$\pm$7.7 | 29.0$\pm$8.0 | -
DiffPool-NOLP | 79.98 | - | 76.22 | - | 61.95 | -
SEG-BERT | - | 90.8 $\pm$6.5 | 77.1$\pm$4.2 | - | - | -
U2GNN | 80.2$\pm$1.5 | 89.9$\pm$3.6 | 78.5$\pm$4.07 | - | - | -
EigenGCN | 78.6 | - | 76.6 | - | 64.5 | -
Graph U-Nets | 82.43 | - | 77.68 | - | - | -
PatchGT-GCN |  83.3$\pm$3.1 |  94.7$\pm$3.5 |  80.3$\pm$2.5 |  62.5$\pm$4.1 |  73.3$\pm$3.3 | 78.3$\pm$2.2
PatchGT-GIN | 79.6$\pm$3.3 | 89.4$\pm$3.2 | 79.5$\pm$3.1 | 58.4$\pm$2.9 | 70.0$\pm$3.5 | 80.4$\pm$1.4
PatchGT-DeeperGCN | 76.1$\pm$2.8 | 89.4$\pm$3.7 | 77.5$\pm$3.4 | 60.0$\pm$2.6 | 56.6$\pm$3.1 |  80.6$\pm$1.5

Baselines. In this section, we compare the performance of PatchGT against
several baselines including GCN [17], GIN [36], as well as recent works Nested
Graph Neural Networks [44] and GraphSNN [33]. To compare with learnable
pooling methods, we also include DiffPool [39], MinCutPool [4] Graph
U-Nets[12], and EigenGCN[23] as baselines for TU datasets. We also include the
Graphormer model, but note that Graphormer needs a large-scale pre-training
and cannot be easily applied to a wider range of datasets. We also compare our
model with other transformer-based models such as U2GNN[26] and SEG-BERT[43].

Settings. We search model hyper-parameters such as the eigenvalue threshold,
the learning rate, and the number of graph neural network layers on the
validation set. Each OGB dataset has its own data split of training,
validation, and test sets. We run ten fold cross-validation on each TU
dataset. In each fold, one-tenth of the data is used as the test set, one-
tenth is used as the validation set, and the rest is used as training. For the
detailed search space, please refer to Section A.12.

Results. Table 1 and Table 2 summarize the performance of PatchGT and other
baselines on OGB datasets and TU datasets. We take values from the original
papers and the OGB website; EXCEPT the performance values of Nested GIN on the
last three OGB datasets – we obtain the three values by running Nested GIN. We
also tried to run the contemporary method GRAPHSNN+VN on the other three OGB
datasets, but we did not find the official implementation at the submission of
this work.

From the results, we see that the proposed method gets good performances on
almost all datasets and often outperforms competing methods with a large
margin. On the ogbg-molhiv dataset, the performance of PatchGT with GCN is
only slightly worse than Graphormer, but note that Graphormer needs large-
scale pre-training, which limits its applications.

PatchGT with GCN outperforms three baselines on the other three OGB datasets.
The improvements on these three OGB datasets are significant. PatchGT with GCN
outperforms baselines on four out of six TU datasets. When it does not
outperform all baselines, its performances are only slightly worse than the
best performance. Similarly, two other configurations, PatchGT-GIN and
PatchGT-DeeperGCN, also perform very well on these two datasets.

Figure 5: Analysis of the key design for the proposed PatchGT. All results are
based on PatchGT GCN. In the left figure, we show how changing the threshold
for eigenvalues affects performance on the ogbg-molclintox and PROTEINS
datasets; The middle figure shows the model performances with the removal of
patch-GNN or Transformer (replaced by mean pool) on DD and ogbg-molhiv
datasets; The right figure shows the effect of the different readout functions
for patch representations.

###  5.2 Ablation study

We perform ablation studies to check how different configurations of our model
affect its performance. The results are shown in Figure 5.

Figure 6: Attention visualization of PatchGT on ogbg-molhiv molecules. The
second and fourth figures show the attention weights of query tokens on the
node patches for the corresponding molecules, which are in the first and third
figures. The molecule in the first figure does not inhibit HIV virus, yet the
molecule in the third figure does.

Effect of eigenvalue threshold. The eigenvalue threshold $\gamma$ influences
how many patches for a graph after the segmentation. Generally speaking,
larger $\gamma$ introduces more patches and patches with smaller sizes. When
$\gamma$ is large enough, the number of patches $k$ equals the number of nodes
$|V|$ in the graph, and the Transformer actually works at the node level. When
the $\gamma$ is $0$, then the whole graph is treated as one patch, and the
model is reduced to a GNN with pooling. The left figure shows that there is a
sweet point (depending on the dataset) for the threshold, which means that
using patches is a better choice than not using patches.

Effect of GNN layer on the coarse graph and Transformer layers. This ablation
study removes either patch-level GNN layers or Transformer layers to check
which part of the architecture is important for the model performance. From
the middle plot in Figure 5, we see that both types of layers are useful, and
Transformer layers are more useful. This is another piece of evidence that
PatchGT can combine the strengths of different models.

Comparison of readout functions. We compare the performance of PatchGT model
using different readout functions when aggregating node representations at
each patch in Equation 4. In the right figure, we observe the remarkable
influence of the readout function on the performance. Empirical studies
indicate max-pooling is the optimal choice under most circumstances.

###  5.3 Understanding the attention

Besides improving learning performances, we are also interested in
understanding how the attention mechanism helps the model identify the graph
property. We train the PatchGT model on the ogbg-molhiv dataset and visualize
the attention weights between query tokens and each patch. Interestingly, the
attention only concentrates on some chemical motifs such as
$\text{Cl\vspace{0.1pt} O}_{3}$ and $\text{CON}_{2}$ but ignores other very
common motifs such as benzene rings. It can be noticed that for the molecule
in the first figure, the two benzene rings are connected to each other by
-C-C-. However, the model does not pay any attention to this part. The two
rings in the molecule of the second molecule are connected by -S-S-;
differently, the model pays attention to this part this time. It indicates
that Transformer can identify which motifs are informative and which motifs
are common. Such property offers better model interpretability compared to the
traditional global pooling. It not only makes accurate predictions but also
provides some insight into why decisions are made. In the two examples shown
above, we can start from motifs $\text{SO}_{3}$ and -S-S- to look for
structures meaningful for the classification problem.

##  6 Conclusion and Limitations

In this work, we show that graph learning models benefit from modeling patches
on graphs, particularly when it is combined with Transformer layers. We
propose PatchGT, a new learning model that uses non-trainable clustering to
get graph patches and learn graph representations based on patch
representations. It combines the strengths of GNN layers and Transformer
layers and we theoretically prove that it helps mitigate the bottleneck of
graphs and limitations of trainable clustering. It shows superior performances
on a list of graph learning tasks. Based on graph patches, Transformer layers
also provides a good level of interpretability of model predictions.

However, the work tested our model mostly on chemical datasets. It is unclear
whether the model still performs well when input graphs do not have clear
cluster structures.

## Acknowledgements

Li-Ping Liu was supported by NSF1908617. Xu Han was supported by Tufts RA
support. The research of J.H. is supported by NSF grant DMS-2054835. J.W and
H.G would like to acknowledge the funds from National Science Foundation under
Award Nos. CMMI-1934300 and OAC-2047127.

## References

  * [1] Amir Hosein Khas Ahmadi  “Memory-based graph networks”, 2020
  * [2] Uri Alon and Eran Yahav  “On the bottleneck of graph neural networks and its practical implications”  In _arXiv preprint arXiv:2006.05205_ , 2020
  * [3] Jinheon Baek, Minki Kang and Sung Ju Hwang  “Accurate learning of graph representations with graph multiset pooling”  In _arXiv preprint arXiv:2102.11533_ , 2021
  * [4] Filippo Maria Bianchi, Daniele Grattarola and Cesare Alippi  “Spectral Clustering with Graph Neural Networks for Graph Pooling”  In _Proceedings of the 37th International Conference on Machine Learning_ 119, Proceedings of Machine Learning Research  PMLR, 2020, pp. 874–883
  * [5] Cristian Bodnar et al.  “Weisfeiler and Lehman go cellular: CW networks”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 2625–2640
  * [6] Dexiong Chen, Leslie O’Bray and Karsten Borgwardt  “Structure-Aware Transformer for Graph Representation Learning”  In _arXiv preprint arXiv:2202.03036_ , 2022
  * [7] D.J. Cook and L.B. Holder  “Mining Graph Data”  Wiley, 2006  URL: https://books.google.com/books?id=bHGy0%5C_H0g8QC
  * [8] James Demmel, Ioana Dumitriu and Olga Holtz  “Fast linear algebra is stable”  In _Numerische Mathematik_ 108.1 Springer, 2007, pp. 59–91
  * [9] Alexey Dosovitskiy et al.  “An image is worth 16x16 words: Transformers for image recognition at scale”  In _arXiv preprint arXiv:2010.11929_ , 2020
  * [10] Charbel Farhat  “Dimensional Reduction of Highly Nonlinear Multiscale Models Using Most Appropriate Local Reduced-Order Bases”, 2020
  * [11] Kimon Fountoulakis et al.  “Graph Attention Retrospective”  In _arXiv preprint arXiv:2202.13060_ , 2022
  * [12] Hongyang Gao and Shuiwang Ji  “Graph u-nets”  In _international conference on machine learning_ , 2019, pp. 2083–2092  PMLR
  * [13] Daniele Grattarola, Daniele Zambon, Filippo Maria Bianchi and Cesare Alippi  “Understanding Pooling in Graph Neural Networks”  In _arXiv preprint arXiv:2110.05292_ , 2021
  * [14] Weihua Hu et al.  “Open graph benchmark: Datasets for machine learning on graphs”  In _Advances in neural information processing systems_ 33, 2020, pp. 22118–22133
  * [15] Dejun Jiang et al.  “Could graph neural networks learn better molecular representation for drug discovery? A comparison study of descriptor-based and graph-based models”  In _Journal of cheminformatics_ 13.1 BioMed Central, 2021, pp. 1–23
  * [16] Katikapalli Subramanyam Kalyan, Ajit Rajasekharan and Sivanesan Sangeetha  “Ammus: A survey of transformer-based pretrained models in natural language processing”  In _arXiv preprint arXiv:2108.05542_ , 2021
  * [17] Thomas N Kipf and Max Welling  “Semi-supervised classification with graph convolutional networks”  In _arXiv preprint arXiv:1609.02907_ , 2016
  * [18] Devin Kreuzer et al.  “Rethinking graph transformers with spectral attention”  In _Advances in Neural Information Processing Systems_ 34, 2021
  * [19] Junhyun Lee, Inyeop Lee and Jaewoo Kang  “Self-attention graph pooling”  In _International conference on machine learning_ , 2019, pp. 3734–3743  PMLR
  * [20] Yujia Li, Daniel Tarlow, Marc Brockschmidt and Richard Zemel  “Gated graph sequence neural networks”  In _arXiv preprint arXiv:1511.05493_ , 2015
  * [21] Derek Lim et al.  “Sign and Basis Invariant Networks for Spectral Graph Representation Learning”  In _arXiv preprint arXiv:2202.13013_ , 2022
  * [22] Ze Liu et al.  “Swin transformer: Hierarchical vision transformer using shifted windows”  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 10012–10022
  * [23] Yao Ma, Suhang Wang, Charu C Aggarwal and Jiliang Tang  “Graph convolutional networks with eigenpooling”  In _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, 2019, pp. 723–731
  * [24] Grégoire Mialon, Dexiong Chen, Margot Selosse and Julien Mairal  “Graphit: Encoding graph structure in transformers”  In _arXiv preprint arXiv:2106.05667_ , 2021
  * [25] Christopher Morris et al.  “TUDataset: A collection of benchmark datasets for learning with graphs”  In _ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020)_ , 2020  arXiv: www.graphlearning.io
  * [26] Dai Quoc Nguyen, Tu Dinh Nguyen and Dinh Phung  “Universal graph transformer self-attention networks”  In _arXiv preprint arXiv:1909.11855_ , 2019
  * [27] Emmanuel Noutahi et al.  “Towards interpretable sparse graph representation learning with laplacian pooling”  In _arXiv preprint arXiv:1905.11577_ , 2019
  * [28] Hoang Nt and Takanori Maehara  “Revisiting graph neural networks: All we have is low-pass filters”  In _arXiv preprint arXiv:1905.09550_ , 2019
  * [29] Yousef Saad  “Iterative methods for sparse linear systems”  SIAM, 2003
  * [30] Jianbo Shi and Jitendra Malik  “Normalized cuts and image segmentation”  In _IEEE Transactions on pattern analysis and machine intelligence_ 22.8 Ieee, 2000, pp. 888–905
  * [31] Asher Trockman and J Zico Kolter  “Patches Are All You Need?”  In _arXiv preprint arXiv:2201.09792_ , 2022
  * [32] Ashish Vaswani et al.  “Attention is all you need”  In _Advances in neural information processing systems_ 30, 2017
  * [33] Asiri Wijesinghe and Qing Wang  “A New Perspective on" How Graph Neural Networks Go Beyond Weisfeiler-Lehman?"”  In _International Conference on Learning Representations_ , 2021
  * [34] Zhanghao Wu et al.  “Representing long-range context for graph neural networks with global attention”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 13266–13279
  * [35] Zonghan Wu et al.  “A comprehensive survey on graph neural networks”  In _IEEE transactions on neural networks and learning systems_ 32.1 IEEE, 2020, pp. 4–24
  * [36] Keyulu Xu, Weihua Hu, Jure Leskovec and Stefanie Jegelka  “How powerful are graph neural networks?”  In _arXiv preprint arXiv:1810.00826_ , 2018
  * [37] Pinar Yanardag and SVN Vishwanathan  “Deep graph kernels”  In _Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining_ , 2015, pp. 1365–1374
  * [38] Chengxuan Ying et al.  “Do Transformers Really Perform Badly for Graph Representation?”  In _Advances in Neural Information Processing Systems_ 34, 2021
  * [39] Zhitao Ying et al.  “Hierarchical graph representation learning with differentiable pooling”  In _Advances in neural information processing systems_ 31, 2018
  * [40] Seongjun Yun et al.  “Graph transformer networks”  In _Advances in neural information processing systems_ 32, 2019
  * [41] Habil Zare, Parisa Shooshtari, Arvind Gupta and Ryan R Brinkman  “Data reduction for spectral clustering to analyze high throughput flow cytometry data”  In _BMC bioinformatics_ 11.1 Springer, 2010, pp. 1–16
  * [42] Aston Zhang, Zachary C. Lipton, Mu Li and Alexander J. Smola  “Dive into Deep Learning”  In _arXiv preprint arXiv:2106.11342_ , 2021
  * [43] Jiawei Zhang  “Segmented graph-bert for graph instance modeling”  In _arXiv preprint arXiv:2002.03283_ , 2020
  * [44] Muhan Zhang and Pan Li  “Nested Graph Neural Networks”  In _Advances in Neural Information Processing Systems_ 34, 2021

##  Appendix A Appendix

###  A.1 Proof of Theorem 1

The proof that DiffPooling cannot distinguish graphs that are colored in the
same way by the 1-WL algorithm.

###### Proof.

The function form of a pooling layer in DiffPooling is

$\displaystyle\mathbf{H}^{\prime}=\mathbf{S}^{\top}\mathbf{A}\mathbf{S}\mathbf{S}^{\top}\mathbf{H},\quad\mathbf{S}=\mathrm{gnn}_{c}(\mathbf{A},\mathbf{X}),\quad\mathbf{H}=\mathrm{gnn}_{r}(\mathbf{A},\mathbf{X})$
(16)

Here $\mathrm{gnn}_{c}(\cdot,\cdot)$ learns a cluster assignment $\mathbf{S}$
of all nodes in the graph, and $\mathrm{gnn}_{r}(\cdot,\cdot)$ learns node
representations.

Note that $\mathrm{gnn}_{r}$ has at most the ability of 1-WL algorithm [36].
Two nodes must get the same representation when they have the same color in
the 1-WL coloring result. We use an indicator matrix $\mathbf{C}$ to represent
the 1-WL coloring of the graph, that is, the node $i$ is colored as $j$ if
$C_{i,j}=1$, then we can write

$\displaystyle\mathbf{S}=\mathbf{C}\mathbf{B}$ (17)

Here the $j$-th row of $\mathbf{B}$ denote the vector representation learned
for color $j$.

If two graphs represented by $\mathbf{A}$ and $\mathbf{\Lambda}$ cannot be
distinguished by the 1-WL algorithm, then they get the same coloring matrix
$\mathbf{C}$ (subject to some node permutation that does not affect our
analysis here). Now we show that:

$\displaystyle\mathbf{C}^{\top}\mathbf{A}\mathbf{C}=\mathbf{C}^{\top}\mathbf{\Lambda}\mathbf{C}$
(18)

Let’s compare the two matrices on both sides of the equation at an arbitrary
entry $(k,t)$. Let $\alpha_{k}$ and $\alpha_{t}$ represent nodes colored in
$k$ and $t$, then the entry at $(k,t)$ is
$\sum_{i\in\alpha_{k}}\sum_{j\in\alpha_{t}}A_{i,j}$, which is the count of
edges that have one incident node colored in $k$ and the other incident node
colored in $t$. Since the coloring is obtained by 1-WL algorithm, each node
$i\in\alpha_{k}$ has exactly the same number of neighbors colored as $t$. The
number of nodes in color $k$ and the number of neighbors in color $t$ are
exactly the same for $\mathbf{\Lambda}$ because $\mathbf{\Lambda}$ receives
the same coloring as $\mathbf{A}$. Therefore,
$\sum_{i\in\alpha_{k}}\sum_{j\in\alpha_{t}}A_{i,j}=\sum_{i\in\alpha_{k}}\sum_{j\in\alpha_{t}}\Lambda_{i,j}$,
and (18) holds.

At the same time, if two graphs cannot be distinguished by 1-WL, they have the
same node representations $\mathbf{H}$, then they have the same
$\mathbf{H}^{\prime}$. ∎

###  A.2 Proof of Theorem 2

We first prove a lemma.

######  Lemma 1.

Suppose two graphs represented by $\mathbf{A}$ and $\mathbf{\Lambda}$ obtain
the same coloring from the 1-WL algorithm, then

  * i)

the resultant two graphs from removal of nodes in the same color still get the
same coloring by the 1-WL algorithm; and

  * ii)

the two multigraphs represented by $\mathbf{A}^{\ell}$ and
$\mathbf{\Lambda}^{\ell}$ still get the same coloring by the 1-WL algorithm.

Here $\mathbf{A}^{\ell}$ and $\mathbf{\Lambda}^{\ell}$ are the $\ell$-th power
of the two adjacency matrices, and they represent multigraphs that may have
self-loops and parallel edges. The 1-WL algorithm is still valid over graphs
with self-loops and multi-edges. A 1-WL style GNN defined in Section 3.1 or
[12] is still bounded by the 1-WL algorithm on such multigraphs.

###### Proof.

i) We first consider updating of 1-WL coloring when nodes in a color is
removed. Suppose we have stable coloring of graphs represented by
$\mathbf{A}$. Let $\alpha_{t}$ and $\alpha_{r}$ denote two groups of nodes in
color $t$ and $r$ respectively. We also assume each node in $r$ has $t$ in its
color set – if there are not such cases, then we can simply remove nodes in a
color and obtain a stable 1-WL coloring.

Suppose we remove nodes in color $t$ from both graphs. Note that all nodes
$\alpha_{r}$ have the same number of neighbors in color $t$. We update the
color set of each $i\in\alpha_{r}$ by removing color $t$ from it. Then all
nodes in $\alpha_{r}$ still get the same color. Therefore, removing the color
$t$ from nodes in all relevant color groups gives at least a stable coloring,
which, however, might not be the coarsest.

Then we merge some colors when nodes share the same color set. If a node in
color $r$ has the same color set as a node in color $r^{\prime}$, then we
assign the same color to both nodes in colors $r$ and $r^{\prime}$. We run
merging steps until no nodes in different colors share the same color set,
then the coloring is a stable coloring of the graph, and the resultant
coloring of the graph can be viewed as the 1-WL coloring of the graph.

In the procedure above, the step of removing a color, and the steps of merging
colors directly operate on nodes’ color sets. Since nodes in $\mathbf{A}$ and
nodes in $\mathbf{\Lambda}$ have the same color sets, therefore, they will
have the same color sets after color updates.

The update procedure above purely runs on color relations between different
colors. Since $\mathbf{A}$ and $\mathbf{\Lambda}$ have exactly the same color
relations because they receive the same 1-WL coloring. Therefore, the update
procedure above still gives the same stable coloring to $\mathbf{A}$ and
$\mathbf{\Lambda}$.

ii) For the second part of the lemma, we first check the coloring of
$\mathbf{A}^{\ell}$. We show that the coloring of $\mathbf{A}$ is a stable
coloring of $\mathbf{A}^{\ell}$. Suppose each node $i$ has a color set
$C_{i}$. In the graph $\mathbf{A}^{\ell}$, $i$’s $\ell$-th neighbors become
direct neighbors of $i$. The color set of $i$ becomes

$\displaystyle C_{i}\cup\left(\cup_{j_{1}\in
N(i)}C_{j}\right)\cup\ldots\cup\left(\cup_{j_{1}\in
N(i)}\ldots\cup_{j_{\ell}\in N(j_{\ell-1})}C_{j_{\ell}}\right)$ (19)

We know that if two nodes $i$ and $i^{\prime}$ have the same color if and only
if their color sets are the same. By using the relation recursively, $i$ and
$i^{\prime}$ have the same color set in $\mathbf{A}^{\ell}$. Therefore, the
stable coloring of $\mathbf{A}$ is also a stable coloring of
$\mathbf{A}^{\ell}$. If necessary, we can also run the merging procedure above
and eventually get 1-WL coloring of $\mathbf{A}^{\ell}$. With the same
argument as above, the operations only run on color sets, therefore,
$\mathbf{A}^{\ell}$ and $\mathbf{\Lambda}^{\ell}$ have the same coloring. ∎

Now we are ready to prove the main theorem that the Graph U-Net variant cannot
distinguish graphs colored in the same way by the 1-WL algorithm.

###### Proof.

In the calculation of Graph U-Net-th, the indicator $\mathbf{b}$ for removing
nodes is obtained by thresholding $\mathbf{v}$, which is computed by a 1-WL
GNN. Therefore, nodes in the same color are always kept or removed all
together in $\mathbf{b}$.

Suppose the inputs to a Graph U-Net layer are $(\mathbf{A},\mathbf{X})$ and
$(\mathbf{\Lambda},\mathbf{X})$ respectively, and $\mathbf{A}$ and
$\mathbf{\Lambda}$ cannot be distinguished by the 1-WL algorithm. The inputs
to next layer are
$(\mathbf{A}^{\ell}[\mathbf{b},\mathbf{b}],\mathbf{X}[\mathbf{b}])$ and
$(\mathbf{\Lambda}^{\ell}[\mathbf{b},\mathbf{b}],\mathbf{X}[\mathbf{b}])$
respectively. By the lemma above, the 1-WL algorithm cannot distinguish
$\mathbf{A}^{\ell}$ and $\mathbf{\Lambda}^{\ell}$, and it cannot be
distinguish $\mathbf{A}^{\ell}[\mathbf{b},\mathbf{b}]$ and
$\mathbf{\Lambda}^{\ell}[\mathbf{b},\mathbf{b}]$ either. Therefore, it still
cannot distinguish the inputs
$(\mathbf{A}^{\ell}[\mathbf{b},\mathbf{b}],\mathbf{X}[\mathbf{b}])$ to the
next layer.

By using the argument above recursively, the network cannot distinguish the
graph at the final outputs if network inputs $(\mathbf{A},\mathbf{X})$ and
$(\mathbf{\Lambda},\mathbf{X})$ cannot be distinguished by the 1-WL algorithm.
∎

######  Remark 1.

For graphs with noise or low homophily ratios, the aforementioned issue may
not be severe and long-distance aggregation is helpful.

###  A.3 Analysis for expressiveness of Graph U-Nets

In this section we use an example in Fig. 7 to understand how to maintain a
graph’s global structure with pooling operations. In a pooling step, DiffPool
and MinCutPool will assign nodes in the same color to the same cluster and
merge them as one node. Clearly it does not maintain the global structure of
the graph and cannot distinguish the two graphs.

Graph U-Net always ranks nodes in one color above nodes of the other color. It
is not always permutation invariant: for example, it may get different
structures when it breaks tie to take two green nodes. In many cases, it
cannot distinguish the two graphs: when it takes three nodes, either three
green nodes or two blue and one green nodes, it cannot distinguish the two
graphs. The Graph U-Net variant considered above always remove blue or green
nodes, thus it cannot distinguish the two graphs. One important observation is
Graph U-Net cannot preserve the global graph structure in its pooling steps.
For example, when it removes three nodes, the structure left is vastly
different from the original graph.

###  A.4 A proof showing that PatchGT is more expressive than the 1-WL
algorithm

Figure 7: Two graphs that cannot be distinguished by the 1-WL algorithm. The
colors illustrate the 1-WL coloring of graph nodes. In comparison, PatchGT can
differentiate them through the patch-level graph.

###### Proof.

From the proof of GIN, we know that the two multi-sets
$\\{\mathbf{h}^{L_{1}}_{i}:i\in G_{1}\\}$ and $\\{\mathbf{h}^{L_{1}}_{i}:i\in
G_{2}\\}$ are already different if the two graphs can be distinguished by the
$L_{1}$-round 1-WL algorithm.

Then we show that the rest of a learned network from Equation 4 to Equation 9
is a bijective operation. We first consider the patch aggregation by the sum-
pooling is bijective. According to Corollary 6 of [36], and assuming the GIN
layers are properly trained, then there is an inverse $\mathrm{inv}(\cdot)$ of
sum-pooling such that $\\{\mathbf{h}^{L_{1}}_{i}:i\in
C_{k^{\prime}}\\}=\mathrm{inv}(\mathbf{z_{0}})$. Then the inverse of patch
aggregation is:

$\displaystyle\\{\mathbf{h}^{L_{1}}_{i}:i\in
G_{1}\\}=\cup_{k^{\prime}=1}^{k}\mathrm{inv}(\mathbf{z}^{0}_{k^{\prime}})$
(20)

If the $L_{2}$ GNN layers on patches are also properly trained, then the
mapping from $\mathbf{Z}_{0}$ to $\mathbf{Z}_{L_{2}}$ is also bijective. At
the same time, we assume vectors in $\mathbf{Z}_{L_{2}}$ are properly
transformed, which will be useful in the following MHA operation.

Finally, we consider MHA layers. We first analyze the case with only one layer
with one attention head. Note that
$\mathbf{q}_{1}=k\cdot\mathrm{softmax}\left(\mathbf{q}_{0}^{\top}\mathbf{Z}_{L_{2}}/\sqrt{d}\right)^{\top}\mathbf{Z}_{L_{2}}$
with $d$ being the dimension of row vectors in $\mathbf{Z}_{1}$. Suppose
PatchGT learns the query $q_{0}$ to be a zero vector, and the linear
transformation in Equation 9 is the identity operation, then
$\mathbf{g}_{1}=\mathbf{q}_{1}=\mathbf{1}^{\top}\mathbf{Z}_{L_{2}}$, which is
the summation of patch vectors $\mathbf{Z}_{L_{2}}$. Combining the last GIN
layer, this summation is a bijective operation according to Corollary 6 of
[36]. If there are multiple MHA layers, then we only need the MLP in Equation
8 to zero out the input, and the layer is equivalent to no operation. If there
are multiple attention heads, the network can always take the first attention
head. Therefore, a general case of MHA layers can also be a summation of input
vectors.

Putting these steps together, there is an inverse mapping $\mathbf{g}_{1}$ to
$\\{\mathbf{h}^{L_{1}}_{i}:i\in G_{1}\\}$ and mapping $\mathbf{g}_{2}$ to
$\\{\mathbf{h}^{L_{1}}_{i}:i\in G_{2}\\}$. Then $\mathbf{g}_{1}$ and
$\mathbf{g}_{2}$ must be different.

We further show that there are cases that cannot be distinguished by the 1-WL
algorithm but can be distinguished by PatchGT. Consider two examples in Figure
7. The two original graphs $G_{1}$ and $G_{2}$, or $G_{3}$ and $G_{4}$, are
non-isomorphic. However, both the 1-WL algorithm cannot differentiate them. In
comparison, by segmenting these graphs into patches, PatchGT can discriminate
$G_{1}$ from $G2$. After segmentation, the two patches from $G_{1}$ and the
pacthes $G_{2}$ can be distinguished by the 1-WL algorithm and also PatchGT.
Note that node degrees of $G_{1}$ patches are already different from node
degrees of $G_{2}$ patches. It is the same for $G_{3}$ and $G_{4}$. These two
examples indicate that the expressiveness of PatchGT is beyond 1-WL algorithm.

∎

###  A.5 Proof of Theorem 4

We prove the theorem 4 through three lemmas below.

######  Lemma 2.

The patches split via $k$-means are invariant to column vectors in
$\mathbf{U}$ from the spans of eigenvectors associated with the multiplicities
of eigenvalues.

$\mathrm{kmeans}(\mathbf{V})=\mathrm{kmeans}(\mathbf{V}\mathbf{Q})$ (21)

where $\mathbf{Q}$ is a standard block-diagonal rotation matrix.

###### Proof.

If we use $N_{u}$ eigenvectors for the graph patch splitting, corresponding to
the first $N_{u}$ smallest eigenvalues, we can write them as
$(\lambda_{1},\mathbf{u}_{1}),...,(\lambda_{N_{u}},\mathbf{u}_{N_{u}})$. If we
have multiplicities in these eigenvalues, we can rotate the eigenvectors by a
block-diagonal rotation matrix $\mathbf{Q}\in\mathbb{R}^{N_{u}\times N_{u}}$
to obtain another set of eigenvectors,

$\mathbf{U}^{\prime}=[\mathbf{u}_{1}^{\prime},...,\mathbf{u}_{k}^{\prime}]=[\mathbf{u}_{1},...,\mathbf{u}_{k}]\mathbf{Q}=\mathbf{U}\mathbf{Q}$
(22)

where $\mathbf{u}_{i},\mathbf{u}_{i}^{\prime}\in\mathbb{R}^{|V|\times 1}$. If
we perform $k$-means on the row vectors of
$[(\mathbf{u}_{1})_{i},...,(\mathbf{u}_{k})_{N_{u}}]$, we can write the nodes’
coordinates as

$[\mathbf{x}_{1};...;\mathbf{x}_{|V|}]=[\mathbf{u}_{1},...,\mathbf{u}_{N_{u}}].$
(23)

Similarly, we can write down the new coordinates after rotation as

$[\mathbf{x}_{1}^{\prime};...;\mathbf{x}_{|V|}]=[\mathbf{u}_{1}^{\prime},...,\mathbf{u}_{N_{u}}^{\prime}].$
(24)

From the above three equations, it holds that

$[\mathbf{x}_{1}^{\prime};...;\mathbf{x}_{|\mathcal{V}|^{\prime}}]=[\mathbf{x}_{1};...;\mathbf{x}_{|\mathcal{V}|}]\mathbf{Q}.$
(25)

So for $i,j\in\\{1,...,|V|\\}$, we have

$\mathbf{x}_{i}^{\prime}=\mathbf{x}_{i}\mathbf{Q}\quad\mathbf{x}_{j}^{\prime}=\mathbf{x}_{j}\mathbf{Q}.$
(26)

The relative distance of new coordinates can be calculated as

$(\mathbf{x}_{i}^{\prime}-\mathbf{x}_{j}^{\prime})(\mathbf{x}_{i}^{\prime}-\mathbf{x}_{j}^{\prime})^{\top}=(\mathbf{x}_{i}\mathbf{Q}-\mathbf{x}_{j}\mathbf{Q})(\mathbf{x}_{i}\mathbf{Q}-\mathbf{x}_{j}\mathbf{Q})^{\top}=(\mathbf{x}_{i}-\mathbf{x}_{j})\mathbf{Q}\mathbf{Q}^{\top}(\mathbf{x}_{i}-\mathbf{x}_{j})^{\top}.$
(27)

From the property of the rotational matrix, we have

$\mathbf{I}=\mathbf{Q}\mathbf{Q}^{\top}.$ (28)

So it holds that

$(\mathbf{x}_{i}^{\prime}-\mathbf{x}_{j}^{\prime})(\mathbf{x}_{i}^{\prime}-\mathbf{x}_{j}^{\prime})^{\top}=(\mathbf{x}_{i}-\mathbf{x}_{j})(\mathbf{x}_{i}-\mathbf{x}_{j})^{\top}.$
(29)

So for any two node pair, the relative distance is preserved, thus it will not
affect the $k$-means results. ∎

######  Lemma 3.

The patches split via $k$-means are invariant to column vectors in
$\mathbf{U}$ with different signs.

###### Proof.

The sign invariance is a special case of rotation invariance by taking
$\mathbf{Q}$ as a diagonal matrix with entry $(\mathbf{Q})_{ii}\in\\{-1,1\\}$
∎

######  Lemma 4.

The patches split via $k$-means are invariant to the permutations of nodes

$\mathrm{kmeans}(\mathbf{U})=\mathrm{kmeans}(\mathbf{P}\mathbf{U})$ (30)

where $\mathbf{P}$ is a permutation matrix.

###### Proof.

We denote $\mathbf{I}_{|V|}=[1,...,1]^{\top}\in\mathbb{R}^{|{V}|\times 1}$ For
a permutation matrix $\mathbf{P}$ of $\mathbf{A}$, we have the corresponding
permutation matrix $P$ such that

$\mathbf{A}^{\prime}=\mathbf{P}^{\top}\mathbf{A}\mathbf{P}$ (31)

where $\mathbf{A}$ and $\mathbf{A}^{\prime}$ are adjacency matrices of $G$ and
$G^{\prime}$ respectively. And the for the degree matrix of $G$ and
$G^{\prime}$

$\mathbf{D}=\mathrm{diag}(\mathbf{A}^{\prime}\mathbf{I}_{|V|}),\;\mathbf{D}^{\prime}=\mathrm{diag}(\mathbf{A}^{\prime}\mathbf{I}_{|V|})$
(32)

Substitute equation31 into equation 32

$\mathbf{D}^{\prime}=\mathrm{diag}(\mathbf{P}^{\top}\mathbf{A}\mathbf{P}\mathbf{I}_{|V|})=\mathrm{diag}(\mathbf{P}^{\top}\mathbf{A}\mathbf{P}(\mathbf{P}^{\top}\mathbf{I}_{|V|}\mathbf{P}))$
(33)

From the symmetry of the permutation matrix, it holds that

$\mathbf{P}^{-1}=\mathbf{P}^{\top}$ (34)

Combine the above three equations, we can get

$\mathbf{D}^{\prime}=\mathbf{P}^{\top}\mathrm{diag}(\mathbf{A}\mathbf{I}_{|V|})\mathbf{P}=\mathbf{P}^{\top}\mathbf{D}\mathbf{P}$
(35)

So the permuted Laplacian matrix is

$\begin{split}\mathbf{L}^{\prime}&=\mathbf{I}-\mathbf{D}^{\prime-0.5}\mathbf{A}^{\prime}\mathbf{D}^{\prime-0.5}=\mathbf{P}^{\top}\mathbf{I}\mathbf{P}-\mathbf{P}^{\top}\mathbf{D}^{-0.5}\mathbf{P}\mathbf{P}^{\top}\mathbf{A}\mathbf{P}\mathbf{P}^{\top}\mathbf{D}^{-0.5}\mathbf{P}\\\
&=\mathbf{P}^{\top}(\mathbf{I}-\mathbf{D}^{-0.5}\mathbf{A}\mathbf{D}^{-0.5})\mathbf{P}=\mathbf{P}^{\top}\mathbf{L}\mathbf{P}\end{split}$
(36)

Substitute into the Laplacian eigen decomposition, we have the equation

$\mathbf{L}^{\prime}-\lambda\mathbf{I}=\mathbf{P}^{\top}\mathbf{L}\mathbf{P}^{\top}-\mathbf{P}^{\top}\lambda\mathbf{I}\mathbf{P}=\mathbf{P}^{\top}(\mathbf{L}-\lambda\mathbf{I})\mathbf{P}$
(37)

and its algebraic form

$\mathrm{det}(\mathbf{L}^{\prime}-\lambda\mathbf{I})=\mathrm{det}(\mathbf{P}^{\top})\mathrm{det}(\mathbf{L}-\lambda\mathbf{I})\mathrm{det}(\mathbf{P})=\mathrm{det}(\mathbf{L}-\lambda\mathbf{I}),$
(38)

so the eigenvalues are remaining invariant.

Next we look at the eigenvector. For a eigenvector of
$bL^{\prime},(\lambda,\mathbf{u}^{\prime})$, we have

$\mathbf{L}^{\prime}\mathbf{u}^{\prime}=\lambda\mathbf{u}^{\prime}$ (39)

Combine with equation 36, we can get

$\mathbf{P}^{\top}\mathbf{L}\mathbf{P}\mathbf{u}^{\prime}=\lambda\mathbf{u}^{\prime}\iff\mathbf{L}(\mathbf{P}\mathbf{u}^{\prime})=\lambda(\mathbf{P}\mathbf{u}^{\prime})$
(40)

So we have the relation of two corresponding eigenvectors as

$\mathbf{u}=\mathbf{P}\mathbf{u}^{\prime}\iff\mathbf{u}^{\prime}=\mathbf{P}^{\top}\mathbf{u}$
(41)

So we have the relation for the node coordinate

$[\mathbf{x}_{1}^{\prime};...;\mathbf{x}_{|V|}^{\prime}]=\mathbf{P}^{T}[\mathbf{x}_{1};...;\mathbf{x}_{|V|}].$
(42)

Thus there is a bijective mapping $\mathcal{B}:n\rightarrow m$ such that
$(\mathbf{P})_{n\mathcal{B}(n)}=1$ and
$\mathbf{x}_{n}=\mathbf{x}^{\prime}_{\mathcal{B}(n)}$. Then for any node pair
$(i,j)$, we can find $(i^{\prime},j^{\prime})=(\mathcal{B}(i),\mathcal{B}(j))$
such that

$\mathbf{x}_{i}=\mathbf{x}^{\prime}_{i^{\prime}},\quad\mathbf{x}_{j}=\mathbf{x}^{\prime}_{j^{\prime}},$
(43)

then it clearly holds that

$(\mathbf{x}_{i}-\mathbf{x}_{j})(\mathbf{x}_{i}-\mathbf{x}_{j})^{\top}=(\mathbf{x}^{\prime}_{i^{\prime}}-\mathbf{x}^{\prime}_{j^{\prime}})(\mathbf{x}^{\prime}_{i^{\prime}}-\mathbf{x}^{\prime}_{j^{\prime}})^{\top}.$
(44)

So for any two node pair, the relative distance is preserved, thus it will not
affect the $k$-means results. ∎

###  A.6 Multi-head attention

Transformer [32] has been proved successful in the NLP and CV fields. The
design of multi-head attention (MHA) layer is based on attention mechanism
with Query-Key-Value (QKV). Given the packed matrix representations of queries
$\mathbf{Q}$, keys $\mathbf{K}$, and values $\mathbf{V}$, the scaled dot-
product attention used by Transformer is given by:

$\displaystyle\mathrm{ATTENTION}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{D_{k}}}\right)\mathbf{V},$
(45)

where $D_{k}$ represents the dimensions of queries and keys.

The multi-head attention applies $H$ heads of attention, allowing a model to
attend to different types of information.

$\displaystyle\mathrm{MHA}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{CONCAT}\left(\mathrm{head}_{1},\ldots,\mathrm{head}_{H}\right)\mathbf{W}$
$\displaystyle\mathrm{where}\quad\mathrm{head}_{i}=\mathrm{ATTENTION}\left(\mathbf{Q}\mathbf{W}_{i}^{Q},\mathbf{K}\mathbf{W}_{i}^{K},\mathbf{V}\mathbf{W}_{i}^{V}\right),i=1,\ldots,H.$
(46)

###  A.7 Proof of proposition 1

Given a $L$ layer GNN with uniform hidden feature and initial feature
$\mathbf{H}_{0}=\mathbf{X}$, for $l=0,...,L$, the recurrent output of a GNN
layer $\mathbf{H}_{l+1}$ follows

$\mathbf{H}_{l+1}=\sigma(\mathbf{H}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbf{H}_{l}\mathbf{W}_{2l}^{\top})$
(47)

where $\mathbf{H}_{l}\in\mathbb{R}^{|V|\times d}$,
$\mathbf{W}_{1l},\mathbf{W}_{2l}\in\mathbb{R}^{d\times d}$. And then we
introduce another recurrent relationship to track the output change of each
layers propagated from an initial perturbation
$\mathbold{\epsilon}_{0}\in\mathbb{R}^{|V|\times d}$ on $\mathbf{H}_{0}$,

$\mathbold{\epsilon}_{l+1}=\sigma(\mathbf{H}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbf{H}_{l}\mathbf{W}_{2l}^{\top}+\mathbold{\epsilon}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbold{\epsilon}_{l}\mathbf{W}_{2l}^{\top})-\sigma(\mathbf{H}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbf{H}_{l}\mathbf{W}_{2l}^{\top}).$
(48)

We denote $|\cdot|$ as an operator to replace a matrix’s ($\cdot$) elements
with absolute values and we write $|\mathbf{J}|\leq|\mathbf{K}|$ if
$|(\mathbf{J})_{ij}|\leq|(\mathbf{K})_{ij}|$. Let
$\mathbf{I}_{S}\in\mathbb{R}^{|V|\times 1}$ is an indicator vector of $S$ such
that $(\mathbf{I}_{S})_{i}=1\;\mathrm{if}\;i\in S\;\mathrm{else}\;0$. We
firstly prove a lemma below.

######  Lemma 5.

Given $\mathbold{\epsilon}_{0}=\mathbold{\alpha}$, it holds that

$|\mathbold{\epsilon}_{l}|\leq
a_{l}\mathbf{I}_{S}\mathbf{V}_{l}^{\top}+\mathbf{r}_{l}\in\mathbb{R}^{|V|\times
d}\;\mathrm{or}\;|(\mathbold{\epsilon}_{l})_{ij}|\leq
a_{l}(\mathbf{V}_{l})_{j}(\mathbf{I}_{S})_{i}+(\mathbf{r}_{l})_{ij}$ (49)

where $a_{l}=\epsilon(\tau+1)^{l}$, $\mathbf{V}_{l}\in\mathbb{R}^{d\times
1}_{+},||\mathbf{V}_{l}||_{1}\leq d$ and $||\mathbf{r}_{l}||_{1}\leq
2r\epsilon m(l+1)(\tau+1)^{l}$.

###### Proof.

We prove by induction. For $l=0$, we can take $a_{0}=\epsilon$,
$\mathbf{V}_{0}=\mathbf{I}_{d}=[\underbrace{1,...,1}_{d}]$ and
$\mathbf{r}_{0}=\mathbf{0}$, then it holds

$|\mathbold{\epsilon}_{0}|\leq
a_{0}\mathbf{I}_{S}\mathbf{V}_{0}^{\top}+\mathbf{r}_{0}.$ (50)

From the recurrent relation in equation 48, it holds that

$\mathbold{\epsilon}_{l+1}=\sigma((\mathbf{H}_{l}+\mathbold{\epsilon}_{l})\mathbf{W}_{1l}^{\top}+\mathbf{A}(\mathbf{H}_{l}+\mathbold{\epsilon}_{l})\mathbf{W}_{2l}^{\top})-\sigma(\mathbf{H}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbf{H}_{l}\mathbf{W}_{2l}^{\top}).$
(51)

From the Lipschitz continuity of $\sigma$, it holds that

$|\mathbold{\epsilon}_{l+1}|\leq|\mathbold{\epsilon}_{l}\mathbf{W}_{1l}^{\top}+\mathbf{A}\mathbold{\epsilon}_{l}\mathbf{W}_{2l}^{\top}|.$
(52)

From the triangle inequality, we have

$|\mathbold{\epsilon}_{l+1}|\leq|\mathbold{\epsilon}_{l}||\mathbf{W}_{1l}^{\top}|+\mathbf{A}|\mathbold{\epsilon}_{l}||\mathbf{W}_{2l}^{\top}|.$
(53)

From the assumption the statement holds at $l$th layer, we have

$(*)\quad|\mathbold{\epsilon}_{l}|\leq
a_{l}\mathbf{I}_{S}\mathbf{V}_{l}^{\top}+\mathbf{r}_{l}.$ (54)

Substitute equation 54 into equation 53, we have,

$|\mathbold{\epsilon}_{l+1}|\leq(a_{l}\mathbf{I}_{S}\mathbf{V}_{l}^{\top}+\mathbf{r}_{l})|\mathbf{W}_{1l}^{\top}|+\mathbf{A}(a_{l}\mathbf{I}_{S}\mathbf{V}_{l}^{\top}+\mathbf{r}_{l})|\mathbf{W}_{2l}^{\top}|$
(55)

Expand the above equation,

$|\mathbold{\epsilon}_{l+1}|\leq
a_{l}(\mathbf{A}\mathbf{I}_{S})(\mathbf{V}_{l}^{\top}|\mathbf{W}_{2l}^{\top}|)+\mathbf{A}\mathbf{r}_{l}|\mathbf{W}_{2l}^{\top}|+a_{l}\mathbf{I}_{S}\mathbf{V}^{\top}_{l}|\mathbf{W}_{1l}^{\top}|+\mathbf{r}_{l}|\mathbf{W}_{1l}^{\top}|$
(56)

Using the property of undirected $\tau$-graph, it holds that

$\mathbf{A}\mathbf{I}_{S}=\tau\mathbf{I}_{S}-\sum_{(i,j)\in E,i\in S,j\in
T}(\mathbf{E}_{i}-\mathbf{E}_{j})=\tau\mathbf{I}_{S}+\mathbf{B}_{S},$ (57)

where we denote

$\mathbf{B}_{S}=-\sum_{(i,j)\in E,i\in S,j\in
T}(\mathbf{E}_{i}-\mathbf{E}_{j}),$ (58)

and $\mathbf{E}_{i},\mathbf{E}_{j}\in\mathbb{R}^{|V|\times 1}$ are unit
vectors with $i$th and $j$th entry equal to 1 respectively. Then it is trivial
to show that

$||\mathbf{B}_{S}||_{1}\leq 2m.$ (59)

Substitute equation 57 into equation 56, we have

$|\mathbold{\epsilon}_{l+1}|\leq
a_{l}\tau\mathbf{I}_{S}\mathbf{V}_{l}^{\top}|\mathbf{W}_{2l}^{\top}|+a_{l}\mathbf{B}_{S}\mathbf{V}_{l}^{\top}|\mathbf{W}_{2l}^{\top}|+\mathbf{A}\mathbf{r}_{l}|\mathbf{W}_{2l}^{\top}|+a_{l}\mathbf{I}_{S}\mathbf{V}^{\top}_{l}|\mathbf{W}_{1l}^{\top}|+\mathbf{r}_{l}|\mathbf{W}_{1l}^{\top}|.$
(60)

Let

$\begin{split}&a_{l+1}=(1+\tau)a_{l},\\\
&\mathbf{V}_{l+1}^{\top}=\frac{\tau}{\tau+1}\mathbf{V}_{l}^{\top}|\mathbf{W}_{2l}^{\top}|+\frac{1}{\tau+1}\mathbf{V}_{l}^{\top}|\mathbf{W}_{1l}^{T}|,\\\
&\mathbf{r}_{l+1}=a_{l}\mathbf{B}_{S}\mathbf{V}_{l}^{\top}|\mathbf{W}_{2l}^{\top}|+\mathbf{A}\mathbf{r}_{l}|\mathbf{W}_{2l}^{\top}|+\mathbf{r}_{l}|\mathbf{W}_{1l}^{\top}|,\end{split}$
(61)

then we rewrite equation 60 as

$|\mathbold{\epsilon}_{l+1}|\leq
a_{l+1}\mathbf{I}_{S}\mathbf{V}_{l+1}^{\top}+\mathbf{r}_{l+1}$ (62)

From the assumption that

$||\mathbf{W}_{1l}||_{1}\leq 1,\quad||\mathbf{W}_{2l}||_{1}\leq 1,$ (63)

we have

$||(|\mathbf{W}_{1l}|)||_{1}=||\mathbf{W}_{1l}||_{1}\leq
1,\quad||(|\mathbf{W}_{2l}|)||_{1}=||\mathbf{W}_{2l}||_{1}\leq 1.$ (64)

So substitute equation 64, equation 59 and equation 54 into equation 61,

$\begin{split}&a_{l+1}=(\tau+1)a_{l}\leq\epsilon(\tau+1)^{l+1}\\\
&||\mathbf{V}_{l+1}^{\top}||\leq\frac{\tau}{\tau+1}||\mathbf{V}_{l}^{\top}||_{1}+\frac{1}{\tau+1}||\mathbf{V}_{l}^{\top}||\leq
d\end{split}$ (65)

and

$\begin{split}||\mathbf{r}_{l+1}||_{1}&\leq
a_{l}||\mathbf{B}_{S}||_{1}||\mathbf{V}_{l}^{\top}||_{1}+||\mathbf{A}||_{1}||\mathbf{r}_{l}||_{1}+||\mathbf{r}_{l}||_{1}\\\
&\leq 2a_{l}md+(\tau+1)||\mathbf{r}_{l}||_{1}\leq
2md\epsilon(\tau+1)^{l}+(\tau+1)||\mathbf{r}_{l}||_{1}\\\ &\leq
2md\epsilon(\tau+1)^{l}+2md\epsilon(l+1)(\tau+1)^{l+1}\\\ &\leq
2md\epsilon(\tau+1)^{l+1}+2md\epsilon(l+1)(\tau+1)^{l+1}=2md\epsilon(l+2)(\tau+1)^{l+1}.\end{split}$
(66)

This finishes the induction. ∎

The above lemma gives

$\underset{\begin{subarray}{c}||\mathbf{W}_{1l}||_{1}\\\
||\mathbf{W}_{2l}||_{1}\\\
\mathbold{\alpha}\end{subarray}}{\max}|\mathbold{\epsilon}_{l}|\leq\epsilon(\tau+1)^{l}\mathbf{I}_{S}\mathbf{V}_{l}^{\top}+\mathbf{r}_{l}$
(67)

where $||\mathbf{V}_{l}^{\top}||\leq d$ and $||\mathbf{r}_{l}||_{1}\leq
2d\epsilon m(l+1)(\tau+1)^{l}$. So when only looking at indices
$\mathbold{\epsilon}_{ij}$ with $i\in T$, the first term vanishes and it holds
that

$\underset{\begin{subarray}{c}||\mathbf{W}_{1l}||_{1}\\\
||\mathbf{W}_{2l}||_{1}\\\ \mathbold{\alpha}\end{subarray}}{\max}\sum_{i\in
T}|\mathbold{\epsilon}_{l}|_{ij}\leq 2d\epsilon m(l+1)(\tau+1)^{l}$ (68)

For the denominator, we simply construct $\mathbf{W}_{1l}=\mathbf{W}_{2l}$ as
both identity matrix and take $\mathbold{\epsilon}_{0}=\mathbold{\beta}$. Then
it simply holds that

$|\mathbold{\epsilon}_{0}|=(1+\tau)^{0}\epsilon\mathbf{I}_{T}\mathbf{I}_{d}^{\top}$
(69)

where $\mathbf{I}_{T}$ is the indicator vector on set $T$. Assume it holds,

$\mathbold{\epsilon}_{l}=(1+\tau)^{l}\epsilon\mathbf{I}_{T}\mathbf{I}_{d}^{\top}$
(70)

then from the Lipschitz continuity (ReLU) of $\sigma$ and standard
$\tau$-graph, it holds that

$\mathbold{\epsilon}_{l+1}=\sigma((\mathbf{I}+\mathbf{A})(\mathbf{H}_{l}+\mathbold{\epsilon}_{l}))-\sigma((\mathbf{I}+\mathbf{A})(\mathbf{H}_{l}))=(1+d)^{l}\epsilon(\mathbf{I}+\mathbf{A})\mathbf{I}_{T}\mathbf{I}_{d}^{\top}=(1+\tau)^{l+1}\epsilon\mathbf{I}_{T}\mathbf{I}_{d}^{\top}$
(71)

So we can get

$\sum_{i\in T}|(\mathbold{\epsilon}_{l})_{ij}|=\epsilon(1+\tau)^{l}\sum_{i\in
T}(\mathbf{I}_{T}\mathbf{I}_{d}^{\top})_{ij}=(1+\tau)^{l}\epsilon|T|d$ (72)

So that it holds that

$\underset{\begin{subarray}{c}||\mathbf{W}_{1l}||_{1}\\\
||\mathbf{W}_{2l}||_{1}\\\ \mathbold{\beta}\end{subarray}}{\max}\sum_{i\in
T}|\mathbold{\epsilon}_{l}|_{ij}\geq(1+\tau)^{l}\epsilon|T|d$ (73)

Combine equation 68 and equation 73, and substitute the last layer number as
$L-1$, we have

$\frac{\underset{\begin{subarray}{c}||\mathbf{W}_{1l}||_{1}\\\
||\mathbf{W}_{2l}||_{1}\\\ \mathbold{\alpha}\end{subarray}}{\max}\sum_{i\in
T}|\mathbold{\epsilon}_{l}|_{ij}}{\underset{\begin{subarray}{c}||\mathbf{W}_{1l}||_{1}\\\
||\mathbf{W}_{2l}||_{1}\\\ \mathbold{\beta}\end{subarray}}{\max}\sum_{i\in
T}|\mathbold{\epsilon}_{l}|_{ij}}\leq\frac{2mL}{|T|}.$ (74)

###  A.8 Proof of Theorem 5

From the proof of proposition 1 in section A.7, by simply constructing
$\mathbf{W}_{1l},\mathbf{W}_{2l}$ in the node-level GNN as identity matrix, we
have

$\begin{split}\sum_{i\in
S}|(\mathbold{\epsilon}_{S})_{ij}|=(1+\tau)^{L}\epsilon|S|d\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\alpha},\\\
\sum_{i\in
T}|(\mathbold{\epsilon}_{S})_{ij}|=(1+\tau)^{L}\epsilon|T|d\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\beta}.\end{split}$
(75)

Then from Lipschitz continuity (ReLU) we have

$\begin{split}\eta_{S\rightarrow
T}&=g_{T}^{\mathrm{GNN}}(\mathbf{X}+\mathbold{\alpha})-g_{T}^{\mathrm{GNN}}(\mathbf{X})\\\
&=\sigma(\mathbf{z}_{T}\mathbf{W}_{1}^{\top}+(\mathbf{z}_{S}+\frac{1}{|V|}\sum_{i\in
S}(\mathbold{\epsilon}_{S})_{ij})\mathbf{W}_{2})-\sigma(\mathbf{z}_{T}\mathbf{W}_{1}^{\top}+\mathbf{z}_{S}\mathbf{W}_{2}^{\top})\\\
&=(\frac{1}{|V|}\sum_{i\in
S}(\mathbold{\epsilon}_{S})_{ij})\mathbf{W}_{2}^{\top}\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\alpha}\end{split}$
(76)

and

$||\eta_{S\rightarrow T}||_{1}=||(\frac{1}{|V|}\sum_{i\in
S}(\mathbold{\epsilon}_{S})_{ij})\mathbf{W}_{2}^{\top}||_{1}=||\mathbf{W}_{2}^{\top}||_{1}(1+\tau)^{L}\epsilon\frac{|S|}{|V|}d\;\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\alpha}$
(77)

Similarly, we can get

$\eta_{T\rightarrow T}=(\frac{1}{|V|}\sum_{i\in
T}(\mathbold{\epsilon}_{T})_{ij})\mathbf{W}_{1}^{\top}.\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\beta},$
(78)

and

$||\eta_{T\rightarrow
T}||_{1}=||\mathbf{W}_{1}^{\top}||_{1}(1+\tau)^{L}\epsilon\frac{|T|}{|V|}d\;\quad\mathrm{if}\;\mathbold{\epsilon}_{0}=\mathbold{\beta},$
(79)

Then we can simply make
$\frac{||\mathbf{W}_{1}||_{1}}{||\mathbf{W}_{2}||_{1}}=\frac{|S|}{|T|}$, so
that the ratio is 1.

######  Remark 2.

The assumption of output norm unification can be achieved by standard
normalization, such as batch and layer normalizations. Lipschitz continuity
exists widely in the activation functions such as ReLU. And most molecules can
be modeled as qusi standard graphs. These assumptions are fair assumptions in
graph learning. Although it is difficult to universally obtain a precise and
tight bound, the existence of such bounds is still helpful for GNN structure
design.

######  Remark 3.

The ratio may become informative if there are information bottlenecks within a
cluster. We can mitigate the problem by having an appropriate, sufficient
number of clusters. However, the number of clusters can not be too large, so
there is a tradeoff between avoiding bottlenecks and computational cost.

Here, we also introduce a heuristic example for possibly extending to a non-
standard graph. Let subgraph $S$ be an cycle ($2$-graph) and subgraph $T$ be a
clique ($n$-graph), approximately. And we assume $|S|=|T|=n$, and node values
are all units with perturbation $\epsilon$. After one propagation of node
level, each node in $S$ has the value $3(1+\epsilon)$, each node in $T$ has
the value $n(1+\epsilon)$. Then at patch level, equations (75), (77), (79) are
modified accordingly as $\eta_{S\rightarrow T}=3\epsilon W_{2}$,
$\eta_{T\rightarrow T}=n\epsilon W_{1}$, then $\frac{\eta_{S\rightarrow
T}}{\eta_{T\rightarrow T}}=\frac{3}{n}\cdot\frac{W_{2}}{W_{1}}$, which
indicates the unevenness may affect the performance. However, if at patch
level $\frac{W_{2}}{W_{1}}\approx O(n)$ can be learned, we can still reach a
sub-optimal balance. Actually, if $\frac{W_{2}}{W_{1}}>1$ can be learned, it
will help mitigate the bottleneck anyway.

###  A.9 Graph segmentation

Molecules | Eigenvectors | Patch Segmentation
---|---|---
|  |
|  |
|  |
|  |
|  |
Figure 8: Examples of eigenvectors, and graph patches for molecules.

As a graph has an irregular structure and contains rich structural
information, forming patches on a graph is not as straightforward as
segmenting images. The previous works [9, 22] generally split an image in the
euclidean space. However, graphs are segmented through spectral clustering
based on its topology. Figure 8 shows the second eigenvector and patch
segmentations based on the algorithm described in Section 3.2. It can be seen
that the eigenvectors change along with the graph structures, and the graphs
are splitted into several function groups. Such patches are useful for
discriminating the property of the given molecule.

###  A.10 More results

Table 7 provides the performance of PatchGT on ogbg-moltox21 and ogbg-
moltoxcast.

Table 3: Results $(\%)$ on OGB datasets | ogbg-moltox21 | ogbg-moltoxcast
---|---|---
GCN +VN | 75.51 $\pm$ 0.86 | 66.33$\pm$0.35
GIN + VN | 76.21 $\pm$ 0.82 | 66.18 $\pm$0.68
GRAPHSNN +VN | 76.78$\pm$ 1.27 | 67.68 $\pm$ 0.92
PatchGT-GCN | 76.49 $\pm$0.93 | 66.58 $\pm$0.47
PatchGT-GIN |  77.26 $\pm$ 0.80 |  67.95 $\pm$0.55

###  A.11 Datasets

Table 4 contains the statistics for the six datasets from Open Graph Bechmark
(OGB) [14], and Table 5 contains the statistics for the six datasets from TU
datasets [25].

Table 4: Statistics of OGB datasets Name |  $\\#$Graphs |  $\\#$Nodes per graphs |  $\\#$Edges per graph |  $\\#$Tasks
---|---|---|---|---
molhiv | 41,127 | 25.5 | 27.5 | 1
molbace | 1,513 | 34.1 | 36.9 | 1
molclintox | 1,477 | 26.2 | 27.9 | 2
molsider | 1,427 | 33.6 | 35.4 | 27
ogbg-moltox21 | 7,831 | 18.6 | 19.3 | 12
ogbg-moltoxcast | 8,576 | 18.8 | 19.3 | 617
Table 5: Statistics of TU datasets Name |  $\\#$Graphs |  $\\#$Nodes per graphs |  $\\#$Edges per graph
---|---|---|---
DD | 1,178 | 284.3 | 715.7
MUTAG | 188 | 17.9 | 19.8
PROTEINS | 1,113 | 39.1 | 72.8
PTC-MR | 344 | 14.3 | 14.7
ENZYMES | 600 | 32.6 | 62.1
Mutagenicity | 4,337 | 30.3 | 30.8

###  A.12 Hyper-parameters selection

We report the detailed hyper-parameter settings used for training PatchGT in
Table 6. The search space for $\lambda$ is $\\{0.1,0.2,0.4,0.5,0.8\\}$.

Table 6: Model Configurations and Hyper-parameters | OGB | TU
---|---|---
$\\#$ GNN layers | 5 | 4
$\\#$ patch-GNN layers | 2 | 2
Embedding Dropout | 0.0 | 0.1
Hidden Dimension $d$ | 512 | 256
$\\#$ Attention Heads | 16 | 4
Attention Dropout | 0.1 | 0.1
Batch Size | 512 | 256
Learning Rate | 1e-4 | 1e-4
Max Epochs | 150 | 50
eigenvalue threshold $\lambda$ | $\\{0.1,0.2,0.4,0.5,0.8\\}$

###  A.13 Visualization of attention on nodes

Figure 9 shows more attention on graphs. We notice that some patches the model
concentrates on are far away from each other. This can help address
information bottleneck in the graph. Also, it provides more model
interpretation.

Figure 9: Attention visualization of PatchGT on ogbg-molhiv molecules.

###  A.14  Analysis of the computational complexity

We compare our computational complexity with the node-level Transformer,
Graphormer [38]. The comptutational complexity for both framework can be
classified into two parts. The first part is extracting graph structure
infromation. For PatchGT, the complexity is $O(|V|^{3})$ for calculating the
eigenvectors and perform kmeans for $k$ patches. For Graphormer, the
complexity is $O(|V|^{4})$ due to node pairwise shortest path computation.

######  Remark 4.

The software and algorithms of eigen-decomposition are being widely developed
in many disciplines [10]. The complexity can be reduced to $O(|V|^{2})$ if a
partial query and approximation of eigenvectors and eigenvalues are allowed
[8, 29]. And spectral clustering does not require all eigenvectors with exact
values. However, we admit that for graphs with eigenvalues that are too close
to each other, the complexity of computing the eigenvectors takes $O(N^{3})$.

The second part is neural network computation. For PatchGT, the complexity is
$O(|E|)$ for GNN if the adjacency matrix is sparse and $O(k^{2})$ for
Transformer. And for Graphormer, the complexity of Transformer is
$O(|V|^{2})$. It shoud be noticed that for a large graph, $k<<|V|$. Overall,
the complexity of patch-level transformer is significantly less than that of
applying transformer directly on the node level.

For other hierarchical pooling methods, they also need $O(L|E|)$ to learn the
segmentation ($L$ is the number of layers used in GNN), which is comparable to
spectral clustering. And spectral clustering is easier for parallel
computation. Specifically, for a $N_{\text{pool}}$-level hierarchical pooling,
it needs $O(\sum_{i=1}^{N_{\text{pool}}}L_{i}|E_{i}|)$ to learn the
segmentation and $O(\sum_{i=1}^{N_{\text{pool}}}|V_{i}|d_{i}k_{i})$ to perform
the segmentation. When training epoch number becomes a large number, the extra
accumulated cost is non-trivial. Our segmentation cost does not scale with the
training iterations.

###  A.15 Frequencies of motifs

There are two classes in ogbg-molhiv, and we record the frequencies of motifs
PatchGT pay attention to. There are an apparent difference between the two
classes. It indicates the model has a better interpretability.

Figure 10: Frequency of motifs PatchGT pay attention to in two classes.

###  A.16 Ablation study for patch level GNN

In PatchGT, we apply patch level GNN to the entire graph. We can also apply it
to each patch so that there would not be any connection between subgraphs.
Here we test the difference of these two designs.

Table 7: Results $(\%)$ on ogbg-molhiv | single GNN | multiple GNNs
---|---|---
PatchGT-GCN | 80.22 $\pm$0.0.84 | 79.13 $\pm$0.47
PatchGT-GIN | 79.99 $\pm$ 1.21 | 78.96 $\pm$0.55


$n_{v}=2m-3$, and the number of vertical edges connecting the two vertices
which go through the temporal boundary is $n^{c}_{v}=d_{m}-2m+2$. As such, to
ensure that MWPM does not map to a temporal boundary, thus incorrectly
flipping the parity of the multi-qubit Pauli measurement, we must choose a
large enough value of $d_{m}$ such that $n^{c}_{v}>n_{v}$ resulting in
$d_{m}>4m-5$. Such an increase in $d_{m}$ has the effect of roughly doubling
the runtime of a quantum algorithm. This increase in $d_{m}$ should be
expected since performing a vertical cleanup is equivalent to adding
additional measurement errors to the system “by hand”, thus requiring a
doubling in the code distance to have the same protection compared to a scheme
which doesn’t perform a vertical cleanup.

Two variations of the vertical cleanup protocol during a lattice surgery merge
may maintain the full timelike distance and thus require fewer syndrome
measurement rounds. The first variation would consist of identifying vertical
pairs prior to applying the local NN decoder, and only removing new vertical
pairs which are created after the local NN decoder is applied. In this case,
vertical pairs due to measurement errors would not be removed, although this
comes at the the cost of a higher syndrome density. Another approach would be
to re-weight vertical edges incident to highlighted vertices which were
removed from a vertical cleanup protocol, so that MWPM would give preferences
to paths which go through such vertices. Lastly, using a TELS protocol
described in Ref. [16] would allow larger timelike failure rates and thus
could be used to avoid having to use a large value of $d_{m}$ when performing
a vertical cleanup. We leave the numerical analysis of such protocols, along
with using TELS alongside a vertical cleanup strategy, to future work.

## References

  * Shor [1994] P. W. Shor, Polynomial-time algorithms for prime factorization and discrete logarithms on a quantum computer, in  _Proceedings of the 35th Annual Symposium on Foundations of Computer Science, Santa Fe, NM_ (1994) quant-ph/9508027 .
  * Grover [1996] L. K. Grover, Rapid sampling though quantum computing, in  _Proceedings of the 28th Annual ACM Symposium on the Theory of Computation_ (ACM, New York, 1996) pp. 212–219.
  * Grover [1997] L. K. Grover, Quantum mechanics helps in searching for a needle in a haystack, Phys. Rev. Lett. 79, 325 (1997).
  * Shor [1996] P. W. Shor, Fault-tolerant quantum computation, in  _Proceedings of the 37th Annual Symposium on Foundations of Computer Science_ (IEEE, 1996) pp. 56–65.
  * Chamberland  _et al._ [2016] C. Chamberland, T. Jochym-O’Connor, and R. Laflamme, Thresholds for universal concatenated quantum codes, Phys. Rev. Lett. 117, 010501 (2016).
  * Chamberland  _et al._ [2017] C. Chamberland, T. Jochym-O’Connor, and R. Laflamme, Overhead analysis of universal concatenated quantum codes, Phys. Rev. A 95, 022313 (2017).
  * Fowler  _et al._ [2012a] A. G. Fowler, M. Mariantoni, J. M. Martinis, and A. N. Cleland, Surface codes: Towards practical large-scale quantum computation, Phys. Rev. A 86, 032324 (2012a).
  * Paetznick and Reichardt [2013] A. Paetznick and B. W. Reichardt, Universal fault-tolerant quantum computation with only transversal gates and error correction, Physical review letters 111, 090505 (2013).
  * Anderson  _et al._ [2014] J. T. Anderson, G. Duclos-Cianci, and D. Poulin, Fault-tolerant conversion between the steane and reed-muller quantum codes, Phys. Rev. Lett. 113, 080501 (2014).
  * Yoder  _et al._ [2016] T. J. Yoder, R. Takagi, and I. L. Chuang, Universal fault-tolerant gates on concatenated stabilizer codes, Phys. Rev. X 6, 031039 (2016).
  * Fowler and Gidney [2018] A. G. Fowler and C. Gidney, Low overhead quantum computation using lattice surgery, arXiv preprint arXiv:1808.06709 (2018).
  * Litinski [2019a] D. Litinski, A Game of Surface Codes: Large-Scale Quantum Computing with Lattice Surgery, Quantum 3, 128 (2019a).
  * Litinski [2019b] D. Litinski, Magic State Distillation: Not as Costly as You Think, Quantum 3, 205 (2019b).
  * Chamberland and Noh [2020] C. Chamberland and K. Noh, Very low overhead fault-tolerant magic state preparation using redundant ancilla encoding and flag qubits, npj Quantum Information 6, 91 (2020).
  * Chamberland  _et al._ [2022] C. Chamberland, K. Noh, P. Arrangoiz-Arriola, E. T. Campbell, C. T. Hann, J. Iverson, H. Putterman, T. C. Bohdanowicz, S. T. Flammia, A. Keller, G. Refael, J. Preskill, L. Jiang, A. H. Safavi-Naeini, O. Painter, and F. G. Brandão, Building a fault-tolerant quantum computer using concatenated cat codes, PRX Quantum 3, 010329 (2022).
  * Chamberland and Campbell [2022a] C. Chamberland and E. T. Campbell, Universal quantum computing with twist-free and temporally encoded lattice surgery, PRX Quantum 3, 010331 (2022a).
  * Gottesman [1998] D. Gottesman, Theory of fault-tolerant quantum computation, Phys. Rev. A 57, 127 (1998).
  * Terhal [2015] B. M. Terhal, Quantum error correction for quantum memories, Rev. Mod. Phys. 87, 307 (2015).
  * Harrington [2004] J. Harrington, Analysis of quantum error-correcting codes: Symplectic lattice codes and toric codes,  (2004).
  * Breuckmann  _et al._ [2017] N. P. Breuckmann, K. Duivenvoorden, D. Michels, and B. M. Terhal, Local Decoders for the 2D and 4D Toric Code, Quantum Information and Computation 17, 0181 (2017).
  * Herold  _et al._ [2017] M. Herold, M. J. Kastoryano, E. T. Campbell, and J. Eisert, Cellular automaton decoders of topological quantum memories in the fault tolerant setting, New Journal of Physics 19, 063012 (2017).
  * Kubica and Preskill [2019] A. Kubica and J. Preskill, Cellular-automaton decoders with provable thresholds for topological codes, Phys. Rev. Lett. 123, 020501 (2019).
  * Vasmer  _et al._ [2021] M. Vasmer, D. E. Browne, and A. Kubica, Cellular automaton decoders for topological quantum codes with noisy measurements and beyond, Scientific Reports 11, 2027 (2021).
  * Duclos-Cianci and Poulin [2010] G. Duclos-Cianci and D. Poulin, Fast decoders for topological quantum codes, Phys. Rev. Lett. 104, 050504 (2010).
  * Duclos-Cianci and Poulin [2014] G. Duclos-Cianci and D. Poulin, Fault-tolerant renormalization group decoder for abelian topological codes, Quantum Information and Computation 14, 721–740 (2014).
  * Delfosse and Nickerson [2021] N. Delfosse and N. H. Nickerson, Almost-linear time decoding algorithm for topological codes, Quantum 5, 595 (2021).
  * Delfosse [2020] N. Delfosse, Hierarchical decoding to reduce hardware requirements for quantum computing, arXiv e-prints , arXiv:2001.11427 (2020).
  * Torlai and Melko [2017] G. Torlai and R. G. Melko, Neural decoder for topological codes, Phys. Rev. Lett. 119, 030501 (2017).
  * Krastanov and Jiang [2017] S. Krastanov and L. Jiang, Deep neural network probabilistic decoder for stabilizer codes, Scientific Reports 7, 11003 (2017).
  * Varsamopoulos  _et al._ [2017] S. Varsamopoulos, B. Criger, and K. Bertels, Decoding small surface codes with feedforward neural networks, Quantum Science and Technology 3, 015004 (2017).
  * Baireuther  _et al._ [2018] P. Baireuther, T. E. O’Brien, B. Tarasinski, and C. W. J. Beenakker, Machine-learning-assisted correction of correlated qubit errors in a topological code, Quantum 2, 48 (2018).
  * Breuckmann and Ni [2018] N. P. Breuckmann and X. Ni, Scalable Neural Network Decoders for Higher Dimensional Quantum Codes, Quantum 2, 68 (2018).
  * Chamberland and Ronagh [2018] C. Chamberland and P. Ronagh, Deep neural decoders for near term fault-tolerant experiments, Quantum Science and Technology 3, 044002 (2018).
  * Sweke  _et al._ [2020] R. Sweke, M. S. Kesselring, E. P. L. van Nieuwenburg, and J. Eisert, Reinforcement learning decoders for fault-tolerant quantum computation, Machine Learning: Science and Technology 2, 025005 (2020).
  * Varsamopoulos  _et al._ [2020a] S. Varsamopoulos, K. Bertels, and C. G. Almudever, Decoding surface code with a distributed neural network–based decoder, Quantum Machine Intelligence 2, 2524–4914 (2020a).
  * Andreasson  _et al._ [2019] P. Andreasson, J. Johansson, S. Liljestrand, and M. Granath, Quantum error correction for the toric code using deep reinforcement learning, Quantum 3, 183 (2019).
  * Wagner  _et al._ [2020] T. Wagner, H. Kampermann, and D. Bruß, Symmetries for a high-level neural decoder on the toric code, Phys. Rev. A 102, 042411 (2020).
  * Varsamopoulos  _et al._ [2020b] S. Varsamopoulos, K. Bertels, and C. G. Almudever, Comparing neural network based decoders for the surface code, IEEE Transactions on Computers 69, 300 (2020b).
  * Fitzek  _et al._ [2020] D. Fitzek, M. Eliasson, A. F. Kockum, and M. Granath, Deep q-learning decoder for depolarizing noise on the toric code, Phys. Rev. Research 2, 023230 (2020).
  * Sheth  _et al._ [2020] M. Sheth, S. Z. Jafarzadeh, and V. Gheorghiu, Neural ensemble decoding for topological quantum error-correcting codes, Phys. Rev. A 101, 032338 (2020).
  * Ni [2020] X. Ni, Neural Network Decoders for Large-Distance 2D Toric Codes, Quantum 4, 310 (2020).
  * Domingo Colomer  _et al._ [2020] L. Domingo Colomer, M. Skotiniotis, and R. Muñoz-Tapia, Reinforcement learning for optimal error correction of toric codes, Physics Letters A 384, 126353 (2020).
  * Meinerz  _et al._ [2022] K. Meinerz, C.-Y. Park, and S. Trebst, Scalable neural decoder for topological surface codes, Phys. Rev. Lett. 128, 080505 (2022).
  * [44] S. Gicev, L. C. L. Hollenberg, and M. Usman, A scalable and fast artificial neural network syndrome decoder for surface codes, arXiv e-prints , arXiv:2110.05854.
  * Tomita and Svore [2014] Y. Tomita and K. M. Svore, Low-distance surface codes under realistic quantum noise, Phys. Rev. A 90, 062320 (2014).
  * Edmonds [1965] J. Edmonds, Paths, trees, and flowers, Canadian Journal of mathematics 17, 449 (1965).
  * Kitaev [2003] A. Y. Kitaev, Fault-tolerant quantum computation by anyons, Annals of Physics 303, 2 (2003).
  * Dennis  _et al._ [2002] E. Dennis, A. Kitaev, A. Landahl, and J. Preskill, Topological quantum memory, Journal of Mathematical Physics 43, 4452 (2002).
  * Calderbank and Shor [1996] A. R. Calderbank and P. W. Shor, Good quantum error-correcting codes exist, Phys. Rev. A 54, 1098 (1996).
  * Steane [1996] A. Steane, Multiple particle interference and quantum error correction, Proc.Roy.Soc.Lond. A 452, 2551 (1996).
  * Huang  _et al._ [2020] S. Huang, M. Newman, and K. R. Brown, Fault-tolerant weighted union-find decoding on the toric code, Phys. Rev. A 102, 012419 (2020).
  * Fowler [2015] A. G. Fowler, Minimum weight perfect matching of fault-tolerant topological quantum error correction in average o(1) parallel time, Quantum Info. Comput. 15, 145–158 (2015).
  * Fowler  _et al._ [2012b] A. G. Fowler, A. C. Whiteside, and L. C. L. Hollenberg, Towards practical classical processing for the surface code: Timing analysis, Phys. Rev. A 86, 042313 (2012b).
  * Das  _et al._ [2020] P. Das, C. A. Pattison, S. Manne, D. M. Carmean, K. M. Svore, M. K. Qureshi, and N. Delfosse, A scalable decoder micro-architecture for fault-tolerant quantum computing, CoRR abs/2001.06598 (2020), 2001.06598 .
  * Chamberland  _et al._ [2020a] C. Chamberland, G. Zhu, T. J. Yoder, J. B. Hertzberg, and A. W. Cross, Topological and subsystem codes on low-degree graphs with flag qubits, Phys. Rev. X 10, 011022 (2020a).
  * Chamberland  _et al._ [2020b] C. Chamberland, A. Kubica, T. J. Yoder, and G. Zhu, Triangular color codes on trivalent graphs with flag qubits, New Journal of Physics 22, 023019 (2020b).
  * Knill [2005] E. Knill, Quantum computing with realistically noisy devices, Nature 434, 39 (2005).
  * Chamberland  _et al._ [2018] C. Chamberland, P. Iyer, and D. Poulin, Fault-tolerant quantum computing in the Pauli or Clifford frame with slow error diagnostics, Quantum 2, 43 (2018).
  * Litinski and Oppen [2018] D. Litinski and F. v. Oppen, Lattice Surgery with a Twist: Simplifying Clifford Gates of Surface Codes, Quantum 2, 62 (2018).
  * Chamberland and Campbell [2022b] C. Chamberland and E. T. Campbell, Circuit-level protocol and analysis for twist-based lattice surgery, Phys. Rev. Research 4, 023090 (2022b).
  * [61] C. Gidney and A. G. Fowler, Flexible layout of surface code computations using AutoCCZ states, arXiv e-prints , arXiv:1905.08916.
  * Krinner  _et al._ [2022] S. Krinner, N. Lacroix, A. Remm, A. Di Paolo, E. Genois, C. Leroux, C. Hellings, S. Lazar, F. Swiadek, J. Herrmann, G. J. Norris, C. K. Andersen, M. Müller, A. Blais, C. Eichler, and A. Wallraff, Realizing repeated quantum error correction in a distance-three surface code, Nature 605, 669–674 (2022).
  * He  _et al._ [2016] K. He, X. Zhang, S. Ren, and J. Sun, Deep residual learning for image recognition,  _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ ,  , 770–778 (2016).
  * Wu  _et al._ [2020] D. Wu, Y. Wang, S.-T. Xia, J. Bailey, and X. Ma, Skip Connections Matter: On the Transferability of Adversarial Examples Generated with ResNets, arXiv e-prints , arXiv:2002.05990 (2020).
  * Fahim  _et al._ [2021] F. Fahim, B. Hawks, C. Herwig, J. Hirschauer, S. Jindariani, N. Tran, L. P. Carloni, G. D. Guglielmo, P. C. Harris, J. D. Krupa, D. S. Rankin, M. B. Valentin, J. D. Hester, Y. Luo, J. Mamish, S. Orgrenci-Memik, T. Aarrestad, H. Javed, V. Loncar, M. Pierini, A. A. Pol, S. Summers, J. M. Duarte, S. Hauck, S. Hsu, J. Ngadiuba, M. Liu, D. Hoang, E. Kreinar, and Z. Wu, hls4ml: An open-source codesign workflow to empower scientific low-power machine learning devices, CoRR abs/2103.05579 (2021), 2103.05579 .
  * Liu  _et al._ [2019] Z. Liu, P. Chow, J. Xu, J. Jiang, Y. Dou, and J. Zhou, A uniform architecture design for accelerating 2d and 3d cnns on fpgas, Electronics 8, 10.3390/electronics8010065 (2019).
  * Geng  _et al._ [2019] T. Geng, T. Wang, C. Wu, C. Yang, S. L. Song, A. Li, and M. Herbordt, LP-BNN: Ultra-low-latency BNN inference with layer parallelism, Proceedings of the International Conference on Application-Specific Systems, Architectures and Processors 2019-July, 9 (2019).
  * Abts  _et al._ [2020] D. Abts, J. Ross, J. Sparling, M. Wong-VanHaren, M. Baker, T. Hawkins, A. Bell, J. Thompson, T. Kahsai, G. Kimmell, J. Hwang, R. Leslie-Hurd, M. Bye, E. R. Creswick, M. Boyd, M. Venigalla, E. Laforge, J. Purdy, P. Kamath, D. Maheshwari, M. Beidler, G. Rosseel, O. Ahmad, G. Gagarin, R. Czekalski, A. Rane, S. Parmar, J. Werner, J. Sproch, A. Macias, and B. Kurtz, Think fast: A tensor streaming processor (TSP) for accelerating deep learning workloads, in  _47th ACM/IEEE Annual International Symposium on Computer Architecture, ISCA 2020, Valencia, Spain, May 30 - June 3, 2020_ (IEEE, 2020) pp. 145–158.
  * Higgott [2021] O. Higgott, PyMatching: A Python package for decoding quantum codes with minimum-weight perfect matching, arXiv e-prints , arXiv:2105.13082 (2021).
  * Edwards and Xie [2016] M. Edwards and X. Xie, Graph Based Convolutional Neural Network,  , arXiv:1609.08965 (2016).
  * Vuillot  _et al._ [2019] C. Vuillot, L. Lao, B. Criger, C. G. Almudéver, K. Bertels, and B. M. Terhal, Code deformation and lattice surgery are gauge fixing, New J. Phys. 21, 033028 (2019).