Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations
Authors: Aahlad Manas Puli
Lily H Zhang
Eric Karl Oermann
Rajesh Ranganath
Venue: ICLR 2022 Poster (2022)
Keywords: spurious correlations
out of distribution generalization
ml for health
representation learning

Abstract: In many prediction problems, spurious correlations are induced by a changing relationship between the label and a nuisance variable that is also correlated with the covariates. For example, in classifying animals in natural images, the background, which is a nuisance, can predict the type of animal. This nuisance-label relationship does not always hold, and the performance of a model trained under one such relationship may be poor on data with a different nuisance-label relationship. To build predictive models that perform well regardless of the nuisance-label relationship, we develop Nuisance-Randomized Distillation (NURD). We introduce the nuisance-randomized distribution, a distribution where the nuisance and the label are independent. Under this distribution, we define the set of representations such that conditioning on any member, the nuisance and the label remain independent. We prove that the representations in this set always perform better than chance, while representations outside of this set may not. NURD finds a representation from this set that is most informative of the label under the nuisance-randomized distribution, and we prove that this representation achieves the highest performance regardless of the nuisance-label relationship. We evaluate NURD on several tasks including chest X-ray classification where, using non-lung patches as the nuisance, NURD produces models that predict pneumonia under strong spurious correlations.

Published as a conference paper at ICLR 2022
OUT-OF-DISTRIBUTION GENERALIZATION IN THE PRESENCE OF NUISANCE-INDUCED SPURIOUS COR- RELATIONS
Aahlad Puli New York University
Lily H. Zhang
Eric K. Oermann
Rajesh Ranganath
ABSTRACT
In many prediction problems, spurious correlations are induced by a changing re- lationship between the label and nuisance variable that is also correlated with the covariates. For example, in classifying animals in natural images, the background. which is nuisance, can predict the type of animal. This nuisance-label relation- ship does not always hold, and the performance of a model trained under one such relationship may be poor on data with different nuisance-labet relationship. To build predictive models that perform well regardless of the nuisance-label rela- tionship, we develop Muisance-Randominnn the nuisance-ranee randomized distribution, a distribution where the nuisance and the label are independent. Under this distribution, we define the set of representa- tions such that conditioning on any member, the nuisance and the label remain independent. We prove that the representations in this set always perform better than chance, while representations outside of this set may not. NURD finds rep- resentation from this set that is most informative of the label under the nuisance- randomized distribution, and we prove that this representation achieves the highest performance regardless of the nuisamce- several tasks including chest X-ray classification where, using non-lung patches as the nuisance, NURD produces models that predict pneumonia under strong spurious correlations.
Distillation (NURD). We introduce
relationship. We evaluate NURD on
INTRODUCTION
Spurious correlations are relationships between the label and the covariates that are prone to change between training and test distributions (Gulrajani and Lopez-Paz, 2020). Predictive models that ex- ploit spurious correlations can perform worse than even predicting without covariates on the test dis- tribution (Arjovsky al., 2019). Discovering spurious correlations requires more than the training distribution because any single distribution has fixed label-covariate relationship. Often, spurious correlations are discovered by noticing different relationships across multiple distributions between the label and muisance factors correlated with the covariates. We call these celsaacce-nnuuce spuri- ous correlations.
For example, in classifying cows vs. penguins, typical images have cows appear on grasslands and penguins appear near snow, their respective natural habitats (Arjovsky et al., 2019; Geirhos et al., 2020), but these animals can be photographed outside their habitats. classifying hair color from celebrity faces on CelebA (Liu et al., 2015), gender is correlated with the hair color. This relation- ship may not hold different countries (Sagawa al., 2020). In language, sentiment of movie review determines the types of words used in the review to convey attitudes and opinions. However, directors' names appear in the reviews and are correlated with positive sentiment in time periods where directors make movies that are well-liked (Wang and Culotta, 2020a). In X-ray classification, conditions like pneumonia are spuriously correlated with non- -on-physiological traits of X-ray images due to the association between the label and hospital X-ray collection protocols (Zech et al., 2018). Such factors are rarely recorded in datasets but produce subtle differences in X-ray images that convolutional networks easily learn (Badgeley et al., 2019).
We formalize quisance-induced spurious correlations in nuisance-varying family of distributions where any two distributions are different only due the differences in the nuisance-label rela- tionship. As the nuisance is informative of the label, predictive models exploit the nuisance-label relationship to achieve the best performance on any single member of the family. However, predic- tive models that perform best on one member can perform even worse than predicting without any covariates on another member, which may out-of-distribution (OOD). We develop NURD to use
*Corresponding email: aahl adenyu..deeee The code is available here.

---

Published as a conference paper at ICLR 2022
data collected under nuisance-labe relationship to build predictive models that perform well on other members of the family regardless of the nuisance-label relationship in that member.
we motivate and develop ideas that help guarantee performance on every member of the family. The first is the wwisance-randomiin independent of the label. An example is the distribution where cows and penguins have equal chances of appearing on backgrounds of grass or snow. The second is an uncorrelating represen- tation: a representation of the covariates such that under the nuisance-randomized distribution, the nuisance remains independent of the label after conditioning on the representation. 'he set of such representations is the uncorrelating set. We show that the nuisance-randomize label given an uncorrelating representation has performance guarantees: such conditionals perform as well or better than predicting without covariates on every member in the family while other con- ditionals may not. Within the uncorrelating set, we characterize one that is optimal on every member of the nuisance-var varinin family simultaneously. We then prove that the same optimal performance can be realized by uncorrelating representations that are most informative of the label under the nuisance-randomized distribution.
In section 2.
distribution: a distribution where the nuisance
conditional of the
Following the insights in section 2, we develop tion 3. NURD finds an uncorrelating representation that is maximally informative of the label un- der the randdmne-r distribution. NURD's first step, muisance-randomizaee nuisance-label dependence to produce nuisance-randomizee data. We provide two nuisance ran- domization methods based on generative models and reweighting. The second step, distillation, maximizes the information representation has with the label on the nuisance-randomizer data over Gaussians, labeling colored MNIST
Distillation (NuRD) in sec-
breaks the
the uncorrelating set. We evaluate NURD on class-conditions images (Arjovsky et al., 2019), distinguishing waterbirds from landbirds, and classifying chest X- rays. In the latter, using the non-lung patches as the nuisance, NURD produces models that predict pneumonia under strong spurious correlations.
2 DUUSANCE-RANDOMINN
AND UNCORRELATING SETS
We formalize nuisance- nuisance-induced spurious correlations via a family of data generating processes. Let y be the label, z be the nuisance, and x be the covariates (i.e. features). The family consists of distributions where the only difference in the members of the family comes from the difference in their nuisance-label relationships. Let D index a family of distributions J = {pD) J ni a member pp in the nuisance-varying family of distributions F takes the following form:
(1) where Pz(z yy is positive and bounded for any where p(y) > O and any z in the family's nuisance space Sy. This family is called the quisance-varying family. Due to changing nuisance- label relationships in this family, the conditional distribution of the label y given the covariates x in one member, e.g. the training distribution, can perform worse than predicting without covariates on another member of the family, e.g. a test distribution with a different nuisance-label relationship. We
Po z. x) p(y) (z y)p(x z,y),
define performance of a model ply x) on a distribution Pte as the negative expected KL-divergence from the true conditional (te x): Perf p.. (p(y x)) -Ep,
(x) KL [pte (y | x) || p(y x)]. Higher is better. This performance equals the expected log-likelihood up to a constant, Ope. =
H pio (y x), that only depends on the Pte: Perfp.. (f(y x)) (p..(x.)) logū(y x) Cpic- Consider the following example family JoER:
y~ (0,, z ~~y 0.5) x = Nxi~N(y - z, 1.5), x2 ~ Ny + 2, 0.5)]
(2)
Given training distribution Per Q1 and test distribution Pte = q-1, the conditional (er (y x)
performs even worse than predicting without covariates, Perf (p(y)) 2 Perfp.. (Pe x));
see appendix A.9 for the proof. The problem is that (y x) utilizes label-covariate relationships that do not hold when the nuisance-label relationships change. When the changing nuisance-label relationship makes the conditional PD(y x) of one member unsuitable for another P'D F, the family exhibits muisance-inducee spurious correlations.
Next, we identify conditional distribution with performance guarantees across all members of the family. We develop two concepts to guarantee performance on every member of the nuisance- varying family: the muisance-randomizee distribution and uncorrelating representations.
Definition 1. The muisance-randomiz distribution is Pa (x,
y, 2) p(x y, 2) (2)ply..
Different marginal distributions P1 (z) produce different distributions where the label and nuisance are independent. The results are insensitive to the choice long as (z) > 0 for any z E Sy. One distribution that satisfies this requirement is p. (z) (z). See lemma 2.
2

---

Published as a conference paper at ICLR 2022
In the cows vs. penguins example, p, is the distribution where either animal has an equal chance to appear on backgrounds of grass or snow. The motivation behind the nuisance-rante randomized distribu- tion is that when the nuisance is independent of the label, (noisy") functions of only the nuisance are not predictive of the label. If the covariates only consist of (noisy) functions of either the nuisance or the label but never mix of the two (an example of mixing X1 y + noise), then the conditional Pa (y x) does not vary with the parts of x that are (noisy) functions of just the nuisance. Thus, P. (y x) ignores the features which have changing relationships with the label.
How about nuisance-vary varying families where the covariates contain functions that mix the label and the nuisance? Equation (2) is one such family, where the covariates x and X2 are functions of both the label and the nuisance. In such nuisange-varying families, the conditional p, (y x) can use functions that mix the label and the nuisance even though the nuisance is not predictive of the label by itself. These mixed functions have relationships with the label which change across the family; for example in eq. (2), the coordinate X is correlated positively with the label under qo but negatively under 9-2- Then, under changes in the nuisance-label relationship, the conditional (y x) can perform worse than predicting without covariates because it utilizes a relationship, via these mixed features, that no longer holds. See appendix A.9 for details.
We address this performance degradation of Pı x) by introducing representations that help avoid reliance on functions that mix the label and the nuisance. We note that when the condi- tional Pa (y x) uses functions that mix the label and the nuisance, knowing the exact value of the nuisance should improve the prediction of the label, i.c. y z / x. Therefore, to avoid reliance on mixed functions, we define uncorrelating representations r(x), where the nuisance does not provide any extra information about the label given the representation:
Definition 2. An uncorrelating set of representations is (p.)) s.t. Vr E R(p.). y.lp, z rxx).
In the example in eq. (2), r(x) X1 + X2 is an uncorrelating representation because it is purely a function of the label and the noise. Conditional distributions p_ (y r(x)) for any uncorrelating r(x) only depend on properties that are shared across all distributions in the muisance-varying family. Specifically, for E Rp..)), the conditional distribution pa (y r(x)) uses p(r(x) y 2) and ply)
which are both shared across all members of the family F. For z' such that (z' r(x)) >
(y r(x))= = (y 2')),,²) = p. (y z²)p, (r()) yy,z) p(y)p(r(x) y,zz)
(3)
Pa (r(x) (z')
Ep(y)p(((x) z')
This fact helps characterize the performance of p_ (y rx))) on any member Pte E F:
Perf pea (p r(x))) Perf (p(y)) + (xr(y.))
KL [p(r(x) yyz) || ((()) ly [y, (4)
As KL-divergence is non-negative, for any uncorrelating representation r. the conditional (y r(x)) does least as well as predicting without covariates for all members Pte E F: Perf.. p.s (p_(x r(x))) 2 perfp)((((y)) See appendix A.3 for the formal derivation. In fact, we show in appendix A.5 that when the identity representation r(x) = x is uncorrelating, then p. x) is minimax optimal for family with sufficiently diverse nuisance-label relationships.
Equation (4) lower bounds the performance of (y r(x)) for any representation in the uncorrelat- ing set across all Pte E F. However, it does not specify which of these representations leads to the best performing conditional. For example, between two uncorrelating representations like the shape of the animal and whether the animal has horns, which predicts better? Next, we characterize un- correlating representations that are simultaneously optimal for all test distributions pte E F.
Optimal uncorrelating representations. As we focus on henceforth, by performance of r(x), we mean the performance of p | r(x)): Perf (r(x) Perff.c (p.(x r(x))). Consider two uncorrelating representations r, r2, where the pair (r, 12) is also uncorrelating. How can (x) dominate r(x) in performance across the quisance-varying fam- ily? Equation (3) shows that
conditionals,
( rx()), x p(y)p(r(x) | (2(x), y, = 2)p(r2(x) y, 2).
If (x() blocks the dependence between the label and r(x), i.c. r(x) y (2(x), z, then knowing does not change the performance when r2 is known, suggesting that blocking relates to perfor- mance. In theorem 1, we show that the maximally blocking uncorrelating representation is simulta- neously optimal: its performance is as good or better than every other uncorrelating representation on every distribution in the nuisance-varying family. We state the theorem first:
Noisy functions of variable are functions of that variable and exagenous noise.
3

---

Published as a conference paper at ICLR 2022
Theorem 1. Let R(p ) be maximally blocking: Vr R(p,), y.lp 22, (x). Then,
(Simultaneous optimality) Ypte E F, Vr E R(p.). Perf per (r" (x)) 2 Perfp. (r(x)).
2. (Information maximality) vr(x)= E R((x)), Ip (x;r((((x))) 2 Ip (y;r(x)))
3. (Information maximality implies simultaneous optimality) Vr' R(p.).
(y: Ip ry:r"(x))
Vpte E F, perf (x)(x) Perf (r))..
The proof is in appendix A.4. The first part of theorem simultaneous optimality, says that maxi- mally blocking uncorrelating representation r' dominates every other r R(p, in performance every test distribution in the family. In the cows vs. penguins example, the segmented foreground that contains only the animal is maximally blocking representation because the animal blocks the dependence between the label and any other semantic feature of the animal.
The second and third parts of theorem 1 are useful for algorithm building. The second part proves that a maximally blocking r" is also maximally informative of the label under P1. indicating how to find a simultaneously optimal uncorrelating representation. What about other ieformmttionmaxiine uncorrelating representations? The third part shows that if an uncorrelating representation r' has the same mutual information with the label (under p.) as the maximally blocking r", then r' achieves the same simultaneously optimal performance as r*. In the cows vs. penguins example, an example of maximally informative uncorrelating representation is the number legs of the animal because the rest of the body does not give more information about the label. The second and third parts of theo- rem together show that finding an uncorrelating representation that maximizes information under the nuidance-randomized distribution finds simultaneously optimal uncorrelating r(x).
DISTILLATION (NURD)
Theorem 1 says representation that maximizes information with the label under the nuisance- randomized distribution has the best performance within the uncorrelating set. We develop rep- resentation learning algorithm maximize the mutual information between the label and rep- resentation in the uncorrelating set under the nuisance-randomizze distribution. We call this al- gorithm Messance-Randmmize Distillation (NURD). NURD has two steps. The first step, called nuisance randomization, creates an estimate of the iissance-randmizze distribution. The second step, called distillation, finds representation in the uncorrelating set with the maximum information with the label under the estimate of the muisance-randomizz. distribution from step one.
Nuisance Randomization. We estimate the muussacceraadommine distribution with generative models or by reweighting existing data. Generative-NURD uses the fact that p(x y, z) is the same for each member of the nuisance-varying family F. With an estimate of this condi- tional denoted p(x y,z), generativeNuRRD" estimate of the muisance-randomizze distribution is 2 Per (z), x ~ p(( z). For high dimensional x, the estimate p(x x y, 2) can be constructed with deep generative models. Reweighting-Nut NURD importance weights the data from Per by pix)/perly 2), making it match the nuisance-randomizz distribution:
Per (y
z)
p(y)
P. y.) (2)(x y. 2) = ((y)).r(z)
p(x
(y,z)
Per (x,
Per (y
z)
Per (y z)
Reweighting-NURD uses a model trained on samples from Ptr estimate
p(y)
Pc-1y x]
Distillation. Distillation seeks to find the representation in the uncorrelating set that maximizes the information with the label under P.. the estimate of the imizing the information translates to maximizing likelihood because the entropy Hp, (y) is constant with respect to the representation Ty parameterized by 7:
distribution. Max-
Ip (x):)) Hp_ (y) = Ep, (x)) log (y (x)) = max Ep, (y.r, log (r, (x))..
Theorem requires the representations be in the uncorrelating set. When conditioning on rep- resentations in the uncorrelating set, the nuisance has zero mutual information with the label: I, (y:²) (x)) O. We operationalize this constraint by adding a conditional mutual information penalty to the maximum likelihood objective with tunable scalar parameter X
max E xx) p rx,xx)) 0,7
ALg (y:z) rx)x)))
(5)

---

Published as a conference paper at ICLR 2022
The objective in eq. (5) can have local optima when the representation is function of the nuisance and exogenous noise (noise that generates the covariates given the nuisance and the label). The in- tuition behind these local optima is that the value of introducing information that predicts the label does not exceed the cost of the introduced conditional dependence. Appendix A.6 gives a formal dis- cussion and example with such local optima. Annealing X, which controls the cost of conditional dependence, can mitigate the local optima issue at the cost of setting annealing schedules.
Instead, we restrict the distillation step to search over representations r-, rx,(x) that are also marginally independent of the nuisance z under Pa. i.e. -_l,,,, "> (x). This additional independence removes representations that depend on the nuisance but are not predictive of the label; in turn, this removes local optima that correspond to functions of the nuisance and exogenous noise. In the cows vs. pen- guins example, representations that are functions of the background only, like the presence of snow, are uncorrelating but do not satisfy the marginal independence. Together, the conditional indepen- dence y -Lpp z rry (x) and marginal independence 2_Ip,, rx((x) hold if and only if the representation and the label are jointly independent of the nuisance: (y, rry (V)r)(()))LL (X) z. Using mutual information to
penalize joint dependence (instead of the penalty in eq. (5)), the distillation step in NURD is
max F. z,,) log Po (y rry (x)) - AI ([x,rxx))): (x)] 0,7
(6)
We show in lemma 4 that within the set of representations that satisfy joint independence, NURD learns a representation that is simultancously optimal in performance on all members of the drisance-varying family. To learn representations using gradients. the mutual information needs to be estimated in a way that is amenable to gradient optimization. To achieve this, we estimate the mutual information in NURD via the classification- bbased density-ratio estimation trick (Sugiyama et al., 2012). We use a critic model Pe to estimate said density ratio. We describe this technique in appendix A. for completeness. We implement the distillation step as bi-level optimization where the outer loop optimizes the predictive model pe (y | rr (x)) and the inner loop optimizes the critic model P& which helps estimate the mutual information.
Algorithm. We give the full algorithm boxes for both reweighting-NuRG and generative-NRRD in appendix 1. In Seeeighting-NRRDD to avoid poor weight estimation due to models memorizing The setup of uuisance-induced spurious correlations in eq. (1) assumes p(y) is fixed across distributions within the nuisance-varying family F. This condition can be relaxed when (y) is known; see appendix A.
the training data, we use fooss-fitting see algorithm
4 RELATED WORK
In table 1, we summarize key differences between NURD and the related work: invariant learning (Arjovsky et al., 2019; Krueger et al., 2020), distribution matching (Mahajan et al., 2020; Guo et al., 2021), shift-stable prediction (Subbaswamy et al., 2019a), group-DRO (Sagawa et al., 2019), and causal regularization (Veitch et al., 2021; Makar et al., 2021). We detail the differences here.
Nuisance versus Environment. In general, environment is distribution with specific spuri- ous correlation (Sagawa et al., 2019). When the training and test distributions are members of the same quisance-varyine family, environments denote specific nuisance-label relationships. In con- trast, nuisances are variables whose changing relationship with the label induces spurious correla- tions. While obtaining data from diverse environments requires data collection from sufficiently dif- ferent sources, one can specify nuisances from a single source of data via domain knowledge.
-imaiiinnvvriant learning, and subgroup robustness We briefly men- tion existing methods that aim to generalize unseen test data and focus on how these methods can suffer in the presence of nuisance-induced spurious correlations; for a more detailed presentation, see appendix A.2. Domain generalization and Nomain-invariann learning methods assume the train- ing data consists of multiple sufficiently different environments to generalize to unseen test data that related to the given environments or subgroups (Li et al., 2018a; Arjovsky et al., 2019; Akuzawa et al., 2019; Mahajan et al., 2020; Krueger et al., 2020; Bellot and van der Schaar, 2020; Li et al., 2018b; Goel et al., 2020; Wald et al., 2021; Ganin et al., 2016; Xie et al., 2017: Zhang et al., 2018; Ghimire et al., 2020; Adeli et al., 2021: Zhang et al., 2021). Due to its focus nuisances, NURD works with data from a single environment. Taking a distributional robustness (Duchi et al., 2021) approach, Sagawa et al. (2019) used group-DRe DRO build models that perform well on every one of finite set of known subgroups. Other work also aims to minimize worst subgroup error with a finite number of fixed but unknown subgroups (Lahoti et al., 2020; Martinez et al., 2021); as subgroups
Domain generalization,
5

---

