Measure the Predictive Heterogeneity
Authors: Jiashuo Liu
Jiayun Wu
Renjie Pi
Renzhe Xu
Xingxuan Zhang
Bo Li
Peng Cui
Venue: ICLR 2023 poster (2023)
Keywords: data heterogeneity
predictive information
predictive heterogeneity

Abstract: As an intrinsic and fundamental property of big data, data heterogeneity exists in a variety of real-world applications, such as in agriculture, sociology, health care, etc. For machine learning algorithms, the ignorance of data heterogeneity will significantly hurt the generalization performance and the algorithmic fairness, since the prediction mechanisms among different sub-populations are likely to differ. In this work, we focus on the data heterogeneity that affects the prediction of machine learning models, and first formalize the Predictive Heterogeneity, which takes into account the model capacity and computational constraints. We prove that it can be reliably estimated from finite data with PAC bounds even in high dimensions. Additionally, we propose the Information Maximization (IM) algorithm, a bi-level optimization algorithm, to explore the predictive heterogeneity of data. Empirically, the explored predictive heterogeneity provides insights for sub-population divisions in agriculture, sociology, and object recognition, and leveraging such heterogeneity benefits the out-of-distribution generalization performance.

Published as a conference paper at ICLR 2023
MEASURING THE PREDICTIVE HETEROGENEITY
Jiashuo Liu', Jiayun Wu', Renjie Pit, Renzhe Xu', Xingxuan Zhang', Bo Leng Cuite TTinghua University, 'Hong Kong University of Science and Technology
(liujiashuo77, jiayun. wu wou.work @egmai1 com, #pi@connect.us ust. .kk xr2199721@gmail com, Langxuanzhanggehotaall.. libo@sem. tainghua hewpiinnuaauuuu.com edu.cn,
ABSTRACT
As an intrinsic and fundamental property of big data, data heterogeneity exists in a variety of real-world applications, such as in agriculture, sociology, health care, etc. For machine learning algorithms, the ignorance of data heterogeneity will sig- nificantly hurt the generalization performance and the algorithmic fairness, since the prediction mechanisms among different sub- populations this work, we focus on the data heterogeneity that affects the prediction of machine learning models, and first formalize the Predictive Heterogeneity, which takes into account the model capacity and computational constraints. We prove that can be reliably estimated from finite data with PAC bounds even in high dimensions. Additionally, we propose the Information Maximization (IM) algorithm, a bi-level optimization algorithm, to explore the predictive heterogeneity of data. Empiri- cally, the explored predictive heterogeneity provides insights for sub-population divisions in agriculture, sociology, and object recognition, and leveraging such heterogeneity benefits the out-of -i-ttribiittttionnnn generalization performance.
likely differ. In
INTRODUCTION
Big data bring great opportunities modern society and promote the development of machine learn- ing, facilitating human life within a wide variety of areas, such as the digital economy. healthcare, scientific discoveries. Along with the progress, the intrinsic heterogeneity of big data introduces new challenges to machine learning systems and data scientists (Fan et al. 2014 He, (017). In general, data heterogeneity, as fundamental property of big data, refers to any diversity inside data, including the diversity of data sources, data generation mechanisms, sub-populations, data struc- tures, etc. When not properly treated, data heterogeneity could bring pitfalls to machine learning systems, especially in high-stake applications, such as precision medicine, autonomous driving, and financial risk management (Dzobo et al. 2018 Breitenstein et al. 2020 Challen et al. 2019), lead- ing to poor out-of duttof-distributton generalization performances and some fairness issues. For example, in supervised learning tasks where machine learning models learn from data to predict the target variable with given covariates, when the whole dataset consists of multiple uub-ooplattions with shifts or different prediction mechanisms, traditional machine learning algorithms will mainly focus on the majority but ignore the minority. It will hurt the generalization ability and compromise the algorithmic fairness, as is shown in (Kearns et al. 2018 Sagawa et al.) 2019 Duchi & Namkoong. 2021. Another well-known example Simpson's paradox, which brings false discoveries to the social research (Wagner 1982 Hernán et al., 2011).
Despite its widespread existence, due its complexity, data heterogeneity has converged to a uniform formulation so far, and has different meanings among different fields. Li & Reynolds (1995) define the heterogeneity in ecology based on the system property and complexity or vari- ability. [ooenbaum (2005) views the uncertainty of the potential outcome as unit heterogeneity in observational studies in economics. More recently, in machine learning, several works of causal learning (Peters et al.) 2016 Arjovsky al. 2019 Koyama & Yamaguchi 2020 Liu et al.] 2021 Creager et al.] 2221)) and robust learning (Sagawa et al. [0199] Liu et al., 2022 leverage hetero- geneous data from multiple environments to improve the
generalization ability. However, previous works have not provided a precise definition or sound quantification. In this
"Corresponding Author.

---

Published as a conference paper at ICLR 2023
work, from the perspective of prediction power, propose the predictive heterogeneity, new rype of data heterogeneity.
From the machine learning perspective, the main concern is the possible negative effects of data heterogeneity on making predictions. Therefore, given the complexity of data heterogeneity, in this work, we focus the data heterogeneity that affects the prediction of machine learning mod- els, which could facilitate the building of machine learning systems, and we name it the predictive heterogeneity. We raise the precise definition of predictive heterogeneity, which is quantified as the maximal additional predictive information that can be gained by dividing the whole data distribution into sub uub-populations. The new measure takes into account the model capacity and computational constraints, and can be reliably estimated from finite samples even in high dimensions with PAC bounds. We theoretically analyze its properties and examine under typical cases of data hetero- geneity (Fan et al., 2014). Additionally, we design the information maximization (IM) algorithm to empirically explore the predictive heterogeneity inside data. Empirically, we find the explored heterogeneity is explainable and it provides insights for sub-population divisions in many fields, including agriculture, sociology, and object recognition. And the explored sub-populations could be leveraged enhance the cuttof-ddisrrbuttion generalization performances of machine learning models, which is verified with both simulated and real-world data.
PRELIMINARIES ON MUTUAL INFORMATION AND PREDICTIVE V-INFORMATION
In this section, we briefly introduce the mutual information and predictive V-information (Xu et al, 2020 which are the preliminaries of our proposed predictive heterogeneity.
Notations. For a probability triple (S, F, IP), define random variables X S + X' and Y S + y where X' is the covariate space and Y' is the target space. Accordingly. x E X' denotes the covariates, and y E y denotes the target. Denote the set of random categorical variables as C = {C S - N|supp(C)) C') finite). Additionally, (x)), P(Y)) denote the set of all probability measures over the Borel algebra on the spaces X respectively. (.) denotes the Shannon entropy of a random variable, and (-I-) denotes the conditional entropy of two random variables.
In information theory, the mutual information of two random variables X, Y measures the depen- dence between the two variables, which quantifies the reduction of entropy for one variable when observing the other:
T(X;Y) = MYY HY|X))
It is known that the mutual information is associated with the predictability of Y (Cover Thomas & Thomas Joy 1991). While the standard definition of mutual information unrealistically assumes the unbounded computational capacity of the predictor, rendering it hard to estimate especially in high dimensions. To mitigate this problem, Xu et al. (2020) raise the predictive V-information under realistic computational constraints, where the predictor is only allowed to use models in the predictive family V to predict the target variable Y. Definition 1 (Predictive Family (Xu et al. (0200). Ler Ω {f XU {0} P())). We say that V V Q is a predictive family if satishes:
vf E V, vP E range(f), Af' s.t. Yo E X, [x] = P. (' P.
(2)
A predictive family contains all predictive models that are allowed use, which forms computa- tional or statistical constraints. The additional condition in Equation 2 means that the predictor can always ignore the input covariates (z) if chooses (only use ()).
Definition (Predictive V- V-information (Xu et al. (2020). Let X, Y be two random variables taking values in X x y and V be a predictive family. The predictive i-information X to Y is defined as:
lv(X +Y- (Y - Hv(Y)X), where HY Hv(Y are the predictive conditional V-entropy defined as:
(3)
Hv(Y|X)
inf
(4)
fev
Hv (0) inf Ey~y log [1](
(5)
Notably that f E V is a function XU [0) + P(V), so f[x] E P(Y) is a probability measure on Y and [g) E TR. is the density evaluated on y V. Hv (0) is also denoted as HY(
2

---

Published as a conference paper at ICLR 2023
Compared with the mutual information, the predictive V-information restricts the computational power and much easier to estimate in high-dimensionss cases. When the predictive family V
contains all possible models, i... V f2, it is proved that Ly(X - Y) I(X;Y) (Xu et al. 2020
3 PREDICTIVE HETEROGENEITY
In this paper, from the machine learning perspective, we quantify the data heterogeneity that affects decision making, named Predictive Heterogeneity, which is easy integrate with machine learning algorithms and could help analyze big data and build more rational algorithms.
3.1 INTERACTION HETEROGENBITY
To formally define the predictive heterogeneity, we begin with the formulation of the interaction heterogeneity. The interaction heterogeneity is defined as:
Definition 3 (Interaction Heterogeneity). Let X, Y be random variables taking values in X' x y C, and take its subset & C C. Then & is
Denote the set of random categorical variables environment set iff there exists & & such that X,Y 1L E. & 8 is called an environment variable. The interaction heterogeneity between X and Y w.r.t. the environment set is defined as:
H. (x EE
sup [X() - I(Y:X).
(6)
Each environment variable & represents a stochastic 'partition' of V, and the condition for the environment set implies that the joint distribution of X, Y could be preserved in each environment. In information theory, I(Y; I[Y; X) is called the interaction information, which measures the influence of the environment variable & the amount of information shared between the target Y and the covariate X. And the interaction heterogeneity defined in Equation 6 quantifies the maximal additional information that be gained from involving uncovering the environment variable E. Intuitively, large 7 (P) indicates that the predictive power from X to Y is enhanced by E, which means that uncovering the latent sub-population associated with the environment partition E will benefit the X - Y prediction.
3.2 PREDICTIVE HETEROGENEITY
Based on the mutual information, the computation of the interaction heterogeneity is quite hard, since the standard mutual information is notoriously difficult to estimate especially in big data sce- narios. Also, even if the mutual information could be accurately estimated, the prediction model may not be able to make good use of it.
Inspired by Xu et al. 2020), we raise the Predictive Heterogeneity, which measures the interaction heterogeneity that can be captured under computational constraints and affects the prediction of models within the specified predictive family. To begin with, we propose the Conditional Predictive V-information. which generalizes the predictive V-information.
Definition 4 (Conditional Predictive V-information). Let X,Y be two random variables taking values in X x y and be an environment variable. The conditional predictive V-information defined as:
Iv(X + Y|E) = Hv Hv where EY and Hu(Y|X, E) are defined as:
Hv(Y|X, Ee-g
inf lev
Be log log fl
(8)
Hv(Y10,E) Eewg
inf E Ey-y [(0 lev
Intuitively, the conditional predictive V-information measures the weighted average of predictive V- information among environments. And here are ready to formalize the predictive heterogeneity measure.
Definition (Predictive Heterogeneity). Let X, Y be random variables taking values in X' x y and & be an environment set. The predictive heterogeneity for the prediction X - Y with respect to & defined as:
H (X Y)
sup (x X - Y18) Tv(X Y).
(10)
3

---

Published as a conference paper at ICLR 2023
where II + Y) is the predictive V-information following from Definition 2
Leveraging the predictive V-information, the predictive heterogeneity defined in Equation 10 char- acterizes the maximal additional information that can be used by the prediction model when involv- ing the environment variable E. It restricts the prediction models in V and the explored additional information could benefit the prediction performance of the model f V, for which it is named predictive heterogeneity. Next, we present some basic properties of the interaction heterogeneity and predictive heterogeneity.
Proposition 1 (Basic Properties of Predictive Heterogeneity). Let X, Y be random variables laking values in X V, V be a functionn family, and E, E,, 8g be environment sets.
1. Monotonicity: 1 - E2. H6 (x +Y) (X
2. Nonnegativity: HX + 1) 20. O.
3. Boundedness: (X - +Y) S Hy((
4. Corner Case: If the predictive family V is the largest possible predictive family that includes all possible models, i.e. V = Ω2 we huve H. (X.Y) Ha (x - Y).
For further theoretical properties of predictive heterogeneity, in Section 3.3 we derive its explicit
forms under endogeneity, a common reflection of data heterogeneity. And we demonstrate in Section 3.4 that our proposed predictive heterogeneity can be empirically estimated with guarantces the complexity of V is bounded (c.g., its Rademacher complexity).
3.3 THEORETICAL PROPERTIES IN LINEAR CASES
In this section, we analyze the theoretical properties of the predictive heterogeneity in multiple linear settings, including (1) homogeneous case with independent noises and (2) heterogeneous cases with endogeneity brought by selection bias and hidden variables. Under these typical settings, could approximate the analytical forms of the proposed measure and the conclusions provide insights for general cases.
Firstly, under homogeneous case with data heterogeneity, Theorem proves that our measure is bounded by the scale of label noises (which is usually small) and reduces to O in linear case under mild assumptions. It indicates that the predictive heterogeneity is insensitive to independent noises. Notably that in the linear case we only deal with the environment variable satisfying X I e E,, since in common prediction tasks, the independent noises are unknown and unrealistic to be exploited for the inference of latent environments E
Theorem (Homogeneous Case with Independent Noises). For a prediction task X + Y where X, Y are random variables taking values in IR=" x IR, consider the data generation process as Y = g(x) + N o²) where g IR" + R is a measurable function. 1) For a function class & such that g E G, define the family as Vg {f][[x] N (((zz) (f d E G,ov E RR} }. With an environment set &, we have H (X Y) S NO². 2) Take = and 9(x) = -3x,3 E R. Assume [[X] and [[X"] exists. Given the function family V. [f]f[=]] N00r, 0²), 0 R. a fixed and the environment set & {E3E E C, [supp(E)] = 2, X 1 <|E}. We have H&S (X + Y) 0.
Secondly, we examine the proposed measure under two typical cases of data heterogeneity (Fan et al. 2014), named endogeneity by selection bias (Heckman, 1979; Winship & Mare 1992 Cui &
Athey 2022) and endogeneity with hidden variables (Fan et al. [2014] AArovsky et al. 2019
To begin with, in Theorem 2 we consider the prediction task X + Y with X, Y taking values in R2 x R. Let X = [S, VIT. The predictive family is specified as:
V == [/1f[x] = NNssSS + V,²), 0s.0v E RR
(1I)
And the data distribution P(X,Y* Y) is a mixture of latent sub-populations, which could be formulated by an environment variable 8 e C such that P(X. = E. Ee P(E PXX, e). For each e E supp(E*), P(X, Y = e) is the distribution of a homogeneous sub-population. Note that the prediction task is to predict Y with covariates X, and the sub- uub-oopuattion structure is latent. That is, P( [X, Y) is unknown for models. In the following, we derive the analytical forms of our measure under the one typical case.
4

---

Published as a conference paper at ICLR 2023
Theorem 2 (Endogeneity with Selection Bias). For the prediction task X x= [S, + Y with latent environment variable E", the data generation process with selection bias is defined as:
= B$ f(5) + 6Y N3), =
ev, NO0,1),
(12)
where f R + R and r,a supp(E") + R are measurable functions. 3 E R. Assume that ([ and there exists L >> such that (E" < (E))))) For the predictive family defined in equation and the environment set & C. the predictive heterogeneity of the prediction task [S, v + approximates to:
Var(r.) 3 (E') ((F((S)) error bounded by 1 max(oy. R(r, (13) Er (E+)]
HC(X- Y) ~
2
And ((r(E"),a(E) f, El( PE7a "B][ + Ee
1 BB| B Matty + E+2+
Intuitively, the data generation process in Theorem 2 introduces the spurious correlation between the spurious feature V and the target Y, which varies across different sub-populations (i.e. r(E*) and (EE)) varies) and brings about data heterogeneity. Here (S) specification since there is a nonlinear term f(S) that could not be inferred by the linear predictive family with the stable feature S. The constant characterizes the strength of the spurious correlation between V and Y. Larger L means V could provide more information for prediction.
indicates a model mis-
From the approximation in Equation 13 we can see that our proposed predictive heterogeneity is dominated by two terms: (1) Var(r(E (ar(r(E))//B)"((((( (E*)) characterizes the variance of r(E") among sub-populations; (2) (S)] reflects the strength of model misspecifications. These two compo- nents account for two sources of the data heterogeneity under selection bias, which validates the rationality of our proposed measure. According to the theorem, the more various r(E*) among the sub- uub-populations and stronger model misspecifications, the larger the predictive heterogeneity.
In general, Theorem D and indicate that (1) our proposed measure is insensitive to the homoge- neous cases and (2) for the typical sources of data heterogeneity, our measure accounts for the key components reflecting the latent heterogeneity. Therefore, the theoretical results validate the rationality of our measure.
3.4 PAC GUARANTEES FOR PREDICTIVE HETEROGENEITY ESTIMATION
Defined under explicit computation constraints, our Predictive Heterogeneity could be empirically estimated with guarantees the complexity of the model family is bounded. In this work, we provide finite sample generalization bounds with the Rademacher complexity. First, describe the definition of the empirical predictive heterogeneity, the explicit formula for which could be found in Definition 7 in Appendix.
Definition 6 (Empirical Predictive Heterogeneity (informal)). For the prediction task X + Y with X, Y taking values in X' XxY, a dataset independently and identically drawn from the population such that D = {(I;, yi); 1 ~ X. Given the predictive family V and the environment set K} where K E N+ is the number of environments, the empirical predictive heterogeneity FL (X + Y:D) with respect to 'D is readily obtained by estimating H (X- +Y, Y) on D with expectations replaced by statistics of finite samples. The formal definition placed in Desinition 7
Ex = ={E E C, supp(E)
Theorem 3 (PAC Bound), Consider the prediction task X + +Y where X. Y are random variables taking values in x V. Assume that the predictive family V satisfies Yz E X, Vy ..Vf E V, log [x] (gy) E [[B, where > O. For given E N. the environment set is defined as Ek = EE E supp(E)= K). Let Q be the set of all probability distributions of X,Y,E where & E Ex. Take an e E supp(E) and define a function class Gv = (ggggg
= log J (E elxx, y), V E ②). Denole the Rademacher complexity of G with N samples by R (9). Then for any d' (0,1/(2K +2K+)2),, +2),, with a probability over 1 2(K 1)8. for dataset D independently and identically drawn from X, Y. we have:
|H (X - Y)- Hi (X- Y;D) S 1(K 1) A (Gv) + 2(K + B)BB/ 2log /ID.
(14)
where Rip (Gv) O(ID)- (Bartett & Mendelson 2002
5

---

Published as a conference paper at ICLR 2023
4 ALGORITHM
To empirically estimate the predictive heterogeneity in Definition 6
we derive the Information Max- to infer the distribution of & that
imization (IM) algorithm from the formal definition in Equation 33
maximizes the empirical predictive heterogeneity H (X + YDD).
Objective Function. Given dataset D {XN. YN)
{( v)) denote supp(E)
(e1,-..,ex* exk we parameterize the distribution of ( N, YN) with weight matrix W E WK. where K' is the pre-defined number of enviromments and WK = (W: : W E R BNNK and W1 = 1x) is the allowed weight space. Each clement wi in W represents P(E = z,, v4) (the probability of the i-th data point belonging to the j-th sub-oopulation). For a predictive family V. the solution to the supremum problem in the Definition 7 is equivalent to the following objective function:
min Rv(W,0((W) (WW)))....((W))))) WeWk
ΣΣ
mobbo(foz(1)). (x.), UV(WYNs)
N
(15)
s.t. (3(W) arg min Co(W,)) Σ wogev((a(e).uu)) for ,,..,,
where So X' + y denotes predicting function parameterized by 0, :) y x y R rep- resents a loss function and (W, YN) is regularizer. Specifically, Se. by and Uy are determined by the predictive family V. Here we provide implementations for two typical and general machine learning tasks, regression and classification.
(1) For the regression task, the predictive family is typically modeled as:
Vi = {g : = oZ) f is the predicting function and 0 is learnable, is a constant).
(16)
The corresponding loss function is Byr (Jo(X),)) (Jo(X) Y)2, and Uv, (W, YN) becomes
1
Uv. (W, Ys) = Varze[k] (YZ)
Wiyl/' N NE wij
N
31
(17)
where Y denotes the mean value of the label Y given = eg and UW, Yx) calculates the variance Y among sub-populations ei ~ ek-
(2) For the classification task, the predictive family is typically modeled as:
1/2 = {g g|ø] fa(w) E A is the classification model and 0 is learnable).
(18)
where è is the class number and Ae denotes the e-dimensional simplex. Here each model in the predictive family V/2 outputs a discrete distribution the form of a c-dimensional simplex. In this case, the corresponding loss function (,,(,)) is the cross entropy loss and the regularizer becomes Uv, (W, YN) - HYZ)). where H(Y)) is the entropy of Y given & = ej-
Optimization. The bi-level optimization in Equation 15 can be solved by performing projected gradient descent w.r.t W. The gradient of W can be calculated by: (we omit the subscript here)
K
VwR
VwU + (E)So, (x)).
NXK
Vo, Rlo VwO
(19)
82C
JC 80, 00,
where
Vo, R
VwO
2 Vo, R
II
(20)
03
00,00T 3
Jcc
for
K.
(21)
00,0wT
where {(fo, (I.). 3:), is an N x K matrix in which the (i, element is {(fe, (2+), 91). Here Equation 20 approximates 07 by 0f from steps of inner loop gradient descent and Equation 21
Nx.
performs Ttep truncated backpropagation (Shaban et al. [2019 [hou et al. (022). Our information maximization algorithm updates W by projected gradient descent
W Projwxx (W - V R), 77 the learning rate of W.
(22)

---

Published as a conference paper at ICLR 2023
Then we prove that minimizing Equation 13 exactly finds the supremum w.r.t. & in the Definition (formal) of the empirical predictive heterogeneity.
Theorem 4 (Justification of the IM Algorithm). For the regression task with predictive family V1 and classification task with V2, the optimization of Equation [5] is equivalent to the supremum problem of the empirical predictive heterogeneity Ne (X - Y; D), H (X - Y;D) respectively in Equation 33 with the pre-defined environment number K (i.e. =upp(E)=K]
Remark (Difference from Expectation Maximization). The expectation maximization (EM) al- gorithm is to infer latent variables of statistic model to achieve the maximum likelihood. Our proposed information maximization (IM) algorithm is to infer the latent variables W which brings the maximal predictive heterogeneity associated with the maximal information. Due the regu- larizer Uy in our objective function, the EM algorithm cannot efficiently solve our problem, and therefore we adopt bi-level optimization techniques.
5
EXPERIMENTS
5.1 PROVIDE INSIGHTS FOR THE SUB-POPULATTIO DIVISION
The predictive heterogeneity could provide insights for the sub-population division and benefit decision-making and we illustrate this prediction tasks of various fields, including agricultural research, sociological research, and object recognition. From the illustrative examples, we show that the learned sub- tobboooulttion division highly explainable and relevant to decision decision-making.
Example: Agriculture It is known that the climate affects crop yields and crop suitability (o-- bell al. 2008). We leverage the data from the NOAA database which contains daily weather from weather stations around the world. Following Zhao et al. (202T), we summarize the weather sequence of the year 2018 into summary statistics, including the average yearly temperature, hu- midity, wind speed and rainy days. The task is to predict the crop yield in each place with weather summary statistics and location covariates (i.e. longitude and latitude) of the place. For easy illus- tration, we focus on the places with crop types of wheat or rice. Notably that the input covariates do not contain the crop type. We use MLP models in this task and set K = for our IM algorithm.
Since the crop yield prediction mechanisms are closely related to the crop type that is unknown in the prediction task, we think this causes data heterogeneity in the entire data and the recognized predictive heterogeneity should relate it. In Figure [ (a), we plot the real distribution map of wheat and rice planting areas, and in Figure (b), we plot the learned two sub-populations of our IM algorithm. From the results, we surprisingly find the division given by algorithm is quite similar to the real division of the two crops, indicating the rationality of our measure. For the areas that are similar (e.g. Tibet Plateau in Asia), we think it is due some missing features (e.g. population density, altitude) that significantly affect the crop yields.
Example: Sociology We use the UCI Adult dataset (Kohavi & Becker 1996). which is derived from the 1994 Current Population Survey conducted by the US Census Bureau and is widely used in the study of algorithmic fairness. The task is to predict whether the income of person greater or less than 50k US dollars according to personal features. We use linear models in this task and set K 2. In this example, we would like to investigate whether there exist sub-population structures inside data that affect the learning of machine learning models.
Cluster
Division wheat areas
(b) Division by algorithm
Figure 1: Results the crop yield data. We color each region according its main crop type, and the shade represents the proportion of the main crop type after smoothing via k-means (k = 3).

---

Published as a conference paper at ICLR 2023
Maritl
Marital
Education
Education
Sez
Set
Bird
822
Caphtal Gate
Capital
Geta
Water
Bird (f - 1)
Capital
Lees
0 The averge kataraa
(8) The festurt coffficians.
Sub patt-pppplliinn
Sul-papulation 2
Figure 2: Results on Adults data. Here we show the average of features and the feature co- efficients of the two learned cub-populations
Figure 3: Results on the Waterbird data. Here randomly sample 50 images for each class and each learned sub-population.
In Figure 2 (a), we plot summary statistics for the two sub-populations, where the main difference lies in the capital gain. In Figure (b), we plot the feature importance given by linear models for the two sub- tobb-populatios, and we find that for people with high capital gain, the prediction model focuses mainly on capital gain, which is fair. However, for people with a low capital gain, models also address some sensitive attributes such as sex and marital status, which tend to cause discrimina- a
tion. Our results correspond with the results in Zhao et al. (202I) and can help us identify potential inequality in cocision-making For example, our results indicate the potential discrimination for low capital gain people, which could further promote algorithm design and improve policy fairness.
the Waterbird dataset (Sagawa et al. 2019). which widely used to evaluate the model's robustness the robust learning field. The task i8 to recognize waterbirds landbirds. However, the image backgrounds label, i.e. for the majority, waterbirds are on the water and landbirds on the land, and for minority, the correlation is reversed. Therefore, the spurious correlation causes predictive heterogeneity in this dataset, since such correlation could affect the machine learning model. In this example, we use the show the learned sub-populations of
Example: Object Recognition Finally,
spuriously correlated with the target
ResNet18 and K in our IM algorithm. In Figure our method, we randomly sample 50 images for each class (waterbird or landbird) and each learned sub-pon popplppoulaion.... In sub-poye aub-populttion the majority of landbirds are on the ground and waterbirds are in the water, while in sub-population 2, the majority of landbirds are in the water and waterbirds are on the ground. Our measure captures such spurious correlation and the two sub- sub-populations, the spurious correlation between the object and background is inverse. And the learned sub- sub-populations could be leveraged by many robust learning methods (Sagawa et al.) 2019 Koyama & Yamaguchi 2020) to learn models with better generalization ability, since they can help eliminate the influence of backgrounds on model predictions.
5.2 BENEFIT OOD GENERALIZATION
The predictive heterogeneity could benefit the out-o--dist (OOD) generalization of machine learning models. Here we investigate the empirical performance of our IM algorithm w.r.t. the OOD generalization performances on simulated data and real-world colored MNIST data.
Baselines First, we compare with empirical risk minimization (ERM) and environment inference for invariant learning (EIIL. (Creager et al., (202I) which infers the environments for learning in- variance. Then compare with the well-known KMeans algorithm, which is the most popular
Table 1: Results of the experiments on ouittf-disttibution generalization, including the simulated data and colored MNIST data.
1. Simulated Data
Colored MNIST
Method
Training sopulattn Error
Test Error
Train Accuracy
Test Accuracy
Major (r 1.9)
Minor (r -1.9)
2.3
-2.7
ERM
0.255(20.024)
0.740 022)
0.737 0.738
0.998 =0.001
0.
0.164
1.428 005:
1.431ca0
1.431
0.812:20 006
Balance
1:20..200022
0. 0.847:30.00 024;
0,846
0.845:20.026
0.999:
0..20
KMcans
IRM
1223..00.022)
0.8455:a..com 024)
=.844:a..ca9)
0.843taa.029)
0.947::0.004:
0.235
0.840::.o
0.839 cus)
0.838
0.997:30.011
0. cm
Balance
0.403
0.423
0.416
0.416:20.010
0.749:
0.692
Ours
IRM
0.391ss-.0599
0.442:++o.oud)
=.400)=
0.430caa.014
0.759::0.014:
0.727(=0.047
IGA
0.449
0.416
0.417:c0.022
0.417
0.713
8

---

Published as a conference paper at ICLR 2023
Jurget
Sparkorrfooerr
Spurisas foence
Spertonn feshume
Sab-pspaleionn
Sul- Suh-papuhiiin 5
(a) KMcans.
(b) EII...
(c) Our IM.
Figure 5: Sub-population di- Figure 4: Sub-po sub-population division the simulated data of three vision the MNIST data of methods, where two colors denote two sub-populations.
our IM algorithm.
clustering algorithm. For our IM algorithm and KMeans, we involve three algorithms as backbones to leverage the learned uub-populations, including sub sub-population balancing and invariant learning methods. The sub uob-populttion balancing simply equally weighs the learned cubb-poulations. in- variant risk minimizution (IRM, (Arjovsky et al. 2019)) and inter-environment grudient alignment (IGA, (Koyama & Yamaguchi] 2020)) are typical methods in OOD generalization, which take the sub-populations as input environments to learn the invariant models.
Data Generation of Simulated Data The input features X |S, VJ" RR10 consist of stable features S E R,, noisy features T E R and the spurious feature V E R:
S S~ 2I,), T ~ (0, 21)), Y = 6 h(S) +V 0.5), V Laplace(sign(()) 1/55 (rD) (23) where Os E IR is the coefficient and (()) S2S is the nonlinear term. |r| > is a factor for each sub-population, and here the data heterogeneity is brought by the endogeneity with hidden variable (Fan et al. 2014]. V is the spurious feature whose relationship with Y is unstable across sub-pon uub-populations and is controlled by the factor r. Intuitively, sign(r) controls whether the spurious correlation between V and Y is positive negative. And ir controls the strength of the spurious correlation, i.e. the larger |r| means the stronger spurious correlation. In training, we generate 10000 points, where the major group contains 80% data with correlation) and the minor group contains 20% data with r = - -1.9 (i.e. strong negative spurious correlation). In testing, we test the performances of the two groups respectively, and we also set r = 2.3 and r = 227 to simulate stronger distributional shifts. We use linear regression and set K- for all methods, and we report the mean mean-square errors (MSE) of all methods.
1.9 (i.e. strong positive spurious
Data Generation of Colored MNIST Following AAjovssy et al. (2019), we design a binary clas- sification task constructed on the MNIST dataset. Firstly, digits 0 4 are labeled Y and Secondly, noisy labels Y are induced by randomly flipping the
digits are labeled Y label Y with a probability of 0.2. Then we sample the colored id V spurious correlated with Y as V v= with probability r, with probability 1 r. In fact, r controls the spurious correlation between Y and V. In training, we randomly sample 10000 data points and set r = 0.85, meaning that for 85% of the data, V is positively correlated with Ý and for the rest 15%, the spurious correlation becomes negative, which causes data heterogeneity w.r.. V and Y. testing, we set = 0 (strong negative spurious correlation), bringing strong shifts between training and testing.
Analysis From the results in Table D for both the simulated and colored MNIST data, the two backbones with our IM algorithm achieve the best OOD generalization performances. Also, for the simulated data, the learned predictive heterogeneity enables backbone algorithms to equally treat the majority and minority inside data (i.e. low-performance gap between 'Major' and 'Minor'), and significantly benefits the OOD generalization. Further, for both experiments, we plot the learned sub-populations of our IM algorithm Figure and And From Figure 4 compared with KMeans and EIIL, our predictive heterogeneity exploits the spurious correlation between V and Y, and enables the backbone algorithms eliminate From Figure 5) the learned sub- uub-populations of our method also reflect the different directions of the spurious correlation between digit labels Y and colors (red or green), which helps backbone methods to avoid using colors to predict digits.
CONCLUSION
We define the predictive heterogeneity, the first quantitative formulation of the data heterogeneity that affects the prediction of machine learning models. We demonstrate its thcoretical properties and show that it benefits the out-of-distribution generalization performances.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We would like to thank Yuting Pan, Jiaming Song, Fan Bao and anonymous reviewers for helpful feedback. Peng Cui's research supported in part by National Key R&D Program of China (No. 2018AAA0102004, No. 2020AAA0106300), National Natural Science Foundation China (No. U1936219, 62141607), Beijing Academy of Artificial Intelligence (BAAI). Bo Li's research was supported by the National Natural Science Foundation of China (No.7217113) 72133002); the Technology and Innovation Major Project of the Ministry of Science and Technology of China under Grants 2020AAA0108400 and 2020AAA0108403.
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprini arXiv: aeee::90700090000 2019.
Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-493 2002.
Jasmin Breitenstein, Jan-Aike Termöhlen, Daniel Lipinski, and Tim Fingscheidt. Systematization of corner cases for visual perception in automated driving. In 2020 IEEE Intelligent Vehicles Symposium (IV), pp. 1257-1264. IEEE, 2020.
Robert Challen, Joshua Denny, Martin Pitt, Luke Gompels, Tom Edwards, and Krasimira Tsaneva- Atanasova. Artificial intelligence, bias and clinical safety. BMJ Quality & Safety, 28(3):231-237, 2019.
M Cover Thomas and A Thomas Joy. Elements of information theory. New York: Wiley, 3:37-38, 1991.
Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. International Conference on Machine Learning, pp. 2189-22000 18000 PMLR, 2021.
Peng Cui and Susan Athey. Stable learning establishes some common ground between causal infer- ence and machine learning. Nature Machine Intelligence, 10-1111
John Duchi, Tatsunori Hashimoto, and Hongseok Namkoong. Distributionally robust losses for latent covariate mixtures. Operations Research, 2022.
John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. The Annals of Statistics, 49(3):1378-144 1406, 2021.
Kevin Dzobo, Dimakatso Alice Senthebane, Nicholas Ekow Thomford, Arielle Rowe, Collet Dan- dara, and M Iqbal Parker. Not everyone fits the mold: Intratumor and intertumor heterogeneity and innovative cancer drug design and development. Omics: a journal of integrative biology. 22 (1):17-34, 2018
Jianqing Fan, Fang Han, and Han Liu. Challenges of big data analysis. National science review, 12::293-314 2014.
Jingrui He. Learning from data heterogeneity: Algorithms and applications. IJCAI, pp- 5126 5130, 2017.
James J Heckman. Sample selection bias as a specification error Econometrica: Journal of the econometric society, pp. 153-161, 1979.
Sumyea Helal. Subgroup discovery algorithms: a survey and empirical evaluation. Journal of computer science and technology, 31 (1))::61-576 2016.
Miguel A Hernán, David Clayton, and Niels Keiding. The simpson's paradox unraveled. Interna- tional journal of epidemiology, 40(3):780-785, 2011.
Michael Keamss, Seth Neel, Aaron Roth, and Zhiwei Steven Wu. Preventing faimess gerryman- dering: Auditing and learning for subgroup fairness. International Conference on Machine Learning, pp. 2564-2572 PMLR, 2018.
10

---

Published as a conference paper at ICLR 2023
R. Kohavi and B. Becker. Adult, 1996.
Masanori Koyama and Shoichiro Yamaguchi. Out-of-distribution generalization with maximal in- variant predictor. 2020.
Michel Ledoux and Michel Talagrand. Probability in Banach Spaces: isoperimetry and processes, volume 23. Springer Science & Business Media, 1991.
Hi and JF Reynolds. On definition and quantification of heterogeneity. Oikos; pp. 280-284, 1995.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In International Conference Machine Learning, pp. 2804-6814, PMLR, 2021.
Jiashuo Liu, Zheyan Shen, Peng Cui, Linjun Zhou, Kun Kuang, and Bo Li. Distributionally robust learning with stable adversarial training. IEEE Transactions on Knowledge and Data Engineer- ing, 2022.
David B Lobell, Marshall Burke, Claudia Tebaldi, Michael Mastrandrea, Walter Falcon, and Rosamond L Naylor. Prioritizing climate change adaptation needs for food security in 2030. Science, 3196586))607-610,, 2008.
Todd K Moon. The expectation--mxxiizaaiii algorithm. IEEE Signal processing magazine, 13(6): 47-60, 1996.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant pre- diction: identification and confidence intervals. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 78(5):947-1022 2016.
Douglas A Reynolds. Gaussian mixture models. Encyclopedia of biometrics, (41(659-663), 2009.
Paul R Rosenbaum. Heterogeneity and causality: Unit heterogeneity and design sensitivity in ob- servational studies. The American Staristician, 1972):447-152, 2005.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint arXiv:/91 1.08731, 2019.
Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated pack-propagation for bilevel optimization. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 1723-1732. PMLR, 2019.
Clifforrd H Wagner. Simpson's paradox in real life. The American Statistician, 36(1):46-48, 1982.
Xiao Wang, Houye Ji, Chuan Shi, Bai Wang, Yanfang Ye, Peng Cui, and Philip S Yu. Heterogeneous graph attention network. The world wide web conference, pp. 2022-2032 22--2032,442 2019.
Christopher Winship and Robert D Mare. Models for sample selection bias. Annual review of sociology, pp. 327-350, 1992.
Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints. In 8th International Conference on Learning Rep- resentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https://cpenrevil .et wttpeeewweeeeooooiiieeeeeeeeeeee
Shengjia Zhao, Abhishek Sinha, Yutong He, Aidan Perreault, Jiaming Song, and Stefano Ermon. Comparing distributions by measuring differences that affect decision making. In International Conference on Learning Representations, 2021.
Xiao Zhou, Yong Lin, Renjie Pi, Weizhong Zhang, Renzhe Xu, Peng Cui, and Tong Zhang. Model agnostic sample reweighting for out-of-distribution learning. In International Conference on Ma- chine Learning, pp. 27203-277 27221. PMLR, 2022.
11

---

Published as a conference paper at ICLR 2023
A FORMAL DEFINITION OF EMPIRICAL PREDICTIVE HETEROGENBITY
In this section, we derive the explicit formula for the empirical estimation of the predictive hetero- geneity which is described in Definition 6
The dataset
[(x; w;)) is independently and identically drawn from the population X, Y. Given function family V and an environment set 6x. let Q be the set of all probability distributions of X,Y,& E where & E 6x. For given ε, denote supp(E)) = {(e))==)) The empirical predictive heterogeneity Hee (X + YD)) is given by:
H (X Y;2D = sup BEER
f + (|ED) - T -Y;
(24)
Σ Q(E== (Y) D2; - -=== (x)x=
sup
QEQ
(25)
[Hv(Y:D) H)))
(26)
Specifically,
QIE ex) Hv) cx:D) D)
(27)
Q0n = cx)
inf Q(E = ex)
Σ
- log (vr)
(28)
FEV
En Q(E,, E = Ex)
ILD
QIE = ex x Q1) v1)
inf Q(E = -kk)
Σ
log [zzi((ii)
(29)
feV
Σ WD Q(E = en
Q(E = ) (I..
)
inf Q(E ex) fEV
Σ
log f(x;((1s)
(30)
Q(E = ess)
FDDiED
innf FEV ED
Σ
log f(x:)(.)) -ells (1) (1 Uz)
(31)
inf fEV
Σ Z..DED
- log (o: = en (xllri,11) vk)-
(32)
The explicit formula for Q(E ex) Hv) = ex; D), Hv(Y|X;; D) and Hy(Y; D) could be similarly derived. Here we are ready formally define the empirical predictive heterogeneity.
Definition (Empirical Predictive Heterogeneity (formal)). For the prediction task X + Y with X, Y taking values in X x )) dalaset 'D is independently and identically drawn from the population such that DD (( y:) ~ X,Y), Given the predictive family V and the environment set EK = =EEE E cuspp(()) = K} where K E N, let Q be the set of all probability distributions of X,Y,E where & E 6x. The empirical predictive heterogeneity H65 (X - Y;D) with respect to 'D is defined as:
R (X +Y;D) sup
= ex: ( (
QEE k=1
- [(v((Y:)) Hv(Y)x:D)),
(33)
12

---

Published as a conference paper at ICLR 2023
where
= ( = inf
Σ
log ʃ (ai = ex ex(21;
(34)
fEV [D]
cx)= ex: D) = inf
log f ( 2x Uo).
(35)
JEV [
F..D:ED
Hx(Y|x:D) = inf
Σ
log ʃ [.)(()))
(36)
JEV
[D]
ED
inf
Σ
log (())))).
(37)
JEV D
F..D:ED
B SENSITIVITY OF K
In the experiments of Section [] we set the K for easy illustrations. In this section, we add the results of choosing different Ks for the simulated experiment in Section 5.2 to show that the OOD generalization performances of some typical algorithms plus our proposed method are not sensitive to the choices of K.
0.7
0.6
0.5
MSE
0.4
0.3
Ours Balance
E
0.2
Ours + IRM
0.1
Ours + IGA
ERM
K K= 2
3=3
K= 4
5=
Figure 6: The out-of-distribution generalization error of our methods with Sub- sub-population Balanc- ing, IRM and IGA as backbones. Here we plot the eπrors of different backbones under -277, which introduces strong distributional shifts with training data.
In Figure 6 we show the -iitti--iitiitttinn generalization error of our methods with Sub-population Balancing. IRM and IGA as backbones. We plot the OOD testing performances under r -2.7. which has strong distributional shift with the training distribution. From the results, we can see that the performances of three OOD generalization methods do not be affected much by the choice of K, and from Table I. our performances significantly outperforms all the baselines.
Also, we add one more experiment to show that (1) when the chosen is smaller than the ground- truth, the performances of our methods will drop but are still better than ERM (2) when the chosen K is larger, the performances are not affected much (consistent with the results in Appendix B).
[S, T, V] E 10 consist of stable features S IR5 noisy features T' I and the spurious feature V E I::
Experiment Setting: The input features X
S N22, 2Tss, T N10 2T44;
Y = 05. + SIS2S + (00,
and we generate the spurious feature via:
V = Y (0..0...)
13

---

Published as a conference paper at ICLR 2023
where By varies across sub col--oollttions and is dependent which sub sub-population the data point belongs to. In training, we sample 8000 data points from é1 with 8. 3.0, 1000 points from 62 with Bi 1.0,, 1000 points from ez with 8, -2.0 and 1000 points from e4 with f -3.0. Therefore, the ground-truth number of sub-populations is 4. In testing, we test the performances on e4 with 0f -3.0, which has strong distributional shifts from training data. The average MSE over 10 runs are shown in Figure
Ours Balance
0.85
Ours IRM
0.80
Ours + IGA
0.75
MSE
0.70
0.65
0.60
0.55
0.50
ERM
K= K=3 K= 4 K=5 K=6
Figure The out-of-distributioo generalization error of our methods with Sub-population Bal- ancing, IRM and IGA as backbones for the added experiments. The ground-truth sub-population number is 4.
From the results, we can that when K is smaller than the ground-truth, increasing K benefits the OOD generalization performance, and when K larger, the performances are not affected much, which is consistent with the results in Figure 6
For our IM algorithm, we think there are mainly two ways to choose K:
According to the predictive heterogeneity index: When the chosen K is smaller than the ground-truth, our measure tends to increase quickly when increasing K; and when K is larger than the ground-truie truth, the increasing speed will slow down, which could direct people to choose appropriate K.
According to the prediction model: Since our IM algorithm aims to learn sub-populations with different prediction mechanisms, one could compare the learned model parameters 0, 01,...,p, 0x to judge whether K is much larger than the ground-truth, i.e., if two resultant models quite similar, K may be large (divide sub-population into two). For linear models, one can directly compare the coefficients. For deep models, we think one can calculate the transfer losses across sub- cobppopulaiion..
For detailed analysis of the best choice of K, leave for future work.
C RELATED WORK
To the best of our knowledge, data heterogeneity has not converged to a uniform formulation so far, and has different meanings among different fields. Li Reynolds (1995) define the heterogeneity in ecology based on the system property and complexity or variability. [Rosenbaum (2005) views the uncertainty of the potential outcome as unit heterogeneity in observational studies in economics. For graph data, the heterogeneity refers various types of nodes and edges (Wang et al. 2019)). More recently, in machine learning, several works of causal learning (Peters et al. Arjovsky
et al. 2019: Koyama & Yamaguchi 2020 Creager et al. 2021) and robust learning (Sagawa al.)
2019 leverage heterogeneous data from multiple environments to improve the of distributio generalization ability. Specifically, invariant learning methods (Arjovsky al. 2019
Koyama &
14

---

Published as a conference paper at ICLR 2023
Yamaguchi 2020
Creager et al. 2021 Zhou et al.
2022) leverage the heterogeneous environment
to learn the invariant predictors that have uniform performances across environments. And in distri- butionally robust optimization field, Sagawa et al. (2019); Duchi et al.) (2022) propose to optimize the worst-group prediction emor to guarantee the OOD generalization performance. However, in machine learning, previous works have not provided a precise definition sound quantification of data heterogeneity, which makes it confusing and hard to leverage to develop more rational machine learning algorithms.
As for clustering algorithms, most algorithms only focus on the covariates X, typified by KMeans and Gaussian Mixture Model (GMM, (Reynolds, 2009)). However, the learned clusters by KMeans cannot reflect the predictive heterogeneity, which is shown by our experiments. And the expectation maximization (EM, (Moon] 1996)) can also be used for clustering. However, our IM algorithm has essential differences from EM, for our IM algorithm infers latent variables that maximizes the predictive heterogeneity but EM maximizes the likelihood. Also, there are methods (Creager al. 2021) from the invariant learning field to infer environments. Though it could benefit the OOD generalization, it lacks the theoretical foundation and only works in some settings.
D PROOF OF PROPOSITION
Proof of Proposition L
Monotonicity.
Because of &1 C E2.
M (X +Y) = sup 1v Y18) Iv(( +))
(38)
EE61
I sup Iy + Y|8) - -v(X -)
(39)
EEll
= (X
(40)
Nonnegativity:
According to the definition of the environment set, there exists & E 6 such that for any E supp(E), X, Y| = e is identically distributed as X, Y. Thus, we have
Hb(x + Y) sup [Y E) - Hv(Y)x, [Hv(Y Hv(Y|x)
(41)
838
(Y En) Hv (Hv -
(42)
Specifically,
Hv(Y)X,Eo) EeWEo
inf E E log [+]
(43)
JEV
=EewEo
inf E En fogg JEV
(44)
=Hv(Y
(45)
Similarly, H,,Y = HY)Y Thus, ( + +Y) 2 0.
3. Boundedness:
First, we have
Hy(Y|X,E) EenE
inf Exyy-x. log f(x [z)] Fev
(46)
Eeng
inf F
Eyny m. f(]
(47)
.ev
0,
(48)
by noticing that Ey-Yiel- log [xr] (y)) is the cross entropy between Y|a, e and f(x)
15

---

Published as a conference paper at ICLR 2023
Next,
(Y = BEVE inf E8 B B log f f10]
(49)
LED
I inf Bewe B log log (1) JeV
(50)
inf By-rl- fEV
(51)
Hy(Y
(52)
where Equation 50 is due to Jensen's 's inequality.
Combing the above inequalities,
H6(X - Y)
sup [Hv - Hv(YIX, -1 Hu(VIX)
(H>
(53)
EEE
I sup Hv - ( - Hv(XI))]
(54)
EEE
I Hy(( ( - Hv(x|x)]
(55)
Hv(YlX).
(56)
4. Corner Case: According to Proposition 2 in Xu al. (2020).
(Y HY).
(57)
Hx = XY
(58)
By taking random variables R, S identically distributed as X, Y for supp(E), we have
H e) H)(S|R)= (())R) H =)e. e).
(59)
Thus,
Ho = Eu- = = = E.e ( = =e]] = =H(1X,E)).
(60)
Similarly, we have 10, 8) = (E)) Thus,
( + Y) sup
(Y 8) - Ho(Y - (Ha(Y Ho(YIX)
(61)
8E6
sup
(H(Y))) - H(YIX, (100)-H(YX)) [((IX)
(62)
EE6
sup ILY; - X
(63)
EEB
H H (X,Y).
(64)
L
E PROOF OF THEOREM 1
Proof of Theorem 7
1)
Hvc (Y(X) = inf SEV9
Ex-x [ log fxr [v]]
(65)
(y) 2
< Ex-x
By-Ylel- log
exp
(66)
V2T /2=
I
2 2x
Exax (yy
(67)
?
(68)
16

---

Published as a conference paper at ICLR 2023
Equation (olatiioon6oooll holds by taking f[a] ((((2)), 2.x)
2)
Given the function family Vo = - a²), 0 E R, fixed J. by expanding the Gaussian probability density function in the definition predictive V-information, it could be shown that
=
Iv. (X - Y)
min E( kx)?] Var(YY)
(69)
kER
where the predictive V-information is proportional to Mean Square Error subtracted by the variance of target, by coefficient completely dependent on a.
The minimization problem is solved by
E h: E
(70)
Substituting & into eq 69
11, (X + Y) x Ble-] - Var(X + E)
(71)
-ar((X -E(X
(72)
Denote supp(E) = E2) Let Q be the joint distribution of (X. E). Let Q(Ei) and Q(E2) = be the marginal E. Abbreviate by e2(X,))
e(E, Ei) by Pi and &( = 82)
Similar to
Iv (X + Y|E) x min B[(Y - - Var(Y]E).
(73)
For & E the minimization problem is solved by
Ep [X]] Ep [X2]
k
(74)
Thus,
l1v, (X Ei) x Ep. [x-
PPP [XY]
27
X
Varp, (Y)
(75)
Ep, X2
13%, [XY]
Ep, 2
(Ep, [ - E3 (Y))
(76)
Ep, [X
E7, [XY]
E3 [Y]
(77)
Ep [x2]
Similarly, we have
E2, [X]
Iv. (X + Y = = B2 [Y]
(78)
Ep. [
Notably, Ep, [X2] and Ep, [ are constrained by a and E[X²).
E[x²1 = E| = a [ + (1 - a EE, (x²).
(79)
Similarly,
= = E|XY] RF XY| + (1- Br, [XY]
(80)
0= E|y] = FIE,, [ + (1 - a y..
(81)
The moments of Pa could thereafter be represented by those of Pt.
X aL [X
Ep, [x
(82)
E 2 AB X E E [XY]
(83)
Ok
17

---

Published as a conference paper at ICLR 2023
aEp, M
Ep M
(84)
1 1-a a
Substituting to eq 78
a²
1 [B[x2] aEp, (XY]] 2
Iv, (X + E2) x (1- a)²
E M
(85)
1- a [| - p1 [X
Thus,
HS, (X + = sup lv: (X - Y)- aIv. (X + 3/E E1 - (1- a (X ---0-800 >1IE E2) EE6
(86)
12 [Y]
a²
E [ BB
[XY])
X sup 333
- [[X] - xlp M + a
E2 [
E [2]
a
E
a E,,, [
(87)
p [X]] EEp (XY])
sup
a
E2, Y + a Ep X [ a.E [X
[[X²
(88)
E38
1 1-a a
Ex [Xe]
sup
E2 X {)
a.P (x21) B[X
(89)
333
a
Ep
[X2
[E
X
Assuming X I. E
HS. (X + Y) = -sup a B [X + c 50.
(90)
1 - a
From Proposition ! we have HS, (X + Y) 2 0. Thus, H. (X + +Y) = -.
F PROOF OF LINEAR CASES (THEOREM 2 AND ??)
Proof of Theorem
For the ease of notion, we denote the r(E*) as Te, (EE) as oe, and (E8)) E. as Eg. And we omit the superscript C of Hb. Firstly, we calculate the Hy Y [0] as:
1
1
Hv[110] = Var(Y) + log o
log 2nT,
(91)
20
2
1
Hy H10 = 20
Re- [Var(Y|8")] + +log a + 2 log 2n..
(92)
Therefore, we have
-
Var (ar(E()))) s
(93)
202
As for Hy[Y we have
Hy[X = inf Ex, [[]Y - (hsS + /o V)
(94)
hs.
20
1
inf hg.hy
Ex, [I]((S) + cx (hoS + by V)
(95)
20
1
inf Eg- [[[] + - (hsS + (r + [.)
(96)
hs.bv
20
where we let hg hg 3 here. Then we have
20 20² Hy [V]X=
inf Ee- [1]] (() hgS hvcel1
(97)
hashv
inf Eg- [[]|[[[ by.
+o + 12
(98)
18

---

Published as a conference paper at ICLR 2023
notably that here for c₁ Cj supp(E*), we assume P (S. PSy (S,Y) (we choose such E" as one possible split). And the solution of hg, hy
hs E E
(99)
E' [(S)
Elr. B
bv + [[ E2
(100)
[SS)S
According to the assumption that [S} 0, we have hg 0,
(101)
hv =
(102)
Therefore, we have
20 H = Eg [ h
[- + of + t 1o.21
(103)
Var(r. )B 18 (E*)
[[) + +
(104)
E|r2]
20 Hv oタ + E r2
? Eg I[
(105)
Ia + g
T.
oz
Note that here simply set 1 in the main body, And we have:
Hy Y) Var(r. [[ (E*)) E[f((S)
(106)
E + (E+)]
The approximation error is bounded by ((())"))((("))BB)))) is defined as:
1 2 max R(r(E*), ((()), ((33)),
and
1
((("))).((E")B))) E[( PE
| + Eg-
)z|
(107)
+
de
+ r E[
Proof of Theorem ??. As proved above, have
E S(
(Var(r))
hs = B +
(108)
( E2 [(S))
SS)
S
by
(109)
-
For the model misspecification case, we further assume that (1) (S) 0 and (2)
E|o"| < (S) and then we have
hg = 3,
(110)
hv r
(111)
E
and for the heterogeneity, we have
Var(re)
(5)) + 1e +oy 20 22 +)
(112)
Var(re)
(E] +
Be ,el
Without the model misspecification, we assume that f = 0, and then we have
hs 8
(113)
E Elr.oo
by E1 1
(114)
19

---

Published as a conference paper at ICLR 2023
and for the heterogeneity we have
2v² + Y) ž of 2 B(r.] + + Bi - E 12
(115)
20 Hv + Y) v -1 B(2) +
B[o.1 Bool
(116)
G PROOF OF THE ERROR BOUND FOR FINITE SAMPLE ESTIMATION (THEOREM
In this section, we will prove the error bound of estimating the predictive heterogeneity with the empirical predictive heterogeneity. Before the proof of Theorem 3 which is inspired by Xu et al. 2020 we will introduce three lemmas.
Lemma 1. Assume Væ E X.Vy J.V V. log f(r(y) e [-B, where B 0. Define a function class g = {9lø(z, = log (ogl(l) ex(z,v), E V.q q E
2). Denote the Rademacher complexity of with N samples by (N( Define fk log [<:))(v))((E ex(ei,
arg inf, [D] Em.RED
Then for any q Q, any 8 (0,1). with probability over 1 - 8, we have
1
Σ
log fa [xs)(()))(
(117)
D
log 1 DD]
< 2 11 + BV
(118)
Proof. Apply McDiarmid's incquality to the function (DD) which is defined as:
((D) = sup SEV
(E = ex [66==x]o
log (x) E = ex]
Σ
llog (1:l()))((E = ex
D
2i.NEDD
(119)
Let D and D' be two identical datasets except for one data point x, + x'y We have:
((D) - ((D")
(120)
sup fEV.gEQ
q(E log ex]
log
1
Σ
llg f(xs((1))9(E cxl Uz)
D
F:NEDD
(121)
q(E = e Ex log f( = ex]
Σ
- log flAl(149 = ex exlll/. 1/6)
D
(122)
sup fEV.gEQ
Σ
- log -logg((;(())))t = ex(ox,
Σ
log log fl13(111)) = ex v/1)
D
D
(123)
sup fEV.geQ
llog f(=,1111,) ex v7) log f11110/)) eal2/, 2/11
(124)
D
2B
(125)
[D]
20

---

Published as a conference paper at ICLR 2023
According to McDiarmid's inequality, for any 8 E (0, 1), with probability over ô, we have:
2 3
(DD) - Ep (D()) + BV [D]
(126)
Next we derive bound for Ep (D)].
Consider dataset DD' independently and identically drawn from g(X = P(X,Y) with the same size D. We notice that
g(E = ex log 1ogg ex] ex)e, ex)
(127)
q(E ex) - log logg/
(128)
E, [E,, lolog [( - M)lE = ex]l
(129)
E, |- f(r]()) =x,
(130)
Ed
Σ
log M ex(x(.n1))
(131)
D
Thus, Ep (D]] could be reformulated as:
Ep [D])
(132)
= LED
sup fEV.nEQ
E
Σ
log f(ad((1)) = (alkf/111)
(133)
D
Σ
log Mora((lii) = ek biz Mii)
(134)
[D]
R.DEED
< EE
sup SEV.ER
Ep'
log fl:l((())) = exlo1.08
(135)
D
1
Σ
log fxi((v))) eulers
(136)
D
ELED
<
ED,
sup SEV.gE
Σ
log forl()))) ex/2rt,11)
(137)
D
BL.VED
Σ log M exloboo1))
(138)
ED.D'
sup fEV
1
Σ
log flggf(ll(u))q = ex exl2x, Vi)
(139)
[
RiWEED
Σ o, log flr1))))((E = ex|z vi1
(140)
< EDD
sup
Σ
a, log [o:l))i)) en exx (c)
(141)
SEV,EE
D
P. ED
+ Ep' . 2% (g)),
sup
os log flz = exl2(,1/)
(142)
Q e [
(143)
21

---

Published as a conference paper at ICLR 2023
where o; independent Rademacher variables. Equation 137 follows from Jensen's inequality and the convexity of sup. Equation 139 holds due to the symmetry of log [ (3x )q(E e, 8c)
log ff ( ex |2 v/) and the argument that Radamacher variables preserve the expected sum of symmetric random variables with convex mapping (Ledoux & Talagrand (1991), Lemma 6.3).
Substituting Equation 143 to Equation 126 we have for any & E (0, 1), with a probability over 1 1-6 6,
Vf E, Vqg 2. the tollowing holds:
q(E = ex -- log [-][()]E =ec]
Σ
-
log [ = ex(o,,
(144)
[D]
FLBEED
2 2oog
s 23 (g B
(145)
D
Let fo = ing / (q(E = ck Bg f(r((u)) = ex]}.
Let f - arg inf o p E D2. - log [v: = ex 14 11))
Now we have
(E =cx) En -log fulzr()))lE ex
Σ
log (ari en( (146)
D
S 9E ()) ex) [D]
Σ
log =
(147)
1 = eb) - log [x] = ex
Σ log (a(tr)((n)) =xlo,,,)) (148)
D
Z.BIED
Combining Equation 144 and Equation
146
148
the lemma is proved.
Lemma 2. Assume Vx X. V.Vf V. log [v) E [[ B) where B 0. The definition of S and AN (N(9) follows from Lemma
Define fo arg inf p En.nE log f'0]()))(EE
ex at, V:).
Then for any q E Q, any (0, 1), with a probability over 8, we have
q(E ex))) =
Σ
log 110](1)) = ex(2s,81)
(149)
D
FLDDEED
log
< 28 () B
(150)
D
Proof. Similar to Lemma]
we could prove that
q(E ex) Hx ex)
Σ -
log fe(0(()))) ex/2is v4)
(151)
D
ED
2 2oo 1
< 20 2 (f + BV
(152)
[D]
where Gt: (,lg(=,)) v) = log 100 (v)n(E
E V. 2 2}.
According to the definition for the predictive family (Xu et al. (2020). Definition 1), Vf E V.
there exists f E V such that Vx E X, f(0] f'[x]. Thus, g C gr, and therefore A (0 < DD) (9t)- Substituting into Equation 151 the lemma is proved.
22

---

Published as a conference paper at ICLR 2023
Lemma (Xu et al. 2020), Theorem 1). Assume Vz X. J.Vf V. log e where > 0. Define a function class Gy = {glg(z y) = log (())()), V}. The definition of
NN follows from Lemma 7
Then for any E (0, 0.5), with a probability over 26, we have
[v(x-))) -) vY) s 4. (2
2 2oog
+ 2B
(153)
|D]
Finally we are prepared to prove Theorem
Proof of Theorem
We first bound the error empirical estimation with the sum of items in
Lemma 22
Hl 6x (X +Y) - - (x++)))) - ,;;
(154)
+ m0]
sup lv(X + Y18) -Ivl--lx----- +
sup
li(X + Y|E;D) - -i(xx - +Y;D)
LEE6K
LEE6K
(155)
5 sup Ilv(X Y18) sup liy(X Y|8:D) -1) Y:D)
(156)
|EE
EEEK
[x + Y18)- - -iv(X + 1lE: D) + Io(X +1) iv(x - v:D)
I sup
(157)
SEćK
sup g
Σ Ia(E = Hy)EE = ex) Q(E 1==
ex)]
(158)
K
Σ 1(E =
= - - =c:
(159)
+ Tv(X + V) -iv(X - vY:
(160)
Σ
sup aEQ
I(E = en) ex))Hv(Y E = =ex) -((E ex ex; D)
(161)
sup 9E@
a(E = (Y), =x) q(E ex) -ex:) ex:D)
(162)
+ Tv(X +V) -1 iv(x + Y:D)
(163)
K
Σ SEQ
sup
q(E = en Hy E =ex))
1
Σ
- log (o( = =xx
(164)
D
EID
+
Σ @
sup
q(== = H cx)
1
Σ
log (((0(()))(E =.l
D
(165)
+ Iv(( + +V- -ix(x + Y:D)
(166)
where fr arginf p En. (lg ex(
and 8% arg inf, T - logg(0(())(( = ex|o,, y.). for any q Q and 1 < k < K.
23

---

Published as a conference paper at ICLR 2023
For simplicity, let
Err, sup 1E
q(E = = = =
Σ
log (oll(waaa - ex vii
[D
(167)
Erry sup
qE ex) =ex)
log (100)))) en ea/rs,
(168)
D
ELEDED
Err* Ilv(X + +) - - ix + Y:D)/-
(169)
Then, by Lemma 123
2 2
Pr
HK Hk
4(K ()v)) + 2(K + +1)B
(170)
[D]
K
2 2og
S
Pr
EIra
+
Erry + Err" > 1) + 2(K + B)BBV
(171)
|D]
i=1
K
K
K
log 3 [D]
< Pr
ETTA
+
Errg + Err" >
4 pa (g + 40 (9 + 2(K + 1)B1
k=1
(172)
K
2 2log 1 |D}
2 log _ [D)
<
U
ErTK 12 + B1
+ U
Erry, >28 22( + B
1=1
1=1
(173)
2
Err* > 49h (9 -2B
(174)
|D]
K
2 log 1 [D]
K
2 2og A [D]
Pr
ErT, 22 + B
Pr
Ery. >2 1) + B
k=1
(175)
2 log D
+ Pr Err" > 4 (9)) 2B
(176)
< 2(K + 1)6.
(177)
Equation
172 is because of Gi = Iv. 96 C By and therefore (g)) & Rlp Rip (9t) < (Gv).
Hence,
2 2 pm
Pr
HHQ (X +Y - Y:D) < 4(K + 11 +(K + 12(K21())))) +1)B1
[D]
(178)
2 1- - 2(K + 1)6.
(179)
24

---

Published as a conference paper at ICLR 2023
H PROOF OF THEOREM 4
Proof of Theorem 4 The objective function of our IM algorithm is directly derived from the defi-
nition of empirical predictive heterogeneity in Definition For the regression task, we assume the predictive family as
V1 = (o: g|ø] = N(fo(z), f is the regression model and d is learnablc, = 1000=100000m 1.0(fixed)
(180)
where we only care about the output of the model and the noise scale of the Gaussian distribution is often ignored, for which we simply set 1.0 as fixed term. Then for each environment
e E supp(E"), the (X + E = =) becomes
Tx(X - E e) min B|I fa(X) = Var(Y\E"),
(181)
which corresponds with the MSE loss and the proposed regularizer in Equation 17 For the classifi cation task, the derivation is similar, and the regularizer becomes the entropy of Y in sub-population and the loss function becomes the cross-entropy loss.
I DISCUSSION ON DIFFERENCES WITH SUB-GROUP DISCOVERY
Subgroup discovery (SD, (Helal 2016)) is aimed at extracting "interesting" relations among differ- ent variables (X) with respect to target variable Y. Coverage and precision of each discovered group is the focus of such method. To be specific, it learns partition on P(X) such that some target label dominates within each group. The most siginficant gap between subgroup discovery and our predictive heterogeneity lies the pattern of distributional shift among clusters: for sub- group discovery, P(X) and P(Y) varies across subgroups but there is a universal P(Y While for predictive heterogeneity P(Y[X) differs across sub-population, which indicates diversified pre- diction mechanism. It is such disparity of prediction mechanism that inhibits the performance of universal predictive model on a heterogeneous dataset, which is the emphasis of OOD problem and group fairness.
We think sub-group discovery is more applicable for settings where the distributional shift is minor while high explainability is required, since generates simplified rules that people understand. Also, sub-group discovery methods suitable for the settings that only involve tabular data (typli- cally from relational database), where the input features have clear semantics. And proposed method could deal with general machine learning settings, including complicated data (e.g., image data) that involves representation learning. Also, when people have handle settings where data heterogeneity w.r.t. prediciton mechanism exists inside data, our method is more applicable. How- ever, both kinds of methods can be used help people understand data and make more reasonable decisions.
J DISCUSSION ON THE POTENTIAL FOR FAIRNESS
We find combining our measure with algorithmic fairness is an interesting and promising direction and we think our measure has the potential to deal with algorithmic bias. Our method could generate uub-populations with possibly different prediction mechanisms, which could do some help in the following aspects:
Risk feature selection: we could select features according to our predictive heterogeneity measure to see what features bring the largest heterogeneity. If they are sensitive features, people should avoid their effects, and if they are not, they could direct people to build better machine learning models.
Examine the algorithmic fairness: we could use the learned sub- aub-populations to examine whether a given algorithm is fair by calculating the performance gap across the sub-populations.
25


Reviews:
[{"review_id": "miJrzN_6jo", "reviewer": "Reviewer_QrpC", "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": "", "summary_of_the_paper": "The paper studies criteria and methods for finding data heterogeneities that affect building predictive models. The authors define a predictive homogeneity criterion and present an algorithm to find predictive heterogeneity in data. The authors further show predictive heterogeneity criterion can provide insights for subpopulation in agriculture, sociology and object recognition applications and its benefits for OOD generalisation.", "summary_of_the_review": "It is a technically sound paper with a modest empirical impact paper. I am not excited by the idea or the experimental results. \n", "strength_and_weaknesses": "S1: The paper is well written and easy to follow.\n\nS2: The theoretical analysis is substantial and looks sound. \n\nS3: The experimental case studies are from different areas.\n\n\nW1: The predictive heterogeneity criterion is interesting, but is incremental to the previous work (Xu et al.). In Definition 5, $\\Epsilon$ is unknown, and the authors' solutions for finding $\\Epsilon$  are in linear settings with independent noise. These assumptions will restrict the empirical impact of the work.  Also, since $\\Epsilon$ is unknown, setting a $k$ is challenging and the authors have not discussed how to set a $k$. \n\nW2: The empirical impact of the proposed method has not been shown to be significant. In three experimental demonstrations in subpopulation division, the number of k is 2 (the minimum possible number in heterogeneity discovery). Finding data heterogeneity is not a new topic and many works have been done. For example, in cancer subtype discovery in bioinformatics, many sophisticated methods have been proposed to deal with complex data types for many subgroups. The paper lacks discussions and comparisons with this type of work.\n\n====\n\nThe authors' detailed explanations and follow-up experiments clarify my doubts. ", "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": "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": "ly-mrp29ZBt", "reviewer": "Reviewer_uqaW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work defines predictive heterogeneity to quantify the heterogeneity in the data which influences the predictive performance. It proposes an algorithm which quantifies the mentioned heterogeneity and use that in examples to show how it can be helpful in understanding the subtleties of the data.", "summary_of_the_review": "This work is significant and it has very good applicability.", "strength_and_weaknesses": "The paper seems interesting and timely. Its implications are general and the experiments show its usefulness in providing insights as well as OOD generalization.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NgOwN6mbDw", "reviewer": "Reviewer_6bwF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studied the problem of measuring the data heterogeneity, i.e., the distributional diversity inside the data. It provided a novel data heterogeneity notion named predictive heterogeneity to explore how it affects the prediction of machine learning models. The crucial idea was to define the environment-conditional predictive V-information, which enabled the empirical estimation from finite samples. Its efficacy was then verified in various prediction tasks.", "summary_of_the_review": "The overall idea of this paper is interesting. It theoretically and empirically shows the quality of predictive heterogeneity in measuring data heterogeneity. I would like to improve my score if my concerns can be solved well.", "strength_and_weaknesses": "Strengths\n(1) It introduces a predictive heterogeneity to define the data heterogeneity quantitatively.\n(2) It theoretically analyzes how predictive heterogeneity works in homogenous and heterogeneous settings, and studies its empirical estimate with PAC bounds.\n(3) The experiments demonstrate the effectiveness of the proposed Information Maximization algorithm in various prediction tasks.\n\nWeaknesses\n(1) One major concern is the tightness of the predictive heterogeneity in Theorem 2. The first term Var[r]/E[r^2] involves both the variance of r and the expectation of r^2. It is unclear how the E[r^2] affects the approximation tightness.\n(2) In the introduction, it mentioned the data heterogeneity in the graph data. Could the proposed predictive heterogeneity be applied to learn the data heterogeneity of graph data? If so, how is the node interdependence reflected in the predictive heterogeneity?\n(3) It mentioned that \"Large H(P) indicates that the correlation between X and Y is enhanced by E\" in Section 3.1. What is H(P) here? Why does it indicate the correlation between X and Y?\n(4) How is the regularize U_V(W, Y_N) defined in the objective function of IM?\n(5) Theorem 4 shows the connection between the optimization of Eq. (15) and the empirical predictive heterogeneity. Would this result consider the approximation of Eq. (20-21)?\n(6) Appendix B showed that the proposed method is not sensitive to the choices of K. But K might be much smaller than the real number of sub-populations. In this case, the model performance might be improved by increasing the value of K.\n(7) In Table 1, it is unclear why EIL (Creager et al., 2021) fails to predict the simulated data.\n\n###########\nUpdate: The authors have addressed most of my concerns, and thus I would like to improve the score accordingly.", "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": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "6jmRrvP0dX", "reviewer": "Reviewer_E69D", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors analyzed the problem of identifying differing prediction mechanisms amongst different sub-populations that are likely to differ. They provided theoretical outlining of the mechanism and proposed a novel algorithm to explore the differences. The provided empirical analysis of the explore differences and provided discussions around the usefulness of extracted insights", "summary_of_the_review": "Overall the paper is well motivated, has strong theoretical under-pinnings, and backed by some level of empirical evidence. While the empirical part of the paper may be improved upon, the proposed method could be an interesting contribution, especially when considering the problems of algorithmic bias and model selection. \nApart from the aforementioned aspects, below are a few other aspects that the authors may consider\n\n- Provide a glossary of symbols used for the theoretical descriptions. Define symbols before using them (e.g. Section 2, Notations the probability triplet is not described)\n- Provide more details around the complexity of the proposed optimization. As the authors have noted, MI is notoriously intractable - providing more details on the final complexity would be useful for consideration of the algorithm to a wide range of problems\n- Considering expanding the discussions around the benefits to quantifying/debiasing w.r.t algorithmic fairness of models\n\nEdit: The authors have addressed a number of concerns regarding validity in their response", "strength_and_weaknesses": "There are several strong aspects of the paper as follows\n\n- The authors proposed a formal definition of predictive heterogeneity by starting from the first principles of probablity and using concepts from information theory. This allows the authors to immediately connect certain strong theoretical results for the proposed definition\n- The theoretical analysis around linear cases is interesting. While not being the most real-world scenario, this analysis allows readers to gain an intuitive association with the properties of the methods and deliberate on its applicability\n- The empirical exploration around the yield datasets is motivating. Similarly the other empirical evaluations lends substance to the otherwise theoretical paper\n\nThe paper could still be improved upon certain manner\n- The presentation of the paper could be improved upon. While the background around information theory is useful, more intuitions around the theoretical propositions could improve the readability of the paper\n- While acknowledging that the paper is theoretical, the empirical results do not sufficiently motivate the \"correctness\" of the method. The OOD experiments are interesting, however, It may be useful to devise some sort of quantitative/qualitative measure of extracted sub-groups from the 3 example tasks. There is a huge body of literature around sub-group discovery that can provide pointers for this (e.g using a synthetic dataset)\n- The applicability of the algorithm is also not well discussed. There are many sub-group discovery methods. Are there situations where other methods are more applicable? Can you provide more details on the computational complexity ", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster