[Paper Review] Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models

Zachary Ankner, Cody Blakeney+4 more|arXiv (Cornell University)|May 30, 2024

Algorithms and Data Compression4 citations

TL;DR

This paper proposes using small language models (125M parameters) to compute perplexity scores for pruning large-scale pretraining datasets, significantly improving downstream performance of larger models (3B parameters). It demonstrates that perplexity-based pruning with small reference models reduces pretraining steps by up to 1.45× and boosts average downstream accuracy by up to 2.04 points across diverse dataset compositions.

ABSTRACT

In this work, we investigate whether small language models can determine high-quality subsets of large-scale text datasets that improve the performance of larger language models. While existing work has shown that pruning based on the perplexity of a larger model can yield high-quality data, we investigate whether smaller models can be used for perplexity-based pruning and how pruning is affected by the domain composition of the data being pruned. We demonstrate that for multiple dataset compositions, perplexity-based pruning of pretraining data can \emph{significantly} improve downstream task performance: pruning based on perplexities computed with a 125 million parameter model improves the average performance on downstream tasks of a 3 billion parameter model by up to 2.04 and achieves up to a $1.45 imes$ reduction in pretraining steps to reach commensurate baseline performance. Furthermore, we demonstrate that such perplexity-based data pruning also yields downstream performance gains in the over-trained and data-constrained regimes.

Motivation & Objective

• To evaluate whether small language models can effectively identify high-quality data subsets for pretraining large language models.
• To investigate how dataset domain composition influences the effectiveness of perplexity-based data pruning.
• To assess the impact of perplexity-based pruning in non-standard training regimes, such as over-training and data-constrained settings.
• To determine whether upstream test set perplexity is a reliable proxy for downstream performance gains in data pruning interventions.
• To compare the efficacy of perplexity-based pruning across diverse datasets with varying domain distributions, such as the Pile and Dolma.

Proposed method

• Train a small reference language model (125M parameters) on a random subset of the pretraining dataset.
• Compute perplexity scores for each sample in the full dataset using the small reference model.
• Prune the dataset by selecting samples within a specific perplexity range—either the lowest or highest perplexity scores.
• Fine-tune a larger target model (3B parameters) on the pruned dataset and evaluate downstream performance.
• Use multiple downstream benchmarks (e.g., MMLU, HellaSwag, BIG-bench) to evaluate model performance beyond upstream metrics.
• Compare pruning outcomes across datasets with different domain compositions, including diverse curated domains (Pile) and web-scraped data (Dolma).

Experimental results

Research questions

• RQ1 Can a small reference model effectively identify high-quality data subsets that improve the performance of a much larger target language model?
• RQ2 How does the domain composition of a pretraining dataset affect the efficacy of perplexity-based data pruning?
• RQ3 Does perplexity-based data pruning yield performance gains in over-trained and data-constrained training regimes?
• RQ4 Is upstream test set perplexity a reliable indicator of downstream performance improvements from data pruning?
• RQ5 How do different pruning strategies (e.g., selecting low vs. high perplexity samples) vary in effectiveness across datasets with different domain distributions?

Key findings

• Perplexity-based pruning using a 125M-parameter reference model improved the average downstream performance of a 3B-parameter model by up to 2.04 points on benchmark evaluations.
• The same pruning strategy reduced the number of pretraining steps required to reach baseline performance by up to 1.45×, indicating significant training efficiency gains.
• Pruning effectiveness varied significantly across datasets: the method worked best on the Pile (diverse, curated domains) and was less effective on Dolma (web-scrape skewed), indicating strong sensitivity to domain composition.
• Perplexity-based pruning yielded measurable gains in both over-trained and data-constrained training regimes, demonstrating robustness beyond standard pretraining settings.
• Upstream test set perplexity was found to be a misleading metric for evaluating data pruning efficacy, as some interventions improved downstream performance without affecting upstream metrics.
• The study confirms that small reference models can serve as effective proxies for data quality filtering, even when the target model is 30× larger in parameter count.