Original: Learning Substructure Invariance for Out-of-Distribution Molecular Representations
Revised: Learning Substructure Invariance for Out-of-Distribution Molecular Representations

Original: AfﬁliationAddress

Original: Abstract
Revised: Abstract

Original: Molecule representation learning (MRL) has been extensively studied and currentmethods have shown promising power for various tasks, e.g., molecular propertyprediction and target identiﬁcation.
Revised: Molecule representation learning (MRL) has been extensively studied and current methods have shown promising power for various tasks, e.g., molecular property prediction and target identiﬁcation.

Original: However, a common hypothesis of existingmethods is that either the model development or experimental evaluation is mostlybased on i.i.d.
Revised: However, a common hypothesis of existing methods is that either the model development or experimental evaluation is mostly based on i.i.d.

Original: data across training and testing.
Revised: data across training and testing.

Original: Such a hypothesis can be violatedin real-world applications where testing molecules could come from new environments, bringing about serious performance degradation or unexpected prediction.
Revised: Such a hypothesis can be violated in real-world applications where testing molecules could come from new environments, bringing about serious performance degradation or unexpected prediction.

Original: We propose a new representation learning framework entitled MoleOOD to enhancethe robustness of MRL models against such distribution shifts, motivated by anobservation that the (bio)chemical properties of molecules are usually invariantlyassociated with certain privileged molecular substructures across different environments (e.g., scaffolds, sizes, etc.).
Revised: We propose a new representation learning framework entitled MoleOOD to enhance the robustness of MRL models against such distribution shifts, motivated by an observation that the (bio)chemical properties of molecules are usually invariantly associated with certain privileged molecular substructures across different environments (e.g., scaffolds, sizes, etc.).

Original: Speciﬁcally, We introduce an environment inference model to identify the latent factors that impact data generation from differentdistributions in a fully data-driven manner.
Revised: Speciﬁcally, We introduce an environment inference model to identify the latent factors that impact data generation from different distributions in a fully data-driven manner.

Original: We also propose a new learning objective to guide the molecule encoder to leverage environment-invariant substructuresthat more stably relate with the labels across environments.
Revised: We also propose a new learning objective to guide the molecule encoder to leverage environment-invariant substructures that more stably relate with the labels across environments.

Original: Extensive experimentson ten real-world datasets demonstrate that our model has a stronger generalizationability than existing methods under various out-of-distribution (OOD) settings,despite the absence of manual speciﬁcations of environments.
Revised: Extensive experiments on ten real-world datasets demonstrate that our model has a stronger generalization ability than existing methods under various out-of-distribution (OOD) settings, despite the absence of manual speciﬁcations of environments.

Original: Particularly, ourmethod achieves up to 5.9% and 3.9% improvement over the strongest baselineson OGB and DrugOOD benchmarks in terms of ROC-AUC, respectively.
Revised: Particularly, our method achieves up to 5.9% and 3.9% improvement over the strongest baselines on OGB and DrugOOD benchmarks in terms of ROC-AUC, respectively.

Revised: Our source code is publicly available at https://github.com/yangnianzu0515/MoleOOD .

Original:  Introduction
Revised:  Introduction

Original: Predicting molecular properties plays an important role in many related applications like drugdiscovery [12] and material design [47].
Revised: Predicting molecular properties plays an important role in many related applications like drug discovery [12] and material design [47].

Original: These professional tasks conventionally take great efforts byexperts e.g. in chemistry and pharmacology.
Revised: These professional tasks conventionally take great efforts by experts e.g. in chemistry and pharmacology.

Original: Recent years have witnessed inspiring breakthroughs onbuilding effective machine learning models for scientiﬁc discovery and solid progress has been madealong the avenue of ML-based molecule representation learning (MRL).
Revised: Recent years have witnessed inspiring breakthroughs on building effective machine learning models for scientiﬁc discovery, and solid progress has been made along the avenue of ML-based molecule representation learning (MRL).

Original: In general, MRL aims atembedding a molecule into a vector in latent space as a foundation model, on top of which the learned representations could be used for a variety of downstream tasks, such as target identiﬁcation [65], retrosynthetic analysis [61], search of antibiotics [50], virtual screening [37] for drug discovery, etc.
Revised: In general, MRL aims at embedding a molecule into a vector in latent space as a foundation model, on top of which the learned representations could be used for a variety of downstream tasks, such as target identiﬁcation [65], retrosynthetic analysis [61], search of antibiotics [50], virtual screening [37] for drug discovery, etc.

Original: The challenge, however, is that existing MRL methods are mostly based on an underlying hypothesis that training and testing molecules are independently sampled from an identical environment, yetreal-world environments are often dynamic and uncertain, which requires the model to effectively handle distribution shifts.
Revised: The challenge, however, is that existing MRL methods are mostly based on an underlying hypothesis that training and testing molecules are independently sampled from an identical environment, yet O H O H O H Cyclopropanol 1,4-Cyclohexanediol

Revised: Scaffold 1 3C-Ring

Revised: Scaffold 2 6C-Ring

Original: In fact, the available experimental molecule data are rather limited while the candidate molecules to be tested are often diverse, coming from unknown environments
Revised: O H O O O H

Revised: O

Revised: O H

Revised: O H

Revised: O

Revised: O H

Revised: Smaller Size

Revised: Larger Size

Revised: Environment: Scaffold

Revised: Environment: Size

Revised: Formic Acid

Revised: Citric Acid real-world environments are often dynamic and uncertain, which requires the model to effectively handle distribution shifts.

Revised: In fact, the available experimental molecule data are rather limited while the candidate molecules to be tested are often diverse, coming from unknown environments.

Original: Taking the virtual screening [37] as an example (which is a common protocol in drug discovery and usually for target identiﬁcation), the prediction model is typically trained on some known target proteins.
Revised: Taking the virtual screening [37] as an example (which is a common protocol in drug discovery and usually for target identiﬁcation), the prediction model is typically trained on some known target proteins.

Original: However, some unpredictable events like COVID-19 may occur, bringing new targets fromunknown distributions.
Revised: However, some unpredictable events like COVID-19 may occur, bringing new targets from unknown distributions.

Original: Similar scenarios where training and testing data are sampled from differentdistributions are common in real world, posing an urgency for strengthening current MRL methodsregarding out-of-distribution (OOD) generalization [38, 6, 41].[Figure]
Revised: Similar scenarios where training and testing data are sampled from different distributions are common in real world, posing an urgency for strengthening current MRL methods regarding out-of-distribution (OOD) generalization [38, 6, 41].

Original: invariantly contributes to the water solubility of the two molecules which contain different scaffolds, i.e. sampled from different environments by deﬁnition.

Original: Existing methods devised for out-of-distribution generalizationmostly focus on Euclidean data such as images, while few endeav- [Figure] ors OOD generalization on non-Euclidean data [57].
Revised: Existing methods devised for out-of-distribution generalization mostly focus on Euclidean data such as images, while few endeavors OOD generalization on non-Euclidean data [57].

Original: In particular, molecules, as a kind of typical non-Euclidean data, i.e., graphstructured data, is different from visual data in nature.
Revised: In particular, molecules, as a kind of typical non-Euclidean data, i.e., graph-structured data, is different from visual data in nature.

Original: The work [25]point out that existing OOD models [3, 51, 16, 66, 46] fail to exhibitsigniﬁcant improvement on MRL tasks against distribution shifts andeven the simple Empirical Risk Minimization (ERM) [52] methodoutperforms these latest methods, which is also empirically veriﬁedby [15].
Revised: The work [25] point out that existing OOD models [3, 51, 16, 66, 46] fail to exhibit signiﬁcant improvement on MRL tasks against distribution shifts and even the simple Empirical Risk Minimization (ERM) [52] method outperforms these latest methods, which is also empirically veriﬁed by [15].

Original: We aim to develop an OOD method tailored for moleculesto solve the OOD generalization problem on MRL in this paper.
Revised: We aim to develop an OOD method tailored for molecules to solve the OOD generalization problem on MRL in this paper.

Original: We incorporate an effective prior in the molecule domain into ourmodel design: the (bio)chemical properties of a molecule are usually associated with a few privileged molecular substructures, whichhas been consistently shown by studies [29, 44, 68, 27] across bioinformatics, pharmacy, and data mining.
Revised: We incorporate an effective prior in the molecule domain into our model design: the (bio)chemical properties of a molecule are usually associated with a few privileged molecular substructures, which has been consistently shown by studies [29, 44, 68, 27] across bio-informatics, pharmacy, and data mining.

Revised: The common practice speciﬁes environments as some prominent information of the molecules e.g. scaffold pattern [30, 21] and molecule size [25].

Original: Fig.
Revised: Fig.

Original: 1 provides an illustrative  example, where two molecules Cyclopropanol ( C 3 H 6 O ) and1,4-Cyclohexanediol ( C 6 H 12 O 2 ) contain different scaffold patterns 1 :the former is 3C-ring and the latter is 6C-ring .The common practice speciﬁes environments as someprominent information of the molecules e.g. scaffold pattern [30, 21].
Revised: 1 provides two illustrative examples. Let’s ﬁrst take a look at the left example, where two molecules Cyclopropanol ( C 3 H 6 O ) and 1,4-Cyclohexanediol ( C 6 H 12 O 2 ) contain different scaffold patterns 2 : the former is 3C-ring and the latter is 6C-ring .

Original: Thus, the data-generatingenvironments and the induced distributions which these two molecules are sampled from can beconsidered different [21].
Revised: Thus, the data-generating environments and the induced distributions which these two molecules are sampled from can be considered different [21].

Original: Though sampled from different distributions, they are both readily solublein water due to the invariant substructure hydroxy [22] shared across different environments.
Revised: Though sampled from different distributions, they are both readily soluble in water due to the invariant substructure hydroxy [22] shared across different environments.

Revised: As for the example on the right of Fig. 1, the sizes of two molecules Formic Acid ( CH 2 O 2 ) and Citric Acid ( C 8 H 8 O 7 ) differ a lot.

Revised: Consequently, they can also be considered as being sampled from different environments.

Revised: Owing to the shared invarint substructure carboxy ( − COOH ) , they are both readily soluable in water, too.

Original: Hence,a promising paradigm would be to learn the causal data-generating invariance from the substructuresacross environments, regarding a certain property, for the OOD generalization purpose.
Revised: Hence, a promising paradigm would be to learn the causal data-generating invariance from the substructures across environments, regarding a certain property, for the OOD generalization purpose.

Original: Another important observation for consideration is that existing speciﬁcations for environments areoften handcrafted or rule-based and not structured, which could provide insufﬁcient informationfor capturing the fundamental relations across domains from the casual data-generating perspective.
Revised: Another important observation for consideration is that existing speciﬁcations for environments are often handcrafted or rule-based and not structured, which could provide insufﬁcient information for capturing the fundamental relations across domains from the casual data-generating perspective.

Original: Besides, some studies [25, 15] show that directly utilizing such environment labels as input whenadapting existing OOD generalization methods to MRL tasks can be problematic.
Revised: Besides, some studies [25, 15] show that directly utilizing such environment labels as input when adapting existing OOD generalization methods to MRL tasks can be problematic.

Original: Furthermore,manual speciﬁcations of environments may be unavailable in reality.
Revised: Furthermore, manual speciﬁcations of environments may be unavailable in reality.

Original: Hence, we aim to develop alabel-free model that does not rely on the above ad-hoc environment labels.
Revised: Hence, we aim to develop a label-free model that does not rely on the above ad-hoc environment labels.

Original: As shown later, our modelcan infer the environment labels in an unsupervised manner, namely for environment clustering.
Revised: As shown later, our model can infer the environment labels in an unsupervised manner, namely for environment clustering.

