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

Original: O H
Revised: Abstract Molecule representation learning (MRL) has been extensively studied and current methods have shown promising power for various tasks, e.g., molecular property prediction and target identiﬁcation. However, a common hypothesis of existing methods is that either the model development or experimental evaluation is mostly based on i.i.d. data across training and testing.

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

Revised: We propose a new representation learning framework entitled

Revised: MoleOOD to enhance the robustness of MRL models against such distribution shifts, motivated by an observation that the (bio)chemical properties of molecules are usually invariantly associated with certain privileged molecular substructures across different environments (e.g., scaffolds, sizes, etc.).

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

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

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

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

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

Revised:  Introduction

Original: O H
Revised: Predicting molecular properties plays an important role in many related applications like drug discovery [13] and material design [51]. These professional tasks conventionally take great efforts by experts e.g. in chemistry and pharmacology. Recent years have witnessed inspiring breakthroughs on building effective machine learning models for scientiﬁc discovery, and solid progress has been made along the avenue of ML-based molecule representation learning (MRL). In general, MRL aims at embedding a molecule into a vector in latent space as a foundation model, on top of which the learned representations could be used for a variety of downstream tasks, such as target identiﬁcation [69], retrosynthetic analysis [65], search of antibiotics [54], virtual screening [40] for drug discovery, etc.

Revised: The challenge, however, is that existing MRL methods are mostly based on an underlying hypothesis that training and testing molecules are independently sampled from an identical environment, yet

Original: O H  Cyclopropanol 1,4-Cyclohexanediol
Revised: O H O H O H Cyclopropanol 1,4-Cyclohexanediol

Original: Scaffold 1 3C-Ring
Revised: Scaffold 1 3C-Ring

Original: Scaffold 2 6C-Ring
Revised: Scaffold 2 6C-Ring

Original: O H
Revised: O H

Original: O O
Revised: O O

Original: O H
Revised: O H

Original: O
Revised: O

Original: O H O H
Revised: O H O H

Original: O
Revised: O

Original: O H
Revised: O H

Original: Smaller Size
Revised: Smaller Size

Original: Larger Size
Revised: Larger Size

Original: Environment: Scaffold
Revised: Environment: Scaffold

Original: Environment: Size
Revised: Environment: Size

Original: Formic Acid
Revised: Formic Acid

Original: Citric Acid real-world environments are often dynamic and uncertain, which requires the model to effectively handle distribution shifts.
Revised: Citric Acid real-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 [40] 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 [40] 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 from unknown distributions.
Revised: However, some unpredictable events like COVID-19 may occur, bringing new targets from unknown distributions.

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

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

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

Original: The work [27] point out that existing OOD models [3, 55, 18, 70, 50] fail to exhibit signiﬁcant improvement on MRL tasks against distribution shifts and even the simple
Revised: The work [27] point out that existing OOD models [3, 55, 18, 70, 50] fail to exhibit signiﬁcant improvement on MRL tasks against distribution shifts and even the simple

Original: Empirical Risk Minimization (ERM) [56] method outperforms these latest methods, which is also empirically veriﬁed by [16].
Revised: Empirical Risk Minimization (ERM) [56] method outperforms these latest methods, which is also empirically veriﬁed by [16].

Original: We aim to develop an OOD method tailored for molecules to solve the
Revised: We aim to develop an OOD method tailored for molecules to solve the

Original: OOD generalization problem on MRL in this paper.
Revised: OOD generalization problem on MRL in this paper.

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

Original: The common practice speciﬁes environments as some prominent information of the molecules e.g. scaffold pattern [32, 23] and molecule size [27].
Revised: The common practice speciﬁes environments as some prominent information of the molecules e.g. scaffold pattern [32, 23] and molecule size [27].

Original: Fig.
Revised: Fig.

Original: 1 provides two illustrative examples.
Revised: 1 provides two illustrative examples.

Original: Let’s ﬁrst take a look at the left example, where two molecules Cyclopropanol ( C 3 H 6 O ) and 1,4-Cyclohexanediol ( C 6 H 12 O 2 ) contain different scaffold patterns 2 : the former is 3C-ring and the latter is 6C-ring .
Revised: Let’s ﬁrst take a look at the left example, where two molecules Cyclopropanol ( C 3 H 6 O ) and 1,4-Cyclohexanediol ( C 6 H 12 O 2 ) contain different scaffold patterns 2 : the former is 3C-ring and the latter is 6C-ring .

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: Then we analyze the essential cause behind the failure of existing MRL models and propose a new learning scheme based on the invariance principle [47, 44, 3, 61].
Revised: Then we analyze the essential cause behind the failure of existing MRL models and propose a new learning scheme based on the invariance principle [47, 44, 3, 61].

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

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

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

