TRIAGE: Characterizing and auditing training data for improved regression

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

Paper

Similar papers

Peer review

Reviewer N4d86/10 · confidence 3/52023-07-03

Summary

This paper present a data characterization method for regression task. The method leverages conformal predictive systems literature and proposed to estimate training data scores by thresholding percentile of calibration data points given their conformity measure. The method is interesting in terms of leveraging calibration dataset CPD to measure the training data score. The threshold of grouping (4.3) is less explained, but empirical results show it works. The paper is self-claimed to be the first data characterization method that is suitable for regression task.

Strengths

The proposed method is novel to me. Leveraging calebration dataset to re-score training data is interesting. I am quite curious how the choice of calibration dataset would impact the method's performance. The paper is very well motivated. And, maybe over motivated, considering justification and explanation missing in the section 4.2 and 4.3. The experimental results show the proposed method works well in benchmark datasets and could be a good tool to do data selection or feature selection.

Weaknesses

The paper takes pretty long paragraph to highlights its novelty and difference between other data characterization method, which looks like an ads. The motivations part should be reduced and focus more on algorithm explanation. There is very less justification on why thresholding CPD works for the data characterization purpose, and where the rules in Eq 3 come from. They need further justification and explanation. Algorithm 1 need to further adjustment and state what the "eval sample" is. I think it is one training data in D_train. The output of Algorithm should be a matrix |X| x |M| right? please state it somewhere to help understanding. Design of 4.2 is not explained. Why mean over training steps? What if we put more weights to higher iteration? What if there are multiple candidate regression methods? would this alter the outcome?

Questions

Why mean over training steps when computing C(x,y)? What if we put more weights to higher iteration? What if there are multiple candidate regression methods? would this alter the outcome? Why does the group assignment threshold designed as Eq 3? Any explanation or justification?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

No concern. But the authors should describe further in terms of mis using this method, where decision maker can remove data from underrepresented population based on the solution.

Reviewer pNo87/10 · confidence 2/52023-07-06

Summary

The task of data characterization aims to address variations in individual-level performance despite achieving good average performance. Existing methodologies have predominantly focused on classification, leaving a gap in data characterization approaches for regression. In this paper, the authors propose the TRIAGE framework to bridge this gap. Extensive experiments demonstrate the framework's superior performance in regression tasks.

Strengths

- The paper is well-written. - This is the first paper to introduce a principled data characterization framework in regression settings, supported by extensive experiments.

Weaknesses

Many contents are in the appendix. It would be helpful more discussions about the appendix in the main text.

Questions

- In simplified setting (e.g., linear (kernel) regression?), is it possible to derive further theoretical guarantees on the behavior of C and V, defined in Section 4.2? - Is the group assignment consistent? (in the sense that under what conditions are the samples accurately assigned their true group label?)

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

refer to weakness

Reviewer pEXt6/10 · confidence 3/52023-07-07

Summary

The authors introduce a new data characterization framework, TRIAGE, for regression models. The method utilizes conformal predictive distributions to compute the training examples' scores. To compute TRIAGE scores, the authors use predictive distributions and conformal prediction. A proper training set is used to train a regressor, and a separate calibration set is used for conformal calibration. Conformity measures the dataset's agreement with the observation. Consequently, the conformal predictive scores are computed at each epoch for each training point. Afterward, TRIAGE measures the mean and standard deviation for each training point. The method analyzes each example's training dynamics at each epoch to group each point into one of the groups: under/over/well estimated by the model based on the thresholds. TRIAGE can reduce the number of samples to train the regressor compared with the baseline methods. Specifically, it is observed that the MSE performance is improved with the selected data. It holds an advantage steadily over four different scoring methods.

Strengths

+ The authors propose a novel data characterization method for any regression models that analyze the training dynamic of each training point. Using conformal predictive systems for this problem is intuitive and effective. + TRIAGE effectively reduces the number of samples to train regressors compared with the baseline methods. It shows that the MSE performance is improved with data keeping and holds an advantage for four different scoring methods. + The paper is well written.

Weaknesses

+ Time complexity can be enormous for large neural networks to compute scores over multiple epochs. + Only simple datasets and simple regressors are performed. The data sculpting experiment is also limited to 500 available samples. + Some more challenging points might be important in the medical field, but if TRIAGE discards those samples, the model would not learn the critical medical case.

Questions