Published as a conference paper at ICLR 2022
Table NURD vs. methods that use nuisances or environments. In this work, the training data comes from single member of the family F, i.e. single environment. For methods that require multiple environments, values of the nuisance can be treated as environment labels. Unlike existing methods, NURD works with high-dimensional nuisances without requiring them at test time.
Invariant
Dist. match
Shift-stable
Group-DRO
Causal reg.
NURD
High-dim z No test-time z
x
x
are unknown, they only find an approximate minimizer of the worst subgroup error in general even with infinite data. While these methods (Lahoti et al., 2020; Martinez et al., 2021) were developed to enforce faimess with respect to a sensitive attribute, they can be applied to OOD generalization with the nuisance treated as the sensitive attribute; see (Creager et al., 2021). Given the nuisance, existence of a finite number of subgroups maps to additional discreteness assumption on the nuisance variable; in contrast, NURD works with general nuisances. Given a high dimensional z, as in our experiments, defining groups based the value of the nuisance like in (Sagawa et al., 2019) typically results in groups with at most one sample; with the resulting groups, methods that minimize worst subgroup error will encourage memorizing the training data.
Nuisance as the environment label for domain generalization. Domain generalization methods are inapplicable when the training data consists only of a single environment. In this work, the training data comes from only one member of the nuisance-varying family, i.e. from single envi- ronment. What if one treats groups defined by nuisance values as environments? Using the nuisance as the environment label can produce non-overlapping supports (over the covariates) between en- vironments. In Colored-N -NISTT for example, splitting based on color produces an environment of green images and an environment of red images. When the covariates do not overlap between envi- ronments, methods such as Arjovsky et al. (2019); Krueger et al. (2020) will not produce invariant representations because the model can segment out the covariate space and learn separate func- tions for each environment. Methods based on conditional distribution matching (Mahajan et al., 2020; Guo et al., 2021) build representations that are conditionally independent of the environment variable given the label. When the training data is split into groups based on the nuisance value, representations built by these methods are independent of the nuisance given the label. However, splitting
digh-ddimensooal nuisance like image patches tends to yield many groups with only single image. Matching distributions of representations for the same label across all environments is not possible when some environments only have one label.
Causal learning and shift-stable prediction. Anticausal learning (Schölkopf et al., 2012) sumes a causal generative process for a class of distributions like the nuisance-varying family in eq. (1). In such interpretation, the label y and nuisance z cause the image x, and under indepen- dence of cause and mechanism, p(x y, z) is fixed - in other words, independent - regardless of the distribution PD z). A closely related idea to NURD is that of Shift-Stable Prediction (Sub- baswamy et al., 201 1a;b). Subbaswamy et al. (2019a) perform graph surgery to learn p. (y x, z) assuming access to z during test time. Shift-stable models are not applicable without nuisances at test time while NURD only requires nuisances during training. However, if the nuisance is available at test time, the combined covariate-set [x, z] 1) uncorrelating because y yll z [x, 2] and 2) max- imally blocking because r([x, 2])_Ly| z, [x, z). and theorem says p_ (y [xx, is optimal.
Two concurrent works also build models using the idea of nuisance variable. Makar et al. (2021) assume that there cxists a stochastic function of y but not z, called x", such that y x x² z; they use a marginal independence penalty r(x) z. Veitch et al. (2021) use counterfactual invariance p1
to derive a conditional independence penalty r(x)_llp., p., z y. The theory in these works requires the nuisance to be discrete, and their algorithms require the nuisance to be both discrete and low- cardinality; NURD and its theory work with general high-dimensional nuisances. Counterfactual invariance promises that a representation will not vary with the nuisance but it does not produce optimal models in general because it rejects models that depend on functions of the nuisance. On the other hand, the uncorrelating property allows using functions of only the nuisance that are in the covariates to extract further information about the label from rest of the covariates; this leads to better performance in some nuisance- nuisance-varying families, as we show using the theory of minimal sufficient statistics (Lehmann and Casella, 2006) in appendix A.7.

---

Published as a conference paper at ICLR 2022
5 EXPERIMENTS
We evaluate the implementations of NURD on class-conditional Gaussians, Colored-MNISTT (Ar- jovsky et al., 2019), Waterbirds (Sagawa et al., 2019), and chest X-rays (Irvin et al., 2019; Johnson et al., 2019). See appendix B for implementation details and further evaluations of NURD.
Model selection, baselines, and metrics. Models in both steps of NURD are selected using held- out subsets of the training data. We split the training data into training and validation datasets with an 80 - 20 split. For raaaaaccerandoiiziioonn this selection uses standard measures of held-out per- formance. Selection in the distillation step picks models that give the best value the distillation objective held out subset of the dassunce-randomizee data from NURD's first step.
We compare against Empirical Risk minimization (ERM) because, as discussed in section 4, exist- ing methods that aim to generalize under spurious correlations require assumptions, such as access to multiple environments or discrete nuisance of small cardinality, that do not hold in the experiments. When possible, report the oracle accuracy or build gold standard models using data that does not have a nuisance-label relationship to exploit. For every method, we report the average accuracy and standard error across 10 runs each with different random seed. We report the accuracy of each model for each experiment on the test data (Pte) and on heldout subsets of the original training data (Ptr) and the estimate of the
distribution (p.). For all experiments, we use A and one two epochs of critic model updates for every predictive model update.
5.1 CLASS-CONDITIONNL GAUSSIANS
We generate data as follows: with (0.55) as the uniform Bernoulli distribution, qa (y, z, x) is
~ B(055) 2 ~ (aa(2 - 1), x x= [x1 ~ .Mly z, 9), ~ Ny + z, 0.01)].
(7)
The training and test sets consist of 10000 samples from Ptr 9o.5 and 2000 samples from Pte Q-0.9 respectively. All models in both NURD methods are parameterized with neural networks.
Table 2: Accuracy of NURD versus ERM on class conditional Gaussians.
Method
Heldout ptr
Heldout p_
Pte
ERM generative-NuRD
84 + 0%
39 + 0%
71 + 0%
+ 0%
58 0%
reweighting-NuRDD
71 + 1%
66 + 0%
58 + 0%
Results. Table reports results. The test accuracy of predicting with the optimal linear uncor- relating representation r (x) Xi + X2, is 62%; appendix B gives the optimality proof. Both generative- and reweighting-NURD achieve close to this accuracy.
5.2 COLORED-MNIST
We construct a colored-MNIST dataset (Arjovsky et al., 2019; Gulrajani and Lopez-Paz, 2020) with images of (Is and 1s. In this dataset, the values in each channel for every pixel are either 0 1. We construct two environments and use one as the training distribution and the other as the test. In training, 90% of red images have label 0; 90% of green images have label 1. In test, the relationship is flipped: 90% of the Os are green, and 90% of the 1s are red. In both training and test, the digit determines the label only 75% of the time, meaning that exploiting the nuisance- labell relationship produces better training accuracies. The training and test data consist of 4851 and 4945 samples respectively. We run NURD with the most intense pixel as the nuisance.
Table Accuracy of NURD versus ERM on ColorTd-MNIST: The oracle accuracy 75%.
Method
Heldout Ptr
Heldout P
Pte
ERM
90 + 0%
10 0%
generative-NuRD
73 + 1%
80 = 1%
68 - 2%
NewDighting-NURD
75 = 0%
74 = 0%
75 = 0%
7

---

Published as a conference paper at ICLR 2022
Results. See table 3 for the results. ERM learns to color, evidenced by the fact that it achieves a test accuracy of only 10%. The oracle accuracy of 75% is the highest achievable by models that do not use color because the digit only predicts the label with 75% accuracy. While Sinerative-NuRD has an average accuracy close to the oracle, SeweighiingNNuRe
matches the oracle at 75%.
5.3 LEARNING TO CLASSIFY WATERBIRDS AND LANDBIRDS
Sagawa al. (2019) consider the task of detecting the type of bird (water or land) from images where the background is nuisance. Unlike Sagawa et al. (2019), we do not assume access to validation and test sets with independence between the background and the label. So, we split their dataset differently to create our own training and test datasets with substantially different nuisance- relationships. The training data has 90% waterbirds on backgrounds with water and 90% landbirds on backgrounds with land. The test data has this relationship flipped. We use the original image size of 224 x 224 x 3. The training and test sets consist of 3510 and 400 samples respectively. We ensure that ply = 1)- 0.5 in training and test data. Thus, predicting the most frequent class achieves an accuracy of 0.5. Cropping out the whole background requires manual segmentation. Instead, we use the pixels outside the central patch of 196 196 pixels nuisance in NURD. This is a high-dimensional nuisance which impacts many existing methods negatively; see section 4. The covariates are the whole image: see fig. 1.
Method ERM
Per
P=
Pte
91 + 0%
66 2%
reweighting-NuBe NURD
85 + 1%
81 + 1%
83 + 2%
Table of results and figure showing an example of the nuisance for Waterbirds. On the left are the accuracies of NURD and ERM on Waterbirds. Gold standard accuracy is 90% (see the results paragraph below). The figure shows an image and the corresponding border nuisance. NURD has both images during training and only the left image at test time.
Figure 1:
Results. Figure reports results. We construct a gold standard model on data where waterbirds and landbirds have equal chances of appearing on backgrounds with water and land; this model achieves test accuracy of 90%. ERM uses the background to predict the label, as evidenced by its test accuracy of 66%. Reweighting-NuRe
uses the background patches to adjust for the spurious correlation to achieve an average accuracy close to the gold standard, 83%. We do not report -iveraativeNuRDD performance as training on the generated images resulted in classifiers that predict as poorly as chance on real images. This may be due to the small training dataset.
5.4 LEARNING TO LABEL. PNEUMONIA FROM X-RAYS
In many problems such classifying cows versus penguins in natural images, the background, which is a nuisance, predicts the label. Medical imaging datasets have similar property, where fac- tors like the device used to take the measurement are predictive of the label but also leave signature on the whole image. Here, we construct a dataset by mixing two chest x-ray datasets, CheXpert and MIMIC, that have different factors that affect the whole image, with or without pneumonia. The training data has 90% pneumonia images from MIMIC and 90% healthy images from CheXpert. The test data has the flipped relationship, with 90% of the pneumonia images from CheXpert and 90% of the healthy images from MIMIC. We resize the X-ray images to 32 x 32. Healthy cases are downsampled to make sure that the training and test sets, healthy and pneumonia cases are equally probable. Thus, predicting the most frequent class achieves an accuracy of 0.5. The training and test datasets consist of 12446 and 400 samples respectively. chest X-rays, image segmentation cannot remove all the nuisances because nuisances like scanners alter the entire image (Pooch et al., 2019; Zech et al., 2018; Badgeley et al., 2019). However, non-lung patches, i.c. pixels outside the central patches which contain the lungs, are nuisance because they do not contain physiological signals of pneumonia. We use the non-lung patches (4-pixel border) as a nuisance in NURD. This is a high- dimensional nuisance which impacts existing methods negatively; see section 4. The covariates are the whole image; see fig. 2.
Results. Figure reports results. Building an oracle model in this experiment requires knowledge of all factors that correlate the label with all the parts of the X-ray. Such factors also exist within each hospital but are not recorded in MIMIC and CheXpert; for example, different departments in the same hospital can have different scanners which correlate the non-lung patches of the X-ray
8

---

Published as a conference paper at ICLR 2022
Method
Ptr
Pa
Pte
ERM
89 + 0%
37 1%
generativeNURRD
70 + 3%
90 + 2
41 - 2%
reweighting-NURD 75 + 1%
68 +
1%
61 1%
Figure 2: Table of results and figure showing an example of the nuisance for chest X-rays. The figure shows an example of a chest X-ray and the corresponding non-lung patches (right). NURD has both images during training and only the left image at test time.
with the label (Zech et al., 2018). ERM uses the nuisance to predict pneumonia, as evidenced by its test accuracy of 37%. Reweighting- NURDD uses the non-lung patches to adjust for the spurious correlation and achieves an accuracy of 61%, a large improvement over ERM.
Generative-NuRI also outperforms ERM's performance on average. Unlike reweighting-NURG which outperforms predicting without covariates, generative-NuRD performs similar to predicting without covariates on average. The few poor test accuracies may be due to two ways generative can be imperfect: 1) little reliance of x on z with y fixed, 2) insufficient quality of generation which leads poor generalization from generated to real images.
DISCUSSION
We develop an algorithm for OOD generalization in the presence of spurious correlations induced by nuisance variable. We formalize nuisance-induced spurious correlations in nuisance-varying family, where changing nuisance-label relationships make predictive models built from samples of one member unsuitable for other members. To identify conditional distributions that have per- formance guarantees on all members of the nuisance-varying family, we introduce the nuisance- randomized distribution and uncorrelating representations. We characterize one uncorrelating rep- resentation that is simultaneously optimal for all members. Then, we show that uncorrelating repre- sentations most informative of the label under the muisance-randomiin same optimal performance. Following this result, we propose to estimate the muisance-randomized distribution and, under this distribution, construct the uncorrelating representation that is most in- formative of the label. We develop an algorithm called NURD and show that it outperforms ERM on synthetic and real data by adjusting for nuisance-induced spurious correlations. Our experiments show that NURD can use casy-to-acquire nuisances (like the border of an image) to do this adjust- ment; therefore, our work suggests that the need for expensive manual segmentation, even if it does help exclude all the nuisances, could be mitigated.
distribution also achieve the
Limitations and the future. Given groups based on pairs of nuisance-label values, Sagawa et al. (2020) suggest that subsampling equally from each group produces models more robust spurious correlations than reweighting (Byrd and Lipton, 2019; Sagawa et al., 2019); however, subsampling ineffective when the nuisance high-dimensional. Instead, as sufficient statisties of the condi- tional (y z) render y, z independent, grouping based on values of sufficient statistics could be promising. The nuisance-rane randomizaanoiiioiioo steps in generative-Nur different distributions in the training distribution: Ptr(x y, z) and (y z) respectively. Methods that combine the two approaches to produce better estimates of the nuisance-randomizee distribution would be interesting. The first step in reweighting- NURD is estimate Pe- z). As deep net- works tend to produce inflated probabilities (Guo et al., 2017), one must take care to build calibrated models for p(y z). Adapting either cllibration-focusee losses (Kumar al., 2018: Goldstein al., 2020) or ensembling (Lakshminarayanan et al., 2016) may produce calibrated probabilities. In our experiments, the training data contains single environment. Methods for invariant representation learning (Arjovsky et al., 2019; Krueger et al., 2020; Mahajan et al., 2020; Guo et al., 2021) typi- cally require data from multiple different environments. Huisance-randomiza nuisance- relationship from the training data, meaning it is a different environment from the training data. Following this insight, using muisance-randomization to produce samples from differ- ent environments using data from only single environment would fruitful direction. The absolute performance for both ERM which exploits spurious correlations and NURD which does not, is too low to be of use in the clinic. Absolute performance could be improved with larger models, more data, using pretrained models, and multi-task learning over multiple lung conditions, all techniques that could be incorporated into learning procedures in general, including NURD.
and reweighting-Ner RU model
data has a different

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
The authors were partly supported by NIH/NHLBI Award ROIHL148248, and by NSF Award 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science. The authors would like to thank Mark Goldstein for helpful comments.
REFERENCES
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01-344 2020.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Pazz Invariant risk minimization. arXiv preprini arXiv: (1t.iv:907.00880) 2019.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A. Wichmann. Shortcut learning deep neural networks, 2020.
Ziwei Liu, Ping Luo. Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference computer vision, pages 3730-3738, 2015.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning, pages 8346-8356. PMLR, 2020.
Zhao Wang and Aron Culotta. Identifying spurious correlations for robust text classification. arXiv preprint arXiv:2010.024558 2020a.
John R Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric Karl Oermann. Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs: a cross-sectional study. PLoS medicine, 15111)e1,02668, 2018.
Marcus A Badgeley, John R Zech, Luke Rayder,-aayee, Benjamin S Glicksberg, Manway Liu, William Gale, Michael V McConnell, Bethany Percha, Thomas M Snyder, and Joel T Dudley. Deep learning predicts hip fracture using confounding patient and healthcare variables. NPJ digital medicine, 2(1):1-10, 2019.
Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density ratio estimation in machine learning. Cambridge University Press, 2012.
David Krueger, Ethan Caballero, Heeen- Jacobsen, Amy Zhang. Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex). arXiv preprint arXiv:2003.0 00688, 2020.
Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. arXiv preprini urXiv:2006.07500, 2020.
Ruocheng Guo, Pengchuan Zhang, Hao Liu, and Emre Kiciman. Out-of-distribution prediction with invariant risk minimization: The limitation and an effective fix. arXiv preprint arX1v:2101.07732. 2021.
Adarsh Subbaswamy, Peter Schulam, and Suchi Saria. Preventing failures due to dataset shift: Learning predictive models that transport. The 22nd International Conference on Artificial Intelligence and Statistics, pages 3118-3127 PMLR, 2019a
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- 2019.
tion. arXiv preprint arXiv:191I.088
Victor Veitch, Alexander D' Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invari- ance to spurious correlations: Why and how to pass stress tests. arXiv preprint 210.:::2000000555, 2021.
Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpern, and Alexander D' Amour. Causally-motte motivated shortcut removal using auxiliary labels. arXiv preprint arXiv:2105.0644 2021.
10

---

Published as a conference paper at ICLR 2022
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Meta-learning for domain generalization. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018a.
Kei Akuzawa, Yusuke Iwasawa, and Yutaka Matsuo. Adversarial invariant feature learning with ac- curacy constraint for domain generalization. In Joint European Conference Machine Learning and Knowledge Discovery in Databases, pages 315-331. Springer, 2019.
Alexis Bellot and Mihaela van der Schaar. Accounting for unobserved confounding in domain generalization. arXiv preprint arXiv:2007 2020.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adver- sarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 5400-5409, 2018b.
Karan Goel, Albert Gu, Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation. arXiv preprint arXiv:2008.06775, 2020.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain gener- alization. arXiv preprint arXiv:2102.10395, 2021.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversaaiill training of neural net- works. The journal of machine learning research, 17(1):2096-0330, 2016.
Qizhe Xie. Zihang Dai, Yulun Du, Eduard Hovy, and Graham Neubig. Controllable invariance 2017.
through adversarial feature learning. arXiv preprint
Brian Hu Zhang, Blake Lemoine, and Margaret Mitchell. Mitigating unwanted biases with adver- sarial learning. In Proceedings of the 2018 AAAI/ACM Conference on AI, Ethics, and Society, pages 335-340, 2018
Sandesh Ghimire, Satyananda Kashyap, Joy T Wu, Alexandros Karargyris, and Mehdi Moradi. Learning invariant feature representation to improve generalization across chest x-ray datasets. In International Workshop on Machine Learning in Medical Imaging, pages 644-655 653. Springer, 2020.
Ehsan Adeli, Qingyu Zhao, Adolf Pfefferbaum, Edith V Sullivan, Li Fei-Fei, Juan Carlos Niebles, and Kilian M Pohl. Representation learning with statistical independence mitigate bias. In Proceedings of the IEEE/CVF Winter Conference Applications of Computer Vision, pages 2513-2523, 2021.
Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, and Aaron Courville. Can subnetwork structure be the key to out-of-distribution generalization? International Conference on Machine Learning. pages 12356-1223 12367. PMLR, 2021.
John C Duchi, Peter Glynn, and Hongseok Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 2021.
Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang. and Ed H Chi. Fairness without demographics through adversarially reweighted learning, arXiv preprint arXiv:2006.1 13114, 2020.
Natalia L Martinez, Martin A Bertran, Afroditi Papadaki, Miguel Rodrigues, and Guillermo Sapiro. Blind pareto fairness and subgroup robustness. In International Conference on Machine Learning, pages 7492-7501. PMLR, R, 2021.
Elliot Creager. fömn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference Machine Learning, pages 2189-2200. PMLR, 2021.
Bernhard Schölkopf. Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. arXiv preprint arXiv: /2066441,
2012. 2012.

---

Published as a conference paper at ICLR 2022
Adarsh Subbaswamy, Bryant Chen, and Suchi Saria. A universal hierarchy of shift-stable distri- butions and the tradeoff between stability and performance. arXiv preprint 2019b.
Erich Lehmann and George Casella. Theory of point estimation. Springer Science & Business Media, 2006.
Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-llcus, Chris Chute, Henrik Marklund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, et al. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 590-597, 2019.
Alistair EW Johnson, Tom Pollard, Nathaniel R Greenbaum, Matthew P Lungren, Chih-ying Deng, Yifan Peng, Zhiyong Lu, Roger G Mark, Seth Berkowitz, and Steven Horng. Mimic-cxr-jpg, large publicly available database of labeled chest radiographs. arXiv preprint 2019.
Eduardo HP Pooch, Pedro Ballester, and Rodrigo C Barros. Can we trust deep learning mod- els diagnosis? the impact of domain shift in chest radiograph classification. arXiv preprint arXiv:1909.01944 2019.
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting in deep learning? In International Conference on Machine Learning, pages 872-881. PMLR, 2019.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, pages 1321-1330. PMLR, 2017.
Aviral Kumar, Sunita Sarawagi, and Ujjwal Jain. Trainable calibration measures for neural networks from kernel mean embeddings. International Conference on Machine Learning, pages 2805- 2814. PMLR, 2018
Mark Goldstein, Xintian Han, Aahlad Puli, Adler Perotte, and Rajesh Ranganath. X-cal: Explicit calibration for survival analysis. Advances in Neural Information Processing Systems, 33, 2020.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. arXiv preprint 1rX1v:16/2.01444, 2016.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In International Conference on Machine Learning, pages 1448-1458. PMLR, 2020.
Fan Zhou, Zhuqing Jiang, Changjian Shui, Boyu Wang, and Brahim Chaib-draa. Domain general- ization with optimal transport and metric learning. arXiv preprint arXiv:2007.10573, 2020.
Hal Daumé III. Frustratingly easy domain adaptation. arXiv preprint arXiv:0907. 2009.
Shai Ben-Daviee John Blitzer, Koby Crammer, Fernando Pereira, et al. Analysis of representations for domain adaptation. Advances neural information processing systems, 19:137, 2007.
Abolfazl Farahani, Sahar Voghoei, Khaled Rasheed, and Hamid R Arabnia. A brief review of domain adaptation. Advances in Data Science and Information Engineering, pages 877-894, 2021.
Zhao Wang and Aron Culotta. Robustness spurious correlations in text classification via automat- ically generated counterfactuals. arXiv preprint arXiv:2012. 2020b.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1eXiv:1112.6944 69800 2014.
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2, arXiv preprint arXiv:/906.000 990600046,, 2019.
12

---