Original: To achieve robust molecule representation for OOD generalization and overcome potentially unreliable environment labels, we devise a new MRL framework without explicitly using the environmentlabel information.
Revised: To achieve robust molecule representation for OOD generalization and overcome potentially unreliable environment labels, we devise a new MRL framework without explicitly using the environment label information.

Original: We ﬁrst formulate OOD generalization for molecular property prediction byintroducing a latent variable for environments that affect the data generation.
Revised: We ﬁrst formulate OOD generalization for molecular property prediction by introducing a latent variable for environments that affect the data generation.

Original: Then we analyzethe essential cause behind the failure of existing MRL models and propose a new learning schemebased on the invariance principle [57].
Revised: Then we analyze the essential cause behind the failure of existing MRL models and propose a new learning scheme based on the invariance principle [57].

Original: The training procedures contain two steps: 1) optimize anenvironment inference model from training data; 2) optimize a molecule encoder and a predictor.
Revised: The training procedures contain two steps: 1) optimize an environment inference model from training data; 2) optimize a molecule encoder and a predictor.

Original: Our general framework can integrate existing GNN backbones and achieve improvements on four
Revised: Our general framework can integrate existing GNN backbones and achieve improvements on four

Original: OGB molecular property prediction tasks [21], as shown in our experimental results.
Revised: OGB molecular property prediction tasks [21], as shown in our experimental results.

Original: As for a newlyreleased benchmark for drug-oriented OOD learning [25], even without access to environment labels,our method can still outperform state-of-the-art models that rely on environment labels for training inﬁve out of six datasets.
Revised: As for a newly released benchmark for drug-oriented OOD learning [25], even without access to environment labels, our method can still outperform state- of-the-art models that rely on environment labels for training in ﬁve out of six datasets.

Original: The contributions of this paper are:
Revised: The contributions of this paper are:

Original: • We formulate the OOD problem for molecule representation learning, by particularly incorporatingan important observation that the substructure of molecule can convey invariant casual informationacross environments, regarding certain property prediction tasks.
Revised: • We formulate the OOD problem for molecule representation learning, by particularly incorporating an important observation that the substructure of molecule can convey invariant casual information across environments, regarding certain property prediction tasks.

Original: To our best knowledge, this isthe ﬁrst work that explicitly models such substructures for OOD molecule representation.
Revised: To our best knowledge, this is the ﬁrst work that explicitly models such substructures for OOD molecule representation.

Original: • Under the environment-invariance principle with speciﬁc substructure invariance priors, we proposea new learning objective to learn robust representations.
Revised: • Under the environment-invariance principle with speciﬁc substructure invariance priors, we propose a new learning objective to learn robust representations.

Original: In particular, our model does not requireenvironment labels which in fact can be noisy and unreliable, but instead achieve environmentinference in an unsupervised manner.
Revised: In particular, our model does not require environment labels which in fact can be noisy and unreliable, but instead achieve environment inference in an unsupervised manner.

Original: To our best knowledge, this is the ﬁrst work for environmentclustering for molecule representation, which has not been fulﬁlled in existing OOD literaturebeyond molecules (thanks to the speciﬁc substructure-property invariance priors).• We conduct extensive experiments on ten public datasets.
Revised: To our best knowledge, this is the ﬁrst work for environment clustering for molecule representation, which has not been fulﬁlled in existing OOD literature beyond molecules (thanks to the speciﬁc substructure-property invariance priors). • We conduct extensive experiments on ten public datasets.

Original: Results demonstrate that our modelyields consistent and signiﬁcant improvements over various existing MRL methods as backbonesand also achieves competitive or even superior prediction compared to state-of-the-art modelstailored to OOD learning with environment labels used as extra inputs in both training and testing.
Revised: Results demonstrate that our model yields consistent and signiﬁcant improvements over various existing MRL methods as backbones and also achieves competitive or even superior prediction compared to state-of-the-art models tailored to OOD learning with environment labels used as extra inputs in both training and testing.

Original: Particularly, our method achieves up to 5.9% higher ROC-AUC on public OGB molecular propertyprediction benchmarks than the counterpart model trained with traditional objective.
Revised: Particularly, our method achieves up to 5.9% higher ROC-AUC on public OGB molecular property prediction benchmarks than the counterpart model trained with traditional objective.

Original: Besides,for drug-oriented benchmarks DrugOOD, when environment labels are not used, our model stilloutperforms several SOTA approaches tailored for general OOD learning (using environmentlabels as extra training information) by up to 3.9% w.r.t.
Revised: Besides, for drug-oriented benchmarks DrugOOD, when environment labels are not used, our model still outperforms several SOTA approaches tailored for general OOD learning (using environment labels as extra training information) by up to 3.9% w.r.t.

Original: ROC-AUC.
Revised: ROC-AUC.

Original:  Backgrounds and Related Works
Revised:  Backgrounds and Related Works

Original: Out-of-Distribution Generalization.
Revised: Out-of-Distribution Generalization.

Original: Deep neural networks are prone to suffering signiﬁcant performance degradation under distribution shifts, motivating a surge of works on OOD generalization.
Revised: Deep neural networks are prone to suffering signiﬁcant performance degradation under distribution shifts, motivating a surge of works on OOD generalization.

Original: Recent studies [45, 3, 8, 57] assume that there is a potential environment variable e accounting forthe distribution shift between the training and testing data.
Revised: Recent studies [45, 3, 8, 57] assume that there is a potential environment variable e accounting for the distribution shift between the training and testing data.

Original: In general cases the goal is to predict thetarget label y given the associated input x .
Revised: In general cases the goal is to predict the target label y given the associated input x .

Original: Then, the OOD problem could be formally formulated as:
Revised: Then, the OOD problem could be formally formulated as:

Original: [Equation]
Revised: [Equation]

Original: (1) where E denotes the support of environments, f ( · ) is the prediction model and l ( · , · ) represents aloss function.
Revised: (1) where E denotes the support of environments, f ( · ) is the prediction model and l ( · , · ) represents a loss function.

Original: Notice that E ( x,y ) ∼ p ( x , y | e = e ) [ l ( f ( x ) , y )
Revised: Notice that E ( x,y ) ∼ p ( x , y | e = e ) [ l ( f ( x ) , y )

Original: | e ] is called the risk function under a givenenvironment e and denoted as R e ( x e , y e ) [31].
Revised: | e ] is called the risk function under a given environment e and denoted as R e ( x e , y e ) [31].

Original: Invariant Learning.
Revised: Invariant Learning.

Original: There is an emerging line of research [45, 3, 9, 11] regarding invariant predictor learning, for solving the OOD generalization problem.
Revised: There is an emerging line of research [45, 3, 9, 11] regarding invariant predictor learning, for solving the OOD generalization problem.

Original: These methods propose to ﬁnd aninvariant predictor that could uncover invariant relationships between inputs and targets across allenvironments [31].
Revised: These methods propose to ﬁnd an invariant predictor that could uncover invariant relationships between inputs and targets across all environments [31].

Original: The invariant predictor should learn an invariant representation satisfying such ainvariance principle : 1) sufﬁciency : shows sufﬁcient predictive power for the target, 2) invariance :contributes to equal (optimal) performance for the downstream tasks across all environments.
Revised: The invariant predictor should learn an invariant representation satisfying such a invariance principle : 1) sufﬁciency : shows sufﬁcient predictive power for the target, 2) invariance : contributes to equal (optimal) performance for the downstream tasks across all environments.

Original: A recentwork [57] leverages invariance principle to handle distribution shifts on graphs, but it is designed fornode-level prediction.
Revised: A recent work [57] leverages invariance principle to handle distribution shifts on graphs, but it is designed for node-level prediction.

Original: In contrast, we focus on molecular graph-level classiﬁcation tasks.
Revised: In contrast, we focus on molecular graph-level classiﬁcation tasks.

Original: Molecule Representation Learning.
Revised: Molecule Representation Learning.

Original: Existing molecule representation learning methods can beclassiﬁed into two categories.
Revised: Existing molecule representation learning methods can be classiﬁed into two categories.

Original: The ﬁrst is SMILES-based methods where SMILES refers to Simpliﬁed Molecular Input Line Entry System [2].
Revised: The ﬁrst is SMILES-based methods where SMILES refers to Simpliﬁed Molecular Input Line Entry System [2].

Original: They use language models to process the textualrepresentation (SMILES) of a molecule, for example, Transformer [53] or BERT [14].
Revised: They use language models to process the textual representation (SMILES) of a molecule, for example, Transformer [53] or BERT [14].

Original: SMILES is alinear encoding for molecules and highly depends on the traverse order of molecule graphs.
Revised: SMILES is a linear encoding for molecules and highly depends on the traverse order of molecule graphs.

Original: Therefore its expressiveness is limited for problems like medication recommendation which we believecalls for ﬁne-grained molecular structure extraction.
Revised: Therefore its expressiveness is limited for problems like medication recommendation which we believe calls for ﬁne-grained molecular structure extraction.

Original: Beyond the above linear encoding protocol,structure-based methods are also developed, which can be further classiﬁed into ﬁngerprint-basedand graph neural networks (GNN)-based methods.
Revised: Beyond the above linear encoding protocol, structure-based methods are also developed, which can be further classiﬁed into ﬁngerprint-based and graph neural networks (GNN)-based methods.

Original: The molecular ﬁngerprint techniques date back tothe Morgan ﬁngerprints [40].
Revised: The molecular ﬁngerprint techniques date back to the Morgan ﬁngerprints [40].

Original: However, those ﬁngerprint-based methods are often handcrafted andnot trained in an end-to-end fashion [24].
Revised: However, those ﬁngerprint-based methods are often handcrafted and not trained in an end-to-end fashion [24].

Original: Since molecules can be viewed as structured graphs, graphneural networks have been widely used to learn molecule representation [26, 17, 22].
Revised: Since molecules can be viewed as structured graphs, graph neural networks have been widely used to learn molecule representation [26, 17, 22].

Original: Existing general OOD methods [3, 51, 16, 66, 46] are not tailored to such non-Euclidean structureddata, i.e. molecules.
Revised: Existing general OOD methods [3, 51, 16, 66, 46] are not tailored to such non-Euclidean structured data, i.e. molecules.

Original: In several recent works on molecule property classiﬁcation tasks [67, 55], theimportance of molecular substructures has been emphasized and such inductive bias is incorporatedinto the design of those models.
Revised: In several recent works on molecule property classiﬁcation tasks [67, 55], the importance of molecular substructures has been emphasized and such inductive bias is incorporated into the design of those models.

Original: However, they are still based on the i.i.d.
Revised: However, they are still based on the i.i.d.

Original: assumption and do notleverage those invariant substructure across different environments to achieve robust representations.
Revised: assumption and do not leverage those invariant substructure across different environments to achieve robust representations.

Original: In this paper, we propose a general framework orthogonal to these MRL studies to bridge OOD and
Revised: In this paper, we propose a general framework orthogonal to these MRL studies to bridge OOD and

Original: MRL, which can adopt any existing MRL methods as the backbone to improve their robustness.
Revised: MRL, which can adopt any existing MRL methods as the backbone to improve their robustness.

Original:  Methodology 
Revised:  Methodology 3.1 Problem Formulation

Original: 3.1 Problem Formulation

Original: We propose a OOD generalization framework tailored for molecule representation learning, entitled
Revised: We propose a OOD generalization framework tailored for molecule representation learning, entitled

Original: MoleOOD.
Revised: MoleOOD.

Original: All the random variables and the corresponding realizations are denoted as bold and thinletters, respectively.
Revised: All the random variables and the corresponding realizations are denoted as bold and thin letters, respectively.

Original: We ﬁrst formulate the OOD generalization problem for MRL.
Revised: We ﬁrst formulate the OOD generalization problem for MRL.

