Enforcing Predictive Invariance across Structured Biomedical Domains
Authors: Wengong Jin
Regina Barzilay
Tommi S. Jaakkola
Venue:  (2021)
Keywords: Domain Generalization
Molecular Property Prediction

Abstract: Many biochemical applications such as molecular property prediction require models to generalize beyond their training domains (environments). Moreover, natural environments in these tasks are structured, defined by complex descriptors such as molecular scaffolds or protein families. Therefore, most environments are either never seen during training, or contain only a single training example. To address these challenges, we propose a new regret minimization (RGM) algorithm and its extension for structured environments. RGM builds from invariant risk minimization (IRM) by recasting simultaneous optimality condition in terms of predictive regret, finding a representation that enables the predictor to compete against an oracle with hindsight access to held-out environments. The structured extension adaptively highlights variation due to complex environments via specialized domain perturbations. We evaluate our method on multiple applications: molecular property prediction, protein homology and stability prediction and show that RGM significantly outperforms previous state-of-the-art baselines.

Under review as a conference paper at ICLR 2021
ENFORCING PREDICTIVE INVARIANCE STRUCTURED BIOMEDICAL DOMAINS
ACROSS
Anonymous authors Paper under double-bling review
ABSTRACT
Many biochemical applications such as molecular property prediction require models to generalize beyond their training domains (environments). Moreover, natural environments in these tasks are structured, defined by complex descriptors such as molecular scaffolds or protein families. Therefore, most environments are either never seen during training, or contain only single training example. To address these challenges, we propose a new regret minimization (RGM) algorithm and its extension for structured environments. RGM builds from invariant risk minimization (IRM) by recasting simultaneous optimality condition in terms predictive regret, finding representation that enables the predictor to compete against an oracle with hindsight access to held-out environments. The structured extension adaptively highlights variation due complex environments via spe- cialized domain perturbations. We evaluate our method on multiple applications: molecular property prediction, protein homology and stability prediction and show that RGM significantly outperforms previous state-of-the baselines.
INTRODUCTION
In many biomedical applications, training data is necessarily limited or otherwise heterogeneous. It is therefore important ensure that model predictions derived from such data generalize substantially beyond where the training samples lie. For instance, in molecule property prediction (Wu et al. [0188). models are often evaluated under scaffold split, which introduces structural separation between the chemical spaces of training and test compounds. In protein homology detection (Rao et al., 2019), the split is driven by protein superfamily where entire evolutionary groups are held out from the training set, forcing models to generalize across larger evolutionary gaps.
The key technical challenge to be able to estimate models that can generalize beyond their training data. The ability generalize implies notion of invariance to the differences between the available training data and where predictions are sought. A recently proposed approach known as invariant risk minimization (IRM) (Arjovsky et al. 2019) seeks to find predictors that simultancously optimal across different such scenarios (called environments). Indeed, one can apply IRM with environments corresponding molecules sharing the same scaffold (Bemis & Murcko 1996) proteins from the same family (EI-Gebali et al.) (019) (see Figure [].. However, this is chall challenging since, for example, scaffolds are structured objects and can often uniquely identify each example in the training set. is not helpful to create single-example environments as the model would see any variation from one example another as scaffold variation.
In this paper, we propose a regret minimization algorithm to handle both standard and structured environments. The basic idea is to simulate unseen environments by using part of the training set as held-out environments E.. We quantify generalization in terms of regret- - the difference between the losses of two auxiliary predictors trained with and without examples in F.,. This imposes a stronger constraint on & and avoids some undesired representations admitted by IRM. For the structured environments like molecular scaffolds, we simulate unseen environments by perturbing the representation c. The perturbation is defined the gradient of auxiliary scaffold classifier with respect to o. The difference between the original and perturbed representation highlights the scaffold variation to the model. Its associated regret measures how well predictor trained without perturbation generalizes to the perturbed examples. The goal is to characterize the scaffold variation without explicitly creating environment for every possible scaffold.

---

Under review as a conference paper at ICLR 2021
scaffold
molecule
100
Histogram of scaffold environments
Scaffold
Training environments
10
75% of environments (one example)
ID
Molecule
Property
Testing
environmentss
OR
Figure 1: Left: Data generation process for molecule property prediction. Training and test envi- ronments are generated by controlling the scaffold variable. Middle: Scaffold is subgraph of molecular graph with its side chains removed. Right: In a toxicity prediction task (Wu et al. 2018). there are 1600 scaffold environments with 75% of them having a single example.
Our methods are evaluated real-world datasets such molecule property prediction and protein classification. We compare our model against multiple baselines including IRM, MLDG (Li et al.) 2018a) and CrossGrad (Shankar et al. 2018). On the QM9 dataset (Ramakrishnan et al. 2014), we outperform the best baseline by wide margin across multiple properties (41.7 v.8 52.3 average MAE) under an extrapolation evaluation. On protein stability dataset (Rocklin et al. (017). we achieve new Shate-of-the--aee art results compared to Rao et al. (2019) (0.79 v.s. 0.73 spearman's p).
RELATED WORK
Generalization challenges in biomedical applications The challenges of generalization have been extensively documented in this area. For instance, Yang et al. (2019); Rao et al. (2019): Hoou et al. (2018) have demonstrated that state-of-the-ar models exhibit drop in performance when tested under seaffold or protein family split. De facto, the scaffold split and its variants (Feinberg al. 2018 used so commonly in cheminformatics as they emulate temporal evaluation adopted in pharmaceutical industry. Therefore, the ability to generalize to new scaffold or protein family environments the key for practical usage of these models. Moreover, input objects in these domains are typically structured (e.g., molecules are represented by graphs (Duvenaud al. 2015 Dai et al. 2016 Gilmer al.. 2017)). This characteristic introduces unique challenges with respect to the environment definition for IRM style algorithms.
Invariance Prior work has sought generalization by enforcing appropriate invariance constraint over learned representations. For instance, domain adversarial network (DANN) (Ganin et al., 2016)
Zhao et al. 2018) enforces the latent representation Z = Φ(X) to have the same distribution across different environments E (i.e, Z 1 E). However, this forces predicted label distribution P(Y 2 to be the same across all the environments (Zhao et al., 2019). Long al. (2018) y Li et al. (20 8c);
Combes et al. (2020) extends the invariance criterion by conditioning on the label in order to address the label shift issue of DANN. Invariant risk minimization (IRM) (Arjovsky et al.) 2019) seeks different notion of invariance. Instead of aligning distributions of Z, IRM requires that the predictor operating on Z = (X) is simultaneously optimal across different environments. The associated independence is Y I E' Z. Various work (Krueger et al. 2020 Chang et al.) (020) has sought to extend IRM. We focus on the structured setting. where most of the environments can uniquely specify X in the training set. As result, E would act similarly to X. In the extreme case, the IRM principle reduces to 1 X Z, which is not the desired invariance criterion. We propose to address this issue by introducing domain perturbation to adaptively highlight the structured variation.
Domain generalization These methods seek to learn models that generalize to new domains (Muan) det et al. 2013 Ghifary et al. 2015: Motiian et al., 2017 Li et al. 2017: 2018b). Domain general- ization methods can be roughly divided into three categories: domain adversarial truining (Ganin
et al.) 2016 Tzeng et al.; 2017 Long et al.) 2018), meta-learning (Li ct al. 2018a] Balaji et al. 2018
Li et al. 2019a by Dou et al., 2019) and domain augmentation (Shankar et al. 2018 Volpi et al. 2018). Our method resembles meta-learning based methods in that create held-out environments to simulate domain shift during training. However, our objective seeks to reduce the regret between predictors trained with or without access to the held-out environments.
Existing domain generalization benchmarks assume that each domain contains sufficient amounts of data. We focus different setting where most of the environments contain only few (or single)
2

