# Multiresolution Graph Transformers and Wavelet Positional Encoding for
Learning Hierarchical Structures

Nhat Khang Ngo  Truong Son Hy  Risi Kondor

###### Abstract

Contemporary graph learning algorithms are not well-defined for large
molecules since they do not consider the hierarchical interactions among the
atoms, which are essential to determine the molecular properties of
macromolecules. In this work, we propose Multiresolution Graph Transformers
(MGT), the first graph transformer architecture that can learn to represent
large molecules at multiple scales. MGT can learn to produce representations
for the atoms and group them into meaningful functional groups or repeating
units. We also introduce Wavelet Positional Encoding (WavePE), a new
positional encoding method that can guarantee localization in both spectral
and spatial domains. Our approach achieves competitive results on two
macromolecule datasets consisting of polymers and peptides. Furthermore, the
visualizations, including clustering results on macromolecules and low-
dimensional spaces of their representations, demonstrate the capability of our
methodology in learning to represent long-range and hierarchical structures.
Our PyTorch implementation is publicly available at
https://github.com/HySonLab/Multires-Graph-Transformer.

Machine Learning, ICML

(a)

(b)

(c)

Figure 1: Overview of Wavelet Positional Encoding (WavePE) and Multiresolution
Graph Transformer (MGT). a) $k$ diffusion matrices of size $N\times N$ are
stacked together to produce a wavelets tensor with size $N\times N\times K$
which are contracted by equivariant encoding methods to yield a tensor of
positional representation $N\times k$. b) Atomic representations are derived
by passing the molecular graph augmented with positional features through $L$
GPS layers. c) A macromolecule is decomposed into several substructures in
which the features are aggregated from the atom-level outputs, resulting in a
set of substructures that are moved to a Transformer encoder.

##  1 Introduction

Macromolecules are long-range and hierarchical structures as they consist of
many substructures. While small molecules in existing datasets (Ruddigkeit et
al., 2012; Ramakrishnan et al., 2014; Sterling & Irwin, 2015) comprise less
than 50 atoms connected by simple rings and bonds, this number in a
macromolecule can be dozens or even hundreds. Substructures such as repeating
units and functional groups are intrinsic parts of macromolecules; they
present unique chemical reactions regardless of other compositions in the same
molecules (Jerry, 1992). Therefore, studying the multiresolution
characteristic of large molecules is imperative to gain comprehensive
knowledge about real-life materials like polymers or proteins (Schmid, 2022).
In recent years, several works (Anand et al., 2022; Gaul & Cuesta-Lopez, 2022;
Depta et al., 2022) have been proposed to apply machine learning algorithms to
learn macromolecules at multiple scales. These approaches, however, rely on
thorough feature selection and extraction, which are inefficient when learning
from large databases of multicomponent materials (Anand et al., 2022).

Message passing is a prevailing paradigm for designing neural networks that
operate on graph-structured data. Previous studies (Gilmer et al., 2017b; Kipf
& Welling, 2016; Veličković et al., 2018; Corso et al., 2020b; Xu et al.,
2019c) have proposed different strategies to perform message passing on graphs
and achieved remarkable results across various domains. However, message-
passing-dominated graph neural networks (GNNs) have some inherent limitations,
such as limited expressiveness capability (Morris et al., 2019; Xu et al.,
2019c), over-smoothing (Chen et al., 2020; Li et al., 2018; Oono & Suzuki,
2020), over-squashing (Alon & Yahav, 2021) issues. Over-smoothing exists in
graph neural networks that consist of a sufficiently large number of layers,
and node representations are likely to converge to a constant after going
through these deep networks. Over-squashing problems occur when messages are
ineffectively propagated and aggregated through bottlenecks on long-range
graph structures. These two shortcomings hinder GNNs from making good
predictions on long-range and hierarchically structured data. Furthermore, the
molecular properties of large molecules are formed not only by interactions
among atoms within neighborhoods but also by distant atoms. Therefore, local
information is not sufficient to model macromolecules.

Transformers are classes of deep learning models that leverage self-attention
mechanisms to handle long-range dependencies in various data domains, such as
natural language processing (Vaswani et al., 2017; Devlin et al., 2019) or
computer vision (Dosovitskiy et al., 2021; Liu et al., 2021). In graph
domains, Transformer-like architectures (Kreuzer et al., 2021; Dwivedi &
Bresson, 2020a; Rampášek et al., 2022) have proved their effectiveness in
learning node representations as they can overcome the over-smoothing and
over-squashing issues by directly measuring the pairwise relationships between
the nodes. Contrary to GNNs, graph transformers do not use the graph structure
as hard-coded information. They, instead, encode positional and structural
information on graphs as soft inductive bias, making them flexible learners in
graph learning problems (Kreuzer et al., 2021). Node positional
representations can be derived based on spectral (Dwivedi & Bresson, 2020b;
Dwivedi et al., 2022a) or spatial (You et al., 2019; Li et al., 2020b)
domains. Most existing spectral-based methods decompose the graph Laplacian
into sets of eigenvectors and eigenvalues. However, these eigenvectors have
sign ambiguity and are unstable due to eigenvalue multiplicities. On the other
hand, spatial-based approaches compute the shortest distances among the nodes;
however, these encoding methods do not consider the structural similarity
between nodes and their neighborhoods (Chen et al., 2022).

##### Contribution

We propose Multiresolution Graph Transformer (MGT) and Wavelet Positional
Encoding (WavePE), using multiresolution analysis on both spectral and spatial
domains for learning to represent hierarchical structures. Our contributions
are four-fold:

  * •

We design Multiresolution Graph Transformer (MGT), a Transformer-like
architecture that can operate on macromolecules at multiple scales. Our
proposed model can learn the atomic representations and group them into
meaningful clusters via a data-driven algorithm. Finally, the substructures,
i.e. clusters, are fed to a Transformer encoder to calculate the
representations of several substructures in macromolecules.

  * •

We introduce Wavelet Positional Encoding (WavePE), a new positional encoding
scheme for graph-structured data. Since wavelet analysis can provide
localization in both spatial and spectral domains, we construct a set of
wavelets to capture the structural information on graphs at different scales.
Then, we apply equivariant encoding methods to project the wavelet tensors
into positional representations for the atoms.

  * •

We show the effectiveness of our methodology by reporting its superior
performance on three molecular property prediction benchmarks. These datasets
contain macromolecules, i.e. peptides and polymers, that are highly
hierarchical and consist of up to hundreds of atoms.

  * •

Our visualization demonstrates the comprehensiveness of our proposed methods
in learning to represent large molecules. In general, we show the
representations of molecules produced by MGT and how MGT determines and groups
the atoms in long-chain molecules.

##  2 Related work

##### Hierachical Learning on Molecule

Functional groups or repeating units are essential phenomena in chemistry.
While functional groups constitute large molecules, repeating units are the
primary parts that produce complete polymer chains. We regard them as
substructures. In particular, similar substructures undergo similar chemical
reactions regardless of the remaining compositions existing in the molecules
(Jerry, 1992). Previous work has leveraged the hierarchical property of
molecules to improve the performance in molecular representation learning and
generation. Fang et al. (2021), Rong et al. (2020), and Chen et al. (2021)
used functional groups as prior knowledge to guide the models to predict
accurate molecular properties. For the molecular generation task, Maziarz et
al. (2022) and Jin et al. (2020) used chemical rules to extract substructures
and constructed a vocabulary of structural motifs to generate the molecules.

##### Graph Transformers

Earlier research efforts have adopted Transformer-like architectures to graph-
structured data. Dwivedi & Bresson (2020b) proposed an early approach to
generalize Transformers to graphs using Laplacian positional encoding and
performing self-attention on one-hop neighbors surrounding center nodes. On
the other hand, Kreuzer et al. (2021) computes attention scores on the entire
graph with differentiation between positive and negative edges, while also
using Laplacian positional encoding. Rong et al. (2020) introduced
GTransformer that utilizes vectorized outputs from local GNNs as inputs for a
Transformer encoder, making up an effective combination between node local and
global information. Rampášek et al. (2022) proposed a general framework that
integrates essential components of Graph Transformers, including positional or
structural encoding, graph feature extraction, local message passing, and
self-attention. Also, Chen et al. (2022) extracted multiple k-hop subgraphs
and fed them to local GNNs to compute their embeddings, which are then moved
to a Transformer encoder. Graphormer proposed in (Ying et al., 2021) uses
attention mechanisms to estimate several types of encoding, such as
centrality, spatial, and edge encodings. In addition, Kim et al. (2022)
treated all nodes and edges as independent tokens augmented with orthonormal
node identifiers and trainable type identifiers, and fed them to a standard
Transformer encoder. Moreover, Yun et al. (2019) generated multiple meta-
paths, i.e. views, of a graph and computed their pairwise attention scores,
before aggregating them into a final representation for the entire graph.

