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

Original: AfﬁliationAddress
Revised: 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 currentmethods have shown promising power for various tasks, e.g., molecular propertyprediction 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 existingmethods is that either the model development or experimental evaluation is mostlybased 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 violatedin 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 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.).

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 differentdistributions 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 substructuresthat 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 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.

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, ourmethod achieves up to 5.9% and 3.9% improvement over the strongest baselineson OGB and DrugOOD benchmarks in terms of ROC-AUC, respectively.

Original:  Introduction
Revised:  Introduction

Original: Predicting molecular properties plays an important role in many related applications like drugdiscovery [10] and material design [41].
Revised: Predicting molecular properties plays an important role in many related applications like drugdiscovery [11] and material design [44].

Original: These professional tasks conventionally take great efforts byexperts e.g. in chemistry and pharmacology.
Revised: These professional tasks conventionally take great efforts byexperts 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 onbuilding effective machine learning models for scientiﬁc discovery and solid progress has been madealong 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 [55], retrosynthetic analysis [52], search of antibiotics [42], virtual screening [33] for drug discovery, etc.
Revised: 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 [60], retrosynthetic analysis [56], search of antibiotics [45], virtual screening [36] 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, yetreal-world environments are often dynamic and uncertain, which requires the model to effectively handle distribution shifts.

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: 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 [33] 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 [36] 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 fromunknown 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 [34, 5, 36].
Revised: 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 [37, 5, 39].

Original: [Figure]
Revised: [Figure]

Original: invariantly contributes to the water solubility of the two molecules which contain different scaffolds, i.e. sampled from different environments by deﬁnition.
Revised: 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-
Revised: Existing methods devised for out-of-distribution generalization [Figure] mostly focus on Euclidean data such as images, while few endeav-

Original: [Figure]
Revised: [Figure]

Original: ors OOD generalization on non-Euclidean data [48].
Revised: ors OOD generalization on non-Euclidean data [52].

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., graphstructured data, is different from visual data in nature.

Original: The work [23]point out that existing OOD models [2, 43, 14, 56, 40] fail to exhibitsigniﬁcant improvement on MRL tasks against distribution shifts andeven the simple Empirical Risk Minimization (ERM) [44] methodoutperforms these latest methods, which is also empirically veriﬁedby [13].
Revised: The work [24]point out that existing OOD models [2, 46, 15, 61, 43] fail to exhibitsigniﬁcant improvement on MRL tasks against distribution shifts andeven the simple Empirical Risk Minimization (ERM) [47] methodoutperforms these latest methods, which is also empirically veriﬁedby [14].

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 moleculesto 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 [27, 38, 58, 25] across bioinformatics, pharmacy, and data mining.
Revised: 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 [28, 41, 63, 26] across bioinformatics, pharmacy, and data mining.

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 .
Revised: 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 .

Original: The common practice speciﬁes environments as someprominent information of the molecules e.g. scaffold pattern [28, 19].
Revised: The common practice speciﬁes environments as someprominent information of the molecules e.g. scaffold pattern [29, 20].

Original: Thus, the data-generatingenvironments and the induced distributions which these two molecules are sampled from can beconsidered different [19].
Revised: Thus, the data-generatingenvironments and the induced distributions which these two molecules are sampled from can beconsidered different [20].

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

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 substructuresacross 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 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.

Original: Besides, some studies [23, 13] show that directly utilizing such environment labels as input whenadapting existing OOD generalization methods to MRL tasks can be problematic.
Revised: Besides, some studies [24, 14] show that directly utilizing such environment labels as input whenadapting 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 alabel-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 modelcan 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 environmentlabel 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 byintroducing 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 [48].
Revised: Then we analyzethe essential cause behind the failure of existing MRL models and propose a new learning schemebased on the invariance principle [52].

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 anenvironment 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 [19], as shown in our experimental results.
Revised: OGB molecular property prediction tasks [20], as shown in our experimental results.