Published as a conference paper at ICLR 2022
A FURTHER DETAILS ABOUT NURD AND PROOFS
A.1 DETAILS ABOUT NURD
The algorithm boxes for reweighting-NuRE and generative-NuRD are given in algorithms 1 and
2.
Estimating and using the weights in reweighting-NuRD. In learning Ptr x) for high- dimensional z, flexible models like deep neural networks can have zero training loss when the model memorizes the training data. For discrete y, such model would output Por (y z) every sample in the training data. Then, the model's weight estimates on the training data are pix)/pe-ly 2) ply). Weighting the training data with such estimates fails to break the nuisance- p(y)
for
label relationship because (t, (,, x) pt-ly - Ptr z, x)p(y) x Ptr y)p(x y, 2). To avoid such poor weight estimation, we employ a cross-fitting procedure: split the data into K disjoint folds, and the weights for each fold are produced by model trained and validated on the rest of the folds. See algorithm for details. In estimating loss for each batch under p. during training, one can either weight the per-sample loss or produce the batches themselves via weighted sampling from the data with replacement.
Density-ratio trick in Distillation. The density-ratio trick for estimating mutual information (Sugiyama et al., 2012) involves Monte Carlo estimating the mutual information using binary classifier. Let e = be the pseudolabel for samples from p. (x), and l = 0 for samples from P. (y, r- (p_p))pp)))) (z). Then,
p (x). Ep, (xxxx) log (x,, ry, Fr (x)
p(e = 1 y y,,,
Bp (v.xx) log
1 p(e = 1
y. 1 (x))
With parameters Φ, we estimate the conditional probability with a critic model, denoted pé-
Accounting for shifts in the marginal label distribution. NURD relies on the assumption in eq. (1) that distributions in the nuisance-varying family F have the same marginal p(y). What happens if Pte comes from a nuisance-varying family with a different marginal? Formally, with Per F. let pte belong to a nuisance-varying family F = (p==lt//pt--x)pp 2, Pte Pte(y) (y)Pp(Z y))( (x( y. z)} where PD E F. Given knowledge of the marginal distribution po.l0l//p..x)pprrry z, x) lives in F'. Run- ning NURD on Pt, produces predictive models that generalize to Pte- To see this, note p' (y,z,x) = Per "rr (z) y, 2) is randomize distribution in F¹, With R(y') as the uncorrelat- ing set of representations defined with respect to p'_. i.e. r'(x) E R(F()) = y z r', lemma and theorem hold. It follows that running NURD on samples from pt.(l)//.--(yPPP. (y, z, x) pro- duces an estimate of p', (y rx()) (r'(x) E R(p()) )) with the maximal performance every Pte E FF' if maximally blocking (x) E R(p'))
Pte(N (y), note that the weighted training distribution ptr"
A.2 EXTENDED RELATED WORK
Domain generalization methods aim to build models with the goal of generalizing to unseen test data different from the training data (Gulrajani and Lopez-Paz, 2020). Recent work uses multiple sufficiently different environments to generalize to unseen test data that lies in the support of the given environments or subgroups (Li et al., 2018a; Arjovsky et al., 2019; Akuzawa et al., 2019; Mahajan et al., 2020; Krueger et al., 2020; Bellot and van der Schaar, 2020; Li et al., 2018b; Goel et al., 2020; Wald al., 2021). Chang al. (2020) develop multi-ennironment objective to interpret neural network predictions that are robust to spurious correlations. Similarly, domain-invariant learning and related methods build representations that are independent of the domain (Ganin et al., 2016; Xie et al., 2017 7; Zhang et al., 3018; Ghimire et al., 2020; Adeli et al., 2021; Zhou et al., 2020).
Due to its focus on nuisances, NURD works with data from a single environment. As in section 4, to split the data into multiple environments, one can split the data into groups based on the value of the nuisance. Then, domain-invariint methods build representations that are independent of the nui- sance and under ienisance-inducee spurious correlations these representations may ignore semantic features because they are correlated with the nuisance. Domain adaptation (Daumé III. 2009; Ben- David et al., 2007; Farahani et al., 2021) methods assume access unlabelled test data which NURD
13

---

Published as a conference paper at ICLR 2022
does not require. We do not assume access to the test data because nuisance- relationships can change over time or geography which, in turn, changes the the target distribution.
Algorithm 1: Reweighting-NRR
Input:
Training data D, specification of the weight model P z) which estimates (r z), representation model (x),), predictive model po(y Iry (x)) and critic model
( y, z, r-, (x)); regularization coefficient A, number of iterations for the weight
model N for the predictive model and representation Np, and the number of critic
model steps Ne. Number of folds K.
Result: Return estimate of Pa (x)) forry r E R(p_)) with maximal information with y.
Nuisance Randomization step;
Estimate the marginal distribution over the label p(y);
Split data into K equal disjoint folds, D = {Al]k for cross-fitting;
for each fold F; < K do (cross-fitting)
Initialize (y z); for Nw iterations
Sample training batch from the rest of the folds (F-i): B ~ F-i
Compute likelihood L(,, = log Pa (y, 1;);;
Update to maximize this likelihood (via Adam for example);
end
Produce weights w, p(xx)/Dalx, I a.) for each (y,, Fú
end Distillation step:
Initialize T,,, Pe-PPeP2 Pai for Np iterations do for Ne ittratiioss do
Sample training batch B ~ D and sample independent copies of z marginally: ži ~ D; Construct batch B {y, x;};
Compute likelihood
w; log y, Zi (x),)
Σ
wa log 0| 2i, rx,x)))
(
Update & to maximize likelihood (via Adam for example); end
Sample training batch B D;
Compute distillation objective using the density-ratio trick
pole yis 2s, (x:))
Σ
W/ log byg (x.)) - A Alog 1- (- Y Z,, (x:))
B
(
Update 0,7 y to maximize objective (via Adam for example).
end
Return (y 1r, (x)).
Taking distributional robustness (Duchi et al., 2021) approach, Sagawa et al. (2019) applied group- DRO to training data where the relative size of certain groups in the training data results in spurious correlations. Given these groups, group-DRO optimizes the worst еттог across distributions formed by weighted combinations of the groups. With high dimensional groups based on the value of the nuisance typically results in groups with at most one sample; with such groups, group-DRO will encourage memorizing the training data. Other work aims to minimize worst subgroup error with finite number of fixed but unknown subgroups (Lahoti et al., 2020; Martinez et al., 2021): as subgroups are unknown, they only find an approximate minimizer of the worst subgroup error in general even with infinite data.
in our experiments, defining
In contrast, NURD builds predictive models with performance guarantees across all test distribu- tions (that factorize as eq. (1)) using knowledge of the nuisance. Given the nuisance, existence of a finite number of subgroups maps to an additional discreteness assumption on the nuisance vari- able; NURD works with general nigh-dimensional nuisances. Wang and Culotta (2020b) focus on
14

---

Published as a conference paper at ICLR 2022
sentiment analysis of reviews and build dataset where the nuisance label relationship is destroyed by swapping words known to be associated with sentiment of the review, with their antonyms. This equivalent to using domain-spectin knowledge sample from p(x y, 2) in generative NURD. NURD requires no domain-specific knowledge about the generative model p(x y.2).
Algorithm 2: Generative-NuRI
Input:
Training data D. specification of the generative model ( y, z) that estimates Ptr(x y. z), representation model (x)), predictive model pe ry rx(x)), and critic model p& (E y, 2, ryy (x)); regularization coefficient X, number of iterations for the weight model Nw. number of iterations for the predictive model and representation Ny-
number of critic steps NN-
Result: Return estimate of Pı (y (x)) for ry E R(p_)) with maximal information with y.
Nuisance Randomization step;
for Nu iterations do
Sample training batch B ~ D;
Compute likelihood Lyy,,, (x,))n e,)nn log (i, (or some generative objective);
Update 3 to maximize objective above;
end
Estimate the marginal distribution over the label p(y);
Sample independent label and nuisance D, 2gj D, and then sample x ~ po(x y;
Construct dataset D using triples {yx = Yo, 2k Zjj 2k x):
Distillation step:
Initialize Poi for Np iterations do
for N₂ iteratiion do
Sample training batch B ~ D and sample independent copies of z marginally: i; D;
Construct batch B (y, Zi, x; using B;
Compute likelihood
w; log = ys, 2i,
Σ
wi 10 (x,,,,x))
(E)
(,,
Update & to maximize likelihood (via Adam for example); end
Sample batch from generated training data B ~ D: Compute distillation objective using the density-ratee trick
Σ
log Po (xx()))) Alog
(e YK Zg. Zg. (xx))
B
p& = y Ziz (xx)
(E)
Update 0,7 maximize objective (via Adam for example).
end
Return (y (x))..
15

---

Published as a conference paper at ICLR 2022
A.3 KEY LEMMAS FOR UNCORRELATING REPRESENTATIONS R(p.)
In this lemma, we derive the performance of the nuisance-randomized conditional (y r(x)) for any r R(,,) and show that is at least as good predicting without covariates on any Pte E F.
Lemma 1. Let F be a nuisance-varying family (eq. (I) and p_ - =p(y)p:r(z)pss y, z) for some Ptr E F. Assume Vpn E F, pn(z y) is bounded. If r(x) E R(p,). then Vpte E F, the performance of P. (y r(x)) is
Perfpie (p r(x))) Perf (p(y))
E (x)(y
KL [p(r(x) 2) (xts)(((x)) ly, |y,)]] (8)
As the Kl-divergence is non-negative, Perf pie (p r(x))) 2 Perfpp pre (p(y)).
Proof. (of lemma 1) Note that the identity Ep(x)g f(x) = [p(((x)) f(x) implies that
Pte (y)
Pte (y)
xx.) log
E (x.c(x)(())) log
P. (y r(x))
(_ r(x))
As p. yy) (z) > 0 on E Sy and y s.t. p(y) > 0 is bounded, lemma 3 implies that This fact implies the following KL terms and expectations of
(y 2, x) > 0 + Pto (y, z, x) >
log-ratios are all well-defined:
-Perf (p r(x))) = (x.,(x)KL KL proly | x) 11 P [x(x))
-Eperyyx)) log Pie (y x))
(y r(x)))pt(())
Pro (y [x]
+ (x log Pyo
(x) log
(y)
(y r(x))
xp),ly.x) log (y |x)
ply)
+ E (x.((x)(()) log p. (y r(())
(y)
-Epr(x)) KL proly x) p_]])] (x())x2))()) log p(y)
P1 (y ((x))
Py()
(x KL Pre (y x) Py)] + Ep Ep (r(x)) I yxx) log
p. (y [(x))
Py z))
Exp.(x) KL Py x) [y)] + E E (2((()) I y.a) log
P_(y r((x), z)
P_ (y x))r(r(x)) (z)
Ep (y x) (y)] + Ep (y EE (r(x)) y.a) log
+
Pa (r,y
Pi (r(x) (2)
= - KL. pre(y x) || [y)]] + Bp E (x
log
(r((x) y,z)
Pr (r(x))) |y,z)
= -Ep KL (y x) || [y)] Emi E (r(x)
yya log
p. (r(x) z)
(x KL (y
x) [y)] (y. KL [p_ (r(x) y,z) Pa (r(x) z))
Here, p. (r(x)
y. p(((x) y, as P. (x y.z) = p(x y. z) by definition of the nuisance-
varying family. The proof follows by noting that the gap in performance of P r(x)) and p(y) equals an expected KL term:
-Perf (ply)) + Perf (p.ly r(x))) Epse (y.x)) KI. (r(x)) ly,2) || (r(x) 2)]
(x. (x. KL [p(((x) (y,z Ep(y)p(((x)) y, z)]
Rearranging these terms completes the proof.
Lemma 2 shows that uncorrelating sets are the same for any muisance-randomized distribution and that the conditional distribution of the label given an uncorrelating representations is the same for all nuisance-rndoomized distributions.
16

---

Published as a conference paper at ICLR 2022
Lemma 2. Let F be muisance-varying family with p(y) and p(x y, z) and muisance space Sy. Consider distributions P1. x) = (Ly)p y. z) and Pı.2(y,2,33 = for Sy, and (y, > O = 2y, x) > 0. Then, the uncorrelating sets are equal R(p1,1) = R((12) and for any r(x) E R(p_n))
ply)p (z))pp( y.z) such that (z) > 0, (2,,,((z)
P_, [r(x)) = Py r(x)).
(r()))
Proof. By the assumption that (y, 2, > 0 + pL.((y, z, x) > 0. there exist some z such that Pi,1 r(x)) and p_,((z r(x)) 0. With such 2, for any R(p,,)),
Py r(x)) 1y ((x),2)
= p(y) -p())<(r(x) p(((x) y2
D (r(x((() ())
p(r(x) yy,)
-(y))
E pa x) (I(),)(((() z.y]]
-p(y) p(r(x) y,,))
((())((((x) y2
p(r(x) y,z)
= =p(y) Epr,ely x)p(r(x) z.
=p(y) p(r(x) [y,
(r,((r(x) 2))
2 |r(x), z)
Taking expectation on both sides with respect to PL.z(z r(x)),
Ep (y r(x)) = E. p2, r(x)))x.22 |rxx), 2) = pm.2(y [x(x))
(10)
Note that E pp (x r(x)) py rx))), which implics
1 (y r(x) = ( rrxx),) 2) = Py.2 rrxx) 2) = PL.2 |r(x)),
completing one part of the proof, p.. (y rxx)) = P r(x)).
Further, we showed y y_lp,.. 2 r(x) = y_lp,, | which means r(x) E R(p2.)). As the
above proof holds with p,,1, P..2 swapped with each other, r(x) E R(p.,)) < r(x) E R(p,2)).
The next lemma shows that every member of the nuisance-varying family is positive over the same set of y, z, x and is used in proposition I and lemma 1.
Lemma 3. Let the nuisance- varying family F be defined with p(y),p(x y, z) and nuisance space SF. Let distributions pD = p(y)pp(z (z y y)p(x z. and P'D p(y)p'((z y)p(x z, be such that PD (z) ()( y) > 0 for all y such that p(y) > and z Sz. Further assume
po(z y),P'(zz y) are bounded. Then, pDly, x) 0 = py,y,,,x x) 0.
Proof. For any E Sy and any y such that p(y) 0, (z y) > 0 and
Ln Iy PD Iy] >0, 0,
pl,y,, = p(x y, ( y)p(y) = p(x y. (z y)p(y) (z y) PD y)
-p z, Pb y) PD y) (11)
Thus, for all z in the nuisance space S and any such that p(y) > 0,
(y z, x) > - py 2, > 0.
As z only takes values in the nuisance space S7 when p(y) = 0,
P'D (0(y,,, x) = PD 1.y,,2x'] x,
Together, the two statements above imply py z, x) + Pbyly, x.x) > 0.
17

---