##### Graph Positional Encoding

Several approaches have been proposed to encode the positional or structural
representations into node features to improve the expressiveness of GNNs and
Graph Transformers. Node positions can be determined via spectral or spatial
domains. Spectral-based methods include Laplacian positional encoding (Dwivedi
& Bresson, 2020b; Kreuzer et al., 2021) and random walk positional encoding
(RWPE) (Dwivedi et al., 2022a). For spatial-based methods, You et al. (2019)
computed distances of sets of nodes to anchor nodes, whereas Li et al. (2020b)
calculated the shortest distances between pairs of nodes.

##### Multiresolution Analysis and Wavelet Theory

Multiresolution Analysis (MRA) has been proposed by (Mallat, 1989, 2008) as a
method to approximate signals at multiple scales in which the signals are
decomposed over elementary waveforms chosen from a family called wavelets
(i.e. mother wavelets and father wavelets), including Haar (Haar, 1910),
Daubechies (Daubechies, 1988), etc., to produce the sparse representations. In
graph and discrete domains, Hammond et al. (2011) introduced spectral graph
wavelets that are determined by applying the wavelet operator on the graph
Laplacian at multi-levels. Coifman & Maggioni (2006) proposed diffusion
wavelet that is a fast multiresolution framework for analyzing functions on
discretized structures such as graphs and manifolds. In the deep learning era,
Rustamov & Guibas (2013) and Xu et al. (2019b) leveraged the power of neural
networks for graph wavelet construction and computation.

##  3 Preliminaries

###  3.1 Notation

A molecule can be represented as an undirected graph in which nodes are the
atoms and edges are the bonds between them. In paticular, we refer to a
molecular graph as
$G=(\mathcal{V},\mathcal{E},\mathbf{A},\mathbf{X},\mathbf{P},\mathcal{V}_{s})$,
where $G$ is an undirected graph having $\mathcal{V}$ ($|\mathcal{V}|=n$) and
$\mathcal{E}$ as sets of nodes and edges respectively; also,
$\mathbf{A}\in\mathbb{R}^{n\times n}$ is the graph’s adjacency matrix. When a
graph is attributed, we augment $G$ with a set of node feature vectors
$\mathcal{X}=\\{x_{1},...,x_{n}\\},x_{i}\in\mathbb{R}^{d}$ and a set of node
positional vectors $\mathcal{P}=\\{p_{1},...,p_{n}\\},p_{i}\in\mathbb{R}^{p}$.
These two types of attributes are stored in $\mathbf{X}\in\mathbb{R}^{n\times
d}$ and $\mathbf{P}\in\mathbb{R}^{n\times p}$ correspondingly. In addition to
the atom-level representation of $G$,
$\mathcal{V}_{s}=\\{v_{s_{1}},...,v_{s_{k}}\\}$ denotes the substructure set
in which $v_{s_{i}}\subset\mathcal{V}$, i.e. $v_{s_{i}}$ is a subset of atoms
of the molecule.

###  3.2 Hierachical Learning on Molecules

Molecular property prediction is regarded as a graph-level learning task. We
need to aggregate node embeddings into graph-level vectors which are then fed
to a classifier to make predictions on graphs. Specifically, a function
$f:\mathcal{V}\xrightarrow{}\mathcal{Z}$ that maps the atom $u\in\mathcal{V}$
to a $d_{o}$-dimensional vector $z_{u}\in\mathcal{Z}\subset\mathbb{R}^{d_{o}}$
should learn to produce atom-level representations. Most existing graph neural
networks compute the vector $z=\zeta(\\{f(u)|u\in\mathcal{V}\\})$ that
indicates a representation for the entire molecular graph, where $\zeta$ can
be sum, mean, max, or more sophisticated operators. For hierarchical learning,
substructure-level representations can be derived in addition to atom-level
representations by aggregating node representations in the same substructures
as $z_{s}=\zeta(\\{f(u)|u\in v_{s}\land v_{s}\in\mathcal{V}_{s}\\})$. Instead
of atom vectors, we aggregate the substructure vectors to represent the entire
graph, i.e. $z=\zeta(\\{z_{s}|z_{s}\in\mathcal{V}_{s}\\})$. Finally, a
classifier $g$ given $z$ as inputs is trained to predict the molecular
properties.

###  3.3 Transformers on Graphs

While GNNs learn node embeddings by leveraging the graph structure via local
message-passing mechanisms, Transformers ignore localities and directly infer
the relations between pairs of nodes using only node attributes. In other
words, the node connectivity is not utilized in pure transformer-like
architectures (Vaswani et al., 2017), reducing the graph conditions to a set
learning problem. Given a tensor of node features
$\mathbf{X}\in\mathbb{R}^{n\times d}$, Transformers compute three matrices
including query ($\mathbf{Q}$), key ($\mathbf{K}$), and value ($\mathbf{V}$)
via three linear transformations $\mathbf{Q}=\mathbf{X}\mathbf{W}_{q}^{T}$,
$\mathbf{K}=\mathbf{X}\mathbf{W}_{k}^{T}$, and
$\mathbf{V}=\mathbf{X}\mathbf{W}_{v}^{T}$. A self-attention tensor
($\mathbf{H}$) can be computed as follows:

$\mathbf{H}=\text{softmax}(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d_{o}}})\mathbf{V}$
(1)

where $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$ are learnable
parameters in $\mathbb{R}^{d_{o}\times d}$, resulting in
$\mathbf{H}\in\mathbb{R}^{n\times d_{o}}$. Furthermore, each $\mathbf{H}$ in
Eq. 1 denotes an attention head. To improve effectiveness, multiple
$\\{\mathbf{H}\\}_{i=1}^{h}$ are computed, which is known as multi-head
attention. All of the attention heads are concatenated to form a final tensor:
$\mathbf{H}_{o}=\text{concat}(\mathbf{H}_{1},...,\mathbf{H}_{h})$, where $h$
is the number of attention heads. Finally, the output $\mathbf{X}^{\prime}$,
i.e. new node representations, can be computed by feeding $\mathbf{H}_{o}$
into a feed-forward neural network (FFN), i.e.
$\mathbf{X}^{\prime}=\text{FFN}(\mathbf{H}_{o})$. It is easy to see that
Transformers operating on inputs without positional encoding are permutation
invariant.

##### Positional Encoding

As pure Transformer encoders only model sets of nodes without being cognizant
of the graph structures, positional or structural information between nodes
and their neighborhoods should be incorporated into node features. In
particular, node positional representations can be added or concatenated with
node features, resulting in comprehensive inputs for Transformer-like
architectures operating on graph-structured data.

##  4 Wavelet Positional Encoding

###  4.1 Spectral Graph Wavelets

Let $\mathcal{A}\in\mathbb{R}^{n\times n}$ be the adjacency matrix of an
undirected graph $G=(\mathcal{V},\mathcal{E})$. The normalized graph Laplacian
is defined as
$\mathcal{L}=\mathcal{I}_{n}-\mathcal{D}^{-1/2}\mathcal{A}\mathcal{D}^{-1/2}$,
where $\mathcal{I}_{n}$ is the identity matrix and $\mathcal{D}$ is the
diagonal matrix of node degrees. $\mathcal{L}$ can be decomposed into a
complete set of orthonormal eigenvectors $U=(u_{1},u_{2},...,u_{n})$
associated with real and non-negative eigenvalues $\\{\lambda\\}_{1}^{n}$.
While graph Fourier transform uses $U$ as a set of bases to project the graph
signal from the vertex domain to the spectral domain, graph wavelet transform
constructs a set of spectral graph wavelets as bases for this projection via:

$\psi_{s}=U\Sigma_{s}U^{T}$