Original: As for a newlyreleased benchmark for drug-oriented OOD learning [23], 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 newlyreleased benchmark for drug-oriented OOD learning [24], 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 incorporatingan important observation that the substructure of molecule can convey invariant casual informationacross 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 isthe ﬁ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 proposea 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 requireenvironment labels which in fact can be noisy and unreliable, but instead achieve environmentinference 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 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.

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 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.

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 propertyprediction 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 stilloutperforms several SOTA approaches tailored for general OOD learning (using environmentlabels 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 [39, 2, 7, 48] assume that there is a potential environment variable e accounting forthe distribution shift between the training and testing data.
Revised: Recent studies [42, 2, 7, 52] assume that there is a potential environment variable e accounting forthe 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 thetarget 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 aloss 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 ) [29].
Revised: | e ] is called the risk function under a givenenvironment e and denoted as R e ( x e , y e ) [30].

Original: Invariant Learning.
Revised: Invariant Learning.

Original: There is an emerging line of research [39, 2, 8, 9] regarding invariant predictor learning, for solving the OOD generalization problem.
Revised: There is an emerging line of research [42, 2, 8, 10] 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 [29].
Revised: These methods propose to ﬁnd aninvariant predictor that could uncover invariant relationships between inputs and targets across allenvironments [30].

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 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.

Original: A recentwork [48] leverages invariance principle to handle distribution shifts on graphs, but it is designed fornode-level prediction.
Revised: A recentwork [52] leverages invariance principle to handle distribution shifts on graphs, but it is designed fornode-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 beclassiﬁed into two categories.

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

Original: They use language models to process the textualrepresentation (SMILES) of a molecule, for example, Transformer [45] or BERT [12].
Revised: They use language models to process the textualrepresentation (SMILES) of a molecule, for example, Transformer [48] or BERT [13].

Original: SMILES is alinear encoding for molecules and highly depends on the traverse order of molecule graphs.
Revised: SMILES is alinear 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 believecalls 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-basedand graph neural networks (GNN)-based methods.

Original: The molecular ﬁngerprint techniques date back tothe Morgan ﬁngerprints [35].
Revised: The molecular ﬁngerprint techniques date back tothe Morgan ﬁngerprints [38].

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

Original: Since molecules can be viewed as structured graphs, graphneural networks have been widely used to learn molecule representation [24, 15, 20].
Revised: Since molecules can be viewed as structured graphs, graphneural networks have been widely used to learn molecule representation [25, 16, 21].

Original: Existing general OOD methods [2, 43, 14, 56, 40] are not tailored to such non-Euclidean structureddata, i.e. molecules.
Revised: Existing general OOD methods [2, 46, 15, 61, 43] are not tailored to such non-Euclidean structureddata, i.e. molecules.

Original: In several recent works on molecule property classiﬁcation tasks [57, 47], 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 [62, 50], theimportance of molecular substructures has been emphasized and such inductive bias is incorporatedinto 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 notleverage 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

Original: 3.1 Problem Formulation
Revised: 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 thinletters, 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 berepresented 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 themolecule 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 ofmolecule representation learning task is to predict the target label y given the associated inputmolecule 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 whilethe 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.
Revised: 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
Revised: 3.2 Understanding the Data-generating Process

Original: [Figure]
Revised: [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
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

Original: Sec. 1, they are sampled from different environments.
Revised: 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.
Revised: 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) [16, 37] in Fig.
Revised: We formalize such a date-generating process of molecule property prediction with a general Structural Causal Model (SCM) [17, 40] in Fig.

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

Original: This SCM illustrates
Revised: This SCM illustrates

Original: the causalities among variables: E as the environment, S as the spurious substructures, C as the
Revised: 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.
Revised: 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.
Revised: • 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.
Revised: 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.
Revised: • 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 .
Revised: • C → Y : Y , the ground-truth label, is only determined by C .

Original: This causation is the focus of our work.
Revised: 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.
Revised: 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 .
Revised: 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 substructures S and the target label Y will be encoded to learnedmolecular representations.
Revised: Hence, the spuriouscorrelations between irrelevant substructures S and the target label Y will be encoded to learnedmolecular representations.

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

Original: With the knowledge that (bio)chemical propertiesof a molecule are usually associated with a few privileged substructures [27, 38, 58, 25], 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 propertiesof a molecule are usually associated with a few privileged substructures [28, 41, 63, 26], we aimto suppress such spurious correlations and leverage environment-invariant substructures that morestably 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: Noticethat the learned invariant molecular representations should satisfy the invariance principle mentionedin Sec. 2.

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.3 Model Formulation

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

