# Graph Rationalization with Environment-based Augmentations

Gang Liu, Tong Zhao, Jiaxin Xu, Tengfei Luo, Meng Jiang  University of Notre
Dame, IN USA gliu7,tzhao2,jxu24,tluo<EMAIL_ADDRESS>
(2022)

###### Abstract.

Rationale is defined as a subset of input features that best explains or
supports the prediction by machine learning models. Rationale identification
has improved the generalizability and interpretability of neural networks on
vision and language data. In graph applications such as molecule and polymer
property prediction, identifying representative subgraph structures named as
graph rationales plays an essential role in the performance of graph neural
networks. Existing graph pooling and/or distribution intervention methods
suffer from lack of examples to learn to identify optimal graph rationales. In
this work, we introduce a new augmentation operation called _environment
replacement_ that automatically creates virtual data examples to improve
rationale identification. We propose an efficient framework that performs
rationale-environment separation and representation learning on the real and
augmented examples in _latent spaces_ to avoid the high complexity of explicit
graph decoding and encoding. Comparing against recent techniques, experiments
on seven molecular and four polymer real datasets demonstrate the
effectiveness and efficiency of the proposed augmentation-based graph
rationalization framework.

Graph Learning, Graph Neural Network, Molecule Property, Data Augmentation,
Rationalization

††journalyear: 2022††copyright: acmcopyright††conference: Proceedings of the
28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining; August
14–18, 2022; Washington, DC, USA††booktitle: Proceedings of the 28th ACM
SIGKDD Conference on Knowledge Discovery and Data Mining (KDD ’22), August
14–18, 2022, Washington, DC, USA††price: 15.00††doi:
10.1145/3534678.3539347††isbn: 978-1-4503-9385-0/22/08††ccs: Applied computing
Chemistry††ccs: Computing methodologies Learning latent representations

##  1\. Introduction

Figure 1. Graph rationalization identifies a rationale subgraph that best
explains or supports the prediction of graph property. Our work makes the
first attempt to improve graph rationalization by graph data augmentations
with _environment subgraphs_ which are the remaining parts after rationale
identification. It proposes new augmentation operations, designs and develops
a novel graph rationalization framework, and conducts experiments on a large
set of molecule and polymer data.

Graph property prediction has attracted attention in different research fields
like chemoinformatics and bioinformatics where small molecules are represented
as labelled graphs of atoms (Hu et al., 2020; Zhou et al., 2020; Guo et al.,
2021). Besides, materials informatics for _polymers_ has emerged in recent
years from property prediction to inverse design (Kim et al., 2018; Chen et
al., 2021). Polymer are materials consisting of macromolecules, composed of
many repeating units. They are ubiquitous in applications ranging from plastic
cups and electronics to aerospace structures. New engineering and
environmental challenges demand that polymers possess unconventional
properties such as high-temperature stability, excellent thermal conductivity,
and biodegradability (Ma et al., 2019; Wei et al., 2021). It’s important to
integrate data science and machine learning into polymer informatics on the
tasks of graph classification and regression.

To automate feature extraction from graph data, graph neural network (GNN)
models learn node representations through nonlinear functions and layers that
aggregate information from node neighborhood (Kipf and Welling, 2017;
Veličković et al., 2018; Hamilton et al., 2017; Zhang et al., 2020; Wu et al.,
2020). Graph pooling is a central component of the GNN architecture that
learns a cluster assignment for nodes and passes cluster nodes and their
representations to the next layer (Ying et al., 2018; Lee et al., 2019). The
final layer returns the representations of entire graphs. Despite the advances
of various GNN models, the limitation of data size makes them easily fall into
_over-fitting and poor generalizability_. For example, the number of graphs in
molecule benchmark datasets is usually in the range of 1,000 and 10,000; and
the size of polymer datasets is even smaller (e.g., $\sim$600) (Ma and Luo,
2020).

Rationalization techniques have been designed to solve the above problem in
vision and language data, where the rationale is defined as a subset of input
features that best explains or supports the prediction by machine learning
models (Chang et al., 2020; Arjovsky et al., 2019; Rosenfeld et al., 2021).
However, graph rationalization has not been extensively studied, which aims at
identifying representative subgraph structures for accurate and interpretable
graph property prediction. Related work mainly focused on advancing graph
pooling methods, but cluster assignment could not reflect the most essential
part that led to accurate prediction (Mesquita et al., 2020; Gao and Ji,
2021). A very recent technique named DIR (Wu et al., 2022) employed two GNN
modules to discover invariant graph rationales: one module separates each
input graph into a rationale subgraph and an environment subgraph; the other
is a graph property predictor based on the rationale subgraph. As shown at the
top in Figure 1, given graph $g_{i}$, the separator $f_{sep}$ identifies
rationale $g^{(r)}_{i}$, and the predictor $f_{pred}$ gives label
$\hat{y}^{(r)}_{i}$ based on the rationale. DIR conducted interventions on
training distribution to improve the invariance. Unfortunately, when the data
size was small, $f_{sep}$ could hardly find good rationales, as reported in
our later experiments.

In this work, we make the first attempt to enhance graph rationalization by
graph data augmentations. Existing augmentation methods were mainly heuristic
modification of graph structure, which could not directly support the
identification of graph rationales (Rong et al., 2019; Wang et al., 2020b, c;
Zhao et al., 2021b). We present two augmentation methods based on _environment
subgraphs_ that are the remaining parts in the graph after rationale
identification. First, rationales are used to train the property predictor,
which can be considered as graph examples augmented by _environment removal_.
Second, we replace the environment of input graph with the environment of
another graph in the batch: to generate an augmented example: this
augmentation method is called _environment replacement_. The idea is that the
rationale can be accurately identified and/or separated from the input graph
when the augmented examples are expected to have the same label of the input
graph example.

Figure 1 presents the idea of generating virtual data for small datasets via
data augmentations. Suppose we have rationale $g^{(r)}_{i}$ separated from
input graph $g_{i}$. We use the same GNN-based separator to find environment
subgraph $g^{(e)}_{j}$ from another graph $g_{j}$ in the batch. The example
augmented by environment replacement is denoted by $g_{(i,j)}=g^{(r)}_{i}\cup
g^{(e)}_{j}$. The model is trained on this example to predict label
$\hat{y}_{i,j}$ to be the same as $y_{i}$ that is the observed label of
$g_{i}$. We compute two losses on the augmented examples, $\mathcal{L}_{rem}$
and $\mathcal{L}_{rep}$ (“rem” for removal and “rep” for replacement), and
jointly optimize $f_{sep}$ and $f_{pred}$ by their combination.

The key challenge in the implementation of this idea is the high computational
complexity of decoding for _explicit graph forms_ of rationales, environment
subgraphs, and augmented examples, as well as encoding them for representation
learning and property prediction. Moreover, it is scientifically and
technically difficult to explicitly combine rationale $g^{(r)}_{i}$ and
environment $g^{(e)}_{j}$ from different graphs, as shown in the three
augmented examples $g_{(i,j)}$ in Figure 1. To address these challenges, our
hypothesis is that the _contextualized representations of (atom) nodes_ play a
significant role in rationales, environment subgraphs, and augmented graphs.
As a result, we can create the representations of all these objects from _one
latent space_ that is created from the input graphs.

In this paper, we propose a novel, efficient framework of Graph
Rationalization enhanced by Environment-based Augmentations (GREA). It
performs rationale-environment separation and representation learning on the
real and augmented examples in one latent space to avoid the high complexity
of explicit subgraph decoding and encoding. Figure 2 presents the architecture
of GREA with a few steps. First, it employs $\text{GNN}_{1}$ and
$\text{MLP}_{1}$ models to infer the probability of nodes being classified
into rationale subgraph $\mathbf{m}$. Second, it employs $\text{GNN}_{2}$ to
create contextualized node representations $\mathbf{H}$. Then, it _directly_
creates the representation vectors of rationales, environment subgraphs and
environment-replaced examples, denoted by $\mathbf{h}^{(r)}_{i}$,
$\mathbf{h}^{(e)}_{i}$, and $\mathbf{h}_{(i,j)}$, respectively. Note that DIR
(Wu et al., 2022) used a GNN to generate a matrix of masks that indicate the
importance of edges and then select the top-$K$ edges with the highest masks
to construct the rationale. Then it had to run GNNs on all the explicit graph
objects. Instead, our GREA uses $\mathbf{m}$ and $\mathbf{H}$ to compute the
representation vectors of the artificial graphs.

We conduct experiments on seven molecule and four polymer datasets. Results
demonstrate the advantages of GREA over baselines. For example, it
significantly reduces the prediction error on oxygen permeability of polymer
membrane with only 595 training examples. The oxygen permeability defines how
easily oxygen passes through a particular material. Accurate prediction will
speed up material discovery for healthcare and energy utilization.

The main contributions of this work are summarized below:

  * •

the first attempt to improve graph rationale identification using data
augmentations, including environment replacement, for accurate and
interpretable property prediction;

  * •

a novel and efficient framework that performs rationale-environment separation
and representation learning on real and augmented examples in one latent
space;

  * •

extensive experiments on more than ten molecule and polymer datasets to
demonstrate the effectiveness and efficiency of the proposed framework.

##  2\. Related Work

Figure 2. The architecture of the proposed graph rationalization framework: It
performs the creation and representation learning of environment-based
augmented examples in a _latent space_ , instead of decoding every example
into a graph form and running a GNN encoder on it. This design aligns graph
representation spaces and avoids high computational complexity.

There are four research topics related to the proposed work. We briefly
present their recent studies and compare with ours.

###  2.1. Graph Property Prediction

Learning representations and predicting properties of entire graphs is
important for chemistry, biology, and material sciences, where molecule and
polymer data can be structured as graphs (Hu et al., 2020). When RDKit is
widely used to generate molecular fingerprints (Landrum, 2013), graph neural
networks (GNNs) such as Graph Convolutional Network (GCN) (Kipf and Welling,
2017), Graph Attention Networks (GAT) (Veličković et al., 2018), and GraphSAGE
(Hamilton et al., 2017) have automated representation learning with nonlinear
functions from graph data (Zhang et al., 2020; Wu et al., 2020; Ma et al.,
2021; Wang et al., 2020a; Wang et al., 2021c; Wang et al., 2021a, b; Zhao et
al., 2021a; Jiang et al., 2022).

In the GNN models, graph pooling is a central component of their architectures
as a cluster assignment function to find local patches in graphs (Mesquita et
al., 2020). For example, DiffPool presented a differentiable graph pooling
module that learned a differentiable soft cluster assignment for nodes at each
layer of a deep GNN, mapped nodes to a set of clusters, and then formed the
coarsened input for the next GNN layer (Ying et al., 2018). Lee et al.
proposed self-attention graph convolution that allows graph pooling to
consider both node features and graph topology (Lee et al., 2019). Gao and Ji
proposed graph pooling and unpooling operations in Graph U-Nets (Gao and Ji,
2021). Xu et al. presented a theoretical framework for analyzing the
representational power of GNNs through the graph pooling functions (Xu et al.,
2019). While graph pooling identifies soft clusters that effectively aggregate
information from nodes (Ying et al., 2019), our work identifies representative
subgraph structures for accurate and interpretable predictions of GNN models.

###  2.2. Graph Rationalization

Most rationalization techniques identify the small subset of input features by
maximizing the predictive performance based only on the subset itself, called
rationale. To rule out spurious correlation between the input features and the
output, Chang et al. proposed the concept of invariant rationalization by
modeling different environments as non-causal input to train predictors (Chang
et al., 2020). Rosefeld et al. offered formal guarantees for improvement of
the invariant causal prediction on out-of-distribution generalization
(Arjovsky et al., 2019; Rosenfeld et al., 2021).

By introducing causal modeling into GNN optimization, Fan et al. presented a
causal representation framework for GNN models to perform on out-of-
distribution graphs (Fan et al., 2021). Li et al. proposed OOD-GNN that
employed a novel nonlinear graph representation decorrelation method that used
random Fourier features to encourage GNNs to eliminate the statistical
dependence between relevant and irrelevant graph representations (Li et al.,
2021). Very recently, Wu et al. proposed the first work called DIR to approach
causal rationales for GNNs to improve the interpretability and predictive
performance on out-of-distribution data (Wu et al., 2022). DIR conducted
interventions on the training distribution to create multiple distributions.
Unfortunately, distribution intervention might not be the optimal solution to
graph rationale identification. Also, the edge selection method suffers from
high computational complexity for rationale creation. Moreover, the studies
were mainly performed on synthetic data. In this paper, we make the first
attempt to define “environment” in graph data, augment data examples by
environment replacement, develop an efficient framework, and conduct
experiments on a large set of real molecule and polymer data. We find that
augmentation-enhanced graph rationalization is more effective than DIR.

###  2.3. Graph Data Augmentation

Graph data augmentation (GDA) techniques (Zhao et al., 2022a, b, 2021c; Chen
et al., 2020) have improved the performance on semi-supervised node
classification, such as DropEdge (Rong et al., 2019), NodeAug (Wang et al.,
2020c), and GAug (Zhao et al., 2021b). Besides, many GDA techniques have been
designed for graph-level tasks, aiming at creating new training examples by
modifying input graph data examples. For example, GraphCrop regularized GNN
models for better generalization by cropping subgraphs or motifs to simulate
real-world noise of sub-structure omission (Wang et al., 2020b). M-Evolve
presented two heuristic algorithms including random mapping and motif-
similarity mapping to generate weakly labeled data for small datasets (Zhou et
al., 2020). MH-Aug adopted the Metropolis-Hastings algorithm to create
augmented graphs from an explicit target distribution for semi-supervised
learning (Park et al., 2021). Meanwhile, graph contrastive learning learned
unsupervised representations of graphs using graph data augmentations to
incorporate various priors (You et al., 2020). Zhu et al. (2021) proposed
adaptive augmentation that incorporated various priors for topological and
semantic aspects of graphs. Specifically, it designed augmentation schemes
based on node centrality measures to highlight important connective structures
and corrupted node features by adding noise to unimportant node features. A
comprehensive survey of GDA is given by Zhao et al. (2022a).

###  2.4. Graph Learning on Polymer Data

Material informatics uses machine learning approaches to fast screen material
candidates or generate new materials meeting certain criteria, so as to reduce
the time of material development. When most related research performed on
molecule data (Guo et al., 2021), polymer researchers have developed a
benchmark database and developed machine learning techniques for polymer data,
called polymer embeddings (Kim et al., 2018; Chen et al., 2021). They can be
used to perform several polymer informatics regression tasks for density,
glass transition temperature, melting temperature, and dielectric constants
(Ma et al., 2019; Ma and Luo, 2020; Wei et al., 2021).

##  3\. Problem Definition

##### Graph Property Prediction

Let $g=(\mathcal{V},\mathcal{E})$ be a graph of $N$ nodes and $M$ edges, where
$\mathcal{V}$ is the set of nodes (e.g., atoms) and
$\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$ is the set of edges (e.g.,
bonds between atoms). We use $y\in\mathcal{Y}$ to denote the graph-level
property of $g$, where $\mathcal{Y}$ is the value space. It can have a
categorical or numerical value, corresponding to the task of classification or
regression, respectively.

A graph property predictor $f_{pred}$ takes a graph $g$ as input and predicts
its label $\hat{y}$. Specifically, a GNN-based predictor employs a GNN encoder
to generate node representations $\mathbf{H}$ from $g$:

(1) $\mathbf{H}=\begin{bmatrix}-\\\ \vec{h}_{v}\\\
-\end{bmatrix}_{v\in\mathcal{V}}=\text{GNN}(g)\in\mathbb{R}^{N\times d},$

where $\vec{h}_{v}\in\mathbb{R}^{d}$ is the representation vector of node $v$
in graph $g$. There are many choices for the GNN encoder $\text{GNN}(\cdot)$,
such as Graph Convolutional Network (GCN) (Kipf and Welling, 2017) and Graph
Isomorphism Network (GIN) (Xu et al., 2019).

Once the node representations are ready, a multilayer perceptron (MLP) can
project them into a one-dimensional space to obtain a scalar for each node:

(2) $m_{v}=\text{MLP}(\vec{h}_{v}).$

As we are more interested in graph-level classification or regression, we
first use a readout operator (e.g., average pooling) to get the graph
representation $\mathbf{h}$:

(3) $\mathbf{h}=\text{READOUT}(\mathbf{H})\in\mathbb{R}^{d},$

and then we apply a MLP to project it to a graph label:

(4) $\hat{y}=\text{MLP}(\mathbf{h})\in\mathcal{Y}.$

##### Graph Rationalization