where
$\Sigma_{s}=\text{diag}(g(s\lambda_{1}),g(s\lambda_{2}),...,g(s\lambda_{n}))$
is a scaling matrix of eigenvalues,
$\psi_{s}=(\psi_{s1},\psi_{s2},...,\psi_{sn})$ and each wavelet $\psi_{si}$
indicates how a signal diffuses away from node $i$ at scale $s$; we choose
$g(s\lambda)=e^{-s\lambda}$ as a heat kernel (Donnat et al., 2018). Since a
node’s neighborhoods can be adjusted by varying the scaling parameter $s$ (Xu
et al., 2019a), using multiple sets of wavelets at different scales can
provide comprehensive information on the graph’s structure. It means that
larger values of $s_{i}$ correspond to larger neighborhoods surrounding a
center node. Figure 2 illustrates how wavelets can be used to determine
neighborhoods at different scales on a molecular graph. In this work, we
leverage this property of graph wavelets to generate node positional
representations that can capture the structural information of a center node
on the graph at different resolutions. We employ $k$ diffusion matrices
$\\{\psi_{s_{i}}\\}_{i=1}^{k}$ in which each $\psi_{s_{i}}$ has a size of
$n\times n$, resulting in a tensor of graph wavelets
$\mathcal{P}\in\mathbb{R}^{n\times n\times k}$. Additionally, WavePE is a
generalized version of RWPE (Dwivedi et al., 2022a) as the random walk process
can be regarded as a type of discrete diffusion. In the following section, we
demonstrate the use of tensor contractions to generate a tensor of node
positional representations $\mathbf{P}\in\mathbb{R}^{n\times k}$ from
$\mathcal{P}$. In general, Fig.1a demonstrates our wavelet positional encoding
method.

(a) s = 3

(b) s = 5

(c) s = 10

Figure 2: Visualization of some of the wavelets with scaling parameters on the
Aspirin $\text{C}_{9}\text{H}_{8}\text{O}_{4}$ molecular graph with 13 nodes
(i.e. heavy atoms). The center node is colored yellow. The colors varying from
bright to dark illustrate the diffusion rate from the center node to the
others, i.e. nodes that are closer to the center node have brighter colors.
Low-scale wavelets are highly localized, whereas the high-scale wavelets can
spread out more nodes on the molecular graphs

###  4.2 Equivariant Encoding

It is important to note that our spectral graph wavelets computed from the
previous section must be further encoded in a permutation-equivariant manner.
For simplicity, that means if we permute (i.e. change the order) the set of
nodes, their position encodings must be transformed accordingly. In this
section, we formally define permutation symmetry, i.e. symmetry to the action
of the symmetric group, $\mathbb{S}_{n}$, and construct permutation-
equivariant neural networks to encode graph wavelets. An element
$\sigma\in\mathbb{S}_{n}$ is a permutation of order $n$, or a bijective map
from $\\{1,\ldots,n\\}$ to $\\{1,\ldots,n\\}$. For example, the action of
$\mathbb{S}_{n}$ on an adjacency matrix $\mathcal{A}\in\mathbb{R}^{n\times n}$
and on a latent matrix $\mathcal{Z}\in\mathbb{R}^{n\times d_{z}}$ are:

$[\sigma\cdot\mathcal{A}]_{i_{1},i_{2}}=\mathcal{A}_{\sigma^{-1}(i_{1}),\sigma^{-1}(i_{2})},\
\ \ \ [\sigma\cdot\mathcal{Z}]_{i,j}=\mathcal{Z}_{\sigma^{-1}(i),j},$

for $\sigma\in\mathbb{S}_{n}$. Here, the adjacency matrix $\mathcal{A}$ is a
second-order tensor with a single feature channel, while the latent matrix
$\mathcal{Z}$ is a first-order tensor with $d_{z}$ feature channels. In
general, the action of $\mathbb{S}_{n}$ on a $k$-th order tensor
$\mathcal{X}\in\mathbb{R}^{n^{k}\times d}$ (i.e. the last index denotes the
feature channels) is defined similarly as:

$[\sigma\cdot\mathcal{X}]_{i_{1},..,i_{k},j}=\mathcal{X}_{\sigma^{-1}(i_{1}),..,\sigma^{-1}(i_{k}),j},\hskip
20.0pt\sigma\in\mathbb{S}_{n}.$

Formally, we define these equivariant and invariant properties in Def. 4.1 and
equivariant neural networks in Def. 4.2.

######  Definition 4.1.

An $\mathbb{S}_{n}$-equivariant (or permutation equivariant) function is a
function $f\colon\mathbb{R}^{n^{k}\times
d}\rightarrow\mathbb{R}^{n^{k^{\prime}}\times{d^{\prime}}}$ that satisfies
$f(\sigma\cdot\mathcal{X})=\sigma\cdot f(\mathcal{X})$ for all
$\sigma\in\mathbb{S}_{n}$ and $\mathcal{X}\in\mathbb{R}^{n^{k}\times d}$.
Similarly, we say that $f$ is $\mathbb{S}_{n}$-invariant (or permutation
invariant) if and only if $f(\sigma\cdot\mathcal{X})=f(\mathcal{X})$.

######  Definition 4.2.

An $\mathbb{S}_{n}$-equivariant network is a function
$f:\mathbb{R}^{n^{k}\times d}\rightarrow\mathbb{R}^{n^{k^{\prime}}\times
d^{\prime}}$ defined as a composition of $\mathbb{S}_{n}$-equivariant linear
functions $f_{1},..,f_{T}$ and $\mathbb{S}_{n}$-equivariant nonlinearities
$\gamma_{1},..,\gamma_{T}$:

$f=\gamma_{T}\circ f_{T}\circ..\circ\gamma_{1}\circ f_{1}.$

On the another hand, an $\mathbb{S}_{n}$-invariant network is a function
$f:\mathbb{R}^{n^{k}\times d}\rightarrow\mathbb{R}$ defined as a composition
of an $\mathbb{S}_{n}$-equivariant network $f^{\prime}$ and an
$\mathbb{S}_{n}$-invariant function on top of it, e.g.,
$f=f^{\prime\prime}\circ f^{\prime}$.

In order to build permutation-equivariant neural networks, we revisit some
basic tensor operations: the tensor product $A\otimes B$ (see Def. 4.3) and
tensor contraction ${A_{\downarrow}}_{x_{1},..,x_{p}}$ (see Def. 4.4). It can
be shown that these tensor operations respect permutation equivariance (Hy et
al., 2018) (Kondor et al., 2018).

######  Definition 4.3.

The tensor product of $A\in\mathbb{R}^{n^{a}}$ with $B\in\mathbb{R}^{n^{b}}$
yields a tensor $C=A\otimes B\in\mathbb{R}^{n^{a+b}}$ where

$C_{i_{1},i_{2},..,i_{a+b}}=A_{i_{1},i_{2},..,i_{a}}B_{i_{a+1},i_{a+2},..,i_{a+b}}.$

######  Definition 4.4.

The contraction of $A\in\mathbb{R}^{n^{a}}$ along the pair of dimensions
$\\{x,y\\}$ (assuming $x<y$) yields a $(a-2)$-th order tensor

$C_{i_{1},..,i_{x-1},j,i_{x+1},..,i_{y-1},j,i_{y+1},..,i_{a}}=\sum_{i_{x},i_{y}}A_{i_{1},..,i_{a}},$

where we assume that $i_{x}$ and $i_{y}$ have been removed from amongst the
indices of $C$. Using Einstein notation, this can be written more compactly as

$C_{\\{i_{1},i_{2},..,i_{a}\\}\setminus\\{i_{x},i_{y}\\}}=A_{i_{1},i_{2},..,i_{a}}\delta^{i_{x},i_{y}},$

where $\delta$ is the Kronecker delta. In general, the contraction of $A$
along dimensions $\\{x_{1},..,x_{p}\\}$ yields a tensor
$C={A_{\downarrow}}_{x_{1},..,x_{p}}\in\mathbb{R}^{n^{a-p}}$ where

${A_{\downarrow}}_{x_{1},..,x_{p}}=\sum_{i_{x_{1}}}\sum_{i_{x_{2}}}...\sum_{i_{x_{p}}}A_{i_{1},i_{2},..,i_{a}},$

or compactly as

${A_{\downarrow}}_{x_{1},..,x_{p}}=A_{i_{1},i_{2},..,i_{a}}\delta^{i_{x_{1}},i_{x_{2}},..,i_{x_{p}}}.$

Based on these tensor contractions and Def. 4.1, we can construct the second-
order $\mathbb{S}_{n}$-equivariant networks encoding a graph with adjacency
matrix $\mathcal{A}\in\mathbb{R}^{n\times n}$, node features
$\mathcal{F}_{v}\in\mathbb{R}^{n\times d_{v}}$ and edge features
$\mathbb{R}^{n\times n\times d_{e}}$ as in Section 4.2.1:

$f=\gamma\circ\mathcal{M}_{T}\circ..\circ\gamma\circ\mathcal{M}_{1}.$

