# Data-IQ: Characterizing subgroups with heterogeneous outcomes in tabular
data

Nabeel Seedat
University of Cambridge
<EMAIL_ADDRESS>
&Jonathan Crabbé
University of Cambridge
<EMAIL_ADDRESS>
Ioana Bica
University of Oxford
The Alan Turing Institute
<EMAIL_ADDRESS>
&Mihaela van der Schaar
University of Cambridge
The Alan Turing Institute
UCLA
<EMAIL_ADDRESS>

###### Abstract

High model performance, on average, can hide that models may systematically
underperform on subgroups of the data. We consider the tabular setting, which
surfaces the unique issue of _outcome heterogeneity_ \- this is prevalent in
areas such as 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 to their outcomes. We do this by analyzing the
behavior of individual examples during 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 distinct subgroups
(${\color[rgb]{0,0.88,0}Easy},{\color[rgb]{1,0.39,0.13}Ambiguous},{\color[rgb]{1,0,0}Hard}$).
We experimentally demonstrate the benefits of Data-IQ on four real-world
medical datasets. We show that Data-IQ’s characterization 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
construct new approaches to both feature acquisition and dataset selection.
Furthermore, we highlight how the subgroups can inform reliable model usage,
noting the significant impact of the ${\color[rgb]{1,0.39,0.13}Ambiguous}$
subgroup on model generalization.

###  1 Introduction

Most machine learning models are optimized using empirical risk minimization
(ERM), to maximize average performance during training [1]. However, in real-
world settings, while models may perform well on average, they might
underperform on specific subgroups of data [2, 3, 4]. Most of the current
literature has focused on this problem in computer vision, where the
underperforming subgroups are typically 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 data is based on relational databases [7, 8].
Specific to the tabular setting, we formalize an understudied source of
underperformance, namely _heterogeneity of outcomes_. This phenomenon is vital
in healthcare, where patients with _similar features_ can have _different
outcomes_ [9, 10, 11]. For example, [12] showed that prognostic models for
risk prediction perform well on average, but underperform on specific cancer
types due to heterogeneity of risk (outcome). Prior works have audited
subgroups belonging to sensitive attributes (e.g. demographics, race or
gender), as it is well-known that ML models generally underperform on these
subgroups [13, 14]. However, this approach is limiting, as it needs the
sensitive attributes to be specified, and it also does not capture the case
where complex feature interactions may lead to underperformance.

Figure 1: Data-IQ systematically characterizes data into subgroups using (i)
any ML model, trained in stages (epochs/iterations) & can be checkpointed. At
training time, Data-IQ (ii) leverages the model’s checkpoints to analyze the
training behavior of individual examples and (iii) characterizes each example
based on its aleatoric (data) uncertainty & prediction confidence. (iv) At
deployment, the subgroups inform model usage by embedding the training set in
a representation space and characterizing the new data points in this
representation space.

We take a different approach to automatically stratify data into subgroups,
usable with any ML model trained in stages (epochs/iterations); e.g., neural
networks, gradient boosting etc. Specifically, we study the behavior of
individual examples during training, called _training dynamics_. This allows
us to formalize that examples can lie on the spectrum from easy to hard to
predict. More concretely, let’s consider the task of patient mortality
prediction. Based on their features, sicker patients more often have a
mortality event. Thus, they are easy to learn for any model and will be
predicted _correctly_ with _high confidence_ (${\color[rgb]{0,0.88,0}Easy}$).
However, a subgroup might have a heterogeneous outcome: survival despite their
poor prognosis. This heterogeneity could result from randomness, making it
practically impossible for a model to learn. These examples will be predicted
_incorrectly_ yet with _high confidence_ (or equivalently have low confidence
for the correct class) (${\color[rgb]{1,0,0}Hard}$). In tabular data, there
are also examples 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, regardless of the model
used [15, 16] (${\color[rgb]{1,0.39,0.13}Ambiguous}$). These subgroups
naturally arise in real data; see Fig.1 (ii).

Identifying these subgroups is practically valuable, as improving accuracy and
robustness often depends on the data’s characteristics and quality [17, 18,
19, 20]. As mentioned in [21, 20], the “data” work is often undervalued as
merely operational, yet failing to account for it can have immense practical
harm [20, 12]. Consequently, our goal is to build a systematic framework with
the following desired properties (P1-P4), motivated by the considerations of
practitioners at various stages of the ML pipeline. In satisfying P1-P4, we
seek to address the “dire need for an ML-aware data quality that is not only
principled, but also practical for a larger collection ($\ldots$) of ML
models” [19]:

(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 dataset selection.
(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 based on 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:
aleatoric uncertainty (i.e. uncertainty inherent to the data). This permits
Data-IQ to provide ML-aware data quality that is principled and practical for
a variety of ML models, making the following contributions:

Contributions: ① Data-IQ models the aleatoric (data) uncertainty, which
permits subgroup identification that is most robust to variation across
different yet similar performing models/parameterizations, compared to other
baselines, i.e. _P1_. ② Data-IQ aids with principled data collection _P2_ in
two ways: Firstly, it permits to quantify the value of an acquired feature by
measuring how the feature reduces the aleatoric uncertainty of the example.
This information enables a more principled approach to feature acquisition.
Secondly, it permits to compare datasets based on the proportion of ambiguous
examples. We demonstrate that the proportions link to how well a model trained
with the dataset generalizes. ③ Experimentally, the subgroups identified by
Data-IQ can inform reliable model deployment, i.e. P3. We highlight cases,
where assessment on average might mask unreliable performance, including data
sculpting, model robustness, and uncertainty estimation methods. ④ Data-IQ by
construction is “plug-and-play” i.e. _P4_ with _any_ ML model that can be
checkpointed, granting practitioners flexibility to apply Data-IQ to their
model of choice.

###  2 Related work

This paper primarily engages with the literature on data characterization and
contributes to the nascent area of data-centric AI [22, 23]. 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 different metrics. 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 outcome
predictions. Furthermore, none of these methods completely addresses all the
desired properties (P1-P4). The closest to our work on data quality is Data
Maps [24]. A key contrast to Data-IQ is that Data Maps use confidence and
prediction variability to flag instances. In Sec. 3, we show that this
prediction 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_) [25]. Epistemic uncertainty is reducible by collecting more
data. In comparison, aleatoric uncertainty is irreducible even with more
samples. This is due to the fact that it captures properties inherent to the
data [26, 25, 27]; only better features can reduce the aleatoric uncertainty
[27]. 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_). This allows practitioners to
characterize their data in such a way that the insights are more consistent.
We further show theoretically in Sec. 3.3, why the characterization by Data-IQ
indeed provides a more principled definition for _Ambiguous_ examples,
compared to Data Maps.

Besides Data Maps, other related methods address specific computer vision
problems: identifying mislabelled images using area under the margin (AUM)
[6], gradient norm to identify “important examples” to aid pruning during
training [28], or underperformance due to spurious image correlations [1]. The
tabular setting considered in this paper requires new methods, due to the
specific problem of heterogeneous outcomes for examples with similar features
(i.e. “feature collision”). The ambiguity in 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 [20]. 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 [20, 17]. Even
when it is considered, the process of assessing datasets is adhoc or artisinal
[20, 22, 29, 30]. The recent growth of the data-centric AI space aims to build
systematic tools for “data collection, labeling, and quality monitoring
processes for datasets to be used in machine learning” [30, 23]. Data-IQ
contributes to this nascent body of work, specifically around ML-aware data
quality monitoring [19].

###  3 Formulation

This section gives a detailed formulation of Data-IQ and motivates our
proposed example stratification that uses aleatoric uncertainty and
confidence. We then describe how Data-IQ stratifies examples into subgroups at
both training and testing time. Finally, we show Data-IQ’s formulation permits
usage with _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\in\mathcal{X}\subseteq\mathbb{R}^{d_{X}}$ to a class
$y\in\mathcal{Y}\subset\mathbb{N}$. We have a dataset $\mathcal{D}$ with
$N\in\mathbb{N}^{*}$ examples, i.e. $\mathcal{D}=\left\\{(x^{n},y^{n})\mid
n\in[N]\right\\}$ drawn IID from an unknown distribution. Our goal is then to
learn a model $f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$, parameterized by
$\theta\in\Theta$. Typically, the parameters $\theta$ are learned to minimize
empirical risk, by minimizing the average training loss , i.e.
$\text{ERM}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\ell(x_{i},y_{i};\theta)$, with a
loss function
$\ell:\mathcal{X}\times\mathcal{Y}\times\Theta\to\mathbb{R}^{+}$.

This brings us to the essence of the problem: “not all examples are created
equally”. e.g. patients with similar features might have heterogeneous
outcomes, reflected in their labels $y$ being different. These correspond to
subgroups within $\mathcal{D}_{\textrm{train}}$ on which a predictive model
might systematically underperform. We formalize this concept of hidden
heterogeneous subgroups by assigning to each example $x^{n}$ a hidden
_subgroup_ label $g^{n}\,{\in}\,\mathcal{G}$, where
$\mathcal{G}=\\{{\color[rgb]{0,0.88,0}Easy},{\color[rgb]{1,0.39,0.13}Ambiguous},{\color[rgb]{1,0,0}Hard}\\}$.
Before giving a precise description of how those group labels are assigned, it
is useful to detail the context. Several works have established that the
training dynamics of a model, contains signal about the quality of the data
itself [31, 32, 33]. 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 and assign a label $g^{n}$ to each
example $x^{n}$ by studying its training dynamic, which is then used to
estimate the aleatoric uncertainty and predictive confidence of each example.
The following sections detail how 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_ ML model that is trained in stages and can be
checkpointed during training, $f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$
parameterized by $\theta$ and on a given example _from the training set_
$(x,y)\in\mathcal{D}_{\textrm{train}}$. Assume that the model $f_{\theta}$
corresponds to a conditional categorical distribution, assigning a probability
to each class given the input $x$: $f_{\theta}(x)=P(Y\mid
X=x,\vartheta=\theta)$. During iterative training, the model parameters
$\theta$ vary, where over $E\in\mathbb{N}^{*}$ epochs/iterations, these
parameters take $E$ different values at each checkpoint, i.e.
$\theta_{1},\theta_{2},\dots,\theta_{E}$. Since our analysis relies on the
model’s training dynamics, we want to take those different parameters into
account. For the sake of notation, we introduce a random variable $\vartheta$
that has an empirical distribution over this set of parameters captured
through the training process $\vartheta\sim
P_{\mathrm{emp}}(\left\\{\theta_{e}\mid e\in[E]\right\\})$. The variability of
the model’s parameters at training time is then reflected by the variance
$\mathbb{V}_{\vartheta}\left[\cdot\right]$.

The uncertainty we model is based on the random variable $Y\mid X=x$ that
represents the possible labels given the input $x$. Since the ground-truth
label $y$ is available for training examples, we would like to distinguish
between 2 cases: ① the predicted label corresponds to the ground-truth label
$Y=y$ and ② the predicted label is different from the ground-truth label
$Y\neq y$. To this end, we introduce a binary random variable $\tilde{Y}$ that
is set to one when the predicted label equals the ground-truth label
($\tilde{Y}=1$ if $Y=y$) and that is zero otherwise ($\tilde{Y}=0$ if $Y\neq
y$). As discussed earlier, we are interested in the uncertainty on the
predictive random variable $\tilde{Y}\mid X=x$. This uncertainty is modeled by
the variance $v(x)=\mathbb{V}_{\tilde{Y}\mid X}\left[\tilde{Y}\mid
X=x\right]$. We will now show that this quantity can be evaluated with the
model predictions.

We start by noting that the definition of $\tilde{Y}$ implies that
$\tilde{Y}\mid X=x,\vartheta=\theta$ is a Bernoulli random variable with
parameter111In this case, $[f_{\theta}(x)]_{y}$ denotes the component $y$ of
the probability vector $f_{\theta}(x)$. $\mathcal{P}(x,\theta)=P(Y=y\mid
X=x,\vartheta=\theta)=[f_{\theta}(x)]_{y}$. From this observation, we can
decompose $v(x)$ with the law of total variance and make each term explicit:

$\displaystyle v(x)$
$\displaystyle=\underbrace{\mathbb{V}_{\vartheta}\left[\mathbb{E}_{\tilde{Y}|X,\vartheta}\left[\tilde{Y}|X=x,\vartheta\right]\right]}_{\text{Bernoulli
Mean:
$\mathcal{P}(x,\vartheta)$}}+\underbrace{\mathbb{E}_{\vartheta}\left[\mathbb{V}_{\tilde{Y}|X,\vartheta}\left[\tilde{Y}|X=x,\vartheta\right]\right]}_{\text{Bernoulli
Var: $\mathcal{P}(x,\vartheta)(1-\mathcal{P}(x,\vartheta))$}}$ (1)
$\displaystyle=\underbrace{\mathbb{V}_{\vartheta}\left[\mathcal{P}(x,\vartheta)\right]}_{\text{Epistemic
uncertainty:
$v_{\mathrm{ep}}(x)$}}+\underbrace{\mathbb{E}_{\vartheta}\left[\mathcal{P}(x,\vartheta)(1-\mathcal{P}(x,\vartheta))\right]}_{\text{Aleatoric
uncertainty: $v_{\mathrm{al}}(x)$}}.$

In Eq. (1), we have split the overall uncertainty into two components:
_epistemic_ and _aleatoric uncertainty_. This type of decomposition is similar
to those in the context of Bayesian neural networks [34, 35]. To understand
the distinction between uncertainties, it is useful to closely examine the
variances in the _first_ equality of Eq. (1). For _epistemic uncertainty_
$v_{\mathrm{ep}}$, variance is evaluated on the _model parameters
$\vartheta$_. Hence, epistemic uncertainty originates from the fact that a
model’s predictions oscillate when we change its parameters. For the
_aleatoric uncertainty_ $v_{\mathrm{al}}$, the variance is evaluated on the
_predicted label $\tilde{Y}\mid X,\vartheta$_. Hence, the variability
originates from the inability to predict the correct label with high
confidence. While existing works use epistemic uncertainty to stratify
examples, we argue that aleatoric uncertainty is a better principled choice to
capture the inherent 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\in\mathcal{G}$ to each training example $x$. First, we use the
empirical distribution $\vartheta\sim P_{\mathrm{emp}}(\left\\{\theta_{e}\mid
e\in[E]\right\\})$ to explicitly write the two types of uncertainties in Eq.
(1),where
$\bar{\mathcal{P}}(x)=\nicefrac{{1}}{{E}}\sum_{e=1}^{E}\mathcal{P}(x,\theta_{e})$:

$\displaystyle
v_{\mathrm{ep}}(x)=\frac{1}{E}\sum_{e=1}^{E}\left[\mathcal{P}(x,\theta_{e})-\bar{\mathcal{P}}(x)\right]^{2}\hskip
14.22636ptv_{\mathrm{al}}(x)=\frac{1}{E}\sum_{e=1}^{E}\mathcal{P}(x,\theta_{e})(1-\mathcal{P}(x,\theta_{e})),$
(2)

Stratification at training time. Before giving a precise definition of the
group labels, let us give an intuitive definition for each group. ①
${\color[rgb]{0,0.88,0}Easy}$: examples that have low data uncertainty that
the model can correctly predict with high confidence, ②
${\color[rgb]{1,0.39,0.13}Ambiguous}$: examples that have high data
uncertainty, hence the model is unable to predict with confidence and ③
${\color[rgb]{1,0,0}Hard}$: examples that have low data uncertainty that the
model is unable to predict (i.e. predicted incorrectly yet with high
confidence or equivalently have low confidence for the correct class). We note
that we need the model’s prediction for the ground-truth class to delineate
${\color[rgb]{0,0.88,0}Easy}$ and ${\color[rgb]{1,0,0}Hard}$ examples. In
practice, we use the model’s average confidence for the ground-truth class
$\bar{\mathcal{P}}(x)$ defined previously for this purpose. We make use of
this concept to detail how labels are assigned to training examples
$(x,y)\in\mathcal{D}_{\textrm{train}}$:

$\displaystyle\vspace{-10mm}g(x,\mathcal{D}_{\textrm{train}})=\begin{cases}${\color[rgb]{0,0.88,0}Easy}$&\text{if}\
\bar{\mathcal{P}}(x)\geq C_{\mathrm{up}}\ \wedge\
v_{\mathrm{al}}(x)<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]\\\
${\color[rgb]{1,0,0}Hard}$&\text{if}\ \bar{\mathcal{P}}(x)\leq
C_{\mathrm{low}}\ \wedge\
v_{\mathrm{al}}(x)<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]\\\
${\color[rgb]{1,0.39,0.13}Ambiguous}$&$\text{otherwise}$\\\
\end{cases}\vspace{-5mm}$ (3)

where $C_{\mathrm{up}}$ and $C_{\mathrm{low}}$ are upper and lower confidence
threshold resp. and $P_{n}$ the n-th percentile. We provide a practical method
to set $C_{\mathrm{up}}$ and $C_{\mathrm{low}}$, applicable to any dataset in
Appendix A.

Figure 2: Example outlining the differences between $v_{\mathrm{al}}$ (Data-
IQ) and $v_{\mathrm{ep}}$ (Data Maps) & showing the type of uncertainty
matters.

In contrast to Data-IQ which uses Aleatoric uncertainty $v_{\mathrm{al}}(x)$;
Data Maps [24] identifies ambiguous training examples
$(x,y)\in\mathcal{D}_{\textrm{train}}$ as those with high epistemic
uncertainty $v_{\mathrm{ep}}(x)$. We consider a typical scenario to see how
this characterization might cause problems – illustrated in Fig. 2. Consider
an example $x$ in which the model cannot classify confidently during the
entire training $\mathcal{P}(x,\theta_{e})=0.5\ \forall e\in[E]$. In this
case, the epistemic uncertainty $v_{\mathrm{ep}}(x)$ vanishes, as the
prediction is consistently unconfident (i.e. low variability of the model
predictions). This implies that Data Maps would consider this example as non-
ambiguous, despite the ambiguous model prediction for this example. This
problem can be traced back to the definition of epistemic uncertainty, which
measures 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 model cannot predict the appropriate label with high confidence
(i.e. data uncertainty). This is precisely what the aleatoric uncertainty
$v_{\mathrm{al}}(x)$ captures (Data-IQ). Furthermore, it is easy to verify
that the previous example $\mathcal{P}(x,\theta_{e})=0.5\ \forall e\in[E]$
maximizes the aleatoric uncertainty (see Fig. 2). Since high 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 quantity to 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. To address this limitation and improve the practical utility of
our method, we also stratify examples into subgroups at deployment time.
However, if we try to apply the above stratification for incoming data at
deployment time, we face a problem: $\bar{\mathcal{P}}(x)$ requires the
ground-truth class $y$.

For this reason, we follow an alternative approach based on representation
learning that does not require access to ground-truth labels. The idea is the
following: we construct a low-dimensional UMAP embedding [36]
$h:\mathcal{X}\rightarrow\mathcal{H}$ of the training set’s examples
$x_{\mathrm{train}}\in\mathcal{D}_{\textrm{train}}$. In doing this, we note
two things (see Appendix C): ① ${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples
have distinctive features and are clustered in embedding space. Thus, it is
possible to distinguish the ${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples
using the embedding. ② It is not possible to reliably distinguish
${\color[rgb]{0,0.88,0}Easy}$ examples from ${\color[rgb]{1,0,0}Hard}$
examples based on the embedding, because ${\color[rgb]{1,0,0}Hard}$ examples
are a minority with outcome randomness that have similar features, as the
${\color[rgb]{0,0.88,0}Easy}$ examples. Combining these observations, we note
it is possible to identify ${\color[rgb]{1,0.39,0.13}Ambiguous}$ test
examples. This label is assigned by computing the related embedding
$h(x_{\mathrm{test}})$ and comparing this embedding to the nearest neighbor
embedding from the training set, i.e.
$d[h(x_{\mathrm{test}}),h(x_{\mathrm{train}})]\forall
x_{\mathrm{train}}\in\mathcal{D}_{\textrm{train}}$. For models like neural
networks with an implicit representation 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, practically in tabular settings (e.g. healthcare/finance
etc), practitioners often use other highly performant iterative learning
algorithms such as Gradient Boost Decision Trees (GBDTs) or variants [37, 7].
Data-IQ’s formulation by construction is naturally adaptable to _any_ ML model
trained in stages, that can be checkpointed. This satisfies _P4_ , which
allows practitioners the flexibility to use Data-IQ with their application-
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.

###  4 Experiments

This section presents a detailed empirical evaluation demonstrating that Data-
IQ 222https://github.com/seedatnabeel/Data-
IQ333https://github.com/vanderschaarlab/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-world applicability with
heterogeneous patient outcomes: (1) Covid-19 dataset of Brazilian patients
[38], (2) Prostate cancer datasets from both the US [39] and UK [40], (3)
Support dataset of seriously ill hospitalized adults [41], (4) Fetal state
dataset of cardiotocography [42]. We describe the datasets in greater detail
in Appendix B, along with further experimental details. _We observe similar
performance across different datasets, but given the space limitations, we
typically show pertinent results for a single dataset, and include results for
the other datasets in Appendix C._

####  4.1 (P1) Robust data characterization

###### Robustness to variation.

Figure 3: Robustness to variation across models based on Spearman correlation,
where Data-IQ has the highest correlation (i.e. consistency) across all
datasets.

As per _P1_ , we desire that Data-IQ identifies subgroups in a manner robust
to variation across different models. This would allow a practitioner to
obtain consistent insights about their data 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 (see Appendix B). To assess robustness to variation, we
compare the consistency of the different characterization metrics, evaluated
on models with different architectures/parameterizations. All models are
trained to convergence, with early stopping on a validation set.

(a) Data-IQ

(b) Data Maps [24]

Figure 4: Data-IQ’s robustness to variation across diff. models (i.e.colors)

Quantitatively, we compute the Spearman rank correlation between all model
combinations, see Fig. 3. We observe that Data-IQ is the most consistent and
robust to variation across different models, having the highest score on all
datasets, satisfying _P1_. Further, the baseline methods themselves are also
not consistent in performance ordering across datasets, which is undesirable.
Ultimately, the robustness means practitioners can feel confident in the
consistency of data insights, derived using Data-IQ.

To further compare Data-IQ and Data Maps[24], we examine 3 distinct models
that achieve similar performance on the Covid-19 [38] tabular dataset, and we
produce a characterization of the training set using each model in Fig. 4. We
note that Data Maps groups can be recovered from (3) by replacing the
aleatoric uncertainty $v_{\mathrm{al}}$ from Data-IQ with its epistemic
counterpart $v_{\mathrm{ep}}$. The y-axis is the same for both methods and
corresponds to $\bar{\mathcal{P}}(x)$. The x-axis corresponds to
$v_{\mathrm{al}}(x)$ for Data-IQ and to $v_{\mathrm{ep}}(x)$ for Data Maps.
Each model is assigned a color in Fig. 4. We note three things ① Data-IQ’s
characterization of the data is significantly more stable across models. ②
Linked to the points in Sec.3, Data Map’s high and low confidence examples in
fact have high epistemic uncertainty $v_{\mathrm{ep}}$, which can lead to
incorrect conclusions when attempting to use Data Maps to characterize data. ③
Data-IQ always distributes the data around a bell shape, which standardizes
its interpretation. We provide a theoretical analysis to explain this bell
shape observation in Appendix A.

Data-IQ: Neural Networks vs Other model classes.  Data-IQ can be used with
_any_ ML model trained in stages linked to _P4_ : Plug and Play. Methods such
XGBoost, LightGBM and CatBoost methods are widely used by practitioners on
tabular data, often more so than neural networks [7]. Ideally, based on _P1_ ,
we desire that the characterization of examples be consistent for similar
performing models, irrespective of whether the model is a neural network or an
XGBoost model.

(a) Data-IQ  (b) Data Maps

Figure 5: NN vs XGBoost: Data-IQ is more consistent

To assess the robustness of both Data-IQ and Data Maps, we train a neural
network, XGBoost, LightGBM and CatBoost models to achieve the same performance
and then perform the characterization for all models. We can clearly see in
Fig. 5 (Support) that Data-IQ has a similar characterization across all four
models. Contrastingly, for Data Maps, the characterizations are significantly
different for the different model classes. The implication of this result is
that by Data-IQ capturing the uncertainty inherent to the data (aleatoric
uncertainty), it leads to a more consistent and stable characterizations of
the data itself. Especially, this highlights that Data-IQ characterizes the
data in a manner that is not as sensitive to the choice of model when compared
to Data Maps. For more, see Appendix C.

Data insights from subgroups. Given the distinct differences between the
subgroups, we seek to understand what factors make these subgroups different
and how they can provide insight into the dataset. Such insights are
especially useful in clinical settings. Results for the prostate cancer
dataset are illustrated in Fig. 6 (with other datasets in Appendix C). To
visualize the different groups of patients within each subgroup, we cluster
each subgroup (${\color[rgb]{0,0.88,0}Easy}$,
${\color[rgb]{1,0.39,0.13}Ambiguous}$ and ${\color[rgb]{1,0,0}Hard}$) using a
Gaussian Mixture Model (GMM) similar to [5], selecting the optimal clusters
based on the Silhouette score. We assess cluster quality vs alternatives in
Appendix C.

In general, across datasets, the subgroups are: (1)
${\color[rgb]{0,0.88,0}Easy}$: Severe patients with a death outcome, and less
severe patients with a survival outcome. (2)
${\color[rgb]{1,0.39,0.13}Ambiguous}$: Patients with similar features, but
different outcomes. This could suggest that the features, we have at hand are
insufficient to separate the differences in outcomes. (3)
${\color[rgb]{1,0,0}Hard}$: Severe patients with a survival outcome, and less
severe patients with a death outcome. i.e. opposite outcomes as expected due
to randomness in the outcomes.

Figure 6: Comparing subgroups identified by Data-IQ (descriptions above).
Colors represent the GMM clusters.

####  4.2 (P2) Principled data collection

###### Principled feature acquisition.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 7: Quantifying the value of feature acquisition based on change in
ambiguity. Only Data-IQ captures this.

As per Sec.4.1, the ${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup has
examples with similar features, yet different outcomes. Recall that this case
of ambiguity in the tabular setting is very different from ambiguity in other
modalities, such as images. The ambiguity is due to insufficient features to
adequately separate the examples. We link this to the concept that the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup has a high aleatoric
uncertainty that is irreducible, even if we collect more data examples.
Rather, aleatoric uncertainty can only be reduced by acquiring better features
[27]. We leverage this idea and show that Data-IQ’s example characterization
provides a principled approach to assessing the benefit of acquiring a
specific feature. This is different from feature selection, where all features
are present and we select the most “important feature”. Additionally, this is
different from active learning which quantifies the value of acquiring
examples, not features.

With the above in mind, a valuable feature should decrease the example
ambiguity (i.e. aleatoric uncertainty). Hence, a decrease in the proportion of
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples can serve as a proxy for the
feature’s potential value to the dataset. Understanding the value of features
is useful in settings such as healthcare, where feature acquisition comes at a
cost. To showcase the potential, we construct a 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 (based on correlation). Fig. 7. shows results for the Support
dataset. For Data-IQ, Fig 7 (a) shows that as we acquire “valuable” features,
the proportion of the ${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup drops,
whilst the ${\color[rgb]{0,0.88,0}Easy}$ subgroup increases, with significant
changes for the important features. This shows that Data-IQ’s subgroup
characterization can be used to quantify a feature’s value, by its ability to
decrease ambiguity. In contrast, Data Maps, Fig 7 (c), shows minimal response
to feature acquisition, suggesting it may not be sensitive enough to capture
the feature’s value.

Further, for Data-IQ we see that the examples that remain as
${\color[rgb]{1,0.39,0.13}Ambiguous}$, after features are collected, maintain
a consistent aleatoric uncertainty. This is desired as it demonstrates for
those examples which remain ${\color[rgb]{1,0.39,0.13}Ambiguous}$, that indeed
the features collected are not informative enough to reduce their inherent
(aleatoric) uncertainty, i.e. those remaining still need better features (see
Fig 7 (b)). While for Data Maps the added features, in fact increase the
variability for _all_ subgroups making it harder to stratify (see Fig 7 (d)).
This links to the fact that Data Maps subgroups can’t capture the value of the
acquired features. We further show experimentally in Appendix C that for
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples, 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 ${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples due
to the increased probability of feature collisions as the dataset size
increases. Ultimately, this motivates the usefulness of principled feature
acquisition (which Data-IQ can guide), as a way to decrease dataset ambiguity.

###### Principled dataset comparison.

Extending beyond feature-level, an understudied scenario involves
systematically selecting between datasets in two cases: (1) purchasing data
from data markets [43, 44, 45] and (2) organizations where the data is siloed,
with lengthy access processes [46, 47]. In both cases, synthetic versions of
the real dataset has begun to be used [46]. For now, we ignore privacy
concerns, and focus on data fidelity and quality, which compares the real and
synthetic datasets using statistical measures. However, as per [48], the
conclusions can vary across different metrics. In practice, 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 might be superior. We now ask whether Data-IQ could permit
us to systematically select between synthetic datasets? We consider the
setting where the real data is _not accessible_. Hence, we can’t use existing
evaluation metrics, yet still wish to compare the synthetic datasets (e.g.
comparing vendors).

Table 1: Comparison of accuracy performance rank and (dataset quality).
Synthetic dataset w/ better quality ($\uparrow$ easy)produces the best real
data test performance.

Dataset | (V1) CTGAN | (V2) Gaussian Copula
---|---|---
Prostate | Rank 1 (63% Easy) | Rank 2 (30% Easy)
Covid | Rank 1 (70% Easy) | Rank 2 (63% Easy)
Support | Rank 1 (59% Easy) | Rank 2 (38% Easy)
Fetal | Rank 2 (40% Easy) | Rank 1 (51% Easy)

We simulate this scenario by generating synthetic data using 2 different
models, representing 2 synthetic data vendors: (V1) CTGAN [49] and (V2)
Gaussian Copula [50]. We then characterize the dataset subgroup proportions
using Data-IQ. We hypothesize that datasets with greater proportions of
${\color[rgb]{0,0.88,0}Easy}$ examples generalize better. As is common, we
validate fidelity by training with synthetic data and testing with real data
[51, 52], where the best fidelity data produces the best model performance on
real data (test set). Table 1 shows that datasets with the highest quality as
measured by Data-IQ indeed produces the best performance on real test data
(i.e. Rank 1). Further, it shows that the same “vendor” does not always
produce the best dataset, highlighting the value of comparative assessment.
Ultimately, these two aspects demonstrate that when the real data is
unavailable, Data-IQ is a useful tool in the hands of practitioners wishing to
assess data quality, especially when selecting between different datasets.

####  4.3 (P3) Reliable Model Deployment

###### Less is more: data sculpting based on subgroups.

Figure 8: Performance improvement via data sculpting based on subgroups

What role do the data subgroups, specifically
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples, play in ensuring model
generalization? Using multi-country prostate cancer data, we train a baseline
model on US data (SEER) and assess generalization when deployed on patients in
the UK (CUTRACT) and vice versa. In Fig. 8, we see that test time
generalization performance, monotonically increases as we decrease the
proportion of ${\color[rgb]{1,0.39,0.13}Ambiguous}$ training data (see
Appendix C.13 for absolute numbers). Ultimately, this illustrates the value of
sculpting the training dataset, by removing ambiguous examples, as a way to
improve the reliability of a deployed model.

###### Group-DRO: Not a silver bullet when used in tabular settings.

Table 2: Comparison of different model improvement/robustness techniques

Dataset | Group | Baseline |  Group-DRO (Data-IQ) |  Group-DRO (George) | JTT
---|---|---|---|---|---
Prostate | Overall | 0.837 |  0.840 $\uparrow$ | 0.565 $\downarrow$ | 0.723 $\downarrow$
Ambiguous | 0.740 |  0.741 $\uparrow$ | 0.598 $\downarrow$ | 0.538 $\downarrow$
The Rest | 0.935 | 0.935 | 0.535 $\downarrow$ | 0.896 $\downarrow$
Covid | Overall | 0.729 |  0.732 $\uparrow$ | 0.687 $\downarrow$ | 0.455 $\downarrow$
Ambiguous | 0.629 |  0.633 $\uparrow$ | 0.609 $\downarrow$ | 0.477 $\downarrow$
The Rest | 0.832 | 0.832 | 0.766 $\downarrow$ | 0.438 $\downarrow$
Support | Overall | 0.734 | 0.734 | 0.660 $\downarrow$ | 0.611 $\downarrow$
Ambiguous | 0.621 |  0.622 $\uparrow$ | 0.576 $\downarrow$ | 0.521 $\downarrow$
The Rest | 0.858 | 0.858 | 0.754 $\downarrow$ | 0.711 $\downarrow$
Fetal | Overall | 0.768 |  0.833 $\uparrow$ | 0.829 $\uparrow$ | 0.829 $\uparrow$
Ambiguous | 0.575 |  0.701 $\uparrow$ | 0.695 $\uparrow$ | 0.698 $\uparrow$
The Rest | 0.970 |  0.974 $\uparrow$ | 0.970 | 0.970

Once subgroups of underperformance are identified, it is generally assumed
that methods such as Group Distributionally Robust Optimization (DRO) [53] can
be applied to improve model performance and robustness. We compare Group-DRO
with groups identified by Data-IQ, and as baselines: George [5] and Just-
Train-Twice (JTT) [1]. As per the previous experiments, the largest
underperforming group (in proportion) is the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples. 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 both
improves overall performance and improves performance on the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ group. Whilst, for the other baselines,
the performance actually degrades.

Nevertheless, while using Data-IQ can boost performance, it is evident that
simply applying Group-DRO is not a silver bullet to equalize subgroup
performance, given the sometimes small improvement. The rationale is our
tabular setting is different, from the spurious correlation setting in
computer vision where Group-DRO typically shines. Ultimately, we believe,
based on the feature acquisition results, that in tabular settings,
practitioners would be well served acquiring better features to improve
performance and reduce ambiguity.

###### Subgroup-informed usage of uncertainty estimation.

Uncertainty estimation methods are essential in safety-critical areas such as
healthcare [54], yet are typically assessed on average. We ask the question:
since subgroups have different performance properties, are uncertainty
estimates equally reliable for each subgroup? As done in the literature [55,
56], if an uncertainty estimate is reliable and informative of predictive
performance, it can be used to defer “uncertain examples”. This is done by
rank sorting examples based on uncertainty and thresholding proportions of
tolerated uncertainty [55, 56, 57, 58]. Ideally, as the threshold proportion
of examples increases (i.e. inclusion of more uncertain examples), we should
see a monotonic decrease of accuracy. We assess this by training a Bayesian
Neural Network (BNN)[59] to obtain uncertainty estimates. We then compute the
performance across different threshold proportions
$\tau\in\\{0.1,0.2,\dots,1\\}$ for the ${\color[rgb]{1,0.39,0.13}Ambiguous}$
subgroup specifically, and as is commonly done, across the entire dataset
(i.e. _average_).

Figure 9: Data-IQ subgroups can unmask unreliable prediction deferral by
uncertainty methods

Fig. 9 shows a _specific_ example, wherein _average_ examples exhibit the
monotonic decreasing relationship, as expected. However, the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples categorized by Data-IQ ,
contrary to expectations, show an increase in accuracy as more uncertain
examples are included. This suggests that uncertainty estimates in this case
are not as informative of predictive performance for the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples. This shows the potential for
practitioners to use Data-IQ at deployment time to understand which examples
require auditing before deferring, as the “average” monotonic decreasing
behavior, based on the uncertainty estimates, may not always hold. Ultimately,
the result further highlights how subgroup characterization via Data-IQ could
assist practitioners in unmasking unreliable performance, not evident on
average.

###  5 Discussion

In this paper, we introduce Data-IQ, a systematic framework that can be used
with _any_ ML model with checkpoints, to characterize examples into subgroups
with respect to the outcome. Through several experiments, we demonstrate that
the usage of aleatoric uncertainty, which captures properties more inherent to
the data, is indeed more principled, being more robust to variation across
models and/or parameterizations. Data-IQ’s consistency is unmatched by any
compared baseline. Data-IQ should 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 scientists
with the “data” work at training time, whilst also guiding reliable model
usage at deployment time.

Data-IQ beyond tabular settings. The main paper has primarily assessed the
utility of Data-IQ in the tabular setting. That said, in Appendices C.7 and
C.8, we evaluate the utility of Data-IQ on text data (_NLP_) and images
(_computer vision_) respectively.

Limitations and future opportunities. ① While Data-IQ characterizes examples;
the current formulation does not allow us to understand which attributes are
responsible for the characterization per example. This would be an interesting
extension around dataset explainability, allowing practitioners to better
probe their data. ② In high-stakes settings such as healthcare, to mitigate
possible adverse effects (e.g. difficulty of ${\color[rgb]{0,0.88,0}Easy}$ vs
${\color[rgb]{1,0,0}Hard}$), Data-IQ should be used with a “human-in-the-
loop”, allowing experts to complement and validate findings with domain
knowledge.

### Acknowledgments

The authors are grateful to Zhaozhi Qian, Yuchao Qin, Evgeny Saveliev and the
anonymous NeurIPS reviewers for their useful comments & feedback. Nabeel
Seedat is supported by the Cystic Fibrosis Trust, Jonathan Crabbe by Aviva,
Ioana Bica by the Alan Turing Institute, EPSRC grant EP/N510129/1 and Mihaela
van der Schaar by the Office of Naval Research (ONR), NSF 1722516.

### References

  * [1] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn.  Just train twice: Improving group robustness without training group information.  In International Conference on Machine Learning, pages 6781–6792. PMLR, 2021.
  * [2] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar.  Do imagenet classifiers generalize to imagenet?  In International Conference on Machine Learning, pages 5389–5400. PMLR, 2019.
  * [3] John C Duchi, Tatsunori Hashimoto, and Hongseok Namkoong.  Distributionally robust losses against mixture covariate shifts.
  * [4] Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang.  Fairness without demographics in repeated loss minimization.  In International Conference on Machine Learning, pages 1929–1938. PMLR, 2018.
  * [5] Nimit Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré.  No subclass left behind: Fine-grained robustness in coarse-grained classification problems.  Advances in Neural Information Processing Systems, 33:19339–19352, 2020.
  * [6] Geoff Pleiss, Tianyi Zhang, Ethan Elenberg, and Kilian Q Weinberger.  Identifying mislabeled data using the area under the margin ranking.  Advances in Neural Information Processing Systems, 33:17044–17056, 2020.
  * [7] Vadim Borisov, Tobias Leemann, Kathrin Seßler, Johannes Haug, Martin Pawelczyk, and Gjergji Kasneci.  Deep neural networks and tabular data: A survey.  arXiv preprint arXiv:2110.01889, 2021.
  * [8] Jinsung Yoon, Yao Zhang, James Jordon, and Mihaela van der Schaar.  Vime: Extending the success of self-and semi-supervised learning to tabular domain.  Advances in Neural Information Processing Systems, 33, 2020.
  * [9] Zepeng Huo, Lida Zhang, Rohan Khera, Shuai Huang, Xiaoning Qian, Zhangyang Wang, and Bobak J Mortazavi.  Sparse gated mixture-of-experts to separate and interpret patient heterogeneity in ehr data.  In 2021 IEEE EMBS International Conference on Biomedical and Health Informatics (BHI), pages 1–4. IEEE, 2021.
  * [10] Gerald Gartlehner, Suzanne L West, Alyssa J Mansfield, Charles Poole, Elizabeth Tant, Linda J Lux, and Kathleen N Lohr.  Clinical heterogeneity in systematic reviews and health technology assessments: synthesis of guidance documents and the literature.  International Journal of Technology Assessment in Health Care, 28(1):36–43, 2012.
  * [11] Jinsung Yoon, Camelia Davtyan, and Mihaela van der Schaar.  Discovery and clinical decision support for personalized healthcare.  IEEE journal of biomedical and health informatics, 21(4):1133–1145, 2016.
  * [12] Eun Jeong Oh, Ravi B Parikh, Corey Chivers, and Jinbo Chen.  Two-stage approaches to accounting for patient heterogeneity in machine learning risk prediction models in oncology.  JCO Clinical Cancer Informatics, 5:1015–1023, 2021.
  * [13] Weijie Jiang and Zachary A Pardos.  Towards equity and algorithmic fairness in student grade prediction.  In Proceedings of the 2021 AAAI/ACM Conference on AI, Ethics, and Society, pages 608–617, 2021.
  * [14] Irene Y Chen, Emma Pierson, Sherri Rose, Shalmali Joshi, Kadija Ferryman, and Marzyeh Ghassemi.  Ethical machine learning in healthcare.  Annual Review of Biomedical Data Science, 4:123–144, 2021.
  * [15] Tin Kam Ho and Mitra Basu.  Complexity measures of supervised classification problems.  IEEE transactions on pattern analysis and machine intelligence, 24(3):289–300, 2002.
  * [16] Ana C. Lorena, Luís P. F. Garcia, Jens Lehmann, Marcilio C. P. Souto, and Tin Kam Ho.  How complex is your classification problem? a survey on measuring classification complexity.  ACM Comput. Surv., 52(5), sep 2019.
  * [17] Abhinav Jain, Hima Patel, Lokesh Nagalapatti, Nitin Gupta, Sameep Mehta, Shanmukha Guttula, Shashank Mujumdar, Shazia Afzal, Ruhi Sharma Mittal, and Vitobha Munigala.  Overview and importance of data quality for machine learning tasks.  In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 3561–3562, 2020.
  * [18] Nitin Gupta, Hima Patel, Shazia Afzal, Naveen Panwar, Ruhi Sharma Mittal, Shanmukha Guttula, Abhinav Jain, Lokesh Nagalapatti, Sameep Mehta, Sandeep Hans, et al.  Data quality toolkit: Automatic assessment of data quality and remediation for machine learning datasets.  arXiv preprint arXiv:2108.05935, 2021.
  * [19] Cedric Renggli, Luka Rimanic, Nezihe Merve Gürel, Bojan Karlas, Wentao Wu, and Ce Zhang.  A data quality-driven view of mlops.  IEEE Data Engineering Bulletin, 2021.
  * [20] Nithya Sambasivan, Shivani Kapania, Hannah Highfill, Diana Akrong, Praveen Paritosh, and Lora M Aroyo.  “everyone wants to do the model work, not the data work”: Data cascades in high-stakes ai.  In proceedings of the 2021 CHI Conference on Human Factors in Computing Systems, pages 1–15, 2021.
  * [21] Kiri Wagstaff.  Machine learning that matters.  arXiv preprint arXiv:1206.4656, 2012.
  * [22] Andrew Ng, Lora Aroyo, Cody Coleman, Greg Diamos, Vijay Janapa Reddi, Joaquin Vanschoren, Carole-Jean Wu, and Sharon Zhou.  Neurips data-centric ai workshop, 2021.
  * [23] Neoklis Polyzotis and Matei Zaharia.  What can data-centric ai learn from data and ml engineering?  arXiv preprint arXiv:2112.06439, 2021.
  * [24] Swabha Swayamdipta, Roy Schwartz, Nicholas Lourie, Yizhong Wang, Hannaneh Hajishirzi, Noah A Smith, and Yejin Choi.  Dataset cartography: Mapping and diagnosing datasets with training dynamics.  In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 9275–9293, 2020.
  * [25] Alex Kendall and Yarin Gal.  What uncertainties do we need in bayesian deep learning for computer vision?  Advances in neural information processing systems, 30, 2017.
  * [26] Armen Der Kiureghian and Ove Ditlevsen.  Aleatory or epistemic? does it matter?  Structural safety, 31(2):105–112, 2009.
  * [27] Efstratios Gavves.  Bayesian deep learning, lecture 10, university of amsterdam, November 2019\.
  * [28] Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite.  Deep learning on a data diet: Finding important examples early in training.  Advances in Neural Information Processing Systems, 34, 2021.
  * [29] Sezal Chug, Priya Kaushal, Ponnurangam Kumaraguru, and Tavpritesh Sethi.  Statistical learning to operationalize a domain agnostic data quality scoring.  arXiv preprint arXiv:2108.08905, 2021.
  * [30] Nabeel Seedat, Jonathan Crabbe, and Mihaela van der Schaar.  Data-suite: Data-centric identification of in-distribution incongruous examples.  In International Conference on Machine Learning. PMLR, 2022.
  * [31] Devansh Arpit, Stanisław Jastrzębski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al.  A closer look at memorization in deep networks.  In International conference on machine learning, pages 233–242. PMLR, 2017.
  * [32] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In International Conference on Machine Learning, pages 322–332. PMLR, 2019.
  * [33] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang.  Learning over-parametrized two-layer neural networks beyond ntk.  In Conference on learning theory, pages 2613–2682. PMLR, 2020.
  * [34] Yongchan Kwon, Joong-Ho Won, Beom Joon Kim, and Myunghee Cho Paik.  Uncertainty quantification using bayesian neural networks in classification: Application to biomedical image segmentation.  Computational Statistics & Data Analysis, 142:106816, 2020.
  * [35] Nabeel Seedat and Christopher Kanan.  Towards calibrated and scalable uncertainty representations for neural networks.  arXiv preprint arXiv:1911.00104, 2019.
  * [36] Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger.  Umap: Uniform manifold approximation and projection.  Journal of Open Source Software, 3(29):861, 2018.
  * [37] Ravid Shwartz-Ziv and Amitai Armon.  Tabular data: Deep learning is not all you need.  Information Fusion, 81:84–90, 2022.
  * [38] Pedro Baqui, Ioana Bica, Valerio Marra, Ari Ercole, and Mihaela van Der Schaar.  Ethnic and regional variations in hospital mortality from covid-19 in brazil: a cross-sectional observational study.  The Lancet Global Health, 8(8):e1018–e1026, 2020.
  * [39] Máire A Duggan, William F Anderson, Sean Altekruse, Lynne Penberthy, and Mark E Sherman.  The surveillance, epidemiology and end results (seer) program and pathology: towards strengthening the critical relationship.  The American journal of surgical pathology, 40(12):e94, 2016.
  * [40] CUTRACT Prostate Cancer UK.  Prostate cancer uk.
  * [41] William A Knaus, Frank E Harrell, Joanne Lynn, Lee Goldman, Russell S Phillips, Alfred F Connors, Neal V Dawson, William J Fulkerson, Robert M Califf, Norman Desbiens, et al.  The support prognostic model: Objective estimates of survival for seriously ill hospitalized adults.  Annals of internal medicine, 122(3):191–203, 1995.
  * [42] Diogo Ayres-de Campos, Joao Bernardes, Antonio Garrido, Joaquim Marques-de Sa, and Luis Pereira-Leite.  Sisporto 2.0: a program for automated analysis of cardiotocograms.  Journal of Maternal-Fetal Medicine, 9(5):311–318, 2000.
  * [43] Pantelis Koutroumpis, Aija Leiponen, and Llewellyn DW Thomas.  Markets for data.  Industrial and Corporate Change, 29(3):645–660, 2020.
  * [44] Ramesh Raskar, Praneeth Vepakomma, Tristan Swedish, and Aalekh Sharan.  Data markets to support ai for all: Pricing, valuation and governance.  arXiv preprint arXiv:1905.06462, 2019.
  * [45] Mohammad Rasouli and Michael I Jordan.  Data sharing markets.  arXiv preprint arXiv:2107.08630, 2021.
  * [46] Khaled El Emam, Lucy Mosquera, and Richard Hoptroff.  Practical synthetic data generation: balancing privacy and the broad availability of data.  O’Reilly Media, 2020.
  * [47] Andre Goncalves, Priyadip Ray, Braden Soper, Jennifer Stevens, Linda Coyle, and Ana Paula Sales.  Generation and evaluation of synthetic patient data.  BMC medical research methodology, 20(1):1–40, 2020.
  * [48] Ahmed Alaa, Boris van Breugel, Evgeny Saveliev, and Mihaela van der Schaar.  How faithful is your synthetic data? sample-level metrics for evaluating and auditing generative models, 2022.
  * [49] Lei Xu, Maria Skoularidou, Alfredo Cuesta-Infante, and Kalyan Veeramachaneni.  Modeling tabular data using conditional gan.  Advances in Neural Information Processing Systems, 32, 2019.
  * [50] Zheng Li, Yue Zhao, and Jialin Fu.  Sync: A copula based framework for generating synthetic data from aggregated sources.  In 2020 International Conference on Data Mining Workshops (ICDMW), pages 571–578. IEEE, 2020.
  * [51] Cristóbal Esteban, Stephanie L Hyland, and Gunnar Rätsch.  Real-valued (medical) time series generation with recurrent conditional gans.  arXiv preprint arXiv:1706.02633, 2017.
  * [52] Michael Platzer and Thomas Reutterer.  Holdout-based empirical assessment of mixed-type synthetic data.  Frontiers in big Data, page 43, 2021.
  * [53] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang.  Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization.  arXiv preprint arXiv:1911.08731, 2019.
  * [54] Nabeel Seedat.  Mcu-net: A framework towards uncertainty representations for decision support system patient referrals in healthcare contexts.  arXiv preprint arXiv:2007.03995, 2020.
  * [55] Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl.  Leveraging uncertainty information from deep neural networks for disease detection.  Scientific reports, 7(1):1–14, 2017.
  * [56] Tong Xia, Jing Han, Lorena Qendro, Ting Dang, and Cecilia Mascolo.  Uncertainty-aware covid-19 detection from imbalanced sound data.  arXiv preprint arXiv:2104.02005, 2021.
  * [57] Ranganath Krishnan and Omesh Tickoo.  Improving model calibration with accuracy versus uncertainty optimization.  Advances in Neural Information Processing Systems, 33:18237–18248, 2020.
  * [58] Jishnu Mukhoti and Yarin Gal.  Evaluating bayesian deep learning methods for semantic segmentation.  arXiv preprint arXiv:1811.12709, 2018.
  * [59] Soumya Ghosh, Jiayu Yao, and Finale Doshi-Velez.  Structured variational learning of bayesian neural networks with horseshoe priors.  In International Conference on Machine Learning, pages 1744–1753. PMLR, 2018.
  * [60] Pang Wei Koh and Percy Liang.  Understanding black-box predictions via influence functions.  In International conference on machine learning, pages 1885–1894. PMLR, 2017.
  * [61] Amirata Ghorbani and James Zou.  Data shapley: Equitable valuation of data for machine learning.  In International Conference on Machine Learning, pages 2242–2251. PMLR, 2019.
  * [62] Gaojie Jin, Xinping Yi, Liang Zhang, Lijun Zhang, Sven Schewe, and Xiaowei Huang.  How does weight correlation affect generalisation ability of deep neural networks?  Advances in Neural Information Processing Systems, 33:21346–21356, 2020.

### Checklist

  1. 1.

For all authors…

    1. (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 Section 1. Furthermore, all claims made about Data-IQ are
verified in Section 4 and in the supplementary material.

    2. (b)

Did you describe the limitations of your work? [Yes] See Section 3.3 and
Section 5.

    3. (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 the negative impacts of ML models. However, as discussed in
Section 5, expert knowledge should still be used to validate our
characterization.

    4. (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.

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]
See Section 3 and Appendix A of the supplementary material

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] See Section
3 and Appendix A of the supplementary material

  3. 3.

If you ran experiments…

    1. (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)? [Yes]
See footnotes 2 and 3

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] See Appendix B, detailing all relevant
information

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes] Included as relevant to the experiments in
Section 4 and Appendix C

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix
B

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes] See
citations [38], [39]. [40], [41],[42]. We cite these in the main paper and
supplementary material

    2. (b)

Did you mention the license of the assets? [Yes] See Appendix B of the
supplementary material.

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [N/A] We only use existing assets

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [Yes] It is discussed in detail in the publications by
the dataset creators, which we have cited

    5. (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.

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

## Appendix: Data-IQ: Characterizing subgroups with heterogeneous outcomes in
tabular data

###  Appendix A Data-IQ details & related work

####  A.1 Extended Related Work

We present a comparison of our framework Data-IQ, and provide further contrast
to related work. Table 3, highlights that the related methods do not satisfy
all the properties (P1-P4). Furthermore even for cases where it is satisfied,
the properties are not naturally satisfied; rather they require our framework
to enable the use-case. _P2_ around data collection and selection, has been
largely unaddressed by prior methods. In the case of _P3_ , the related
methods focus on training time, without taking into account end-to-end
deployment scenarios. Finally, these prior methods are typically applicable to
only neural networks, which limits the broad applicability related to _P4_ ,
where in high-stakes settings such as healthcare and finance, which often have
tabular data, practitioners typically use non-neural methods, such as Gradient
Boosting or XGBoost.

Table 3: Comparison of related work, highlights that _ONLY_ Data-IQ naturally
address all desiderata.
($\heartsuit$): satisfied naturally by construction, ($\diamondsuit$): needs
our framework

|  Assessment metric |  (P1) Data characterization consistent across models |  (P2) Principled data collection & selection |  (P3) Applicable at deployment time |  (P4) Applicable to _any_ ML model
---|---|---|---|---|---
Data-IQ (Ours) | Aleatoric uncertainty | V. High | ✓ | ✓($\heartsuit$) | ✓($\heartsuit$)
Data Maps [24] | Training variability | Medium | ✗ | ✓($\diamondsuit$) | ✓($\diamondsuit$)
AUM [6] | Logit difference | High | ✗ | ✗ | ✓
GraNd [28] | Gradient Norm | Medium | ✗ | ✗ | ✗
JTT [1] | Training errors | Low | ✗ | ✓($\diamondsuit$) | ✓

The properties of _P1_ were evaluated in Section 2. We delve into each
subsequent property below and show that the related works have different goals
and hence do not satisfy the properties:

  * •

_P2_ : Data Maps [24] does not fulfill this property as was shown
experimentally in Section 4.2. The method is shown to be insufficiently
sensitive for the task. AUM [6] uses the differences between logits to flag
mislabelled instances. By virtue, AUM solves a different task of label errors,
whereas the task of improving the data by collecting better features, etc., is
focussed on the inputs rather than labels. GraNd [28] looks at the gradient
norm to identify important examples. Their task, rather than feature
acquisition, is to assist in the pruning of a dataset. JTT [1] looks at
training errors to detect samples with spurious correlations. Not only is the
task different, but flagged errors could also link to both ambiguity (needing
more features) or hardness (example being difficult). Therefore, it is not
sufficiently fine-grained to be leveraged for feature acquisition. In the best
case, we might want to augment the dataset with such flagged examples.

  * •

_P3_ : Data Maps naturally operates only at training time (as per the original
paper). However, if it leverages Data-IQ’s representation space, then it could
be extended to be used at test time. AUM operates to flag mislabelled data
(i.e. outcome labels). These do not exist at test time, hence AUM is a
training-time ONLY method. GraNd uses gradients to flag examples. Hence, it is
only applicable at training time and not at test time, since we do not compute
gradients at test time as we do not have access to labels. JTT surfaces
training errors and hence we would not know test time errors as we do not have
labels. However, if JTT leverages Data-IQ’s representation space idea, then it
could also be extended to be used at test time by flagging examples relative
to a training time embedding.

  * •

_P4_ : Data Maps was only studied for usage with neural networks. For use with
_any_ ML model, it can be extended using our framework (see Section 3.4 and
Appendix A.2). AUM could of course be used with any ML model with output
logits. GraNd does not facilitate this, as only ML models where gradients can
be computed are usable. Finally, since JTT, simply captures model errors it is
easily usable with any ML model.

Tangential related work. We wish to highlight that while methods such as
Influence functions [60] and Data Shapley [61] may seem related, firstly they
have fundamentally different goals around capturing the contributing role of
examples to the models decision boundary and predictions. Hence, they are
considered as tangential to our work of characterizing data into subgroups.
Additionally, these methods assign an instance importance, related to its
contribution to the decision boundary of a specific model. This is different
from our task of characterizing the properties inherent in the data and then
using it to stratify the examples into subgroups.

Another key difference lies in the computational requirements, where both
Influence functions and Data Shapley are highly compute intensive, on top of
the model training. By contrast, Data-IQ’s characterization can be obtained
simply by virtue of training an ML model.

####  A.2 Data-IQ adapted to any model

Figure 10: Pseudo-ensemble for GBDTs for usage with Data-IQ, where each sub-
model is a checkpoint

Data-IQ’s formulation allows us to use it with _any_ ML model training in
stages (i.e. iterative learning) such as gradient boosting decision trees
(GBDTs) methods - such as XGBoost, LightGBM etc. We simply need a set of
checkpoints through training.

This property is incredibly important for practical utility, since methods
such as GBDTs or XGBoost are widely used by practitioners due to their
performant nature on tabular data (often outperforming neural networks). GBDTs
iteratively combines weak models and has shown great success on tabular data,
often outperforming neural networks. Hence, having our method applicable for
GBDTs in addition to neural networks adds to the broad utility.

Before outlining why methods such as GBDT’s fit the Data-IQ pardigm, we
provide a brief overview. Formally, given a dataset
$\mathcal{D}_{\textrm{train}}$, GBDT iteratively constructs a model
$F:X\rightarrow\mathbb{R}$ to minimize the empirical risk. At each iteration
$e$ the model is updated as:

$F^{(e)}(\mathbf{}{x})=F^{(e-1)}(\mathbf{x})+\epsilon h^{(e)}(\mathbf{x}),$
(4)

We now provide guidance on how to apply Data-IQ to such methods. Naively, we
could construct the checkpoints as an ensemble of multiple independent GBDT’s.
However, this is inefficient as the space and time complexity scales with $N$
models. To avoid increasing the overhead of a single model (from a
practitioner perspective), we create a pseudo-ensemble using a single GBDT,
see Figure 10.

Similar to neural networks, the iterative learning nature of a GBDT means that
the sequential submodels can be considered as checkpoints. Formally, each sub-
model has parameters $\mathbf{\theta}^{(i)}$, hence the ensemble of
checkpoints can be described as $\Theta=\\{\mathbf{\theta}^{(i)},1\leq i\leq
N\\}$.

We then apply Data-IQ as normal to the checkpoints
($\theta_{1},\theta_{2}...\theta_{E}$). The flexibility of this approach is
that it applies both to training a new model, but interestingly, we can also
apply this to an ALREADY trained model by looping through the structure to
create the pseudo-ensemble.

####  A.3 Why does Data-IQ produce a bell-shape?

We now theoretically analyze why Data-IQ has a bell-shaped curve. We show if
the x-axis is aleatoric uncertainty and if the y-axis is predictive
confidence, that we always obtain a bell-shaped curve.

Recall two terms from our formulation: (1) $\mathcal{P}(x,\theta)$: the
model’s predictive confidence for an input $x$ applied to a model with
parameters $\theta$, (2) $\bar{\mathcal{P}}(x)$: the model’s average
confidence for the ground-truth class. Now assume $\mathcal{P}(x,\theta)$ is
peaked around $\bar{\mathcal{P}}(x)$, as shown in Figure 11 below.

Figure 11: $\mathcal{P}(x,\theta)$ peaked around $\bar{\mathcal{P}}(x)$

Based on this, when we compute Data-IQ and obtain the following axes. This
results in Figure 12 (a) with the bell-shape, which matches what we observe in
practice (Figure 12 (b))

  * •

y-axis:
$\nicefrac{{1}}{{E}}\sum_{e=1}^{E}\mathcal{P}(x,\theta_{e})\approx\bar{\mathcal{P}}(x)$

  * •

x-axis:
$\frac{1}{E}\sum_{e=1}^{E}\mathcal{P}(x,\theta_{e})(1-\mathcal{P}(x,\theta_{e}))\approx\bar{\mathcal{P}}(x)[1-\bar{\mathcal{P}}(x)]$

(a) Theoretical bell-shape

(b) Real bell-shape

Figure 12: The bell-shapes derived theoretically and practically match.

####  A.4 Evaluating Data-IQ and Data Maps over stages

###### Goal.

Recall that we utilize _any_ ML model trained in stages (e.g. epochs or
iterations), for both Data-IQ and Data Maps. In the main paper, we train our
models to convergence, with early stopping on a validation set. As is common
practice, early stopping is done to prevent the model from overfitting. Once
early stopping kicks in and we stop model training, we then compute the Data-
IQ and Data Maps, averaged over all previously checkpointed models.

This brings two questions to mind, that we seek to answer:

  1. 1.

How do Data-IQ and Data Maps differ at various stages (epochs)?

  2. 2.

What happens if we don’t do early stopping and let the model overfit?

Using Figure 13, we answer both questions. For this we train a neural network
and around 10 epochs early stopping would normally kick in. However, instead
of using early stopping, we continue training the neural network until 20
epochs, thereby letting the model overfit.

Figure 13: Evaluation of Data-IQ and Data Maps over various epochs (stages),
even letting the model overfit. Data-IQ is more stable, whereas Data Maps
changes shape after overfitting which is undesirable.

###### Takeaway.

  1. 1.

Differences across epochs: Data Maps changes significantly over the course of
training, This is due to its sensitivity as the model parameters are
changed/updated through training. This sensitivity is reflected in the
variability (epistemic uncertainty) on the x-axis. In contrast, Data-IQ is
less sensitive and rapidly converges. The rationale is that Data-IQ models the
inherent uncertainty in the data (aleatoric), thereby being more stable and
less sensitive.

  2. 2.

Model overfits: An interesting phenomenon occurs after the early stopping
phase and into the overfitting stage. DataIQ remains consistent/stable, as
expected. It is during this overfitting stage that the Data Maps shape
changes, and begins to look similar to Data-IQ. The reason for this is that
after $\approx$10 epochs (when the model begins to overfit), the model
parameters do not get updated much (the loss doesn’t change). The small
changes of model parameters are then reflected by the epistemic/model
uncertainty not changing much (x-axis of Data Maps). The minimal change means
that when we average at the end of a longer training - the effects get washed
out. This washing out of the effect causes the change in shape by Epoch 20. We
will delve into this point in greater detail in the next sub-section as it
explains the difference of Data Maps in our setting as compared to other
settings.

####  A.5 Data Maps: difference in our setting

###### Goal.

We now take a deep dive, based on the interesting findings at the model
ovefitting stage, where the Data Maps shape changes. We believe that the
difference of our setting from the original NLP setting of Data Maps provides
a clear explanation for why the shape we observe for Data Maps is different
from the shape observed in the original paper [24].

In typical NLP settings, the models are larger and highly parameterized.
Hence, the models are trained for a large number of epochs (sometimes
hundreds). However, in the tabular setting, models need to train for
significantly fewer (as is evident by the early stopping around $\approx 10$
epochs. Further, our models are much shallower compared to highly
parameterized NLP models, motivating why our setting requires early stopping
to prevent overfitting.

###### Takeaway.

We now begin our deep dive. As we saw in Figure 13, after around 10 epochs,
the training loss plateaus. This means that subsequently the model parameters
do not get updated much over training. The result of this is the epistemic
uncertainty would have much less variability after 10 epochs. Thus, when we
average the epistemic uncertainty for longer epochs these minimal changes
would wash out effects from earlier in training.

This effect over longer training, explains why the shape of Data Maps is
different in the NLP settings with large number of epochs and, assuming, we
let the model overfit in our setting.

What implications does this have? For this we assess the evolution of three
examples where we know their ground truth subgroup
(${\color[rgb]{0,0.88,0}Easy},{\color[rgb]{1,0.39,0.13}Ambiguous},{\color[rgb]{1,0,0}Hard}$).

First, we assess Data Maps which uses epistemic uncertainty. Figure 14 shows
the evolution of epistemic uncertainty (Data Maps) for the three subgroup
examples during training. We see that during model training (before early
stopping), that the ${\color[rgb]{0,0.88,0}Easy}$ and
${\color[rgb]{1,0,0}Hard}$ examples, in fact, have higher epistemic
uncertainty than the ${\color[rgb]{1,0.39,0.13}Ambiguous}$ example. This would
lead us to the incorrect conclusion that the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ example is Easy.

In practice, we do not want to let our model knowingly overfit. However, we
can only obtain the correct characterization for Data Maps if we let the model
overfit and wash out the earlier effects. This highlights a limitation of Data
Maps vs Data-IQ.

Figure 14: Data Maps: Epistemic Uncertainty (variability) evolution over
epochs per subgroup

We contrast this with Data-IQ which uses the Aleatoric uncertainty. We not
only have consistent behavior for the ${\color[rgb]{1,0.39,0.13}Ambiguous}$
examples, but also irrespective of whether we use early stopping or let the
model overfit, we can obtain the correct characterization of the examples.

Figure 15: Data-IQ: Aleatoric Uncertainty evolution over epochs per subgroup

Finally, we look at the Data Maps shapes at the point where we perform early
stopping (Figure 16(a)) and at the end of training if we let the model
severely overfit (Figure 16(b)).

(a) Data Maps @ 10 epochs (Correct model) (b) Data Maps @ 30 epochs (OVERFIT
model)

Figure 16: Comparison of Data Maps when model is just right and overfits,
illustrating how the shape changes

Figure 16 shows we can only obtain the expected shape from the original paper
[24], if we let the model overfit. This is of course not an ideal solution, as
in practice we do not want to overfit our models by training for extra epochs.
Especially in our tabular setting with less parameterized models, on smaller
tabular datasets.

This motivates further why indeed we do require Data-IQ and why it’s
characterization of examples based on aleatoric uncertainty is more useful and
principled.

####  A.6 Data-IQ thresholds

As outlined in Section 3.3, we stratify samples based on the Aleatoric
uncertainty and predictive confidence.

Typically, for aleatoric uncertainty we set the percentile $P_{n}=50$th
percentile. The rationale for this is any example where the aleatoric
uncertainty is greater than the 50th percentile should be considered
${\color[rgb]{1,0.39,0.13}Ambiguous}$. This is based on the definition of
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples. Of course, this threshold can
be updated by practitioners.

This means that the examples below $P_{50}$ of the aleatoric uncertainty are
then ${\color[rgb]{0,0.88,0}Easy}$ or ${\color[rgb]{1,0,0}Hard}$. To
differentiate the two, we look at the predictive confidence. We consider
examples with high confidence, on the true label as
${\color[rgb]{0,0.88,0}Easy}$ and those with low confidence, on the true label
as ${\color[rgb]{1,0,0}Hard}$. This means that we need to practically set
$C_{\mathrm{up}}$ and $C_{\mathrm{low}}$.

We define a threshold $thresh\in\\{0,0.5\\}$ such that
$C_{\mathrm{up}}=1-thresh$ and $C_{\mathrm{low}}=thresh$. An obvious
threshold=0.25 such that $C_{\mathrm{up}}=0.75$ and $C_{\mathrm{low}}=0.25$.
However, we present a practical method that could be applied to determine this
threshold. We note that from our characterizations, examples with high
aleatoric uncertainty are often uncertain about their predicitive confidence.
Hence, they never have very high or low uncertainty, rather the uncertainty is
around 0.5. However, since it is never exactly 0.5 this motivates a band
namely $C_{\mathrm{up}}$ and $C_{\mathrm{low}}$.

Now assume for any dataset, we train a model and apply Data-IQ. We can then
sweep $thresh\in\\{0,0.5\\}$ and assess the proportion of examples in each
group.

We show results in Figure 17 below as we sweep the threshold. For low
threshold values, the proportions are, of course, stable (selecting the most
confident samples). However, there is a threshold value beyond which where the
curve rises until a threshold knick point. After that knick-point there is
stability once again. We propose that this knick-point threshold of stability
be used as the threshold in practice. The rationale is that this will allow
for better consistency of data characterization.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 17: SUPPORT Quantifying the value of feature acquisition based on
change in ambiguity. Only Data-IQ captures this relationship.

####  A.7 Data-IQ subgroup characterization

With Data-IQ, we can obtain a plot of the data where the x-axis represents
aleatoric uncertainty and the y-axis predicitive confidence for the true
label. As discussed this means we obtain a bell-shaped characterization. Our
goal is highlight where on the plot each subgroup lies. i.e. illustrate the
location of ${\color[rgb]{0,0.88,0}Easy},{\color[rgb]{1,0.39,0.13}Ambiguous}$
and ${\color[rgb]{1,0,0}Hard}$ on the plot.

We highlight in Figure 18, that ${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples
more likely to lie on the right side with higher aleatoric uncertainty. We
note that based on our characterization this region of the plot never has high
and low predictive uncertainty. Rather, it typically only has uncertainty
approximately around 0.5.

Figure 18: $\mathcal{P}(x,\theta)$ peaked around $\bar{\mathcal{P}}(x)$

###  Appendix B Benchmarks & Experimental Details

####  B.1 Benchmarks, Datasets and Implementation details

#####  B.1.1 Data-IQ

The description of Data-IQ is detailed in the main paper, where we stratify
examples based on aleatoric uncertainty and predictive confidence.

###### Implementation details.

We implement Data-IQ via a class that is versatile and can plug into any
framework including Pytorch, Tensorflow/Keras, Skorch and Scikit-learn style
APIs. The rationale is that since Data-IQ can be used with any model, the
functionality should be equally plug and play from a usability perspective.
Note that for libraries including Pytorch, Tensorflow/Keras, Skorch, we can
easily include Data-IQ simply by adding it to the training callbacks of the
respective training loops.

#####  B.1.2 Data-IQ algorithm

Algorithm 1 below presents the Data-IQ algorithm in more detail.

Data: Dataset $\mathcal{D}$

Result: Subgroup assignment $g$ per sample $x$, i.e.
$g(x,\mathcal{D}_{\textrm{train}})$

Inputs: $C_{\mathrm{up}}=0.75$, $C_{\mathrm{low}}=0.25$, $P_{50}$ = 50th
percentile

Training: Train model ($f$) for $E$ epochs. At each epoch $e$, compute
$\mathcal{P}(x,\theta_{e})$ for datapoint $x$, where $\theta_{e}$ is the model
parameters at epoch $e$.

Characterize samples into subgroups ($g(x,\mathcal{D}_{\textrm{train}})$): for
_$(x_{i},y_{i})$ in $\mathcal{D}_{\textrm{train}}$_ do

Metrics:
$\bar{\mathcal{P}}(x_{i})=\nicefrac{{1}}{{E}}\sum_{e=1}^{E}\mathcal{P}(x_{i},\theta_{e})$;

$v_{\mathrm{al}}(x_{i})=\frac{1}{E}\sum_{e=1}^{E}\mathcal{P}(x_{i},\theta_{e})(1-\mathcal{P}(x_{i},\theta_{e}))$

if  _$\bar{\mathcal{P}}(x_{i})\geq C_{\mathrm{up}}\ \wedge\
v_{\mathrm{al}}(x_{i})
<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]$_ then

Assign ${\color[rgb]{0,0.88,0}Easy}$ subgroup

else if  _$\bar{\mathcal{P}}(x_{i})\leq C_{\mathrm{low}}\ \wedge\
v_{\mathrm{al}}(x_{i})
<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]$_ then

Assign ${\color[rgb]{1,0,0}Hard}$ subgroup

else

Assign ${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup

end if

end for

Algorithm 1 Data-IQ Algorithm

$\displaystyle\vspace{-10mm}g(x,\mathcal{D}_{\textrm{train}})=\begin{cases}${\color[rgb]{0,0.88,0}Easy}$&\mathrm{if}\
\bar{\mathcal{P}}(x)\geq C_{\mathrm{up}}\ \wedge\
v_{\mathrm{al}}(x)<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]\\\
${\color[rgb]{1,0,0}Hard}$&\mathrm{if}\ \bar{\mathcal{P}}(x)\leq
C_{\mathrm{low}}\ \wedge\
v_{\mathrm{al}}(x)<P_{50}\left[v_{\mathrm{al}}(\mathcal{D}_{\textrm{train}})\right]\\\
${\color[rgb]{1,0.39,0.13}Ambiguous}$&\mathrm{otherwise}\\\
\end{cases}\vspace{-5mm}$ (5)

#####  B.1.3 Data Maps

Similar to Data-IQ; Data Maps has been detailed in the main paper.

###### Implementation details.

Our implementation of Data Maps can be used with the flexible approach we
built for Data-IQ. We adapt the implementation of Data Maps from
444https://github.com/allenai/cartography.

#####  B.1.4 AUM

The Area-under-the-Margin (AUM) metric first computes the margin for data
$x,y$ at epoch $t$.

$M^{(t)}(x,y)=\overbrace{z^{(t)}_{y}(x)}^{\text{assigned
logit}}-\overbrace{\textstyle\max_{i\neq y}z^{(t)}_{i}(x)}^{\text{largest
other logit}}.$ (6)

Over all epochs the AUM is then computed as the average of these margin
calculations.

$AUM(x,y)=\textstyle{\frac{1}{T}\sum_{t=1}^{T}}\>M^{(t)}(x,y),$ (7)

###### Implementation details.

The benchmark is based on [6] and we use the implementation from
555https://github.com/asappresearch/aum.

#####  B.1.5 GraNd

The gradient norm at epoch $t$ for an input $x,y$ is computed as :
$\chi_{t}(x,y)=\mathbb{E}\left\|\sum_{k=1}^{K}\gradient_{f^{(k)}}\ell\left(f_{t}(x),y\right)^{T}\psi_{t}^{(k)}(x)\right\|_{2}$
where $\psi_{t}^{(k)}(x)=\gradient_{\mathbf{w}_{t}}f_{t}^{(k)}(x)$.

###### Implementation details.

The benchmark is based on [28]. We adapt the Jax implementation from
666https://github.com/mansheej/data_diet to Pytorch with the help of
777https://github.com/cybertronai/autograd-lib.

#####  B.1.6 JTT

The Just-Train-Twice (JTT) method is a two-stage method. For a model, trained
via ERM, it first identifies errors in the training examples during training
$i.e.y\neq\hat{y}$.

The second step is to train a final model, where the error samples are
increased via weighting.

$\displaystyle L_{\text{up-ERM}}(\theta,E)=(\lambda\sum\limits_{(x,y)\in
E}\ell(x,y;\theta)+\sum\limits_{(x,y)\not\in E}\\!\\!\ell(x,y;\theta)),$ (8)

where $\lambda$ is the weighting factor and $E$ errors.

The goal is that the weighting factor will upsample the samples with errors
and by training on this augmented dataset it will aid to improve worst-group
performance.

###### Implementation details.

Our benchmark is based on [1] and we adapt the implementation from
888https://github.com/anniesch/jtt.

####  B.2 Datasets

We describe the four real-world medical datasets in greater detail.

###### SEER Dataset

The SEER dataset is a publicly available dataset consisting of 240,486
patients enrolled in the American SEER program [39]. The dataset consists of
features used to characterize prostate cancer: including age, PSA (severity
score), Gleason score, clinical stage, treatments etc. A summary of the
covariate features can be found in Table 4. The classification task is to
predict patient mortality, which is binary label $\in\\{0,1\\}$.

The dataset is highly imbalanced, where $~{}94\%$ of patients survive. Hence,
we extract a balanced subset of of 20,000 patients (i.e. 10,000 with label=0
and 10,000 with label=1).

Table 4: Summary of features for the SEER Dataset [39]

Feature | Range
---|---
Age | $37-95$
PSA | $0-98$
Comorbidities | $0,1,2,\geq 3$
Treatment | Hormone Therapy (PHT), Radical Therapy - RDx (RT-RDx),
Radical Therapy -Sx (RT-Sx), CM
Grade | $1,2,3,4,5$
Stage | $1,2,3,4$
Primary Gleason | $1,2,3,4,5$
Secondary Gleason | $1,2,3,4,5$

###### CUTRACT Dataset

The CUTRACT dataset is a private dataset consisting of 10,086 patients
enrolled in the British Prostate Cancer UK program [40]. Similar, to the SEER
dataset, it consists of the same features to characterize prostate cancer. In
addition, it has the same task to predict mortality. A summary of the
covariate features can be found in Table 5.

Once again, the dataset is highly imbalanced, hence we then choose extract a
balanced subset of of 2,000 patients (i.e. 1000 with label=0 and 1000 with
label=1).

Table 5: Summary of features for the CUTRACT Dataset [40]

Feature | Range
---|---
Age | $44-95$
PSA | $1-100$
Comorbidities | $0,1,2,\geq 3$
Treatment | Hormone Therapy (PHT), Radical Therapy - RDx (RT-RDx),
Radical Therapy -Sx (RT-Sx), CM
Grade | $1,2,3,4,5$
Stage | $1,2,3,4$
Primary Gleason | $1,2,3,4,5$
Secondary Gleason | $1,2,3,4,5$

###### Fetal Dataset

The Fetal dataset [42] is a publicly available data set consisting of 2126
patients who underwent fetal cardiotograms (CTG). The patients had specific
diagnostic features measured and the CTGs were classified by 3 expert
obstertricians (with a consensus label) to assign three fetal states (normal,
suspect and pathologic). Our goal is to predict these three fetal states,
hence the task is a multi-class classification problem. A summary of the
covariate features can be found in Table 6.

Table 6: Summary of features for the Fetal Dataset [42]

Feature | Range
---|---
LB (FHR baseline) | $106-160$
AC (Acceleration/sec) | $0-260$
FM (fetal movements/sec) | $0-564$
UC (uterine contractions/sec) | $0-23$
ASTV (% of time with abnormal short term variability) | $12-87$
ALTV (% of time with abnormal long term variability) | $0-91$
MLTV (mean value of long term variability) | $0-51$
DS (severe deceleration/sec) | $0-1$
DP (prolonged deceleration/sec) | $0-4$
Nzeros (n histogram zeros) | $0-10$
Variance (histogram variance) | $0-269$
Tendency (histogram tendency) | $0-1$
NSF (fetal class code) | $1(Normal),2(suspect),3(pathologic)$

###### Covid-19 Dataset

The Covid-19 dataset [38] consists of Covid patients from Brazil. The dataset
is publicly available as it is based on SIVEP-Gripe data. The dataset consists
of 6882 patients from Brazil recorded between Februrary 27-May 4 2020. The
dataset captures risk factors including comorbidities, symptoms, and
demographic characteristics. There is a mortality label from Covid-19 making
it a binary classification task. A summary of the characteristics of the
covariates can be found in Table 7.

Table 7: Summary of features for the Covid-19 Dataset [38]

Feature | Range
---|---
Sex | 0 (Female), 1(Male)
Age | $1-104$
Fever | $0,1$
Cough | $0,1$
Sore throad | $0,1$
Shortness of breath | $0,1$
Respiratory discomfort | $0,1$
SPO2 | $0-1$
Diharea | $0,1$
Vomitting | $0,1$
Cardiovascular | $0,1$
Asthma | $0,1$
Diabetes | $0,1$
Pulmonary | $0,1$
Immunosuppresion | $0,1$
Obesity | $0,1$
Liver | $0,1$
Neurologic | $0,1$
Branca (Region) | $0,1$
Preta (Region) | $0,1$
Amarela (Region) | $0,1$
Parda (Region) | $0,1$
Indigena (Region) | $0,1$

###### Support Dataset

The Support [41] is a private dataset which examines outcomes for seriously
ill hospitalized patients consisting of 9105 patients. The features include
demographics, medical history, clinical variables, and laboratory variables.
Our goal is to predict patient mortality. A summary of the covariate features
can be found in Table 8.

Table 8: Summary of features for the Support dataset [41]

Feature | Range
---|---
Sex | 0 (Female), 1(Male)
Age | $18-102$
Number of comorbities | $0-9$
ARF/MOSF w/Sepsis | $0,1$
COPD | $0,1$
CHF | $0,1$
Cirrhosis | $0,1$
Coma | $0,1$
Colon Cancer | $0,1$
Lung cancer | $0,1$
ARF/MOSF w/Malig | $0,1$
Cancer | $0,1$
under $11k | $0,1$
$11k-$25k | $0,1$
$25k-$50k | $0,1$
>$50k | $0,1$
white | $0,1$
black | $0,1$
asian | $0,1$
hispanic | $0,1$
yeas of education | $0-31$
avg TISS | $1-83$
diabetes | $0,1$
dementia | $0,1$
mean bp | $0-195$
white blood cell count | $0-200$
heart rate | $0-300$
respiratory rate | $0-90$
temperature | $31-42$
pafi | $12-890$
albumnin | $0.39-29$
bilirubin | $0.09-63$
creatinine | $0.09-63$
sodium | $6-181$
ph | $6.8-7.81$
glucose | $0-1092$
bun | $1-300$
urine output | $0-9000$
adlp | $0-7$
adls | $0-7$

####  B.3 Additional experiment details

We note that all experiments were performed using a single Nvidia Tesla P100
GPU.

#####  B.3.1 Section 4.1. Robustness to variation experiment details.

We assess robustness to variation across different models and/or
parameterizations. The rationale for this is that these are model changes or
perturbations that different practitioners might easily make. However, we
desire consistency and stability in the metrics such that the data insights
remain consistent. For layers we evaluate a: 3 layer MLP, 4 layer MLP and 5
layer MLP. For number of hidden units we define two types: type 1 has 64 units
in layer one, which reduce by half per layer and type 2 has 256 units in layer
one, which reduce by half per layer. Finally, we evaluate these combinations
with both Adam as the optimizer during training. These models perform at a
similar level (i.e. accuracy). We then assess the consistency of the metrics
for all examples using these different models/parameterizations based on
Spearman rank correlation.

#####  B.3.2 Section 4.1. Data insights from subgroups experiment details.

For the data insights experiment, we perform clustering on each subgroup
identified by Data-IQ. We make use of Gaussian Mixture Models (GMM), similar
to [5]. The difference is the sub-space in which the clustering is applied. In
Data-IQs case, we cluster within each subgroup identified by Data-IQ. For each
supergroup, we search over $k\in 2,\ldots,10$. We choose the $k$ value that
yields the highest average Silhouette score. Our implementation is based on
999https://scikit-learn.org

#####  B.3.3 Section 4.2. Principled feature acquisition experiment details.

We assess the potential for principled feature acquisition using a semi-
synthetic experiment. We assess how valuable each feature is based on the
Pearson correlation of the feature with the target. We then rank sort the
correlations from lowest to highest. This ordering is then used to determine
the order in which to “acquire” and assess the features.

#####  B.3.4 Section 4.2. Principled dataset comparison experiment details.

As described in the main paperm we compare two synthetic datasets using two
models representing two synthetic data vendors: namely CTGAN and Gaussian
Copula. We use the implementations of 101010https://github.com/sdv-dev/SDV.

#####  B.3.5 Section 4.3. Less is more experiment details.

For this experiment we note that we train multiple baseline models to evaluate
the generalization based on dataset sculpting. We evaluated three different
XGBoost models with 100,150 and 200 estimators each. Our XGBoost
implementation uses the Python version from
111111https://github.com/dmlc/xgboost.

#####  B.3.6 Section 4.3. Group-DRO experiment details.

Group DRO uses training group annotations to directly minimize the worst-group
error on the training set. Hence, we compare different methods to obtain the
group annotations. This experiment compares Group-DRO applied to groups
identified by Data-IQ and as a baseline George [5]. We also compare this to
Just-Train-Twice (JTT) [1]. We use the implementation of Group-DRO from
121212https://github.com/HazyResearch/hidden-stratification, which is based on
[53]. Our George benchmark is based on the implementation from
131313https://github.com/HazyResearch/hidden-stratification and the JTT
benchmark is based on the implementation from
141414https://github.com/anniesch/jtt . All code is in Pytorch.

The group DRO objective can then be written as:

$\displaystyle L_{\text{group-DRO}}(\theta)=\max_{g\in
G}\frac{1}{n_{g}}\sum_{i\mid g_{i}=g}\ell(x_{i},y_{i};\theta)$ (9)

where $n_{g}$ is the number of training points with group $g_{i}=g$.

#####  B.3.7 Section 4.3. Subgroup informed usage of uncertainty estimation
experiment details.

As mentioned in the main paper we obtain uncertainty estimates using a
Bayesian Neural Network based on [59]. Specifically, we train a 5-layer MLP
model. A Gaussian prior is placed over the weights and we optimize the KL
divergence during training. Our benchmark is based on [59] and we use the
implementation from 151515https://github.com/IBM/UQ360

.

###  Appendix C Additional Experiments

####  C.1 Understanding the representation space

###### Goal.

At deployment/testing time, we stratify the examples into subgroups using the
representation space. of the training examples. Our goal is to illustrate the
points made in Section 3.3 with respect to the representation space

Figure 19: Representation space of training examples broken down by sub-group.
We then zoom into the Ambiguous subgroup and illustrate the heterogeneous
outcomes, where examples with similar features can have different outcomes.

###### Takeaway.

We recall the two main points about the representation space made in the main
paper and illustrate them using the representation space (i.e. UMAP embedding)
shown in Figure 19.

  1. 1.

${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples have distinctive features and
are clustered in embedding space. Thus, we can distinguish the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples using the embedding. We can
clearly see this with the green points in Figure 19.

  2. 2.

It is not possible to reliably distinguish ${\color[rgb]{0,0.88,0}Easy}$
examples from ${\color[rgb]{1,0,0}Hard}$ examples based on the embedding,
because ${\color[rgb]{1,0,0}Hard}$ examples are a minority with outcome
randomness that have similar features, as the ${\color[rgb]{0,0.88,0}Easy}$
examples. Naturally, since the Easy and Hard example embeddings are similar,
they also have similar features. We can clearly see this where the orange
examples (Hard) are randomly scattered with the Blue (Easy) examples, in
Figure 19. This makes these two groups difficult to distinguish.

We also zoom in to the ${\color[rgb]{1,0.39,0.13}Ambiguous}$ region (sub-
figure of Figure 19). This illustrates the specific issue of outcome
heterogeneity, where similar patients (with similar features) have different
outcomes. This is illustrated by the different colors, which represent the
different outcomes (i.e. label=0 is survive and label=1 is death).

####  C.2 Data-IQ: Neural Networks vs Other model classes

###### Goal.

As discussed in the main text and Appendix A.2, Data-IQ can be used with _any_
ML model trained in stages. Methods such XGBoost, LightGBM and CatBoost
methods are widely used by practitioners on tabular data, often more so than
neural networks [7]. While Data-IQ can naturally handle any such model, as
long as there are checkpoints (over iterations), Data Maps naturally cannot
(or at least it was never formulated to do so). However, to enable comparison,
we can use Data Maps within our framework, so that it can be applied to a more
general class of models.

Ideally, based on _P1_ , we desire that the characterization of examples be
consistent for similar performing models, irrespective of whether the model is
a neural network or, for example, an XGBoost model.

To assess this the robustness of both Data-IQ and Data Maps, we train a neural
network, XGBoost, LightGBM and CatBoost models to achieve the same performance
and then perform the characterization for all models. The evaluation of all
data sets is shown in Figures 20-23.

###### Takeaway.

We can clearly see that Data-IQ has a similar characterization for across all
four models. Contrastingly, for Data Maps, the characterizations are
significantly different for the different model classes. The key difference is
Data-IQs usage of Aleatoric uncertainty, rather than the Data Maps usage of
Epistemic uncertainty. The implication of this result is that by Data-IQ
capturing the uncertainty inherent to the data (aleatoric uncertainty), it
leads to a more consistent and stable characterizations of the data itself.
Especially, this highlights that Data-IQ characterizes the data in a manner
that is not as sensitive to the choice of model when compared to Data Maps.

(a) Data-IQ  (b) Data Maps

Figure 20: NN vs XGBoost comparison on the Covid-19 dataset, showing Data-IQ
is more consistent

(a) Data-IQ  (b) Data Maps

Figure 21: NN vs XGBoost comparison on the Fetal dataset, showing Data-IQ is
more consistent

(a) Data-IQ (b) Data Maps

Figure 22: NN vs XGBoost comparison on the Prostate dataset, showing Data-IQ
is more consistent

(a) Data-IQ  (b) Data Maps

Figure 23: NN vs XGBoost comparison on the Support dataset, showing Data-IQ is
more consistent

To further augment these results, we quantitatively assess the consistency
across model types. Similar to Section 4.1, we have compared the Spearman rank
correlation between the metrics of all model combinations (neural networks,
XGBoost, LightGBM, CatBoost). Note that all four models are trained to perform
similarly on a held-out dataset.

Below, we present Table 9 showing the average and standard deviation. We see
that Data-IQ has higher scores for all datasets, highlighting the better
consistency compared to Data Maps.

Table 9: Comparison of robustness/consistency across different models on the
basis of the Spearman rank correlation

Dataset | (Ours) Data-IQ | Data Maps
---|---|---
Covid | $0.85\pm 0.05$ | $0.53\pm 0.07$
Support | $0.84\pm 0.06$ | $0.48\pm 0.04$
Prostate | $0.84\pm 0.11$ | $0.44\pm 0.17$
Fetal | $0.72\pm 0.07$ | $0.45\pm 0.06$

####  C.3 Collecting more data - it might make things worse

###### Goal.

Typically it is assumed that simply by collecting more data that most problems
will be mitigated. We aim to assess the influence of adding more data samples,
in contrast to previous experiments in the main paper on acquiring features.
Especially, we want to show how, in the case of heterogeneous outcomes,
blindly adding more data might not always be the solution.

###### Experiment.

We sub-sample proportions of the dataset from 0-100% of the original size of
the dataset. We then recompute the Data-IQ subgroups on each proportion, as
more data examples are added. We show two examples in Figure 24, in which the
subgroups evolve as more data is added.

(a) covid more samples  (b) support more samples

Figure 24: The ${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup proportion in
fact increases as the dataset size increases, i.e. more examples are added.
This is due to the increased probability of feature collisions.

###### Takeaway.

As the number of data examples increases, the proportion covered by the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup increases. This is due to the
increased probability of “feature collisions” as more examples are added.
Therefore, the result shows that the simple solution, in which we naively
collect more data, can sometimes do more harm.

####  C.4 Identifying subsets in an outcome driven manner: assessing quality

###### Goal.

Data-IQ identifies subgroups in an outcome-driven manner, compared to the
conventional unsupervised manner. We compare the cluster quality on the
outcome-driven subgroups as obtained by Data-IQ compared to other benchmarks,
namely: George [5], clustering in the latent space, clustering in a PCA space,
and clustering the raw feature space. All methods use a Gaussian Mixture model
for consistency.

We evaluated these methods based on two well-known cluster quality metrics,
namely, (1) Silhouette score (SIL) and (2) Davies-Bouldin (DB) score. The
results for two example datasets can be seen in Table 10.

Table 10: Clustering results

Dataset | Metric | Data-IQ (Ours) | Outcome latent space (George) | Latent space | PCA | Feature space
---|---|---|---|---|---|---
Covid-19 | SIL ($\uparrow$) | 0.68 | 0.48 | 0.47 | 0.51 | 0.02
| DB ($\downarrow$) | 0.42 | 0.87 | 0.78 | 0.81 | 5.7
Prostate | SIL ($\uparrow$) | 0.80 | 0.6 | 0.56 | 0.57 | 0.54
| DB ($\downarrow$) | 0.3 | 0.6 | 0.63 | 0.57 | 0.61

###### Takeaway.

We can clearly see that Data-IQ has the highest SIL score and the lowest DB
score. These results demonstrate that the outcome-driven clustering approach
taken by Data-IQ, where the clusters are within the Data-IQ subgroups, results
in higher quality clusters.

####  C.5 Data Insights: additional results

###### Goal.

The main text (Section 4.1) illustrated the data insights that can be obtained
for a single dataset (Prostate). Here, we include the results for all other
datasets (using the same setup). See Figures 25-27.

###### Takeaway.

In general, for each dataset, the subgroups represent: (1)
${\color[rgb]{0,0.88,0}Easy}$: Severe patients with a death outcome, and less
severe patients with a survival outcome, (2)
${\color[rgb]{1,0.39,0.13}Ambiguous}$: Patients with similar features, but
different outcomes. This could suggest that the features we have at hand are
insufficient to separate the differences in outcomes and (3)
${\color[rgb]{1,0,0}Hard}$: Severe patients with a survival outcome, and less
severe patients with a death outcome. i.e. opposite outcomes as expected due
to randomness in the outcomes.

(a) Easy  (b) Ambiguous  (c) Hard

Figure 25: Fetal subgroups identified by Data-IQ (descriptions above). Colors
represent the GMM clusters

(a) Easy  (b) Ambiguous  (c) Hard

Figure 26: Support subgroups identified by Data-IQ (descriptions above).
Colors represent the GMM clusters

(a) Easy  (b) Ambiguous  (c) Hard

Figure 27: Covid subgroups identified by Data-IQ (descriptions above). Colors
represent the GMM clusters

####  C.6 Principled feature acquisition (P2): additional results

###### Goal.

The main text (Section 4.2) illustrated principled feature acquisition using
Data-IQ for a single dataset. Here, we include the results for all other
datasets (using the same semi-synthetic setup, based on correlation of the
feature with the target). Recall our formulation, that a valuable feature
should decrease the ambiguity of an example, overall decreasing the proportion
of ${\color[rgb]{1,0.39,0.13}Ambiguous}$ examples in the dataset. This permits
a principled approach to feature acquisition.

###### Takeaway.

We see similar results to the main text for all datasets. For Data-IQ, we see
that as we acquire “valuable” features, the proportion of the
${\color[rgb]{1,0.39,0.13}Ambiguous}$ subgroup drops, while the proportion of
the ${\color[rgb]{0,0.88,0}Easy}$ subgroup increases. There are significant
changes for the important features. This shows that Data-IQ’s subgroup
characterization can be used to quantify a feature’s value, by its ability to
decrease ambiguity. In contrast, Data Maps, shows minimal response to feature
acquisition, suggesting that it may not be sensitive enough to capture the
feature’s value.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 28: COVID Quantifying the value of feature acquisition based on change
in ambiguity. Only Data-IQ captures this relationship.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 29: FETAL Quantifying the value of feature acquisition based on change
in ambiguity. Only Data-IQ captures this relationship.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 30: Prostate Quantifying the value of feature acquisition based on
change in ambiguity. Only Data-IQ captures this relationship.

(a) Data-IQ subgroup ambiguity proportion is reduced as more informative
features are acquired. (b) Data-IQ aleatoric uncertainty remains stable for
Ambiguous, reduces for others as features are acquired.

(c) Data Maps subgroup proportions largely unaffected as features acquired.
(d) Data Maps variability increases across subgroups as features acquired.

Figure 31: Support Quantifying the value of feature acquisition based on
change in ambiguity. Only Data-IQ captures this relationship.

(a) Data-IQ: TEXT  (b) Data Maps: TEXT (c) Data-IQ: IMAGE  (d) Data Maps:
IMAGE

Figure 32: Data-IQ vs Data Maps compared on text and image modalities

####  C.7 Data-IQ on text data (NLP)

###### Goal.

We wish to assess the generality of Data-IQ. Whilst, in the main paper we
mainly tackled tabular data, we wish to showcase the potential utility of
Data-IQ on text data for NLP tasks.

###### Experiment.

We train a bidirectional RNN on the IMBDb 50K dataset, where the task is to
predict the sentiment on the basis of the reviews. We then use Data-IQ to
analyze the training dynamics and categorize the data into subgroups -
${\color[rgb]{0,0.88,0}Easy}$, ${\color[rgb]{1,0.39,0.13}Ambiguous}$ and
${\color[rgb]{1,0,0}Hard}$

###### Takeaway.

We notice that the ${\color[rgb]{0,0.88,0}Easy}$,
${\color[rgb]{1,0.39,0.13}Ambiguous}$ and ${\color[rgb]{1,0,0}Hard}$ groups
represent different phenomena in the text data. We summarize these below.

EASY: Average review length - 163 words

\- Easy to predict review. Very negative sentiment with a negative label

NEGATIVE (0): This movie is pathetic in every way possible. Bad acting,
horrible script (was there one?), terrible editing, lousy cinematography,
cheap humor. Just plain horrible. I had seen ’The Wishmaster’ a couple weeks
before this movie and I thought it was a dead-ringer for worst movie of the
year. Then, I saw ’The Pest’ and suddenly ’The Wishmaster’ didn’t seem so bad
at all. Bad Bad Bad. Excruciatingly bad.

POSITIVE (1) When I saw this movie for the first time I was both surprised and
a little shocked by the blatant vibrance of the story. It is a very artistic
drama with incredible special effects, spectacular acting, not to mention a
very excellent job in the makeup department. Jennifer Lopez has pulled herself
out of past roles that dug into her career with this movie, portraying a very
sensitive child psychologist who works with a team of engineers to enter the
minds of comatose patients to treat them. Vincent D’onofrio played amazingly
well. His portrayal of a sadist serial killer was perfect to a T. The sheer
emotion conveyed by his performance is astounding. Vince Vaughn isn’t my
favorite, but still performed exceptionally well. The symbolism and artistry
was intriguing and titillating, sometimes surprising, and other times
shocking. Overall, I say this is a wonderful movie, with excellent acting and
beautiful artwork

NEGATIVE (0): I’ve tried to watch this so-called comedy, but it’s very hard to
bear. This is a bad, narrow-minded, cliché-ridden movie. Definitively not
funny, but very much boring and annoying, indeed. Bad script, bad acting. It’s
a complete waste of time \- and there remains nothing more to say, I’m afraid
1 out of 10 points

AMBIGUOUS: Average review length - 283 words \- Linguistic ambiguity. A mix of
positive and negative comments. Other reviews of AMBIGUOUS even have text
which simply describes the story without any positive/negative sentiment. \-
Ambiguous reviews also are much longer (in words).

NEGATIVE (0) I saw this movie on the film festival of Rotterdam (jan ’06) and
followed the discussion between director and public afterwards. Many people
reacted shocked and protesting. He will get a lot of negative critics. But:
the world is cruel like this, and it’s not funny. People don’t like it. That
itself doesn’t mean that the movie is bad. I can see that difference. Don’t
shoot the messenger that shows us the world outside our ’hubble’! Nevertheless
I think this a bad movie. Film-technically it’s a good one. Nice shots and
script, most good fitting music, great actors. The director pretends to make a
psychological movie, - the psychology however is of poor quality. Describing
such a powerful violence itself is not the art. The art would be a powerful
description of the psychological process behind that violence. How does a shy
boy come to such a cruelty? The director pretends to describe that, - but is
not good in that. The director used several times the word the ’selfishness’
of people, mentioning for instance the teacher. Only: this teacher wasn’t
selfish,- just someone in several roles, caring for his pupils, ánd worried
about his script. I think it’s a simplification to call him selfish. The
atmosphere in the village is creepy, and the mother made awful mistakes (’you
terribly let me down’) but it doesn’t become believable for me, that there is
caused súch a lot of pain, that the shyest boy comes to such terrible things.
In fact, reality is far more complex than the way, this film describes and it
needs far better descriptions. The interesting thing would be: how does it
work? Describe that process for me please, so that we understand.With the
written phrase on the end, the director said to point to an alternative way of
life. It was the other extreme, and confirmed for me that director and
scriptwriter are bad psychologists, promoting black/white-thinking. The
connection between violence in films and in society has been proved. Use such
a violence gives the responsibility to use it right. There are enough
black/white-thinkers in the world, causing lots of war and misery. I hope,
this movie won’t be successful.

NEGATIVE (0) When the film started I got the feeling this was going to be
something special. The acting and camera work were undoubtedly good. I also
liked the characters and could have grown to empathise with them. The film had
a good atmosphere and there was a hint of fantasy. However, as the film went
on, the plot never appeared to takeoff and just rolled on scene by scene. I
was unable to understand the connection between the stories. All I could see
was the characters occasionally bumping into each other and references to
ships in bottles. Without that connection, I was just left with a few
unremarkable short stories.

POSITIVE (1) Artistically speaking, this is a beautiful movie–the
cinematography, music and costumes are gorgeous. In fact, this movie is
prettier than those directed by Akira Kurasawa himself. In this case, he only
wrote the movie as it was made several years after his death.So, as far as the
writing goes, the dialog was well-written and the story, at times, was
interesting. However, the story was also rather depressing yet uninvolving in
some ways–after all, it’s the story of a group of women who work in a brothel.
It’s interesting that although prostitution has been seen as a much more
acceptable business in Japan, the women STILL long for a better life. This
reminds me a lot of the movie Streets Of Shame, though Streets Of Shame’s
characters are a lot less likable and more one-dimensional.So, overall it gets
a 7–mostly due to everything BUT the writing. It’s too bad that the weakest
link in this movie is the story by the great Kurasawa.

HARD: Average review length - 124 words \- Hard to predict review. Very
positive sentiment with a negative label. Might reflect mislabeling.

NEGATIVE (0) this one of the best celebrity’s reality shows a ever saw. we can
see the concerts we can see the life of Britney, i love the five episodes. i
was always being surprised by Britney and the subjects of the show i think
that some people don’t watch the show at all we can how a great person she
his. she his really funny really gentle and she loves her fans and we can see
how she loves her work. i just don’t give a 10 because of k-fed he his a real
jerk he doesn’t seem to like Britney at all. I they make a second season of
this great show because it shows at some people how Britney really is. Go
Britney your the best and you will never leave our hearts.

POSITIVE (1) Okay, first of all I got this movie as a Christmas present so it
was FREE! FIRST - This movie was meant to be in stereoscopic 3D. It is for the
most part, but whenever the main character is in her car the movie falls flat
to 2D! What!!?!?! It’s not that hard to film in a car!!! SECOND - The story
isn’t very good. There are a lot of things wrong with it. THIRD - Why are they
showing all of the deaths in the beginning of the film! It made the movie suck
whenever some was going to get killed!!! Watch it for a good laugh , but don’t
waste your time buying it. Just download it or something for cheap

NEGATIVE (0) this is a great movie. I love the series on tv and so I loved the
movie. One of the best things in the movie is that Helga finally admits her
deepest darkest secret to Arnold!!! that was great. i loved it it was pretty
funny too. It’s a great movie! Doy!

####  C.8 Data-IQ on images (computer vision)

###### Goal.

We wish to assess the generality of Data-IQ. Whilst, in the main paper we
mainly tackled tabular data, we wish to showcase the potential utility of
Data-IQ on image data for computer vision tasks.

###### Experiment.

We train a convolutional neural network on the CelebA dataset, where the task
is to predict the gender based on the image of the celebrity. We then use
Data-IQ to analyze the training dynamics and categorize the data into
subgroups - ${\color[rgb]{0,0.88,0}Easy}$,
${\color[rgb]{1,0.39,0.13}Ambiguous}$ and ${\color[rgb]{1,0,0}Hard}$

###### Takeaway.

We notice that the ${\color[rgb]{0,0.88,0}Easy}$,
${\color[rgb]{1,0.39,0.13}Ambiguous}$ and ${\color[rgb]{1,0,0}Hard}$ groups
represent different phenomena in the image data. We summarize the findings in
Figure 33

Figure 33: Subgroups highlighted by Data-IQ on the CelebA dataset

####  C.9 Effect of averaging over the uncertainty metric

###### Goal.

We wish to assess the impact of averaging/taking an expectation over the
metric such as aleatoric uncertainty over training. Specifically, we assess
the case where an example has high variance at the beginning of training and
flattens with time, compared to an example that remains at a medium level
throughout training. These represent the EASY/HARD (high variance in the
beginning and then flattening) and AMBIGUOUS (medium level throughout)
subgroups respectively. The question is, by taking an average will the
uncertainties be similar, despite the training dynamics being different? If
the average uncertainties are not similar that means we can distinguish the
groups with different training dynamics.

###### Experiment.

We conduct an experiment to assess this by sampling 100 examples from each of
these groups and plot the following three metrics per sample:

  * •

Training dynamic (i.e. the aleatoric uncertainty over training epochs).

  * •

The distribution over aleatoric uncertainty scores across all training steps.

  * •

The “average” aleatoric uncertainty.

###### Takeaway.

The result as shown in Figure 34 below shows that even by taking an average,
there is still a distinct separation to be able to distinguish between the two
types of samples with different training dynamics. This suggests that
averaging over the uncertainty is indeed a useful metric to evaluate the
uncertainty of an individual data sample.

Figure 34: Assessment of whether we can distinguish groups with different
training dynamics when we average over the metric

####  C.10 Effect of additional sources of randomness

###### Goal.

We wish to assess the impact of incorporating additional sources of randomness
- for example dropout.

###### Experiment.

We have test the model parameterizations from Section 4.1, but additionally
trained an additional 3 models with the same architecture, yet with the
inclusion of different dropout masks with dropout probability:
$p=0.1,p=0.2,p=0.3$.

We then assess the Spearman correlation of the metrics across all model
combinations, as was done in the experiment on robustness to variation from
Section 4.1. The higher correlation, implies better robustness to variations.

###### Takeaway.

Data-IQ remains most robust to variation, even with this additional
randomness. We highlight the results on all four datasets below in Table 11.

Table 11: Comparison of robustness/consistency across different models on the
basis of the Spearman rank correlation

Dataset | (Ours) Data-IQ | Data Maps
---|---|---
Covid | $0.93\pm 0.07$ | $0.66\pm 0.16$
Support | $0.92\pm 0.02$ | $0.81\pm 0.11$
Prostate | $0.97\pm 0.01$ | $0.96\pm 0.01$
Fetal | $0.86\pm 0.06$ | $0.81\pm 0.13$

####  C.11 Effect of hyper-parameters on the recovery of estimated groups

###### Goal.

We wish to assess how model choices/parameterizations might affect the
recovery of the different subgroups. i.e. we desire that if a sample is
categorized as EASY with one parameterization, that similarly it should be
categorized as EASY for another parameterization. Similarly, for the other
subgroups - we desire that their subgroup characterization stays the same
across different model parameterizations.

###### Experiment.

We use the same sets of model parameterizations as assessed in Section 4.1
(Robustness to variation experiment). We then compute the overlap of sample
ids for each subgroup (${\color[rgb]{0,0.88,0}Easy}$,
${\color[rgb]{1,0.39,0.13}Ambiguous}$, ${\color[rgb]{1,0,0}Hard}$) between the
different parameterizations. This can be thought of as accuracy or sample
level stability for assignment to a specific subgroup.

###### Takeaway.

We present the results across all four datasets below in Table 12 which
suggest that recovery of subgroups is stable, with more than 90% of samples
being assigned to the same subgroup between models.

Table 12: Assessment of stability in group assignment across different
parameterizations. The score reflects the average overlap of id
characterization

Dataset | (Ours) Data-IQ
---|---
Covid | $0.94\pm 0.03$
Support | $0.91\pm 0.01$
Prostate | $0.95\pm 0.01$
Fetal | $0.77\pm 0.14$

####  C.12 Correlation of weights through training

###### Goal.

We wish to assess the correlation of weights through model training.

###### Experiment.

To provide this assessment we follow Jin et al [62], where weight correlation
is defined as the average cosine similarity between weight vectors of neurons.
We conduct experiment to assess the correlation of model weights over epochs
$e\in\\{1,2...E\\}$. i.e. we assess the average correlation of model weights
$(e_{1},e_{2}),(e_{2},e_{3})$ and so on.

###### Takeaway.

We note that [62], shows that LOW correlation is less than 0.2. Our validation
shows that the weights have LOW correlation, with an average weight
correlation across the 4 datasets of $0.1\pm 0.02$. Thus, even though our goal
is different from a typical sampling setting (i.e. we just want to measure the
dispersion of observed variables) - we still wish to highlight the low
correlation of weights between epochs.

However, as a recommendation for practitioners if this is a concern: a
practical solution is that the practitioner could decide to simply increase
the intervals for which the dynamics are sampled, thereby decreasing this
correlation even further.

####  C.13 Data sculpting: Absolute number of samples

###### Goal.

When performing data sculpting as in Section 4.3 we desire to know the
absolute number of samples removed in order to quantify the importance of the
effect of removing such datapoints on model performance.

###### Experiment.

We present this analysis for two scenarios:

  * •

Balanced SEER-CUTRACT: 1000 samples each

  * •

Unbalanced SEER-CUTRACT: 10000 samples in SEER and 1000 in CUTRACT

###### Takeaway.

For the balanced experiment the full Ambiguous subgroup (i.e. proportion=1)
has a total of 500 samples. Hence, when we sweep the x-axis representing the
proportions of Ambiguous samples; $p=\\{0,0.2,0.4,0.6,0.8,1\\}$, this is
reflective of $\\{0,100,200,300,400,500\\}$ absolute samples. This result
shows that the number of samples that can be excluded is fairly substantial.

We compare this to the unbalanced experiment. The results are in Figure 35. It
shows that a substantial number of Ambiguous samples can be removed to improve
performance. However, in the case of (b) when there are a lot of samples,
removing too many can also be harmful.

(a) US-UK Balanced

(b) US-UK Unbalanced

(c) UK-US Balanced

(d) UK-US Unbalanced

Figure 35: Data sculpting - absolute number of
${\color[rgb]{1,0.39,0.13}Ambiguous}$ samples removed vs test accuracy