+ How much does the calibration dataset affect the data sculpting performance? + How exactly are distances computed in KNN for residuals of the calibration dataset? Does it use some embedding space? + What exactly do authors mean by model-agnostic? It seems that conformity scores are computed on residuals which are based on trained regression model predictions.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer oVcJ6/10 · confidence 3/52023-08-01

Summary

The problem studied in this work is the following: Given a dataset $\lbrace (x_i, y_i ) \rbrace_{i=1}^M$ and a regressor $f_\theta$ trained on this dataset, assign a group label $g_i$ to each sample that specifies whether the regressor under or overestimates on the sample. Such group labels can be used to identify and remove outliers in the dataset. The key formula of the proposed method is Eqn. (1), which is similar to Eqn. (5) in [23] that takes into account the "prediction accuracy" over each sample. The authors proposed to first put each sample into one of the several bins defined on a calibration set, and then assign a CPD score accordingly. The group label is then assigned using the CPD score as well as its variance. In the experiments, the authors use these group labels to remove outliers in the training set, and observe that it can improve the regression performance for a variety of regressors.

Strengths

- The problem studied in this work is very important for real applications since most datasets contain outliers. - The desiderata for data characterization, as listed in (P1)-(P4) after line 44, are very clear and define the overall goal of this work. - The experimental results presented in Section 5 support that the proposed method satisfies the desiderata.

Weaknesses

- The overall presentation is not good enough and there are many confusing points in this work, which I will discuss in the Questions section. - The authors claim TRIAGE to be "the first data characterization framework tailored to regression settings" (lines 63-64), yet in [23] cited in this work, a similar approach for regression conformal prediction was proposed. In fact, the proposed method in this work largely resembles the method in [23], including the CPD score definition Eqn. (1) (versus Eqn. (5) in [23]), the use of calibration scores $\lbrace \alpha_1,\cdots,\alpha_q \rbrace$ (versus Eqn. (1) in [23]), and the use of KNN for estimating $\sigma$ (versus Eqn. (11) in [23]). Thus, I am not sure about the contributions of this work. However, I am sure that "TRIAGE is the first data characterization framework tailored to regression" is an overclaim. - I feel that the writing is too wordy, especially in Section 5, and the authors use some LaTeX tricks which make the manuscript looks more compact than necessary. I think the authors are able to make Section 5 more concise, so that the main text could easily fit in 9 pages and the paper would have a much better shape.

Questions

1. There are several confusing points in the manuscript that require clarification: (a) In Algorithm 1, line 3, what do you mean by "nearest neighbor residuals of KNN"? Is it the average distance to the k nearest neighbors? (b) In Eqn. (11) of [23], the definition of $\sigma$ considers two factors: (i) Whether a sample is close to its k nearest neighbors; (ii) Whether the k nearest neighbors have consistent labels. Does this work use the same definition for $\sigma$? If not, what are the differences? (c) In the definition of the TRIAGE score $T(x,y,\theta)$ in lines 161-162, what do you mean by "$P(y \le f_\theta(x))$"? What is this probability taken over? If $x, y, f_\theta$ are all deterministic, why should there be a probability? 2. Can the authors make a thorough comparison between this work and [23], including the methodology, the algorithm, the experimental setup and results? This could help clarify the contributions of this work. 3. (Bonus) It would be great to make a theoretical connection between the proposed method and robust statistics. Consider a simple case, where $x$ is 1-dimensional, and the goal is to fit $y=f(x)$ on a dataset $(X,Y)$ with outliers. The dataset could follow Huber's contamination model. In such a scenario, I am curious about whether TRIAGE is provably better than the error baseline (as empirically compared in lines 244-260), and whether there is any guarantee on the performance of the predictor with the outliers removed using TRIAGE. A good starter for robust statistics could be the thesis of Jacob Steinhardt: https://cs.stanford.edu/~jsteinhardt/publications/thesis/paper.pdf. This is a bonus point, but such theoretical analysis could greatly enhance the contributions of this work. **Post rebuttal note:** I have read the rebuttal and had a discussion with the authors. The authors have addressed most of my questions. I have raised my rating from 4 (original) to 6 (current).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

2 fair

Limitations

Limitations are discussed in Section 6. **Review summary:** I think this work is interesting overall, and there are some good takeaways. However, the writing is confusing at times, and it also seems to me that a large part of this work is very similar to [23]. I believe many people will find this work useful, but I also believe that this work still needs improvement. So I give a borderline rating to this submission though I rarely give borderline ratings in my reviews. If the authors could address my questions during the rebuttal, I am willing to raise my score to 6 or 7, depending on the quality of the rebuttal.

Reviewer ckq87/10 · confidence 3/52023-08-06

