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

Original: Abstract
Revised: Abstract

Original: 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.
Revised: Molecule representation learning (MRL) has been extensively studied and current methods have shown promising power for various tasks, e.g., molecular property prediction and target identiﬁcation.

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

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

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

Original: We propose a new representation learning framework entitled
Revised: We propose a new representation learning framework entitled

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

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

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

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

Original: Particularly, 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: 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.

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

Original:  Introduction
Revised:  Introduction

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

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

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

Original: In general, MRL aims 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 [67], retrosynthetic analysis [62], search of antibiotics [52], virtual screening [38] for drug discovery, etc.
Revised: In general, MRL aims at embedding a molecule into a vector in latent space as a foundation model, on top of which the learned representations could be used for a variety of downstream tasks, such as target identiﬁcation [66], retrosynthetic analysis [62], search of antibiotics [52], virtual screening [38] 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, yet
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
Revised: O H

Original: O H
Revised: O H

Original: O H Cyclopropanol 1,4-Cyclohexanediol
Revised: 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 [38] 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 [38] 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 [39, 7, 42].
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 [39, 7, 42].

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 [59].
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 [59].

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 [26] point out that existing OOD models [3, 53, 17, 68, 48] fail to exhibit signiﬁcant improvement on MRL tasks against distribution shifts and even the simple
Revised: The work [26] point out that existing OOD models [3, 53, 17, 67, 48] fail to exhibit signiﬁcant improvement on MRL tasks against distribution shifts and even the simple

Original: Empirical Risk Minimization (ERM) [54] method outperforms these latest methods, which is also empirically veriﬁed by [16].
Revised: Empirical Risk Minimization (ERM) [54] 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 [30, 46, 70, 28] 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 [30, 46, 69, 28] 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 [31, 22] and molecule size [26].
Revised: The common practice speciﬁes environments as some prominent information of the molecules e.g. scaffold pattern [31, 22] and molecule size [26].

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 [22].
Revised: Thus, the data-generating environments and the induced distributions which these two molecules are sampled from can be considered different [22].

Original: Though sampled from different distributions, they are both readily soluble in water due to the invariant substructure hydroxy [23] shared across different environments.
Revised: Though sampled from different distributions, they are both readily soluble in water due to the invariant substructure hydroxy [23] 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 [26, 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 [26, 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 [45, 42, 3, 59].
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 [45, 42, 3, 59].

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

Original: As for a newly released benchmark for drug-oriented OOD learning [26], 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 [26], 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 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 explicitly models such substructures for OOD molecule representation.
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: To our best knowledge, this is the ﬁrst work for environment clustering for molecule representation, which has not been fulﬁlled in existing OOD literature beyond molecules (thanks to the speciﬁc substructure-property invariance priors).
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 [47, 3, 9, 59] assume that there is a potential environment variable e accounting for the distribution shift between the training and testing data.
Revised: Recent studies [47, 3, 9, 59] 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 ) [32].
Revised: | e ] is called the risk function under a given environment e and denoted as R e ( x e , y e ) [32].

Original: Invariant Learning.
Revised: Invariant Learning.