Original: OOD Generalization Problem on Molecule Representation Learning.
Revised: OOD Generalization Problem on Molecule Representation Learning.

Original: A molecular graph can berepresented as G =
Revised: A molecular graph can be represented as G =

Original: ( V, E ) , where V is the graph
Revised: ( V, E ) , where V is the graph

Original: ’s node set corresponding to atoms constituting themolecule and E denotes the graph’s edge sets corresponding to chemical bonds.
Revised: ’s node set corresponding to atoms constituting the molecule and E denotes the graph’s edge sets corresponding to chemical bonds.

Original: The training and testing molecule graph datasets are denoted as G train =
Revised: The training and testing molecule graph datasets are denoted as G train =

Original: { ( G i , y i ) } N train i = and G test =
Revised: { ( G i , y i ) } N train i = and G test =

Original: { ( G i , y i ) } N test i =1 .
Revised: { ( G i , y i ) } N test i =1 .

Original: Notice that the test dataset is drawn outside the distribution of the training dataset.
Revised: Notice that the test dataset is drawn outside the distribution of the training dataset.

Original: The goal ofmolecule representation learning task is to predict the target label y given the associated inputmolecule G .
Revised: The goal of molecule representation learning task is to predict the target label y given the associated input molecule G .

Original: Based on Eq.
Revised: Based on Eq.

Original: 1, we can formulate the OOD problem on MRL tasks as:
Revised: 1, we can formulate the OOD problem on MRL tasks as:

Original: [Equation]
Revised: [Equation]

Original: The difﬁculty of this problem is that the training data only cover very limited environments in E whilethe model is expected to perform well on all the environments.
Revised: The difﬁculty of this problem is that the training data only cover very limited environments in E while the model is expected to perform well on all the environments.

Original: Before delving into our solution, weﬁrst provide a causal perspective to understand the data-generating process and shed light on thelimitations of existing MRL models.

Original: 3.2 Understanding the Data-generating Process

Original: [Figure]

Original: We elaborate our approach in the context of molecule property classiﬁ- cation tasks in this paper.
Revised: We elaborate our approach in the context of molecule property classiﬁcation tasks in this paper.

Original: Recalling the two molecules Cyclopropanol( C 3 H 6 O ) and 1,4-Cyclohexanediol ( C 6 H 12 O 2 ) used for illustration in

Original: Sec. 1, they are sampled from different environments.

Original: Because both ofthem contain the hydroxy ( − OH ) , which we can call invariant or causalsubstructure in this case, these two molecules are readily soluble in water.

Original: We formalize such a date-generating process of molecule property prediction with a general Structural Causal Model (SCM) [18, 43] in Fig.

Original: The abstract data variablesare denoted by the nodes and the directed arrows represent the causalities

Original: This SCM illustrates

Original: the causalities among variables: E as the environment, S as the spurious substructures, C as the

Original: invariant/causal substructures w.r.t Y , G as the instance molecule graph, Y as the ground-truth label.

Original: • S ← E → C : the environmental variable impacts the underlying data generating distribution.

Original: Furthermore, substructures could be divided into causal and spurious ones across all environments.

Original: • S → G ← C : an instance molecule graph is made up of the causal and spurious substructures.

Original: • C → Y : Y , the ground-truth label, is only determined by C .

Original: This causation is the focus of our work.

Original: Taking Cyclopropanol ( C 3 H 6 O ) as an example, we can specify E as the 3C-ring scaffold, C as thesubstructure hydroxy ( − OH ) , S as the substructures aside from hydroxy, G as Cyclopropanol , Y asgood water solubility.

Original: The good water solubility Y is only attributed to the invariant substructurehydroxy, i.e. C , rather than other spurious substructures S .

Original: Existing MRL methods do not differentiate invariant and spurious substructures.
Revised: Existing MRL methods do not differentiate invariant and spurious substructures.

Original: Hence, the spuriouscorrelations between irrelevant substructuresS and the target labelY will be encoded to learnedmolecular representations.
Revised: Hence, the spurious correlations between irrelevant substructures and the target label will be encoded to learned molecular representations.

Original: When tested on unseen environments, the downstream classiﬁer will beeasily misled by these spurious correlations [58].
Revised: When tested on unseen environments, the downstream classiﬁer will be easily misled by these spurious correlations [58].

Original: With the knowledge that (bio)chemical propertiesof a molecule are usually associated with a few privileged substructures [29, 44, 68, 27], we aimto suppress such spurious correlations and leverage environment-invariant substructures that morestably relate with the labels across environments to learn invariant molecular representations.
Revised: With the knowledge that (bio)chemical properties of a molecule are usually associated with a few privileged substructures [29, 44, 68, 27], we aim to suppress such spurious correlations and leverage environment-invariant substructures that more stably relate with the labels across environments to learn invariant molecular representations.

Original: Noticethat the learned invariant molecular representations should satisfy the invariance principle mentionedin Sec.2.
Revised: Notice that the learned invariant molecular representations should satisfy the invariance principle mentioned in Sec.

Original: We next introduce our method formally and then give the instantiation of our model.
Revised: We next introduce our method formally and then give the instantiation of our model.

Original: 3.3 Model Formulation
Revised: 3.2 Model Formulation

Original: The framework contains two parts, the fronted molecule encoder Φ for learning an “invariantrepresentation” of theinput molecule graph and the back-end predictor ω for ﬁnal prediction. Solvingthe formulation in Eq.
Revised: The framework contains two parts, the fronted molecule encoder Φ for learning an “invariant representation” the formulation in Eq.

Original: 2 directly is intractable in practice since we cannot know all the environments, i.e, obtain a complete support set E .
Revised: 2 directly is intractable in practice since we cannot know all the environments, i.e, obtain a complete support set E .

Original: We resort to minimizing the expectation of risks from differentenvironments known in the training data,
Revised: We resort to minimizing the expectation of risks from different environments known in the training data,

Original: [Equation]
Revised: [Equation]

Original: where f = ω ◦ Φ and ⫫ denotes probabilistic independence.
Revised: where f = ω ◦ Φ and ⫫ denotes probabilistic independence.

Original: All learnable parameters of the moleculeencoder Φ and the predictor ω are included in θ .
Revised: All learnable parameters of the molecule encoder Φ and the predictor ω are included in θ .

Original: Different from Eq.
Revised: Different from Eq.

Original: 2, we add an extra invarianceconstraint y ⫫ e | Φ( G ) , which is used to suppress spurious correlations [9].
Revised: 2, we add an extra invariance constraint y ⫫ e | Φ( G ) , which is used to suppress spurious correlations [9].

Original: Since assessingcausality is challenging, we could rethink the problem on the basis of information theory.
Revised: Since assessing causality is challenging, we could rethink the problem on the basis of information theory.

Original: Recallthat we hope to let the molecule encoder leverage environment-invariant substructures and learn amolecular representation Φ( G ) given a molecule G .
Revised: Recall that we hope to let the molecule encoder leverage environment-invariant substructures and learn a molecular representation Φ( G ) given a molecule G .

Original: Our goal is to maximize the predictive powerof Φ( G ) on y , which can be measured by mutual information between Φ( G ) and y .
Revised: Our goal is to maximize the predictive power of Φ( G ) on y , which can be measured by mutual information between Φ( G ) and y .

Original: Meanwhile,probabilistic independence between y and e given Φ( G ) can be achieved via minimizing their mutualinformation.
Revised: Meanwhile, probabilistic independence between y and e given Φ( G ) can be achieved via minimizing their mutual information.

Original: For convenience, we denote Φ( G ) as z and Eq.
Revised: For convenience, we denote Φ( G ) as z and Eq.

Original: 3 can be approximately solved by:
Revised: 3 can be approximately solved by:

Original: [Equation]
Revised: [Equation]

Original: Treating the outputs of ω and Φ as distribution q θ ( z | G ) and q θ ( y | z ) respectively, Eq. 4 can bespeciﬁed as:
Revised: Treating the outputs of ω and Φ as distribution q θ ( z | G ) and q θ ( y | z ) respectively, Eq. 4 can be speciﬁed as:

Original: [Equation]
Revised: [Equation]

Original: Now, we have arrived at a clearer but still intractable optimization objective.
Revised: Now, we have arrived at a clearer but still intractable optimization objective.

Original: Before specifying thepractical instantiation of Eq. 5, let’s discuss on the environment variable e ﬁrst.
Revised: Before specifying the practical instantiation of Eq. 5, let’s discuss on the environment variable e ﬁrst.

Original: In practice, due to the non-trivial efforts to label the molecular environments, manual speciﬁcationsof the environments may be unavailable in many cases.
Revised: In practice, due to the non-trivial efforts to label the molecular environments, manual speciﬁcations of the environments may be unavailable in many cases.

Original: We may directly label molecules to differentenvironments in terms of their scaffolds when the environment label is unavailable.
Revised: We may directly label molecules to different environments in terms of their scaffolds when the environment label is unavailable.

Original: But this isunreasonable in practice, because the ﬁnal total environment number will be too large.
Revised: But this is unreasonable in practice, because the ﬁnal total environment number will be too large.

Original: Taking thedataset HIV for molecule property prediction tasks released by Open Graph Benchmark [21] forn example, OGB uses scaffold to split the molecules into different environments.
Revised: Taking the dataset HIV for molecule property prediction tasks released by Open Graph Benchmark [21] as an example, OGB uses scaffold to split the molecules into different environments.

Original: Assuming thatwe regard each scaffold as an environment directly, 41 , 127 molecules in HIV are partitioned into19 , 076 environments (see details in Appendix E).
Revised: Assuming that we regard each scaffold as an environment directly, 41 , 127 molecules in HIV are partitioned into 19 ,environments (see details in Appendix E).

Original: This environment count is much larger than other OOD datasets from other domains, e.g. Camelyon17 2 [4], CivilComments 3 [7], etc.
Revised: This environment count is much larger than other OOD datasets from other domains, e.g. Camelyon17 3 [4], CivilComments 4 [7], etc.

Original: Even thoughsome datasets may provide manual speciﬁcations of environments, the environment counts are alsotoo large, which is unfriendly to existing OOD models [25, 15].
Revised: Even though some datasets may provide manual speciﬁcations of environments, the environment counts are also too large, which is unfriendly to existing OOD models [25, 15].

Original: Therefore, we propose to design anenvironment-inference model ψ to partition the molecule into different environments with a relativelysmaller environment count.
Revised: Therefore, we propose to design an environment-inference model ψ to partition the molecule into different environments with a relatively smaller environment count.

Original: We denote the environment count as a hyper-parameter k .
Revised: We denote the environment count as a hyper-parameter k .

Original: Given prior p ( e | G ) , we need to maximize the log likelihood of p τ ( y | G ) and then obtain the posteriorp τ ( e | G , y ) , which are parameterized by τ .
Revised: Given prior p ( e | G ) , we need to maximize the log likelihood of p τ ( y | G ) and then obtain the posterior p τ ( e | G , y ) , which are parameterized by τ .

Original: Since there is no analytical solutions to the true posterior,here we use variational inference (VI) to approximate it.
Revised: Since there is no analytical solutions to the true posterior, here we use variational inference (VI) to approximate it.

Original: Speciﬁcally, we introduce a variationaldistribution q κ ( e | G , y ) parameterized by κ to approximate p τ ( e | G , y ) .
Revised: Speciﬁcally, we introduce a variational distribution q κ ( e | G , y ) parameterized by κ to approximate p τ ( e | G , y ) .

Original: Proposition 1.
Revised: Proposition 1.

Original: The Evidence Lower BOund (ELBO) of the observed molecule graph and corresponding label tuple ( G, y ) :
Revised: The Evidence Lower BOund (ELBO) of the observed molecule graph and corresponding label tuple ( G, y ) :