Original: As for a newly released benchmark for drug-oriented OOD learning [27], even without access to environment labels, our method can still outperform state-of-the-art models that rely on environment labels for training in ﬁve out of six datasets.
Revised: As for a newly released benchmark for drug-oriented OOD learning [27], 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 out-of-distribution (OOD) generalization problem for molecule representation learning (MRL), by particularly incorporating an important observation that the substructure of molecule can convey invariant casual information across environments, regarding certain property prediction tasks.
Revised: • We formulate the out-of-distribution (OOD) generalization problem for molecule representation learning (MRL), by particularly incorporating an important observation that the substructure of molecule can convey invariant casual information across environments, regarding certain property prediction tasks.

Original: To our best knowledge, this is the ﬁrst work that formulates the OOD problem in MRL background and proposes to leverage the invariance principle which opens a new perspective for handling substructure-aware distribution shifts.
Revised: To our best knowledge, this is the ﬁrst work that formulates the OOD problem in MRL background and proposes to leverage the invariance principle which opens a new perspective for handling substructure-aware distribution shifts.

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

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

Original: This design endows our model with practical applicability for molecular OOD learning where the manual speciﬁcations of the environments are often unavailable.
Revised: This design endows our model with practical applicability for molecular OOD learning where the manual speciﬁcations of the environments are often unavailable.

Original: • We conduct extensive experiments on ten public datasets.
Revised: • We conduct extensive experiments on ten public datasets.

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

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

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

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

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Invariant Learning.
Revised: Invariant Learning.

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

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

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

Original: Recent works adopt the invariance principle as a cornerstone for handling various distribution shifts on semistructured data like graphs [61, 6] and sequences [66].
Revised: Recent works adopt the invariance principle as a cornerstone for handling various distribution shifts on semistructured data like graphs [61, 6] and sequences [66].

Original: To our knowledge, our work is a pioneering attempt that leverages the invariance principle and incorporate useful domain knowledge for handling molecular graph classiﬁcation tasks under distribution shifts.
Revised: To our knowledge, our work is a pioneering attempt that leverages the invariance principle and incorporate useful domain knowledge for handling molecular graph classiﬁcation tasks under distribution shifts.

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

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

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

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

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

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

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

Original: The molecular ﬁngerprint techniques date back to the
Revised: The molecular ﬁngerprint techniques date back to the

Original: Morgan ﬁngerprints [43].
Revised: Morgan ﬁngerprints [43].

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

Original: Since molecules can be viewed as structured graphs, methods tailored for graph data, such as graph neural networks [28, 19, 24] or graph transformers [17, 62] can be used to learn molecule representation.
Revised: Since molecules can be viewed as structured graphs, methods tailored for graph data, such as graph neural networks [28, 19, 24] or graph transformers [17, 62] can be used to learn molecule representation.

Original: Existing general OOD methods [3, 55, 18, 70, 50] are not tailored to such non-Euclidean structured data, i.e. molecules.
Revised: Existing general OOD methods [3, 55, 18, 70, 50] are not tailored to such non-Euclidean structured data, i.e. molecules.

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

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

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

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

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

Original:  Methodology 3.1 Problem Formulation
Revised:  Methodology 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 thin letters, respectively.
Revised: All the random variables and the corresponding realizations are denoted as bold and thin letters, respectively.

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

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

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

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

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

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

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: 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: Existing MRL methods do not differentiate invariant and spurious substructures.
Revised: Existing MRL methods do not differentiate invariant and spurious substructures.

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

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

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

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

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

Original: 3.2 Model Formulation
Revised: 3.2 Model Formulation

Original: The framework contains two parts, the fronted molecule encoder Φ for learning an “invariant representation” 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 “invariant representation” of the input molecule graph and the back-end predictor ω for ﬁnal prediction.

Original: Solving the formulation in Eq.
Revised: Solving the formulation in Eq.

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: Treating the outputs of ω and Φ as distribution q θ ( z | G ) and q θ ( y | z ) respectively, Eq.
Revised: Treating the outputs of ω and Φ as distribution q θ ( z | G ) and q θ ( y | z ) respectively, Eq.

Original: 4 can be speciﬁed as:
Revised: 4 can be speciﬁed as:

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: Proposition 1.
Revised: Proposition 1.

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

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

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

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

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

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Figure]
Revised: [Figure]

Original: ec o m po s e
Revised: ec o m po s e

Original: P r e d i c t o r s ub s t r u c t u r e s attentive pooling
Revised: P r e d i c t o r s ub s t r u c t u r e s attentive pooling

Original: OH
Revised: OH

Original: Environment
Revised: Environment

Original: Classifier
Revised: Classifier

Original: Conditional
Revised: Conditional

Original: GNN
Revised: GNN

Original: O H
Revised: O H

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 instantiation, which involves the environment-inference model deﬁned above:
Revised: 5 and give an equivalent tractable objective in practical instantiation, which involves the environment-inference model deﬁned above:

Original: [Equation]
Revised: [Equation]

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

Original: Theorem 1.
Revised: Theorem 1.

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

Original: ( y ; e | z ) , letting z show equal performance for the downstream tasks across all environments, i.e. p ( y | z , e ) = p ( y | z ) .
Revised: ( y ; e | z ) , letting z show equal performance for the downstream tasks across all environments, i.e. p ( y | z , e ) = p ( y | z ) .

Original: Theorem 2.
Revised: Theorem 2.

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

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

Original: I
Revised: 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 forces the learned representation z to satisfy the invariance principle mentioned in Sec. 2, thus ensuring a valid solution for OOD problem deﬁned in Eq.
Revised: 7 forces the learned representation z to satisfy the invariance principle mentioned in Sec. 2, thus ensuring a valid solution for OOD problem deﬁned in Eq.

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

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

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

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

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

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

Original: We could set the prior p τ ( e | G ) to Uniform distribution or Gaussian distribution .
Revised: We could set the prior p τ ( e | G ) to Uniform distribution or Gaussian distribution .

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

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

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

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

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

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

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

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

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, we operate attentive pooling on these substructure representations to obtain a new substructure-aware molecular representation.
Revised: Then, considering r G as a query with regard to substructures, we operate attentive pooling on these substructure representations to obtain a new substructure-aware molecular representation.

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

Original: Guided by our proposed learning objective, we can encode some invariant relationships between certain substructures and target properties into this representation.
Revised: Guided by our proposed learning objective, we can encode some invariant relationships between certain 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 for the Predictor ω , we implement it with a multi-layer perceptron, followed by a softmax function.
Revised: As for the Predictor ω , we implement it with a multi-layer perceptron, followed by a softmax function.

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

Original: Training.
Revised: Training.

Original: We adopt a simple yet efﬁcient two-stage training strategy to search for optimal parameters and the training procedure of our method is summarized in Algorithm 1:
Revised: We adopt a simple yet efﬁcient two-stage training strategy to search for optimal parameters and the training procedure of our method is summarized in Algorithm 1:

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: Algorithm 1: The training procedure.
Revised: Algorithm 1: The training procedure.

Original: inference
Revised: inference

Original: predictor
Revised: predictor

Original:  Experiments
Revised:  Experiments

Original: Experiments are performed on 10 benchmark datasets and repeated 5 times with mean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPU and 128G RAM.
Revised: Experiments are performed on 10 benchmark datasets and repeated 5 times with mean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPU and 128G RAM.

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) [23].
Revised: Graph Benchmark (OGB) [23].

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 split contains a set of scaffolds (almost) different to each other.
Revised: Each split contains a set of scaffolds (almost) different to each other.

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

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

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 manual speciﬁed environment labels.
Revised: Notice that only the six datasets from DrugOOD provide manual speciﬁed environment labels.

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

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 (c)
Revised: Env3 Env (c)

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: GraphSAGE + ours.
Revised: GraphSAGE + ours.

Original: [Table]
Revised: [Table]

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 [68, 64, 58].
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 [68, 64, 58].

Original: Baselines.
Revised: Baselines.

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

Original: We adapt three backbones: GCN [30], GIN [63] and GraphSAGE [21] into our method.
Revised: We adapt three backbones: GCN [30], GIN [63] and GraphSAGE [21] 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 ” [19, 25, 38].
Revised: We also compare against another augmented version “ + virtual node ” [19, 25, 38].

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

Original: DeepCoral [55], DANN [18], MixUp [70] and GroupDro [50].
Revised: DeepCoral [55], DANN [18], MixUp [70] and GroupDro [50].

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

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

Original: For the sake of fairness, the embedding size of all methods are set to be equal in comparison.
Revised: For the sake of fairness, the embedding size of all methods 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 consistent improvements after adapted to our methods across all the four datasets released by OGB in terms of ROC-AUC.
Revised: As demonstrated in Table 1, baselines obtain consistent improvements after adapted to our methods across all the four datasets released by OGB in terms of ROC-AUC.

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

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

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

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

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

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

Original: [Table]
Revised: [Table]

Original: which minimizes the average empirical loss on training data, for comparison.
Revised: which minimizes the average empirical loss on training data, for comparison.

Original: Notice that ERM and MixUp don’t require manual speciﬁed environments labels.
Revised: Notice that ERM and MixUp don’t require manual speciﬁed environments labels.

Original: We also present the results of DANN, which requires manual speciﬁcations of environment and obtains competitive results with MixUp.
Revised: We also present the results of DANN, which requires manual speciﬁcations of environment and obtains competitive results with MixUp.

Original: All methods use GIN [63] as backbone.
Revised: All methods use GIN [63] as backbone.

Original: [Table]
Revised: [Table]

Original: Our
Revised: Our

Original: [Table]
Revised: [Table]

Original: GIN + new learning objective
Revised: GIN + new learning objective

Original: [Table]
Revised: [Table]

Original: DANN using our inferred environment label
Revised: DANN using our inferred environment label

Original: [Table]
Revised: [Table]

Original: Our model using given environment label
Revised: Our model using given environment label

Original: [Table]
Revised: [Table]

Original: Our full model
Revised: Our full model

Original: [Table]
Revised: [Table]

Original: tions.
Revised: tions.

Original: Except on IC50-size, our method outperforms all baselines across all datasets due to its ability to enforce the molecule encoder to leverage environment-invariant substructures that more stably relate with the labels across environments.
Revised: Except on IC50-size, our method outperforms all baselines across all datasets due to its ability to enforce the molecule encoder to leverage environment-invariant substructures that more stably relate with the labels across environments.

Original: 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 more data to train the model.
Revised: Different from the other methods, MixUp constructs more training exmaples and uses more data to train the model.

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

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

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

Original: 4.3 Ablation Study of Components
Revised: 4.3 Ablation Study of Components

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 introduced in Sec.
Revised: We study the impact of the attention-based architecture introduced in Sec.

Original: 3.3 by assembling this architecture with ERM loss.
Revised: 3.3 by assembling this architecture with ERM loss.

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

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

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

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

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

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

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

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

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

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

Original: Environment inference.
Revised: Environment inference.

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

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

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

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

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

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

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

Original: MixUp.
Revised: MixUp.

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

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

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

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

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

Original: 4.4 Hyper-parameter Sensitivity Study
Revised: 4.4 Hyper-parameter Sensitivity Study

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

Original: Fig.
Revised: Fig.

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

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

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

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

Original: Fig.
Revised: Fig.

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

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

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

Original: = 0 .
Revised: = 0 .

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

Original: 4.5 Risk Dynamics
Revised: 4.5 Risk Dynamics

Original: Additionally, to shed insights of the ability of our method to lower the risks of different environments, we visualize the risk dynamic curve of some environments in Fig. 3(c).
Revised: Additionally, to shed insights of the ability of our method to lower the risks of different environments, we visualize the risk dynamic curve of some environments in Fig. 3(c).

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

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

Original:  Conclusion
Revised:  Conclusion

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

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

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

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

Original: Additionally, our model achieves competitive or even superior performance compared to state-of-the-art models 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-art models designed for OOD learning that require manual speciﬁed environment labels as extra inputs.

Original: Acknowledgement
Revised: Acknowledgement

Original: This work was partly supported by National Key Research and Development Program of China (2020AAA0107600), National Natural Science Foundation of China (61972250, 72061127003), and
Revised: This work was partly supported by National Key Research and Development Program of China (2020AAA0107600), National Natural Science Foundation of China (61972250, 72061127003), and

Original: Shanghai Municipal Science and Technology (Major) Project (22511105100, 2021SHZDZX0102).
Revised: Shanghai Municipal Science and Technology (Major) Project (22511105100, 2021SHZDZX0102).