Published as a conference paper at ICLR 2022
A.4 OPTIMAL UNCORRELATING REPRESENTATIONS
Theorem 1. Let E R(p, )) be maximally blocking: Vr E R(p,), y.lpa r(x) (2,2r((x) Then,
1. (Simultaneous optimality) Vpte E F, Vr R(P2). perfpr, (r*(x)) 2 Perf (r(x)).
2. (Information maximality) Vr(x)= E R(p.), Ip, 2 lp (y;r(x)).
3. (Information maximality implies simultaneous optimality) Vr E R(p.),
Ip_ (x:) Ip (x)r"(x))
Vpte F, Perf (r)(x)) Perf (r)()))
Proof. (proof for theorem 1)
We first prove that for any pair r,72 R((p)) such that blocks r, rx)) r2)2) domi- nates the performance of r on every Ptc E F. The simultaneously optimality of the maximally block- ing representation will follow. For readability, let E(r2) (x) KL [Pte(y x) pı [x()))] We will show that
(x) KL [Prc(y x) || PL (rx))) 2 E(r2).
We will use the following identity which follows from the fact that p(x y, 2) does not change between distributions in the data generating process eq. (1):
PD(12(x) I 2. r(x)) = pD(r2(x), y,2)
(D(r(x) y.z)
plr2(x), r(x) y.z)
p(r(x) y222)
p(r((x) y,z, r(x)).
Next, we will show that
Epi (x) KL (Pre(y x) pl (y [(x))]
[(r2) + E Ep. KL [p (x2(x) v,z,r(x)) (r((x)) (r()) rx)x)
The steps similar lemma I's proof
Ptely x) + E Excc log P. r2(x))
Exc (x) KL [pre(y| xx) P2 r(x))] (xixx log
P. (y rr(xx))
( r(x))
(y r2(x))
=((22) + I (x)) log
p- (y r(x))
P, (y (2()))
= =((r2) + (xx)))))) log
(as R(p,)}
P1 (y r(x),z)
P1 r2 (x), rx)))
-((22) + E (xxx()))x))))))))())))))) log
{y (yIlpa rxx) 2, r2()))
p. (y r(x),z) z)
p, (y.r((x) [(()),))
-((22) +
log
P1 (y
(rz(() r(x),2)
P1 (r2(x) y, rx))),,
[((2)) + (x.. log
p, (r2(x) r(x),
-((r2) + Ep (y.,, (px(t) yz.((x)) log P. (re(() y, rx),
P_ (r2()) r(x), z)
=((2)) + (er2)+Epp..ttee (y.... KL (p. (r2(x) y, 2, rxx)) p. (r2(x) [(xx,, z)]
Noting that KL is non nog-neeative and that Perf is negative KL proves the theorem:
Ex KL [pte (y x) || p_ rx(x))) E (x) KL (y (y x) P2 r(x))
(12)
It follows that for maximally blocking r*
Vr E R(p,)
E (p KL (y xx p. (y (x)) < (x) KL Kpte (y x) 1.plyy rxx))].
18

---

Published as a conference paper at ICLR 2022
As performance is negative KL, the proof follows that r" dominates in performance. This concludes the first part of the proof.
For the second part, we prove information maximality of a maximally blocking r" (x) E R(p,) ). The proof above shows that the model p. (y rx)x)) performs at least as well as p. (y r(x)) for any r(x) E R((p,) on any Pte E F. We characterize the gap in performance between Pa (y r (x)x)
and P1 r(x)) for any r(x) E RR(p_) as the following conditional mutual information term:
(x.x.((x KL [p_ (r"(x) y,2,r(x)) (r (x) rxxx,, z)] = Ip_ (x):x z, r(x)).
The entropy decomposition of conditional mutual information (with (-) the entropy under distribution 4) gives two mutual information terms.
Ep, (x)).((x)) KL [p_ (x) y, 2, r(x)) (r(""(" (r" (x) rxxx), 2)] = Ip. (r- (x)::y z. r(x));
= Hp1 (y 2, rxx) Hp (y 2, r(xx),""(x)
Hp, (y rx)) - Hp, (y (x))
Hpp (y |(xx) - Hp (y (x)))
fr,r" E R(p.))
Hp, (y [xx)) - Hp, + Hp, (y) - -p, (r
Ip_ (y:r"(x)) - ry,x)x(
This difference is non-negative for any r R(p_) which proves the second part of the theorem:
1p_ (y; (x)) - ln_ (y, rx)) = I2_ (r"(x))) z, r())) 2
For the third part, note that any representation which satisfies Ip rx))) = Ip, (x)) eniformation-equuialleee)) also satisfies
Ip (x))"(x) z, r' (xx) = = y.l.. r" r'(x) 2, (xx).
Under this condition, eq. (12) implies
EPc (x) KI. KI. (pte (y x) PL (y (x(())] 2 (x) KL. Ppally x) || (y [x)]]]
However, as r' R(p,)) and that r (x) is maximally blocking, which (by the proof above) implies
E Pp (x) KL [Preey x) (y rx(x)) S Epcc(x)
KL (Ptoly x) || (y r((x))]
The only way both these conditions hold is
E Ep.. KL (y x) ( | rx())] Ex. KL (y x) l| (_ |r(())]]
This completes the proof that for any r' R(p,)) that is information-equuivale. to r (x) under P1. the model p_(y (y (x)) has the same performance as py | (x)) for every pte E F, and consequently, y' is also optimal.
A.5 MINIMAX OPTIMALITY
Proposition 1. Consider a muisance-varyinn family F (eq. (1)) such that for some Ptr E F there exists a distribution Pa_ E F such that P1 = (()))xx())p(( y,2) E F. Let F satisfy
y Ydho z
3pD Fs.t. KL x) PD(y x)] - (x:)) > 0. (13)
If y flllp z x, then Paly x) is minimax optimal: P. x) arg arg max KL) (y x) PD x)].
Poly P
Proof. (of proposition 1) By lemma 3, as PD x) > 0 + P'D (y, x, > 0. performance is well defined for any (y x) on any P'D E F. First, lemma 1 with Pte P/D and r(x) = x gives
Ibp (x:)) (x KL by x) || (p.... (y x)]
(x KL PPD (y x) | (y)) E (x KL Wbly xx) ( x)] (14)
=Epi x) KL [p(x ly.z) Ex y,z)].
-Rpb(ra) KI. (p_ (x y,z) (x 2)) 2
19

---

Published as a conference paper at ICLR 2022
Thus, unlike any pD F such that y:hpp 2,
max FDEF
(x) KL (y x) Py x)] - I (x: y) - 0.
(15)
For any pD such that "hipo pp 2, let p/D such that Eyg (x) KL (y x) PD x)] - 1o (x: y) >
(
As eq. (15) implies B KL (Pb (y x) P. x]] 1/b yx - 0, follows that vpp such that po
PhEF Bob(x) KL (y x) PD(y x)]
max Bpg (x) KL (y x) Pi x)] FTE
By lemma 3, any E F is positive over the same set of y, x and if y y 2, then (D x)
Pi (y x) (see lemma 2 for proof with instantiation r(x) = x). This means
( x) = arg min max E (x) KL PP (y x) Po(y x)] pDEF DDEF
See proposition 3 for example nuisance- family where the information criterion in eq. (13)
holds.
A.6 DISTILLATION DETAILS AND LOCAL OPTIMA EXAMPLE FOR EQ. (5)
0.5
Objective landscape with =0
0.0
00.5
0.5
-1.0
-1.5
-2.0
005
-2.5
-3.0
1.0
1.0
-3.5
Figure 3: Landscape of the objective in eq. (5) for the example in eq. (2) for linear representations Fu,o (x) = uX1 + UX2. Local maxima correspond to representations r-n,a and global maxima to representations Tu,a-
The objective in eq. (5) can have local optima when the representation function of the nui- sance and the exogenous noise in the generation of the covariates given the nuisance and the label. Formally, let the exogenous noise e satisfy (E, 2) ll, y. Then,
(E,z) - (f(E,))
2 f(z,e).
Such representation r(x) = f(c,z) is both in the uncorrelating set and independent of the label under the nuisance-randomizze distribution p. meaning it does not predict the label.
Local optima example for conditional information regularization eq. (5). Figure plots the value of the objective in eq. (5) computed analytically for À 20, over the class of linear represen- tations indexed by u. E R. (x) = ux1 + UX2, under the data generating process in eq. (2). Representations of the kind (x) u(x; - X1) functions of and some noise independent of the label and, as fig. 3 shows, are local maxima on the landscape of the maximization objective in eq. (5). Global maxima correspond to representations Tu..--
20

---

Published as a conference paper at ICLR 2022
Performance characterization for jointly independent representations
Lemma 4. Let F. be nuisance varying family. For any jointly independent representation r, i.e. [r(x), y 2
Ptt F Perf Pce (px(y [r(x))) = Opcc + Ip (r(x);y),
where Cp. is Pte dependent constant that does vary with r(x).
NURD maximizes the information term I, (x)r-(x)) and, therefore, maximizes performance on every member of F simultaneously. follows that within the set of jointly independent represen-
tations. NURD, at optimality, produces representation that is simultaneously optimal on every Pte F.
Proof. Lemma says that for any uncorrelating representation (R(p.) and VDte E F,, E KL [p(r(x) y,z) (pp))(((())) y. z)]
Perf (p_(y r(x))) Perf (p(y))
(t)(y
However, the joint independence [r(x), yل z implies both the uncorrelating property and r(x)_pp y, the second term in the RHS above can be expressed as Ip (r(x);y):
KL. [p(((x)
y l.2)) (pp())(((()) y.z)]
Pre (x)))
E KL (r(x) yy, 2) || Ep (y)P, (r(x) v,
(v)()
KI. (r(xx) yy) || PPp (r(x)) yy) (to(y
E KL (p_(((()) yy) (x(r(x))]
(ie(y,a)
E KL [p, (r(x) Iy) (_((x)))]
=
Peo(y)
=
E KL [p. (r(x) y) (r(x))
(y(y)
Lp, (r(x);y).
Noting Cp.. Perf pio (p(y)) does not vary with r(x) completes the proof.
Performance gaps between jointly independent representations and uncorrelating representa- tions. The joint independence [r(x), y. 22_ implies the uncorrelating property but uncorrelating representations only satisfy this joint independence when they are independent of the nuisance. Thus, representations that satisfy joint independence form subset of uncorrelating representations. This begs a question: is there loss in performance by restricting NURD representations that satisfy said joint independence? In appendix A.7.I, we use the theory of minimal sufficient statistics (Lehmann and Casella, 2006) show that there exists nuisance-varying family where the best un- correlating representation dominates every representation that satisfies joint independence on every member distribution, and is strictly better in at least one.
A.7 COUNTERFACTUAL INVARIANCE VS. THE UNCORRELATING PROPERTY
We A) show that counterfactually invariant representations tions by reducing counterfactual invariance to the joint independence [r(x). y'',,, z and B) give an example nuisance-varying family F where the best uncorrelating representation strictly dominates every jointly independent representation in performance every test distribution Pte F: least as good on all Pte E F and strictly better on at least one.
subset of uncorrelating representa-
We show A by proving counterfactually invariant representations satisfy joint independence [r(x),y] z which implies the uncorrelating property, but not vice versa. Counterfactual invari- ance implies that for all PD F, the conditional independence r(x) in Veitch et al. (2021). As y y_I.p. it follows that [xxx,,]] this joint independence implies the uncorrelating property, y yllp, z r(x). But, uncorrelating representations only satisfy the said joint independence when they are independent of the nuisance.
holds by theorem 3.2
We show B in appendix A.7.1 1. by constructing a nuisance- nuisance-varying family where the optimal perfor- mance is achieved by uncorrelating representation that is dependent the nuisance.
21

---

Published as a conference paper at ICLR 2022
A.7.1 JOINT INDEPENDENCE VS. THE UNCORRELATING PROPERTY
Here, we discuss the performance gap between representations that are uncorrelating (y_l,, 2 r(x)) and those that satisfy the joint independence (y, r(x)) We construct data generating process where optimal performance every member of J is achieved only by uncorre- lating representations that do not satisfy joint independence.
Define muisance-varying family F {pp(y, x) p(y)pp(z y)p(x y, 2)) Let R. {r(x); [r(x);[((x), [rxx). that, under the muisance-ranen randomizd distribution, satisfy joint independence and conditional inde- pendence respectively. Then there exists nwisance-varring family F such that
Theorem 2.
z} and Rc = fr(x): r(x)} be the set of representations
Ype E E
max Perf Pie (r(x)) S max Perf (r()))
(16)
rERJ
rERc
and Spre E F for which the inequality is strict
max Perf prc < max perf (r(x)),
(17)
FERJ
reRe
Proof. In this proof we will build a nuisance-varying family F such that p2 E F and y y.l.p. z x. This makes x maximally blocking uncorrelating representation because r(x) y x, z. Thus has optimal performance on every Ptc. E F within the class of uncorrelating representations. We let y be binary, Pa E F. The structure of the rest of the proof is as follows:
1. The representation f(x) x on every member of the family F.
(y = 1 x) is optimal in that it performs exactly as well as
2. Any representation T(x) that matches the performance of x every Pte E F satisfies y y_ x T(x).
3. All functions T(x) such that y x T(x) determine f(x). This is shown in lemma 5.
4. We construct family where {x) which, by the point above, means that every opti- mal representation Txx) dependent on z: T(x) z. But every representation r E R, satisfies r(x)_lp, z and, therefore. is strictly worse in performance than f(x) on p_, mean- ing that they perform also strictly worse than x (because Perf pr. (f(x)) Perf (x)) Noting x E Re completes the proof.
For 1, let f(x) = (y = 1 x). However, we show here that (, f(x)) = Pa (y x).
Pa (y = 1 1 fxx) = Ep_( f5 = (x=))py==11xx 11x, f(x))
Epp( (6)))Pa( = 1|
Epp (xx)
= =((x)
= -y
(= P= x, (x)))
This means fxx performs exactly as well as x every Pte F and xاI fxx).
For 2, recall Perfp.. (r(x)) Ex KL (y x) p, (y r(x)) and note that I implies
Perf p_ (x) = Perfp, (f(x)) = O.
Let T(x) be any function that performs as well as [xx on every Pie As FF.
0 = Perf,, (f(x)) = Perf,, (T(x))
-E)(()) KL (p, (y ix) II p. [(x)]
(x) KL [p_ xx (x)) 1l p [(x)]
-II,, (y; x T(x))
y_lp2 x T(x).
We leave 3 to lemma 5 and show 4 here.
22

---

0.1. Published as a conference paper at ICLR 2022
The example data generating process. We give data generating process where f(x) p(y
1 x) is dependent on z : f(x)iLp,2 z. We assume Pa E F... With a binary y and a normal z, let
(y, 2, x) = p(y)p(z)p(s ( y, 2) be generated as follows: with p {0, 1 x {0, - (0, 1), let
p(y = 1) = 0.5, z ~ 1). p(b = 1 y = y, = 2) = ply, 1[ 2 O)), x = [b. 1 2 OJ].
We will drop the subscript in _L, for readability next. Throughout the next part, we use a key property of independence: [a, b|llcc = bille| a, alc..
As z is a standard normal random variable 1|z 2 O].L|zz| y. meaning we can write (y, 1/z > 0)) [z) because z is generated independently of y. Thus, as the distribution of b only depends on 1 2 0] and y due to the data generating process, it holds that (b, y. 1[z 2 LL. Then
(b.y, 1 2 2 12 = ylllz| b, 1 22 2 0J
yllz b. 1 1z 2 0]
y x
As x only depends on 1 2 0) and b. for readability, we define a= - 1[2 2 0j. Then P_ p_yy, a, b) =
ay,a,))
Pa (a))=()))(b a., where p(b = 1 y = y, a =a - ply, a) and x = [b.a].
We overload the notation for f: expanding x = [b, a]. we wet fxx f(b p(y = 1 x= = [b,a]).
We write f(b, a) for different values of b here,
((1,a) = ply x = ply = 1 1= a) p(b = 1,y
aa
p(b = a a=
ply 11==1pbb=1yy110
Ey -
= 1bpbb=1y==,,a21
0.5 0.5 EvE(11)) (p(b= y.a= 0))
p(1,a)
p(0 a) - a(L,a)'
f(0,a) ply 1 1|x x=
= -
p(b = 0,y = a=
p(b = | a=
ply = = b(y --a=
0.5(1 1b[1-pbb=1y=-1a-
0.5
1- p(1,a)
2 21-p(0,0) () p(L,a)
We let ply, 1) = 0.5 for y E (0,1), 0.1, and p(1, 0.9. Then, with a
1(1,1)
0.5
f(L.a) p(0, 1) + p(1,
0.5,
(18)
0.5 0.5+0.5 0.5
1- 11,11
1- 0.5
s(0,a) = 2 p(0 1) - 1))1))))a))
= 0.5,
(19)
0.5 0.5
and with O.
[(1,a) = 0(1,0)
0.9
= 0.9,
(20)
p(0, 0) + p(1,0)
0.1 - 0.9
1 p()
1 0.9
f().a) = 2 p(0, - p(1,0) -(100)))p(110
(21)
2 - 0.1 - 0.9
Thus, the distribution fxx | a = a changes with a meaning that f(x) which implies fxx) HLz as 1/2 a) is function of
23

---

Published as a conference paper at ICLR 2022
Note that f(x)),, then fxx & R, as f(x) is implication of joint independence. Any function T(x that achieves the same performance as (y f(x)) (by 3 and lemma 5), determines fxx). It follows that, T(x) of R, because
f(x) - T(x)=
So every RJ must perform worse than f(x), and consequently x, on PL. Finally, the indepen- dence xthpp 2 implies x & RR but y z x and so x E Re. In this example we constructed, x is the maximally blocking uncorrelating representation which means that it is optimal Rg on every Pte F. As R, is subset, any R. can at best match the performance of and we already showed that every r E RJ is worse than (x) and consequently x on PL. This completes the proof.
Lemma 5. Consider joint distribution ply,x) x) with binary Assume that p(x = y) has the same support for y E 1). Then, for any function T(x) such that y.lxx T(x). the function f(x) = ply = x) is meaasuaabee (x) determines f(x)).
Proof. We use the notion of sufficient statistics from estimation theory, which are defined for a family of distributions, to define the set of functions (x) for the joint distribution
Sufficient statistics in estimation theory. Consider family of distributions P (x) 0 E 52}.
Assume θ is discrete and that f2 is finite. A function T(x) is a sufficient statistic of family of distributions P if the conditional distribution pa T(x) = 1) does not vary with 0 for (almost) any value of L. A minimal sufficient statistic a sufficient statistic (xx such that for any sufficient
statistic T(X) T(x) = T(x') = M(x) (zz). Any bijective transform of M is also
minimal sufficient statistic.
The rest of the proof will follow from relying on theorem 6.12 from Lehmann and Casella (2006) which constructs a minimal sufficient statistic for a finite family of distributions P = {pr; i E
{0, K - 13} as
(x() P2(x) (u(x)) Po(x))
PK PX-1(x)) } Po(x) S
((x
Defining the family with conditionals. Now let the family P = {p) E {0,1]} where (x) p(x y = y) which are conditionals of the joint distribution p(x, y) were given in the theorem statement. Next, we show that the set of functions T(x) such that y - [xx) is exactly
the set of the sufficient statistics for the family {p(x y = y);y E {0,1))) 1}).
By definition of sufficiency where (, T'xx)= = f) does not vary with y for any value t,
Vt, P1 (x| [(x) t) px Txx) t) = Vt, p(x [x) y 1) p(x [xx) t,y 0),
where the last statement is equivalent to the conditional independence x y [(x))
Minimality of p(y = 1 x). By definition (x(x) = p(x y = y). As this family contains only two elements, the minimal sufficient statistic is
(() = (x) p(x 1
ply = 0) ply = 1
(x) p(x 0) ply = 1) 1-1--pppp 1
Thus, M(x is a bijective transformation of the function p(y = 1 x) (when ply = 1 1 E (0, which in turn implies that p(y = 1 x) is minimal sufficient statistic for the family P.
Conclusion. We showed that the set of functions that satisfy y_lx T(x) are sufficient statis- ties for the family P. In turn, because only sufficient statistics T(x) of the family P satisfy y T(x), it follows by definition that p(y = 1 x) is determined by every T(x), complet- ing the proof.
A.8 GAUSSIAN EXAMPLE OF THE INFORMATION CRITERION
Proposition 3. Consider the following family of distributions qu indexed by a R.
CyyE ~ 10, y ~ (0,
z z~ ~~Nl 1/2) x= ly + Ey. 2 / VTze.l
In this family. for any PD = q zz) where y-lhpp z. there exists a P'D quly z) such that
E KL KL (pbly x) poly x)] - 1n6 (x:)) > 0.
24