Following the existing literature on graph rationalization (Ying et al., 2018;
Lee et al., 2019; Gao and Ji, 2021; Fan et al., 2021; Wu et al., 2022) and GNN
explanation (Ying et al., 2019), we use rationale
$g^{(r)}=(\mathcal{V}^{(r)},\mathcal{E}^{(r)})$ to indicate the causal
subgraph of the property $y$, where $g^{(r)}$ is a subgraph of $g$ such that
$\mathcal{V}^{(r)}\subseteq\mathcal{V}$ and
$\mathcal{E}^{(r)}\subseteq\mathcal{E}$. We use $g^{(e)}$ to denote the
environment subgraph, which is the complementary subgraph of $g^{(r)}$ in $g$.
In contrast with the rationale subgraph $g^{(r)}$, the environment subgraph
$g^{(e)}$ corresponds to the non-causal part of the graph data, which has no
causal relationship or only spurious correlation with the target graph
property (Chang et al., 2020; Wu et al., 2022).

Let $f_{sep}$ be a GNN-based graph rationalization model that splits an input
graph $g$ into a rationale subgraph $g^{(r)}$ and an environment subgraph
$g^{(e)}$. Existing graph rationalization methods used only the rationale
subgraph as input for property prediction (Lee et al., 2019; Ying et al.,
2018; Gao and Ji, 2021; Wu et al., 2022):

(5) $\hat{y}=\hat{y}^{(r)}=f_{pred}\big{(}g^{(r)}\big{)},$

where $f_{pred}(\cdot)=\text{MLP}(\text{READOUT}(\text{GNN}(\cdot)))$ and
$\hat{y}^{(r)}$ denotes the predicted property of the rationale subgraph
$g^{(r)}$.

Unfortunately, when suffering from lack of training examples, these methods
chose to discard environment subgraphs at the training stage. Our idea is that
the environment subgraphs can provide natural noise through data augmentation
to improve graph rationalization. In the next section, we present a novel
framework to implement this idea.

##  4\. Proposed Framework

In this section, we introduce a novel graph rationalization framework GREA.
The key idea is to augment the rationale subgraph by removing its own
environment subgraph and/or combining it with different environment subgraphs.
Figure 2 shows the overall architecture of GREA: $\text{GNN}_{1}$ and
$\text{MLP}_{1}$ first separate input graph $g$ into rationale subgraph
$g^{(r)}$ and environment subgraph $g^{(e)}$ using Eq.(1-2); $\text{GNN}_{2}$
next generates node representations $\mathbf{H}$ using Eq.(1); the rationale
subgraph’s representation $\mathbf{h}^{(r)}_{i}$ is then combined with
different environment subgraph’s representations $\mathbf{h}^{(e)}_{j}$ for
the augmented graph’s representations $\mathbf{h}_{(i,j)}$; finally, both
$\mathbf{h}^{(r)}_{i}$ and $\mathbf{h}_{(i,j)}$ are fed into $\text{MLP}_{2}$
for the prediction of $y_{i}$ during training as Eq.(4).

###  4.1. Rationale-Environment Separation

To separate input graph $g$ into rationale subgraph $g^{(r)}$ and environment
subgraph $g^{(e)}$, the rationale-environment separator consists of two
components: a GNN encoder ($\text{GNN}_{1}$) that generates latent node
representations and a MLP decoder ($\text{MLP}_{1}$) that maps the node
representations to a mask vector $\mathbf{m}\in{(0,1)}^{N}$ on the nodes in
the set $\mathcal{V}$. $m_{v}=Pr(v\in\mathcal{V}^{(r)})$ is the node-level
mask that indicates the probability of node $v\in\mathcal{V}$ being classified
into the rationale subgraph. The mask can be on either a node or an edge (Wu
et al., 2022). we choose to learn masks on the nodes to avoid the
computational complexity of edge selection. Hence, $\mathbf{m}$ can be
calculated as

(6) $\mathbf{m}=\sigma(\text{MLP}_{1}(\text{GNN}_{1}(g))),$

where $\sigma$ denotes the sigmoid function. Based on $\mathbf{m}$, we have
$(\mathbf{1}_{N}-\mathbf{m})$ that indicates the probability of nodes being
classified into the environment subgraph. $\text{GNN}_{1}$ and
$\text{MLP}_{1}$ make up the GNN-based graph rationalization model $f_{sep}$
mentioned in Section 3.

GREA uses another GNN encoder to generate contextualized node representations
$\mathbf{H}$: $\mathbf{H}=\text{GNN}_{2}(g)$. With $\mathbf{m}$ and
$\mathbf{H}$, the rationale subgraph and environment subgraph can be easily
separated in the _latent space_. Using sum pooling, we have

(7)
$\displaystyle\mathbf{h}^{(r)}=\mathbf{1}^{\top}_{N}\cdot(\mathbf{m}\times\mathbf{H}),\quad\mathbf{h}^{(e)}=\mathbf{1}^{\top}_{N}\cdot((\mathbf{1}_{N}-\mathbf{m})\times\mathbf{H}),$

where $\mathbf{1}_{N}$ denotes the $N$-size column vector with all entries as
1, and $\mathbf{h}^{(r)},\mathbf{h}^{(e)}\in\mathbb{R}^{d}$ are the
representation vectors of graph $g^{(r)}$ and $g^{(e)}$, respectively.

###  4.2. Environment-based Augmentations

Suppose $g_{1},g_{2},\dots,g_{B}$ are the input graphs in one batch for
training, where $B$ is known as batch size. The rationale-environment
separator has generated the graph representations of rationale and environment
subgraphs for each graph $g_{i}$. That is, we have
$\\{(\mathbf{h}_{1}^{(r)},\mathbf{h}_{1}^{(e)}),(\mathbf{h}_{2}^{(r)},\mathbf{h}_{2}^{(e)}),\dots,(\mathbf{h}_{B}^{(r)},\mathbf{h}_{B}^{(e)})\\}$.
We design environment-based augmentations in the latent space of graph
representations.

####  4.2.1. Environment Removal Augmentation

As graph rationalization aims to find the rationale subgraph which is regarded
as the causal factor of graph property, the rationale itself should be good
for property prediction. As in the graph pooling methods (Lee et al., 2019;
Gao and Ji, 2021) and the graph rationalization as defined in Eq. 5, the
environment removal augmentation uses the rationale subgraph only for training
the graph property predictor. That is, given the rationale subgraph
representation $\mathbf{h}_{i}^{(r)}$ of graph $g_{i}$, the predicted label is

(8) $\hat{y}_{i}^{(r)}=\text{MLP}_{2}\big{(}\mathbf{h}_{i}^{(r)}\big{)}.$

####  4.2.2. Environment Replacement Augmentation

As aforementioned in Section 3, the environment subgraphs can be viewed as
natural noises on the rationale subgraphs. Hence, in order to enhance the
model’s robustness against the noise signal brought by the environment
subgraphs, for each graph $g_{i}$, we combine its rationale subgraph
$g_{i}^{(r)}$ not only with its own environment subgraph $g_{i}^{(e)}$, but
also with all other environment subgraphs
$g_{j}^{(e)},j\in\\{1,2,\dots,B\\}\setminus\\{i\\}$ in the batch. By replacing
the environment subgraph with other environment subgraphs in the batch, the
environment replacement augmentation generates $B-1$ augmented data samples
for each graph during training. As the environment replacement happens on the
latent space, an aggregation function $\text{AGG}(\cdot,\cdot)$ is used to
combine the rationale subgraph representation $\mathbf{h}_{i}^{(r)}$ and
environment subgraph representation $\mathbf{h}_{j}^{(e)}$. The aggregation
function can be any combining/pooling functions such as concatenation, sum
pooling, and max pooling. Taking the element-wise sum pooling as an example,
the graph representation $\mathbf{h}_{(i,j)}$ of a combined graph of rationale
subgraph $g_{i}^{(r)}$ and environment subgraph $g_{j}^{(e)}$ can be
calculated as below:

(9)
$\mathbf{h}_{(i,j)}=\text{AGG}\big{(}\mathbf{h}_{i}^{(r)},\mathbf{h}_{j}^{(e)}\big{)}=\mathbf{h}_{i}^{(r)}+\mathbf{h}_{j}^{(e)}.$

For the graph representations $\mathbf{h}_{(i,j)}$ generated by the
environment replacement augmentation, the MLP property predictor is trained to
predict $y_{i}$. That is,