The “raw” graph wavelets can be treated as a second-order tensor of size
$n\times n\times k$ where $k$ is the number of scales, similarly as the edge
features. We employ the higher-order permutation-equivariant message passing
proposed by (Maron et al., 2019), (Hy et al., 2018) and (Kondor et al., 2018)
to encode the “raw” graph wavelets from size $n\times n\times k$ into $n\times
k$ that will be further used as nodes/tokens’ embeddings of our Transformer
architecture (see Fig. 1).

####  4.2.1 Higher-order message passing

The second order message passing has the message
$\mathcal{H}_{0}\in\mathbb{R}^{|\mathcal{V}|\times|\mathcal{V}|\times(d_{v}+d_{e})}$
initialized by promoting the node features $\mathcal{F}_{v}$ to a second order
tensor (e.g., we treat node features as self-loop edge features), and
concatenating with the edge features $\mathcal{F}_{e}$. Iteratively,

$\mathcal{H}_{t}=\gamma(\mathcal{M}_{t}),\ \ \ \
\mathcal{M}_{t}=\mathcal{W}_{t}\bigg{[}\bigoplus_{i,j}(\mathcal{A}\otimes{\mathcal{H}_{t-1})_{\downarrow}}_{i,j}\bigg{]},$

where $\mathcal{A}\otimes\mathcal{H}_{t-1}$ results in a fourth order tensor
while $\downarrow_{i,j}$ contracts it down to a second order tensor along the
$i$-th and $j$-th dimensions, $\oplus$ denotes concatenation along the feature
channels, and $\mathcal{W}_{t}$ denotes a multilayer perceptron on the feature
channels. We remark that the popular MPNNs (Gilmer et al., 2017a) is a lower-
order one and a special case in which
$\mathcal{M}_{t}=\mathcal{D}^{-1}\mathcal{A}\mathcal{H}_{t-1}\mathcal{W}_{t-1}$
where $\mathcal{D}_{ii}=\sum_{j}\mathcal{A}_{ij}$ is the diagonal matrix of
node degrees. The message $\mathcal{H}_{T}$ of the last iteration is still
second order, so we contract it down to the first order latent
$\mathcal{Z}=\bigoplus_{i}{{\mathcal{H}_{T}}_{\downarrow}}_{i}$.

##  5 Multiresolution Graph Transformers

In this section, we present Multiresolution Graph Transformers (MGT), a neural
network architecture for learning hierarchical structures. MGT uses
Transformers to yield the representations of macromolecules at different
resolutions. While previous work either neglects the hierarchical
characteristics of large molecules or fails to model global interactions
between distant atoms, our proposed approach can satisfy these two properties
via multiresolution analysis.

Figs. 1b and 1c show an overview of our framework. MGT consists of three main
components: an atom-level encoder, a module to extract substructures, and a
substructure-level encoder. We use a graph transformer to generate the atomic
embeddings. Then, substructures present in molecules are extracted by a
learning-to-cluster algorithm. The molecular graph is coarsened into a set of
substructures, and we use a pure Transformer encoder to learn their relations.

###  5.1 Atom-Level Encoder

To utilize the proposed wavelet positional encoding demonstrated in Section 4,
we leverage the design of the graph transformer proposed in (Rampášek et al.,
2022), which is a general, powerful, and scalable Graph Transformer (GraphGPS)
for graph representation learning. Let $\mathbf{A}\in\mathbb{R}^{n\times n}$
be the adjacency matrix of a graph with $n$ nodes and $e$ edges;
$\textbf{X}^{l}$ and $\textbf{E}^{l}$ are node and edge features at layer
$l$-th, respectively. In addition, $\textbf{X}^{0}\in\mathbb{R}^{n\times d}$
and $\textbf{E}^{0}\in\mathbb{R}^{e\times d}$ are initial atom and bond
features embedded in $d$-dimensional spaces created by two embedding layers.
The wavelet positional vectors $\textbf{p}\in\mathbb{R}^{n\times k}$ are fed
to an encoder (e.g., a feed-forward neural network or a linear
transformation), yielding a tensor of positional features
$\textbf{P}\in\mathbb{R}^{n\times d_{p}}$. We let
$\textbf{X}^{0}:=\text{concat}(\textbf{X}^{0},\textbf{P})$ to produce new node
features $\textbf{X}^{0}\in\mathbb{R}^{n\times(d+d_{p})}$. From here, we
define $d:=d+d_{p}$, and for convenience, the output dimensions of all layers
are equal to $d$.

Each layer of GraphGPS uses a message-passing neural network
($\text{MPNN}^{l}$) to exchange information (i.e., messages) within the
neighborhood and a self-attention layer ($\text{SA}^{l}$) described in Eq. (1)
to compute global interactions among distant nodes:

$\displaystyle\textbf{X}_{L}^{l+1},\textbf{E}^{l+1}$
$\displaystyle=\text{MPNN}^{l}(\textbf{X}^{l},\textbf{E}^{l},\textbf{A})$ (2)
$\displaystyle\textbf{X}_{G}^{l+1}$
$\displaystyle=\text{SA}^{l}(\textbf{X}^{l})$ (3)
$\displaystyle\textbf{X}^{l+1}$
$\displaystyle=\text{FFN}^{l}(\textbf{X}_{L}^{l+1}+\textbf{X}_{G}^{l+1})$ (4)

where $\textbf{X}_{L}^{l+1}$ and $\textbf{X}_{G}^{l+1}$ are node local and
global representations; they are unified into $\textbf{X}^{l+1}$ via Eq. 4.
Popular techniques such as Dropout (Srivastava et al., 2014) and normalization
(Ioffe & Szegedy, 2015; Ba et al., 2016) are omitted for the sake of clarity.
By feeding the molecular graph through $L$ layers, we attain two tensors
$\textbf{X}_{a}:=\textbf{X}^{L}$ and $\textbf{E}_{a}:=\textbf{E}^{L}$
indicating the node and edge embeddings, respectively.

###  5.2 Learning to Cluster

In this work, we use a message-passing neural network augmented with
differentiable pooling layers (Ying et al., 2018) to cluster the atoms into
substructures automatically:

Z
$\displaystyle=\text{MPNN}_{\text{e}}(\textbf{X}_{a},\textbf{E}_{a},\textbf{A})$
(5) S
$\displaystyle=\text{Softmax}(\text{MPNN}_{\text{c}}(\textbf{X}_{a},\textbf{E}_{a},\textbf{A}))$
(6)

where $\text{MPNN}_{\text{e}}$ and $\text{MPNN}_{\text{c}}$ are two-layer
message-passing networks that learn to generate node embeddings
($\textbf{Z}\in\mathbb{R}^{n\times d}$) and a clustering matrix
($\textbf{S}\in\mathbb{R}^{n\times C})$, respectively; $C$ denotes the number
of substructures in molecules. A tensor of features
$\textbf{X}_{s}\in\mathbb{R}^{C\times d}$ for the substructures is computed:

$\textbf{X}_{s}=\textbf{S}^{T}\textbf{Z}$ (7)

This learning-to-cluster module is placed after the atom-level encoder.
Intuitively, atom nodes updated with both local and global information should
be classified into accurate substructures.

###  5.3 Substructure-level Encoder

Given a set of substructures $\mathcal{V}_{s}$ with a tensor of features
$\mathbf{X}_{s}\in\mathbb{R}^{C\times d}$, we forward $\mathbf{X}_{s}$ to $L$
conventional Transformer encoder layers (Vaswani et al., 2017) to capture
their pairwise semantics:

$\displaystyle\textbf{H}_{\text{1}}^{l+1}$
$\displaystyle=\text{Norm}(\text{SA}^{l}(\textbf{H}^{l})+\textbf{H}^{l})$ (8)
$\displaystyle\textbf{H}^{l+1}$
$\displaystyle=\text{Norm}(\text{FFN}(\textbf{H}_{\text{1}}^{l+1})+\textbf{H}_{\text{1}}^{l+1})$
(9)

where SA refers to (multi-head) self-attention described in Eq. (1), and
$\textbf{H}^{0}$ is equal to $\textbf{X}_{s}$. Additionally, we add a long-
range skip connection to alleviate gradient vanishing as:

$\textbf{H}_{s}=\text{FFN}(\text{concat}(\textbf{H}^{0},\textbf{H}^{L}))$ (10)

$\textbf{H}_{s}\in\mathbb{R}^{C\times d}$ is the output indicating the
representations for the substructures. Finally, we aggregate all $C$ vectors
$h_{s}\in\textbf{H}_{s}$ to result in a unique representation
$z\in\mathbb{R}^{d}$ for the molecules (refer to Section 3.2), before feeding
it to a feed-forward network to compute the final output $y\in\mathbb{R}^{c}$
for property prediction:

