Data-IQ: Characterizing subgroups with heterogeneous outcomes in tabular data
Anonymous Author(s)
AffiliationAddress email
Abstract
High model performance, on average, can hide that models may systematicallyunderperform on subgroups of the data.
We consider the tabular setting, whichsurfaces the unique issue of outcome heterogeneity - this is prevalent in areas suchas healthcare, where patients with similar features can have different outcomes ,thus making reliable predictions challenging.
To tackle this, we propose Data-IQ ,a framework to systematically stratify examples into subgroups with respect totheir outcomes.
We do this by analyzing the behavior of individual examplesduring training, based on their predictive confidence and, importantly, the aleatoric(data) uncertainty.
Capturing the aleatoric uncertainty permits a principled characterization and then subsequent stratification of data examples into three distinctsubgroups ( Easy , Ambiguous , Hard ).
We experimentally demonstrate the benefits of Data-IQ on four real-world medical datasets.
We show that Data-IQ’scharacterization of examples is most robust to variation across similarly performant(yet different) models, compared to baselines.
Since Data-IQ can be used with any
ML model (including neural networks, gradient boosting etc.), this property ensures consistency of data characterization, while allowing flexible model selection.
Taking this a step further, we demonstrate that the subgroups enable us to constructnew approaches to both feature acquisition and dataset selection.
Furthermore, wehighlight how the subgroups can inform reliable model usage, noting the significantimpact of the Ambiguous subgroup on model generalization.
Introduction
Most machine learning models are optimized using empirical risk minimization (ERM), to maximizeaverage performance during training [1].
However, in real-world settings, while models may performwell on average, they might underperform on specific subgroups of data [2–4].
Most of the currentliterature has focused on this problem in computer vision, where the underperforming subgroups aretypically associated with data examples that have spurious correlations [1, 5] or mislabelling [6].
In this paper, we focus on tabular data, the most ubiquitous format in medicine and finance, where datais based on relational databases [7, 8].
Specific to the tabular setting, we formalize an understudiedsource of underperformance, namely heterogeneity of outcomes .
This phenomenon is vital inhealthcare, where patients with similar features can have different outcomes [9, 10].
For example,[11] showed that prognostic models for risk prediction perform well on average, but underperform onspecific cancer types due to heterogeneity of risk (outcome).
Prior works have audited subgroupsbelonging to sensitive attributes (e.g. demographics, race or gender), as it is well-known that MLmodels generally underperform on these subgroups [12, 13].
However, this approach is limiting, as itneeds the sensitive attributes to be specified, and it also does not capture the case where complexfeature interactions may lead to underperformance.
[Figure]
We take a different approach to automatically stratify data into subgroups, usable with any ML modeltrained in stages (epochs/iterations) ; e.g., neural networks, gradient boosting etc.
Specifically, westudy the behavior of individual examples during training, called training dynamics .
This allows us toformalize that examples can lie on the spectrum from easy to hard to predict.
More concretely, let’sconsider the task of patient mortality prediction.
Based on their features, sicker patients more oftenhave a mortality event.
Thus, they are easy to learn for any model and will be predicted correctlywith high confidence ( Easy ).
However, a subgroup might have a heterogeneous outcome: survivaldespite their poor prognosis.
This heterogeneity could result from randomness, making it practicallyimpossible for a model to learn.
These examples will be predicted incorrectly yet with high confidence(or equivalently have low confidence for the correct class)
( Hard ).
In tabular data, there are alsoexamples with inherent ambiguity where the predicted probability for the correct class remains low.
They appear where the current features are insufficient to distinguish the example correctly, regardlessof the model used [14, 15] ( Ambiguous ).
These subgroups naturally arise in real data; see Fig.1 (ii).
Identifying these subgroups is practically valuable, as improving accuracy and robustness oftendepends on the data’s characteristics and quality [16–19].
As mentioned in [19, 20], the “data” workis often undervalued as merely operational, yet failing to account for it can have immense practicalharm [11, 19].
Consequently, our goal is to build a systematic framework with the following desiredproperties ( P1-P4 ), motivated by the considerations of practitioners at various stages of the MLpipeline.
In satisfying P1-P4 , we seek to address the “dire need for an ML-aware data quality that isnot only principled, but also practical for a larger collection ( . . . )
of ML models”
(P1) Robust data characterization : the characterization of data examples should be robust, such that
it is consistent across similar performing models, that have different architectures/parametrizations.
(P2) Principled data collection : the characterization should be informative and actionable, providing
practitioners insights that enable both quantitative feature collection and selection between datasets.
(P3) Reliable model deployment : the characterization should enable reliable model usage, both by
unmasking unreliable subgroups or using the subgroups to tailor the data for better performance.
(P4) Plug & play : the characterization should be applicable to a variety of ML models widely used
on tabular data, including neural networks, gradient boosting (and variants) etc.
To fulfill P1-P4 , we propose Data-IQ , a systematic framework that characterizes examples basedon the inherent qualities (IQ) of the data; at both training and deployment time.
As outlined in
Fig.1, Data-IQ leverages confidence and in the “data-centric AI” spirit focused on the data: aleatoricuncertainty (i.e. uncertainty inherent to the data).
This permits Data-IQ to provide ML-aware dataquality that is principled and practical for a variety of ML models, making the following contributions:
Contributions: 1 ⃝ Data-IQ models the aleatoric (data) uncertainty
, which permits subgroup identification that is most robust to variation across different but similar performing models/parameterizations,compared to baselines i.e. P1 (Sec.4.1).
2 ⃝ Data-IQ aids with principled data collection P2 in twoways: Firstly, it permits to quantify the value of an acquired feature by measuring how the featurereduces an example’s aleatoric uncertainty.
This information enables a more principled approach tofeature acquisition.
Secondly, it permits to compare datasets based on the proportion of ambiguousexamples.
We demonstrate that the proportions link to how well a model trained with the datasetgeneralizes (Sec.4.2).
3 ⃝ Experimentally, the subgroups identified by Data-IQ can inform reliablemodel deployment, i.e. P3.
We highlight cases, where evaluation on average might mask unreliable performance, including data sculpting, model robustness, and uncertainty estimation methods(Sec.4.3).
4 ⃝ Data-IQ by construction is “plug-and-play” i.e. P4 with any ML model that can becheckpointed, granting practitioners flexibility to use Data-IQ with their model of choice (Sec.3.4).
Related work
This paper primarily engages with the literature on data characterization and contributes to the nascentarea of data-centric AI [21, 22].
An extended discussion of related work is found in Appendix A.1.
Data characterization.
The literature to characterize data samples has used a myriad of differentmetrics.
However, their goals have typically been different, such as spurious correlation or mislabelling, compared to Data-IQ, whose goal is to characterize subgroups with respect to the outcomepredictions.
Furthermore, none of these methods completely addresses all the desired properties(P1-P4).
The closest to our work on data quality is Data Maps [23].
A key contrast to Data-IQ is that
Data Maps use confidence and prediction variability to flag instances.
In Sec. 3, we show that thisprediction variability corresponds to the model uncertainty (i.e. epistemic uncertainty).
Alternatively,
Data-IQ takes a different and more principled approach, capturing the inherent data uncertainty(known as aleatoric uncertainty )
Epistemic uncertainty is reducible by collecting more data.
In comparison, aleatoric uncertainty is irreducible even with more samples.
This is due to the factthat it captures properties inherent to the data [24–26]; only better features can reduce the aleatoricuncertainty [26].
Later in Fig.
4, we show on real data that capturing the aleatoric uncertainty allows
Data-IQ to be more robust to variation across different models, compared to Data Maps ( P1 ).
Thisallows practitioners to characterize their data in such a way that the insights are more consistent.
Wefurther show theoretically in Sec.
3.3, why the characterization by Data-IQ indeed provides a moreprincipled definition for Ambiguous examples, compared to Data Maps.
Besides Data Maps, other related methods address specific computer vision problems: identifyingmislabelled images using area under the margin (AUM) [6], gradient norm to identify “importantexamples” to aid pruning during training [27], or underperformance due to spurious image correlations[1].
The tabular setting considered in this paper requires new methods, due to the specific problem ofheterogeneous outcomes for examples with similar features (i.e. “feature collision”).
The ambiguityin the tabular, “feature collision” sense, is different or non-existent in modalities such as images.
Data-Centric AI.
The assessment of data quality is a critical but often overlooked problem in ML[19].
While the focus in ML is typically on optimizing models, the task of ensuring high quality data(or even improving one’s data) can be equally valuable to improving performance [16, 19].
Evenwhen it is considered, the process of assessing datasets is adhoc or artisinal [19, 21, 28].
The recentgrowth of the data-centric AI space aims to build systematic tools for “data collection, labeling, andquality monitoring processes for datasets to be used in machine learning” [22].
Data-IQ contributesto this nascent body of work, specifically around ML-aware data quality monitoring [18].
Formulation
This section gives a detailed formulation of Data-IQ and motivates our proposed example stratificationthat uses aleatoric uncertainty and confidence.
We then describe how Data-IQ stratifies examples intosubgroups at both training and testing time.
Finally, we show Data-IQ’s formulation permits usagewith any ML model trained in stages, e.g. neural networks, GBDTs etc, unlike other approaches.
3.1 Preliminaries
We consider the typical supervised learning setting, where the aim is to assign an input x ∈ X ⊆
R
d Xto a class y ∈ Y ⊂ N .
We have a dataset D with N ∈ N ∗ examples, i.e. D = {( x n , y n ) ∣ n ∈ [ N ]}drawn IID from an unknown distribution.
Our goal is then to learn a model f θ ∶ X → Y , parameterized by θ ∈ Θ .
Typically, the parameters θ are learned to minimize empirical risk, by minimizing theaverage training loss , i.e. ERM ( θ ) = 1 n ∑ ni = 1 ℓ ( x i , y i ; θ ) , with a loss function ℓ ∶ X × Y × Θ → R + .
This brings us to the essence of the problem: “not all examples are created equally”.
e.g. patientswith similar features might have heterogeneous outcomes, reflected in their labels y being different.
These correspond to subgroups within D train on which a predictive model might systematicallyunderperform.
We formalize this concept of hidden heterogeneous subgroups by assigning to eachexample x n a hidden subgroup label g n ∈ G , where G =
{ Easy , Ambiguous , Hard } .
Before giving aprecise description of how those group labels are assigned, it is useful to detail the context.
Severalworks have established that the training dynamics of a model, contains signal about the quality ofthe data itself [29–31].
For instance, it takes more epochs/iterations for a model to assign the correct label to noisier/more difficult training examples.
With Data-IQ, we build on those observations andassign a label g n to each example x n by studying its training dynamic, which is then used to estimatethe aleatoric uncertainty and predictive confidence of each example.
The following sections detailhow this is done and how this contrasts with existing approaches.
3.2 Uncertainty decomposition during training
Recall that practitioners desire flexibility in the choice of the model.
Hence, we focus on any MLmodel that is trained in stages and can be checkpointed during training, f θ ∶ X → Y parameterizedby θ and on a given example from the training set ( x, y ) ∈ D train .
Assume that the model f θcorresponds to a conditional categorical distribution, assigning a probability to each class giventhe input x : f θ ( x ) =
P ( Y ∣ X
= x, ϑ = θ ) .
During iterative training, the model parametersθ vary, where over E ∈ N ∗ epochs/iterations, these parameters take E different values at eachcheckpoint, i.e. θ 1 , θ 2 , . .
, θ E .
Since our analysis relies on the model’s training dynamics, we wantto take those different parameters into account.
For the sake of notation, we introduce a randomvariable ϑ that has an empirical distribution over this set of parameters captured through the trainingprocess ϑ ∼ P emp ({ θ e ∣ e ∈
[ E ]}) .
The variability of the model’s parameters at training time is thenreflected by the variance V ϑ [ ⋅ ] .
The uncertainty we model is based on the random variable Y ∣ X = x
that represents the possiblelabels given the input x .
Since the ground-truth label y is available for training examples, we wouldlike to distinguish between 2 cases: 1 ⃝ the predicted label corresponds to the ground-truth label Y = yand 2 ⃝ the predicted label is different from the ground-truth label Y ≠ y .
To this end, we introduce abinary random variable ˜ Y that is set to one when the predicted label equals the ground-truth label( ˜ Y =
if Y = y ) and that is zero otherwise ( ˜ Y =
if Y ≠ y ).
As discussed earlier, we are interestedin the uncertainty on the predictive random variable ˜ Y ∣ X
= x .
This uncertainty is modeled by thevariance v ( x ) =
V ˜ Y ∣ X
[ ˜ Y ∣ X = x ] .
We will now show that this quantity can be evaluated with themodel predictions.
We start by noting that the definition of ˜ Y implies that
˜ Y ∣ X
= x, ϑ = θ is a Bernoulli randomvariable with parameter 1 P ( x, θ ) =
P ( Y =
y ∣ X
= x, ϑ = θ ) =
[ f θ ( x )]
y .
From this observation,we can decompose v ( x ) with the law of total variance and make each term explicit:
[Equation]
] (cid:205)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:209)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:207)
[Table]
Bernoulli Var: P ( x,ϑ )( 1 − P ( x,ϑ )) =
V ϑ
[ P ( x, ϑ )] (cid:205)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:209)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:207)
Epistemic uncertainty: v ep ( x )
[Table]
[ P
[Table] [Table]
[Table]
P
[Table]
))] (cid:205)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:209)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:210)(cid:207)
Aleatoric uncertainty: v al ( x ) .
In Eq.
(1), we have split the overall uncertainty into two components: epistemic and aleatoricuncertainty .
This type of decomposition is similar to those in the context of Bayesian neuralnetworks [32].
To understand the distinction between uncertainties, it is useful to closely examinethe variances in the first equality of Eq.
For epistemic uncertainty v ep , variance is evaluatedon the model parameters ϑ .
Hence, epistemic uncertainty originates from the fact that a model’spredictions oscillate when we change its parameters.
For the aleatoric uncertainty v al , the varianceis evaluated on the predicted label ˜
Y ∣ X, ϑ .
Hence, the variability originates from the inabilityto predict the correct label with high confidence.
While existing works use epistemic uncertaintyto stratify examples, we argue that aleatoric uncertainty is a better principled choice to capture theinherent data uncertainty.
3.3 Stratification based on data uncertainty
We now explain how the above notion of uncertainty permits to assign a group label g ∈ G to eachtraining example x .
First, we use the empirical distribution ϑ ∼ P emp ({ θ e ∣ e ∈
[ E ]}) to explicitlywrite the two types of uncertainties in Eq.
(1),where ¯ P ( x ) = 1 / E ∑
[Equation]
[Equation]
Stratification at training time.
Before giving a precise definition of the group labels, let us give anintuitive definition for each group.
1 ⃝ Easy : examples that have low data uncertainty that the modelcan correctly predict with high confidence, 2 ⃝ Ambiguous : examples that have high data uncertainty,hence the model is unable to predict with confidence and 3 ⃝ Hard : examples that have low datauncertainty that the model is unable to predict (i.e. predicted incorrectly yet with high confidence orequivalently have low confidence for the correct class).
We note that we need the model’s predictionfor the ground-truth class to delineate Easy and Hard examples.
In practice, we use the model’saverage confidence for the ground-truth class ¯ P ( x ) defined previously for this purpose.
We make useof this concept to detail how labels are assigned to training examples ( x
, y ) ∈ D train :
[Equation]
Easy
[Equation]
Hard
[Equation]
Ambiguous otherwise
[Equation]
where C up and C low are upper and lower confidence threshold resp.
and P n the n-th percentile.
We 184 provide a practical method to set C up and C low , applicable to any dataset in Appendix A.
[Figure]
In contrast to Data-IQ which uses Aleatoric uncertainty v al ( x ) ; Data Maps [23] identifies ambiguoustraining examples ( x
, y ) ∈ D train as those with highepistemic uncertainty v ep ( x ) .
We consider a typicalscenario to see how this characterization might causeproblems – illustrated in Fig.
Consider an example x in which the model cannot classify confidentlyduring the entire training P ( x, θ e ) = 0 .
5 ∀ e ∈ [ E ] .
In this case, the epistemic uncertainty v ep ( x ) vanishes, as the prediction is consistently unconfident(i.e.
low variability of the model predictions).
Thisimplies that Data Maps would consider this exampleas non-ambiguous, despite the ambiguous model prediction for this example.
This problem can be tracedback to the definition of epistemic uncertainty, whichmeasures the sensitivity of a model prediction with respect to the model’s parameters.
A more principled definition for ambiguous examples should capture examples for which the modelcannot predict the appropriate label with high confidence (i.e. data uncertainty).
This is preciselywhat the aleatoric uncertainty v al ( x ) captures (Data-IQ).
Furthermore, it is easy to verify that theprevious example P ( x, θ e )
= 0 .
5 ∀ e ∈
[ E ] maximizes the aleatoric uncertainty (see Fig. 2).
Sincehigh aleatoric uncertainty captures ambiguous examples for various values of the model’s parameters,we believe that it better reflects the inherent quality of the data.
In that sense, we expect this quantityto be more stable and robust to variation for different ML model parameters/architecture changes( P1 ).
We experimentally validate the consistency in Sec. 4.
Stratification at inference time.
Most previous methods are only applicable at training time.
Toaddress this limitation and improve the practical utility of our method, we also stratify examples intosubgroups at deployment time.
However, if we try to apply the above stratification for incoming dataat deployment time, we face a problem: ¯ P ( x ) requires the ground-truth class y .
For this reason, we follow an alternative approach based on representation learning that does notrequire access to ground-truth labels.
The idea is the following: we construct a low-dimensional
UMAP embedding [33] h ∶ X → H of the training set’s examples x train ∈ D train .
In doing this,we note two things (see Appendix C): 1 ⃝ Ambiguous examples have distinctive features and areclustered in embedding space.
Thus, it is possible to distinguish the Ambiguous examples usingthe embedding.
2 ⃝ It is not possible to reliably distinguish Easy examples from Hard examplesbased on the embedding, because Hard examples are a minority with outcome randomness thathave similar features, as the Easy examples.
Combining these observations, we note it is possibleto identify Ambiguous test examples.
This label is assigned by computing the related embeddingh ( x test ) and comparing this embedding to the nearest neighbor embedding from the training set,i.e. d [ h ( x test ) , h ( x train )] ∀ x train ∈ D train .
For models like neural networks with an implicitrepresentation space, the same analysis can be done using the model’s representation space.
3.4 Using Data-IQ with a variety of models, beyond Neural Networks (P4)
The baseline methods discussed are primarily applicable only to neural networks.
However, practicallyin tabular settings (e.g. healthcare/finance etc), practitioners often use other highly performant iterativelearning algorithms such as Gradient Boost Decision Trees (GBDTs) or variants [7, 34].
Data-IQ’sformulation by construction is naturally adaptable to any ML model trained in stages, that can becheckpointed.
This satisfies P4 , which allows practitioners the flexibility to use Data-IQ with theirapplication-specific model of choice.
Appendix A provides guidelines, space and time considerations,as well as, discussing the specifics of how Data-IQ is easily adapted, for example to GBDTs.
Experiments
This section presents a detailed empirical evaluation demonstrating that Data-IQ satisfies ( P1 )
Robust data characterization, ( P2 ) Principled data collection and ( P3 ) Reliable Model Deployment,introduced in Sec.1.
Recall that ( P4 ) Plug and play is satisfied by construction of Data-IQ.
Datasets.
We conduct experiments on four real-world medical datasets, with diverse characteristics(different sizes, binary/multiclass, varying degrees of task difficulty etc) and highlight real-worldapplicability with heterogeneous patient outcomes: (1) Covid-19 dataset of Brazilian patients [35],(2)
Prostate cancer datasets from both the US [36] and UK [37], (3) Support dataset of seriously illhospitalized adults [38], (4) Fetal state dataset of cardiotocography [39].
We describe the datasets ingreater detail in Appendix B, along with further experimental details.
We observe similar performanceacross different datasets, but given the space limitations, we typically show pertinent results for asingle dataset, and include results for the other datasets in Appendix C.
4.1 (P1) Robust data characterization
[Figure]
to variation.
As per P1 , wedesire that Data-IQ identifies subgroups ina manner robust to variation across different models.
This would allow a practitionerto obtain consistent insights about theirdata even when using different model architectures/parameterizations.
When comparing the different methods from Sec.
2,we note that each method has its own specific metric used to characterize examples.
To assess robustness to variation, we compare the consistency of the different characterization metrics, evaluatedon models with different architectures (number of layers or hidden units) and parameterizations(optimizers).
Specifics of the differences are outlined in Appendix B. Note that all models achievesimilar training performance and are trained to convergence, with early stopping on a validationset.
The main paper does not explicitly compare across model types (neural networks vs gradientboosting).
The reason is that besides Data-IQ, the majority of methods are constrained to neuralnetworks.
For this comparison, see Appendix C.
[Figure]
Quantitatively, we compute the Spearman rank correlation between allmodel combinations, see Fig.
We observe that Data-IQ is the mostconsistent and robust to variation across different models, having thehighest score on all datasets, satisfying P1 .
Further, the baseline methodsthemselves are also not consistent in performance ordering across datasets,which is undesirable.
Ultimately, the robustness means practitioners canfeel confident in the consistency of data insights, derived using Data-IQ.
To further compare Data-IQ and Data Maps[23], we examine 3 distinctmodels that achieve similar performance on the Covid-19 [35] tabulardataset, and we produce a characterization of the training set using eachmodel in Fig.
We note that Data Maps groups can be recovered from (5)by replacing the aleatoric uncertainty v al from Data-IQ with its epistemiccounterpart v ep .
The y-axis is the same for both methods and correspondsto ¯ P ( x ) .
The x-axis corresponds to v al ( x ) for Data-IQ and to v ep ( x ) for
Data Maps.
Each model is assigned a color in Fig.
We note three things
1 ⃝ Data-IQ’s characterization of the data is significantly more stable across models.
2 ⃝ Linked to thepoints in Sec.3, Data Map’s high and low confidence examples in fact have high epistemic uncertaintyv ep , which can lead to incorrect conclusions when attempting to use Data Maps to characterize data.3 ⃝ Data-IQ always distributes the data around a bell shape, which standardizes its interpretation.
Weprovide a theoretical analysis to explain this bell shape observation in Appendix A.
Data insights from subgroups.
Given the distinct differences between the subgroups, we seek tounderstand what factors make these subgroups different and how they can provide insight into thedataset.
Such insights are especially useful in clinical settings.
Results for the prostate cancer datasetare illustrated in Fig.
5 (with other datasets in Appendix C).
To visualize the different groups ofpatients within each subgroup, we cluster each subgroup ( Easy , Ambiguous and Hard ) using a
Gaussian Mixture Model (GMM) similar to [5], selecting the optimal clusters based on the Silhouettescore.
We assess cluster quality vs alternatives in Appendix C.
In general, across datasets, the subgroups are: (1) Easy : Severe patients with a death outcome, andless severe patients with a survival outcome.
(2) Ambiguous : Patients with similar features, butdifferent outcomes.
This could suggest that the features, we have at hand are insufficient to separatethe differences in outcomes.
(3) Hard : Severe patients with a survival outcome, and less severepatients with a death outcome.
i.e. opposite outcomes as expected due to randomness in the outcomes.
[Figure]
2 (P2) Principled data
[Figure]
Principled feature acquisition.
As shown in
Sec.4.1, the Ambiguous subgroup has exampleswith similar features, yet different outcomes.
Recall that this case of ambiguity in the tabularsetting is very different from ambiguity in othermodalities such as images.
Rather, in this setting,the ambiguity is due to insufficient features to adequately separate the examples.
We link this to theconcept that the Ambiguous subgroup has a highaleatoric uncertainty that is irreducible, even ifwe collect more data examples.
Rather, aleatoricuncertainty can only be reduced by acquiring better features [26].
We leverage this idea and showthat Data-IQ’s example characterization providesa principled approach to assessing the benefit ofacquiring a specific feature.
This is a distinctlydifferent task from feature selection, where all features are present and we select the most “importantfeature”.
Additionally, this is different from activelearning which quantifies the value of acquiringexamples, rather than features.
With the above in mind, a valuable feature should decrease the example ambiguity (i.e. aleatoricuncertainty).
Hence, a decrease in the proportion of Ambiguous examples can serve as a proxy forthe feature’s potential value to the dataset.
Understanding the value of features is useful in settingssuch as healthcare, where feature acquisition comes at a cost.
To showcase the potential, we constructa semi-synthetic experiment, where we rank sort the features based on correlation with the target.
We then train different models, where we sequentially “acquire” features of increasing value (basedon correlation).
Fig.
shows results for the Support dataset.
For Data-IQ, Fig 30 (a) shows that aswe acquire “valuable” features, the proportion of the Ambiguous subgroup drops, whilst the Easysubgroup increases, with significant changes for the important features.
This shows that Data-IQ’ssubgroup characterization can be used to quantify a feature’s value, by its ability to decrease ambiguity.
In contrast, Data Maps, Fig 30 (c), shows minimal response to feature acquisition, suggesting it maynot be sensitive enough to capture the feature’s value.
Further, for Data-IQ we see that the examples that remain as Ambiguous , after features are collected,maintain a consistent aleatoric uncertainty.
This is desired as it demonstrates for those exampleswhich remain Ambiguous , that indeed the features collected are not informative enough to reducetheir inherent (aleatoric) uncertainty, i.e. those remaining still need better features (see Fig 30 (b)).
While for Data Maps the added features, in fact increase the variability for all subgroups making itharder to stratify (see Fig 30 (d)).
This links to the fact that Data Maps subgroups can’t capture thevalue of the acquired features.
We further show experimentally in Appendix C that for Ambiguousexamples, it is not simply a case of increasing the size of the dataset (i.e. more examples).
In fact,this can increase the proportion of Ambiguous examples due to the increased probability of featurecollisions as the dataset size increases.
Ultimately, this motivates the usefulness of principled featureacquisition (which Data-IQ can guide), as a way to decrease dataset ambiguity.
Principled dataset comparison.
Extending beyond feature-level, an understudied scenario involvessystematically selecting between entire datasets in two cases: (1) purchasing data from data markets[40–42] and (2) organizations where the data is siloed, with lengthy access processes [43, 44].
In bothcases, synthetic versions of the real dataset has begun to be used [43].
For now, we ignore privacyconcerns, and focus on data fidelity and quality, which compares the real and synthetic datasets usingstatistical measures.
However, as per [45], the conclusions can vary across different metrics.
Inpractice, competing “synthetic” datasets can be generated by different ML models or vendors.
Thus,while they model the same underlying distribution, depending on the process used, one version mightbe superior.
We now ask whether Data-IQ could permit us to systematically select between syntheticdatasets?
We consider the setting where the real data is not accessible .
Hence, we can’t use existingevaluation metrics, yet still wish to compare the synthetic datasets (e.g. comparing vendors).
[Table]
We simulate this scenario by generating synthetic datausing 2 different models, representing 2 synthetic data vendors: (V1) CTGAN [46] and (V2) Gaussian Copula [47].
We then characterize the dataset subgroup proportions using Data-IQ.
We hypothesize that datasets with greaterproportions of Easy examples generalize better.
As is common, we validate fidelity by training with synthetic dataand testing with real data [48, 49], where the best fidelitydata produces the best model performance on real data (testset).
Table 1 shows that datasets with the highest quality as measured by Data-IQ indeed producesthe best performance on real test data (i.e. Rank 1).
Further, it shows that the same “vendor” does notalways produce the best dataset, highlighting the value of comparative assessment.
Ultimately, thesetwo aspects demonstrate that when the real data is unavailable, Data-IQ is a useful tool in the handsof practitioners wishing to assess data quality, especially when selecting between different datasets.
4.3 (P3) Reliable Model Deployment
[Figure]
Less is more: data sculpting based on subgroups.
What role dothe data subgroups, specifically Ambiguous examples, play in ensuringmodel generalization?
Using multi-country prostate cancer data, we traina baseline model on US data (SEER) and assess generalization whendeployed on patients in the UK (CUTRACT) and vice versa.
In Fig.
7,we see that test time generalization performance, monotonically increasesas we decrease the proportion of Ambiguous training data (see Appendix
C.13 for absolute numbers).
Ultimately, this illustrates the value of sculpting the training dataset, byremoving ambiguous examples, as a way to improve the reliability of a deployed model.
Group-DRO: Not a silver bullet in tabular settings.
Once subgroups of underperformance areidentified, it is generally assumed that methods such as Group Distributionally Robust Optimization
(DRO) [50] can be applied to improve model performance and robustness.
We compare Group-DROwith groups identified by Data-IQ, and as baselines: George [5] and Just-Train-Twice (JTT) [1].
Asper the previous experiments, the largest underperforming group (in proportion) is the Ambiguousexamples.
Similar to the literature, we evaluate the performance change from a baseline model, after Group-DRO is applied.
The results in Table 2 show that Group-DRO using Data-IQ’s groups bothimproves overall performance and improves performance on the Ambiguous group.
Whilst, for theother baselines, the performance actually degrades
[Figure]
[Table]
Nevertheless, while using Data-IQ can boost performance, it is evident that simply applying Group-DROis not a silver bullet to equalize subgroup performance,given the sometimes small improvement.
The rationaleis our tabular setting is different, from the spurious correlation setting in computer vision where Group-DROtypically shines.
Ultimately, we believe, based on thefeature acquisition results, that in tabular settings, practitioners would be well served acquiring better featuresto improve performance and reduce ambiguity.
Subgroup-informed usage of uncertainty estimation.
Uncertainty estimation methods are essentialin safety-critical areas such as healthcare, yet are typically assessed on average.
We ask the question:since subgroups have different performance properties, are uncertainty estimates equally reliablefor each subgroup?
As done in the literature [51, 52], if an uncertainty estimate is reliable andinformative of predictive performance, it can be used to defer “uncertain examples”.
This is doneby rank sorting examples based on uncertainty and thresholding proportions of tolerated uncertainty[51–54].
Ideally, as the threshold proportion of examples increases (i.e. inclusion of more uncertainexamples), we should see a monotonic decrease of accuracy.
We assess this by training a Bayesian
Neural Network (BNN)[55] to obtain uncertainty estimates.
We then compute the performance acrossdifferent threshold proportions τ ∈ { 0 . 1 , 0 . 2 , . . .
, 1 } for the Ambiguous subgroup specifically, andas is commonly done, across the entire dataset (i.e. average ).
[Figure]
[Figure]
Fig.
8 shows an example in which the average examples exhibitthe monotonic decreasing relationship, as expected.
However, the
Ambiguous examples categorized by Data-IQ , contrary to expectations, show an increase in accuracy as more uncertain examplesare included.
This suggests that the uncertainty estimates are not asinformative of predictive performance for the Ambiguous examples.
Data-IQ could thus, be leveraged by practitioners at deployment timeto understand which examples require auditing before deferring, asthe “average” monotonic decreasing behavior, based on the uncertainty estimates might not alwayshold.
Ultimately, the result further highlights how subgroup characterization via Data-IQ could assistpractitioners in unmasking unreliable performance, which is not evident on average.
Discussion
In this paper, we introduce Data-IQ, a systematic framework that can be used with any ML modelwith checkpoints, to characterize examples into subgroups with respect to the outcome.
Throughseveral experiments, we demonstrate that the usage of aleatoric uncertainty, which captures propertiesmore inherent to the data, is indeed more principled, being more robust to variation across modelsand/or parameterizations.
Data-IQ’s consistency is unmatched by any compared baseline.
Data-IQshould not automate and replace the intuition of a data scientist.
Rather, as we have demonstrated,
Data-IQ should serve as a systematic “data-centric ML” tool that assists and empowers data scientistswith the “data” work at training time, whilst also guiding reliable model usage at deployment time.
Finally, we outline limitations and future opportunities.
1 ⃝ While Data-IQ characterizes examples;the current formulation does not allow us to understand which attributes are responsible for thecharacterization per example.
This would be an interesting extension around dataset explainability,allowing practitioners to better probe their data.
2 ⃝ In high-stakes settings such as healthcare, tomitigate possible adverse effects (e.g. difficulty of Easy vs Hard ), Data-IQ should be used with a“human-in-the-loop”, allowing experts to complement and validate findings with domain knowledge.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] Please refer to our “contributions” paragraph in Section1.
Furthermore, all claims made about Data-IQ are verified in Section 4 and in the
supplementary material.
(b) Did you describe the limitations of your work?
[Yes] See Section 3.3 and Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Data-IQ
provides practioners with a tool to help them with their “data” work and to better
understand the characteristics of their data.
In fact, Data-IQ may help reveal thenegative impacts of ML models.
However, as discussed in Section 5, expert knowledge
should still be used to validate our characterization.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes] We have carefully read the ethics review guideline and confirm that
our
paper respects the guidelines.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 3and Appendix A of the supplementary material
(b) Did you include complete proofs of all theoretical results?
[Yes] See Section 3 and
Appendix A of the supplementary material
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] Will be released
upon acceptance
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix B, detailing all relevant information
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] Included as relevant to the experiments in Section 4 andAppendix C
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix B
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See citations [35],
We cite these in the main paper and supplementary material
(b) Did you mention the license of the assets?
[Yes] See Appendix B of the supplementary material.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
[Table]
We only use existing assets
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[Yes] It is discussed in detail in the publications by the dataset creators,
which we have cited
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[Yes] See Appendix B of the supplementary material.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]