(10) $\hat{y}_{(i,j)}=\operatorname{MLP}_{2}\big{(}\mathbf{h}_{(i,j)}\big{)}.$

Table 1. Statistics of eleven datasets for graph property prediction: The four
top rows are polymer datasets. The prediction tasks are graph regression. The
seven bottom rows are molecule datasets. Their tasks are graph classification.

Dataset # Graphs Avg./Max # Nodes Avg./Max # Edges GlassTemp 7,174 36.7 / 166
79.3 / 362 MeltingTemp 3,651 26.9 / 102 55.4 / 212 PolyDensity 1,694 27.3 / 93
57.6 / 210 O2Perm 595 37.3 / 103 82.1 / 234 ogbg-HIV 41,127 25.5 / 222 54.9 /
502 ogbg-ToxCast 8,576 18.8 / 124 38.5 / 268 ogbg-Tox21 7,831 18.6 / 132 38.6
/ 290 ogbg-BBBP 2,039 24.1 / 132 51.9 / 290 ogbg-BACE 1,513 34.1 / 97 73.7 /
202 ogbg-ClinTox 1,477 26.2 / 136 55.8 / 286 ogbg-SIDER 1,427 33.6 / 492 70.7
/ 1010

The graph representations generated by both environment removal augmentation
and environment replacement augmentation (i.e., $\mathbf{h}_{i}^{(r)}$ and
$\mathbf{h}_{(i,j)}$) are fed into the same property predictor
$\text{MLP}_{2}$. The GNN-based property predictor $f_{pred}$ defined in
Section 3 includes $\text{MLP}_{2}$ and $\text{GNN}_{2}$ that generates the
contextualized node representation $\mathbf{H}$.

####  4.2.3. Optimization

During training, the type of loss function on the observed graph property
($y_{i}$) and predicted labels ($\hat{y}_{i}^{(r)}$ and $\hat{y}_{(i,j)}$)
depends on the type of the property label. For example, when the graph
property $y$ has binary values in the binary classification task, we use the
standard binary cross-entropy loss. When the graph property $y$ has real
values in the graph regression task, we use the mean squared error (MSE) loss.
Without loss of generality, suppose we focus on the binary classification
task. Given a batch of $B$ graphs $g_{1},g_{2},\dots,g_{B}$, the loss
functions for each graph example $g_{i}$ and its label $y_{i}$ are defined as

(11)
$\displaystyle\mathcal{L}_{rem}=y_{i}\cdot\log\hat{y}_{i}^{(r)}+\left(1-y_{i}\right)\cdot\log\big{(}1-\hat{y}_{i}^{(r)}\big{)},$
(12)
$\displaystyle\mathcal{L}_{rep}=\frac{1}{B}\sum_{j=1}^{B}\big{(}y_{i}\cdot\log\hat{y}_{(i,j)}+\left(1-y_{i}\right)\cdot\log(1-\hat{y}_{(i,j)})\big{)},$

where $\mathcal{L}_{rem}$ is the loss for the examples created by environment
removal augmentation, and $\mathcal{L}_{rep}$ is the loss for the examples
created by the environment replacement augmentation.

Moreover, the following regularization term is used to control the size of the
selected rationale subgraph:

(13)
$\mathcal{L}_{reg}=\Big{(}\frac{\mathbf{1}^{\top}_{N}\mathbf{m}}{N}-\gamma\Big{)}+\Big{(}\frac{\sum_{k:\mathbf{m}_{k}>0}1}{N}-\gamma\Big{)},$

where $\gamma\in[0,1]$ is a hyperparamter to control the expected size of the
rationale subgraph $g^{(r)}$. The node probabilities classified as the
rationale should only exist on certain nodes with the expected size.
Therefore, we use the first term to penalize the number of nodes in the
rationale when it deviates from our expectations and the second term to
encourage an uneven distribution for $\mathbf{m}$.

Table 2. Results on polymer property prediction: GREA consistently achieves
the highest R2 and smallest RMSE.

GlassTemp MeltingTemp PolyDensity O2Perm R2 $\uparrow$ RMSE $\downarrow$ R2
$\uparrow$ RMSE $\downarrow$ R2 $\uparrow$ RMSE $\downarrow$ R2 $\uparrow$
RMSE $\downarrow$ GCN (Kipf and Welling, 2017) as encoder U-NetsPool (Gao and
Ji, 2021) 0.839$\pm$0.005 44.9$\pm$0.7 0.685$\pm$0.012 63.4$\pm$1.2
0.615$\pm$0.053 0.100$\pm$0.007 0.833$\pm$0.084 865$\pm$214 SelfAttnPool (Lee
et al., 2019) 0.848$\pm$0.007 43.5$\pm$1.0 0.709$\pm$0.008 61.0$\pm$0.9
0.688$\pm$0.019 0.090$\pm$0.003 0.656$\pm$0.135 1251$\pm$266 StableGNN (Fan et
al., 2021) 0.809$\pm$0.013 48.8$\pm$1.6 0.635$\pm$0.033 70.0$\pm$4.5
0.667$\pm$0.070 0.093$\pm$0.009 0.676$\pm$0.127 1219$\pm$241 OOD-GNN (Li et
al., 2021) 0.852$\pm$0.006 43.0$\pm$0.9 0.714$\pm$0.025 60.4$\pm$2.6
0.676$\pm$0.010 0.092$\pm$0.001 0.921$\pm$0.059 576$\pm$212 IRM (Arjovsky et
al., 2019) 0.830$\pm$0.008 46.1$\pm$1.1 0.677$\pm$0.006 64.2$\pm$0.6
0.690$\pm$0.016 0.090$\pm$0.002 0.871$\pm$0.043 770$\pm$141 DIR (Wu et al.,
2022) 0.697$\pm$0.061 61.2$\pm$6.0 0.380$\pm$0.214 87.8$\pm$14.
0.656$\pm$0.036 0.094$\pm$0.005 0.135$\pm$0.068 2028$\pm$80 DIR$+$RepAug
0.800$\pm$0.006 56.5$\pm$3.2 0.520$\pm$0.101 77.8$\pm$8.2 0.671$\pm$0.033
0.092$\pm$0.005 0.915$\pm$0.031 626$\pm$115 GREA$-$RepAug 0.685$\pm$0.172
60.6$\pm$16.5 0.679$\pm$0.034 64.0$\pm$3.3 0.686$\pm$0.007 0.090$\pm$0.001
0.459$\pm$0.254 1556$\pm$395 GREA (ours) 0.855$\pm$0.003 42.6$\pm$0.5
0.716$\pm$0.016 60.2$\pm$1.6 0.717$\pm$0.023 0.086$\pm$0.003 0.941$\pm$0.018
524$\pm$91 GIN (Xu et al., 2019) as encoder U-NetsPool (Gao and Ji, 2021)
0.852$\pm$0.006 42.9$\pm$0.9 0.703$\pm$0.009 61.6$\pm$0.9 0.635$\pm$0.029
0.097$\pm$0.004 0.868$\pm$0.085 753$\pm$250 SelfAttnPool (Lee et al., 2019)
0.848$\pm$0.003 43.5$\pm$0.4 0.726$\pm$0.009 59.2$\pm$1.0 0.654$\pm$0.024
0.095$\pm$0.003 0.601$\pm$0.267 1265$\pm$546 StableGNN (Fan et al., 2021)
0.794$\pm$0.007 50.8$\pm$0.9 0.535$\pm$0.061 76.9$\pm$5.0 0.642$\pm$0.045
0.096$\pm$0.006 0.501$\pm$0.266 1487$\pm$404 OOD-GNN (Li et al., 2021)
0.862$\pm$0.007 41.6$\pm$1.1 0.721$\pm$0.006 59.7$\pm$0.6 0.666$\pm$0.025
0.093$\pm$0.003 0.917$\pm$0.029 620$\pm$109 IRM (Arjovsky et al., 2019)
0.842$\pm$0.004 44.5$\pm$0.5 0.681$\pm$0.008 63.8$\pm$0.8 0.682$\pm$0.031
0.091$\pm$0.004 0.890$\pm$0.042 709$\pm$146 DIR (Wu et al., 2022)
0.594$\pm$0.070 71.0$\pm$6.0 0.287$\pm$0.121 95.1$\pm$7.9 0.617$\pm$0.045
0.099$\pm$0.006 0.501$\pm$0.309 1446$\pm$537 DIR$+$RepAug 0.744$\pm$0.029
56.4$\pm$3.2 0.542$\pm$0.083 76.2$\pm$7.0 0.647$\pm$0.058 0.095$\pm$0.008
0.743$\pm$0.150 1054$\pm$338 GREA$-$RepAug 0.494$\pm$0.110 79.0$\pm$9.3
0.660$\pm$0.107 65.2$\pm$9.5 0.717$\pm$0.022 0.086$\pm$0.003 0.400$\pm$0.286
1623$\pm$474 GREA (ours) 0.864$\pm$0.005 41.2$\pm$0.8 0.736$\pm$0.012
58.0$\pm$1.2 0.723$\pm$0.030 0.085$\pm$0.005 0.930$\pm$0.020 569$\pm$86