Original: Solvingthe formulation in Eq.
Revised: Solvingthe 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 differentenvironments 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 moleculeencoder Φ 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 [8].
Revised: 2, we add an extra invarianceconstraint y ⫫ e | Φ( G ) , which is used to suppress spurious correlations [8].

Original: Since assessingcausality is challenging, we could rethink the problem on the basis of information theory.
Revised: Since assessingcausality 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: Recallthat we hope to let the molecule encoder leverage environment-invariant substructures and learn amolecular 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 powerof Φ( 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 mutualinformation.

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 bespeciﬁ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 thepractical 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ﬁcationsof 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 differentenvironments 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 isunreasonable 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 [19] forn example, OGB uses scaffold to split the molecules into different environments.
Revised: Taking thedataset HIV for molecule property prediction tasks released by Open Graph Benchmark [20] forn 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 thatwe regard each scaffold as an environment directly, 41 , 127 molecules in HIV are partitioned into19 , 076 environments (see details in Appendix E).

Original: This environment count is much larger than other
Revised: This environment count is much larger than other

Original: OOD datasets from other domains, e.g. Camelyon17 2 [3], CivilComments 3 [6], etc.
Revised: OOD datasets from other domains, e.g. Camelyon17 2 [3], CivilComments 3 [6], 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 [23, 13].
Revised: Even thoughsome datasets may provide manual speciﬁcations of environments, the environment counts are alsotoo large, which is unfriendly to existing OOD models [24, 14].

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

Original: We denote the environment countas 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 posteriorp τ ( 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 variationaldistribution 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 κ [Equation] 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 practical 230 instantiation, which involves the environment-inference model deﬁned above:
Revised: 5 and give an equivalent tractable objective in practical 229 instantiation, which involves the environment-inference model deﬁned above:

Original: [Equation]
Revised: [Equation]

Revised: where G e consists of the pairs of molecular graph G and corresponding label y under environment e .

Original: [Figure]
Revised: [Figure]

Original: Theorem 1.
Revised: Theorem 1.

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

Original: 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: 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 ) .

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 equalsto 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 forcesthe learned representation z to satisfy the invariance principle mentioned in Sec. 2, thus ensuring avalid 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 befound in Appendix B.

Original: 3.4 Model Instantiations and Training
Revised: 3.4 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 weshould design a module, entitled Environment Classiﬁer , that takes ( G , y ) as the input and outputsthe probabilistic distribution of e .

Original: We use a Graph Isomorphism Network (GIN) [50], to learn a graphrepresentation given G .
Revised: We use a Graph Isomorphism Network (GIN) [54], to learn a graphrepresentation 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 fedto 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 theprior p τ ( e | G ) to Uniform 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 followed by a softmax activation function to modelit.

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 inputand 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 representationlearning 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) [11], which is available as an API in RDKit [30].
Revised: Meanwhile, we decompose the input molecule into a set of chemicalsubstructures using a molecule segmentation method, e.g. breaking retrosynthetically interestingchemical substructures (BRICS) [12], which is available as an API in RDKit [31].

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, weoperate attentive pooling on these substructure representations to obtain a new substructure-awaremolecular 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 betweencertain substructures and target properties into this representation.

Original: The Complete Encoder, the
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 forthe Predictor ω , we implement it with a multi-layer perceptron, followed by a softmax function.

Original: Theoverview of our model is demonstrated in Fig.
Revised: Theoverview 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 , τ ) .
Revised: 2) optimizing the molecule encoder and the predictor: θ ∗ ← arg min θ L inv ( θ ; G train , τ ) .

Original: 15 k 0 .0 .0 .0 .0 .
Revised: 15 k 0 .0 .0 .0 .0 .

Original: R O CAU C (a) 0.
Revised: R O CAU C (a) 0.

Original: 2 β 0 .0 .0 .0 .0 .0 .0 .0 .
Revised: 2 β 0 .0 .0 .0 .0 .0 .0 .0 .

Original: R O CAU C (b)
Revised: R O CAU C (b)

Original:     100 Epoch 0 .0 .0 .