---

Published as a conference paper at ICLR 2022
Proof. (of proposition 3) First, write ay v/22 where 6 ~ W/ 1). Let v +c); this is normal variable with mean and variance 1. Then, write x ty+ ay tc where Ey. e are Gaussian random variables with joint distribution q ((c)) Therefore, qa (y,yx, x) is multivariate Gaussian distribution, with the following covariance matrix (over y, x2))
a
1 a² +2
a2 + -0
Σ
22,
a² +1
-a
The conditional mean and variance are:
E. y x = Σ1, 22.2x
22+211,,0x
11,
a² 22
a² +1
1
4a ly x) Σ1,1 Σ1, Σ1
a² +
a² +2
Rewrite the quantity in the theorem statement as a single expression:
(x) KL 1ga(y x) || @6(y xx] Lg (x:y)
-1Eg (x) KL Ma (y x) 9b x)] Eg (x) KL [ga(y x) [(y)).
(xy)) log (y x)
Bag (x.y) log 9a (y x)
(22)
(b x)
g(y)
(xy)) (log q - log (y x))
Expand (log g(y) - log (y x)) in terms of quantities that vary with and those that do not:
y²
(y E ly x))
log q(y y) log (y = y x)
log 2n
+ log 2TT 96 (y x)
2
20 (y x)
y 842 [1,1x 2
y²
2x
loog 2 + (b² + 2)
+ log V
2
2
6a +
y
y
1
[1 b|x
+ (b² + 2)
+ log
2
2
12 + 2
As only the first two terms vary with y.x, compute the expectations Ey. over these:
(y
2
y²
म [1, b]x
y²
y
8242
[1. b|x
E (x) xx)
+ (b² + 2)
Bgly))
(6² B)))) (x
2
2
2
2
1
((b2 +)y [1, b|))
2
+ (b² (
2)
y)
2(12 2)2
1
(1b2 + 2)y b1,
+ Egly)a v)
2
2(12 +
(122 + 2)y - - aby - be) be)²
1
2
2(b2 +2)
1
(6b + - ab)y Ey be)
En(yr)((s,)
2
2(b2 +2)
(112 + -bbby) ) + var(e,,) ) var(be)
+
2
2(b² + 2)
(6² + 1 ab) var (y) + var(ey, (2aa((e)
+
2
2 + 2)
(b2 + - ab)² + 1 b²
+
2
2(1² +
25

---

Published as a conference paper at ICLR 2022
(B² +1 ab)²
2(6² + 2)
The proof follows for any a such that
(b + + -b)² -1 + log 1/8-22 (b² + 1 ab)²
-1 -1
1
log (b² + 0
2(b² +2)
2(b² 2)
2
Let b+ 1+2 for some scalar v. Then, |2| > 1 (b² + 2) loo((b² + 2),
b
(62 +1
1²-1
log (b² + 22) = 2(h² 2) 2
log (12 +2) > 0.
2(b 2)
2
A.9 EXAMPLE WHERE Ptr (y x) x) PERFORM WORSE THAN p(y) UNDER Pte
In this section, we motivate nuisance-randomizatto and the uncorrelating property. Consider the following data generating process for a family JaER and fixed positive scalar a²:
y~ 10) z~~ Nay, 0.5)
[x] ~ NN 2, 1 0.5), x2 Ny z, 0.5)] (23)
Letting o² recovers the example in eq. (2). We keep o² for ease of readability. We first derive the performance of ply) = (y)) q(y) relative to ( x) under Ja-
Performance of q(y) relative to goly
x) under
Rewrite x
(1 (1-a) a) +y VoLa, (1 a xy + ca)]. where 61, C2 ~ (), Let the joint distribution over E1, be ((y)g(())((())) Then, (y.x)) is a multivariate Gaussian distribution with the following covariance matrix (over y, X2):
1
(1-a)
(1+ a)
Σ (1 a1 a)
(1-0)²
(1-a²)
(1+a) a) a2
(1+a)²
(1 (1+ +1
-(1-a²) (1- +
Σ12 = [1 - a, a,, u4,
h-1
2,2
d² + a)²
(1 2) t
1
Eg,ly x x Σn 2
[11 (1-a)),0(14 +1+
0² (1 + a)²
+ (1 a)² + 0²
0² (1+ + a)²
σ²
= ( Σn, Σn =
a² 02 + a) +
o2 + a)² + - a)² + 02
The performance of g(y) (recall Perf is negative KL) relative to qoly x) on ga x) can be written
(g KL Iga (y x) || (y x)] Be, KL (ga(y| x) || 4(y)] = E (x.y) (log g(y) log (o x))
Expand (log g(y) - log (y x)) in terms that vary with y, x and those that do not:
logg(y = y) - lo = y x)
y²
loog / + (y - Emp [y] x))
+ logg / 22TO2
qo (y x)
2
20 41
(y x)
(y
(1 (1-b))=2(1+6)))
y²
log V2T + (0= + b)2 +(1 - b)2 +02) (y (1+b)
o (1+b) -11-b)+++>
+ log v 2TO (y x)
2
20
82 (0²(1+ + (1 (1- 6)2 + 0²) o2(1+6)* -(
+ log
02 (y x)
2
20
As 0 (y terms:
x) does not vary with x,y, we will compute the expectation over the first two
(x.() (ogg(y) log (y x))
26

---

Published as a conference paper at ICLR 2022
+ (1-b)) + +02) (y 011+0) 111 - b), 02 (1++b)xx
y²
= (x) xx)
2
20
(y
02(1+b)2+(1 83 tra [11 - b), 02) (1)x) 2
1
y²)
=By()) (-x)
+ (02 + b)2 + +(1 (1- + B2 (x)
20²
(02 +1 b)2 o [11 b), b)]xx) 2
- +
+1 b)2 + (1 (1=b)² + +22) y)
20 + (1+ + (1 b)² + 02)2
2
(o²(1 + (1 b)² + + - 10 0 b1)x
+ + (xy))x lx)
202 - b) + b)2 0²)
(recall==[(0) x = (1
0017++11100 ((
(0((0((+00*(1100)++tt))))))) -0000 - 2
+ () -1--b-o((a+++d)) a1 b)² + 02)
+ (oyyso(e)aaee)
2
20²
+ (1-
flei C(a, b) =
(1 (0
+b+ - - a)(1 (1) b))
var b)y) + (am((1-1)a) b) var (1 b)c2)
(1
1 (C(a, var ()) var(e2)
0 +
1 + (0
(02 + b)² + (1- (1 +
b)a,b)2 - (a²)2
(o" + (1 b)² o²)
= 1/4 log (02(++02+((1008)
Note that log v o2 (y| x) log
With this,
a²
we have the ability to bound the performance of q(y) relative to (o x) under qa-
The conditional qı (y x) performs worse than qy = p(y) on 9-1- We show here that (y x) = a (y x) does worse than (y) p(y) on Ptc q-1. Letting = -1
0, - a) = and b- = = (1 (1+b) = 2, (1 - b) = 0 then
C(a,b) = 02(1 02 b)2 + (1- b)2 (a + b b ab) (1 a)(1 - b) = 0² * 22 + 0² = 502.
(o2(1+ + b)2 + (1 - bij² 2 + 0²
2 - (02)2
1 log
20" 6)2 + (1 (1- b)² + 02
2
σ2
25(02)2- (02)²
5 0
log
2 02 x2) 2
d2
12 log 5 > O.
5
Predictive failure when x is not uncorrelating Note that go | x) = Py | x) corresponds to the nuisance randomized conditional of the label given the covariates. We show that go x)
27

---

Published as a conference paper at ICLR 2022
performs worse than p(y) on infinitely many qa- Letting b 0 in C(a. helps compute the performance of g(y) over (y x) for every 9u:
b) = -((=)))=((-))(((0) -
-- - 0² pa-a(1---
Then, the performance of g(y) relative p. (y x)
(1 + (1- + +02
b(a, 20 202 b)2 (1 b) b)² +
log
2
g²
(a²(1 - 02)) + (a²)² 20 0²) )a) (o²)²
(2 (2*0² o² + 1)
log
* 0² 0² +1)
2
(letting a² = 2.}
a² log 2
5
20
This performance difference between gyy) and p. (y x) is positive for a 5r < -1. Thus, for every test distribution qu such that > 5 or -1, p. (y x) performs worse than marginal prediction.
B FURTHER EXPER IMENTAL DETAILS
Implementation details In section 5, the label y is a binary variable and, consequently, we use the Bernoulli likelihood in the predictive model and the weight model. In reweighting- tice, the estimate of the nuisance-randomizee distribution p. (y, z, x) x Pe-(Y)/B.. syptr 2, x) with estimated Per (y z) may have different marginal distribution p. (y) # Per (y). To ensure that ptr (y) = p_ (y), we weight our preliminary estimate P_ again as per(y)
in prac-
Py(y) P1 (y, 2, x).
In all the experiments, the distribution po(y (x)) is Bernoulli distribution parameterized by rh and a scaling parameter 0. In general, when the family of (y rx)( is unknown, learning pre- dictive models requires a parameterization p r- (x)). When the family known, for example when y is categorical, the parameters @ are not needed because the distribution p(y (x)) can be parameterized by the representation itself. For the critic model (f y, z, rx,x)) in the distillation step, we use two layer neural network with 16 hidden units and ReLU activations that takes as input y, ryy (x), and scalar representation s, (z); the critic model's parameters are o, eb. The repre- sentation s (z) is different in the different experiments and we give these details below.
In generative-NURDI we select models for p(x y. 2) by using the generative objective's value on heldout subset of the training data. For model selection, class-conditional Gaussian experiment, binary likelihood in the colored-MNIST experiment, and squared-loss reconstruction error in the Waterbirds and chest X-ray experiments. In reweighting- NURD, we use cross- fittinn procedure where the training data is split into K folds, and K models are trained: for each fold, we produce weights using a model trained and validated on the other K - folds. Hyperparameter selection for the distillation step is done using the distillation loss from eq. (6) evaluated on heldout validation subset of the dnisance-randomiie data from the first step.
Gaussian likelihood in the
In all experiments, we report results with the distillation step optimized with a fixed = and with 1 2 epochs worth of critic model updates per every representation update. In setting the hyperparameter X, a practitioner should choose the largest À such that optimization is still stable for different seeds and the validation loss is bounded away from that of marginal prediction. Next, we give details about each individual experiment.
Optimal linear uncorrelating representations in Class Conditional Gaussians. Here, we show that r rx) xi X2 is the best linear uncorrelating representation in terms of performance. First let the Gaussian noises in the two coordinates of x (given y, z) be e NN0 9) and E2 ~ ~0, 0.01) respectively. Define Fa (x) ux, + eX2 = (u - r)y + (v u)z + uei + vEg. We will show that qo (z) Fx,v((), = 1) F 4o (x), = 0) when u * v and u of -v. First, (z(z, zz Fn (x) y y) is a bivariate Gaussian with the following covariance matrix:
1 Eny (vV -u) u) (w - - 9u² 0.012²
(v-u)
1
E 1 u, Ey.2.
(w - u)² + 9n 0.
28

---

Published as a conference paper at ICLR 2022
The conditional mean is:
By [z] (x) y. = Elz + (r
Blr ly =])
= [R[z] + E1,2> (r (u + u)y)
(u (w-w)( ur - (u ww)+
(w- u)² + 2 0.01
which is independent of y if and only if u + = 0 or u - v = 0. The conditional vari- ance does not change with y because it is determined by Σy which does not change with y- Thus (z y)).) = 9a(z Fu. if and only if u = u or u = -.. When u = v, Tu,o 2uy + noise and ru. meaning that Tu,u helps predict y. In contrast, when u = -. Tu,e 2rz + noise and ro = (o (r-..)) qo(y), meaning that go /r-,,))
has the same performance as the marginal. However, for all u * 0, Tu,u = 1.1,1 is a bijective transform of r1.1 and, therefore, qo (y Fu,u (x)) = go rr,, (x)). Thus, within the set of linear uncorrelating representations, T1, is the best because its performance dominates all others on every
Pte F.
Implementation details for Class Conditional Gaussians. In reweighting-NuRRR the model for Ptr(y] z) is Bernoulli distribution parameterized by neural network with 1 hidden layer with 16 units and ReLU activations. In Nenerative-NRRD. the model for p(x y.z) is an isotropic Gaussian whose mean and covariance are parameterized with neural network with one layer with 16 units and ReLU activations. We use 5 cross-fitting folds in estimating the weights in reweighting-Nute We use weighted sampling with replacement in computing the distillation objective.
In the distillation step in both reweighting and generative-Nute NURRD, the representation r- (x) is a neural network with one hidden layer with 16 units and ReLU activations. The critic model (f( y, z, F.y (x)) consists of a neural network with 2 hidden layers with 16 units each and ReLU activations that takes as input y; (x)),, and scalar representation (z) which is again neural network with a single hidden layer of 16 units and ReLU activations.
We use cross entropy to train (y | z). ( I (x)),, and ps(e | y. r.z,),(x)) using the Adam (Kingma and Ba, 2014) optimizer with a learning rate of 10-2 We optimized the model for (y z) for 100 epochs and the model for per (x y, y,z) for 300 epochs. We ran the distillation step for 150 epochs with the Adam optimizer with the default learning rate. We use batch size of 1000 in both stages of NURD. We run the distillation step with fixed >= 1 and two epoch's 's worth of gradient steps (16) for the critic model for each gradient step of the predictive model and the rep- resentation. In this experiment, we do not re-initializ á th after predictive model update.
Implementation details for CoNOTEd-MNIST. For reweighting-NURD. to use the same architec- ture for the representation r.y(x) and for Pt- (y z). we construct the nuisance as a 28 28 image with each pixel being equal to the most intense pixel in the original image. In Aenerative-NuRD. we use a PixelCNN model for p(x y, z) with 10 masked convolutional layers each with 64 filters. The model was trained using a Bernoulli likelihood with the Adam optimizer and a fixed learning rate of 10-3 and batch size 128. We parameterize multiple models in this experiment with the follow- ing neural network: 4 convolutional layers (with 32, 64, 128, 256 channels respectively) with ReLU activations followed by a fully connected linear layer into a single unit. Both r., (x), 8 (z) are pa- rameterized by this network. Both pr-(y z) in reweighting- NURD and P x) for ERM are Bernoulli distributions parameterized by the network described above. We use 5 cross- fottss- folds in estimating the weights in reweighting-Note
For the critic model pb(E y, z, (x)) in the distillation step. we use a two wwo-hidden-layee neural network with 16 hidden units and ReLU activations that takes as input y, (x), and the scalar representation Su (z); the parameters @ contain vú and the parameters for the two hidden-layer neural network. The predictive model po (y (x)) is a Bernoulli distribution parameterized by r,,xx) multiplied by scalar 0.
We use cross entropy to train P | z), ( (x)), and (e I y, v, (x)) using the Adam (Kingma and Ba, 2014) optimizer with a learning rate of 10-3 We optimized the model for per (y z) for 20 epochs and ran the distillation step for 20 epochs with the Adam optimizer with the default learning rate. We use a batch size of 300 in both stages of NURD. We run the distillation step with fixed X = 1 and one epoch's worth of gradient steps (14) for the critic model for each gradient step of the predictive model and the representation. In this experiment, we do not re-initialize o vi after predictive model update.
29