$\displaystyle z$ $\displaystyle=\zeta(\\{h_{s}\\}_{s=1}^{C})$ (11)
$\displaystyle\hat{y}$ $\displaystyle=\text{FFN}(z)$ (12)

##### Training Objective

We train MGT by minimizing $\mathcal{L}$:

$\mathcal{L}=\mathcal{L}_{1}+\lambda_{1}\mathcal{L}_{LP}+\lambda_{2}\mathcal{L}_{E}$
(13)

where $\mathcal{L}_{1}=l(\hat{y},y)$ denotes the loss function between
predicted values and ground truths (e.g., cross-entropy or mean-squared
error), $\mathcal{L}_{LP}=||\textbf{A}-\textbf{S}^{T}\textbf{S}||_{F}$
indicates auxiliary link prediction loss ($||\cdot||_{F}$ denotes the
Frobenius norm), and $L_{E}=\frac{1}{n}\sum_{i=1}^{n}H(\textbf{S}_{i})$
denotes the entropy regularization of the cluster assignment, i.e. each atom
should be assigned into a unique cluster. Additionally, $\lambda_{1}$ and
$\lambda_{2}$ are hyperarameters.

Table 1: Results on peptides property prediction
Model | No.Params | Peptides-struct | Peptides-func
---|---|---|---
MAE $\downarrow$ | AP $\uparrow$
GCN | 508k | 0.3496 $\pm$ 0.0013 | 0.5930 $\pm$ 0.0023
GINE | 476k | 0.3547 $\pm$ 0.0045 | 0.5498 $\pm$ 0.0079
GatedGCN | 509k | 0.3420 $\pm$ 0.0013 | 0.5864 $\pm$ 0.0077
GatedGCN + RWPE | 506k | 0.3357 $\pm$ 0.0006 | 0.6069 $\pm$ 0.0035
Transformer + LapPE | 488k | 0.2529 $\pm$ 0.0016 | 0.6326 $\pm$ 0.0126
SAN + LapPE | 493k | 0.2683 $\pm$ 0.0043 | 0.6384 $\pm$ 0.0121
SAN + RWPE | 500k | 0.2545 $\pm$ 0.0012 | 0.6562 $\pm$ 0.0075
MGT + LapPE (ours) | 499k | 0.2488 $\pm$ 0.0014 | 0.6728 $\pm$ 0.0152
MGT + RWPE (ours) | 499k | 0.2496 $\pm$ 0.0009 | 0.6709 $\pm$ 0.0083
MGT + WavePE (ours) | 499k | 0.2453 $\pm$ 0.0025 | 0.6817 $\pm$ 0.0064

##  6 Experiments

(a)

(b)

Figure 3: Examples of two macromolecules. a) An example of a peptide that
consists of many functional groups. b) An example of a polymer that consists
of many repeating units

We empirically validate our proposed approach in two types of macromolecules
including peptides and polymers. Figure 3 illustrates two examples of
macromolecules in the datasets. Our PyTorch implementation is publicly
available at https://github.com/HySonLab/Multires-Graph-Transformer.

###  6.1 Peptides Property Prediction

Peptides are small chains of amino acids found in nature that serve a variety
of critical biological roles (Singh et al., 2016); however, they are far
shorter than proteins. Because each amino acid is made up of several heavy
atoms, a peptide’s molecular graph is significantly greater than that of a
tiny drug-like molecule. Since peptides are formed by sequences of amino
acids, they are naturally hierarchical and long-range dependencies (Dwivedi et
al., 2022b), i.e. a peptide should be ideally segmented into an exact set of
amino acids. Therefore, we evaluate our method on peptide structures to
demonstrate its superiority.

##### Experimental Setup

We run experiments on two real-world datasets including (1) Peptides-struct
and (2) Peptides-func (Dwivedi et al., 2022b). The two datasets are multi-
label graph classification problems and share the same peptide molecular
graphs, but with different tasks. While the former consists of 10 classes
based on peptides function, the latter is used to predict 11 aggregated 3D
properties of peptides at the graph level. For a fair comparison, we follow
the experimental and evaluation setting of (Dwivedi et al., 2022b) with the
same train/test split ratio. We use mean absolute error (MAE) and average
precision (AP) to evaluate the method’s performance for Peptides-struct and
Peptides-func, respectively.

##### Baselines and Implementation Details

We compare our proposed approach with the baselines taken from (Dwivedi et
al., 2022b). The local message-passing network class involves GCN (Kipf &
Welling, 2016), GCNII (Bresson & Laurent, 2017), GINE (Xu et al., 2019c; Hu et
al., 2020), and GatedGCN (Bresson & Laurent, 2017). For Transformer-based
architectures, we compare our method with vanilla Transformer (Vaswani et al.,
2017) with Laplacian PE (Dwivedi et al., 2020; Dwivedi & Bresson, 2020b) and
SAN (Kreuzer et al., 2021). Since all baselines are limited to approximately
500k learnable parameters, we also restrict MGT to roughly the same number of
parameters. Additionally, we use GatedGCN (Bresson & Laurent, 2017) for local
message passing customized with the PEG technique to stabilize the positional
features (Wang et al., 2022) 111The implementation of GPS is adapted from
https://github.com/vijaydwivedi75/lrgb.git.. We experiment with each task in
four different random seeds. We provide further implementation details of MGT
for this task in Appendix A.1.

##### Results

Table 1 shows that our proposed MGT + WavePE achieves the best performances in
two peptide prediction tasks. In addition to WavePE, MGT + RWPE also attains
the second-best performances. The superiority of WavePE to RWPE can be
explained as mentioned in Section 4 that WavePE is a generalized version of
RWPE. In particular, our proposed MGT outperforms all the baselines in the
Petides-func task by a large margin and decreases the MAE score to less than
0.25 in the Petides-struct task.

Table 2: Experimental results on the polymer property prediction task. All the methods are trained in four different random seeds and evaluated by MAE $\downarrow$. Our methods are able to attain better performance across three DFT properties of polymers while having less number of parameters. All the properties are measured in eV. Model | No. Params | Property
---|---|---
GAP | HOMO | LUMO
GCN + Virtual Node | 557k | 0.1130 $\pm$ 0.0009 | 0.0100 $\pm$ 0.0016 | 0.0871 $\pm$ 0.0019
GINE + Virtual Node | 557k | 0.1669 $\pm$ 0.0078 | 0.1233 $\pm$ 0.0130 | 0.0945 $\pm$ 0.0028
GPS | 600k | 0.0895 $\pm$ 0.0032 | 0.0703 $\pm$ 0.0045 | 0.0695 $\pm$ 0.0011
Transformer + LapPE | 700k | 0.5653 $\pm$ 0.0923 | 0.2621 $\pm$ 0.0450 | 0.2791 $\pm$ 0.0229
MGT + LapPE (ours) | 499k | 0.0724 $\pm$ 0.0008 | 0.0590 $\pm$ 0.0025 | 0.0533 $\pm$ 0.0006
MGT + RWPE (ours) | 499k | 0.0736 $\pm$ 0.0031 | 0.0600 $\pm$ 0.0011 | 0.0523 $\pm$ 0.0015
MGT + WavePE (ours) | 499k | 0.0741 $\pm$ 0.0023 | 0.0619 $\pm$ 0.0009 | 0.0539 $\pm$ 0.0020

###  6.2 Polymer Property Prediction

Polymers are long chains of repetitive substructures known as repeating units.
They, as a result, are also hierarchically structured and contain various
types of long-range dependencies among the atoms. Since polymers have a wide
range of applications in daily life, it is essential to understand their
molecular properties. In this work, we experiment on a dataset of polymers to
verify the superior performance of MGT in predicting molecular properties, as
well as detecting repeating units existing in polymer structures.

##### Experimental Setup

We use a polymer dataset proposed in (St. John et al., 2019). Each polymer is
associated with three types of Density Functional Theory (DFT) (Hohenberg &
Kohn, 1964) properties including the first excitation energy of the monomer
calculated with time-dependent DFT (GAP), the energy of the highest occupied
molecular orbital for the monomer (HOMO), and the lowest unoccupied molecular
orbital of the monomer (LUMO). The dataset is split into train/validation/test
subsets with a ratio of 8:1:1, respectively. For training, we normalize all
learning targets to have a mean of 0 and a standard deviation of 1.

##### Baselines and Implementation Details