Revised:     100 Epoch 0 .0 .0 .

Original: R i s k
Revised: R i s k

Original: Env1 Env
Revised: Env1 Env

Original: Env3 Env
Revised: Env3 Env

Original: [Figure]
Revised: [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 withmean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPUand 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) [19].
Revised: Graph Benchmark (OGB) [20].

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 splitcontains 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 certaindegree, it provides an OOD test-bed as different scaffold often suggest different data-generationenvironments.

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

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 manualspeciﬁ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 [54, 51, 46].
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 [59, 55, 49].

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 theirgeneralization ability against distribution shifts.

Original: We adapt three backbones: GCN [26], GIN [50]and GraphSAGE [17] into our method.
Revised: We adapt three backbones: GCN [27], GIN [54]and GraphSAGE [18] 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 ” [15, 21, 32].
Revised: We also compare against another augmented version “ + virtual node ” [16, 22, 35].

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

Original: DeepCoral [43], DANN [14], MixUp [56] and GroupDro [40].
Revised: DeepCoral [46], DANN [15], MixUp [61] and GroupDro [43].

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 thesemethods require the manual speciﬁcation of environments in dataset, we report this comparison ondatasets 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 reportedin the original paper or selected by grid search.

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

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

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 consistentimprovements after adapted to our methods across all the four datasets released by OGB in termsof 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 onall 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 improvetheir 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 inmolecular properties predictions under OOD setting to that, our proposed learning objective enforcesthe 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 comparsionwith 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
Revised: Except on IC50-size, our method outperforms all baselines across all datasets due to its ability

Original: [Table]
Revised: [Table]

Original: except ERM and MixUp all require environment labels.
Revised: except ERM and MixUp all require environment labels.

Original: All methods including ours use GIN [50] as backbones.
Revised: All methods including ours use GIN [54] as backbones.

Original: [Table]
Revised: [Table]

Original: to enforce the molecule encoder to leverage environment-invariant substructures that more stablyrelate with the labels across environments.
Revised: 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 moredata 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-sizein our analysis.

Original: As for ERM, [23, 13] have pointed out that simple ERM shows better performancecompared to subsequent OOD methods when datasets have relatively large environment counts.
Revised: As for ERM, [24, 14] have pointed out that simple ERM shows better performancecompared 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: 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.

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 introducedin Sec. 3.4 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 thisarchitecture 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 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.

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 performanceof 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, theimpact 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.
Revised: Combined withthe architecture discussed above, the new objective is able to better guide the molecule encoder tolearn environment-invariant molecular representations against distribution shifts.

Original: [Table]
Revised: [Table]

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 proposedenvironment-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ﬁcationsof environments may be unavailable due to the high price for labeling environments by experts.

Original: Butwhen environment labels are available
Revised: Butwhen environment labels are available

Original: [Table]
Revised: [Table]

Original: how will be performance be like if directly utilizing the givenenvironment partition?
Revised: how will be performance be like if directly utilizing the givenenvironment 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 anexample, 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 utilizethe 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 obtainsinferior 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 environmentinference method, we relabel the environment for each molecule for DANN according to our inferredenvironment 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 betterperformance 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 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.

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 numberof environments k , the trading-off parameter β in Eq.

Original: Fig.
Revised: Fig.

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