---

Under review as a conference paper at ICLR 2021
examples since they defined by structured descriptors. This setting often arises in chemical and biological applications (see Figure D. Similar to data augmentation method in Shankar et al. (2018), our structured RGM also creates perturbed examples based domain domain-guided perturbations. However, our method operates over learned representations since our inputs are discrete. Moreover, the perturbed examples are only used regularize the feature extractor Φ via the regret term.
REGRET MINIMIZATION
To introduce method, start with standard setting where the training set 7> is comprised of environments & {Ei.... En} (Arjovsky et al. 2019). Each environment E, consists examples (xr, y) randomly drawn from distribution P.. Assuming that new environments at test time exhibit similar variability as the training environments, our goal is to train a model that generalizes to such new environments Erest Suppose our model consists of components f ø. where the predictor f operates on the feature extractor Φ. Let L (ʃ o ()) = (aaa)e f (((())) be its empirical loss in environment Eg and [(f の) E. (f Φ). IRM learns o and such that f is simultaneously optimal in all training environments:
may encounter
min L(f o σ) 0.
s.t. Ve: f arg [" Φ)
(1)
h
One possible way solve this objective is through Lagrangian relaxation:
min 6.8
C(f o Ø) Σ de (f
L (h o))
h.
The regularizer (f - min, [ o) measures the performance gap between f and the best predictor à E Fe(Ø)= arg min, (h o) specific environment Ee. Note that both and are trained and evaluated on examples from environment Eg. This motivates us to replace the regularizer with predictive regret. Specifically, for each environment Ee, define the associated regret R (Ø) as the difference between the losses of two auxiliary predictors trained with and without access to examples (2r, y) E E,:
(d) = (f-. - min =
L"- L"e o)
(3)
lef
where the two auxiliary predictors are obtained from (assuming bounded and closed):
f-e F-e(o) arg min fe E F.( arg ( o)
Σ ch
(4)
hET
hET kyff
The oracle predictor fe is trained on environment Ec, while f-e uses the rest of the environments E\ EMEE, E2)) for training but is tested F., Note that R (Ø) does not depend the predictor f we seeking to estimate; it is function of the representation @ as well as the two auxiliary predictors f-c and fe. For notational simplicity, predictors are evaluated on the same set of training examples in E., we immediately have
have omitted R (ø)'s dependence on f-e and fe. Since both
Proposition The regret R (Ø) always non-negative for any representation cb
The proof is straightforward since fe is the minimizer of (f' (f' o Φ) and both fe and f-e are drawn from the same parametric family J. The overall regret R(Ø) = Σ. (²) expresses our stated goal of finding representation é that generalizes each held-out environment. Our regret minimization (RGM) objective regularizes the empirical loss with a regret term weighted by X:
LRGM = x) + RT)
(5)
3.1 COMPARISON WITH IRM
Compared to IRM, the proposed RGM objective imposes stronger constraint on Ø since f-e is not trained on Ec. To show this formally, let (o). (o) be the set of optimal predictors E. and E\{E, respectively defined in Eq. Since R" (²)Ø) = 0 + f-e F. (Ø) and f-e is chosen arbitrarily from (o), the constrained form the RGM objective can be stated as
min L(f o ②)) 6,
s... Ve (o) S Fe(
(6)
3

---

Under review as a conference paper at ICLR 2021
Φ
F-2(4)
Figure 2: A counterexample illustrating that FIRM PRGM- The environments are generated by different translations of X1- For the identity mapping ((X) (X1, X2) and the true hypothesis is 2 > 0j. There exists a predictor fIRM which is simultaneously optimal in all environments. In contrast, ø is not feasible under RGM because there is a linear classifier (-) that is optimal in environment E₁ but performs poorly in environment E2.
E F2(()
FIRM
$1
E
The analogous IRM constraints are f ( and (gFe(6)= + 0. Suppose both IRM and RGM constraints are feasible and let LiRM: CROM be their optimal loss respectively. Consider the set of optimal features under both objectives:
PIRM = {ol
min C(f ))) LiRM (eFe(Ø) + 0}
(en)F.(d)
RGM = { min L(J o ()) = fiRM. Ve: F-o(d) S Fo())} fEF
(8)
Proposition 2. Assuming two environments, if CROM LiRM then PRGM C 1IRM- The converse PIRM S PRGM does not hold in general.
While limited to two environments, the proposition suggests that RGM imposes stronger constraints on Φ. Figure 2 shows a counterexample illustrating that IIRM & PRGM. Suppose there are two environments generated by translation of X1 and the true hypothesis is I[X² > 0]. The identity mapping ((X) (X1. X2) is not translation invariant, but ø E PIRM because there exists predictor fIRM that is simultaneously optimal in all environments. On the other hand, is not feasible under RGM because there is linear classifier h. E F (1) that is optimal in F but suboptimal in Fez, violating the RGM constraint (0) C F2(()) Thus & & P
To see why it would be helpful to add a stronger constraint on 6, consider the following data generation process where the environment can be inferred from input x alone:
p(a,y, e) (xle)((ale)((ull,ce)) e);
plulæ, = p(yl e(x:))
For molecules and proteins, this assumption is often valid because the environment labels (scaffolds, protein families) typically depend on z only. We call o label-preserving if retains all the information about the label: (()l(=)) = p(y]a, e). Such representation may not generalize to new environments given the dependence associated ERM optimal predictor also satisfies the IRM constraints:
through o. However,
can show that for any label-preserving o, its
Proposition For any label-pree preserving é with ply ((x)) = p(ylx,,) e). associated ERM optimal predictor f satisfies the IRM constraint. Moreover, if ø E PIRM. f 6 is optimal under IRM.
While IRM constraints are vacuous for any label paeee-ire Φ, this is not necessarily the case with RGM constraints. Consider, for example, the counterexample in Figure 4 The identity mapping x(X) = (x) is label preseevig since retains all the input information. However, o is infeasible under RGM.
3.2 STRUCTURED ENVIRONMENTS
Now let us consider a more challenging setting, where the environments {Ex} are structured (i.e., k is a structured object rather than an integer). Formally, the training set comes in the form {1ri, Vi, s;)), where 8; is the structured environment label of (I,, ys) E Es;- For instance, in molecule property prediction, 8; is defined as the Murcko scaffold (i.e., subgraph) of molecule I;- It is hard to model scaffolds standard environments because they structured descriptors and often uniquely identify each molecule in the training set (Figure D.. When an environment has only molecule, the model cannot decide which subgraph of that molecule is the right scaffold. Thus, creating mangle-example environments is not helpful for domain generalization.
describe scaffold variation by perturbation in the representation d. The idea is to create a perturbed instance ž; for each example (Iz; Vı. 8;) so that the difference between æ; and ž; highlights how scaffold information has changed in the representation. Specifically, the perturbation
Alternativel,,
4

---

Under review as a conference paper at ICLR 2021
a
(
(t))
b
scaffolds
1 regret
y
1
regret
scaffold
f-]
f
perturbation
x(((x))
((x) + 8(x)
Molecule
Input.x
Virtual environment
Figure 3: a) Structured RGM: introduce additional oracle predictors fe for the perturbed inputs; b) In molecule tasks, the scaffold classifier g is trained by negative sampling.
Algorithm Structured RGM: Forward Pass
1: for each environment Eg & do
2: Sample a minibatch В, from environment E.
3: Compute scaffold classification loss Cg(g Φ) over Be.
4:
Construct perturbed examples Be from Be via gradient perturbation (see Eq. 10).
5:
Compute empirical loss [(f o の) Be.
6:
Compute auxiliary predictor loss [- Φ) on B-c.
7:
Compute oracle predictor losses (Je () and ( o (sb + 6)) on B, and B..
8:
Compute regret terms (6),, (6+ + on Be and Be.
9:
end for
8(x;) is defined through parametric scaffold classifier built top of the representation of The associated scaffold classification loss (ss, g ((xx))). Given that our inputs are discrete, we define the perturbation ō the gradient with respect the continuous representation d:
()t.) = <(x,) + 6(x,) = <(x,) + a (si, (())
(10)
where a is step size parameter. The perturbation is specifically designed to contain less information about the scaffold 8,, and we require that the model should not be affected by this variation in the representation. Since these perturbations introduce additional simulated test scenarios that wish to generalize to, we propose to regularize our model also based on regret associated with perturbed inputs. Similar to 33. the regret corresponding to perturbed inputs is defined as R (Ø + 8):
R((
L (6+ min, ( (0 ( 6))
(11)
L (
Σ E E(y, h(o(z,()())))))))))))))(())))()) + 8(x:)))
(12)
This introduces new oracle predictor fe arg min, ( (Ø + 8)) for each environment E. (see Figure 3d). Note that f-e is the same auxiliary predictor before. It minimizes separate objective C-- o Φ). which does not include the perturbed examples.
The structured RGM (SRGM) objective LSRCM augments the basic RGM with additional regret terms as well as the scaffold classification loss Lg ()::
LSRGM
E.E. (0
R' + vh)
(13)
Lg(go0) Σ En (( ((+()))))))
(14)
The forward pass of SRGM is shown in Algorithm I Since is structured object with a large number of possible values, we train the classifier g with negative sampling (Figure 3b). Note that is also updated to partially optimize Ly. This is necessary to ensure that the scaffold classifier operating on á has enough information to introduce a reasonable gradient perturbation 5(x). This trade-off keeps some scaffold information in while ensuring, via the associated regret terms, that this information is not strongly relied upon. The effect of this design choice is studied in our experiments.
'Or method is introduced using scaffolds examples. It can be applied other structured environments like protein families by simply replacing the scaffold classifier with protein family classifier.
5

---

Under review as a conference paper at ICLR 2021
minimatch B
Feature extractor A
Environment E,
L9 .b)
fee
Environment Fj
minimatch Bj
Regret R(())
fe
minimatch
al.-
L'(
Environment E B
a
alc
Gradient
Reversal
of
Figure 4: the RGM forward pass, we sample a minibatch B. from each environment Eg and compute regret Rº (ő). In the backward pass, the gradient of (fe ø) goes through a gradient reversal layer (Ganin et al. 2016) which negates the gradient during back-propagatton.
3.3 OPTIMIZATION
The standard RGM objective in Eq. can be viewed as finding a stationary point of a multi- player game between f. well the auxiliary predictors {f-e) and {fe}. Our predictor f and representation o find their best response strategies by minimizing
min
{C( o X E. (c= -(o f
(15)
while the auxiliary predictors minimize
min ( and min L Ø) Ve fe
(16)
f-c
This multi-player game can be optimized by stochastic gradient descent. Since fe and øi optimizes (ffe o σ) opposite directions, we introduce a gradient reversal layer (Ganin et al. 2016) between 6 and fe. This allows us to update all the players in a single forward-backwe pass (see Figure 4 In each step, we simultaneously update all the players with learning rate ny:
- o Φ) nd Σ. (o)
- -- o ()) Ve
where L. (f-c() - E' -o-eo00 x). In each step, we sample minibatches BI, B, from each environment E,,... E.. The loss L(f d) is computed over all the minibatches Up BR, while "f o Φ) is computed over minibatches B = Ukte BR. The regret term R (Φ) is evaluated based on examples in B. only.
(f-co
For structured RGM, its optimization rule is analogous to RGM, with additional gradient updates for the oracle predictors fc and scaffold classifier (see Appendix A.4). While the perturbation 8 defined on the basis of & and g. we do not include the dependence during back- propaproppgttion as incorporating this higher order gradient does not improve empirical results.
4
EXPERIMENTS
Our methods (RGM and SRGM) are evaluated on real-world applications such as molecular property prediction, protein homology and stability prediction. Our baselines include:
Standard empirical risk minimization (ERM) trained on aggregated environments;
Domain adversarial training methods including DANN (Ganin et al. 2016) and CDAN (Long et al.
2018), which seek to learn
features;
IRM (Arjovsky et al. 2019) requiring the model to be simultaneously optimal in all environments:
MLDG (Li et al. 201 8a). meta-learning method which simulates domain shift by dividing training
environments into meta-training and meta-testing:
CrossGrad (Shankar et al. 2018)) which augments the training set with domain domain-guided perturbations of inputs. Since our inputs discrete, we perform perturbation the representation instead.

---

Under review as a conference paper at ICLR 2021
Table Mean absolute error (MAE) on the QM9 dataset. Models are evaluated under scaffold split. Due to space limit, we only show standard deviation for the top three methods in subscripts.
Categorical methods
Structured methods
Property
ERM
DANN
CDAN
IRM
MLDG
RGM
CrossGrad
SRGM
0.736
0.709
0.748
1.059
0.745
0.682.057)
0.745
0.720
alpha
3.455
3.525
3.668
3.711
3.261
2.600
3.563
2.694
HOMO
0.011
0.011
0.012
0.011
0.010
0.01 (,00))
(.012
0,01
LUMO
0.020
0.020
0.021
0.021
0.020
0.020, 002)
0.017
0.019
gap
0.021
0.020
0.021
0.022
0.021
0.020 002)
0.019
(.01
R2
119.5
117.1
120.4
174.2
110.8
113.9 (13.9
112.2
100.5
ZPVE
0.008
0.008
0.009
0.009
0.009
0.008
0.008
0.007
Cv
1.917
1.960
2.093
2.344
2.029
2.268(417)
1.702 330)
2.133
17.50
18.50
18.25
16.21
16.24
14.93
20. (55000
13.97
20.11
20.51
20.41
16.72
17.65
14.39:2.57
(4.52
12.67(0822)
H
17.40
17.34
18.11
16.53
14.77
13.97(1.m)
(8.55
12.80
G
17.67
18.63
19.09
17.68
16.14
13.53
17. 17.95
13.1 (3.15
Training scaffolds
Validation scaffolds
Test scaffolds
OB
NB
HO XH
RQ
HO
106. 319 120 330 140 150
Figure 5: Examples of scaffolds in the QM9 dataset (highlighted in grey). We split the data based on scaffold complexity. Thus, the test scaffolds are structurally distinct from scaffolds in the training set. As shown in the right figure, the molecular weight distribution of training, validation and test sets are similar. This shows that scaffold complexity split is more realistic than molecular weight split.
These methods fall into two categories. SRGM and CrossGrad are structured methods as they can leverage the structural information of the environment (e.g., scaffold). RGM and other methods categorical methods since they do not utilize the structure and simply treat each environment as a set.
4.1 MOLECULAR PROPERTY PREDICTION
Data The training data consists of {(x;, ys, s:)). where x; a molecular graph, y/ its property and s; is its scaffold. We adopt four datasets from the MoleculeNet benchmark (Wu et al. 2018):
QM9 is a regression dataset of 134K organic molecules with up to 9 heavy atoms. Each molecule is labeled with 12 quantum mechanical properties.
HIV is a classification dataset of 42K molecules. Each molecule is associated with binary label indicating whether it is HIV inhibitor.
Tox21 is a classification dataset of 8.8K molecules. Each compound has 12 binary labels for toxicity measurements.
The blood-brain barrier penetration (BBBP) dataset contains 2K molecules. Each molecule is labeled with a binary permcability label.
Data split To test whether a model generalizes to new domains, important to create a test set that is distributionally distinct from the training set. Scaffold split (Wu et al. 2018) is common framework for this purpose. Molecules are clustered based on its Bemis-Murcko scaffold (Bemis & Murcko, 1996) and random subset of scaffolds selected into test set. However, this approach degenerates to random split when most scaffold clusters contain only one molecule (see Figure I). To address this issue, Feinberg al. (2019) proposed molecular weight split, where test molecules
7

---

Under review as a conference paper at ICLR 2021
Table Left: Results on molecule and protein datasets. CrossGrad and SRGM use the structure of environments (scaffolds or protein superfamily) while others do not. Right: Ablation study of SRGM. Detach= means we do not update & to optimize the scaffold (or protein superfamily) classification loss Lg Accs stands for the scaffold/protein superfamily classification accuracy. Property is the property prediction performance (AUROC for molecules, top-1 accuracy for protein).
HIV
Tox21
BBBP
Protein
detach
Accs
property
ERM
(.715
(.64|(000)
0285
20.9%
88.1%
0.736
HIV
DANN
(.727
(.639.006)
0.857c016)
22.3%
99.4%
0.751
CDAN
(o735
0.639(008)
(.85
21.9%
V
92.2%
0,640
Tox21
IRM
(.74
0.632
0,86
21.0%
X 97.1 %
0,649
MLDG
0.724
0.637
(.149
22.0%
V
73.7%
0.871
BBBP
RGM
(.79)
(.637
028
23.4%
94.0%
0.891
CrossGrad
0.746)015)
0.644.005)
0. (038)
20.9%
29.3%
21.9%
Protein
SRGM
(.75)
(064
(.85)
23.8%
X
33.5%
23.8%
are much bigger than molecules in the training set. While this creates strong structural distinction between the training and test sets, it is not realistic as the scaffold split.
Given these observations, we propose a variant of scaffold split called scaffold complexity split. We define the complexity of a scaffold the number of cycles in the scaffold graph. Specifically, we put a scaffold in the test set its scaffold complexity is greater than T and the training set if it is less than т. We set 7 2 for QM9 and 4 for other datasets. As shown in Figure S this forces the test scaffolds to be structurally different from the training scaffolds. It is also more realistic than the molecular weight split since the molecular weight distribution of training and test sets similar.
Model The molecule encoder & is a graph convolutional network (Yang et al. 2019) which translates molecular graph into continuous vector. The predictor is two-layer MLP that takes ((x) as input and predicts the label. The scaffold classifier g is also a two-layer MLP trained by negative sampling since scaffold is combinatorial object with a large number of possible values. Specifically, for a given molecule z; with scaffold s;. we randomly sample K other molecules and take their associated scaffolds { as negative classes. Details of model architecture and dyper-parametent are discussed in the appendix.
Results Following Wu et al. (2018). we report mean absolute error (MAE) for QM9 and AUROC for other datasets. All the results are averaged across five independent runs. Our results on the QM9 dataset are shown in T RGM outperforms other categorical methods and demonstrates clear improvement on six properties (mu, alpha, UO, U, H, G). SRGM outperforms all baselines on seven properties, with a significant error reduction on R2, UO, U, H and (3-10%). Compared to and alpha. On the three classification
RGM, SRGM performs better on all properties except datasets, SRGM also achieves statc-of-the- compared to all the baselines (see Table 2 These results confirm the advantage of exploiting the structure of environments.
4.2 PROTEIN HOMOLOGY PREDICTION
Data The protein homology dataset (Fox et al. 2013 Rao et al. 2019) consists of tuples (Z, Vi. 8:)]; where F; is protein represented as sequence of amino acids, V/4 its fold label and s; its superfamily label. The task is to predict the fold label y/. There are 1195 fold classes and 1823 protein superfamilies in total. Around 1200 superfamilies have less than 10 instances in the training set.
Data split Provided by Rao et al. (201)), the dataset consists of 12K instances for training, 736 for validation and 718 for testing. The dataset is split based protein superfamilies. As a result, proteins in the test set are structurally distinct from the training set, requiring models to generalize across large evolutionary gaps.
Model Our protein encoder & is a pre-trained BERT model (Rao et al. 2019). To generate a sequence-length invariant protein embedding, simply take the mean of all the vectors output by BERT. The predictor ff is a linear function that takes ø(z) as input and predicts its fold class. The superfamily classifier g is a two-layer ML.P The hyperparameters listed in the appendix.
8

---

Under review as a conference paper at ICLR 2021
Table 3: Comparison between SRGM with different perturbations on the QM9 dataset. "Scaffold" means perturbation via the gradient of the scaffold classifier. "Random" means random perturbation.
SRGM
mu
alpha
homo
lumo
gap
zpve
U0
H
Scaffold
0.72
2.69
0.011
0.019
0.019
100.5
0.007
2.13
13.97
12.67
12.80
13.15
Random 0.77
2.90
0.010
0.017
0.020
115.9
0.008
2.46
18.57
13.80
16.75
18.55
Table 4: SRGM performance on the QM9 dataset with different number of graph convolutional layers in ø. Adding more layers increases model complexity.
mu
alpha
homo
lumo
gap
zpve
Cv
U0
U
H
2 layer
0.69
3.06
0.010
0.014
0.018
106.1
0.014
3.33
18.72
17.72
17.43
16.67
3 layer
0.72
2.69
0.011
0.019
0.019
100.5
0.007
2.13
13.97
12.67
12.80
13.15
4 layer
0.83
3.15
0.014
0.016
0.019
111.3
0.012
2.02
21.54
30.23
17.89
21.45
Results Following Rao et al. (2019), report the top-1 accuracy for homology prediction. Our ERM baseline matches their transformer model performance. As shown in Table 2 both RGM and SRGM outperforms all the baselines (23.8% v.s. 22.3%). The difference between RGM and SRGM is relatively small due to inaccurate superfamily classifier. The top-1 and top- superfamily classification accuracy is around 33.5% and 51.0%. Nevertheless, SRGM can still give performance improvement because the gradient perturbation computed based on the ground truth superfamily label during training. This teacher forcing step helps SRGM to be robust superfamily variability despite the inaccurate superfamily classifier.
4.3 ABLATION STUDY OF SRGM
Updating o for Ly In section 3.2 we mentioned that the feature extractor ø is updated to optimize the scaffold (or superfamily) classification loss L. To study the effect of this design choice, we evaluate variant of SRGM called SRGM-detach, where classification loss, As shown in Table 2 (right), the performance of SRGM-detach is worse than SRGM across the four datasets. This because the scaffold classifier performs better in SRGM and the gradient 5(x) reflects the change of scaffold information more accurately.
updated to optimize the scaffold
Random perturbation In Table we report the performance SRGM under random perturbation on the QM9 dataset. Random perturbation performs significantly worse for most of the properties. This shows the importance of the scaffold classifier SRGM.
Model complexity To study how the model complexity of & affects the performance of SRGM, we train SRGM under different number of graph convolutional layers on QM9 dataset. As shown in Table SRGM performs the best when there are three graph convolutional layers, which is adopted in all experiments. In short, SRGM underfits the data when the model is too simple (layer=2) and overfits when the model is too complex (layer=4).
CONCLUSION
In this paper, propose regret minimization for generalization across structured biomedical domains such as molecular scaffolds or protein families. We seek to find a representation that enables the predictor to compete against oracle with hindsight access to unseen domains. Our method significantly outperforms all baselines on real-world biomedical tasks.
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Pazz Invariant risk minimization. arXiv preprint arXiv: 1907: 0028999,93 2019.

---

Under review as a conference paper at ICLR 2021
Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain gener- alization using metu--egularization. In Advances in Neural Information Processing Systems, pp. 998-1008, 2018.
Guy W Bemis and Mark A Murcko. The properties of known drugs. 1. molecular frameworks. Journal of medicinal chemistry. (9(15):2877-2893, 1996.
Shiyu Chang, Yang Zhang. Mo Yu, and Tommi S Jaakkola. Invariant rationalization. arXiv preprint arXiv:2003.09772. 2020.
Guangyong Chen, Pengfei Chen, Chang- Yu Hsieh, Chee-Kong Lee, Benhen Liao, Renjie Liao, Weiwen Liu, Jiezhong Qiu, Qiming Sun, Jie Tang, et al. Alchemy: A quantum chemistry dataset for benchmarking models. arXiv preprint
0942777777 2019.
Remi Tachet des Combes, Han Zhao, Yu-X Ni.aiang Wang, and Geoff Gordon. Domain adaptation with conditional distribution matching and generalized label shift. arXiv preprint arXiv:2003.044 04445, 2020.
Hanjun Dai, Bo Dai, and Le Song. Discriminative embeddings of latent variable models for structured data. In International Conference on Machine Learning, pp. 2702-2711, 2016.
Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. Advances in Neural Information Processing Systems, pp. 6447-6458, 2019.
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P Adams. Convolutional networks graphs for learning molecular fingerprints. In Advances in neural information processing systems, pp. 2224-2232, 2015.
Sara El-Gebali. Jaina Mistry, Alex Bateman, Sean R Eddy, Aurélien Luciani, Simon C Potter, Matloob Qureshi, Lorna J Richardson, Gustavo A Salazar, Alfredo Smart, Erik L. L Sonnhammer, Layla Hirsh, Lisanna Paladin, Damiano Piovesan, Silvio C E Tosatto, and Robert D Finn. The Pfam protein families database in 2019, Nucleic Acids Research, 472D1)) 2019. ISSN 0305- 1048. doi: 10. 1093//ar///yy9955.. D1/D4277 /
URLhttps: /ttp:://////amimmcccc/m oup. com http:///weemmc..o.ccccc///aaaaat
Evan N Feinberg, Debnil Sur, Zhenqin Wu, Brooke E Husic, Huanghao Mai, Yang Li, Saisai Sun, Jianyi Yang, Bharath Ramsundar, and Vijay S Pande. Potentialnet for molecular property prediction. ACS central science, 4(11):1520-1500, 2018.
Evan Feinberg, Robert Sheridan, Elizabeth Joshi, Vijay Pande, and Alan Cheng. Step change improvement in admet prediction with potentialnet deep fcaturization. arXiv preprint 193iv:1903.11789, 2019.
Naomi K Fox, Steven E Brenner, and John-Marc Chandonia. Scope: Structural classification of mooteins-ettended integrating scop and astral data and classification of new structures. Nucleic acids research, D2301): 2013.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The Journal of Machine Learning Research, 17()):099-20300 2016.
Muhammad Ghifary, W Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE international conference on computer vision, pp. 2551-2559, 2015.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. arXiv preprint urXiv: 2017.
Jie Hou. Badri Adhikari, and Jianlin Cheng. Deepsf: deep convolutional neural network for mapping protein sequences to folds. Bioinformatics, 34(8): 2018.
10

---

Under review as a conference paper at ICLR 2021
David Krueger, Ethan Caballero, Joemn-Henrle Jacobsen, Amy Zhang, Jonathan Binas, Remi Le Priol, and Aaron Courville. -iit----isttrbutiinn generalization via risk extrapolation (rex). arXiv preprint 27XY::2003.00688, 2020.
Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5542-5550, 2017.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Learning to generalize: Meta- learning for domain generalization. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018a.
Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy Hospedales. Episodic training for domain generalization. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1446-1455, 2019a.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex Kot. Domain generalization with adver- sarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5400-5409, 2018b.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings the European Conference on Computer Vision (ECCV), pp. 624-639, 2018c.
Yiying Li, Yongxin Yang, Wei Zhou, and Timothy M Hospedales. Feature-critic networks for heterogeneous domain generalization. arXiv preprint 1rXiv:1901.114 2019b.
Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Michael Jordan. Conditional adversarial domain adaptation. In Advances in Neural Information Processing Systems, pp. 1640-1650, 2018.
Saeid Motian, Marco Piccirilli, Donald A Adjeroh, and Gianfranco Doretto. Unified deep supervised domain adaptation and generalization. In Proceedings of the IEEE International Conference on Computer Vision, pp. 5715 5725, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pp. 10-18, 2013.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177,001 2014.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with tape. In Advances in Neural Information Processing Systems, 2019.
Gabriel Rocklin, Tamuka M Chidyausiku, Inna Goreshnik, Alex Ford, Scott Houliston, Alexander Lemak, Lauren Carter, Rashmi Ravichandran, Vikram K Mulligan, Aaron Chevalier, et al. Global analysis of protein folding using massively parallel design, synthesis, and testing. Science, 357 (6347):168-177)) 2017.
Shiv Shankar, Vihari Piratla, Soumen Chakrabarti, Siddhartha Chaudhuri, Preethi Jyothi, and Sunita Sarawagi. Generalizing across domains via cross-gradient training. arXiv preprint arXiv: 10745, 2018.
Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 7167-7176, 2017.
Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing to unseen domains via adversarial data augmentation. In Advances in Neural Information Processing Systems, pp. 5334-5344, 2018.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9()):133-330 2018
11

---

Under review as a conference paper at ICLR 2021
Kevin Yang. Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):
3370-3388, 2019.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costeira, and Geoffrey Gordon. Adversarial multiple source domain adaptation. Advances in neural information processing systems, pp. 8559-8570, 2018.
Han Zhao, Remi Tachet des Combes, Kun Zhang, and Geoffrey J Gordon. On learning invariant representation for domain adaptation. arXiv preprint 1901, 09453, 2019.
12

---

Under review as a conference paper at ICLR 2021
A TECHNICAL DETAILS
A.1 PROOF OF PROPOSITION
Note that (f Φ) is defined set of fixed examples in Eg. Since fe arg arg (f o0) and fer f- are in the same parametric family F, we have (*) ( (S. 00) 2 0.
A.2 PROOF OF PROPOSITION
Proof. Consider any representation o² PRGM. When there only two environments {E,, Eg J. we have F' 2=22(0)) = Fi( and (L) = F2(Ø") by definition. Thus the RGM constraint implies
F2 = F-1 (F('")
Fi(6") F-2(0") S F2(6")
Therefore Fi()") ). Since the loss function is non-negative and F is bounded and closed, Fi(6") # Thus, 0 (0)) = Fi(Ø') + Now consider any f E (ef"""") By definition,
Ve (f o Φ²) min
( 6*) hEF
By summing the above inequality over all environments, we have
her (s o o³) Σ L o²) min
Σ (h)
her
Since E. C L(ʃ x²), ). the above inequality implies
f(f x²) S min L(h Φ²) LRGM LiRM
KEF
Thus, fo o² an optimal solution under IRM and o² E PIRM-
A.3 PROOF OF PROPOSITION 3
Proof. Let recall our assumption of the data generation process:
p(z, y, e) = p(e)p((le)pp(lc, e); p(ulz,e) = p(ulz,
Under this assumption, we can rephrase the IRM objective as
min Ee Cly, (((()))
(17)
f
s.t.
En By (((())) I min E ((y. Ve
(18)
Given any label-pres representation ((x), its ERM optimal predictor is ("((()) = arg min E(,, ((((x))
(19)
To see that f* is ERM optimal, consider
min E, En E E Clv, f(o(=)))
2
E, Enje min Cl f(((x))
(20)
Eg Enje min E a (g, f((()))
(21)
f
E, Eg (v. (o(z)))
(22)
where 21 holds because d(x) is label- bebel-preserving. Note that f satisfies the IRM constraint because it is simultaneously optimal across all environments:
Ye min E., E.j. (l f((())))
2 Exje min e fe ((()))))
(23)
Je
f.
= Enje min E. | f(.
(24)
/
Eyye ((n)
( (g. (((x))
(25)
Moreover, if 0 PIRM is optimal representation, f di is optimal solution of IRM.
13

---

Under review as a conference paper at ICLR 2021
Table 5: Dataset statistics
QM9
HIV
Tox21
BBBP
Homology
Training
67K
27K
7.6K
1275
12.3K
Validation
36K
7.7K
776
519
736
Testing
30K
6.3K
483
248
718
A.4 STRUCTURED RGM UPDATE RULE
Since fe and optimizes L(fe o, Eg) in different directions, also introduce gradient reversal layer between @ and fe- The SRGM update rule is the following:
+ Φ v Nollf n,, Vo, p E.E.
Σ
(oR
vE(0,5)
f V
9 -
+ fe V o)
-- o (0 ()) Ve
f- f-e- NV- ( Ve
B EXPERIMENTAL DETAILS
B.1 MOLECULAR PROPERTY PREDICTION
Data The four property prediction datasets are provided in the supplementary material, along with the
splits. The size of each training environment, validation and test set are listed Table 5 The QM9, HIV, Tox21 and BBBP dataset are downloaded from Wu et al. (2018).
Model Hyperparameters For the feature extractor é, we adopt the GCN implementation from Yang al. 2019). We use their default hyperparameters across all the datasets and baselines. Specifically, the GCN contains three convolution layers with hidden dimension 300. The predictor is two-layer MLP with hidden dimenion 300 and ReLU activation. The model is trained with Adam optimizer for 30 epochs with batch size 50 and learning rate linearly annealed from 10-3 10 4 For RGM, explore À E {0.01, 0.1} for each dataset. For SRGM, we explore >g E {0.1. 1} for the classification datasets while X E {0.01, 0.1) for the QM9 dataset as X, causes gradient explosion.
The scaffold classifier trained by negative sampling since scaffolds are structured objects. Specifically, for each molecule x; in minibatch B, the negative samples the scaffolds {sk] of other molecules in the minibatch. The probability that xy is mapped to its correct
Scaffold Classification
scaffold s is then defined as
p(s: Iz. B) = exp(e(o(x:) g(((s)))
(26)
EKEB exp(g(o(x.)) o(((s=))) (
The scaffold classification loss is - Σ, logp(s; 34 B) for minibatch B. We choose the classifier g to be two-layer MLP with hidden dimension 300 and ReLU activation.
B.2 PROTEIN MODELING
Data The homology and stability dataset are downloaded from Rao et al. (2019). The size of each training environment, validation and test set are listed Table
Model hyperparameters For both tasks, our protein encoder is pre-trained BERT (Rao al. 2019). The predictor is a linear layer and the superfamilyytopology classifier is two-layer MLP whose hidden layer dimension is 768. The model is fine-tuned with an Adam optimizer with learning rate 10- 4 and linear warm up schedule. The batch size 16 and 20 for the homology and stability task. For RGM and SRGM, explore À E {0.01, 0.1} and Ag E {0.1, respectively.
14

---

Under review as a conference paper at ICLR 2021
Table 6: Mean absolute error (MAE) on the QM9 dataset under molecular size split. Models are trained on molecules with no more than 7 atoms and tested on molecules with 9 atoms. Due to space limit, we only show standard deviation for the top three methods subscripts.
Categorical methods
Structured methods
Property ERM
DANN
CDAN
IRM
MLDG
RGM
CrossGrad
SRGM
mu
0.658
0.655
0.655
0.690
0.654
0.656, 004)
0.664.000)
0.666
alpha
13.08
13.17
13.19
13.16
14.13
12.99.025
12.79 379)
(1.54
HOMO
0.008
0.008
0.008
0.009
0.008
0.008
0.008
0.009.000
LUMO
0.011
0.011
0.011
0.011
0.011
0.010
0.01 (000)
3.01
gap
0.014
0.013
0.014
0.015
0.014
0.012
0.014
0.016
352.8
355.7
357.3
368.6
381.2
328.4
351.7
279.9
ZPVE
0.025
0.024
0.025
0.025
0,026
0.022 (000)
0,024(001)
0.019
Cv
5.336
5.351
5.369
5.327
5,756
4.860,
5.235
3,909
U0
67.18
67.57
67.34
67.67
71.83
60.25(2.62)
(3.82
51.32
U
66.67
67.00
67.24
68.55
71.60
28.74
64.30
51.554/.099)
H
67.00
67.39
67.27
68.23
71.47
59.72
(4.39(2.19)
50.17 (2.56)
65.92
65.95
66.02
68.16
70.70
29.40
64.63
(6.23
Table 7: SRGM performance under different molecular size split.
Train
mu
alpha
homo
lumo
gap
R2
zpve
Cv
UO
H
6 atoms
0.92
21.5
0.010
0.016
0,018
521.6
0.035
7.73
100.1
99.4
102.0
100.1
7 atoms
0.67
11.5
0.009
0.013
0.016
279.9
0.019
3.91
51.3
51.5
50.17
51.23
8 atoms
0.69
4.00
0.007
0.009
0.011
119.0
0.009
1.59
20.2
20.1
19.7
20.5
B.3
ADDITIONAL EXPERIMENTS
For the quantum chemistry dataset (QM9), prior work (Chen al. 2019) has proposed measure domain generalization via molecular size split. To show that our method also works well under this evaluation setup, we split the dataset based on the number of heavy atoms. The training set contains molecules with more than 7 heavy atoms. The validation and test set consist of molecules with 8 and heavy atoms respectively. This setup is much harder than random split as it requires models to extrapolate to new chemical space.
Our results on the QM9 dataset are shown in Table 6 Among the categorical methods, RGM outperforms all the baselines (except for property mu), with significant improvement on six properties (R2, Cv. UO U, H. G) with 7-10% relative error reduction. SRGM outperforms all the baselines on eight properties (out of 12). While CrossGrad utilizes scaffold information, its performance is worse than RGM in general. Compared to RGM, SRGM shows significant error reduction (10-20%) on seven properties (alpha, R2, Cv. UO, U, H. G). This validates the ad advantage of exploiting structures of the environments (scaffolds).
Model parformance the severity of domain shift
Enor Ratio witt. ERM
0.9
RGM
CrossGrad
SRGM
We further conduct additional experiments to study the performance of RGM/SRGM with respect the severity domain shift. Fixing the test set to molecules with 9 atoms, we construct three progres- sively harder training sets: molecules with no more than 8, and 6 atoms. We report the MAE ratio (averaged over 12 properties) between SRGM//GMCCooSCraa and ERM. As shown in Figure CrossGrad and RGM across different setups.
Training set number atoms
Figure QM9 ablation study
SRGM consistently outperforms
15


Reviews:
[{"review_id": "zQQzNXoJtdE", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In bioinformatics and biochemistry, we may need to generalize the model beyond the training distribution. The authors propose a new method, the regret minimization (RGM) algorithm, to handle such a problem. The method is inspired by invariant risk minimization (IRM). It regulates the empirical loss with the regret regularizer. Such a regularizer encourages the model to generalize to a new unseen environment. The authors further extend the RGM to the structured inputs, leading to structured RGM (SRGM), which can be trained using gradient perturbation. Overall, the manuscript is interesting, which may be suitable for the general audience of ICLR. \n\nPro:\n1. Interesting problem and a new self-contained solution.\n2. Promising performance across different tasks from different domains.\n3. The manuscript is easy to follow although I believe it could be improved further.\n\nConcerns:\n1. The authors should further proofread the manuscript. I guess on page 5, the authors wanted to say that \"the model would not see any change from one example to another as scaffold variation\", right? Some symbols in the paper are not defined as well. For example, in Figure 2, I did not find the definition for $B_i$, $B_j$, and $B_e$. I am willing to increase the score if the manuscript is refined further.\n2. On page 6, my feeling is that the model is not very difficult to train. Are there any tricks to train such a model? Is it similar to adversarial training?\n3. In Table 2, why we have both RGM and SRGM? What's the difference between the two methods within this table? If I understand correctly, all the inputs are structured, right?\n4. In Figure 4, wow much would SRGM decrease for the different number of atoms? I mean the absolute value, not the relative value to ERM.\n5. For the HOMO, the authors include data from different protein families. How can the method generalize across different families?\n6. Regarding the feature extractor $\\phi$, how would the model complexity affect the performance of the proposed method? Any suggestions on the selection of the parametric family?", "extra_scores": {}, "extra_text": {}}, {"review_id": "ogTBvd4LwuB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nSummary: The paper proposed a new  regret minimization (RGM) algorithm on structured environment. It is based on one existing SOTA work --- invariant risk minimization (IRM), which assumes the predictor is simultaneously optimal across various environments. To relax the restriction, the authors propose to introduce held-out environment to perform domain perturbation. The experiments are conducted on real-world biomedical applications including molecule property prediction and protein homology and stability prediction. The proposed method outperform baseline methods.\n\n\n\nStrength\n\nThe paper proposed a novel Regret minimization (RGM) algorithm that extends IRM by relaxing the restrictive assumption on IRM, reformulate the problem as a joint optimization problem including regret minimization. The usage of held-out environment to encourage domain generalization looks very interesting. \n\nThe paper has solid theoretical analysis, and also demonstrate that the predictor obtained by RGM satisfies IRM constraints in the general case.  Empirical studies are convincing. Experiments on two real-world biomedical applications are of high value. The selected baselines are thorough and all from very recent paper. The performance gain looks significant, consistently outperforming baseline methods. \n\n\nWeakness:\n\n1. The author argues that their methods works in both standard and structured setting in Introduction Section, but only structured prediction is done. \n\n2. MLDG is very relevant to the proposed method. It would be great if the author can discuss its difference with RGM in details, instead of mentioning it only in Related Work Section. \n\n3. The presentation of Section 3 could be improved - notations could be better explained.\n\n4. In theoretical analysis, can authors add the discussion about \\phi_*? Now all the discussion are based on \\phi \\in \\Phi_{IRM/RGM}. \n\n5. I am not sure if the assumptions made in Section 3 are restrictive. For example, “new environments we may encounter at test time exhibit similar variability as the training environments”. Can you add some reference? \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "jG007KzK5mq", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\nThe authors proposed regret minimization (RGM) and structured RGM (SRGM) algorithms for generalization across biomedical domains. They quantified generalization to held-out environments by using two auxiliary predictors trained with and without each environment. Furthermore, for structured environments, they simultaneously trained an auxiliary scaffold classifier and used it to generate perturbations highlighting the environmental variation. The authors evaluated the proposed method on several tasks from molecular property prediction, protein homology, and stability prediction. They showed significant performance improvement over current state-of-the-art algorithms.\n\n##########################################################################\n\nMajor comments:\nWhile the paper has its own merits, unfortunately, it has several issues that need to be addressed.\n-\tMy main concern is that I am not sure the experiment settings are truly realistic. I agree with the authors’ motivation that generalizing beyond training domains and environments holds great importance. However, in my view, the current experiment setting does not quite represent real-world applications. For molecular property prediction, the authors used the number of atoms or molecular weights to determine the scaffold of data. However, in real-world settings, I do not think it is common to have training data with small molecules and test data with larger molecules. Wouldn’t it be more realistic to determine scaffolds based on sharing a subgraph of a molecular graph as described in Figure 1? \n-\tIn the case of stability prediction, the train/validation/test splits are not based on the topology. As stated by the authors, the test split contains Hamming distance-1 neighbors of around those from the training set. I am not sure whether the dataset is appropriate to be used to evaluate the proposed method for generalization beyond training domains.\n-\tThe key component of SRGM compared to RGM would be the adoption of a scaffold classifier. Then, I think the performance of SRGM should be inevitably dependent on the performance of the scaffold classifier. Can you provide the classification performance of the scaffold classifier? Assuming that the top-1 accuracy for superfamily classification is similar to that in TAPE (which is only about 40%), can you explain how can SRGM provide performance improvement with such an inaccurate classifier? Furthermore, if most scaffolds contain only few examples, wouldn’t the training of the scaffold classifier be unstable, thus, lower the performance of SRGM as well?\n-\tI think the ablation studies with SRGM-detach showed some related results, can you show how random perturbations instead of using the scaffold classifier would affect the performance of the proposed method? \n\n##########################################################################\nMinor comments:\n-\tIn protein modeling experiments, how did the authors generate sequence-length invariant protein embeddings from the BERT representations? Did you adopt an additional layer to compute attention-weighted mean protein embeddings as used in Rao et al.?\n-\tI understand that the authors omitted detailed explanations of each task and data due to space limitations. However, to be self-contained, I would like to recommend to include more information in the appendix. For example, the number of labels for the HOMO task is not stated in the paper.\n-\tCan you provide a more concise definition of standard and structured environments? \n-\tIt is obvious that SRGM stands for structured RGM, but anyhow it appeared in the manuscript without proper abbreviation explanation.\n\n##########################################################################\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "e1gN9nglljD", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThis work is built on top of IRM, aiming to solve an issue in the biomedical domain where there are very few instances in a given environment. The authors first proposed a new scheme, RGM, to apply tighter constraints on the generalization. And then they described a gradient perturbation scheme in the representation space to corrupt the environment information carried in the representation. The authors claimed improved performance over existing methods on several empirical experiments using biomedical benchmarks. \n\nStrength: \n1) RGM and a perturbation scheme for structured environments with few instances. Some theoretical analyses showing RGM imposes stricter constraints than IRM. \n2) Pretty comprehensive evaluation on several benchmark datasets to show the model performance/generalization  \n3) The flow of the paper is clear, not hard to follow\n\nWeakness \n1) The contribution and significance are not clear to me. Both RGM and perturbation (esp. compared to CrossGrad) do not seem very novel or significant. Maybe the authors could elaborate on this more.   \n2) While the authors described the hyperparameter search in the appendix, it would be nice to know more about the auxiliary models and losses. For example, how well does the scaffold classifier perform in the experiments as well as in the ablation studies (i.e. fig 5)? \n3) The data split based on (heavy) atom numbers, scaffold molecule weights and #instances in a protein superfamily look a bit weird. Does this truly probe the model generalization? In other words, does a big molecule weight difference necessarily correspond to two very distinct scaffolds?\n4) Figures are not in the correct order (fig 2). Understandable this is likely to accommodate the page limits but is causing some confusion. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper has some interesting points in extending IRM to regret minimization, and extending to structured environments.  I can see the writing has been improved in the revision.  The main criticism arises from the experiment, which can be improved in several aspects.   The reviews have been quite detailed and helpful.
Decision: Reject