---

Published as a conference paper at ICLR 2022
Implementation details for the Waterbirds experiment. For generativeNNuRD, we VQ- VAE (Razavi et al., 2019) to model Ptr(x 2). For multiple latent sizes and channels in the encoder and the decoder, we saw that the resulting generated images were insufficient to build clas- sifiers that predict better than chance on real data. This may be because of the small training dataset that consists of only 3000 samples. The model for p_ (y (x)( is two feedforward layers stacked on top of the representation T- (x). The model ( | z) in Soweighting-NuRe
is the same model as p. (y rx)(x)) as a function of x. The model for po(E y, 2, Try (x)) consists of neural network with two feedforward layers that takes as input y." (x), and a representation s& (z). Both ry and se are Resnet-18 18 models initialized with weights pretrained on Imagenet; the parameters & contain e and the parameters for the two hidden-layer neural network. The model in ERM for ptr x) uses the same architecture as p. (y rr, (x)) as a function of x. We use cross-fitting folds in estimating the weights in reweighting-NuRI.
We use binary cross entropy as the loss in training Per (y 2), ( Fry (x)), and pa(l y. z, Tyy (x)))
using the Adam (Kingma and Ba, 2014) optimizer with a learning rate of 10-3, We optimized the model for (y z)) for 10 epochs and ran the distillation step for 5 epochs with the Adam optimizer with the default learning rate for all parameters except y, which parameterizes the representation ry; for 7 used 0.0005. The predictive model, the critic model, and the weight model are all optimized with a weight decay of 0.01. We use a batch size of 300 for both stages of NURD. We run the distillation step with a fixed À and two epoch's 's worth of gradient steps (16) for the critic model for each gradient step of the predictive model and the representation. To prevent the critic model from overfitting, we re-initialize Φ, & after every gradient step of the predictive model.
Implementation details for the chest X-ray experiment. To help with generative modeling, when creating the dataset, we remove X-ray samples from MIMIC that had all white or all black borders. We use VQ-VAE2 (Razavi et al., 2019) to model p(x y. 2) using code from here both train and sample. The encoder takes the lung patch as input, and the decoder takes the quantized embeddings and the non-lung patch as input. VQ-VAE2 is hierarchical with a top latent code and bottom latent code which are both vector-quantized and fed into the decoder to reconstruct the image. Both latents consist of 8 x 8 embeddings each of dimension 64. The VQ-VAE is trained for 200 epochs with Adam (Kingma and Ba, 2014) with batch size of 256 and dropout rate of 0.1. Generating samples from the VQ-2AE2 involves sampling the top latent code conditioned on the label, followed by sampling the bottom latent code conditioned on the label and the top latent code, and passing both latent codes to the decoder. To generate from the latent codes, we build a PixelSNAIL. generate the top latent code given the label and PixelCNN generate the bottom latent code given the label and the top latent code. These models have 5 residual layers with 128 convolutional channels. All other details were default in here. We train these models for 450 epochs with batch size of 256 with a learning rate of 5 x 10-5.
For reweighting-Nut representation (x). The model in ERM for Pt- x) uses the same architecture as 9. (y rr (x)) as a function of x. Next we use a single architecture to parameterize multiple parts in this experi- ment: 3 convolutional layers (each 64 channels) each followed by batch norm, and dropout with rate of 0.5 and followed by a linear fully-connected layer into a single unit. We parameterize the two representations r. (x), (z) with this network. To build ( y, 3,2,r, (x),,, we stack two feed- forward layers of 16 hidden units with ReL.U activations on top of a concatenation of y, r, (x). and the scalar representation s. (z) as described above; the parameters contain &b and the parameters for the two hidden-layer neural network. We use 5 cross-fitting folds in estimating the weights in reweighting-NuRE
the model p_ (y (x)) is two feedforward layers stacked on top of the
We use binary cross entropy the loss in training Per 22), po [r, [x)),. and Po(f y T,, (x))) using the Adam (Kingma and Ba, 2014) optimizer with a learning rate of 10 We use a batch size of 1000 for both stages of NURD. We optimized the model for Py- z) for 150 epochs and ran the distillation step for 100 epochs with the Adam optimizer with the default learning rate. Only the optimization for per 2) has a weight decay of le 2. We run the distillation step with a fixed X = 1 and two epoch's 's worth of gradient steps (20) for the critic model for each gradient step of the predictive model and the representation. To prevent the critic model from overfitting, we re- inittiallize é se after every gradient step of the predictive model.
30

---

Published as a conference paper at ICLR 2022
NuRD vs. ERM on classifying Waterbirds
NuRD vs. ERM on Pneumonia detection
70
Avg. accuracy
Avg. accuracy
80
60
75
50
70
NuRD
NuRD
ERM
40
ERM
0.5
0.7
0.9
0.5
0.7
0.9
Test o (strength of the spurious correlation)
Test (strength of the spurious correlation)
Figure 4: Plots of average accuracy vs. test for classifying Waterbirds and Pneumonia. A larger implies larger difference between the nuisance-label relationship the test data used for evaluation and the training data, which has a = 0.1. quisance-randoming data corresponds to p = 0.5. Unlike NURD, ERM's performance quickly degrades as the difference between the train and the test distributions increases.
B.I ADDITIONAL EXPERIMENTS
Excluding the boundary from the images (covariates) does not improve ERM in general. In both Waterbirds and chest X-rays, we use the casy-to-acquire border as a nuisance in NURD. Models trained on the central (non-border) regions of the image exploit the nuisances in the center and consequently fail to generalize when the nuisance-label relationship changes. In fact, classifiers produced by ERM M border-less images do not generalize well to the test data, producing test accuracies of 39 + 0.5% on chest X-rays and 65 + 2.3% on Waterbirds averaged over 10 seeds. However, as independence properties that hold for the border also hold for nuisances in the central region that are determined by the border, NURD can use the border to control for certain nuisances in the center of the image.
MudelModels
Additional experiments with NURD, We evaluate reweighting Boiighting-uRDDDD NURD further in the following ways:
Run NURD on data from the training data distribution defined in section and evaluate on data from test distributions Pte with different nuisance-labe relationships.
2. Train NURD with different-sized borders as nuisances.
3. Train NURD without a nuisance where the training and the test data have the same nuisance-label relationship; we implement this by setting the nuisance = 0 wherever it is passed as input in the weight model or critic model.
4. Run the distillation step with different X.
Different test distributions. For this experiment, compute the test accuracies of the models trained in the experiments in section 5 on data with different nuisance-labe relationships. For both classifying Waterbirds and Pneumonia, scalar parameter p controls nuisance- label relationships in the data generating process. In waterbirds, = p(y waterbird background = land) p(y = landbird background water). In chest X-rays, p corresponds the fraction of Pneumonia cases that come from CheXpert and normal cases that come from MIMIC in the data; in this task, hospital differences are one source of iriaance-nddueed spurious correlations. In both tasks, P 0.1 in the training data; as test increases, the quisance-labed relationship changes and becomes more different from the training data. We plot the average and standard error of accuracies aggregated over 10 seeds for different test p E (0.5, 0.7, 0.9) in fig. 4.
Nuisance specification with different borders. For Waterbirds, we ran NuRD with the pixels outside the central 168x 168 patch (a 56 pixel border) as the nuisance. Averaged over 10 seeds, seweighting-NuRe produced model with 81% test accuracy which is similar to the accuracy achieved by NURD using a 28-pixel border as the nuisance. In comparison, ERM achieves an accuracy of 66%.
NURD without nuisance and no nuisance-induced spurious correlations in classifying Wa- terbirds. We performed an additional experiment on classifying Waterbirds where NURD is given constant nuisance which is equivalent not using the nuisance. We generated training and test
31

---

Published as a conference paper at ICLR 2022
data with independence between the nuisance and the label; the nuisance-label relationship does not change between training and test. Averaged over 10 seeds, ERM achieved a test accuracy of 89 + 0.4% and NuRD achieved test accuracy of 88 = 1%.
Reweighting-NuRD with different X. Large Äs may make optimization unstable by penalizing even small violations of joint independence. Such instabilities can lead NURD to build predic- tive models that do not do hetter than marginal prediction, resulting in large distillation loss (log- likelihood + information loss) on the validation subset of the training data. However, small À may result in NURD learning non-uncorrelating representations which can also perform worse than chance.
We ran NURD on the waterbirds and class-conditional Gaussians experiments with À (instead of A like in section 5) and found that, on few seeds, NURD produces models with close to 50% accuracy (which the same as majority prediction) or large information loss or both. Excluding seeds with large validation loss, reweighting-Nute NuRD achieves average test accuracy of 76% on waterbirds and 61% on cladsicond Gaussians. Annealing the during training could help stabilize optimization.
In setting the hyperparameter À in general, a practitioner should choose the largest À such that optimization still stable over different seeds and the validation loss bounded away from that of predicting without any features.
32