Original: L ( τ, κ ; ( G, y )) =
Revised: L ( τ, κ ; ( G, y )) =

Original: E q κ [log p τ ( y | G, e )] − D KL ( q κ ( e | G, y ) (cid:107) p τ ( e | G )) .
Revised: E q κ [log p τ ( y | G, e )] − D KL ( q κ ( e | G, y ) (cid:107) p τ ( e | G )) .

Original: Our goal is to minimize the Kullback-Leibler (KL) divergence between q κ ( e | G , y ) and p τ ( e | G , y ) ,i.e. D KL ( q κ ( e | G , y ) (cid:107) p τ ( e | G , y )) , which is equivalent to maximizing the ELBO in Proposition 1.
Revised: Our goal is to minimize the Kullback-Leibler (KL) divergence between q κ ( e | G , y ) and p τ ( e | G , y ) , i.e. D KL ( q κ ( e | G , y ) (cid:107) p τ ( e | G , y )) , which is equivalent to maximizing the ELBO in Proposition 1.

Original: Then, the objective used to train this environment-inference model is transformed to:
Revised: Then, the objective used to train this environment-inference model is transformed to:

Original: [Equation]
Revised: [Equation]

Original: Let’s look back to the objective given in Eq.
Revised: Let’s look back to the objective given in Eq.

Original: 5 and give an equivalent tractable objective in practical229 instantiation, which involves the environment-inference model deﬁned above:
Revised: 5 and give an equivalent tractable objective in practical instantiation, which involves the environment-inference model deﬁned above:

Original: [Equation]
Revised: [Equation]

Revised: [Figure]

Original: where G e consists of the pairs of molecular graph G and corresponding label y under environment e .
Revised: d ec o m po s e P r e d i c t o r s ub s t r u c t u r e s attentive pooling OH Environment

Original: [Figure]
Revised: Classifier

Original:  Theorem 1.
Revised: Conditional GNN O H where G e consists of the pairs of molecular graph G and corresponding label y under environment e . Theorem 1.

Original: With q θ ( y | z ) treated as a variational distribution, minimizing term 1 in Eq. 7contributes to min q θ ( y | z ) ,q θ ( z | G ) I( y ; e | z ) , letting z show equal performance for the downstreamtasks across all environments, i.e. p ( y | z , e ) = p ( y | z ) .
Revised: With q θ ( y | z ) treated as a variational distribution, minimizing term 1 in Eq.contributes to min q θ ( y | z ) ,q θ ( z | G ) I ( y ; e | z ) , letting z show equal performance for the downstream tasks across all environments, i.e. p ( y | z , e ) = p ( y | z ) .

Original: Theorem 2.
Revised: Theorem 2.

Original: Regarding q θ ( y | z ) as a variational distribution, minimizing term 2 in Eq.
Revised: Regarding q θ ( y | z ) as a variational distribution, minimizing term 2 in Eq.

Original: 7 equalsto max q θ ( y | z ) ,q θ ( z | G ) I
Revised: 7 equals to max q θ ( y | z ) ,q θ ( z | G ) I

Original: ( z ; y ) , letting z show sufﬁcient predictive power for downstream tasks.
Revised: ( z ; y ) , letting z show sufﬁcient predictive power for downstream tasks.

Original: Serving as theoretical justiﬁcations, Th. 1 and Th.
Revised: Serving as theoretical justiﬁcations, Th. 1 and Th.

Original: 2 reveal that optimizing the objective in Eq.
Revised: 2 reveal that optimizing the objective in Eq.

Original: 7 forcesthe learned representation z to satisfy the invariance principle mentioned in Sec. 2, thus ensuring avalid solution for OOD problem deﬁned in Eq.
Revised: 7 forces the learned representation z to satisfy the invariance principle mentioned in Sec. 2, thus ensuring a valid solution for OOD problem deﬁned in Eq.

Original: Due to the limited space, the detailed proofs can befound in Appendix B.
Revised: Due to the limited space, the detailed proofs can be found in Appendix B.

Original: 3.4 Model Instantiations and Training
Revised: 3.3 Model Instantiations and Training

Original: Environment-inference Module.
Revised: Environment-inference Module.

Original: For the approximate posterier model q κ ( e | G , y ) , in principle weshould design a module, entitled Environment Classiﬁer , that takes ( G , y ) as the input and outputsthe probabilistic distribution of e .
Revised: For the approximate posterier model q κ ( e | G , y ) , in principle we should design a module, entitled Environment Classiﬁer , that takes ( G , y ) as the input and outputs the probabilistic distribution of e .

Original: We use a Graph Isomorphism Network (GIN) [59], to learn a graphrepresentation given G .
Revised: We use a Graph Isomorphism Network (GIN) [59], to learn a graph representation given G .

Original: Then, the concatenation of this graph representation and label vector is fedto a feed-forward network to obtain a probabilistic distribution with regard to e .
Revised: Then, the concatenation of this graph representation and label vector is fed to a feed-forward network to obtain a probabilistic distribution with regard to e .

Original: We could set theprior p τ ( e | G ) to Uniform distribution  , encouraging the learned environment-partition to be uniform.
Revised: We could set the prior p τ ( e | G ) to Uniform distribution or Gaussian distribution , encouraging the learned environment-partition to be uniform.

Original: As for p τ ( y | G , e ) , we also choose a GNN model  followed by a softmax activation function to modelit.
Revised: As for p τ ( y | G , e ) , we also choose a GNN model (e.g. GIN) followed by a softmax activation function to model it.

Original: We call this module Conditional GNN because it conditions on e .
Revised: We call this module Conditional GNN because it conditions on e .

Original: It takes ( G , e ) as the inputand outputs the probabilistic distribution of y .
Revised: It takes ( G , e ) as the input and outputs the probabilistic distribution of y .

Original: The Molecule Encoder & The Final Predictor.
Revised: The Molecule Encoder & The Final Predictor.

Original: Recall that we aim to learn an invariant substructureaware molecular representation.
Revised: Recall that we aim to learn an invariant substructureaware molecular representation.

Original: Given a molecule G , we can choose any molecule representationlearning method to learn a representation r G for the complete molecular graph.
Revised: Given a molecule G , we can choose any molecule representation learning method to learn a representation r G for the complete molecular graph.

Original: This part is entitled Complete Encoder .
Revised: This part is entitled Complete Encoder .

Original: Meanwhile, we decompose the input molecule into a set of chemicalsubstructures using a molecule segmentation method, e.g. breaking retrosynthetically interestingchemical substructures (BRICS) [13], which is available as an API in RDKit [32].
Revised: Meanwhile, we decompose the input molecule into a set of chemical substructures using a molecule segmentation method, e.g. breaking retrosynthetically interesting chemical substructures (BRICS) [13], which is available as an API in RDKit [32].

Original: For each substructure, we consider using a simple GNN to learn a corresponding representation.
Revised: For each substructure, we consider using a simple GNN to learn a corresponding representation.

Original: We call this
Revised: We call this

Original: GNN Substructure Encoder .
Revised: GNN Substructure Encoder .

Original: Then, considering r G as a query with regard to substructures, weoperate attentive pooling on these substructure representations to obtain a new substructure-awaremolecular representation.
Revised: Then, considering r G as a query with regard to substructures, we operate attentive pooling on these substructure representations to obtain a new substructure-aware molecular representation.

Original: We then use this substructure-aware representation for downstream task.
Revised: We then use this substructure-aware representation for downstream task.

Original: Guided by our proposed learning objective, we can encode some invariant relationships betweencertain substructures and target properties into this representation.
Revised: Guided by our proposed learning objective, we can encode some invariant relationships between 15 k 0 .0 .0 .0 .0 . R O CAU C (a) 0. 2 β 0 .0 .0 .0 .0 .0 .0 .0 .

Original: The Complete Encoder, the{+R O CAU C (b)     100 Epoch 0 .0 .0 .
 R i s k Env1 Env+}
Revised: {+R O CAU C (b)     100 Epoch 0 .0 .0 .
 R i s k Env1 Env+}

Revised: Env3 Env (c) certain substructures and target properties into this representation.

Revised: The Complete Encoder, the

Original: Substructure encoder and the attentive pooling operation constitute our Molecule Encoder Φ .
Revised: Substructure encoder and the attentive pooling operation constitute our Molecule Encoder Φ .

Original: As forthe Predictor ω , we implement it with a multi-layer perceptron, followed by a softmax function.
Revised: As for the Predictor ω , we implement it with a multi-layer perceptron, followed by a softmax function.

Original: Theoverview of our model is demonstrated in Fig.
Revised: The overview of our model is demonstrated in Fig.

Original: Training.
Revised: Training.

Original: We adopt a simple yet efﬁcient two-stage training strategy to search for optimal parameters:
Revised: We adopt a simple yet efﬁcient two-stage training strategy to search for optimal parameters:

Original: 1) optimizing the environment-inference model: κ ∗ , τ ∗ ← arg max κ,τ L elbo ( τ, κ ; G train ) .
Revised: 1) optimizing the environment-inference model: κ ∗ , τ ∗ ← arg max κ,τ L elbo ( τ, κ ; G train ) .

Original: 2) optimizing the molecule encoder and the predictor: θ ∗ ← arg min θ L inv ( θ ; G train , τ ) .15 k 0 .0 .0 .0 .0 . R O CAU C (a) 0. 2 β 0 .0 .0 .0 .0 .0 .0 .0 .
Revised: 2) optimizing the molecule encoder and the predictor: θ ∗ ← arg min θ L inv ( θ ; G train , τ ) .

Original: R O CAU C (b)

Original: [-100 Epoch 0 .0 .0 .
-]
Revised: [-100 Epoch 0 .0 .0 .
-]

Original: R i s k

Original: Env1 Env

Original: Env3 Env

Original: [Figure]

Original:  Experiments
Revised:  Experiments

Original: Experiments are performed on 10 benchmark datasets.
Revised: Experiments are performed on 10 benchmark datasets.

Original: Each experiment is repeated 5 times withmean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPUand 128G RAM.
Revised: Each experiment is repeated 5 times with mean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPU and 128G RAM.

Original: The source code for reproducing the results will be made publicly available.
Revised: The source code for reproducing the results will be made publicly available.

Original: 4.1 Datasets and Setups
Revised: 4.1 Datasets and Setups

Original: Datasets and protocols.
Revised: Datasets and protocols.

Original: The four datasets BACE , BBBP , SIDER and HIV , are from by Open
Revised: The four datasets BACE , BBBP , SIDER and HIV , are from by Open

Original: Graph Benchmark (OGB) [21].
Revised: Graph Benchmark (OGB) [21].

Original: We use the default train/val/test split with ratio 8:1:1.
Revised: We use the default train/val/test split with ratio 8:1:1.

Original: Each splitcontains a set of scaffolds (almost) different to each other.
Revised: Each split contains a set of scaffolds (almost) different to each other.

Original: Hence we believe that to a certaindegree, it provides an OOD test-bed as different scaffold often suggest different data-generationenvironments.
Revised: Hence we believe that to a certain degree, it provides an OOD test-bed as different scaffold often suggest different data-generation environments.

Original: The other six datasets are generated by the dataset curator provided by DrugOOD [25].
Revised: The other six datasets are generated by the dataset curator provided by DrugOOD [25].

Original: DrugOOD provides more diverse splitting indicators than OGB, including assay, scaffold and size.
Revised: DrugOOD provides more diverse splitting indicators than OGB, including assay, scaffold and size.

Original: To comprehensively evaluate the performance of our method under different environment deﬁnitions,we adopt these three different splitting schemes on categories IC50 and EC50 provided in DrugOOD.
Revised: To comprehensively evaluate the performance of our method under different environment deﬁnitions, we adopt these three different splitting schemes on categories IC50 and EC50 provided in DrugOOD.

Original: Then we obtain six datasets, EC50- ∗ and IC50- ∗ , where the sufﬁx ∗ speciﬁes the splitting scheme i.e.
Revised: Then we obtain six datasets, EC50- ∗ and IC50- ∗ , where the sufﬁx ∗ speciﬁes the splitting scheme i.e.

Original: IC50/EC50-assay/scaffold/size .
Revised: IC50/EC50-assay/scaffold/size .

Original: Notice that only the six datasets from DrugOOD provide manualspeciﬁed environment labels.
Revised: Notice that only the six datasets from DrugOOD provide manual speciﬁed environment labels.

Original: Refer to Appendix E for more details of datasets.
Revised: Refer to Appendix E for more details of datasets.

Original: Metric.
Revised: Metric.

Original: As the concerned property prediction tasks all relate to classiﬁcation, we report the ROCAUC score which is also in line with previous MRL works [64, 60, 54].
Revised: As the concerned property prediction tasks all relate to classiﬁcation, we report the ROCAUC score which is also in line with previous MRL works [64, 60, 54].

Original: Baselines.
Revised: Baselines.

Original: Ideally, any MRL method can be adapted into our method as backbone to improve theirgeneralization ability against distribution shifts.
Revised: Ideally, any MRL method can be adapted into our method as backbone to improve their generalization ability against distribution shifts.

Original: We adapt three backbones: GCN [28], GIN [59]and GraphSAGE [19] into our method.
Revised: We adapt three backbones: GCN [28], GIN [59] and GraphSAGE [19] into our method.

Original: We compare the adapted version with the original method.
Revised: We compare the adapted version with the original method.

Original: We also compare against another augmented version “ + virtual node ” [17, 23, 36].
Revised: We also compare against another augmented version “ + virtual node ” [17, 23, 36].

Original: Furthermore,we compare our method with six OOD generalization methods on MRL tasks: ERM [52], IRM [3],
Revised: Furthermore, we compare our method with six OOD generalization methods on MRL tasks: ERM [52], IRM [3],

Original: DeepCoral [51], DANN [16], MixUp [66] and GroupDro [46].
Revised: DeepCoral [51], DANN [16], MixUp [66] and GroupDro [46].

Original: Due to the fact that most of thesemethods require the manual speciﬁcation of environments in dataset, we report this comparison ondatasets from DrugOOD only.
Revised: Due to the fact that most of these methods require the manual speciﬁcation of environments in dataset, we report this comparison on datasets from DrugOOD only.

Original: Each of the method is conﬁgured using the same parameters reportedin the original paper or selected by grid search.
Revised: Each of the method is conﬁgured using the same parameters reported in the original paper or selected by grid search.

Original: For the sake of fairness, the embeddingsize of allmethods are set to be equal in comparison.
Revised: For the sake of fairness, the embedding

Original: We specify the training details in the Appendix C.
Revised: [Table] [Table]

Original: 4.2 Performance Comparison
Revised: 4.2 Performance Comparison

Original: Improvements to existing MRL methods.
Revised: Improvements to existing MRL methods.

Original: As demonstrated in Table 1, baselines obtain consistentimprovements after adapted to our methods across all the four datasets released by OGB in termsof ROC-AUC.
Revised: As demonstrated in Table 1, baselines obtain consistent improvements after adapted to our methods across all the four datasets released by OGB in terms of ROC-AUC.

Original: Our method also beats the augmented version, “+ virtual node” , of baselines onall datasets, i.e. adding a virtual node.
Revised: Our method also beats the augmented version, “+ virtual node” , of baselines on all datasets, i.e. adding a virtual node.

Original: The results indicate that, orthogonal to prior studies on
Revised: The results indicate that, orthogonal to prior studies on

Original: MRL, our method is a general framework which can incorporate existing MRL methods and improvetheir generalization ability for OOD data.
Revised: MRL, our method is a general framework which can incorporate existing MRL methods and improve their generalization ability for OOD data.

Original: We attribute the superior performance of our method inmolecular properties predictions under OOD setting to that, our proposed learning objective enforcesthe model to learn environment-invariant representations against distribution shifts.
Revised: We attribute the superior performance of our method in molecular properties predictions under OOD setting to that, our proposed learning objective enforces the model to learn environment-invariant representations against distribution shifts.

Original: Superiority to other OOD generalization methods.
Revised: Superiority to other OOD generalization methods.

Original: Table 2 summarizes the results in comparsionwith six state-of-the-art methods tailored for OOD learning, where we obtain the following observations.
Revised: Table 2 summarizes the results in comparsion with six state-of-the-art methods tailored for OOD learning, where we obtain the following observations.

Original: Except on IC50-size, our method outperforms all baselines across all datasets due to its ability [Table] except ERM and MixUp all require environment labels. All methods including ours use GIN [59] as backbones.
Revised: Except on IC50-size, our method outperforms all baselines across all datasets due to its ability to enforce the molecule encoder to leverage environment-invariant substructures that more stably relate with the labels across environments.

Original: [Table]

Original: to enforce the molecule encoder to leverage environment-invariant substructures that more stablyrelate with the labels across environments.

Original: Our method ranks the third on IC50-size after MixUp and
Revised: Our method ranks the third on IC50-size after MixUp and

Original: ERM.
Revised: ERM.

Original: Different from the other methods, MixUp constructs more training exmaples and uses moredata to train the model.
Revised: Different from the other methods, MixUp constructs more training exmaples and uses more data to train the model.

Original: That’s why MixUp obtains best performance among all methods on IC50-sizein our analysis.
Revised: That’s why MixUp obtains best performance among all methods on IC50-size in our analysis.

Original: As for ERM, [25, 15] have pointed out that simple ERM shows better performancecompared to subsequent OOD methods when datasets have relatively large environment counts.
Revised: As for ERM, [25, 15] have pointed out that simple ERM shows better performance compared to subsequent OOD methods when datasets have relatively large environment counts.

Original: Eventhough we have set the environment number k to a smaller value than the ground-truth number givenby the dataset, we still need to prevent k from being too small (see discussion in Sec. 4.4), leading toour poorer performance than ERM on IC50-size. 
Revised: Even though we have set the environment number k to a smaller value than the ground-truth number given by the dataset, we still need to prevent k from being too small (see discussion in Sec. 4.4), leading to our poorer performance than ERM on IC50-size. [Table]

Original: 4.3 Ablation Study
Revised: 4.3 Ablation Study

Original: We analyze the contributions of different model components to the ﬁnal performance in this section.
Revised: We analyze the contributions of different model components to the ﬁnal performance in this section.

Original: Table 3 reports detailed ablation experimental results on EC50-assay, EC50-scaffold and EC50-size.
Revised: Table 3 reports detailed ablation experimental results on EC50-assay, EC50-scaffold and EC50-size.

Original: Attention-based architecture.
Revised: Attention-based architecture.

Original: We study the impact of the attention-based architecture introducedin Sec. 3.4 by assembling this architecture with ERM loss.
Revised: We study the impact of the attention-based architecture introduced in Sec. 3.3 by assembling this architecture with ERM loss.

Original: We beat ERM and MixUp only with thisarchitecture on three datasets.
Revised: We beat ERM and MixUp only with this architecture on three datasets.

Original: The results show that learning a representation for each substructureand then attentively aggregating these learned representations to obtain a ﬁnal substructure-awarerepresentation performs better than learning a representation for a complete molecular graph directly.
Revised: The results show that learning a representation for each substructure and then attentively aggregating these learned representations to obtain a ﬁnal substructure-aware representation performs better than learning a representation for a complete molecular graph directly.

Original: This veriﬁes our assumption that the substructure perspective is of importance to boosting performanceof existing MRL methods.
Revised: This veriﬁes our assumption that the substructure perspective is of importance to boosting performance of existing MRL methods.

Original: With the aid of such a substructure-grained learning architecture, theimpact of our learning objective can be further strengthened.
Revised: With the aid of such a substructure-grained learning architecture, the impact of our learning objective can be further strengthened.

Original: New learning objective.
Revised: New learning objective.

Original: To evaluate the impact of our proposed new learning objective, we equip
Revised: To evaluate the impact of our proposed new learning objective, we equip

Original: GIN with this new objective.
Revised: GIN with this new objective.

Original: We can see compared to using the substructure-grained learning architecture only, only using the proposed new learning objective can bring more signiﬁcant improvement.
Revised: We can see compared to using the substructure-grained learning architecture only, only using the proposed new learning objective can bring more signiﬁcant improvement.

Original: Thus, we can attribute the main superiority of our full model to this new objective.
Revised: Thus, we can attribute the main superiority of our full model to this new objective.

Original: Combined withthe architecture discussed above, the new objective is able to better guide the molecule encoder tolearn environment-invariant molecular representations against distribution shifts.[Table]
Revised: Combined with the architecture discussed above, the new objective is able to better guide the molecule encoder to learn environment-invariant molecular representations against distribution shifts.

Original: Environment inference.
Revised: Environment inference.

Original: Now we turn to investigate the performance with respect to our proposedenvironment-inference module.
Revised: Now we turn to investigate the performance with respect to our proposed environment-inference module.

Original: One motivation for this module is that in reality manual speciﬁcationsof environments may be unavailable due to the high price for labeling environments by experts.
Revised: One motivation for this module is that in reality manual speciﬁcations of environments may be unavailable due to the high price for labeling environments by experts.

Original: Butwhen environment labels are available [Table] how will be performance be like if directly utilizing the givenenvironment partition?
Revised: But when environment labels are available, how will be performance be like if directly utilizing the given environment partition?

Original: An ablation study is targeted on this.
Revised: An ablation study is targeted on this.

Original: Taking the EC50-assay dataset as anexample, it has given the environment partition and it speciﬁes 47 environments in total.
Revised: Taking the EC50-assay dataset as an example, it has given the environment partition and it speciﬁes 47 environments in total.

Original: We utilizethe given environment partition directly and keep the remaining parts in line with our full model.
Revised: We utilize the given environment partition directly and keep the remaining parts in line with our full model.

Original: The results show that utilizing the given environment label, our method still can beat ERM and
Revised: The results show that utilizing the given environment label, our method still can beat ERM and

Original: MixUp.
Revised: MixUp.

Original: But compared to our full model where we set the environment number k to 20 , it obtainsinferior performance.
Revised: But compared to our full model where we set the environment number k to 20 , it obtains inferior performance.

Original: Additionally, to further examine the effectiveness of our proposed environmentinference method, we relabel the environment for each molecule for DANN according to our inferredenvironment partition.
Revised: Additionally, to further examine the effectiveness of our proposed environment inference method, we relabel the environment for each molecule for DANN according to our inferred environment partition.

Original: We can see that based on the new environment partition, DANN obtains betterperformance than using the initial given environment labels across three datasets.
Revised: We can see that based on the new environment partition, DANN obtains better performance than using the initial given environment labels across three datasets.

Original: The reason whyinferring environment instead can outperform directly using the given environment label is mainlydue to the existing given partitions are often handcraftedrule-based and not structured.
Revised: The reason why inferring environment instead can outperform directly using the given environment label is mainly due to the existing given partitions are often handcraftedrule-based and not structured.

Original: In contrast,letting the model learn a environment partition by itself may be more effective to some degree.
Revised: In contrast, letting the model learn a environment partition by itself may be more effective to some degree.

Original: 4.4 Hyper-parameter Sensitivity & Risk Dynamics
Revised: 4.4 Hyper-parameter Sensitivity & Risk Dynamics

Original: We investigate the sensitivity of our method to these two hyper-parameters: the speciﬁed numberof environments k , the trading-off parameter β in Eq.
Revised: We investigate the sensitivity of our method to these two hyper-parameters: the speciﬁed number of environments k , the trading-off parameter β in Eq.

Original: Fig.
Revised: Fig.

Original: 4(a) shows the performance regardingdifferent environment number k .
Revised: 3(a) shows the performance regarding different environment number k .

Original: It shows that the performance of our methods degrades when k istoo small (e.g. k = 1 , 5 ) or too large (e.g. k = 15 ).
Revised: It shows that the performance of our methods degrades when k is too small (e.g. k = 1 , 5 ) or too large (e.g. k = 15 ).

Original: When k = 1 i.e. we regard all training dataas from only one environment, the performance is the poorest.
Revised: When k = 1 i.e. we regard all training data as from only one environment, the performance is the poorest.

Original: This justiﬁes that partitioning thetraining samples into different environments is necessary.
Revised: This justiﬁes that partitioning the training samples into different environments is necessary.

Original: Fig. 4(b) shows the results of our methodby varying the trade-off parameter β .
Revised: Fig. 3(b) shows the results of our method by varying the trade-off parameter β .

Original: Our method obtains the worst performance when β = 0 .
Revised: Our method obtains the worst performance when β = 0 .

Original: Thisis mainly because Eq. 7 is reduced to the ﬁrst term when β
Revised: This is mainly because Eq. 7 is reduced to the ﬁrst term when β

Original: = 0 .
Revised: = 0 .

Original: According to Theorem 2, without thesecond term of Eq. 7, the sufﬁciency condition of invariance principle cannot be satisﬁed, resultingin the performance degradation.
Revised: According to Theorem 2, without the second term of Eq. 7, the sufﬁciency condition of invariance principle cannot be satisﬁed, resulting in the performance degradation.

Original: Additionally, to shed insights of the ability of our method to lowerthe risks of different environments, we visualize the risk dynamic curve of some environments in
Revised: Additionally, to shed insights of the ability of our method to lower the risks of different environments, we visualize the risk dynamic curve of some environments in

Original: Fig. 4(c).
Revised: Fig. 3(c).

Original: As is shown in Fig. 4(c), the difﬁculties of decreasing the risk on different environmentsare different.
Revised: As is shown in Fig. 3(c), the difﬁculties of decreasing the risk on different environments are different.

Original: Though the risks of some environments vibrate violently at the beginning of trainingprocess (e.g. Env 1 and Env 2), with time elapsing, risks on all environments can decrease stably.
Revised: Though the risks of some environments vibrate violently at the beginning of training process (e.g. Env 1 and Env 2), with time elapsing, risks on all environments can decrease stably.

Original:  Conclusion
Revised:  Conclusion

Original: We have proposed a general framework which can incorporate any existing MRL method as backboneto improve their generalization ability against distribution shifts.
Revised: We have proposed a general framework which can incorporate any existing MRL method as backbone to improve their generalization ability against distribution shifts.

Original: Speciﬁcally, we devise a newlearning scheme with its equivalent practical instantiation.
Revised: Speciﬁcally, we devise a new learning scheme with its equivalent practical instantiation.

Original: We also develop an environment inferencemodel to identify each molecule’s corresponding environment without need of manual speciﬁcationsof environments.
Revised: We also develop an environment inference model to identify each molecule’s corresponding environment without need of manual speciﬁcations of environments.

Original: Extensive experimental results on ten datasets demonstrate that our model yieldsconsistent and signiﬁcant improvements over various existing MRL methods as backbones.
Revised: Extensive experimental results on ten datasets demonstrate that our model yields consistent and signiﬁcant improvements over various existing MRL methods as backbones.

Original: Additionally, our model achieves competitive or even superior performance compared to state-of-the-artmodels designed for OOD learning that require manual speciﬁed environment labels as extra inputs.Checklist For all authors...
Revised: Additionally, our model achieves competitive or even superior performance compared to state-of-the-art models designed for OOD learning that require manual speciﬁed environment labels as extra inputs.

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes] See Appendix I.

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] SeeAppendix J.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results? [Yes] (b) Did you include complete proofs of all theoretical results? [Yes] See Appendix Aand B.
Revised: Acknowledgement

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [Yes] See Appendix C.

Original: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes] See Section 4 and Appendix C.

Original: (c) Did you report error bars (e

Original: g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [Yes] See Section 4.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Section 4.

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators? [Yes] See Section 4 and
Revised: This work was partly supported by National Key Research and Development Program of China (2020AAA0107600), National Natural Science Foundation of China (61972250, 72061127003), and

Original: Appendix E.

Original: (b) Did you mention the license of the assets?

Original: [Yes] See Section 4 and

Original: Appendix E.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [Yes] See Appendix E.

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?

Original: [Yes] See Appendix E.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
Revised: Shanghai Municipal Science and Technology Project (22511105100, 2021SHZDZX0102).

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]

Original: A Proof for Proposition1
Revised: A Proof for Proposition

Original: Proof.
Revised: Proof.

Original: Our goal is to minimize the Kullback-Leibler (KL) divergence between q κ ( e | G , y ) andp τ ( e | G , y ) .
Revised: Our goal is to minimize the Kullback-Leibler (KL) divergence between q κ ( e | G , y ) and p τ ( e | G , y ) .

Original: For the observed molecule graph and corresponding label tuple ( G, y ) ,
Revised: For the observed molecule graph and corresponding label tuple ( G, y ) ,

Original: [Equation]
Revised: [Equation]

Original: Rearrange Eq.
Revised: Rearrange Eq.

Original: and we can get,
Revised: and we can get,

Original: [Equation]
Revised: [Equation]

Original: The deﬁned L ( τ, κ ; G, y ) is called Evidence Lower BOund (ELBO)
Revised: The deﬁned L ( τ, κ ; G, y ) is called Evidence Lower BOund (ELBO)

Original: According to Eq. 9,maximizing this ELBO is equivalent to minimizing the KL divergence and maximizing log p τ ( y | G ) .
Revised: According to Eq. 9, maximizing this ELBO is equivalent to minimizing the KL divergence and maximizing log p τ ( y | G ) .

Original: For the observed molecule graph and corresponding label tuple ( G, y ) , we obtain the ELBO:[Equation]
Revised: For the observed molecule graph and corresponding label tuple ( G, y ) , we obtain the ELBO:

Original: [Equation]
Revised: [Equation]

Original: We thus conclude the proof.
Revised: We thus conclude the proof.

Original: B Proofs for Theorems
Revised: B Proofs for Theorems

Original: In this paper, we extend the invariance assumption [45, 3] to molecule representation learning:
Revised: In this paper, we extend the invariance assumption [45, 3] to molecule representation learning:

Original: Assumption 1.
Revised: Assumption 1.

Original: Given a molecular graph G , there exists an encoder Φ yielding a graph-levelrepresentation r G ∈ R d .
Revised: Given a molecular graph G , there exists an encoder Φ yielding a graph-level representation r G ∈ R d .

Original: Deﬁne r as a random variable of r G and it satisﬁes: 1) (Invariancecondition): p ( y | r , e ) = p ( y | r ) , and 2) (Sufﬁciency condition): y = h ( r ) + n , where h is a nonlinear function, n is a independent noise.
Revised: Deﬁne r as a random variable of r G and it satisﬁes: 1) (Invariance condition): p ( y | r , e ) = p ( y | r ) , and 2) (Sufﬁciency condition): y = h ( r ) + n , where h is a nonlinear function, n is a independent noise.

Original: With the terminology of information theory, we present a useful lemma [57] that interprets theinvariance and sufﬁciency conditions in Assumption 1:
Revised: With the terminology of information theory, we present a useful lemma [57] that interprets the invariance and sufﬁciency conditions in Assumption 1:

Original: Lemma 1.
Revised: Lemma 1.

Original: In terms of information theory, the two conditions in Assumption 1 can be equivalentlyexpressed as 1) invariance: I( y ; e | r ) = 0 and 2) sufﬁciency: I( y ; r ) is maximized.
Revised: In terms of information theory, the two conditions in Assumption 1 can be equivalently expressed as 1) invariance: I( y ; e | r ) = 0 and 2) sufﬁciency: I( y ; r ) is maximized.

Original: Proof.
Revised: Proof.

Original: For the invariance, it can be obtained by the fact that
Revised: For the invariance, it can be obtained by the fact that

Original: [Equation]
Revised: [Equation]

Original: For the sufﬁciency, we ﬁrst prove that every triplet ( G , r , y ) satisfying that y = h ( r ) + n would alsosatisfy r = arg max r I( y ; r ) .
Revised: For the sufﬁciency, we ﬁrst prove that every triplet ( G , r , y ) satisfying that y = h ( r ) + n would also satisfy r = arg max r I( y ; r ) .

Original: We prove it by contradiction.
Revised: We prove it by contradiction.

Original: Assume that r (cid:54) =
Revised: Assume that r (cid:54) =

Original: arg max r I( y ; r ) andthere exists r (cid:48) with r (cid:48) 
Revised: arg max r I( y ; r ) and there exists r (cid:48) with r (cid:48) =

Original: =

Original: arg max r I( y ; r ) with r (cid:54) =
Revised: arg max r I( y ; r ) with r (cid:54) =

Original: r (cid:48) .
Revised: r (cid:48) .

Original: Then there exists another random variable˜ r and a mapping function f m such that r (cid:48)
Revised: Then there exists another random variable ˜ r and a mapping function f m such that r (cid:48)

Original: =
Revised: =

Original: f m ( r , ˜ r ) .
Revised: f m ( r , ˜ r ) .

Original: Then we will have I( y ; r (cid:48) )
Revised: Then we will have I( y ; r (cid:48) )

Original: =
Revised: =

Original: I( y ; r , ˜ r ) = I( h ( r ); r , ˜ r ) =
Revised: I( y ; r , ˜ r ) = I( h ( r ); r , ˜ r ) =

Original: I( h ( r ); r ) =
Revised: I( h ( r ); r ) =

Original: I( y ; r ) , which leads to contradiction.
Revised: I( y ; r ) , which leads to contradiction.

Original: Then we prove that every triplet ( G , r , y ) satisfying that r = arg max r
Revised: Then we prove that every triplet ( G , r , y ) satisfying that r = arg max r

Original: I( y ; r ) would also satisfyy = h ( r )+ n by contradiction.
Revised: I( y ; r ) would also satisfy y = h ( r )+ n by contradiction.

Original: Suppose that y (cid:54) = h ( r )+ n and there exists r (cid:48) (cid:54) =
Revised: Suppose that y (cid:54) = h ( r )+ n and there exists r (cid:48) (cid:54) =

Original: r with y = h ( r (cid:48) )+ n .
Revised: r with y = h ( r (cid:48) )+ n .

Original: Then the inequality I( h ( r (cid:48) ); r ) <
Revised: Then the inequality I( h ( r (cid:48) ); r ) <

Original: I( h ( r (cid:48) ); r (cid:48) ) holds.
Revised: I( h ( r (cid:48) ); r (cid:48) ) holds.

Original: That means r (cid:48)
Revised: That means r (cid:48)

Original: =
Revised: =

Original: arg max r I( y ; r ) , leading tocontradiction.
Revised: arg max r I( y ; r ) , leading to contradiction.

Original: B.1 Proof for Theorem
Revised: B.1 Proof for Theorem

Original: Proof.
Revised: Proof.

Original: According to the dependency relationship z ← G → y , we have
Revised: According to the dependency relationship z ← G → y [Equation] we have

Original: [Equation]
Revised: [Equation]

Original: Next, we have
Revised: Next, we have

Original: [Equation]
Revised: [Equation]

Original: Based on
Revised: Based on

Original: Jensen
Revised: Jensen

Original: Inequality and
Revised: Inequality and

Original: Triangle
Revised: Triangle

Original: Inequality, we can obtain that
Revised: Inequality, we can obtain that