Original: There is an emerging line of research [47, 3, 10, 12] regarding invariant predictor learning, for solving the OOD generalization problem.
Revised: There is an emerging line of research [47, 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 [32].
Revised: These methods propose to ﬁnd an invariant predictor that could uncover invariant relationships between inputs and targets across all environments [32].

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 structured data like graphs [6, 59] and sequences [63].
Revised: Recent works adopt the invariance principle as a cornerstone for handling various distribution shifts on semistructured data like graphs [59, 6] and sequences [63].

Original: To our knowledge, we are the ﬁrst to leverage the invariance principle  for handling molecular property prediction 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 [55] or BERT [15].
Revised: They use language models to process the textual representation (SMILES) of a molecule, for example, Transformer [55] 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 Morgan ﬁngerprints [41].
Revised: The molecular ﬁngerprint techniques date back to the Morgan ﬁngerprints [41].

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

Original: Since molecules can be viewed as structured graphs, graph neural networks have been widely used to learn molecule representation [27, 18, 23].
Revised: Since molecules can be viewed as structured graphs, graph neural networks have been widely used to learn molecule representation [27, 18, 23].

Original: Existing general OOD methods [3, 53, 17, 68, 48] are not tailored to such non-Euclidean structured data, i.e. molecules.
Revised: Existing general OOD methods [3, 53, 17, 67, 48] are not tailored to such non-Euclidean structured data, i.e. molecules.

Original: In several recent works on molecule property classiﬁcation tasks [69, 57], 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 [68, 57], 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 [30, 46, 70, 28], 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 [30, 46, 69, 28], 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 [22] 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 [22] 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 [26, 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 [26, 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]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: d 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

Revised: O H

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

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

Revised: [Equation]

Original: O H 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) [60], to learn a graph representation given G .
Revised: We use a Graph Isomorphism Network (GIN) [60], 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 , encouraging the learned environment-partition to be uniform.
Revised: We could set the prior p τ ( e | G ) to Uniform distribution or Gaussian distribution , encouraging the learned environment-partition to be uniform.

Original: As for p τ ( y | G , e ) , we also choose a GNN model (e.g.
Revised: As for p τ ( y | G , e ) , we also choose a GNN model (e.g.

Original: GIN) followed by a softmax activation function to model it.
Revised: 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 [33].
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 [33].

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:

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

Original:  the parameters θ ; 
Revised: 2) optimizing the molecule encoder and the predictor: θ ∗ ← arg min θ L inv ( θ ; G train , τ ) . Algorithm 1: The training procedure. inference predictor

Original:  Experiments
Revised:  Experiments

Original: Experiments are performed on 10 benchmark datasets. Each experiment is repeated 5 times with mean and standard deviation reported, running on a machine with i9-10920X CPU, RTX 3090 GPU and 128G RAM.
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) [22].
Revised: Graph Benchmark (OGB) [22].

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 [26].
Revised: The other six datasets are generated by the dataset curator provided by DrugOOD [26].

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.Metric.
Revised: Refer to Appendix D for more details of datasets.

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 [66, 61, 56].

Original: Baselines.

Original: 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 [29], GIN [60]

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:  + 
Revised: GraphSAGE + ours.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: Metric. 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 [65, 61, 56]. Baselines. Ideally, any MRL method can be adapted into our method as backbone to improve their generalization ability against distribution shifts.

Original:  and GraphSAGE [20] into our method.
Revised: We adapt three backbones: GCN [29], GIN [60] and GraphSAGE [20] 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 ” [18, 24, 37].
Revised: We also compare against another augmented version “ + virtual node ” [18, 24, 37].

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

Original: DeepCoral [53], DANN [17], MixUp [68] and GroupDro [48].
Revised: DeepCoral [53], DANN [17], MixUp [67] and GroupDro [48].

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 observations.
Revised: Table 2 summarizes the results in comparsion with six state-of-the-art methods tailored for OOD learning, where we obtain the following observa-

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

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 [60] as backbone.
Revised: All methods use GIN [60] as backbone.

Revised: [Table]

Revised: Our

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original:  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:  full 
Revised: Our full model

Original: [Table]
Revised: [Table]

Original:  relate with the labels across environments.
Revised: tions. 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, [26, 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, [26, 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).

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

Original: Proof.
Revised: Proof.

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

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

Original: [Equation]
Revised: [Equation]

Original: Rearrange Eq.
Revised: Rearrange Eq.

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Assumption 1.
Revised: Assumption 1.

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

Original: Deﬁne r as a random variable of r G and it satisﬁes: 1) (Invariance condition): p ( y | r , e ) =
Revised: Deﬁne r as a random variable of r G and it satisﬁes: 1) (Invariance condition): p ( y | r , e ) =

Original: p ( y | r ) , and 2) (Sufﬁciency condition): y = h ( r ) + n , where h is a nonlinear function, n is a independent noise.
Revised: p ( y | r ) , and 2) (Sufﬁciency condition): y = h ( r ) + n , where h is a nonlinear function, n is a independent noise.

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

Original: Lemma 1.
Revised: Lemma 1.

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

