M EASURING THE P REDICTIVE H ETEROGENEITY
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
Big data bring great opportunities to modern society and promote the development of machine learning, 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, 2017).
In general, data heterogeneity, as a fundamental property of big data, refers to any diversity inside data, including the diversity of data sources, data generation mechanisms, sub-populations, data structures, 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), leading to poor out-of-distribution 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 sub-populations 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 is Simpson’s paradox, which brings false discoveries to the social research (Wagner, 1982; Hern´an et al., 2011).
Despite its widespread existence, due to its complexity, 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, a heterogeneity graph refers to various types of nodes and edges in the graph (Wang et al.
More recently, in machine learning, several works of causal learning (Peters et al., 2016; Arjovsky et al., 2019; Koyama & Yamaguchi, 2020; Creager et al., 2021) and robust learning (Sagawa et al., 2019) leverage heterogeneous data from multiple environments to improve the out-of-distribution generalization ability.
However, previous works have not provided a precise definition or sound quantification. In this work, from the perspective of prediction power , we propose the predictive heterogeneity, a new type 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 on the data heterogeneity that affects the prediction of machine learning models, 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-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 it under typical cases of data heterogeneity (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 to enhance the out-of-distribution generalization performances of machine learning models, which is verified with both simulated and real-world data.
P RELIMINARIES ON M UTUAL I NFORMATION AND P REDICTIVE V -I NFORMATION
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 , P ) , define random variables X : S → X and Y : S → Y where X is the covariate space and Y is the target space.
Accordingly.
x ∈ X denotes the covariates, and y ∈ Y denotes the target.
Denote the set of random categorical variables as C =
{ C : S
→ N | supp ( C ) is finite } .
Additionally, P ( X ) , P ( Y ) denote the set of all probability measures over the Borel algebra on the spaces X , Y respectively.
H ( · ) denotes the Shannon entropy of a random variable, and H ( ·|· ) 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:
[Equation]
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., 2020)) .
Let Ω =
{ f : X ∪ {∅} → P ( Y ) } .
We say that V ⊆ Ω is a predictive family if it satisfies:
[Equation]
A predictive family contains all predictive models that are allowed to use, which forms computational or statistical constraints.
The additional condition in Equation 2 means that the predictor can always ignore the input covariates ( x ) if it chooses to (only use ∅ ).
Definition 2 (Predictive V -information (Xu et al., 2020)) .
Let X, Y be two random variables taking values in X × Y and V be a predictive family.
The predictive V -information from X to Y is defined as:
[Equation]
where H V ( Y |∅ ) , H V ( Y | X ) are the predictive conditional V -entropy defined as:
[Equation]
[Equation]
Notably that f ∈ V is a function X ∪ {∅} → P ( Y ) , so f [ x ] ∈ P ( Y ) is a probability measure on Y , and f [ x ]( y ) ∈ R is the density evaluated on y ∈ Y .
H V ( Y |∅ ) is also denoted as H V ( Y ) .
Compared with the standard mutual information, the predictive V -information restricts the computational power and is much easier to estimate in high-dimensional cases.
When the predictive family
V contains all possible models, i.e. V = Ω , it is proved that I V ( X → Y )
=
I ( X ; Y ) (Xu et al., 2020).
P REDICTIVE H ETEROGENEITY
In this paper, from the machine learning perspective, we propose a quantitative definition of the data heterogeneity that affects decision making, named Predictive Heterogeneity, which is easy to integrate with machine learning algorithms and could help analyze big data and develop more rational algorithms.
3. I NTERACTION H ETEROGENEITY
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 × Y .
Denote the set of random categorical variables as C , and take its subset E ⊆ C .
Then E is an environment set iff there exists E ∈ E such that X, Y ⊥⊥ E .
E ∈ E is called an environment variable.
The interaction heterogeneity between X and Y w.r.t.
the environment set E is defined as:
[Equation]
Each environment variable E represents a stochastic ‘partition’ of X × Y , 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 ; X |E ) − I ( Y ; X ) is called the interaction information , which measures the influence of the environment variable E on 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 can be gained from involving or uncovering the environment variable E .
Intuitively, large H E ( 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. P REDICTIVE H ETEROGENEITY
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 scenarios.
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 × Y and E be an environment variable.
The conditional predictive V -information is defined as:
[Equation]
where H V ( Y |∅ , E ) and H V ( Y | X, E ) are defined as:
[Equation]
Intuitively, the conditional predictive V -information measures the weighted average of predictive Vinformation among environments.
And here we are ready to formalize the predictive heterogeneity measure.
Definition 5 (Predictive Heterogeneity) .
Let X , Y be random variables taking values in X ×Y and E be an environment set.
The predictive heterogeneity for the prediction X
→ Y with respect to E is defined as:
[Equation]
where I V ( X → Y ) is the predictive V -information following from Definition 2.
Leveraging the predictive V -information, the predictive heterogeneity defined in Equation 10 characterizes the maximal additional information that can be used by the prediction model when involving 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 taking values in X × Y , V be a function family, and E , E 1 , E 2 be environment sets.
Monotonicity :
If E 1 ⊆ E 2 , H E 1 V ( X → Y ) ≤ H E 2 V ( X → Y ) .
Nonnegativity : H E V ( X → Y ) ≥ 0 .
Boundedness : H E V ( X → Y ) ≤ H V
( Y | X ) .
Corner Case : If the predictive family V is the largest possible predictive family that includes all possible models, i.e. V = Ω , we have H E ( X, Y )
=
H E Ω ( 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 guarantees if the complexity of V is bounded (e.g., its Rademacher complexity).
3. T HEORETICAL P ROPERTIES IN L INEAR C ASES
In this section, we analyze the theoretical properties of the predictive heterogeneity in multiple linear settings, including (1) a homogeneous case with independent noises and (2) heterogeneous cases with endogeneity brought by selection bias and hidden variables.
Under these typical settings, we could approximate the analytical forms of the proposed measure and the conclusions provide insights for general cases.
Firstly, under a homogeneous case with no data heterogeneity, Theorem 1 proves that our measure is bounded by the scale of label noises (which is usually small) and reduces to 0 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 ⊥ ϵ |E , since in common prediction tasks, the independent noises are unknown and unrealistic to be exploited for the inference of latent environments E .
In Appendix D, we add some numerical results under the setting of Theorem 1 to show that the empirical estimated value approaches 0 as the sample size grows.
Theorem 1 (Homogeneous Case with Independent Noises) .
For a prediction task X → Y where
X , Y are random variables taking values in R n × R , consider the data generation process as Y =
g
( x ) + ϵ, ϵ ∼ N (0 , σ 2 ) where g : R n → R is a measurable function.
For a function class G such that g ∈ G , define the function family as V G =
{ f | f [ x ] = N ( ϕ ( x ) , σ 2 V ) , ϕ ∈ G , σ V ∈ R + } .
With an environment set E , we have H E V G ( X → Y ) ≤ πσ 2 . 2) Take n = 1 and g ( x ) = βx , β ∈ R .
Assume
E [ X ] = 0 and E [ X 2 ] exists.
Given the function family V σ = { f | f [ x ] = N ( θx, σ 2 ) , θ ∈ R , σ fixed } and the environment set E = {E|E ∈ C , | supp ( E ) | = 2 , X ⊥ ϵ |E} .
We have H E V σ ( 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; Arjovsky et al., 2019).
To begin with, in Theorem 2, we consider the prediction task X
→ Y with X , Y taking values in R× R .
Let X =
[ S, V ] T .
The predictive family is specified as:
[Equation]
And the data distribution P ( X, Y ) is a mixture of latent sub-populations, which could be formulated by an environment variable E ∗ ∈ C such that P ( X, Y )
= (cid:80) e ∈ supp ( E ∗ ) P ( E ∗ =
e ) P ( X, Y |E ∗ = e ) .
For each e ∈ supp ( E ∗ ) , P ( X, Y |E ∗ = e ) is the distribution of a homogeneous sub-population.
Note that the prediction task is to predict Y with covariates X , and the sub-population structure is latent.
That is, P ( E ∗ | X, Y ) is unknown for models.
In the following, we derive the analytical forms of our measure under the one typical case.
Theorem 2 (Endogeneity with Selection Bias) .
For the prediction task X
=
[ S, V ] T → Y with a latent environment variable E ∗ , the data generation process with selection bias is defined as:
[Equation]
where f : R → R and r, σ : supp ( E ∗ ) → R are measurable functions.
β ∈ R .
Assume that
E [ f ( S ) S ] = {+
+} and there exists L > 1 such that Lσ 2 ( E ∗ ) < r 2 ( E ∗ ) E [ f 2 ] .
For the predictive family defined in equation 11 and the environment set E = C , the predictive heterogeneity of the prediction task [ S, V ] T → Y approximates to:
[Equation]
bounded
[Equation]
And R ( r ( E
[Equation]
) , σ ( E
[Equation]
) , f
[Equation]
=
E [
[Equation]
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 σ ( E ∗ ) varies) and brings about data heterogeneity.
Here E [ f ( S ) S ] = 0 indicates a model misspecification, 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 K characterizes the strength of the spurious correlation between V and Y . Larger K means V could provide more information for prediction.
From the approximation in Equation 13, we can see that our proposed predictive heterogeneity is dominated by two terms: (1) Var [ r ( E ∗ )] / E [ r 2 ( E ∗ )] characterizes the variance of r ( E ∗ ) among subpopulations; (2) E [ f 2 ( S )] reflects the strength of model misspecifications.
These two components account for the two sources of the data heterogeneity in the setting of selection bias, which validates the rationality of our proposed measure.
According to the theorem, the more various r ( E ∗ ) among the sub-populations and stronger model misspecifications, the larger the predictive heterogeneity should be.
In general, Theorem 1 and 2 indicate that (1) our proposed measure is insensitive to the homogeneous cases and (2) for the two 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.
We also derive the analytical form for the case of endogeneity with hidden variable in Theorem 5, which has the similar conclusions and is left in Appendix.
3. PAC G UARANTEES FOR P REDICTIVE H ETEROGENEITY E STIMATION
Defined under explicit computation constraints, our Predictive Heterogeneity could be empirically estimated with guarantees if the complexity of the model family V is bounded.
In this work, we provide finite sample generalization bounds with the Rademacher complexity.
First, we 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 ×Y , a dataset D is independently and identically drawn from the population
such that D
= { ( x i , y i )
Ni =1 ∼ X, Y } .
Given the predictive family V and the environment set E
K = {E|E ∈ C , supp ( E ) =
K } where K ∈ N + is the number of environments, the empirical predictive heterogeneity
ˆ H E K V
( X → Y ; D ) with respect to D is readily obtained by estimating
H E K V ( X → Y ) on D with expectations replaced by statistics of finite samples.
The formal definition is placed in Definition 7.
Theorem 3 (PAC Bound) .
Consider the prediction task X
→ Y where X , Y are random variables taking values in X × Y .
Assume that the predictive family V satisfies ∀ x ∈ X , ∀ y ∈ Y , ∀ f ∈ V , log f
[ x ]( y ) ∈ [ − B, B ] where B > 0 .
For given K ∈ N , the environment set is defined as E K = {E|E ∈ C , supp ( E ) = K } .
Let Q be the set of all probability distributions of X , Y , E where E ∈ E K .
Take an e ∈ supp ( E ) and define a function class G V =
{ g | g ( x, y )
=
log f
[ x ]( y ) Q ( E = e | x, y ) , f ∈
V , Q ∈ Q} .
Denote the Rademacher complexity of G with N samples by R N ( G ) .
Then for any δ ∈ (0 , 1 / (2 K + 2)) , with a probability over 1 − 2( K + 1) δ , for dataset D independently and identically drawn from X , Y , we have:
[Equation]
where R |D| ( G V ) =
O ( |D| − 1 2 ) (Bartlett & Mendelson, 2002).
A LGORITHM
To empirically estimate the predictive heterogeneity in Definition 6, we derive the Information Maximization (IM) algorithm from the formal definition in Equation 33 to infer the distribution of E that maximizes the empirical predictive heterogeneity ˆ H E K V
( X → Y ; D ) .
Objective Function.
Given dataset D =
{ X N , Y N } = { ( x i , y i ) } Ni =1 , denote supp ( E ) = { e 1 , . .
, e K } , we parameterize the distribution of E| ( X N , Y N ) with weight matrix W ∈ W K , where
K is the pre-defined number of environments and W K =
{ W : W ∈ R N × K + and W 1 K = 1 N } is the allowed weight space.
Each element w ij in W represents P ( E =
e j | x i , y i ) (the probability of the i -th data point belonging to the j -th sub-population).
For a predictive family V , the solution to the supremum problem in the Definition 7 is equivalent to the following objective function:
[Equation]
where f θ : X → Y denotes a predicting function parameterized by θ , ℓ V ( · , · ) :
Y × Y → R represents a loss function and U V ( W, Y N ) is a regularizer.
Specifically, f θ , ℓ V and U V 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:
[Equation]
predicting
[Equation]
learnable
[Equation]
The corresponding loss function is ℓ V 1 ( f θ ( X ) , Y ) =
( f θ ( X ) − Y ) 2 , and U V 1 ( W, Y N ) becomes
[Equation]
where Y jN denotes the mean value of the label Y given E = e j and U ( W, Y N ) calculates the variance of Y jN among sub-populations e 1 ∼ e K .
(2) For the classification task , the predictive family is typically modeled as:
[Equation]
where c is the class number and ∆ c denotes the c -dimensional simplex.
Here each model in the predictive family V 2 outputs a discrete distribution in the form of a c -dimensional simplex.
In this case, the corresponding loss function ℓ V 2 ( · , · ) is the cross entropy loss and the regularizer becomes
U V 2 ( W, Y N ) =
[Equation]
j
[Equation]
1 N ( (cid:80) N i =
[Equation]
w ij )
H ( Y j N ) , where H ( Y j N ) is the entropy of Y given E = e j .
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 V here)
[Equation]
where ∇ θ j R (cid:12)(cid:12) θ ∗ j ∇ W θ ∗ j ≈ ∇ θ j R (cid:12)(cid:12) θ tj (cid:88) h ≤ t 
[Equation]
, for j = 1 , . . . ,
K.
(21) where ℓ ( f θ j ( x i ) , y i )]
N × K i,j is an N × K matrix in which the ( i, j ) -th
element is ℓ ( f θ j ( x i ) , y i ) .
Here
Equation 20 approximates θ ∗ j by θ tj from t steps of inner loop gradient descent and Equationperforms 1-step truncated backpropagation (Shaban et al., 2019; Zhou et al., 2022).
Our information maximization algorithm updates W by projected gradient descent as:
[Equation]
W
[Equation] [Equation]
Theoretical Analysis.
Then we prove that minimizing Equation 15 exactly finds the supremum w.r.t.
E in the Definition 7 (formal) of the empirical predictive heterogeneity.
And we also demonstrate its relationship with the expectation maximization (EM) algorithm.
Theorem 4 (Justification of the IM Algorithm) .
For the regression task with predictive family V 1 and classification task with V 2 , the optimization of Equation 15 is equivalent to the supremum problem of the empirical predictive heterogeneity ˆ H E K V 1 ( X → Y ; D ) , ˆ H E K V 2 ( X → Y ; D ) respectively in
Equation 33 with the pre-defined environment number K (i.e. supp ( E ) = K ).
Remark 1 (Difference from Expectation Maximization) .
The expectation maximization (EM) algorithm is to infer latent variables of a 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 to the regularizer U V in our objective function, the EM algorithm cannot efficiently solve our problem, and therefore we adopt bi-level optimization techniques.
E XPERIMENTS 5. P ROVIDE INSIGHTS FOR THE SUB - POPULATION DIVISION
The predictive heterogeneity could provide insights for the sub-population division and benefit decision making, and we illustrate this in prediction tasks of various fields, including agricultural research, sociological research, and object recognition.
From the illustrative examples, we show that the learned sub-population division is highly explainable and relevant to decision-making.
Example: Agriculture It is known that the climate affects crop yields and crop suitability (Lobell et al., 2008).
We leverage the data from the NOAA database which contains daily weather from weather stations around the world.
Following Zhao et al.
(2021), we summarize the weather sequence of the year 2018 into summary statistics, including the average yearly temperature, humidity, 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 illustration, 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 = 2 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 to it.
In Figure 1 (a), we plot the real distribution map of wheat and rice planting areas, and in Figure 1 (b), we plot the learned two sub-populations of our
IM algorithm.
From the results, we surprisingly find the division given by our algorithm is quite similar to the real division of the two crops, indicating the rationality of our measure.
For the areas that are not similar (e.g. Tibet Plateau in Asia), we think it is due to 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 a person is 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.
In Figure 2 (a), we plot summary statistics for the two sub-populations, where the main difference lies in the capital gain.
In Figure 2 (b), we plot the feature importance given by linear models for the two sub-populations, and we find that for people with a 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 discrimination.
Our results correspond with the results in Zhao et al.
and can help us identify potential inequality in decision-making.
For example, our results indicate the potential discrimination for low capital gain people, which could further promote algorithm design and improve policy fairness.
Example: Object Recognition Finally, we use the Waterbird dataset (Sagawa et al., 2019), which is widely used to evaluate the model’s robustness in the robust learning field.
The task is to recognize waterbirds or landbirds.
However, the image backgrounds are spuriously correlated with the target 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
ResNet18 and set K = 2 in our IM algorithm.
In Figure 3, to show the learned sub-populations of our method, we randomly sample 50 images for each class (waterbird or landbird) and each learned sub-population.
In sub-population 1, 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 in the two sub-populations, the spurious correlation between the object and background is inverse.
And the learned 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 to eliminate the influence of backgrounds on model predictions.
5. B ENEFIT OOD G ENERALIZATION
The predictive heterogeneity could benefit the out-of-distribution (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., 2021)) which infers the environments for learning invariance.
Then we compare with the well-known KMeans algorithm, which is the most popular
[Table]
[Figure]
clustering algorithm.
For our IM algorithm and KMeans, we involve three algorithms as backbones to leverage the learned sub-populations, including sub-population balancing and invariant learning methods.
The sub-population balancing simply equally weighs the learned sub-populations.
invariant risk minimization (IRM, (Arjovsky et al., 2019)) and inter-environment gradient 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, T, V ] T ∈ R 10 consist of stable features S ∈ R 5 , noisy features T ∈ R 4 and the spurious feature V ∈ R :
S ∼ N (0 , 2 I 5 ) , T ∼ N (0 , 2 I 4 ) , Y =
θ TS S + h ( S )+ N (0 , 0 . 5) , V ∼ Laplace ( sign ( r ) · Y, 1 / (5 ln | r | ))
(23) where θ S ∈ R 5 is the coefficient and h ( S ) =
S 1 S 2 S 3 is the nonlinear term.
| r | > 1 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-populations and is controlled by the factor r .
Intuitively, sign ( r ) controls whether the spurious correlation between V and Y is positive or negative.
And | r | 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 r = 1 .
(i.e. strong positive spurious correlation) and the minor group contains 5% data with r = − 1 .
9 (i.e. strong negative spurious correlation).
In testing , we test the performances on the two groups respectively, and we also set r = − 2 .
3 and r = − 2 .
7 to simulate stronger distributional shifts.
We use linear regression and set
K = 2 for all methods, and we report the mean-square errors (MSE) of all methods.
Data Generation of Colored MNIST Following Arjovsky et al.
(2019), we design a binary classification task constructed on the MNIST dataset.
Firstly, digits 0 ∼ 4 are labeled Y = 0 and digits 5 ∼ 9 are labeled Y = 1 .
Secondly, noisy labels ˜
Y are induced by randomly flipping the label Y with a probability of 0.2.
Then we sample the colored id V spurious correlated with ˜ Y as
V = (cid:26) +˜ Y , with probability r, − ˜ Y , 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 ˜ Y and for the rest 15%, the spurious correlation becomes negative, which causes data heterogeneity w.r.t.
V and ˜ Y .
In testing , we set r =
( strong negative spurious correlation ), bringing strong shifts between training and testing.
Analysis From the results in Table 1, 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 in Figure 4 and 5.
From Figure 4, compared with KMeans and
EIIL, our predictive heterogeneity exploits the spurious correlation between V and Y , and enables the backbone algorithms to eliminate it.
From Figure 5, the learned sub-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.
C ONCLUSION
We define the predictive heterogeneity, as the first quantitative formulation of the data heterogeneity that affects the prediction of machine learning models.
We demonstrate its theoretical properties and show that it benefits the out-of-distribution generalization performances.
A F ORMAL D EFINITION OF E MPIRICAL P REDICTIVE H ETEROGENEITY
In this section, we derive the explicit formula for the empirical estimation of the predictive heterogeneity which is described in Definition 6.
The dataset D = { ( x i , y i ) } |D| i =1 is independently and identically drawn from the population X, Y .
Given a function family V and an environment set E K , let Q be the set of all probability distributions of X , Y , E where E ∈ E K .
For given E , denote supp ( E ) = { ( e k ) Kk =1 } .
The empirical predictive heterogeneity ˆ H E K V
( X → Y ; D ) is given by:
[Equation] [Equation]
Specifically,
[Equation]
The explicit formula for ˆ Q ( E = e k ) ˆ H V ( Y |E =
e k ; D ) , ˆ H V ( Y | X ; D ) and ˆ H V
( Y ; D ) could be similarly derived.
Here we are ready to formally define the empirical predictive heterogeneity.
Definition 7 (Empirical Predictive Heterogeneity (formal)) .
For the prediction task X → Y with X ,
Y taking values in X × Y , a dataset D is independently and identically drawn from the population
such that D
= { ( x i , y i )
Ni =1 ∼ X, Y } .
Given the predictive family V and the environment set E
K = {E|E ∈ C , supp ( E ) =
K } where K ∈ N , let Q be the set of all probability distributions of
X , Y , E where E ∈ E K .
The empirical predictive heterogeneity ˆ H E K V
( X → Y ; D ) with respect to D is defined as:
[Equation] [Equation]
where
[Equation]
B S ENSITIVITY OF K
In the experiments of Section 5, we set the K = 2 for easy illustrations.
In this section, we add the results of choosing different K s 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 .
In Figure 6, we show the out-of-distribution 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 1 , our performances significantly outperforms all the baselines.
Also, we add one more experiment to show that (1) when the chosen K is smaller than the groundtruth, 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).
Experiment Setting: The input features X =
[ S, T, V ] ∈ R 10 consist of stable features S ∈ R 5 , noisy features T ∈ R 4 and the spurious feature V ∈ R :
[Equation]
and we generate the spurious feature via: [Equation] where θ eV varies across sub-populations and is dependent on which sub-population the data point belongs to. In training, we sample 8000 data points from e 1 with θ 1 V = 3 . 0 , 1000 points from ewith θ 2 V = − 1 . 0 , 1000 points from e 3 with θ 3 V = − 2 . 0 and 1000 points from e 4 with θ 4 V = − 3 . 0 .
Therefore, the ground-truth number of sub-populations is 4.
In testing, we test the performances on e 4 with θ 4 V = − 3 . 0 , which has strong distributional shifts from training data.
The average MSE over 10 runs are shown in Figure 7.
From the results, we can see that when K is smaller than the ground-truth, increasing K benefits the OOD generalization performance, and when K is 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-truth, the increasing speed will slow down, which could direct people to choose an 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 θ 1 , . .
, θ K to judge whether K is much larger than the ground-truth, i e., if two resultant models are quite similar, K may be too large (divide one 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-populations.
For detailed analysis on the best choice of K , we leave it for future work.
C A DDITIONAL T HEORETICAL R ESULTS
Theorem 5 (Endogeneity with Hidden Variables) .
For the prediction task X
=
[ S, V ] T → Y with a latent environment variable E ∗ , the data generation process with hidden variables is defined as:
[Equation]
where f : R → R and r, σ : supp ( E ∗ ) → R are measurable functions.
β ∈ R .
Assume that (1)
E [ f ( S ) S ] = 0 and (2) E [ σ 2 ( E ∗ )] ≪ E [ f 2 ( S )]
E [ S 2 ] .
For the predictive family defined in equation 11 and the environment set E = C , the predictive heterogeneity of the prediction task [ S, V ] T → Y approximates to:
[Equation]
And the approximation error is bounded by (1 / 2) · max( σ 2 Y , E [ σ ( E ∗ ) 2 /r
( E ∗ ) 2 ]) .
D A DDITIONAL NUMERICAL RESULTS
Here we provide some additional numerical results of our linear examples in Section 3.3.
In the left sub-figure, we plot the estimated predictive heterogeneity under the setting of Theorem 1 where the analytical solution is equal to 0.
From the results, we can see that with the growing of sample size, the estimated value of our IM algorithm is approaching to 0 (note that the y -axis is ln( estimated value ) ).
In the middle sub-figure, for the setting in Theorem 2, we plot the theoretical approximation, empirical approximation (finite sample case) and the estimated value of the predictive heterogeneity under different ratios between the majority and the minority (which controls the Var [ r ( E ∗ )] in Equation 13).
And the right sub-figure plots the same values under the setting in Theorem 5.
From these two figures, we can see that (1) the empirical approximation under finite samples lies closely to the theoretical approximation, which is supported by our generalization bounds in Theorem 3; (2) the estimated value of our IM algorithm is closely to the theoretical approximation,, which demonstrates the accuracy of our approximation algorithm in Equation 20 and 21.
Also, as the ratio changes from 4 : 1 to 1 : 1 , the data heterogeneity is increasing, and our predictive heterogeneity is also increasing, which is controlled by the term Var ( r ( E ∗ )) in Equation 13 and 39.
E R ELATED W ORK
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 to various types of nodes and edges (Wang et al.
More recently, in machine learning, several works of causal learning (Peters et al., 2016; Arjovsky et al., 2019; Koyama & Yamaguchi, 2020; Creager et al., 2021) and robust learning (Sagawa et al., 2019) leverage heterogeneous data from multiple environments to improve the out-of-distribution generalization ability.
Specifically, invariant learning methods (Arjovsky et al., 2019; Koyama &
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 distributionally robust optimization field, Sagawa et al.
; Duchi et al.
(2022) propose to optimize the worst-group prediction error to guarantee the OOD generalization performance.
However, in machine learning, previous works have not provided a precise definition or 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 et 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.
F P ROOF OF P ROPOSITION
Proof of Proposition 1.
Monotonicity :
Because of E 1 ⊆ E 2 ,
[Equation]
Nonnegativity :
According to the definition of the environment set, there exists E 0 ∈ E such that for any e ∈ supp ( E ) ,
X, Y |E =
e is identically distributed as X, Y .
Thus, we have
[Equation]
Specifically,
[Equation]
Similarly, H V ( Y |∅ , E 0 ) =
H V ( Y |∅ ) .
Thus, H E V ( X → Y ) ≥ 0 .
3. Boundedness :
First, we have
[Equation]
by noticing that E y ∼ Y | x [ − log f
[ x ]( y )] is the cross entropy between Y | x, e and f [ x ] .
Next,
[Equation]
where Equation 52 is due to Jensen’s inequality.
Combing the above inequalities,
[Equation]
4. Corner Case :
According to Proposition 2 in Xu et al.
[Equation]
By taking random variables R
[Equation]
S identically distributed as X
[Equation]
Y |E =
e for e ∈ supp ( E )
[Equation]
we have
[Equation]
Thus,
[Equation]
Similarly, we have H Ω ( Y |∅ , E )
=
H ( Y |E ) .
Thus,
H E Ω ( X → Y ) =
sup E∈ E
[ H Ω ( Y |∅ , E ) − H Ω
( Y | X, E )] − [ H Ω ( Y |∅ )
− H Ω ( Y | X )] (63) =
sup E∈ E [ H ( Y |E ) − H
( Y | X, E )] − [ H ( Y ) − H
( Y | X )] (64) =
sup E∈ E I ( Y ; X |E ) − I
( Y ; X ) (65) =
H E ( X, Y ) .
G P ROOF OF T HEOREM
Proof of Theorem 1.
[Equation]
[Equation] [Equation]
(69) =
πσ
[Equation]
Equation 68 holds by taking f [ x ] = N ( g ( x ) , 12 π ) .
[Equation]
Given the function family V σ = { f | f [ x ] = N ( θx, σ 2 ) , θ ∈ R , σ fixed } , by expanding the Gaussian probability density function in the definition of predictive V -information, it could be shown that
[Equation]
where the predictive V -information is proportional to Mean Square Error subtracted by the variance of target, by a coefficient completely dependent on σ .
The minimization problem is solved by
[Equation]
Substituting k into eq.71,
[Equation]
Denote supp ( E ) =
{E 1 , E 2 } .
Let Q be the joint distribution of ( X, ϵ, E ) .
Let Q ( E 1 ) =
α and
Q ( E 2 ) =
1 − α be the marginal of E .
Abbreviate Q ( X, ϵ |E =
E 1 ) by P 1 ( X, ϵ ) and Q ( X, ϵ |E
=
E 2 ) by P 2 ( X, ϵ ) .
Similar to 71,
[Equation]
For E = E 1 , the minimization problem is solved by
[Equation]
Thus,
[Equation]
Similarly, we have
I V
[Equation] [Equation]
Notably, E P 1 [ X 2 ] and E P 2 [ X 2 ] are constrained by α and E
[ X 2 ] .
[Equation]
Similarly,
[Equation]
The moments of P 2 could thereafter be represented by those of P 1 .
[Equation]
[Equation]
Substituting to eq.80,
[Equation]
Thus,
[Equation]
Assuming X ⊥ ϵ | E ,
[Equation] [Equation]
From Proposition
[Equation]
we have H E V σ ( X → Y )
[Equation]
Thus, H E V σ ( X → Y ) = 0 .
H P ROOF OF L INEAR C ASES (T HEOREM 2 AND 5)
Proof of Theorem 2.
For the ease of notion
[Equation]
we denote the r ( E ∗ ) as r e , σ ( E ∗ ) as σ e , and σ ( E ∗ ) · ϵ v as ϵ e .
And we omit the superscript C of H CV .
Firstly, we calculate the H V [ Y |∅ ] as:
[Equation]
[Equation] [Equation]
Therefore
[Equation]
we have
[Equation]
As for H V [ Y | X ]
[Equation]
we have
[Equation] [Equation]
[Equation]
(98) where we let h S
= h S − β here.
Then we have
[Equation] [Equation]
notably that here for e i , e j ∈ supp ( E ∗ ) , we assume P e
i
S, Y )
=
P e j ( S, Y ) (we choose such E ∗ as one possible split).
And the solution of h S , h V is
[Equation]
According to the assumption that E [ f ( S ) S ] =
[Equation]
we have
[Equation]
Therefore
[Equation]
we have
[Equation]
Note that here we simply set σ = 1 in the main body.
And we have:
[Equation]
The approximation error is bounded by 12 max( σ 2 Y , R ( r ( E ∗ ) , σ ( E ∗ ) , E [ f 2 ])) , and R ( r ( E ∗ ) , σ ( E ∗ ) , E [ f 2 ]) is defined as:
[Equation]
Proof of Theorem 5.
As proved above, we have
[Equation]
For
the model
misspecification case , we further assume that (1) E [ f ( S ) S ] = [-
-]0 and (2) E [ σ 2 e ] ≪ E [ f 2 ( S )]
E [ S 2 ] , and then we have
[Equation]
and for the heterogeneity, we have
[Equation]
• Without the model misspecification, we assume that f ≡ [Equation] and then we have
[Equation]
and for the heterogeneity we have
[Equation]
I P ROOF OF THE E RROR B OUND FOR F INITE S AMPLE E STIMATION (T HEOREM 3)
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 ∀ x ∈ X , ∀ y ∈ Y , ∀ f ∈ V , log f
[ x ]( y ) ∈ [ − B, B ] where B > 0 .
Define a function class G k V = { g | g ( x, y )
=
log f
[ x ]( y ) q ( E = e k | x, y ) , f ∈ V , q ∈
Q} .
Denote the Rademacher complexity of G with N samples by R N ( G ) .
Define ˆ f k =
[Equation]
Then for any q ∈ Q , any δ ∈ (0 , 1) , with a probability over 1 − δ , we have
[Equation]
Proof.
Apply McDiarmid’s inequality to the function Φ( D ) which is defined as:
[Equation]
Let D and D ′ be two identical datasets except for one data point x j ̸ = x ′ j .
We have:
[Equation]
x ′ i ,y ′ i ∈D ′
− log f [ x ′ i ]( y ′ i ) q ( E = e k | x ′ i , y ′ i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (124) ≤ sup f ∈V ,q ∈Q (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:12) |D| (cid:88) x i ,y i ∈D − log f [ x i ]( y i ) q
( E = e
[Equation]
| x i , y
[Equation]
) − 1 |D ′ | (cid:88) x ′ i ,y ′ i ∈D ′ − log f [ x ′ i ]( y ′ i )
q ( E = e k | x ′ i , y ′ i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:12) (125) = sup f ∈V ,q ∈Q |D| (cid:12)(cid:12)
[Equation]
[Equation]
[Equation]
( E =
[Equation]
[Equation]
f
[Equation]
[Equation]
]( y ′
[Equation]
[Equation]
( E = e
[Equation]
x ′
[Equation] , y ′ j ) (cid:12)(cid:12) (126) ≤ 2 B |D| . According to McDiarmid’s inequality, for any δ ∈ (0 , 1) , with a probability over 1 − δ , we have:
[Equation]
Next we derive a bound for E D [Φ( D )] .
Consider a dataset D ′ independently and identically drawn from q ( X, Y )
=
P ( X, Y ) with the same size as D .
We notice that
[Equation]
Thus
[Equation]
E D [Φ( D )] could be reformulated as:
[Equation]
f [ x i ]( y i ) q ( E = e k | x i , y i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (136) ≤ E D  
[Equation]
E D ′ (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) − |D ′
| (cid:88) x ′ i ,y ′ i ∈D ′ − log f [ x ′ i ]( y ′ i )
q ( E = e k | x ′ i , y ′ i ) (137) −|D| (cid:88) x i ,y i ∈D − log f
[ x i ]( y i ) q
( E = e k | x i , y i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (138) ≤ E D , D ′ 
[Equation]
∈Q |D| (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:88) x i ,y i ∈D log f
[ x i ]( y i ) q ( E = e k | x i , y i ) (139) − (cid:88) x ′ i ,y ′ i ∈D ′ log f [ x ′ i ]( y ′ i )
q ( E = e k | x ′ i , y ′ i ) (cid:12)(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) 
 (140) = E [Equation] , D ′ ,σ 
[Equation]
q ∈Q |D| (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:88) x i ,y i ∈D σ i log f [ x i ]( y i ) q
( E = e k | x i , y i ) (141) − (cid:88) x ′ i ,y ′ i ∈D ′ σ i log f
[ x ′ i ]( y ′ i )
q ( E = e k | x ′ i , y ′ i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (142) ≤
[Equation]
[Equation]

[Equation]
|D| (cid:12)(cid:12)(cid:12)(cid:12) (cid:12)(cid:12) (cid:88) x i ,y i ∈D
[Equation]
log f
[ x i ]( y i ) [Equation]
( E = e
[Equation]
| x i , y i ) (cid:12)(cid:12)(cid:12)(cid:12) (cid:12)(cid:12)   (143) +
[Equation]
[Equation]

[Equation]
|D ′ | (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:88) x
′ i ,y ′ i ∈D ′
σ i log f
[ x ′ i ]( y ′ i )
q ( E = e
[Equation]
| x ′ i , y ′ i ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (144) =
[Equation]
R |D| ( G k
[Equation]
where σ i are independent Rademacher variables.
Equation 139 follows from Jensen’s inequality and the convexity of sup .
Equation 141 holds due to the symmetry of log f
[ x i ]( y i ) q
( E = e k | x i , y i ) − log f
[ x ′ i ]( y ′ i )
q ( E = e k | x ′ i , y ′ i ) and the argument that Radamacher variables preserve the expected sum of symmetric random variables with a convex mapping (Ledoux & Talagrand (1991), Lemma 6.3).
Substituting Equation 145 to Equation 128, we have for any δ ∈ (0 , 1) , with a probability over 1 − δ , ∀ f ∈ V , ∀ q ∈ Q , the following holds:
[Equation]
Let ˜ f k = arg inf f { q ( E = e k ) E q [ − log f
[ x ]( y ) |E
= e k ] } .
Let ˆ f k = arg inf f { 1 |D| (cid:80) x i ,y i ∈D − log f
[ x i ]( y i ) q
( E = e k | x i , y i ) } .
Now we have
[Equation]
Combining Equation 146 and Equation 148-150, the lemma is proved.
Lemma 2.
Assume ∀ x ∈ X , ∀ y ∈ Y , ∀ f ∈ V , log f [ ∅ ]( y ) ∈ [ − B, B ] where B > 0 .
The definition of G k V and R N ( G ) follows from Lemma 1.
Define ˆ f k = arg inf f 1 |D| (cid:80) x i ,y i ∈D − log f [ ∅ ]( y i )
q ( E = e k | x i , y i ) .
Then for any q ∈ Q , any δ ∈ (0 , 1) , with a probability over 1 − δ , we have
[Equation]
Proof.
Similar to Lemma 1, we could prove that
[Equation]
where G k V ∅ = { g | g ( x, y )
=
log f
[ ∅ ]( y ) q ( E = e k | x, y ) , f ∈ V , q ∈ Q} .
According to the definition for the predictive family V (Xu et al.
(2020), Definition 1), ∀ f ∈ V , there exists f ′ ∈ V such that ∀ x ∈ X , f [ ∅ ]
= f ′
[ x ] .
Thus, G k V ∅ ⊂ G k V , and therefore R |D| ( G k V ∅ ) ≤
R |D| ( G k V ) .
Substituting into Equation 153, the lemma is proved.
Lemma 3 ((Xu et al., 2020), Theorem 1) .
Assume ∀ x ∈ X , ∀ y ∈ Y , ∀ f ∈ V , log f
[ x ]( y ) ∈ [ − B, B ] where B > 0 .
Define a function class G ∗V = { g | g ( x, y ) = log f
[ x ]( y ) , f ∈ V} .
The definition of R
N ( G ) follows from Lemma 1.
Then for any δ ∈ (0 , 0 . 5) , with a probability over 1 − 2 δ , we have
[Equation]
Finally we are prepared to prove Theorem 3.
Proof of Theorem 3.
We first bound the error of empirical estimation with the sum of items in
Lemma 1,2,3.
[Equation]
where ˆ f k = arg inf f 1 |D| (cid:80) x i [Equation] y i ∈D
− log f
[ x i ]( y i ) q
[Equation] E =
e k | x i [Equation] y i [Equation]
For simplicity, let
[Equation]
Err
[Equation]
Then, by Lemma 1,2,3,
[Equation]
[Equation]
[Equation]
[Equation]
≤ 2( K + 1) δ
[Equation]
Equation 174 is because of G k V =
G V , G ∗V ⊂ G V and therefore R |D| ( G k V ) ≤ R |D| ( G V ) , R |D| ( G ∗V ) ≤
R |D| ( G V ) .
Hence,
Pr
[Equation]
[Equation]
J P ROOF OF T HEOREM
Proof of Theorem 4.
The objective function of our IM algorithm is directly derived from the definition of empirical predictive heterogeneity in Definition 6.
For the regression task, we assume the predictive family as
[Equation]
the regression model and θ is learnable
[Equation]
[Equation]
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 a fixed term.
Then for each environment e ∈ supp ( E ∗ ) , the I V ( X → Y |E ∗ = e ) becomes
[Equation]
which corresponds with the MSE loss and the proposed regularizer in Equation 17.
For the classification task, the derivation is similar, and the regularizer becomes the entropy of Y in sub-population e and the loss function becomes cross-entropy loss. K D ISCUSSION ON DIFFERENCES WITH SUB - GROUP DISCOVERY
Subgroup discovery (SD, (Helal, 2016)) is aimed at extracting ”interesting” relations among different variables ( X ) with respect to a target variable Y .
Coverage and precision of each discovered group is the focus of such method.
To be specific, it learns a partition on P ( X ) such that some target label y dominates within each group.
The most siginficant gap between subgroup discovery and our predictive heterogeneity lies in the pattern of distributional shift among clusters: for subgroup discovery, P ( X ) and P ( Y ) varies across subgroups but there is a universal P ( Y | X ) .
While for predictive heterogeneity P ( Y | X ) differs across sub-population, which indicates diversified prediction mechanism.
It is such disparity of prediction mechanism that inhibits the performance of a 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 it generates simplified rules that people can understand.
Also, sub-group discovery methods is suitable for the settings that only involve tabular data (typlically from a relational database), where the input features have clear semantics.
And our proposed method could deal with general machine learning settings, including complicated data (e.g., image data) that involves representation learning.
Also, when people have to handle settings where data heterogeneity w.r.t.
prediciton mechanism exists inside data, our method is more applicable.
However, both kinds of methods can be used to help people understand data and make more reasonable decisions.
L D ISCUSSION ON THE P OTENTIAL 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 sub-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-populations to examine whether a given algorithm is fair by calculating the performance gap across the sub-populations.