Reviews:
[{"review_id": "khs7QkYENQ", "reviewer": "Reviewer_Vct3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "I find the mathematical formalization to be a bit too much and it makes it harder to understand the core proposed ideas.\n\nEq 1 is a contribution of this work and it presents the problem of spurious correlations broken down into the contributions of a nuisance variable z conditioned on the label and the input features conditioned on both the label and z.\n\nFrom what I understand of the NuRD method you estimate the nuisance variable using a neural network and then use this to reweight the training samples based on this. How exactly do you learn p_{tr}(z) and p_{tr}(y | z) ? As this is a core contribution of this paper please explain this very clearly.\n\nI find the evaluations lacking. Only one type of spuriously correlated feature was used. Why not evaluate on the CelebA evaluation in Sagawa et al. (2019)? And for chest X-rays why not use the same experimental setup as (Zech et al., 2018) instead of correlating the border? Note: A version of the Zech experiment was done using all publicly available data in https://arxiv.org/abs/1910.00199\n\nAlso in the non-mnist experiments, only a single size border thickness was used. In order to explore the utility of this approach please run experiments with multiple types of spurious correlations. Possibly also varying the amount of correlation. Please show an experiment where this method doesn't work so that we may understand its limitations.\n", "summary_of_the_paper": "The paper presents a formalization of spuriously correlated distractor features into a family of distributions. One goal here is to train a model to be robust against these distribution shifts. The paper presents results comparing against ERM and argues that it is unique in its problem formulation compared to other methods that make assumptions about environments or specific spurious correlations.", "summary_of_the_review": "The proposed formalization and method are interesting. However, the presentation is not focused on the core contribution, the method is lacking detail, and the experiments do not characterize the method well.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "hqfhFyqmPVF", "reviewer": "Reviewer_frmn", "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": "In general, this paper studies a long-standing problem of spurious association in fitting a predictive function. The intuition of NuRD is clear. If we have data from a distribution where the nuisance variables z are randomized and do not cause the label y, the prediction will not suffer from the spurious association. Another strength of this paper is that it studies a variety of representative data sets. \n\nThe following are my major questions/comments about this paper:\n\n1. A major concern is about the setting of the data. In the literature of this problem (e.g., Invariant causal prediction (Peters+, 2016), (Invariant Risk Minimization, Arjovsky+, 2019)), the input is O=(x,z) and we do not know which are x and which are z. For example, with a image of cow on the grass, we do not know which pixels are of cow (x) and which pixels are of grass (z). But this paper assumes we know this separation a priori. Then why not just exclude all nuisance z from the predictor and only use x as the input? \n\n2. The paper lacks sufficient clarity (e.g, Sec. 2.). A main reason might be that the paper focus on the problem of spurious association, which is a causal problem. But it does not use a causal language and notations (e.g., do-calculus / potential outcome). \n\n3. The simulation results might need further check. In the color MNIST example, the NuRD reaches the oracle accuracy 75%. However, this is an ideal upper bound of accuracy. Referring to [IRM, Arjovsky+, 2019], the prediction accuracy on the grey-scale images (without any spurious association) is 73%. The author may need to check/explain why NuRD outperforms this practical oracle.\n\n4. NuRD does not compare with any methods that improves over classic ERM. \n\n5. Some implementation details are not presented, such as how to separate x and z, how to choose the hyperparameter $\\lambda$. \n\n\n\n\n\n\n\n", "summary_of_the_paper": "The paper studies the prediction problem under spurious association. The idea is to construct a \"nuisance randomized distribution\" based on the observed distribution. The nuisance randomized distribution is constructed by reweighting the observed data. Then the samples from the nuisance randomized distribution are used to learn a predictive function. ", "summary_of_the_review": "In sum, the reviewer think the setting of this paper, the writing clarity, and the simulation study may need further revision and improvement before getting published. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YBAcirS3zs8", "reviewer": "Reviewer_S6jJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "**Strengths**\n\nNuisance-induced spurious correlations are a major issue in many applications, since they lead to models that cannot generalize well to unseen data.\n\nThis paper introduces a novel idea to address this issue (to the best of my knowledge) which is theoretically principled and well motivated. As such, I found the paper a very interesting read.\n\nThe proposed method allows to remove many of the strong assumptions that limit the application of competing ones, for example it can be used in high dimensional image classification tasks. \nThe experimental results show that despite highly different training and test distributions, NuRD is able to greatly reduce the influence of nuisance variables.\n\n\n**Main weaknesses**\n\nWhile the presented theory and results are very compelling, there are two main issues with the exposition that would greatly limit its impact in the ICLR community:\n1. The theoretical exposition is unnecessarily abstract and complex.\n    1. Section 2 and 3 are much heavier from the theoretical point of view than the average ICLR paper. Given the nature of the presented method, this is of course not a problem per se. However, to have an impact on the community these sections should be made much more accessible. What I find that is lacking the most, is intuition on what is being presented. One way to convey it, is to use throughout these sections a running example, e.g. the very clear cows vs penguin example of the introduction. Most of the propositions, theorems and lemmas would be way easier to understand if the authors provided intuition of what they mean if we consider x being the input image and z being its background in the running example.\n    2. This is particularly important when presenting NuRD in section 3. The authors should give way more intuition to the reader using an actual example on what the randomization and distillation step do. For example, what is the practical difference of using (6) vs (7) in the cows vs penguin classification task.\n    3. The introduction in page 2 is very hard to understand without more concrete examples. For being the introduction of the paper, the discussion is way too abstract for readers unfamiliar with the field.\n\n2. From the main text, there is close to no information on how this method would be implemented in practice.    \n    1. After reading the main text I had an understanding of the theory, but I had no clue on what I would need to do to implement this, as well as the level of additional complexity this method introduces if I were to apply it in a standard neural network classification task. Many of the info in the appendix should be moved or at least better addressed in the main text. One possible way could be presenting how the reweighting/generative NuRD and critic model would look in a standard neural network classification task, and what optimizing (6) and (7) imply.\n    2. Without the code being available, I doubt many people would know how to implement this method.\n\n\nSome of the above issues might be due to the lack of space in the main text. Some ideas to gain some extra space to implement the above changes could be:\n* moving lemma 1, theorem 1 and the discussion around them to the appendix. The key component for NuRD is theorem 2: lemma 1 and theorem 1 can be just seen as necessary steps to prove theorem 2. This would also allow to simplify the theoretical part of the paper.\n* mostly focusing in the main text on reweighting NuRD, which seems more broadly applicable (see below). Most of the generative NuRD discussion could be covered in the appendix.  \n* moving parts of the related work to the appendix\n\n**Comments/question on experiments**\n\n1. For the more realistic experiments in 5.3 and 5.4 generative NuRD seems to struggle. Is generative NuRD feasible in practice for image classification tasks? Generating good high dimensional images remains a challenge, what if the estimate of nuisance randomized distribution is poor?\n2. In real world applications assuming that the object to classify in the images are centered might not be realistic. What happens if in many images from the dataset the background z contains useful info for the classification task? How robust is NuRD to this? This could happen for example with pathologies that are in the outer region of the lungs in chest x-ray images.\n3. How well does NuRD scale w.r.t standard classification? What are in the experiments the training times of ERM, generative NuRD and reweighting NuRD?\n4. Section 5.1. is missing a reference to table 1\n\n\n", "summary_of_the_paper": "This paper introduces a method to reduce the influence of nuisance variables in predictive tasks. This is done by fitting a nuisance-randomized distribution to the data, and finding a data representation under this distribution that makes the nuisance and the label independent. Experiments show that by using a classifier learned on this representation, the method is able to improve classification performances by limiting the impact of nuisance variables.\n", "summary_of_the_review": "This paper presents a very interesting idea with competitive performances to a very relevant problem in real world applications.\nThe exposition of the paper however needs to be greatly simplified by giving a lot more intuition and implementation details. Due to this, in this current state I am afraid that the paper would not have the impact it deserves in the ICLR community.\n\n\n\n=====================\n\n*Updated review after the author's rebuttal and revision of the paper*:\nI have read the paper once more, the authors did a great job improving the exposition of the paper, which is now way easier to understand and accessible to a wider audience. I have therefore increased my score to an accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "RnEw47tC_ED", "reviewer": "Reviewer_mFHY", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strength\n1. A great list of examples is provided that describe how spurious correlations arise.\n2. Mathematic details and proofs are provided in full details for all theorems etc.\n3. NuRD works extremely well on problems where nuisance can be clearly defined and its relationship with label is strong and changing a lot between train and test data.\n\nWeaknesses\n1. Too much technical details in the Introduction without adequate context, e.g. too early for “We show that the distribution of the label given the covariates under the nuisance-randomized distribution is minimax optimal for the nuisance-varying family, if the label and the nuisance are independent given the covariates” and all the text after. Hard to get intuition when the second half of Intro is just listing math properties without some high level descriptions of how to attain those properties.\n2. Same for the Methods section, which repeats much of the Introduction (with some math interleaved). Would be nice to have a concrete example to illustrate ALL steps involved in NuRD and then go into the theorems, definitions, etc. As currently written, it is hard to get an integrated picture when Section 2 is just a bunch of math definitions.\n3. Perhaps consider putting related work up front to provide context and state contributions with respect to existing works. Any reasons why domain adaptation and optimal transport are not mentioned?\n4. Would the choice of nuisance distribution have a major effect on the results? How sensitive are results to the many parameter choices?\n5. The experiments with Gaussian, color MNIST, and waterbirds seem quite extreme in terms of changing relationship between nuisance and labels. Would we see same results for less extreme scenarios? A huge problem in medicine is using genomic or imaging data to predict neurodegenerative disease status with age being nuisance. Would be great if NuRD works on this problem. \n6. NuRD does not perform as well as ERM when the nuisance-label relationships do not change mainly because NuRD does not use the nuisance variables to its advantage. Is it possible to create experiments where there are no nuisance variables that can be used? Would like to see if NuRD performs as least on par with ERM in that case, i.e. not underperforming when there is no change in nuisance-label relationships.\n7. For the real word problem examined in this paper, pneumonia cases are upsampled, thus might not provide realistic classification assessment, i.e. samples would be correlated. Would like to see results with controls downsampled. \n8. Quantitative comparison with existing DL methods, e.g. if have z, can find representation of x that predicts y well but not z. \n", "summary_of_the_paper": "This paper proposes Nuisance-Randomized Distillation (NuRD) to deal with spurious correlations induced by a changing relationship between labels and nuisance variables, where covariates are also correlated with nuisance. The paper first introduces the nuisance-randomized distribution under which labels and nuisance are independent, and shows that NuRD can find uncorrelating representations of covariates where labels remain independent from nuisance when conditioned on these covariate representations. The paper further describes how to find the optimal representation that has maximum information with labels. NuRD is evaluated on various classification tasks involving class-conditional Gaussians, colored MNIST images, waterbirds, and chest X-rays.", "summary_of_the_review": "Overall, the paper is very rigorous in proving the mathematical guarantees of the proposed method, but the method itself is not performing that well for harder real word problems. My evaluation is based on the simplicity of experiments for scenarios where NURD worked. For the harder case (even when classification bias is introduced to the data), NURD does not perform particularly well.\n\nPost-revision Summary\nThe authors have nicely addressed most of my comments. The text is still heavy on math over intuition, but improved. I am thus increasing my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0THK9yicEZ0", "reviewer": "Reviewer_2F1e", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Strengths:\n\nThe paper proposes a method to solve an important and challenging problem of nuisance-label correlations in supervised learning. The paper provides a good theoretical analysis of the proposed method and interesting simulation results to demonstrate the applications of the proposed method. The related work section is well-structured and comprehensive including previous works using different ideas. Overall, the paper is well written and easy to follow.\n\nWeaknesses and Suggestions:\n\nSections 2 and 3 only introduce the major component of the proposed method. However, it is immediately clear to me what kind of datasets we can apply the proposed method. Importantly, how do we extract the nuisance variable z out a given arbitrary labelled dataset?\n\nEquation (1) is better to be introduced using some examples of x,y and z in the second paragraph of the introduction. For an image, why the data generating process of x is conditional on z and y. I thought both x and z are some pixels on an image.\n\nAt least in statistics, \"predictor\" often refers to a variable a model uses to predict the target. Here \"predictor\" refers to the predictive model or distribution itself. I encourage the authors to reconsider their terminology.  \n\nMeasuring the performance of a predictor by KL divergence seems more reasonable for a classification task. If the label is continuous, then comparing KL divergence does not tell us exactly which model is closer to the ground truth.  \n\nIn proposition 3, in the distribution p_\\indep, x is generated by y and z. Then why we would have y is independent of z conditional on x? I am thinking of x is a collider of y and z in a directed graphical model. \n\nIs it true that for any p_\\indep, we can find an uncorrelated set of representations R(p_\\indep)?\n\nWhy the first equality in Equation (4)?\n\nThe high-level idea of nuisance randomization distillation (NURD) is clear in the paper. I wonder if there are any distributions p(x,y,z) for which we do not want to use distillation. In order words, if there are any cases that we should optimize Equation (6) or (7) without the second mutual information term. For example, if y is a function of x, maximizing the expected log-likelihood should give the best predictor.\n\nYou mentioned quite some related methods in Section 4. Is it possible to compare with some of them, or could you provide some suggestions on when to use these related methods and when we should use  NURD?\n\n\n", "summary_of_the_paper": "This paper develops a representation method called nuisance-randomized distillation (NuRD) for building predictive models with better generalizability (i.e. testing performance) without using spurious nuisance-label correlations in observed data. NuRD breaks the nuisance-label dependence and finds the most informative representation to predict the label for every distribution in the nuisance -varying family which is a set of distributions that differ only in the nuisance-label relationship. The NURD method is evaluated on several application datasets; it produces accurate predictive models on the datasets with strong spurious correlations between the label and some nuisance variables.", "summary_of_the_review": "This paper proposes an interesting idea for an important problem. It provides a good theoretical and experimental analysis of the proposed method. I recommend accepting this paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies how to build predictive models that are robust to nuisance-induced spurious correlations present in the data.  It introduces nuisance-randomized distillation (NuRD), constructed by reweighting the observed data, to break the nuisance-label dependence and find the most informative representation to predict the label. Experiments on several datasets show that by using a classifier learned on this representation, NuRD is able to improve the classification performance by limiting the impact of nuisance variables. The main concerns were about the presentation and organization of the paper, which was heavily focused on the theoretical justifications but fell short in explaining the intuitions and implementation details. The revision and rebuttal have addressed some of these concerns and improved the overall exposition of the paper, based on which two reviewers raised their scores to 8. While there is still room to further improve the paper by providing more detailed discussions about the proposed algorithms, the AC considers the paper ready for publication under its current form.
Decision: Accept (Poster)