As there are no existing baselines on this dataset, we perform experiments
with four different models for comparisons. For local GNNs, we use GCN (Kipf &
Welling, 2016) and GINE (Xu et al., 2019c; Hu et al., 2020) augmented with
virtual nodes as the baselines 222The implementation of local GNN models is
taken from https://github.com/snap-stanford/ogb/tree/master/ogb. Moreover, we
use standard Transformer (Vaswani et al., 2017) with Laplacian positional
encoding (Dwivedi & Bresson, 2020b) and GPS (Rampášek et al., 2022) as the
baselines for Transformer-based architectures. The implementation of MGT is
similar to the Peptide tasks. Please refer to Appendix 4 for the baseline
details.

##### Results

As shown in Table 2, our MGT models achieve the lowest MAE scores across three
properties. In addition, WavePE can attain comparable results with LapPE and
RWPE for this task. We observe that the vanilla Transformer has the poorest
performance. This demonstrates that computing global information without the
awareness of locality is not sufficient for macromolecular modeling. As
described in Section 5, MGT is an extended version of GPS. In particular, a
learning-to-cluster module and a substructure-level Transformer encoder are
extensions to GPS. The better performance of MGT, as a result, indicates that
our methodology in modeling hierarchical structures is appropriate and
reasonable.

###  6.3 Visualization

In addition to numerical results, we provide some visualizations to
demonstrate the comprehensiveness of MGT in learning to represent large
molecules. For visual quality purposes, we place the figures in the Appendix.
Figs. 4 and 5 show clear and smooth clustering patterns in low-dimensional
spaces that indicate our proposed approaches are able to learn meaningful
molecular representations for hierarchical structures such as peptides and
polymers. Furthermore, according to Figs. 6 and 7, our learning-to-cluster
algorithm and multiresolution analysis can pick up functional groups (for
proteins/peptides) and repeating units (for polymers) via back-propagation.

##  7 Conclusion

In this work, we introduce the MGT model and WavePE method, which perform
multiresolution analysis in both spatial and spectral domains. Our proposed
methods can learn to represent hierarchical and long-range structures. In
addition to numerical performance, MGT can provide interpretable clustering
results and produce discernible spaces of the representations of large
molecules.

## Acknowledgements

We would like to thank Dao Quang Huy at FPT AI, Andrew Hands at University of
Chicago, and Le Duy Dung (Andrew) at VinUniversity for their valuable
suggestions.