Original: D KL (cid:0) q ( y | z ) (cid:107) E p ( e | G )
Revised: D KL [Equation] e [Equation]

Revised: G )

Revised: [Equation]

Revised: p

Original: [ p ( y | G , e )] (cid:1) is upper bounded by:
Revised: [Equation] y | G [Equation] e )

Original: [Equation]
Revised: [Equation]

Revised: is upper bounded by:

Original: [Equation]
Revised: [Equation]

Original: Thus we can prove that minimizing term 1 in Eq.
Revised: Thus we can prove that minimizing term 1 in Eq.

Original: 7 is equivalent to min q θ ( y | z ) ,q θ ( z | G ) I
Revised: 7 is equivalent to min q θ ( y | z ) ,q θ ( z | G ) I

Original: ( y ; e | z ) .
Revised: ( y ; e | z ) .

Original: B.2 Proof for Theorem
Revised: B.2 Proof for Theorem

Original: Proof.
Revised: Proof.

Original: Given the dependency relationship z ← G → y , we hold max q ( z | G ) I( y ; z ) is equivalent tomin q ( z | G ) I( y ; G | z ) .
Revised: Given the dependency relationship z ← G → y [Equation] we hold max q ( z | G ) I( y ; z ) is equivalent to min q ( z | G ) I( y ; G | z ) .

Original: Also we have
Revised: Also we have

Original: [Equation]
Revised: [Equation]

Revised: Based on this

Original: [Equation]
Revised: [Equation]

Original: Based on this, we will have
Revised: we will have

Original: [Equation]
Revised: [Equation]

Original: Then we can also derive the following inequality via Jensen Inequality:
Revised: Then we can also derive the following inequality via Jensen Inequality:

Original: [Equation]
Revised: [Equation]

Original: where C is a constant.
Revised: where C is a constant.

Original: Then the problem min q ( y | z ) D KL ( p [Equation] y | G , e [Equation]
Revised: Then the problem min q ( y | z ) D KL ( p ( y | G , e ) (cid:107) q ( y | z )) can be solve by

Original: q ( y | z )) can be solve by

Original: [Equation]
Revised: [Equation]

Original: which means minimizing term 2 in Eq.
Revised: which means minimizing term 2 in Eq.

Original: 7 contributes to max q θ ( y | z ) ,q θ ( z | G ) I
Revised: 7 contributes to max q θ ( y | z ) ,q θ ( z | G ) I

Original: ( z ; y ) .
Revised: ( z ; y ) .

Original: C Implementation DetailsC.1 Baselines
Revised: C Implementation Details C.1 Baselines

Original: This section describes training conﬁgurations for all baselines, which are compared in this paper.
Revised: This section describes training conﬁgurations for all baselines, which are compared in this paper.

Original: Three backbones.
Revised: Three backbones.

Original: We adapt three backbones into our method, namely, GCN [28], GraphSAGE [19] and GIN [59].
Revised: We adapt three backbones into our method, namely, GCN [28], GraphSAGE [19] and GIN [59].

Original: We also emphasize the comparison with their augmented versions,i.e. “ + virtual node ” [17, 23, 36].
Revised: We also emphasize the comparison with their augmented versions, i.e. “ + virtual node ” [17, 23, 36].

Original: For GCN and GIN, we use the implementations provided by Open
Revised: For GCN and GIN, we use the implementations provided by Open

Original: Graph Benchmark [21] 4 .
Revised: Graph Benchmark [21] 5 .

Original: We implement GraphSAGE and its corresponding augmented version byourselves.
Revised: We implement GraphSAGE and its corresponding augmented version by ourselves.

Original: For these baselines, grid search of learning rate over { 1 e − 2 , 5 e − 3 , 1 e − 3 , 5 e − 4 , 1 e − 4 }and dropout rate over { 0 . 1 , 0 . 3 , 0 . 5 } is performed to select the best parameters.
Revised: For these baselines, grid search of learning rate over { 1 e − 2 , 5 e − 3 , 1 e − 3 , 5 e − 4 , 1 e − 4 } and dropout rate over { 0 . 1 , 0 . 3 , 0 . 5 } is performed to select the best parameters.

Original: The embedding sizeof all methods including ours are all set to 256 for the sake of fairness.
Revised: The embedding size of all methods including ours are all set to 256 for the sake of fairness.

Original: • GCN [28] is a scalable approach on graph-structured data that is based on an efﬁcient variant of convolutional neural networks.
Revised: • GCN [28] is a scalable approach on graph-structured data that is based on an efﬁcient variant of convolutional neural networks.

Original: • GIN [59] generalizes the Weisfeiler-Lehman (WL) graph isomorphism test [34] and hence achievesmaximum discriminative power among GNNS.• GraphSAGE [19] learns a function that generates embeddings by sampling and aggregatingfeatures from a node’s local neighborhood.• GCN/GIN/GraphSAGE + virtual node [17, 23, 36] is a variant of the original method augmentedby an additional node connecting to all nodes in the raw graph.
Revised: • GIN [59] generalizes the Weisfeiler-Lehman (WL) graph isomorphism test [34] and hence achieves maximum discriminative power among GNNS. • GraphSAGE [19] learns a function that generates embeddings by sampling and aggregating features from a node’s local neighborhood. • GCN/GIN/GraphSAGE + virtual node [17, 23, 36] is a variant of the original method augmented by an additional node connecting to all nodes in the raw graph.

Original: Models tailored for OOD learning.
Revised: Models tailored for OOD learning.

Original: We compare our method against six state-of-the-art methods:
Revised: We compare our method against six state-of-the-art methods:

Original: ERM [52], IRM [3], DeepCoral [51], DANN [16], MixUp [66] and GroupDro [46].
Revised: ERM [52], IRM [3], DeepCoral [51], DANN [16], MixUp [66] and GroupDro [46].

Original: We use theimplementations of these six method provided by DrugOOD 5 .
Revised: We use the implementations of these six method provided by DrugOOD 6 .

Original: We search for the optimal hyperparameters by ranging learning rate over { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 } and dropout rateover { 0 . 1 , 0 . 3 , 0 . 5 } .
Revised: We search for the optimal hyperparameters by ranging learning rate over { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 } and dropout rate over { 0 . 1 , 0 . 3 , 0 . 5 } .

Original: The embedding size of all models including ours are all set to 128 for fairness.• ERM [52] minimizes the average empirical loss on training data.• IRM [3] penalizes feature distributions for environments that have different optimum predictors.
Revised: The embedding size of all models including ours are all set to 128 for fairness. • ERM [52] minimizes the average empirical loss on training data. • IRM [3] penalizes feature distributions for environments that have different optimum predictors.

Original: We set the penalty weight and the penalty anneal iteration to 10 and 500 , respectively.• DeepCoral [51] penalizes differences in the means and covariances of the feature distributions foreach environment, which are exactly the distribution of last layer activations in a neural network.
Revised: We set the penalty weight and the penalty anneal iteration to 10 and 500 , respectively. • DeepCoral [51] penalizes differences in the means and covariances of the feature distributions for each environment, which are exactly the distribution of last layer activations in a neural network.

Original: The penalty weight is set to 0 .
Revised: The penalty weight is set to 0 .

Original:  DANN [16] encourages feature representations to be consistent across domains.
Revised: • DANN [16] encourages feature representations to be consistent across domains.

Original: We set to theinverse factor to 0 .
Revised: We set to the inverse factor to 0 .

Original: 2 .• MixUp [66] constructs additional virtual samples for training from two examples which arerandomly sampled from the training data.
Revised: • MixUp [66] constructs additional virtual samples for training from two examples which are randomly sampled from the training data.

Original: We set the probability and interpolate strength to 0 .
Revised: We set the probability and interpolate strength to 0 .

Original:  GroupDro [46] minimizes the worst-case training loss over a set of pre-deﬁned environments.
Revised: 1 . • GroupDro [46] minimizes the worst-case training loss over a set of pre-deﬁned environments.

Original: The step size is set to 0 .
Revised: The step size is set to 0 .

Original: C.2 Our Method
Revised: C.2 Our Method

Original: We implement our method in Pytorch.
Revised: We implement our method in Pytorch.

Original: As for experiments on OGB datasets, we implement the
Revised: As for experiments on OGB datasets, we implement the

Original: Environment Classiﬁer, the Conditional GNN and the Substructure Encoder which are mentionedin Sec.
Revised: Environment Classiﬁer, the Conditional GNN and the Substructure Encoder which are mentioned in Sec.

Original: 3.4 all in Graph Isomorphism Network (GIN) [59].
Revised: 3.3 all in Graph Isomorphism Network (GIN) [59].

Original: We use grid search on validation set forhyper-parameter tuning by ranging learning rate from { 1 e − 2 , 5 e − 3 , 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e −5 , 1 e − 5 } , dropout rate from { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 } , the trading-off parameter β from { 0 . 5 , 1 , 2 , 4 } ,the environment count k from { 5 , 10 , 15 , 20 , 40 , 80 } .
Revised: We use grid search on validation set for hyper-parameter tuning by ranging learning rate from { 1 e − 2 , 5 e − 3 , 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 } , dropout rate from { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 } , the trading-off parameter β from { 0 . 5 , 1 , 2 , 4 } , the environment count k from { 5 , 10 , 15 , 20 , 40 , 80 } .

Original: As for the prior p ( e | G ) , we set it to a Uniformdistribution or a discrete Gaussian distribution.
Revised: As for the prior p ( e | G ) , we set it to a Uniform distribution or a discrete Gaussian distribution.

Original: We use CrossEntropyLoss for all models and the
Revised: We use CrossEntropyLoss for all models and the

Original: Adam optimizer is used for gradient-based optimization.
Revised: Adam optimizer is used for gradient-based optimization.

Revised: 5 https://github.com/snap-stanford/ogbhttps://github

Revised: com/tencent

Original: D Training Procedure
Revised: D Training Procedure

Original: The training procedure of our method is summarized in Algorithm 1.Algorithm 1: The training procedure. parameters do for do environment
Revised: The training procedure of our method is summarized in Algorithm 1.

Original: Compute
Revised: Algorithm

Original: Backpropagate L

Original: 14 Output

Original: E More Details of Datasets
Revised: E More Details of Datasets

Original: In this paper, we use ten publicly available benchmark datasets in total.
Revised: In this paper, we use ten publicly available benchmark datasets in total.

Original: Four of them, namely,
Revised: Four of them, namely,

Original: BACE, BBBP, SIDER and HIV are released by Open Graph Benchmark (OGB) [21].
Revised: BACE, BBBP, SIDER and HIV are released by Open Graph Benchmark (OGB) [21].

Original: The rest six arereleased by DrugOOD [25], i.e. IC50-assay, IC50-scaffold, IC50-size, EC50-assay, EC50-scaffoldand EC50-size.
Revised: The rest six are released by DrugOOD [25], i.e. IC50-assay, IC50-scaffold, IC50-size, EC50-assay, EC50-scaffold and EC50-size.

Original: We provide detailed descriptions for them as below.
Revised: We provide detailed descriptions for them as below.

Original: • BBBP is a dataset of Brain-Blood Barrier Penetration.
Revised: • BBBP is a dataset of Brain-Blood Barrier Penetration.

Original: Each molecule has a label indicatingwhether it can penetrate through brain cell membrane to enter central nervous system.
Revised: Each molecule has a label indicating whether it can penetrate through brain cell membrane to enter central nervous system.

Original: • BACE is a dataset of binding afﬁnity against human beta-secretas 1.
Revised: • BACE is a dataset of binding afﬁnity against human beta-secretas 1.