Original: Proof.
Revised: Proof.

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

Original: =
Revised: =

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

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

Original: =
Revised: =

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

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

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

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

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

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

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

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

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

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

Original: =
Revised: =

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

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

Original: Proof.
Revised: Proof.

Original: According to the dependency relationship z ← G → y
Revised: According to the dependency relationship z ← G → y

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: Based on
Revised: Based on

Original: Jensen
Revised: Jensen

Original: Inequality and
Revised: Inequality and

Original: Triangle
Revised: Triangle

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

Original: D KL
Revised: D KL

Original: [Equation]
Revised: [Equation]

Original: e
Revised: e

Original: [Equation]
Revised: [Equation]

Original: G )
Revised: G )

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: y | G
Revised: y | G

Original: [Equation]
Revised: [Equation]

Original: e )
Revised: e )

Original: [Equation]
Revised: [Equation]

Original: is upper bounded by:
Revised: is upper bounded by:

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: Proof.
Revised: Proof.

Original: Given the dependency relationship z ← G → y
Revised: Given the dependency relationship z ← G → y

Original: [Equation]
Revised: [Equation]

Original: we hold max q ( z | G ) I( y ; z ) is equivalent to min q ( z | G ) I( y ; G | z ) .
Revised: we hold max q ( z | G ) I( y ; z ) is equivalent to min q ( z | G ) I( y ; G | z ) .

Original: Also we have
Revised: Also we have

Original: [Equation]
Revised: [Equation]

Original: Based on this
Revised: Based on this

Original: [Equation]
Revised: [Equation]

Original: we will have
Revised: we will have

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: Three backbones.
Revised: Three backbones.

Original: We adapt three backbones into our method, namely, GCN [29], GraphSAGE [20] and GIN [60].
Revised: We adapt three backbones into our method, namely, GCN [29], GraphSAGE [20] and GIN [60].

Original: We also emphasize the comparison with their augmented versions, i.e. “ + virtual node ” [18, 24, 37].
Revised: We also emphasize the comparison with their augmented versions, i.e. “ + virtual node ” [18, 24, 37].

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

Original: Graph Benchmark [22] 5 .
Revised: Graph Benchmark [22] 5 .

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

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

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

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

Original: • GIN [60] generalizes the Weisfeiler-Lehman (WL) graph isomorphism test [35] and hence achieves maximum discriminative power among GNNS.
Revised: • GIN [60] generalizes the Weisfeiler-Lehman (WL) graph isomorphism test [35] and hence achieves maximum discriminative power among GNNS.

Original: • GraphSAGE [20] learns a function that generates embeddings by sampling and aggregating features from a node’s local neighborhood.
Revised: • GraphSAGE [20] learns a function that generates embeddings by sampling and aggregating features from a node’s local neighborhood.

Original: • GCN/GIN/GraphSAGE + virtual node [18, 24, 37] is a variant of the original method augmented by an additional node connecting to all nodes in the raw graph.
Revised: • GCN/GIN/GraphSAGE + virtual node [18, 24, 37] is a variant of the original method augmented by an additional node connecting to all nodes in the raw graph.

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

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

Original: ERM [54], IRM [3], DeepCoral [53], DANN [17], MixUp [68] and GroupDro [48].
Revised: ERM [54], IRM [3], DeepCoral [53], DANN [17], MixUp [67] and GroupDro [48].

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

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

Original: The embedding
Revised: The embedding

Original: size of all models including ours are all set to 128 for fairness.
Revised: size of all models including ours are all set to 128 for fairness.

Original: • ERM [54] minimizes the average empirical loss on training data.
Revised: • ERM [54] minimizes the average empirical loss on training data.

Original: • IRM [3] penalizes feature distributions for environments that have different optimum predictors.
Revised: • IRM [3] penalizes feature distributions for environments that have different optimum predictors.

Original: We set the penalty weight and the penalty anneal iteration to 10 and 500 , respectively.
Revised: We set the penalty weight and the penalty anneal iteration to 10 and 500 , respectively.