## References

  * Alon & Yahav (2021) Alon, U. and Yahav, E.  On the bottleneck of graph neural networks and its practical implications.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=i80OPhOCVH2.
  * Anand et al. (2022) Anand, G., Ghosh, S., Zhang, L., Anupam, A., Freeman, C. L., Ortner, C., Eisenbach, M., and Kermode, J. R.  Exploiting machine learning in multiscale modelling of materials.  _Journal of The Institution of Engineers (India): Series D_ , Nov 2022\.  ISSN 2250-2130.  doi: 10.1007/s40033-022-00424-z.  URL https://doi.org/10.1007/s40033-022-00424-z.
  * Ba et al. (2016) Ba, L. J., Kiros, J. R., and Hinton, G. E.  Layer normalization.  _CoRR_ , abs/1607.06450, 2016.  URL http://arxiv.org/abs/1607.06450.
  * Bresson & Laurent (2017) Bresson, X. and Laurent, T.  Residual gated graph convnets.  _CoRR_ , abs/1711.07553, 2017.  URL http://arxiv.org/abs/1711.07553.
  * Chen et al. (2020) Chen, D., Lin, Y., Li, W., Li, P., Zhou, J., and Sun, X.  Measuring and relieving the over-smoothing problem for graph neural networks from the topological view.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 34(04):3438–3445, Apr. 2020.  doi: 10.1609/aaai.v34i04.5747.  URL https://ojs.aaai.org/index.php/AAAI/article/view/5747.
  * Chen et al. (2022) Chen, D., O’Bray, L., and Borgwardt, K.  Structure-aware transformer for graph representation learning.  In _Proceedings of the 39th International Conference on Machine Learning (ICML)_ , Proceedings of Machine Learning Research, 2022.
  * Chen et al. (2021) Chen, F., Park, J., and Park, J.  A hypergraph convolutional neural network for molecular properties prediction using functional group.  _CoRR_ , abs/2106.01028, 2021.  URL https://arxiv.org/abs/2106.01028.
  * Coifman & Maggioni (2006) Coifman, R. R. and Maggioni, M.  Diffusion wavelets.  _Applied and Computational Harmonic Analysis_ , 21(1):53–94, 2006.  ISSN 1063-5203.  doi: https://doi.org/10.1016/j.acha.2006.04.004.  URL https://www.sciencedirect.com/science/article/pii/S106352030600056X.  Special Issue: Diffusion Maps and Wavelets.
  * Corso et al. (2020a) Corso, G., Cavalleri, L., Beaini, D., Liò, P., and Velickovic, P.  Principal neighbourhood aggregation for graph nets.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020a. Curran Associates Inc.  ISBN 9781713829546.
  * Corso et al. (2020b) Corso, G., Cavalleri, L., Beaini, D., Liò, P., and Veličković, P.  Principal neighbourhood aggregation for graph nets.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 13260–13271. Curran Associates, Inc., 2020b.  URL https://proceedings.neurips.cc/paper/2020/file/99cad265a1768cc2dd013f0e740300ae-Paper.pdf.
  * Daubechies (1988) Daubechies, I.  Orthonormal bases of compactly supported wavelets.  _Communications on Pure and Applied Mathematics_ , 41:909–996, 1988.
  * Depta et al. (2022) Depta, P. N., Dosta, M., Wenzel, W., Kozlowska, M., and Heinrich, S.  Hierarchical coarse-grained strategy for macromolecular self-assembly: Application to hepatitis b virus-like particles.  _International Journal of Molecular Sciences_ , 23(23), 2022\.  ISSN 1422-0067.  doi: 10.3390/ijms232314699.  URL https://www.mdpi.com/1422-0067/23/23/14699.
  * Devlin et al. (2019) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.  URL https://aclanthology.org/N19-1423.
  * Donnat et al. (2018) Donnat, C., Zitnik, M., Hallac, D., and Leskovec, J.  Learning structural node embeddings via diffusion wavelets.  In _Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery $\&$ Data Mining_, KDD ’18, pp. 1320–1329, New York, NY, USA, 2018. Association for Computing Machinery.  ISBN 9781450355520.  doi: 10.1145/3219819.3220025.  URL https://doi.org/10.1145/3219819.3220025.
  * Dosovitskiy et al. (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Dwivedi & Bresson (2020a) Dwivedi, V. P. and Bresson, X.  A generalization of transformer networks to graphs.  _CoRR_ , abs/2012.09699, 2020a.  URL https://arxiv.org/abs/2012.09699.
  * Dwivedi & Bresson (2020b) Dwivedi, V. P. and Bresson, X.  A generalization of transformer networks to graphs.  _CoRR_ , abs/2012.09699, 2020b.  URL https://arxiv.org/abs/2012.09699.
  * Dwivedi et al. (2020) Dwivedi, V. P., Joshi, C. K., Laurent, T., Bengio, Y., and Bresson, X.  Benchmarking graph neural networks.  _CoRR_ , abs/2003.00982, 2020.  URL https://arxiv.org/abs/2003.00982.
  * Dwivedi et al. (2022a) Dwivedi, V. P., Luu, A. T., Laurent, T., Bengio, Y., and Bresson, X.  Graph neural networks with learnable structural and positional representations.  In _International Conference on Learning Representations_ , 2022a.  URL https://openreview.net/forum?id=wTTjnvGphYj.
  * Dwivedi et al. (2022b) Dwivedi, V. P., Rampasek, L., Galkin, M., Parviz, A., Wolf, G., Luu, A. T., and Beaini, D.  Long range graph benchmark.  In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_ , 2022b.  URL https://openreview.net/forum?id=in7XC5RcjEn.
  * Fang et al. (2021) Fang, Y., Yang, H., Zhuang, X., Shao, X., Fan, X., and Chen, H.  Knowledge-aware contrastive molecular graph learning.  _CoRR_ , abs/2103.13047, 2021.  URL https://arxiv.org/abs/2103.13047.
  * Gaul & Cuesta-Lopez (2022) Gaul, C. and Cuesta-Lopez, S.  Machine learning for screening large organic molecules.  _arXiv preprint arXiv:2211.15415_ , 2022.
  * Gilmer et al. (2017a) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _Proceedings of the 34th International Conference on Machine Learning - Volume 70_ , ICML’17, pp. 1263–1272. JMLR.org, 2017a.
  * Gilmer et al. (2017b) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In Precup, D. and Teh, Y. W. (eds.), _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 1263–1272. PMLR, 06–11 Aug 2017b.  URL https://proceedings.mlr.press/v70/gilmer17a.html.
  * Haar (1910) Haar, A.  Zur theorie der orthogonalen funktionensysteme.  _Mathematische Annalen_ , 69:331–371, 1910.
  * Hammond et al. (2011) Hammond, D. K., Vandergheynst, P., and Gribonval, R.  Wavelets on graphs via spectral graph theory.  _Applied and Computational Harmonic Analysis_ , 30(2):129–150, 2011.  ISSN 1063-5203.  doi: https://doi.org/10.1016/j.acha.2010.04.005.  URL https://www.sciencedirect.com/science/article/pii/S1063520310000552.
  * Hohenberg & Kohn (1964) Hohenberg, P. and Kohn, W.  Inhomogeneous electron gas.  _Phys. Rev._ , 136:B864–B871, Nov 1964.  doi: 10.1103/PhysRev.136.B864.  URL https://link.aps.org/doi/10.1103/PhysRev.136.B864.
  * Hu et al. (2020) Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovec, J.  Strategies for pre-training graph neural networks.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=HJlWWJSFDH.
  * Hu et al. (2021) Hu, W., Fey, M., Ren, H., Nakata, M., Dong, Y., and Leskovec, J.  Ogb-lsc: A large-scale challenge for machine learning on graphs.  _arXiv preprint arXiv:2103.09430_ , 2021.
  * Hy et al. (2018) Hy, T. S., Trivedi, S., Pan, H., Anderson, B. M., , and Kondor, R.  Predicting molecular properties with covariant compositional networks.  _The Journal of Chemical Physics_ , 148, 2018.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  _CoRR_ , abs/1502.03167, 2015.  URL http://arxiv.org/abs/1502.03167.
  * Jerry (1992) Jerry, M.  Advanced organic chemistry: reactions, mechanisms and structure, 1992\.
  * Jin et al. (2020) Jin, W., Barzilay, R., and Jaakkola, T.  Hierarchical generation of molecular graphs using structural motifs.  In _Proceedings of the 37th International Conference on Machine Learning_ , ICML’20. JMLR.org, 2020.
  * Kim et al. (2022) Kim, J., Nguyen, D. T., Min, S., Cho, S., Lee, M., Lee, H., and Hong, S.  Pure transformers are powerful graph learners.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=um2BxfgkT2_.
  * Kipf & Welling (2016) Kipf, T. N. and Welling, M.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ , 2016.
  * Kondor et al. (2018) Kondor, R., Hy, T. S., Pan, H., Trivedi, S., and Anderson, B. M.  Covariant compositional networks for learning graphs.  _Proc. ICLR Workshop_ , 2018.  URL https://openreview.net/forum?id=S1TgE7WR-.
  * Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W., Létourneau, V., and Tossou, P.  Rethinking graph transformers with spectral attention.  _Advances in Neural Information Processing Systems_ , 34:21618–21629, 2021.
  * Li et al. (2020a) Li, G., Xiong, C., Thabet, A. K., and Ghanem, B.  Deepergcn: All you need to train deeper gcns.  _CoRR_ , abs/2006.07739, 2020a.  URL https://arxiv.org/abs/2006.07739.
  * Li et al. (2020b) Li, P., Wang, Y., Wang, H., and Leskovec, J.  Distance encoding: Design provably more powerful neural networks for graph representation learning.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020b. Curran Associates Inc.  ISBN 9781713829546.
  * Li et al. (2018) Li, Q., Han, Z., and Wu, X.-M.  Deeper insights into graph convolutional networks for semi-supervised learning.  In _Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence_ , AAAI’18/IAAI’18/EAAI’18. AAAI Press, 2018.  ISBN 978-1-57735-800-8.
  * Liu et al. (2021) Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 2021.
  * Mallat (1989) Mallat, S.  A theory for multiresolution signal decomposition: the wavelet representation.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 11(7):674–693, 1989.  doi: 10.1109/34.192463.
  * Mallat (2008) Mallat, S.  _A Wavelet Tour of Signal Processing, Third Edition: The Sparse Way_.  Academic Press, Inc., USA, 3rd edition, 2008.  ISBN 0123743702.
  * Maron et al. (2019) Maron, H., Ben-Hamu, H., Shamir, N., and Lipman, Y.  Invariant and equivariant graph networks.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Syx72jC9tm.
  * Maziarz et al. (2022) Maziarz, K., Jackson-Flux, H. R., Cameron, P., Sirockin, F., Schneider, N., Stiefl, N., Segler, M., and Brockschmidt, M.  Learning to extend molecular scaffolds with structural motifs.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=ZTsoE8G3GG.
  * Morris et al. (2019) Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, J. E., Rattan, G., and Grohe, M.  Weisfeiler and leman go neural: Higher-order graph neural networks.  In _Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence_ , AAAI’19/IAAI’19/EAAI’19. AAAI Press, 2019\.  ISBN 978-1-57735-809-1.  doi: 10.1609/aaai.v33i01.33014602.  URL https://doi.org/10.1609/aaai.v33i01.33014602.
  * Oono & Suzuki (2020) Oono, K. and Suzuki, T.  Graph neural networks exponentially lose expressive power for node classification.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1ldO2EFPr.
  * Ramakrishnan et al. (2014) Ramakrishnan, R., Dral, P., Rupp, M., and von Lilienfeld, A.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific Data_ , 1, 08 2014.  doi: 10.1038/sdata.2014.22.
  * Rampášek et al. (2022) Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D.  Recipe for a General, Powerful, Scalable Graph Transformer.  _arXiv:2205.12454_ , 2022.
  * Rong et al. (2020) Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., and Huang, J.  Self-supervised graph transformer on large-scale molecular data.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.  ISBN 9781713829546.
  * Ruddigkeit et al. (2012) Ruddigkeit, L., Deursen, R., Blum, L., and Reymond, J.-L.  Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17.  _Journal of chemical information and modeling_ , 52, 10 2012.  doi: 10.1021/ci300415d.
  * Rustamov & Guibas (2013) Rustamov, R. M. and Guibas, L.  Wavelets on graphs via deep learning.  In _Proceedings of the 26th International Conference on Neural Information Processing Systems - Volume 1_ , NIPS’13, pp. 998–1006, Red Hook, NY, USA, 2013. Curran Associates Inc.
  * Schmid (2022) Schmid, F.  Understanding and modeling polymers: The challenge of multiple scales.  _ACS Polymers Au_ , 0(0):null, 2022.  doi: 10.1021/acspolymersau.2c00049.  URL https://doi.org/10.1021/acspolymersau.2c00049.
  * Singh et al. (2016) Singh, S., Chaudhary, K., Dhanda, S. K., Bhalla, S., Usmani, S. S., Gautam, A., Tuknait, A., Agrawal, P., Mathur, D., and Raghava, G. P.  Satpdb: a database of structurally annotated therapeutic peptides.  _Nucleic acids research_ , 44(D1):D1119–D1126, 2016.
  * Srivastava et al. (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R.  Dropout: A simple way to prevent neural networks from overfitting.  _Journal of Machine Learning Research_ , 15(56):1929–1958, 2014.  URL http://jmlr.org/papers/v15/srivastava14a.html.
  * St. John et al. (2019) St. John, P. C., Phillips, C., Kemper, T. W., Wilson, A. N., Guan, Y., Crowley, M. F., Nimlos, M. R., and Larsen, R. E.  Message-passing neural networks for high-throughput polymer screening.  _The Journal of chemical physics_ , 150(23):234111, 2019.
  * Sterling & Irwin (2015) Sterling, T. and Irwin, J.  Zinc 15 - ligand discovery for everyone.  _Journal of chemical information and modeling_ , 55, 10 2015.  doi: 10.1021/acs.jcim.5b00559.
  * van der Maaten & Hinton (2008) van der Maaten, L. and Hinton, G.  Visualizing data using t-sne.  _Journal of Machine Learning Research_ , 9(86):2579–2605, 2008.  URL http://jmlr.org/papers/v9/vandermaaten08a.html.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I.  Attention is all you need.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , NIPS’17, pp. 6000–6010, Red Hook, NY, USA, 2017. Curran Associates Inc.  ISBN 9781510860964.
  * Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y.  Graph Attention Networks.  _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=rJXMpikCZ.
  * Wang et al. (2022) Wang, H., Yin, H., Zhang, M., and Li, P.  Equivariant and stable positional encoding for more powerful graph neural networks.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=e95i1IHcWj.
  * Xu et al. (2019a) Xu, B., Shen, H., Cao, Q., Qiu, Y., and Cheng, X.  Graph wavelet neural network.  In _International Conference on Learning Representations_ , 2019a.  URL https://openreview.net/forum?id=H1ewdiR5tQ.
  * Xu et al. (2019b) Xu, B., Shen, H., Cao, Q., Qiu, Y., and Cheng, X.  Graph wavelet neural network.  In _International Conference on Learning Representations_ , 2019b.  URL https://openreview.net/forum?id=H1ewdiR5tQ.
  * Xu et al. (2019c) Xu, K., Hu, W., Leskovec, J., and Jegelka, S.  How powerful are graph neural networks?  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net, 2019c.  URL https://openreview.net/forum?id=ryGs6iA5Km.
  * Yang et al. (2022) Yang, M., Shen, Y., Li, R., Qi, H., Zhang, Q., and Yin, B.  A new perspective on the effects of spectrum in graph neural networks.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 25261–25279. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/yang22n.html.
  * Ying et al. (2021) Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y.  Do transformers really perform badly for graph representation?  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=OeWooOxFwDa.
  * Ying et al. (2018) Ying, R., You, J., Morris, C., Ren, X., Hamilton, W. L., and Leskovec, J.  Hierarchical graph representation learning with differentiable pooling.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , NIPS’18, pp. 4805–4815, Red Hook, NY, USA, 2018. Curran Associates Inc.
  * You et al. (2019) You, J., Ying, R., and Leskovec, J.  Position-aware graph neural networks.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 7134–7143. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/you19b.html.
  * Yun et al. (2019) Yun, S., Jeong, M., Kim, R., Kang, J., and Kim, H. J.  Graph transformer networks.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/9d63484abb477c97640154d40595a3bb-Paper.pdf.

##  Appendix A Implementation Details

###  A.1 Multiresolution Graph Transformer

Table 3: The hyperparameters for MGT
Hyperparameters | Values
---|---
No. Epoch | 200
Embedding Dimension | 84
Batch size | 128
Learning rate | 0.001
Dropout | 0.25
Attention Dropout | 0.5
Diffusion Step (K) | [1, 2, 3, 4, 5]
No. Head | 4
Activation | ReLU
Normalization | Batchnorm
No. Cluster | 10
$\lambda_{1}$ | 0.001
$\lambda_{2}$ | 0.001

In this section, we elaborate on the architecture and hyperparameters used to
train and evaluate our MGT to achieve the above numerical results. Table 3
show details of the hyperparameters used for MGT in all the experiments. In
particular, we use the atom and bond encoder modules provided by OGB (Hu et
al., 2021) to attribute the molecular graph. We use two GPS layers to compute
the atom-level representations and two Transformer layers for calculating the
substructure-level representations. For learning to cluster, we use a 2-layer
message-passing network to compute Z and S mentioned in Eq. (5) (6) as
follows:

$\displaystyle\textbf{Z}_{a}^{1},\textbf{E}_{a}^{1}$
$\displaystyle=\text{GatedGCN}^{1}(\textbf{X}_{a},\textbf{E}_{a},\textbf{A})$
(14) $\displaystyle\textbf{Z}_{a}^{1}$
$\displaystyle=\text{Batchnorm}(\text{ReLU}(\textbf{Z}_{a}^{1}))$ (15)
$\displaystyle\textbf{Z}_{a}^{2},\textbf{E}_{a}^{2}$
$\displaystyle=\text{GatedGCN}^{2}(\textbf{Z}_{a}^{1},\textbf{E}_{a}^{1},\textbf{A})$
(16) $\displaystyle\textbf{Z}_{a}^{2}$
$\displaystyle=\text{Batchnorm}(\text{ReLU}(\textbf{Z}_{a}^{2}))$ (17) Z
$\displaystyle=\text{concat}(\textbf{Z}_{a}^{1},\textbf{Z}_{a}^{2})$ (18) Z
$\displaystyle=\text{FFN}(\textbf{Z})$ (19)

S is computed similarly with an auxiliary Softmax operation on the output to
produce a probabilistic clustering matrix.

###  A.2 Baselines used in Polymer Property Prediction

Table 4 shows the implementation of the baselines we used in the polymer
experiments. All the models are designed to have approximately 500 to 700k
learnable parameters. For fair comparisons, all the models are trained in 50
epochs with a learning rate of 0.001 and batch size of 128.

Table 4: The detailed settings of baselines for polymer property prediction
Model | No. Layer | Embed Dim | No. Params
---|---|---|---
GCN + Virtual Node | 5 | 156 | 557k
GINE + Virtual Node | 5 | 120 | 557k
GPS | 3 | 120 | 600k
Transformer + LapPE | 6 | 120 | 700k

##  Appendix B Visualization

We use the t-SNE algorithm (van der Maaten & Hinton, 2008) to project the
representations produced by MGT (with WavePE) of peptides and polymers of the
test datasets into two-dimensional spaces for visualization. Also, we take the
probabilistic clustering matrix $\mathbf{S}$ in Eq. (6) to visualize the
clustering results on the molecules. Specifically, we use the RDKit package
333RDKit: Open-source cheminformatics. https://www.rdkit.org to draw the
molecules.

###  B.1 Representations of Peptides

(a)  (b)

Figure 4: Visualization of low-dimensional spaces of peptides on two property
prediction tasks: Peptides-func and Peptides-struct. All the vectors are
normalized to range $[0,1]$. a) t-SNE projection of peptides taken from the
Peptides-func testing dataset. We take four random peptide functions, and each
figure corresponds to one of the properties with positive (1) and negative (0)
ground truths. b) Similarly, we plot the figures of four random peptide
properties taken from the Peptides-struct testing dataset. The spectrums
represent continuous ground truths, where lower values correspond to cooler
colors.

###  B.2 Representations of Polymers

Figure 5:  t-SNE projection of representations of the test polymers in the
dataset. We plot the figures of three properties, including GAP, HUMO, and
LUMO. As the labels are continuous values, points are color-coded by spectrums
wherein higher values correspond to warmer palettes.

###  B.3 Clustering on Peptides

Figure 6: The clustering result on a peptide. MGT can group the atoms of a
long peptide into different substructure types. Specifically, the groups NH3
and OH are recognized even though the atoms are located distantly. Also, local
rings or segments are also detected.

###  B.4 Clustering on Polymers

Figure 7: The clustering result on a polymer. By learning to cluster and using
a substructure-level Transformer encoder, MGT can model repetitive patterns
existing in polymers. In this example, the model can recognize repeating units
in a long-chain polymer or even symmetries.

##  Appendix C Experiments on Small Molecules

Although MGT is intentionally designed for learning to represent hierarchical
structures, we report its experimental results on the ZINC-12K dataset
(Sterling & Irwin, 2015), which consists of small drug-like molecules, in this
section. We train MGT to predict the solubility (LogP) of the molecules with
up to 50 heavy atoms on a subset of the ZINC dataset. We follow the split of
10K/1K/1K for training/validation/testing proposed in (Dwivedi et al., 2020).
Baseline results include GCN (Kipf & Welling, 2016), GINE (Xu et al., 2019c;
Dwivedi et al., 2020), , GAT (Veličković et al., 2018), Spec-GN (Yang et al.,
2022), PNA (Corso et al., 2020a), GatedGCN (Bresson & Laurent, 2017), GPS
(Rampášek et al., 2022), MPNN (Gilmer et al., 2017a), SAN (Kreuzer et al.,
2021), DGN (Li et al., 2020a), and Graphormer (Ying et al., 2021). Our MGT +
WavePE outperforms 7 out of 10 other baselines.

Table 5: Experimental results on the ZINC-12K dataset
Method | No. Params | MAE $\downarrow$
---|---|---
GCN | 505k | 0.367 $\pm$ 0.011
GINE | 510k | 0.526 $\pm$ 0.051
GAT | 531k | 0.384 $\pm$ 0.007
PNA | 387k | 0.142 $\pm$ 0.010
MPNN | 418k | 0.145 $\pm$ 0.007
GatedGCN | 505k | 0.214 $\pm$ 0.006
SAN | 509k | 0.139 $\pm$ 0.006
Graphormer | 489k | 0.122 $\pm$ 0.006
GPS | - | 0.070 $\pm$ 0.004
Spec-GN | 503k | 0.0698 $\pm$ 0.002
MGT + WavePE (ours) | 499k | 0.131 $\pm$ 0.003