SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2: Training dynamics of samples (above) are summarized in a characteristic curve (below-grey)
4.4 Using TRIAGE with a variety of regressors (P4)
The general-purpose characterization methods discussed are largely compatible with differentiable models (neural networks). However, in real-world settings (e.g., healthcare, finance), practitioners often employ other iterative algorithms, like XGBoost/GBDTs [30]. TRIAGE addresses this limitation, offering a model-agnostic approach using CPS — as the conformity scores can be computed for any regressor. Of course, our use of training dynamics assumes regressor  $f$  can be checkpointed. This still confers compatibility with a broader class of models: neural nets, XGBoost, GBDTs, linear regression — satisfying P4. We underscore TRIAGE's versatility beyond neural networks, allowing practitioners to use TRIAGE with their preferred model for a specific application. Appendix A.2 details the ease of usage and space and time considerations.
5 Empirical investigation
This section presents an empirical evaluation demonstrating TRIAGE<sup>2</sup> satisfies (P1) Consistency, (P2) Improved regression performance and (P3) Informative for data collection. Recall (P4) Regressor versatility is satisfied by TRIAGE's design. Additional use-cases and results are in Appendix C.
Datasets. We conduct experiments on 10 real-world regression datasets with varying characteristics. i.e. different sample sizes (500-100k), dimensionality (8-85), and task difficulty. The datasets are drawn from diverse domains, including safety-critical medical regression: (i) Prostate cancer from the US [31] and UK [32], (ii) Hospital Length of Stay [33] and (iii) MIMIC Antibiotics [34]. Additionally, we analyze general UCI regression datasets [35], including Bike, Boston Housing, Bio, Concrete, Protein and Star. The datasets are detailed in Appendix B, along with further experimental details. We observe similar performance across different datasets. However, due to space limitations, we sometimes highlight results for a subset and include results for the remainder in Appendix C.
5.1 (P1) Consistent characterization
A key aspect of data characterization is: the ordering of samples matter [9, 10]. This is important when the scores are utilized for data sculpting/filtering. For example, if we rank sort samples by their scores and filter a subset of samples. Consequently, it is crucial to ensure the scoring method used for data characterization is stable and consistent. This consistency is desired across different seeds, but also different practitioners may adopt diverse model architectures/parameterizations for data characterization.
Consistency of TRIAGE. We compare the consistency and robustness to variation of TRIAGE vs baseline general purpose methods: VoG [11], GradN [13], Loss [18]. We also consider simply computing the end of training Residual Errors (used in the
Figure 3: Stability of scoring methods. TRIAGE is most consistent (highest Spearman correlation of 0.91 & lowest variation) averaged across datasets. Results on all 10 datasets are in Appendix C.1.
$^{2}$ Code: https://github.com/seedatnabeel/TRIAGE or https://github.com/vanderschaarlab/TRIAGE
TRIAGE conformity score). We assess the consistency of the scores across different model architectures/parameterizations (described in Appendix B.3). All models are trained to convergence with early stopping. Similar to prior classification-based characterization methods [9, 10], we assess consistency based on the Spearman rank correlation of the scores evaluated for all model combinations.
Fig. 3, shows TRIAGE is the most consistent compared to baselines, with the highest Spearman correlation (0.91), satisfying  $P1$ . This consistency means insights derived using TRIAGE scores will also remain consistent. Analyzing each of the 10 datasets (see Appendix C.1), further shows that TRIAGE is stable in magnitude and rank order, whereas the baselines do not have such consistency both in magnitude or ordering across datasets, making the baselines challenging to use in practice.
Takeaway 1. TRIAGE has the most consistent scoring method compared to baselines over different datasets and model combinations, satisfying P1.
TRIAGE v Errors/Residuals. We note that TRIAGE and Errors/Residuals (considered an ablation of TRIAGE's dynamics component) have similar correlations. This raises two questions: (i) are the samples identified the same? We stratify the samples in  $\mathcal{D}_{\mathrm{train}}$  into three equal sized groups ordered by TRIAGE or Error scores. The overlap is  $\sim 0.77$ , highlighting a difference. This leads us to ask (ii) what added value does TRIAGE offer? Consider the case of different samples that have the same error magnitude. Errors/Residuals treat these samples the same, based solely on magnitude. Can TRIAGE offer greater precision to further differentiate samples with the same magnitude?
Figure 4: Samples with the same error, have different Triage scores, highlighting the potential to differentiate samples
As shown in Fig.4, TRIAGE scores evidently offer a way to handle samples differently, showing that samples with similar residuals/errors are associated with different TRIAGE scores. We show the phenomena holds across other general-purpose scoring rules in Appendix C.10. This result demonstrates the added value of the training dynamics aspect, to differentiate samples with the same error magnitude at a more granular level. This contrasts, errors/residuals which are an ablation of TRIAGE, computing the conformity score only at the final training checkpoint, instead of for all checkpoints. We explore the value of this differentiation to improve regressor performance next.
5.2 (P2) Improve regression performance
We now evaluate the ability of TRIAGE to improve regression performance by sculpting/filtering the data — by keeping only well-estimated samples. We study two setups using TRIAGE: (1) fine-grained filter (as described above) and (2) sculpting the data based on a deployment purpose.
5.2.1 Fine-grained filtering
In the previous section, we showed TRIAGE's ability to differentiate samples with more precision. Does this differentiation with TRIAGE scores benefit model performance? To find out, we sort the training samples from least to most challenging based on the baseline computed scores: Residuals/Errors, Loss, VoG, GradN. We then evaluate test MSE performance, as we increase the proportion  $p$  of samples retained for training. i.e. increasingly retain more challenging samples with higher loss/VoG etc. As expected, the baselines (shown in Fig 5, blue) have lower MSE with more sam
(a) Errors
(c) Grad
(d) VoG
Figure 5: TRIAGE improves test MSE, as a fine-grained filter. For each baseline proportion retained, TRIAGE is applied on top — such that TRIAGE has less data ( $D_{Triage} \subset D_{Baseline}$ ), only keeping well-estimated samples for each proportion  $p$ .
plies retained for training. However, the elbow around  $p = 0.5$ , after which MSE does not decrease, shows that the added high-magnitude samples are not beneficial.
We now assess the benefit of differentiating these scores using TRIAGE. To do so, we apply TRIAGE as an additional filter on top of the baseline retained samples — only keeping the well-estimated subset of samples. i.e. at each baseline proportion  $p$ , TRIAGE reduces the number of samples retained compared to the baseline s.t.  $D_{\text{Triage}} \subset D_{\text{Baseline}}$ . Fig 5 shows, for all proportions  $p$ , we improve MSE performance over baselines simply by differentiating between samples using TRIAGE and filtering to keep only well-estimated samples. This performance benefit holds for all 4 general-purpose scoring methods, illustrating the value of TRIAGE as a fine-grained filter, where even with less data than baselines, the careful selection of samples by TRIAGE leads to improved regression performance. Results on more datasets mirroring the above are in Appendix C.11.
Takeaway 2. TRIAGE's differentiation of samples can be used to augment other approaches as a fine-grained filter; improving regressor test performance, simply by keeping well-estimated samples.
5.2.2 Data sculpting to fit a deployment purpose.
In many situations, we might want to take a regressor trained on a large  $\mathcal{D}_{\mathrm{train}}$  and deploy it in an area with limited labeled data. For instance, repurposing a regressor in a low-resource healthcare setting. This raises the question, instead of acquiring new data - which in some cases might be impossible - could we sculpt the existing large  $\mathcal{D}_{\mathrm{train}}$  to be more "fit for purpose". The goal is to enable better regressor performance. Our setup assumes limited access to a small amount of data from the deployment environment, which is insufficient in sample size to directly train a highly performant regression model, yet we could use this data for calibration  $\mathcal{D}_{\mathrm{cal}}$  purposes. Can TRIAGE help?
As a case-study, we consider a regression task similar to [36-38], using patient covariates to predict PSA blood levels, a known predictor of prostate cancer mortality. This task has utility in low-resource areas with limited medical funds [38]. We use multi-country prostate cancer data in which their is a distribution shift and train a regressor on data from the US (SEER)[31]. We evaluate performance on data from the UK (CUTRACT)[32], where the smaller  $\mathcal{D}_{\mathrm{cal}}$  is also drawn from. In addition to the normal train-large, test-small baseline, we evaluate the utility of sculpting  $\mathcal{D}_{\mathrm{train}}$  from the US using TRIAGE — only keeping well-estimated samples ( $D_{WE} = D_{Triage}$ ), s.t  $D_{Triage} \subset \mathcal{D}_{\mathrm{train}}$ . Conformal calibration uses  $\mathcal{D}_{\mathrm{cal}}$  from the UK.
We compare TRIAGE to several methods, as detailed in Appendix B.3, including: (i) data-driven training: training a model directly on the  $\mathcal{D}_{\mathrm{cal}}$  or combining the two datasets  $(\mathcal{D}_{\mathrm{train}} \cup \mathcal{D}_{\mathrm{cal}})$ . (ii) prediction-based sculpting: train on  $\mathcal{D}_{\mathrm{cal}}$  (UK data) and then subsequently filter  $\mathcal{D}_{\mathrm{train}}$  (US data) based on predictive uncertainty scores<sup>3</sup>. We assess NGBoost, Bayesian Neural Network, Gaussian Process, Bayesian ridge regression, conformal intervals and residuals/errors. We only consider general purpose methods whose criteria allow us to leverage  $\mathcal{D}_{\mathrm{cal}}$  and are compatible beyond neural nets for flexibility.
We then evaluate MSE test performance of regressors trained on the different "sculpted"  $\mathcal{D}_{\mathrm{train}}$ . We test sensitivity to varying sizes of  $\mathcal{D}_{\mathrm{cal}}$ , as in many settings, we might not have access to many calibration samples. Table 2 shows that TRIAGE consistently outperforms model-driven methods across the board. This suggests that predictive uncertainty scores are insufficient to select the samples with the highest utility in  $\mathcal{D}_{\mathrm{train}}$ . This contrast implies TRIAGE's scoring method better selects samples in  $\mathcal{D}_{\mathrm{train}}$  with greater utility, leading to improved performance on  $\mathcal{D}_{\mathrm{test}}$ .
Compared to data-driven approaches, we find that TRIAGE has the best MSE, specifically, in the small sample regime, where  $\mathcal{D}_{\mathrm{cal}}$  only contains a limited number of samples  $n = 10 - 200$ . Of course, when sufficient data is available ( $\geq 300$  samples), sculpting has a reduced benefit. Rather, training directly on the sufficient  $\mathcal{D}_{\mathrm{cal}}$  is preferable. The result highlights the viability of sculpting a larger  $\mathcal{D}_{\mathrm{train}}$  with respect to a small  $\mathcal{D}_{\mathrm{cal}}$  to improve regression in settings with limited data access — with TRIAGE offering the best performance benefits.
In addition, sculpting especially in healthcare settings should account for both (i) imbalanced feature subgroups and (ii) long-tailed outcomes. We assess this in Appendix C.7 and show that TRIAGE does retain strong performance on minority subgroups and long-tailed outcomes by virtue of the calibration dataset.
3Model-driven methods are fit on  $\mathcal{D}_{\mathrm{cal}}$  , such that we assess how uncertain samples in  $\mathcal{D}_{\mathrm{train}}$  are with respect to  $\mathcal{D}_{\mathrm{cal}}$
Table 2: Comparison of test MSE for different approaches, changing the size of  $\mathcal{D}_{\mathrm{cal}}$ . TRIAGE's approach to data sculpting of the large  $\mathcal{D}_{\mathrm{train}}$  with respect to a small  $\mathcal{D}_{\mathrm{cal}}$  in shown to outperform alternatives with lower test MSE, especially when  $\mathcal{D}_{\mathrm{cal}}$  from the target domain has small sample sizes. (↓ better). Note  $|\mathcal{D}_{\mathrm{train}}| \gg |\mathcal{D}_{\mathrm{cal}}|$
Dcal sample size	10	20	30	40	50	100	200	300	400	500
Ours (Sculpting)	TRIAGE (WE)	0.051+0.003	0.050+0.003	0.047+0.002	0.046+0.002	0.046+0.001	0.046+0.002	0.046+0.001	0.045+0.001	0.045+0.002	0.046+0.002
Not TRIAGE (OE UUE)	0.088+0.014	0.068+0.005	0.066+0.003	0.066+0.005	0.068+0.006	0.080+0.011	0.087+0.015	0.093+0.015	0.082+0.012	0.078+0.009
Data-driven	Baseline (Drain)	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.064+0.002	0.111+0.009
Baseline (Cal)	0.092+0.023	0.070+0.012	0.070+0.010	0.064+0.003	0.066+0.009	0.052+0.004	0.047+0.001	0.043+0.001	0.043+0.001	0.111+0.009
Baseline (Drain ∪ Cal)	0.060+0.002	0.059+0.002	0.058+0.002	0.056+0.002	0.055+0.001	0.049+0.001	0.047+0.001	0.044+0.001	0.044+0.001	0.111+0.009
Prediction based sculpting of Dtrain	Error Sculpt	0.066+0.010	0.058+0.010	0.060+0.010	0.060+0.010	0.057+0.010	0.058+0.010	0.055+0.011	0.055+0.011	0.057+0.010	0.111+0.009
CP Intervals Sculpt	0.064+0.006	0.081+0.023	0.082+0.008	0.106+0.043	0.073+0.013	0.059+0.005	0.055+0.002	0.063+0.011	0.054+0.002	0.111+0.009
NGBoost	0.066+0.006	0.121+0.041	0.085+0.015	0.080+0.008	0.094+0.008	0.196+0.030	0.132+0.010	0.096+0.011	0.095+0.006	0.103+0.008
Bayesian ridge	0.080+0.008	0.118+0.023	0.111+0.012	0.129+0.011	0.129+0.016	0.111+0.017	0.106+0.006	0.110+0.010	0.114+0.009	0.103+0.009
BNN	0.068+0.005	0.063+0.006	0.066+0.004	0.066+0.004	0.064+0.005	0.057+0.006	0.072+0.010	0.056+0.006	0.064+0.003	0.111+0.009
GP	0.051+0.006	0.066+0.008	0.069+0.005	0.077+0.008	0.072+0.006	0.071+0.009	0.085+0.011	0.085+0.003	0.095+0.004	0.111+0.009
Data insights from sculpting. We seek to understand what types of samples are "sculpted". Such insights are especially useful in clinical settings. We see that US patients  $(\mathcal{D}_{\mathrm{train}})$  typically have higher cancer stages and higher PSA scores than their UK counterparts  $(\mathcal{D}_{\mathrm{cal}} / \mathcal{D}_{\mathrm{test}})$ . It is precisely these dichotomous samples that are filtered by TRIAGE to improve predictive performance. These insights can be illustrated to stakeholders via a radial diagram, as in Appendix C.5.
Additionally, given the potential data shift (around exchangeability), Appendix C.6 analyzes CPS calibration, along with continuous ranked probability score (CRPS) quantifying the quality of the predictive distribution.
Beyond predictive performance to fairness. We also show that sculpting the data has potential beyond simply improving predictive performance. In Appendix C.4, we illustrate how TRIAGE can be used to sculpt data to improve fairness of regressors based on access to a limited set of samples.
Takeaway 3. TRIAGE can help to improve regression performance by sculpting a larger dataset to be fit-for-purpose, simply by leveraging a small number of calibration samples.
5.3 (P3) Informative for data collection
We now demonstrate the value of TRIAGE for dataset selection and feature collection/acquisition.
5.3.1 Dataset selection: From sample-level to dataset level.
In real-world applications, we might want entire dataset characterization, rather than just sample-level characterization. An understudied scenario is how to compare and select between different datasets. This need arises in organizations where data is stored in isolated silos and it is difficult to access; with synthetic data a common solution [39, 40] or alternatively when purchasing data via data marketplaces [41]. We study the former scenario, where we have multiple 'synthetic' data versions, produced by different ML models or data vendors. However, our assessment could easily apply to the latter of data marketplaces.
We explore the setting where the real  $\mathcal{D}_{\mathrm{train}}$  is not accessible to calculate statistical similarity measures, yet we still wish to compare synthetic datasets and select a version to train on. The underlying assumption is that the synthetic datasets reflect the true underlying distribution. However, depending on the generation process, some synthetic datasets might prove superior to others.
The question we address is: Can TRIAGE's data characterization allow us to compare and select the synthetic training dataset that yields the best test performance on real data?
We simulate this scenario by generating synthetic training data using (V1) CTGAN and (V2) TVAE [42], representing 2 synthetic data vendors. We employ TRIAGE to analyze and compare them, based on the percentage of well-estimated examples recommended for retention. As is standard, we train a regressor on synthetic data and test it on real data (TSTR) [43, 44]. The best synthetic data should yield the lowest MAE on real test data. Table 3 shows that where TRIAGE's recommendation has a higher retention percentage (i.e. more well-estimated samples), that training with that dataset, leads to lower MAE on real test data. This illustrates that even if the real data is
Table 3: Comparing performance & (Triage retention). Training on synthetic data w/ better utility (↑ retained) has the lowest real test data MAE across datasets.
Dataset	(V1) CTGAN	(V2) TVAE
Bike	0.13 (45% Retained)	0.09 (58% Retained):
Bio	0.24 (45% Retained)	0.21(52% Retained)
Boston	0.16 (38% Retained)	0.13(60% Retained)
Concrete	0.19 (46% Retained)	0.11 (48% Retained)
LOS	0.16 (47% Retained)	0.09(56% Retained)
MIMIC	0.05 (41% Retained)	0.08 (48% Retained)
Prostate	0.21 (46% Retained)	0.18 (52% Retained)
Protein	0.22 (52% Retained)	0.20 (54% Retained)
Star	0.18 (38% Retained)	0.14 (51% Retained)
unavailable, TRIAGE's characterization can be used by practitioners to select the synthetic training dataset that will give the best real test performance.
Takeaway 4. TRIAGE scores allow us to assess utility between different synthetic datasets in a principled manner. We show that synthetic datasets that TRIAGE suggests retaining more samples (i.e. more well-estimated), allow for better generalization when testing on real-data (i.e. MAE).
5.3.2 Value of feature collection/acquisition.
Previous settings considered assume a fixed dataset. What if we can actively improve the data's utility by acquiring features? e.g., a doctor runs extra tests to improve a regressor. Understanding the value of features one might acquire is useful, especially if acquisition is costly. We note feature acquisition differs from feature selection (e.g. [45]), where all features are present and we select useful ones. It also differs from active learning, which quantifies sample acquisition value, not features.
We now show that TRIAGE enables a principled approach to assess the benefit of acquiring a specific feature. We posit that acquiring a valuable feature wrt. the task should increase the proportion of well-estimated samples (which we've shown is linked to better performance).
To illustrate the potential of TRIAGE scores, we set up an experiment where we rank the features based on their correlation to the target (i.e. increasing utility). We then simulate the sequential "acquisition" of features with increasing utility (based on correlation). Concurrently, we observe how the groups change as features are acquired.
Fig. 6 shows that as we acquire "valuable" features, the proportion of the well-estimated samples increases, with the over- and under-estimated proportions dropping. This result, repeated on other datasets in Appendix C.12, shows
Figure 6: The value of acquiring a feature can be quantified by its ability to increase the proportion of well-estimated samples.
TRIAGE scores are sensitive to the value of a feature, quantified by the increase in the proportion of well-estimated samples.
Takeaway 5. TRIAGE offers a principled way to assess active dataset improvement when we acquire new features. The feature's value links to the ability to increase the well-estimated proportion.
We present TRIAGE, a regression-focused data characterization framework, an understudied area in data-centric AI. TRIAGE offers a new scoring method to overcome the limitations of classification-centric scores. We show TRIAGE's characterization can unlock a variety of use cases to enhance regression, with minimal effort, beyond standard training on a dataset and exchangeability assumptions. We highlight that a key aspect of ensuring good performance of TRIAGE is careful construction of the calibration set — which should be as representative as possible. Rather than automating or replacing a data scientist's role, TRIAGE as a "data-centric AI" tool aims to empower data scientists to perform the "data" work in a more principled manner.
Limitations and future opportunities: (1) understanding the attributes contributing to characterization, (2) provide theoretical guarantees taking into account the challenging interaction between the CPD and Triage trajectory (for more see Appendix C.9), (3) in high-stakes settings, to prevent harm we could leverage domain expertise, where a human-in-the-loop could audit the samples surfaced by TRIAGE, before sculpting.
The authors are grateful to Fergus Imrie, Boris van Breugel, Paulius Rauba and the anonymous NeurIPS reviewers for their useful comments & feedback. The authors would also like to thank Richard Samworth, Marco Scutari, Qingyuan Zhao and Yao Zhang for their insightful discussion on robust statistics linked to TRIAGE. Nabeel Seedat is supported by the Cystic Fibrosis Trust, Jonathan Crabbe by Aviva, Zhaozhi Qian by Cancer Research UK.
[54] Michael Feldman, Sorelle A Friedler, John Moeller, Carlos Scheidegger, and Suresh Venkata-subramanian. Certifying and removing disparate impact. In proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pages 259-268, 2015.
[55] Equal Employment Opportunity Commission et al. Uniform guidelines on employee selection procedures. Fed Register, 1:216-243, 1990.
[56] Hans Hersbach. Decomposition of the continuous ranked probability score for ensemble prediction systems. Weather and Forecasting, 15(5):559-570, 2000.
[57] Jacob Steinhardt. Robust learning: Information theory and algorithms. Stanford University, 2018.
[58] Ricardo A Maronna, R Douglas Martin, Victor J Yohai, and Matías Salibián-Barrera. Robust statistics: theory and methods (with R). John Wiley & Sons, 2019.
[59] Mengjie Chen, Chao Gao, and Zhao Ren. A general decision theory for huber's contamination model. 2016.