[Paper Review] BERT on a Data Diet: Finding Important Examples by Gradient-Based Pruning

Mohsen Fayyaz, Ehsan Aghazadeh+4 more|arXiv (Cornell University)|Nov 10, 2022

Topic Modeling4 citations

TL;DR

This paper adapts gradient-based data pruning metrics GraNd and EL2N—originally developed for computer vision—to NLP, demonstrating that fine-tuning BERT on a subset of the most important examples (selected via these metrics after at least one epoch) can maintain or even surpass performance on the full dataset. The key finding is that pruning the highest-scoring (potentially noisy) examples improves generalization, especially on MNLI, where 66% of data yields higher accuracy than the full set.

ABSTRACT

Current pre-trained language models rely on large datasets for achieving state-of-the-art performance. However, past research has shown that not all examples in a dataset are equally important during training. In fact, it is sometimes possible to prune a considerable fraction of the training set while maintaining the test performance. Established on standard vision benchmarks, two gradient-based scoring metrics for finding important examples are GraNd and its estimated version, EL2N. In this work, we employ these two metrics for the first time in NLP. We demonstrate that these metrics need to be computed after at least one epoch of fine-tuning and they are not reliable in early steps. Furthermore, we show that by pruning a small portion of the examples with the highest GraNd/EL2N scores, we can not only preserve the test accuracy, but also surpass it. This paper details adjustments and implementation choices which enable GraNd and EL2N to be applied to NLP.

Motivation & Objective

• To adapt gradient-based data importance metrics GraNd and EL2N, originally used in computer vision, to the NLP domain for identifying important training examples.
• To investigate whether these metrics can identify high-value training examples in pre-trained language models (PLMs) like BERT.
• To evaluate whether pruning low-importance examples—especially those with high GraNd/EL2N scores—can improve or maintain model performance.
• To determine the optimal timing for computing these scores during fine-tuning, given the unique dynamics of PLMs.
• To assess whether removing high-scoring examples (possibly noisy) can lead to better generalization than training on the full dataset.

Proposed method

• Compute GraNd scores as the expected L2 norm of the loss gradient with respect to model weights, estimated over multiple random weight initializations.
• Approximate GraNd using EL2N, defined as the expected L2 norm of the prediction error (prediction - one-hot label), computed across multiple training runs.
• Use five independent training runs with different random seeds to compute stable, averaged EL2N and GraNd scores for each training example.
• Select and fine-tune BERT-base on the top k% of training examples ranked by EL2N or GraNd scores after one full fine-tuning epoch.
• Evaluate performance on AG News (topic classification) and MNLI (natural language inference) tasks using accuracy as the metric.
• Compare results against random pruning baselines and full-dataset fine-tuning to assess gains from structured pruning.

Experimental results

Research questions

• RQ1 Can GraNd and EL2N, originally designed for vision, be effectively adapted to NLP and pre-trained language models like BERT?
• RQ2 Is it necessary to perform at least one full fine-tuning epoch before computing reliable GraNd/EL2N scores in the context of PLMs?
• RQ3 Does pruning the highest-scoring examples (as measured by GraNd/EL2N) lead to better test performance than training on the full dataset?
• RQ4 How does the performance vary with the preserved fraction of data when using EL2N/GraNd-based selection?
• RQ5 Can removing a small portion of the highest-scoring examples (possibly noisy) improve model generalization in NLP tasks?

Key findings

• At least one full fine-tuning epoch is required before GraNd and EL2N scores become reliable for data subset selection in NLP, unlike in vision where early scores are informative.
• On the MNLI dataset, fine-tuning on 66% of the data (after removing the top 4% of highest-scoring examples) achieves higher accuracy than fine-tuning on the full dataset.
• On AG News, pruning the highest-scoring examples yields comparable performance to full training but with only 67% of the data, indicating data efficiency gains.
• EL2N and GraNd scores computed from a single random seed correlate highly (Spearman r = 0.9311) with scores from five seeds, suggesting that one epoch per seed is sufficient for reliable pruning.
• The performance of models trained on top-k% of examples selected by EL2N/GraNd improves monotonically with increasing preserved fraction, and exceeds random pruning when at least 70% of data is retained.
• In some cases, such as MNLI, removing the highest-scoring examples leads to better generalization than including them, suggesting that high-scoring examples may correspond to noisy or hard-to-learn instances.