Original: It shows that the performance of our methods degrades when k istoo small (e.g. k
Revised: It shows that the performance of our methods degrades when k istoo small (e.g. k

Original: = 1 , 5 ) or too large (e.g. k = 15 ).
Revised: = 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 dataas 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 thetraining samples into different environments is necessary.

Original: Fig. 4(b) shows the results of our methodby varying the trade-off parameter β .
Revised: Fig. 4(b) shows the results of our methodby 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: Thisis 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 thesecond term of Eq. 7, the sufﬁciency condition of invariance principle cannot be satisﬁed, resultingin 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 lowerthe risks of different environments, we visualize the risk dynamic curve of some environments in

Original: Fig. 4(c).
Revised: Fig. 4(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. 4(c), the difﬁculties of decreasing the risk on different environmentsare 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 trainingprocess (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 backboneto 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 newlearning 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 inferencemodel to identify each molecule’s corresponding environment without need of manual speciﬁcationsof 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 yieldsconsistent 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.
Revised: 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.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

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

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes] See Appendix G.
Revised: [Yes] See Appendix I.

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

Original: [Yes] SeeAppendix H.(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Revised: [Yes] SeeAppendix J. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See Appendix Aand B.3.
Revised: [Yes] See Appendix Aand B.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (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)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes] See Appendix C.
Revised: [Yes] See Appendix C.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

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

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?

Original: [Yes] See Section 4.
Revised: [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)?
Revised: (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.
Revised: [Yes] See Section 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] See Section 4 and
Revised: [Yes] See Section 4 and

Original: Appendix E.
Revised: Appendix E.

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

Original: [Yes] See Section 4 and
Revised: [Yes] See Section 4 and

Original: Appendix E.
Revised: Appendix E.

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

Original: [No]
Revised: [No]

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

Original: using/curating?
Revised: using/curating?

Original: [Yes] See Appendix E.
Revised: [Yes] See Appendix E.

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

Original: information or offensive content?
Revised: information or offensive content?

Original: [Yes] See Appendix E.
Revised: [Yes] See Appendix E.

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

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

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

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

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Revised: A Proof for Proposition 1

Revised: Proof.

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

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

Revised: [Equation]

Revised: Rearrange Eq.

Revised: and we can get,

Revised: [Equation]

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

Revised: According to Eq.

Revised: 9,maximizing this ELBO is equivalent to minimizing the KL divergence and maximizing log p τ

Revised: ( y | G ) .

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

Revised: [Equation]

Revised: [Equation]

Revised: We thus conclude the proof.

Revised: B Proofs for Theorems

Revised: In this paper, we extend the invariance assumption [42, 2] to molecule representation learning:

Revised: Assumption 1.

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

Revised: 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: With the terminology of information theory, we present a useful lemma [52] that interprets theinvariance and sufﬁciency conditions in Assumption 1:

Revised: Lemma 1.

Revised: 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: Proof.

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

Revised: [Equation]

Revised: 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

Revised: I( y ; r ) .

Revised: We prove it by contradiction.

Revised: Assume that r (cid:54) =

Revised: arg max r I( y ; r ) andthere exists r (cid:48) with r (cid:48)

Revised: =

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

Revised: r (cid:48) .

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

Revised: =

Revised: f m ( r , ˜ r ) .

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

Revised: =

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

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

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

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

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

Revised: I( y ; r ) would also satisfyy =

Revised: h ( r )+ n by contradiction.

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

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

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

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

Revised: That means r (cid:48)

Revised: =

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

Revised: B.1 Proof for Theorem 1

Revised: Proof.

Revised: According to the dependency relationship z ← G → y , we have

Revised: [Equation]

Revised: Next, we have

Revised: [Equation]

Revised: Based on

Revised: Jensen

Revised: Inequality and

Revised: Triangle

Revised: Inequality, we can obtain that

Revised: D KL

Revised: [Equation]

Revised: y | z )

Revised: [Equation]

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

Revised: [Equation]

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

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

Revised: ( y ; e | z ) .

Revised: B.2 Proof for Theorem

Revised: Proof.

Revised: 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: Also we have

Revised: [Equation]

Revised: [Equation]

Revised: Based on this, we will have

Revised: [Equation]

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

Revised: [Equation]

Revised: where C is a constant.

Revised: Then the problem min q ( y | z ) D KL ( p

Revised: [Equation]

Revised: y | G , e

Revised: [Equation]

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

Revised: [Equation]

Revised: which means minimizing term 2 in Eq.

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

Revised: ( z ; y ) .

Revised: C Implementation Details

Revised: [Figure]

Revised: C.1 Baselines

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

Revised: Three backbones.

Revised: We adapt three backbones into our method, namely, GCN [27], GraphSAGE [18] and GIN [54].

Revised: We also emphasize the comparison with their augmented versions,i.e. “ + virtual node ” [16, 22, 35].

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

Revised: Graph Benchmark [20] 4 .

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

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.

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

Revised: • GCN [27] is a scalable approach on graph-structured data that is based on an efﬁcient variant ofconvolutional neural networks.• GIN [54] generalizes the Weisfeiler-Lehman (WL) graph isomorphism test [33] and hence achievesmaximum discriminative power among GNNS.• GraphSAGE [18] learns a function that generates embeddings by sampling and aggregatingfeatures from a node’s local neighborhood.• GCN/GIN/GraphSAGE + virtual node [16, 22, 35] is a variant of the original method augmentedby an additional node connecting to all nodes in the raw graph.

Revised: Models tailored for OOD learning.

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

Revised: ERM [47], IRM [2], DeepCoral [46], DANN [15], MixUp [61] and GroupDro [43].

Revised: We use theimplementations of these six method provided by DrugOOD 5 .

Revised: We search for the optimal hyperparameters by ranging learning rate over { 1 e − 3 , 5 e −

Revised: 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 } and dropout rateover { 0 . 1 , 0 . 3 , 0 . 5 } .

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

Revised: We set the penalty weight and the penalty anneal iteration to 10 and 500 , respectively.• DeepCoral [46] 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: The penalty weight is set to 0 .

Revised: DANN [15] encourages feature representations to be consistent across domains.

Revised: We set to theinverse factor to 0 .

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

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

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

Revised: The step size is set to 0 .

Revised: C.2 Our Method

Revised: We implement our method in Pytorch.

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

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

Revised: 3.4 all in Graph Isomorphism Network (GIN) [54].

Revised: 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: As for the prior p ( e | G ) , we set it to a Uniformdistribution or a discrete Gaussian distribution.

Revised: We use CrossEntropyLoss for all models and the

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

Revised: D Training Procedure

Revised: The training procedure of our method is summarized in Algorithm 1.

Revised: Algorithm 1: The training procedure.

Revised: do data

Revised: for j ← 1 to k do

Revised: Compute batch loss L elbo τ κ j

Revised: Backpropagate −L elbo and optimize parameters τ κ ; 7 Freeze the parameters κ, τ ; 8 for i ← to

Revised: 2 do data with

Revised: for j ← to do

Revised: Determine the environment of each

Revised: G k by

Revised: Compute batch loss L inv θ k according

Revised: Backpropagate L inv and optimize parameters

Revised: ; 14 Output the parameters θ ;

Revised: E More Details of Datasets

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

Revised: Four of them, namely,

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

Revised: The rest six arereleased by DrugOOD [24], i.e. IC50-assay, IC50-scaffold, IC50-size, EC50-assay, EC50-scaffoldand EC50-size.

Revised: We provide detailed descriptions for them as below.

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

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

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

Revised: 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: Molecules are groupedinto 27 system

Revised: organ

Revised: classes.•

Revised: HIV is a dataset of HIV antiviral activity.

Revised: 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.

Revised: The sufﬁx speciﬁesthe splitting scheme.

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

Revised: Eachmolecule has an active or inactive label.

Revised: Notice that all these ten datasets do not contain personally identiﬁable information or offensivecontent.

Revised: Table 4 shows the detailed statistics of datasets.

Revised: For all datasets, we adopt the defaulttraining-validation-test split as shown in Table 4.

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

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

Revised: [Table]

Revised: [Table]

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

Revised: Open Graph Benchmark (OGB) [20].

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

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.

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

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

Revised: 45% of the total scaffoldcount.

Revised: 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 [24].

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

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

Revised: F Notations

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

Revised: G Sensitivity to Molecule Segmentation Method

Revised: 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 [51, 57, 9].

Revised: To investigate the sensitivity of our method to different decomposingstrategies, we adopt another molecule segmentation method called retrosynthetic combinatorialanalysis procedure (RECAP)

Revised: [34], which is also available as an API in RDKit package.

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

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

Revised: We can see that

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

Revised: H Future Direction

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

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

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

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

Revised: Weadd such a SAB after the Substructure Encoder.

Revised: For each molecule, we feed the representions of

Revised: [Table]

Revised: [Table]

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

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

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

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

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

Revised: I Limitations

Revised: Some studies [24, 14] have empirically shown that existing models designed for OOD learning mayfail to outperform the simple ERM [47] model when the environment count is large.

Revised: 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 [3], which onlycontains ﬁve environments.

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

Revised: [Table]

Revised: J Potential Negative Impacts

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