Summary

This paper investigates the problem of training data characterization for regression problems. The authors noted that existing research on training data characterization mostly focuses on classification problems and there remains an absence of research for regression problems. This work proposes TRIAGE, a novel framework designed exclusively for regression settings and applies to a variety of tasks. TRIAGE leverages conformal predictive distributions to provide a model-agnostic scoring method for evaluating how the model performs on each training sample. This framework also enables analysis into the training dynamics from checkpoints, visualizing the samples being under-/well-/over- estimated by the model during the training process. The proposed framework is useful for a variety of data curation tasks, such as improving the model performance by throwing poorly-predicted training samples. Beyond sample-wise data analysis, the framework also applies to broader tasks such as dataset selection/feature acquisition. The work validates the proposed method with a range of empirical studies on 10 tableau datasets with 500-100k samples, showcasing its value in data characterization in practical cases.

Strengths

The problem being studied and data-centric perspectives are attractive and important. The work is well-motivated and nicely written with easy-to-follow illustrations. The narrative of the paper is highly structured, skillful, intriguing and informative. I did enjoy reading this paper. The importance of the work is highlighted. References are substantial and high quality. Objectives P1-P4 and the comparison table with baselines are clear and precise which can be informative and beneficial for a broader audience. And the improvements in these objectives are later validated in empirical studies, facilitating the verification of this work. The technical body is clear. The analysis of training dynamics from checkpoints is interesting. Novel insights are provided in the takeaways. The work is practical with many interesting use cases.

Weaknesses

I am combining this part with the questions. 1. Since this work is situated in the context of data-centric AI and also studies data characterization problems, it is a bit unnatural to miss out on the comparisons with data valuation methods. For example, Shapley-based methods are a common benchmark for the general interpretability of ML applications and also apply to characterizing the effects of training data on regression tasks. Also, I was wondering how model-agnostic data valuation ([1]) pipelines work in this setting. I would like to see discussions both conceptually and empirically on how those methods perform in the context this paper studies and how they differentiate from the proposed approach. [1] Lava: Data valuation without pre-specified learning algorithms, ICLR 2023 2. What is the rationality for throwing out data that is over/under-predicted? It could be outliers, but would simply throwing out such data hurt the generalizability of the model or its robustness against distributional shifts, which is prevalent in real-world applications? Do the authors have any empirical study results on such cases? 3. Data selection is an established field of research. This work showcases the proposed approach apply to dataset selection tasks. It would be beneficial if the authors could provide a direct comparison with data selection baselines and visualize their performance. It won't harm the contribution of this paper if its data selection performance is suboptimal compared to methods designed exclusively for data selection tasks. But I think it is important to know how they compare exactly and benchmark the gap, which would be helpful for future works to compare with or improve over. 4. What is the computational overhead of the proposed framework? Can the authors provide the computation time? How scalable is the proposed framework? What is the largest scale the authors are able to apply it to? The text in Figure 1 can be made larger.

Questions

See Weaknessnes.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

No major limitations. Questions for discussion are listed above.

Reviewer oVcJ2023-08-10

Thank you for this rebuttal. Here is a list of changes I suggest the authors make to the paper: - (A): Add this comparison to the paper. - (B): Change line 3 of Algorithm 1 to "the average distance to k-nearest-neighbors". - (C): Improve the writing of Section 5. - (D): Add this discussion to the paper. I have one follow-up question: In rebuttal (B), (iii), I still cannot see on which distribution the probability is taken. Let me pose this question in a clearer way. In line 176 of the submission, you wrote: $C(x^i, y^i) = \frac{1}{E} \sum_{e=1}^E T(x^i, y^i, \theta _ e) $, which is equal to $\frac{1}{E} \sum_{e=1}^E P(y^i \le f _{\theta _e}(x^i))$. Could you tell me on which distribution is the probability $P(y^i \le f _{\theta _e}(x^i))$ taken over?

Authorsrebuttal2023-08-10

Paper updates & clarification

