SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nabeel Seedat
Data quality is crucial for robust machine learning algorithms, with the recent interest in data-centric AI emphasizing the importance of training data characterization. However, current data characterization methods are largely focused on classification settings, with regression settings largely understudied. To address this, we introduce TRIAGE, a novel data characterization framework tailored to regression tasks and compatible with a broad class of regressors. TRIAGE utilizes conformal predictive distributions to provide a model-agnostic scoring method, the TRIAGE score. We operationalize the score to analyze individual samples' training dynamics and characterize samples as under-, over-, or well-estimated by the model. We show that TRIAGE's characterization is consistent and highlight its utility to improve performance via data sculpting/filtering, in multiple regression settings. Additionally, beyond sample level, we show TRIAGE enables new approaches to dataset selection and feature acquisition. Overall, TRIAGE highlights the value unlocked by data characterization in real-world regression applications.
Data characterization, an important problem. Over a decade ago, it was recognized that "more data beats clever algorithms, but better data beats more data" [1], emphasizing the key role of training data in the performance and robustness of ML algorithms [2-4]. In particular, training data from the "wild" may have various errors or lack coverage, rendering it unsuitable for the ML task [5].
These observations motivate the need for data characterization, which aims to score each training sample based on its impact on the performance of the ML task. The training data can then be partitioned or rank ordered according to the score to give a high-level description of the sample and its utility. Data characterization provides a natural way for data sculpting [6, 7], which filters training samples that have a low or even detrimental effect on the ML task. The goal is that models built on the sculpted/filtered (and smaller) training data have improved performance.
In this work, we address the understudied data characterization problem in regression settings. In contrast, existing works on data characterization have primarily focused on the classification setting [8-13]. These methods leverage the probability distribution over discrete labels (softmax or logits) for scoring. However, in the equally important regression setting, the probability distribution over the continuous labels are not typically given by standard regression models. In fact, the majority of the regression models only output a point estimate, e.g. the predicted mean. Without access to probabilistic estimates, one cannot simply apply the existing classification-tailored methods for
Figure 1: TRIAGE systematically characterizes data in regression settings via a 3-step pipeline. (i) Compute TRIAGE scores: of each sample in the dataset, for each regression checkpoint, via a Conformal Predictive System (CPS). Over the training epochs/iterations, we (ii) Analyze the per-sample training dynamics, studying the trajectory of CPS probability scores for each training sample, with the scoring method operationalized to characterize the samples. (iii) Characterize samples, assigning them to different groups (over-, under- and well-estimated) based on the TRIAGE scores. These groups inform sculpting & other use-cases.
discrete labels to the regression setting. We provide further details why classification-tailored methods are inapplicable to the regression setting in Sec. 2 and Table 1. This non-applicability motivates the need for regression-tailored methods for data characterization, which are "not only principled, but also practical for a larger collection (...) of ML models" [4].
To anchor the need for data characterization in regression, let's consider a few concrete examples. In predicting house prices, an outlier might be a house with an extremely high price due to a unique or under-represented feature in the data. This could affect the model's performance, causing it to overestimate prices for similar houses without the unique feature. Alternatively, imagine predicting patient length of stay based on clinical features. If the dataset contains some erroneous features, the regression predictions may be suboptimal, potentially leading to incorrect clinical decisions.
In designing a data characterization framework for regression, we outline the following desired properties, drawing motivation from the capabilities of works in the classification setting [8-13]:
(P1) Consistency: The characterization of data examples should be stable and consistent across runs and for similar performing models with different parametrizations.
(P2) Improve performance: The characterization should guide the improvements to model performance. For instance, via data sculpting/filtering.
(P3) Informative for data collection: The characterization score should be informative to enable selection between datasets, as well as, feature collection/acquisition.
(P4) Regressor versatility: The characterization should be principled and practical for a large collection of ML models (e.g. not just neural networks).
To fulfill  $P1 - P4$ , we propose a framework to characterize data in regression settings, called TRIAGE (s.f. TRaining Identification & Auditing of Good Examples) — see Fig.1. TRIAGE introduces a novel regression scoring method—TRIAGE score, building on the recent development of Conformal Predictive Systems (CPS) [14]. CPS is a model-agnostic method that can turn the output of any regressor into a probability distribution for continuous labels, termed conformal predictive distribution (CPD). Unlike point estimates such as errors/residuals or conventional conformal prediction intervals, the CPD provides richer information through a full predictive distribution. This enables us to assess the probability of a prediction being over- or under- estimated.
We then operationalize the CPD (and TRIAGE score) in a novel way for data characterization, computing the trajectory of TRIAGE scores over training checkpoints, as a measure of how samples are learned differently during training. Such training dynamics, shown in Fig 1 and have been shown to contain a salient signal about the value of the sample to a learner [15-17]. By analyzing the training dynamics, we then categorize training samples as: under-, over- or well-estimated by the model. This is useful going back to our motivating example — where erroneous values could cause the model to overestimate house price. Using the TRIAGE score, we can identify such a sample, as it would likely have a low probability of being correct based on the model's predictive distribution. Such identification unlocks a variety of use cases. For instance, sculpting/filtering such samples can improve model performance (even with less data), as we empirically show in Sec. 5.
Contributions. ① Conceptually, we present TRIAGE, the first data characterization framework tailored to regression settings, addressing the need for an ML-aware data quality that is principled and practical across a large collection of ML regressors. ② Technically, we introduce a principled regression scoring method called TRIAGE score, that leverages conformal predictive distributions in a novel way for data characterization. We operationalize the score to analyze how different samples are
Table 1: Comparing data characterization for classification to TRIAGE (a tailored regression method)
Type of method	Approaches	Scoring method	Can be repurposed for regression	Applicable to any ML model trained iteratively	Reason for non-applicability
Tailored regression	TRIAGE (Ours)	CPS probability	✓	✓	-
Tailored classification	Data Maps [8]	Model uncertainty	✗	✓	Needs label probability distribution
Data-IQ [9]	Data uncertainty	✗	✓	Needs label probability distribution
SSFT [10]	Forgetting scores	✗	✓	Needs discrete class changes
AUM [12]	Logit margin	✗	✗	Needs logit output
General-purpose	Metadata Archaeology [18]	Per-sample losses	✓	✓	-
VoG [11]	Variance of gradient	✓	✗	Needs differentiable model (e.g. NN)
GraNd [13]	Gradient Norm	✓	✗	Needs differentiable model (e.g. NN)
learned during training, thereby permitting characterization of individual samples as: under-, over- or well-estimated by the model.  $③$  Empirically, we demonstrate the utility of TRIAGE across multiple use cases satisfying P1-P4, including consistent characterization, sculpting to improve performance in a variety of settings, as well as, guiding dataset selection and feature acquisition.
This work engages with recent efforts in data-centric AI [6, 7, 19, 20] on data characterization. However, prior methods have focused on classification settings. In contrast, TRIAGE addresses the regression setting where tailored data characterization has not been explored. We now discuss related works and outline the challenges of applying prior classification focussed methods to regression, for an overview see Table 1. Prior classification scoring methods can be divided into two groups:
■ Tailored methods, assign scores based on aspects specific to classification, including logits [12], probabilities over labels [8, 9] or changes to predictions on discrete labels [10, 21]. However, these methods are inapplicable to our regression setting because these aspects are not present in regression.
■ General purpose methods, while focussed on classification, they assign scores using generic measures and can be repurposed for regression. However, these general methods based on backpropagated gradients [11, 13] or losses [18] have two limitations. Firstly, scores like gradients can only be used with differentiable models (e.g. neural networks) and are not compatible with many widely used regressors, such as XGBoost — i.e. does not satisfy (P4) Versatility. Secondly, even if they can be used, general-purpose scores may not distinguish between samples that have similar magnitude scores (e.g. high loss scores), even if the underlying reasons for these scores may be different with respect to the task. We show the value of such a differentiation between samples in Sec.5.2.
Data valuation. A related area to data characterization is that of data valuation, for example: Shapley based valuation methods [22, 23]. We first outline some conceptual differences: (a) TRIAGE unlocks other data-centric tasks beyond just data selection, e.g. feature-level acquisition/collection and selection between datasets. These are out of scope for data "sample" selection; (b) Unlike TRIAGE, these methods are not tailored to regression; (c) method differences: computationally Shapley-based methods need to assess multiple sample permutations and need to retrain the model many times. Hence, they struggle to scale to very high-samples sizes (e.g. 100k), unlike TRIAGE where the cost is cheap comparatively. Recently, a valuation method LAVA [24] has been introduced to address this which uses an additional embedding model to reduce dimensionality before the optimal transport step. Beyond conceptual differences, we compare TRIAGE experimentally and in terms of computation time in Appendix C.
3 Data Characterization Preliminaries
In this section, we formulate the data characterization problem for regression settings. Consider the standard regression setting, where the aim is to assign an input  $x \in \mathcal{X} \subseteq \mathbb{R}^{d_X}$  to a continuous valued label  $y \in \mathbb{R}$  (often a point estimate for  $y$ ). This is done using a regressor  $f_{\theta}: \mathcal{X} \to \mathcal{Y}$ , parameterized by  $\theta \in \Theta$ . We assume  $f_{\theta}$  is trained on a dataset  $\mathcal{D}$  with  $M \in \mathbb{M}^*$  observations, i.e.  $\mathcal{D} = \{(x^m, y^m) | m \in [M]\}$ . For the purposes of later notation, we consider a sample space  $\mathcal{Z} := \mathcal{X} \times \mathbb{R}$ , where each observation  $z^m = (x^m, y^m) \in \mathcal{Z}$ . Typically, parameters  $\theta$  are learned in regression, via empirical risk minimization (ERM) with a mean squared error (MSE) loss, i.e.  $\mathrm{ERM}(\theta) = \arg \min_{\theta \in \Theta} \frac{1}{M} \sum_{m=1}^{M} \ell(x^m, y^m; \theta)$ , with loss function  $\ell: \mathcal{X} \times \mathcal{Y} \times \Theta \to \mathbb{R}^+$ .
ERM ensures regressor  $f_{\theta}$  performs well on average, however it goes without saying that the model is unlikely to perform equally well for all samples in the dataset. Consequently, we desire to characterize such types of samples. Formally, for each sample  $x^{m}$ , we aim to assign a group label  $g^{m} \in \mathcal{G}$ , where
$\mathcal{G} = \{\text{Over-estimated, Well-Estimated, Under-estimated}\}$ . We then use these groups for downstream tasks, or to sculpt/filter samples in the dataset  $\mathcal{D}$  based on their group assignment. Before giving a precise description of how the group labels are assigned, we provide some context.
Prior works have shown the behavior of samples over the course of model training, called training dynamics contains signal about the nature of a sample itself [15-17]. For instance, some samples are easily learned and stable through training, whereas other samples might be challenging and have high variance. Many classification-based methods leverage this concept and assign  $g^{m}$  by analyzing training dynamics e.g. Data-IQ, Data Maps, AUM etc. However, as discussed in Secs. 1 and 2, the aforementioned classification methods do not apply to regression. Consequently, this requires us to formalize a new principled scoring method for regression. After that, we can operationalize the scoring method to analyze the training dynamics and assign group labels  $g^{m}$  to each  $x^{m}$ .
For regression, such a principled scoring method could evaluate the probability that the true target  $y$  is less than or equal to observation  $f(x)$ , where the categories of samples are: ① Under-estimated:  $P(y \leq f(x)) \gg 0.5$ , ② Well-estimated:  $P(y \leq f(x)) \sim 0.5$ , ③ Over-estimated:  $P(y \leq f(x)) \ll 0.5$ . In the next section, we discuss how we design TRIAGE to achieve this, with the final goal in mind — assigning group labels  $g^{m}$  to each sample  $x^{m}$ .
4 TRIAGE: characterizing data for regression settings
We now introduce TRIAGE, a regression-focused data characterization framework, compatible with any iteratively trained ML regressor (Sec. 4.4). TRIAGE follows a three-step procedure:
1. Compute TRIAGE scores (Sec 4.1): Compute the TRIAGE score for every sample  $x^{m}$  using conformal predictive distributions. Do this for every regression checkpoint to obtain a trajectory.
2. Analyze TRIAGE score training dynamics (Sec 4.2): Operationalize the TRIAGE score for data characterization by analyzing the TRIAGE score's trajectory for each training sample.
3. Characterize & assign groups (Sec 4.3): Characterize samples based on their training dynamics, assigning each sample  $x^{m}$  to a group  $g^{m}$ . The groups inform filtering/sculpting or other use cases.
4.1 Computing TRIAGE scores via Conformal Predictive Systems (Step 1).
Recall, we want to compute  $P(y \leq f(x))$  as a principled scoring method to characterize samples. We want to do so without changing the chosen regressor type, i.e. P4: Regressor versatility. Of course, the full information of an uncertain label exists in its probability distribution, allowing us to compute probabilities of various outcomes. Unfortunately, regular regressors (point predictions) or conventional conformal regressors (prediction intervals) do not provide this distribution. While Bayesian methods offer distributions, they require fundamentally changing the regressor.
These challenges motivate defining TRIAGE around Conformal Predictive Systems (CPS)[14]. CPS combines predictive distributions [25, 26] with conformal prediction [27], crucially enabling us to compute the required probability distribution of a continuous variable—called Conformal Predictive Distribution (CPD). CPS also fulfills our versatility goal (i.e.  $P4$ ) as it applies to any regressor. To clarify Conformal Predictive Systems output valid cumulative distribution functions, which are termed Conformal Predictive Distributions (CPD). This is the cumulative probability with respect to a label  $y$ , given some  $x$  and regressor  $f$ . With CPDs denoted as  $Q$ , the conformal p-values get arranged into a probability distribution which has the properties of a CDF — thus essentially becoming probabilities, see [28]. Since the CPD has the properties of a CDF, we use the CPD to estimate probabilities that the true target is less than or equal to a specified threshold/value.
Computing CPDs. TRIAGE formalizes the CPD computation in the framework of the computationally efficient Split Conformal Predictive Systems [28]. A proper training set  $\mathcal{D}_{\mathrm{tr}} = \{(x_{\mathrm{tr}}^i, y_{\mathrm{tr}}^i) \mid i \in [M]\}$  is used to train regressor  $f$ . A separate, often smaller, calibration set  $\mathcal{D}_{\mathrm{cal}} = \{(x_{\mathrm{cal}}^i, y_{\mathrm{cal}}^i) \mid i \in [q]\}$ , is used for conformal calibration. Finally, we define a conformity measure  $\mu$  where  $\mu: \mathcal{Z} \to \mathbb{R} \cup \{-\infty, \infty\}$ , quantifies the dataset's agreement with the observation. For each label  $y_{\mathrm{cal}}^i \mid i \in [q]$ , we compute  $q$  conformity scores  $\alpha_i = \mu(z_{\mathrm{cal}}^i)$ . We define the conformity measure  $(\mu)$  as per Eq. 1, which is a balanced isotonic conformity measure. We discuss this choice in Appendix A.3, a necessary condition s.t the CPD can be interpreted as a probability distribution [28]
$$
\mu (x, y) = \left(y - f (x)\right) / \sigma (x) \tag {1}
$$
<sup>1</sup>Under the standard conformal exchangeability assumption (see Appendix A.3), the CPDs are valid.
where  $y$  is the label,  $f(x)$  a predicted label,  $\sigma$  an estimate of prediction quality, computed with a separate normalization model, enabling CPDs to reflect task difficulty.  $\sigma$  is kNN fitted on the residuals of  $\mathcal{D}_{\mathrm{cal}}$ , similar to [29].
We compute the CPD denoted  $Q^{m}$  (Eq. 2) for each sample  $x^{m}$  via Algorithm 1, where the CPD is defined for  $y^{*} \in \mathbb{R}$  [28]. We extract the TRIAGE score using the CPD. We repeat the CPD computation for all samples. Next, we describe our novel approach to operationalize the TRIAGE scores, for the purpose of data characterization.
Algorithm 1 Computing a CPD
Require: Training set  $(x_{\mathrm{tr}}^i,y_{\mathrm{tr}}^i)\in \mathcal{D}_{\mathrm{train}},i\in [M]$  and Calibration set  $(x_{\mathrm{cal}}^i,y_{\mathrm{cal}}^i)\in \mathcal{D}_{\mathrm{cal}},i\in [q]$
1: Regressor  $f$  is trained on  $D_{\mathrm{train}}$ .
2: Compute residuals on  $\mathcal{D}_{\mathrm{cal}}$  ..  $R_{i} = y_{\mathrm{cal}}^{i} - f(x_{\mathrm{cal}}^{i}),i\in [q]$
3: Define  $\sigma (x)$ ; which finds the average residuals of
the k-nearest neighbors  $x$ , where KNN is fit on  $\mathcal{D}_{\mathrm{cal}}$ .
4: Calculate conformity scores  $\alpha_{1},\alpha_{2}\dots \alpha_{q}$  over  $\mathcal{D}_{\mathrm{cal}}$
using  $\mu$ ; i.e.  $\alpha_{i} = \mu (x_{\mathrm{cal}}^{i},y_{\mathrm{cal}}^{i}),i\in [q]$ .
5: Sort  $\alpha_{1},\ldots \alpha_{q}$  in ascending order; s.t  $\alpha_{(1)} < \dots < \alpha_{(q)}$
for eval sample  $x\in \mathbf{X}_{train}$  do
Let  $S_{(i)}\coloneqq f(x) + \alpha_i\sigma (x)$  for  $i\in [q]$
Let  $S_{(0)}\coloneqq -\infty$  and  $S_{(q + 1)}\coloneqq \infty$
Return the CPD as per Equation (2).
$$
Q \left(y ^ {*}, \tau\right) := \left\{ \begin{array}{l l} \frac {i + \tau}{q + 1} & \text {i f} y ^ {*} \in \left(S _ {(i)}, S _ {(i + 1)}\right); i \in [ 0: q ] \\ \frac {i ^ {\prime} - 1 + \left(i ^ {\prime \prime} - i ^ {\prime} + 2\right) \tau}{q + 1} & \text {i f} y ^ {*} = S _ {(i)}; i \in [ 1: q ] \end{array} \right. \tag {2}
$$
where  $\tau \in [0,1]$ ,  $i' := \min \{j \mid S_{(j)} = S_{(i)}\}$  and  $i'' := \max \{j \mid S_{(j)} = S_{(i)}\}$  to account for ties.
4.2 Analyzing TRIAGE score trajectories — training dynamics (Step 2)
Novel use of CPD's for data characterization. The TRIAGE score, denoted as  $T(x,y,\theta) = P(y \leq f_{\theta}(x))$ , represents a novel use of CPDs for data characterization, diverging from their original predictive uncertainty use. Additionally, while CPDs are normally calculated post-model fitting, we observe from Fig.1(ii) that samples' learning varies, leading to different convergence rates of scores. As discussed in Sec. 3, such dynamics reveal sample characteristics [15-17]. Hence, we propose to examine TRIAGE score trajectories to capture the learning dynamics of different samples.
To enable this analysis, we utilize the model checkpoints of any ML regressor trained over iterations (e.g. neural networks, XGBoost etc). Specifically, during iterative training, the model parameters  $\theta$  vary over  $E\in \mathbb{N}^*$  iterations, taking  $E$  different parameter values at each checkpoint, i.e.  $\theta_{1},\theta_{2},\dots,\theta_{E}$ . We wish to account for the impact of these changing parameters. At each checkpoint  $\theta_{e}$ , we compute the CPD to obtain the TRIAGE score for each sample  $x^{m}$ . We repeat the computation for every checkpoint to obtain a TRIAGE score trajectory over checkpoints per-sample of dimensionality  $E$ . i.e. we evaluate the function  $Q$  for a specific  $f_{\theta_e}(x)$  to get the estimated probability  $P(y\leq f_{\theta_e}(x))$ . We repeat for all  $f_{\theta_e}(x)$  checkpoints where  $e\in E$  to get the trajectory of TRIAGE scores for sample.
To quantify and contrast the behavior of the different sample trajectories, we compute two metrics to summarize the per-sample trajectory (see Fig. 2), namely (1) Confidence: reflects the degree of confidence wrt to  $P(y \leq f_{\theta}(x))$ . It is computed as the mean of the TRIAGE score over training, i.e.  $C(x^{i},y^{i}) = \frac{1}{E}\sum_{e = 1}^{E}T(x^{i},y^{i},\theta_{e})$ . (2) Variability: reflects the consistency or fluctuation of the score for a specific sample. It is computed as the standard deviation of the TRIAGE score over training, i.e.  $V(x^{i},y^{i}) = \sqrt{\frac{\sum_{e = 1}^{E}(T(x^{i},y^{i},\theta_{e}) - C(x^{i},y^{i}))^{2}}{E}}$ . We describe next how these two metrics form the basis of sample characterization and group assignment.
4.3 Characterization & group assignment (Step 3)
We now discuss how the two characterizing metrics,  $C$  and  $V$ , are used to assign a group label  $g \in \mathcal{G}$  to each training sample  $x^{m}$ . Let's intuitively define each group. ① Under-estimated (UE): low variability  $V$  (i.e. decisive about) and HIGH average TRIAGE score  $C$ , i.e.  $P(y \leq f(x)) \gg 0.5$ , ② Over-estimated (OE): low variability  $V$  (i.e. decisive about) and LOW average TRIAGE score  $C$ , i.e.  $P(y \leq f(x)) \ll 0.5$ , and ③ Well-estimated (WE): average TRIAGE score  $C$  is uncertain  $P(y \leq f(x)) \sim 0.5$ . In addition, these samples are often linked to high variability, implying uncertainty about  $P(y \leq f(x))$ . We empirically observe samples with these different traits per group, depicted in the characteristic curve, see Fig. 2.
We take a threshold-based approach to assign  $g^{m}$ , similar to previous characterization works [8, 9]:
$$
g ^ {m} = \left\{ \begin{array}{l l} \text {U E} & \text {i f} C \left(x ^ {m}, y ^ {m}\right) \geq C _ {\mathrm {u p}} \wedge V \left(x ^ {m}, y ^ {m}\right) <   P _ {n} ^ {V \left(\mathcal {D} _ {\text {t r a i n}}\right)} \\ \text {O E} & \text {i f} C \left(x ^ {m}, y ^ {m}\right) \leq C _ {\text {l o w}} \wedge V \left(x ^ {m}, y ^ {m}\right) <   P _ {n} ^ {V \left(\mathcal {D} _ {\text {t r a i n}}\right)} \\ \text {W E} & \text {o t h e r w i s e} \end{array} \right. \tag {3}
$$
where  $C_{\mathrm{up}}$  and  $C_{\mathrm{low}}$  are upper and lower confidence threshold resp. and  $P_{n}$  the n-th percentile. We discuss the selection of  $C_{\mathrm{up}}$  and  $C_{\mathrm{low}}$  in Appendix A.
In Sec. 5, we empirically show how the groups assigned to each sample can be used to sculpt/filter the training dataset or for other downstream tasks such as dataset selection or feature acquisition.