Table 3. Results on molecule property prediction: GREA consistently achieves
the highest AUC ($\uparrow$).

ogbg-HIV ogbg-ToxCast ogbg-Tox21 ogbg-BBBP ogbg-BACE ogbg-ClinTox ogbg-SIDER
GCN (Kipf and Welling, 2017) as encoder U-NetsPool (Gao and Ji, 2021)
0.7527$\pm$0.0104 0.6507$\pm$0.0086 0.7492$\pm$0.0093 0.6709$\pm$0.0176
0.7757$\pm$0.0173 0.8450$\pm$0.0403 0.6181$\pm$0.0121 SelfAttnPool (Lee et
al., 2019) 0.7733$\pm$0.0187 0.6510$\pm$0.0076 0.7563$\pm$0.0080
0.6602$\pm$0.0220 0.7383$\pm$0.0541 0.8291$\pm$0.0791 0.5718$\pm$0.0219
StableGNN (Fan et al., 2021) 0.7218$\pm$0.0099 0.6520$\pm$0.0109
0.7454$\pm$0.0059 0.6552$\pm$0.0184 0.6607$\pm$0.0500 0.7681$\pm$0.0778
0.5644$\pm$0.0274 OOD-GNN (Li et al., 2021) 0.7580$\pm$0.0176
0.6613$\pm$0.0046 0.7673$\pm$0.0109 0.6795$\pm$0.0165 0.8096$\pm$0.0132
0.8874$\pm$0.0143 0.6133$\pm$0.0095 IRM (Arjovsky et al., 2019)
0.7702$\pm$0.0107 0.6599$\pm$0.0063 0.7654$\pm$0.0072 0.6892$\pm$0.0053
0.7947$\pm$0.0186 0.8819$\pm$0.0231 0.6035$\pm$0.0195 DIR (Wu et al., 2022)
0.7466$\pm$0.0093 0.5954$\pm$0.0154 0.4727$\pm$0.0129 0.6559$\pm$0.0298
0.6751$\pm$0.0323 0.6251$\pm$0.0956 0.5331$\pm$0.0216 DIR$+$RepAug
0.7494$\pm$0.0225 0.6632$\pm$0.0098 0.7437$\pm$0.0054 0.6630$\pm$0.0118
0.7677$\pm$0.0226 0.8606$\pm$0.0144 0.5934$\pm$0.0170 GREA$-$RepAug
0.7377$\pm$0.0210 0.6614$\pm$0.0048 0.7808$\pm$0.0061 0.6736$\pm$0.0077
0.7655$\pm$0.0529 0.8708$\pm$0.0514 0.6222$\pm$0.0166 GREA (ours)
0.7794$\pm$0.0065 0.6662$\pm$0.0041 0.7822$\pm$0.0093 0.6986$\pm$0.0175
0.8191$\pm$0.0240 0.8961$\pm$0.0150 0.6316$\pm$0.0151 GIN (Xu et al., 2019) as
encoder U-NetsPool (Gao and Ji, 2021) 0.7375$\pm$0.0362 0.6524$\pm$0.0126
0.7560$\pm$0.0093 0.6809$\pm$0.0163 0.8026$\pm$0.0105 0.8146$\pm$0.0703
0.5929$\pm$0.0114 SelfAttnPool (Lee et al., 2019) 0.7533$\pm$0.0247
0.6351$\pm$0.0137 0.7507$\pm$0.0110 0.6624$\pm$0.0167 0.7348$\pm$0.0194
0.7912$\pm$0.0995 0.5702$\pm$0.0137 StableGNN (Fan et al., 2021)
0.7218$\pm$0.0078 0.6485$\pm$0.0025 0.7381$\pm$0.0123 0.6695$\pm$0.0120
0.7229$\pm$0.0122 0.8559$\pm$0.0224 0.5593$\pm$0.0172 OOD-GNN (Li et al.,
2021) 0.7799$\pm$0.0078 0.6697$\pm$0.0051 0.7646$\pm$0.0038 0.6710$\pm$0.0188
0.7800$\pm$0.0228 0.8416$\pm$0.0496 0.5916$\pm$0.0169 IRM (Arjovsky et al.,
2019) 0.7817$\pm$0.0120 0.6641$\pm$0.0065 0.7542$\pm$0.0084 0.6835$\pm$0.0071
0.7977$\pm$0.0208 0.8485$\pm$0.0215 0.5778$\pm$0.0206 DIR (Wu et al., 2022)
0.7533$\pm$0.0117 0.5927$\pm$0.0097 0.5078$\pm$0.0313 0.5843$\pm$0.0443
0.6115$\pm$0.0587 0.6911$\pm$0.0810 0.5406$\pm$0.0127 DIR$+$RepAug
0.7725$\pm$0.0249 0.6454$\pm$0.0061 0.7453$\pm$0.0080 0.6813$\pm$0.0203
0.7590$\pm$0.0642 0.8561$\pm$0.0159 0.5730$\pm$0.0115 GREA$-$RepAug
0.7770$\pm$0.0178 0.6681$\pm$0.0066 0.7690$\pm$0.0117 0.6737$\pm$0.0235
0.7997$\pm$0.0380 0.8574$\pm$0.0442 0.5988$\pm$0.0169 GREA (ours)
0.7932$\pm$0.0092 0.6750$\pm$0.0067 0.7723$\pm$0.0119 0.6970$\pm$0.0128
0.8237$\pm$0.0237 0.8789$\pm$0.0368 0.6014$\pm$0.0204

We use the alternate training schema in Chang et al. (2020) to train GREA.
That is, we iteratively train $f_{sep}$ ($\text{GNN}_{1}$ and
$\text{MLP}_{1}$) and $f_{pred}$ ($\text{GNN}_{2}$ and $\text{MLP}_{2}$) for a
fixed number of epochs $T_{sep}$ and $T_{pred}$, respectively. The loss
functions for training GREA are

(14)
$\displaystyle\mathcal{L}_{pred}=\mathcal{L}_{{rem}}+\alpha\cdot\mathcal{L}_{{{rep}}},$
(15)
$\displaystyle\mathcal{L}_{sep}=\mathcal{L}_{{rem}}+\alpha\cdot\mathcal{L}_{{{rep}}}+\beta\cdot\mathcal{L}_{reg},$

where $\mathcal{L}_{pred}$ in Eq. 14 and $\mathcal{L}_{sep}$ in Eq. 15 are
used to train $f_{sep}$ ($\text{GNN}_{1}$ and $\text{MLP}_{1}$) and $f_{pred}$
($\text{GNN}_{2}$ and $\text{MLP}_{2}$), respectively. $\alpha$ and $\beta$
are hyperparameters that control the weights of $\mathcal{L}_{rep}$ and
$\mathcal{L}_{reg}$, respectively. During inference, $\hat{y}_{i}^{(r)}$ is
used as the final predicted property of input graph $g_{i}$.

##  5\. Experiments

We conduct experiments to answer the following questions:

  * •

Q1) Effectiveness: Does the proposed GREA make more accurate prediction on
molecule and polymer properties than existing graph classification/regression
methods?

  * •

Q2) Ablation study: Do the environment-based augmentations make positive
effect on the performance?

  * •

Q3) Case study: Based on domain expertise, are the polymer rationale examples
identified by GREA representative?

  * •

Q4) Efficiency: Does the _latent space-based design_ for augmentations perform
faster than explicit graph decoding and encoding? Can we empirically analyze
the complexity?

  * •

Q5) Sensitivity analysis: Is the performance of GREA sensitive to
hyperparameters such as $\alpha$, $\beta$, and $\text{AGG}(\cdot)$?

###  5.1. Experimental Settings

####  5.1.1. Datasets

We conduct experiments on four polymer datasets and seven molecule datasets.
The statistics of the datasets are given in Table 1, such as number of graphs
and average size of graphs. The four datasets GlassTemp, MeltingTemp,
PolyDensity, and O2Perm are used to predict different properties of polymers
such as _glass transition temperature_ (∘C), _polymer density_ g/cm3, _melting
temperature_ (∘C), and _oxygen permeability_ (Barrer). For all the polymer
datasets, we randomly split by 60%/10%/30% for training, validation, and test.
Besides polymer datasets, we use seven molecule datasets from the graph
property prediction task on Open Graph Benchmark or known as OGBG. For all
molecule datasets, we use the scaffold splitting procedure as OGBG adopted (Hu
et al., 2020). It attempts to separate structurally different molecules into
different subsets, which provides a more realistic estimate of model
performance in experiments (Wu et al., 2018). Dataset descriptions with
details are presented in the Appendix A.

####  5.1.2. Evaluation Metrics

On the polymer datasets, we perform the tasks of graph regression. We use the
coefficient of determination (R2) and Root Mean Square Error (RMSE) as
evaluation metrics according to previous works (Ma and Luo, 2020; Hu et al.,
2020). On the molecule datasets, we perform the tasks of graph binary
classification using the Area under the ROC curve (AUC) as the metric. To
evaluate model efficiency, we use the computational time per training batch
(in seconds).

####  5.1.3. Baseline Methods

There are three categories of related methods that we can compare GREA with.
The first category is _graph pooling_ methods that aim at finding (soft)
cluster assignment of nodes towards aggregated representations of graph. They
are U-NetsPool (Gao and Ji, 2021) and SelfAttnPool (Lee et al., 2019). The
second category improves the _optimization and generalization_ of learned
representations. They include StableGNN (Fan et al., 2021), OOD-GNN (Li et
al., 2021), and IRM (Arjovsky et al., 2019). The third is DIR for _graph
rationale identification_ that was proposed in a very recent work by Wu et al.
(2022). To investigate the effect of _environment replacement augmentation_
(denoted by RepAug as a module that may be used or not in the methods), we
implement two method variants: (1) DIR$+$RepAug: We add environment-replaced
augmentation to DIR (Wu et al., 2022) to identify rationales, however, it has
to explicitly decode and encode the rationales; (2) GREA$-$RepAug: We disable
the environment replacement augmentation and use only the environment removal
augmentation, i.e., rationale subgraphs in GREA. In the experiments, we study
two types of GNN models (GCN (Kipf and Welling, 2017) and GIN (Xu et al.,
2019)) as graph encoders for all the methods. Please refer to Appendix B for
details of implementation.

###  5.2. Results on Effectiveness (Q1)

Table 2 presents the results on polymer property regression with R2 and RMSE
metrics. Table 3 presents the results on molecule property classification
using AUC. Underlined are for the best baseline(s). The best baseline is OOD-
GNN for its elimination of the statistical dependence between property-
relevant graph representation and property-irrelevant graph representation.
The first graph rationalization method DIR was evaluated on synthetic data (Wu
et al., 2022); unfortunately, it performs poorly on real polymer and molecule
datasets because it selects edges to create rationale subgraphs and thus loses
the original contextual information of atoms in the the rationale
representations. Compared to them, our GREA with either GCN or GIN
consistently achieves the best performance on all the polymer and molecule
datasets. On the PolyDensity dataset, GREA with GCN improves R2 over OOD-GNN
relatively by +3.91%. On MeltingTemp, GREA with GIN produces $1.56\times$ R2
over DIR.

###  5.3. Ablation Study on GREA (Q2)

Tables 2 and 3 have presented the results of DIR$+$RepAug and GREA$-$RepAug.
DIR$+$RepAug is a variant of baseline method DIR by enabling _environment
replacement augmentations_ for training. GREA$-$RepAug is a variant of our
GREA that disables the replacement augmentations and uses _environment
removal_ only for training. Clearly, DIR$+$RepAug outperforms DIR, showing
positive effect of the replacement augmentations. And the performance of
GREA$-$RepAug is not satisfactory. Environment replacement augmentations are
effective for training graph rationalization methods.

Figure 3. Three polymer examples in O2Perm test set to compare graph
rationales and property predictions by DIR (Wu et al., 2022) and our GREA. DIR
selects _edges_ to decode rationale subgraphs. Our GREA estimates the
probability of _nodes_ being classified into rationales in latent space. The
red boxes indicate incoherent edges that DIR selects. The blue boxes indicate
coherent node sets that contribute to accurate predictions on oxygen
permeability of polymer membrane.

###  5.4. Case Study on Polymer Data (Q3)