Dear ``R-oVcJ`` Thank you for your feedback on the rebuttal. ---- ### (1) Incorporating changes into the updated paper We will definitely include your suggested changes which have come from our discussions on (A)-(D). These points will be integrated into the revised paper in the sections identified within the **UPDATE** blocks of our response. Thank you for your help in improving the paper! ----- ### (2) Clarification 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 [22] for more details. *Appendix A.3* outlines the conditions necessary for $Q$ to be related to a CDF. Since the CPD has the properties of a CDF, we use the CPD to estimate probabilities that the true target $y$ is less than or equal to a specified threshold/value. Thus, when you ask about the distribution over which the probability is calculated, it's the CPD that provides the probability estimation. To be precise, 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 then do this for all $f_{\theta_{e}}$ checkpoints where $e \in E$ to get the trajectory of TRIAGE scores for sample $x$. We hope this response clarifies and we will incorporate this more detailed explanation into the revised paper. ---- We are grateful for the reviewer's time and suggestions, which have strengthened the paper. We hope these changes address the reviewer's concerns. If you have any other comments or concerns, please let us know. We would be happy to do our utmost to address them. Paper 11604 Authors [22] Vladimir Vovk, Ivan Petej, Ilia Nouretdinov, Valery Manokhin, and Alexander Gammerman. Computationally efficient versions of conformal predictive distributions. Neurocomputing, 2020

Reviewer oVcJ2023-08-13

If I understand clearly, in $T(x^i, y^i, \theta)$, $y^i$ is a random variable, and you assume that its distribution follows the CPD, and that is how you estimate $P(y^i \le f _ {\theta _ e} (x^i))$. I suggest the authors make this point very clear. I don't have any further questions. I will raise my rating to 6.

Authorsrebuttal2023-08-14

Dear ``R-oVcJ`` We are glad our response clarified matters and would like to thank you for raising your score! In the revised paper, we will incorporate the more detailed explanation to ensure greater clarity. Thanks again for your time and suggestions. Regards Paper 11604 Authors

Reviewer pEXt2023-08-13

I appreciate authors for their great effort for responses. I have two small questions: + Why does Figure 16 shows the calibration size only up to 0.5? + I believe the neural network investigated by the authors is the Bayesian NN. What is the architecture of that network and how was it trained?

Authorsrebuttal2023-08-14

Clarifications

Dear ``R-pEXt`` Thank you for your feedback — we are glad our response has helped to address your comments. We clarify your two questions below. (1) Figure 16 calibration size: To clarify, this experiment corresponds to the calibration sample sizes of Table 2 in the main paper, where $D_{cal}$ sizes range from 10-500 samples. For instance, as mentioned on L929: 0.3 (300 samples). In Table 2, we show that at 500 samples (corresponding to 0.5 in Fig 16), that training on $D_{cal}$ directly is better than sculpting $D_{train}$. i.e. TRIAGE sculpting was beneficial at small sample sizes of $D_{cal}$. Thus, we capped the x-axis of Fig 16 to match Table 2. To prevent confusion, we will adjust Figure 16's x-axis to display the raw calibration size, aligning it better with Table 2. (2) Neural network clarification: To clarify, the BNN is only used for the baseline “BNN sculpt” --- comparing prediction-based sculpting with uncertainty to TRIAGE sculpting. The BNN is a MLP regressor where we learn a distribution over the weights using variational inference, in the same way as [R1] --- allowing integration of uncertainty. In contrast, when using TRIAGE with an MLP, we *do not* learn a distribution over the weights. Instead, TRIAGE wraps a conventional MLP regressor that uses point estimate weights. The MLP architecture however is the same between the BNN and the TRIAGE MLP. We are grateful for the reviewer's time and suggestions, which have strengthened the paper. We hope these clarifications address your questions. Paper 11604 Authors [R1] Ghosh, Soumya, Jiayu Yao, and Finale Doshi-Velez. “Structured variational learning of Bayesian neural networks with horseshoe priors.” International Conference on Machine Learning. PMLR, 2018.

Reviewer pEXt2023-08-21

Thank you for your clarification. I will keep my positive score for the paper!

Reviewer N4d82023-08-15

Thank you

Thank you for the clarification. I would keep my current review with higher confidence.

Authorsrebuttal2023-08-16

Dear ``R-N4d8`` Thank you for your feedback and suggestions that helped us strengthen the paper! Regards Paper 11604 Authors

Reviewer ckq82023-08-18

Thanks for the rebuttal

I appreciate the authors for their dedicated work and thanks for the response to my comments. My questions have been adequately discussed and I have no further comments at this moment. I would keep my positive rating in support of this work. I hope the authors compile the new results and additional discussions into the paper or its Appendix. Nice work and good luck, Reviewer ckq8

Authorsrebuttal2023-08-20

Dear ``R-ckq8`` We are glad our response addressed your comments. In the revised paper, we will definitely include these new discussions and results. Thanks for your positive feedback and suggestions which have helped us improve the paper! Regards Paper 11604 Authors

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC