[Paper Review] Fix your Models by Fixing your Datasets

Atindriyo Sanyal, Vikram Chatterji+4 more|arXiv (Cornell University)|Dec 15, 2021

Big Data and Business Intelligence4 citations

TL;DR

This paper proposes a data-centric framework to improve machine learning model performance by systematically identifying noisy labels and selecting the most informative samples for labeling. Using confidence-certainty and certainty-based methods, the approach reduces labeling costs and boosts accuracy—achieving up to 99.7% F1 on toxicity datasets and 92.7% on enterprise data after pruning mislabeled samples.

ABSTRACT

The quality of underlying training data is very crucial for building performant machine learning models with wider generalizabilty. However, current machine learning (ML) tools lack streamlined processes for improving the data quality. So, getting data quality insights and iteratively pruning the errors to obtain a dataset which is most representative of downstream use cases is still an ad-hoc manual process. Our work addresses this data tooling gap, required to build improved ML workflows purely through data-centric techniques. More specifically, we introduce a systematic framework for (1) finding noisy or mislabelled samples in the dataset and, (2) identifying the most informative samples, which when included in training would provide maximal model performance lift. We demonstrate the efficacy of our framework on public as well as private enterprise datasets of two Fortune 500 companies, and are confident this work will form the basis for ML teams to perform more intelligent data discovery and pruning.

Motivation & Objective

• Address the lack of systematic data quality tooling in ML workflows, which leads to inefficient, manual data cleaning and labeling.
• Reduce model training costs and improve generalization by identifying and removing mislabeled data instead of blindly increasing dataset size.
• Optimize limited labeling budgets by selecting the most informative samples that maximize model performance lift.
• Enable scalable, automated data observability in production ML systems through data-centric techniques.
• Demonstrate practical applicability across public and private enterprise datasets to validate real-world impact.

Proposed method

• Use a data-dependent confidence-certainty metric combining model confidence (μ) and prediction margin (δ) to score each sample’s likelihood of mislabeling.
• Apply the 90th percentile of lowest confidence-certainty scores as candidates for noisy labels, based on penultimate training epoch outputs.
• Implement a data-independent approach using Confident Learning to estimate class-conditional joint distributions (Q) and identify mislabeled samples via high-confidence scores.
• Apply core-set selection to unlabelled data to identify diverse, high-impact samples that improve model performance with minimal labeling budget.
• Use model embeddings from a pre-trained baseline to compute core-set diversity, ensuring coverage of decision boundary and rare clusters.
• Integrate both methods into a pipeline that enables iterative data pruning and selective labeling, minimizing human effort while maximizing performance gains.

Experimental results

Research questions

• RQ1 How can noisy or mislabeled samples be systematically detected in machine learning datasets without modifying the model architecture?
• RQ2 What is the relative effectiveness of data-dependent vs. data-independent methods for identifying label errors in real-world datasets?
• RQ3 How can a limited labeling budget be optimally allocated to maximize model performance improvement?
• RQ4 Can core-set selection outperform certainty-based sampling in capturing diverse, informative samples for model training?
• RQ5 To what extent can data-centric techniques reduce labeling costs while improving model accuracy and robustness in production settings?

Key findings

• The confidence-certainty method identified 2,032 mislabeled samples in the Newsgroup20 dataset, reducing error-induced performance degradation and improving F1 to 98.9%.
• The certainty-based method achieved 99.7% F1 on the Toxicity dataset, outperforming baseline models and demonstrating strong generalization on noisy real-world data.
• Core-set sampling achieved 0.52 F1 on the Newsgroup20 dataset with only 300 selected samples, outperforming both random and certainty-based sampling in capturing diverse data regions.
• On private enterprise datasets, the framework reduced mislabeling by 382 and 117 samples respectively, improving accuracy to 92.7% and 86.3% after pruning.
• The core-set method captured previously unseen clusters in the upper-right region of the embedding space, demonstrating superior diversity compared to certainty-based sampling.
• The framework reduced model training time and annotation cost by enabling smaller, higher-quality datasets without sacrificing performance.