(a) Our GREA runs much faster than DIR when batch size (# graphs) increases.

(b) GREA spends comparable amount of training time to deliver the highest AUC.

Figure 4. Efficiency analysis on the ogbg-HIV dataset.

(a) PolyDensity

(b) O2Perm

Figure 5. On two polymer datasets, the performance of GREA is _not_ sensitive
to rationale size $\gamma$ with wide ranges for tuning. Table 4. Effect of
$\text{AGG}\big{(}\mathbf{h}_{i}^{(r)},\mathbf{h}_{j}^{(e)}\big{)}$ in Eq. 9.
We use Sum Pooling by default because it generally performs the best.

MeltingTemp (R2) O2Perm (R2) ogbg-HIV (AUC) Sum Pooling 0.7362$\pm$0.0115
0.9304$\pm$0.0202 0.7932$\pm$0.0092 Mean Pooling 0.7328$\pm$0.0068
0.9288$\pm$0.0331 0.7810$\pm$0.0117 Max Pooling 0.7164$\pm$0.0094
0.8984$\pm$0.0494 0.7809$\pm$0.0137 Concatenation 0.7145$\pm$0.0127
0.9240$\pm$0.0143 0.7771$\pm$0.0096

Given test polymer examples in the O2Perm dataset, we visualize and compare
the rationale subgraphs that are identified by from DIR (Wu et al., 2022) and
our GREA in Figure 3. We have three observations.

First, the rationales identified by GREA have more _coherent structures of
atom nodes_ than those identified by DIR. The red boxes show that quite a few
edges in the rationales by DIR are far separated. This is because DIR
explicitly decodes the subgraphs by selecting edges. Our GREA estimates the
probability of _nodes_ being included in the rationales and uses the
_contextualized representations_ of atoms in the input graphs to create the
representations of rationales. So the rationales have coherent structures of
nodes.

Second, the rationales from GREA are _more interpretable and beneficial_ than
the ones from DIR, based on domain expertise in polymer science. Take a look
at the first polymer example in Figure 3. The rationale from GREA includes
non-aromatic rings and methyl groups. The former group allows larger free
volume elements and lower densities (i.e., enlarge microporousity) in the
polymer’s repeating units, which positively contributes to the gas
permeability (Sanders et al., 2013; Yang et al., 2021). The latter group is
hydrophobic and contributes to steric frustration between polymer chains (Yang
et al., 2021), inducing a positive correlation to the permeability. On the
other hand, the rationale from DIR would make property predictor overestimate
the oxygen permeability, because it suggests that the double-bonded oxygens,
ethers, and nitrogen atoms are positively correlated with the property.
However, it conflicts with observations and conclusions from chemical
experiments in previous literature (Yang et al., 2021) where researchers argue
that the double-bonded oxygens, ethers, and nitrogen atoms are negatively
correlated with gas permeability. For the second and third examples, DIR also
predicts through double-bonded oxygens, ethers, and nitrogen atoms, and it
overestimates the permeability. Our GREA realizes and employs the true
relationship between the functional groups and property and successfully
suppresses the representations of non-aromatic rings and methyl groups in the
prediction. GREA intrinsically discovers correct relationships between
rationale subgraphs and the property.

Third, the rationales from GREA are _commonly observed across different
polymers_. We expect rationales to have universal indication on the polymer
properties. The rationales identified in the second and third examples both
have the fused heterocyclic rings (at the right end of the monomers and
highlighted by blue boxes).

###  5.5. Results on Efficiency (Q4)

We conduct efficiency analysis using the ogbg-HIV dataset without losing the
generality. Results are presented in Figure 4. When batch size increases, in
other words, when a batch has more and more graphs, the time cost per batch of
DIR increases significantly; our proposed GREA spends much less time than DIR.
Empirically we show that our GREA is more efficient than DIR. This is because
GREA does not explicitly decode or encode the subgraphs but directly creates
their representations in latent space. Figure 4(b) shows that compared to
three most competitive baselines, GREA delivers the highest AUC by learning
augmented examples, while spending comparable amount of time.

###  5.6. Sensitivity Analysis (Q5)

(a) GlassTemp

(b) MeltingTemp

(c) PolyDensity

(d) O2Perm

Figure 6. On four polymer datasets, the performance of GREA (in R2) is _not_
sensitive to hyperparameters $\alpha$ and $\beta$ in Eq. 15.

Without losing the generality, we conduct three series of sensitivity
analyses. First, Figure 6 shows that on four polymer datasets, the performance
of GREA in terms of R2 is insensitive to the hyperparameters $\alpha$ and
$\beta$ in Eq. 15. Second, Figure 5 shows that the performance is insensitive
to rationale size $\gamma$ in Eq. 13. Third, on two polymer datasets and one
of the most popular molecule datasets, Table 4 compares the effects of
different choices of $\text{AGG}(\cdot)$ function that aggregates the
representations of rationale and environment subgraphs. Sum pooling is
generally the best choice.

##  6\. Conclusions

In this work, we made the first attempt to improve graph rationale
identification using data augmentations, including environment replacement,
for accurate and interpretable graph property prediction. We proposed an
efficient framework that performs rationale-environment separation and
representation learning on real and augmented examples in one latent space.
Experiments on molecule and polymer datasets demonstrated its effectiveness
and efficiency.

###### Acknowledgements.

This research was supported in part by NSF Grants IIS-1849816, IIS-2142827,
IIS-2146761, and CBET-2102592.

## References

  * (1)
  * Arjovsky et al. (2019) Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. 2019.  Invariant risk minimization. In _arXiv:1907.02893_.
  * Chang et al. (2020) Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. 2020\.  Invariant rationalization. In _ICML_. 1448–1458.
  * Chen et al. (2020) Deli Chen, Yankai Lin, Wei Li, Peng Li, Jie Zhou, and Xu Sun. 2020.  Measuring and relieving the over-smoothing problem for graph neural networks from the topological view. In _AAAI_ , Vol. 34. 3438–3445.
  * Chen et al. (2021) Lihua Chen, Ghanshyam Pilania, Rohit Batra, Tran Doan Huan, Chiho Kim, Christopher Kuenneth, and Rampi Ramprasad. 2021.  Polymer informatics: Current status and critical next steps.  _Materials Science and Engineering: R: Reports_ 144 (2021), 100595\.
  * Fan et al. (2021) Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. 2021.  Generalizing Graph Neural Networks on Out-Of-Distribution Graphs. In _arXiv:2111.10657_.
  * Gao and Ji (2021) Hongyang Gao and Shuiwang Ji. 2021.  Graph U-Nets.  _IEEE TPAMI_ (2021).
  * Guo et al. (2021) Zhichun Guo, Chuxu Zhang, Wenhao Yu, John Herr, Olaf Wiest, Meng Jiang, and Nitesh V Chawla. 2021.  Few-Shot Graph Learning for Molecular Property Prediction. In _WWW_. 2559–2567.
  * Hamilton et al. (2017) William L Hamilton, Rex Ying, and Jure Leskovec. 2017\.  Inductive representation learning on large graphs. In _NeurIPS_. 1025–1035.
  * Hu et al. (2020) Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. 2020\.  Open Graph Benchmark: Datasets for Machine Learning on Graphs. In _NeurIPS_.
  * Jiang et al. (2022) Meng Jiang, Taeho Jung, Ryan Karl, and Tong Zhao. 2022\.  Federated Dynamic Graph Neural Networks with Secure Aggregation for Video-based Distributed Surveillance.  _TIST_ 13, 4 (2022), 1–23.
  * Kim et al. (2018) Chiho Kim, Anand Chandrasekaran, Tran Doan Huan, Deya Das, and Rampi Ramprasad. 2018\.  Polymer genome: a data-powered polymer informatics platform for property predictions.  _The Journal of Physical Chemistry C_ 122, 31 (2018), 17575–17585.
  * Kipf and Welling (2017) Thomas N Kipf and Max Welling. 2017.  Semi-supervised classification with graph convolutional networks. In _ICLR_.
  * Landrum (2013) Greg Landrum. 2013\.  RDKit: A software suite for cheminformatics, computational chemistry, and predictive modeling.
  * Lee et al. (2019) Junhyun Lee, Inyeop Lee, and Jaewoo Kang. 2019.  Self-attention graph pooling. In _ICML_. 3734–3743.
  * Li et al. (2021) Haoyang Li, Xin Wang, Ziwei Zhang, and Wenwu Zhu. 2021\.  OOD-GNN: Out-of-Distribution Generalized Graph Neural Network. In _arXiv:2112.03806_.
  * Ma et al. (2019) Ruimin Ma, Zeyu Liu, Quanwei Zhang, Zhiyu Liu, and Tengfei Luo. 2019.  Evaluating polymer representations via quantifying structure–property relationships.  _Journal of chemical information and modeling_ 59, 7 (2019), 3110–3119.
  * Ma and Luo (2020) Ruimin Ma and Tengfei Luo. 2020.  PI1M: a benchmark database for polymer informatics.  _Journal of Chemical Information and Modeling_ 60, 10 (2020), 4684\.
  * Ma et al. (2021) Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang, and Neil Shah. 2021\.  A unified view on graph neural networks as graph signal denoising. In _CIKM_. 1202–1211.
  * Mesquita et al. (2020) Diego Mesquita, Amauri Souza, and Samuel Kaski. 2020\.  Rethinking pooling in graph neural networks. In _NeurIPS_.
  * Otsuka et al. (2011) Shingo Otsuka, Isao Kuwajima, Junko Hosoya, Yibin Xu, and Masayoshi Yamazaki. 2011.  PoLyInfo: Polymer database for polymeric materials design. In _International Conference on Emerging Intelligent Data and Web Technologies_. 22.
  * Park et al. (2021) Hyeonjin Park, Seunghun Lee, Sihyeon Kim, Jinyoung Park, Jisu Jeong, Kyung-Min Kim, Jung-Woo Ha, and Hyunwoo J Kim. 2021\.  Metropolis-Hastings Data Augmentation for Graph Neural Networks. In _NeurIPS_.
  * Rong et al. (2019) Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. 2019\.  DropEdge: Towards Deep Graph Convolutional Networks on Node Classification. In _ICLR_.
  * Rosenfeld et al. (2021) Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. 2021.  The Risks of Invariant Risk Minimization. In _ICLR_.
  * Sanders et al. (2013) David F Sanders, Zachary P Smith, Ruilan Guo, Lloyd M Robeson, James E McGrath, Donald R Paul, and Benny D Freeman. 2013.  Energy-efficient polymeric gas separation membranes for a sustainable future: A review.  _Polymer_ 54, 18 (2013), 4729–4761.
  * Thornton et al. (2012) A Thornton, L Robeson, B Freeman, and D Uhlmann. 2012\.  Polymer Gas Separation Membrane Database.
  * Veličković et al. (2018) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. 2018.  Graph Attention Networks. In _ICLR_.
  * Wang et al. (2020a) Daheng Wang, Meng Jiang, Munira Syed, Oliver Conway, Vishal Juneja, Sriram Subramanian, and Nitesh V Chawla. 2020a.  Calendar graph neural networks for modeling time structures in spatiotemporal user behaviors. In _KDD_. 2581–2589.
  * Wang et al. (2021a) Daheng Wang, Zhihan Zhang, Yihong Ma, Tong Zhao, Tianwen Jiang, Nitesh Chawla, and Meng Jiang. 2021a.  Modeling co-evolution of attributed and structural information in graph sequence.  _IEEE TKDE_ (2021).
  * Wang et al. (2021b) Daheng Wang, Zhihan Zhang, Yihong Ma, Tong Zhao, Tianwen Jiang, Nitesh Chawla, and Meng Jiang. 2021b.  Modeling co-evolution of attributed and structural information in graph sequence.  _IEEE TKDE_ (2021).
  * Wang et al. (2021c) Daheng Wang, Tong Zhao, Nitesh V Chawla, and Meng Jiang. 2021c.  Dynamic Attributed Graph Prediction with Conditional Normalizing Flows. In _ICDM_. IEEE, 1385–1390.
  * Wang et al. (2020b) Yiwei Wang, Wei Wang, Yuxuan Liang, Yujun Cai, and Bryan Hooi. 2020b.  Graphcrop: Subgraph cropping for graph classification. In _arXiv:2009.10564_.
  * Wang et al. (2020c) Yiwei Wang, Wei Wang, Yuxuan Liang, Yujun Cai, Juncheng Liu, and Bryan Hooi. 2020c.  Nodeaug: Semi-supervised node classification with data augmentation. In _KDD_. 207–217.
  * Wei et al. (2021) Xingfei Wei, Zhi Wang, Zhiting Tian, and Tengfei Luo. 2021\.  Thermal Transport in Polymers: A Review.  _Journal of Heat Transfer_ 143, 7 (2021), 072101\.
  * Wu et al. (2022) Yingxin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. 2022.  Discovering Invariant Rationales for Graph Neural Networks. In _ICLR_.
  * Wu et al. (2020) Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. 2020\.  A comprehensive survey on graph neural networks.  _IEEE TNNLS_ 32, 1 (2020), 4–24.
  * Wu et al. (2018) Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. 2018\.  MoleculeNet: a benchmark for molecular machine learning.  _Chemical science_ (2018), 513–530.
  * Xu et al. (2019) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. 2019\.  How Powerful are Graph Neural Networks?. In _ICLR_.
  * Yang et al. (2021) Jason Yang, Lei Tao, Jinlong He, Jeffrey McCutcheon, and Ying Li. 2021.  Discovery of Innovative Polymers for Next-Generation Gas-Separation Membranes using Interpretable Machine Learning. In _chemrxiv-2021-p4g7z_.
  * Ying et al. (2019) Rex Ying, Dylan Bourgeois, Jiaxuan You, Marinka Zitnik, and Jure Leskovec. 2019.  Gnnexplainer: Generating explanations for graph neural networks. In _NeurIPS_.
  * Ying et al. (2018) Rex Ying, Jiaxuan You, Christopher Morris, Xiang Ren, William L Hamilton, and Jure Leskovec. 2018\.  Hierarchical graph representation learning with differentiable pooling. In _NeurIPS_. 4805–4815.
  * You et al. (2020) Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. 2020\.  Graph contrastive learning with augmentations. In _NeurIPS_. 5812–5823.
  * Zhang et al. (2020) Ziwei Zhang, Peng Cui, and Wenwu Zhu. 2020.  Deep learning on graphs: A survey.  _IEEE TKDE_ (2020).
  * Zhao et al. (2021a) Tong Zhao, Tianwen Jiang, Neil Shah, and Meng Jiang. 2021a.  A synergistic approach for graph anomaly detection with pattern mining and feature learning.  _IEEE TNNLS_ (2021).
  * Zhao et al. (2022a) Tong Zhao, Gang Liu, Stephan Günnemann, and Meng Jiang. 2022a.  Graph Data Augmentation for Graph Machine Learning: A Survey.  _arXiv preprint arXiv:2202.08871_ (2022).
  * Zhao et al. (2022b) Tong Zhao, Gang Liu, Daheng Wang, Wenhao Yu, and Meng Jiang. 2022b.  Learning from Counterfactual Links for Link Prediction.  _ICML_ (2022).
  * Zhao et al. (2021b) Tong Zhao, Yozen Liu, Leonardo Neves, Oliver Woodford, Meng Jiang, and Neil Shah. 2021b.  Data Augmentation for Graph Neural Networks. In _AAAI_. 11015.
  * Zhao et al. (2021c) Tong Zhao, Bo Ni, Wenhao Yu, Zhichun Guo, Neil Shah, and Meng Jiang. 2021c.  Action Sequence Augmentation for Early Graph-based Anomaly Detection. In _CIKM_. 2668–2678.
  * Zhou et al. (2020) Jiajun Zhou, Jie Shen, and Qi Xuan. 2020.  Data Augmentation for Graph Classification. In _CIKM_. 2341–2344.
  * Zhu et al. (2021) Yanqiao Zhu, Yichen Xu, Feng Yu, Qiang Liu, Shu Wu, and Liang Wang. 2021.  Graph contrastive learning with adaptive augmentation. In _WWW_. 2069–2080.

##  Appendix A Dataset Details

##### Polymer datasets

The four datasets GlassTemp, MeltingTemp, PolyDensity, and O2Perm are used to
predict different properties of polymers such as _glass transition
temperature_ (∘C), _polymer density_ g/cm3, _melting temperature_ (∘C), and
_oxygen permeability_ (Barrer). GlassTemp, MeltingTemp, and PolyDensity are
collected from PolyInfo, which is the largest web-based polymer database
(Otsuka et al., 2011). The O2Perm dataset is created from the Membrane Society
of Australasia portal, consisting of a variety of gas permeability data
(Thornton et al., 2012). However, the limited size (i.e., 595 polymers) brings
great challenges to rationale identification and property prediction. Since a
polymer is built from repeated monomer units, researchers use monomers as
polymer graphs to predict properties. Different from molecular graphs, the
monomer graphs have two special nodes (see “$*$” in the molecular structures
in Figure 1), indicating the polymerization points of monomers (Ma and Luo,
2020). For all the polymer datasets, we randomly split by 60%/10%/30% for
training, validation, and test.

##### Molecule datasets

Besides polymer datasets, we use seven molecule datasets from the graph
property prediction task on Open Graph Benchmark or known as OGBG. They were
originally collected by MoleculeNet (Wu et al., 2018) and used to predict the
properties of molecules, including (1) inhibition to HIV virus replication in
ogbg-HIV, (2) toxicological properties of 617 types in ogbg-ToxCast, (3)
toxicity measurements such as nuclear receptors and stress response in ogbg-
Tox21, (4) blood–brain barrier permeability in ogbg-BBBP, (5) inhibition to
human $\beta$-secretase 1 in ogbg-BACE, (6) FDA approval status or failed
clinical trial in ogbg-ClinTox, and (7) having drug side effects of 27 system
organ classes in ogbg-SIDER. For all molecule datasets, we use the scaffold
splitting procedure as OGBG adopted (Hu et al., 2020). It attempts to separate
structurally different molecules into different subsets, which provides a more
realistic estimate of model performance in experiments (Wu et al., 2018).

##  Appendix B Implementation Details

All the experiments in this work are conducted on an Linux server with Intel
Xeon Gold 6130 Processor (16 Cores @2.1Ghz), 96 GB of RAM, and a single RTX
2080Ti card (11 GB of RAM). Our method is implemented with `Python 3.9.9` and
`PyTorch 1.10.1`. We manually tune the hyperparameters over the following
ranges:

  * •

$\gamma\in\\{0.05,0.1,0.15,\dots,0.75,0.8\\}$,

  * •

$T_{sep}\in\\{1,2\\}$,

  * •

$T_{pred}\in\\{2,3\\}$,

  * •

Learning rate $\in\\{0.001,0.005,0.01\\}$,

  * •

Batch size $\in\\{32,128,256,512\\}$,

  * •

Representation dimensions $d_{1}$, $d_{2}\in\\{64,128,300\\}$,

  * •

Number of GNN1 layer $L_{1}=\\{2\\}$,

  * •

Number of GNN2 layers $L_{2}\in\\{2,3,4,5\\}$.

We use sum pooling as the default $\text{AGG}(\cdot)$ in GREA for the
experiments in Tables 2 and 3. We set GIN as the default encoder for all
ablation studies, case studies, and efficiency analysis. We employ the virtual
node trick (Hu et al., 2020) for all methods on the ogbg-HIV, ogbg-Tox21,
ogbg-BBBP, and all polymer datasets. For PolyDensity, we train and evaluate
the models using the logarithm of the property (Ma and Luo, 2020). We report
the mean and standard deviation of the test performance over 10 runs with
different random initialization of the parameters.

Our Codes and data are available on the
GitHub111https://github.com/liugangcode/GREA. To implement the baseline
methods, we use the official code package222https://github.com/Wuyxin/DIR-GNN
from the authors for DIR (Wu et al., 2022). For U-NetsPool (Gao and Ji, 2021)
and SelfAttnPool (Lee et al., 2019), we use the public implementation provided
by the `PyG`333https://github.com/pyg-team/pytorch_geometric package. For IRM
(Arjovsky et al., 2019), we implement it’s graph version based on its official
repository.444https://github.com/facebookresearch/InvariantRiskMinimization As
source codes of OOD-GNN (Li et al., 2021) and StableGNN (Fan et al., 2021) are
not publically available, we implement then with the official code package of
StableNet555https://github.com/xxgege/StableNet and the `PyG` package.