Original: Each molecule has a labelindicating whether it binds to human beta-secretase 1.• SIDER is a dataset of marked drugs and adverse drug reactions (ADRs).
Revised: Each molecule has a label indicating whether it binds to human beta-secretase 1. • SIDER is a dataset of marked drugs and adverse drug reactions (ADRs).

Original: Molecules are groupedinto 27 system organ classes.•
Revised: Molecules are grouped into 27 system organ classes.

Original:  HIV is a dataset of HIV antiviral activity.
Revised: • HIV is a dataset of HIV antiviral activity.

Original: Each molecule has an active or inactive label.• IC50/EC50-scaffold/assay/size are datasets generated by the automated dataset curator providedby DrugOOD from the large-scale bioassay deposition website ChEMBL [39].
Revised: Each molecule has an active or inactive label. • IC50/EC50-scaffold/assay/size are datasets generated by the automated dataset curator provided by DrugOOD from the large-scale bioassay deposition website ChEMBL [39].

Original: The sufﬁx speciﬁesthe splitting scheme.
Revised: The sufﬁx speciﬁes the splitting scheme.

Original: These six datasets target on ligand-based afﬁnity prediction (LBAP).
Revised: These six datasets target on ligand-based afﬁnity prediction (LBAP).

Original: Eachmolecule has an active or inactive label.
Revised: Each molecule has an active or inactive label.

Original: Notice that the phenomenon that there exist a few invariant substructures w.r.t.
Revised: Notice that the phenomenon that there exist a few invariant substructures w.r.t.

Original: certain propertyindeed exists in the datasets we use in our paper.
Revised: certain property indeed exists in the datasets we use in our paper.

Original: Taking HIV dataset as an example, salicylhydrazidesubstructure displays potent HIV-1 integrase (IN) inhibitory activity, which has been identiﬁed byprevious studies [1, 42].
Revised: Taking HIV dataset as an example, salicylhydrazide substructure displays potent HIV-1 integrase (IN) inhibitory activity, which has been identiﬁed by previous studies [1, 42].

Original: Additionally, for BBBP dataset, as pointed out by recent studies [48, 49],some substructures are closely related to brain-blood barrier penetration.
Revised: Additionally, for BBBP dataset, as pointed out by recent studies [48, 49], some substructures are closely related to brain-blood barrier penetration.

Original: All these ten datasets do not contain personally identiﬁable information or offensive content.
Revised: All these ten datasets do not contain personally identiﬁable information or offensive content.

Original: Table 4shows the detailed statistics of datasets.
Revised: Tableshows the detailed statistics of datasets.

Original: For all datasets, we adopt the default training-validation-testsplit as shown in Table 4.
Revised: For all datasets, we adopt the default training-validation-test split as shown in Table 4.

Original: We use all molecules in the training set to optimize the model parameters.
Revised: We use all molecules in the training set to optimize the model parameters.

Original: Then, we select hyper-parameters using the validation set, and we report the results on test moleculeset for the model that achieves the best results on the validation set.
Revised: Then, we select hyper-parameters using the validation set, and we report the results on test molecule set for the model that achieves the best results on the validation set.

Original: [Table]
Revised: [Table]

Revised: number of scaffolds that contain 1, 2,

Original: [Table]
Revised: [Table]

Revised: [Figure]

Original: Next, let’s discuss on the details of HIV dataset, which is released by
Revised: Next, let’s discuss on the details of HIV dataset, which is released by

Original: Open Graph Benchmark (OGB) [21].
Revised: Open Graph Benchmark (OGB) [21].

Original: OGB adopts scaffold splittingscheme to split the HIV into train/validation/test set.
Revised: OGB adopts scaffold splitting scheme to split the HIV into train/validation/test set.

Original: We count the number of scaffolds that only contain 1, 2, 3, 4 and 5 molecules, respectively,and summarize the statistics in Table 5.
Revised: We count the number of scaffolds that only contain 1, 2, 3, 4 and 5 molecules, respectively, and summarize the statistics in Table 5.

Original: Notice that HIV has 19 , 076scaffolds in total.
Revised: Notice that HIV has 19 ,scaffolds in total.

Original: We can see there are 18 , 191 scaffolds containing lessor equal to 5 molecules, accounting for 95 .
Revised: We can see there are 18 , 191 scaffolds containing less or equal to 5 molecules, accounting for 95 .

Original: 45% of the total scaffoldcount.
Revised: 45% of the total scaffold count.

Original: HIV has a great deal of environments that contains few samples,which poses great challenge to directly applying some existing OODgeneralization methods to datasets like HIV [25].
Revised: HIV has a great deal of environments that contains few samples, which poses great challenge to directly applying some existing OOD generalization methods to datasets like HIV [25].

Original: Thus, for datasetsreleased by OGB, partitioning the mocecules into different environmentsaccording to their scaffolds may not be suitable in practice.
Revised: Thus, for datasets released by OGB, partitioning the mocecules into different environments according to their scaffolds may not be suitable in practice.

Original: Such aobservation motivates us to propose the environment-inference model.
Revised: Such a observation motivates us to propose the environment-inference model.

Revised: F Understanding the Data-generating Process

Revised: [Equation]

Revised: We provide a causal perspective to understand the data-generating process.

Revised: Recalling the two molecules Cyclopropanol ( C 3 H 6 O ) and 1,4-Cyclohexanediol ( C 6 H 12 O 2 ) used for illustration in Sec. 1, they are sampled from different environments.

Revised: Because both of them contain the hydroxy ( − OH ) , which we can call invariant or causal substructure in this case, these two molecules are readily soluble in water.

Original: F Notations
Revised: We formalize such a date-generating process of molecule property prediction with a general Structural Causal Model (SCM) [18, 43] in Fig. The abstract data variables are denoted by the nodes and the directed arrows represent the causalities. This SCM illustrates the causalities among variables: E as the environment, S as the spurious substructures, C as the invariant/causal substructures w.r.t Y , G as the instance molecule graph, Y as the ground-truth label. • S ← E → C : the

Revised: environmental variable impacts the underlying data generating distribution.

Revised: Furthermore, substructures could be divided into causal and spurious ones across all environments.

Revised: • S → G ← C : an instance molecule graph is made up of the causal and spurious substructures.

Revised: • C → Y : Y , the ground-truth label, is only determined by C .

Revised: This causation is the focus of our work.

Revised: Taking Cyclopropanol ( C 3 H 6 O ) as an example, we can specify E as the 3C-ring scaffold, C as the substructure hydroxy ( − OH ) , S as the substructures aside from hydroxy, G as Cyclopropanol , Y as good water solubility.

Revised: The good water solubility Y is only attributed to the invariant substructure hydroxy, i.e. C , rather than other spurious substructures S .

Revised: Existing MRL methods do not differentiate invariant and spurious substructures.

Revised: Hence, the spurious correlations between irrelevant substructures S and the target label Y will be encoded to learned molecular representations.

Revised: When tested on unseen environments, the downstream classiﬁer will be easily misled by these spurious correlations [58]

Revised: [Table]

Revised: G Notations

Original: We summarize the notations used in this paper in Table 6.
Revised: We summarize the notations used in this paper in Table [Table] [Table]

Original: G Sensitivity to Molecule Segmentation Method
Revised: H Sensitivity to Molecule Segmentation Method

Original: For all experiments in our original paper, we all adopt breaking retrosynthetically interesting chemicalsubstructures (BRICS) to segment molecule into substructures, which is widely used in other worksrelated to molecules [56, 62, 10].
Revised: For all experiments in our original paper [Table] we all adopt breaking retrosynthetically interesting chemical substructures (BRICS) to segment molecule into substructures, which is widely used in other works related to molecules [56, 62, 10].

Original: To investigate the sensitivity of our method to different decomposingstrategies, we adopt another molecule segmentation method called retrosynthetic combinatorialanalysis procedure (RECAP) [35], which is also available as an API in RDKit package.
Revised: To investigate the sensitivity of our method to different decomposing strategies, we adopt another molecule segmentation method called retrosynthetic combinatorial analysis procedure (RECAP) [35], which is also available as an API in RDKit package.

Original: RECAPand BRICS decompose molecules based on two different rules.
Revised: RECAP and BRICS decompose molecules based on two different rules.

Original: We conduct experiments on EC50assay/scaffold/size three datasets and the comparisions are summarized in Table 7.
Revised: We conduct experiments on EC50assay/scaffold/size three datasets and the comparisions are summarized in Table 7.

Original: We can see that
Revised: We can see that

Original: RECAP and BRICS show competitive performance on our model and both outperform the baselinesby large margins.
Revised: RECAP and BRICS show competitive performance on our model and both outperform the baselines by large margins.

Original: H Future Direction
Revised: I Future Direction

Original: Sometimes, bio-chemical properties are affected by interactions between substructures.
Revised: Sometimes, bio-chemical properties are affected by interactions between substructures.

Original: To encodesuch interactions between substructures into the ﬁnal learned molecular representation, we utilize thepermutation equivariant Set Attention Block (SAB) proposed in Set Transformer [33].
Revised: To encode such interactions between substructures into the ﬁnal learned molecular representation, we utilize the permutation equivariant Set Attention Block (SAB) proposed in Set Transformer [33].

Original: SAB takesa representation set of any size as input and outputs a representation set of equal size.
Revised: SAB takes a representation set of any size as input and outputs a representation set of equal size.

Original: SAB is ableto encode pairwise and higher-order interactions between elements in input sets into outputs.
Revised: SAB is able to encode pairwise and higher-order interactions between elements in input sets into outputs.

Original: Weadd such a SAB after the Substructure Encoder.
Revised: We add such a SAB after the Substructure Encoder.

Original: For each molecule, we feed the representions of 
Revised: For each molecule, we feed the representions of its substructures to SAB to obtain new substruture representations. In this way, the ﬁnal molecule

Original: [Table] [Table]
Revised: [Table] [Table]

Original: its substructures to SAB to obtain new substruture representations.

Original: In this way, the ﬁnal moleculerepresentation could model interactions between substructures.
Revised: representation could model interactions between substructures.

Original: We conduct experiments on EC50assay/scaffold/size to examine the performance of adding such a SAB.
Revised: We conduct experiments on EC50assay/scaffold/size to examine the performance of adding such a SAB.

Original: As demonstrated in Table 8,we can see that adding such a SAB further improves our model on EC50-scaffold.
Revised: As demonstrated in Table 8, we can see that adding such a SAB further improves our model on EC50-scaffold.

Original: This design is anaive attempt but brings us some valuable insights.
Revised: This design is a naive attempt but brings us some valuable insights.

Original: I Limitations
Revised: J Limitations

Original: Some studies [25, 15] have empirically shown that existing models designed for OOD learning mayfail to outperform the simple ERM [52] model when the environment count is large.
Revised: Some studies [25, 15] have empirically shown that existing models designed for OOD learning may fail to outperform the simple ERM [52] model when the environment count is large.

Original: Though we canrelabel the environment for each molecule according to the new environment partition inferred by ourdevised environment-inference module, we still need to set the environment count k to a relativelylarger value than that of other OOD datasets from other domain, e.g. Camelyon17 [4], which onlycontains ﬁve environments.
Revised: Though we can relabel the environment for each molecule according to the new environment partition inferred by our devised environment-inference module, we still need to set the environment count k to a relatively larger value than that of other OOD datasets from other domain, e [Table] g. Camelyon17 [4], which only contains ﬁve environments.

Original: Thus, using our inferred environment partition, existing models designedfor OOD learning might still be inferior to ERM in some cases.[Table]
Revised: Thus, using our inferred environment partition, existing models designed for OOD learning might still be inferior to ERM in some cases.

Original: J Potential Negative Impacts
Revised: K Potential Negative Impacts

Original: As far as we are concerned, we have not identiﬁed any negative social impact of this work.
Revised: As far as we are concerned, we have not identiﬁed any negative social impact of this work.