Original: • DeepCoral [53] penalizes differences in the means and covariances of the feature distributions for each environment, which are exactly the distribution of last layer activations in a neural network.
Revised: • DeepCoral [53] penalizes differences in the means and covariances of the feature distributions for each environment, which are exactly the distribution of last layer activations in a neural network.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: com/tencent
Revised: com/tencent

Original: D More Details of Datasets
Revised: D More Details of Datasets

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

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

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

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

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

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

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

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

Original: Each molecule has a label indicating whether it binds to human beta-secretase 1.
Revised: Each molecule has a label indicating whether it binds to human beta-secretase 1.

Original: • SIDER is a dataset of marked drugs and adverse drug reactions (ADRs).
Revised: • SIDER is a dataset of marked drugs and adverse drug reactions (ADRs).

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

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

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

Original: • IC50/EC50-scaffold/assay/size are datasets generated by the automated dataset curator provided by DrugOOD from the large-scale bioassay deposition website ChEMBL [40].
Revised: • IC50/EC50-scaffold/assay/size are datasets generated by the automated dataset curator provided by DrugOOD from the large-scale bioassay deposition website ChEMBL [40].

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

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

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

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

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

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

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

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

Original: Tableshows the detailed statistics of datasets.
Revised: Tableshows the detailed statistics of datasets.

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

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

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

Original: [Table]
Revised: [Table]

Original: Size (cid:33)
Revised: Size (cid:33)

Original: [Table]
Revised: [Table]

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

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

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

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

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

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

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

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

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

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

Original: E Understanding the Data-generating Process
Revised: E Understanding the Data-generating Process

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: The abstract data variables are denoted by the nodes and the directed arrows represent the causalities.
Revised: The abstract data variables are denoted by the nodes and the directed arrows represent the causalities.

Original: This SCM illustrates the causalities among variables: E as the environment, S as the spurious substructures
Revised: This SCM illustrates the causalities among variables: E as the environment, S as the spurious substructures

Original: , C as the invariant/causal substructures w.r.t Y , G as the instance molecule graph, Y as the ground-truth label.
Revised: , C as the 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 the substructure hydroxy ( − OH ) , S as the substructures aside from hydroxy, G as Cyclopropanol , Y as good water solubility.
Revised: Taking Cyclopropanol ( C 3 H 6 O ) as an example, we can specify E as the 3C-ring scaffold, C as the substructure hydroxy ( − OH ) , S as the substructures aside from hydroxy, G as Cyclopropanol , Y as good water solubility.

Original: The good water solubility Y is only attributed to the invariant substructure hydroxy, i.e. C , rather than other spurious substructures S .
Revised: The good water solubility Y is only attributed to the invariant substructure hydroxy, 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 spurious correlations between irrelevant substructures S and the target label Y will be encoded to learned molecular representations.
Revised: Hence, the spurious correlations between irrelevant substructures S and the target label Y 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: F Notations
Revised: F Notations

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

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

Original: For all experiments in our original paper, we all adopt breaking retrosynthetically interesting chemical substructures (BRICS) to segment molecule into substructures, which is widely used in other works related to molecules [58,64, 11].
Revised: For all experiments in our original paper, we all adopt breaking retrosynthetically interesting chemical substructures (BRICS) to segment molecule into substructures, which is widely used in other works related to molecules, e.g., [58, 11].

Original: To investigate the sensitivity of our method to different decomposing strategies, we adopt another molecule segmentation method called retrosynthetic combinatorial analysis procedure (RECAP)
Revised: To investigate the sensitivity of our method to different decomposing strategies, we adopt another molecule segmentation method called retrosynthetic combinatorial analysis procedure (RECAP)

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

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

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

Original: We cansee that
Revised: We can

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

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

Original: H Future Direction
Revised: H Future Direction

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

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

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

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

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

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

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

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

Original: As demonstrated in Table 8,
Revised: As demonstrated in Table 8,

Original: [Table]
Revised: [Table]

Original: we can see that adding such a SAB further improves our model on EC50-scaffold.
Revised: we can see that adding such a SAB further improves our model on EC50-scaffold.

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

Original: I Limitations
Revised: I Limitations

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

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

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

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

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