Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models — part 1/2

Cleaned-text code points [0, 15228). Bibliography omitted here; source.md includes it.

# Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models

###### 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 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 × 1.45\times 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.

## 1 Introduction

A large focus of the machine learning community has been improving the performance of large language models (LLMs) while reducing their training costs. In this work, we consider how to improve the quality of an LLM by improving the quality of its pretraining data. Although there are many techniques to improve data quality, such as augmenting training samples with additional information ( Li et al., 2024 ; Korbak et al., 2023 ) , in this work we focus on the predominant method of data pruning : intelligently selecting a high-quality subset of a larger dataset to train on.

Data pruning is commonly used for quality filtering of noisy text data. Simple approaches include using symbolic rules ( Bane et al., 2022 ; Raffel et al., 2020 ) or using simple classifiers to determine high-quality samples ( Wenzek et al., 2020 ) . However, in addition to basic quality filtering, more complex data pruning techniques are also applied to datasets to further improve their quality. Xie et al. (2023b) perform importance resampling where importance scores are calculated based on feature similarity to a target text. Tirumala et al. (2023) prune datasets by deduplicating and diversifying data based on a pretrained language model’s embeddings of the text samples. Xie et al. (2023a) re-weight domain proportions based on learnability as determined by a smaller proxy model. Marion et al. (2023) investigate data pruning based on multiple neural heuristics of sample difficulty, ultimately concluding that the perplexity of a sample under a reference language model is the best pruning metric.

In this work, we thoroughly investigate the impact that data pruning based on sample perplexity ( Marion et al., 2023 ) has on LLM pretraining. In particular, we focus on the interplay between pretraining dataset composition and pruning methodology. We further evaluate perplexity pruning in the over-trained and data-constrained regimes. We also investigate whether evaluating the quality of data interventions based on upstream test set perplexity is a sound methodology for gauging downstream performance. To perform perplexity-based data pruning, we train a small language model on a random subset of the given pretraining dataset and then evaluate its perplexity on each sample in the dataset. We then prune the dataset to only include samples within some range of perplexities (i.e., sub-sample to the highest or lowest perplexity samples). We demonstrate that for two vastly different pretraining data compositions, a small language model can be used to effectively prune the pretraining dataset of a significantly larger model, leading to significant gains in the final model’s downstream performance.

Our work differs from previous work on perplexity-based data pruning for LLM pretraining in three key ways: (i) our emphasis on downstream model quality evaluation, (ii) our exploration of different pretraining dataset domain compositions, and (iii) our analysis of pruning in non-standard training regimes. While previous works evaluate the resulting LLM’s quality based on upstream metrics such as perplexity on the test split of the pretraining dataset, we evaluate data pruning’s impact based on downstream evaluation benchmarks (e.g. mmlu ( Hendrycks et al., 2021 ) , hellaswag ( Zellers et al., 2019 ) , etc.). Evaluating on more meaningful benchmarks enables us to make stronger, more rigorous conclusions about the impact of perplexity-based data pruning, as we find that some techniques that significantly improve downstream performance have no, or even adverse, effects on upstream performance. This difference in metrics enables us to conclude that smaller models can prune the data for larger models, which was not observed in previous perplexity-basd pruning works. Secondly, while previous work only investigates pruning on datasets composed of just one domain (CommonCrawl 1 1 1 https://data.commoncrawl.org/ ), we consider two datasets with different domain compositions: the Pile ( Gao et al., 2020 ) and Dolma ( Soldaini et al., 2024 ) . The Pile is composed of many diverse curated domains, with only 15.61% of the data being derived from general web-scrapes, while Dolma is a web-scrape skewed dataset, with 81.31% of its data being derived from the CommonCrawl. We find that successful pruning techniques vary greatly for different dataset compositions to the point that the best technique for one dataset composition may degrade performance for a different composition. Finally, we also evaluate perplexity-based data pruning in the less standard regimes of over-training and data-constrained training. This investigation provides a broader understanding for when practitioners should use perplexity pruning for their data.

##### Contributions

Our work makes the following contributions: • We demonstrate that, across three datasets of varying domain compositions, a small reference model can effectively prune the pretraining dataset of a significantly larger language model ( 30 × 30\times greater parameters), providing both a significant increase in downstream performance and decrease in pretraining steps ( Table 1 and Figure 1 ).

• We show that data pruning techniques can be highly sensitive to the domain composition of the dataset, suggesting the need to evaluate multiple distinct dataset compositions when conducting data pruning research ( Table 1 and Table 4 ).

• We investigate perplexity-based data pruning in multiple non-standard settings demonstrating that it can still lead to gains when over-training and when data-constrained ( Section 3.4 and Section 3.5 ).

• We find that test set perplexity can be a misleading metric for evaluating the efficacy of data pruning techniques, as interventions that result in significantly higher test set perplexity can still achieve better performance on downstream tasks ( Table 3 ).

## 2 Perplexity-Based Data Pruning

We start by training a reference model that will be used to calculate the perplexity of all samples in our dataset. First, we partition the original dataset into two splits: one for training the reference model and one for training the final model. After training the reference model on the standard next-token prediction objective, we compute the reference model’s perplexity on each of the samples in the final model’s training split. We then prune the final model’s dataset split to a fraction of its original size, referred to as the selection rate ( r s r_{s} ), by selecting samples according to a selection criteria which can be one of low, medium, or high. In low selection, samples with the lowest perplexity are selected. In medium selection, we select samples whose perplexity is close to the median perplexity, that is, samples with perplexity in the [ 50 − r s 2 , 50 + r s 2 ] [50-\frac{r_{s}}{2},50+\frac{r_{s}}{2}] percentiles of all perplexities. In high selection, samples with the highest perplexity are selected. After pruning our dataset, we train a final model using the standard next token prediction objective on the pruned version of the final model training split. We present a pseudocode for pruning based on perplexity in Algorithm 1 .

We consider the setting in which the reference model is significantly smaller than the final model. While this assumption is not strictly necessary, we believe that it is the most practically relevant setup, as it best reflects a data pruning paradigm that would be used for the next generation of LLMs where the models being trained are larger than any existing models.

## 3 Experiments

### 3.1 Setup

##### Models.

All models are based on the MPT family of transformer models ( Vaswani et al., 2017 ; MosaicML, 2023c ) . All reference models have 125 million parameters, and we consider final models with 1 billion and 3 billion parameters.

##### Data.

We consider two datasets in this work. The Pile ( Gao et al., 2020 ) is composed of 22 different domains that range from general web scrapes to legal text. Dolma ( Soldaini et al., 2024 ) is composed of 7 different domains and is derived mainly from general web scrapes. We tokenize all datasets using the GPT-4 tokenizer ( OpenAI, 2022 ) .

##### Training and hyperparameters.

All reference models are trained for a fixed duration of 26 billion tokens. Unless otherwise specified, all final models are trained to Chinchilla optimal ( Hoffmann et al., 2022 ) , meaning that each final model’s training duration in tokens is 20 times its parameter count. All models are trained using the decoupled Lion optimizer ( Chen et al., 2024 ) with a cosine learning rate schedule. All reference models and 1B parameter models are trained with a maximum learning rate and weight decay of 2e-4 and all 3B models are trained with a maximum learning rate and weight decay of 1.6e-4 . Training is conducted using llm-foundry ( MosaicML, 2023b ) and using both Nvidia A100s and H100s . We perform two trials for each experiment.

##### Evaluation.

We evaluate models on 33 different downstream question-answering tasks using the MosaicML evaluation gauntlet ( MosaicML, 2023a ) . Before averaging the accuracy across tasks, we normalize each task by the baseline of random guessing 2 2 2 Not to be confused, the random accuracy normalization we use is different from the normalized accuracy reported by the EleutherAI LM Evaluation Harness , which normalizes based on the Byte-length of the response. . Specifically, we normalize the accuracy of each individual task as a n = a m − a r 1 − a r a_{n}=\frac{a_{m}-a_{r}}{1-a_{r}} , where a m a_{m} is the accuracy of the model and a r a_{r} is the expected accuracy of random guessing. We report the average normalized accuracy for each task category as well as the average normalized accuracy across all task categories. More details on tasks and task categories are listed in Section 8 .

### 3.2 Perplexity-Based Data Pruning Improves Downstream Performance

If a certain range of perplexities is a good heuristic for data quality, training on that perplexity-pruned subset should improve downstream performance. We sweep across pruning selection criteria and selection rates ( Section 7 ) and find that the best settings are to select high-perplexity samples at a 50% rate for the Pile and to select medium-perplexity samples at a 50% rate for Dolma. We compare the most performant pruning settings to baseline models trained on the original datasets without pruning in Table 1 . Across all datasets and model sizes, models pretrained on the perplexity pruned version of the dataset significantly outperform the baseline model on average. Specifically, perplexity-based data pruning outperforms the average downstream performance of no pruning for 1B models by 1.89 and 1.51 for the Pile and Dolma respectively, and improves the performance of 3B models by 2.04 and 0.59 for the Pile and Dolma respectively. These results suggest that the perplexity of a small model provides a strong signal of data quality for a much larger model, as training on the data selected by the small model leads to significant downstream performance improvements.

### 3.3 Perplexity-Based Data Pruning Improves Training Efficiency

Since perplexity-based data pruning improves the final performance of models, we also investigate how pruned data affects the training dynamics of models. Specifically, we investigate whether training on perplexity pruned data enables models to achieve the same downstream performance as models trained on the unpruned data in training fewer steps. We plot the average downstream performance of partially trained checkpoints from the 1B baseline and perplexity pruned models in Figure 1 . Perplexity pruning outperforms the baseline model for all intermediate pretraining durations evaluated. Furthermore, perplexity pruned models reach the same average normalized accuracy as the baseline models in 1.31 × 1.31\times and 1.45 × 1.45\times fewer steps for Pile 1B and 3B respectively and in 1.29 × 1.29\times and 1.14 × 1.14\times fewer steps for Dolma 1B and Dolma 3B respectively. These results demonstrate that the resulting high-quality data from perplexity-based data pruning enables faster learning which can be leveraged to achieve the same downstream performance as training on unpruned data with fewer pretraining steps.

### 3.4 Perplexity-Based Data Pruning for Over-Trained Models

A recent trend with LLMs has been to over-train models by training them on more tokens than the Chinchilla optimal number of tokens ( Touvron et al., 2023 ; Gadre et al., 2024 ) . As our work targets the data component of LLM pretraining, we investigate the hypothesis that over-training would be more beneficial for models trained on perplexity pruned datasets as the data is of higher quality. We test this hypothesis by training a 1B parameter model for 130B tokens, which is 5 × 5\times the Chinchilla optimal number of tokens. We evaluate the downstream performance of each over-trained model in Table 2 . The main observation is that while the absolute gain in average downstream normalized accuracy from perplexity-based data pruning on the Pile is similar for both compute optimal and over-trained models, the gain decreases for Dolma when over-training. On the Pile, we find that the gain from perplexity pruned data is similar in the compute optimal regime and the over-trained regime: we see a gain in average performance of 1.89 when training compute optimal and a gain of 1.74 when over-training. On Dolma, the gain from perplexity pruned data decreases in the over-trained regime: we see a gain of 1.51 when training for a compute optimal duration but this decreases to a gain of 0.84 when over-training. These results show that while the higher quality data resulting from perplexity-based data pruning does still lead to an improvement in downstream performance in the over-trained regime, there is not a relative increase in downstream improvement over the baseline when over-training.

### 3.5 Perplexity-Based Data Pruning for the Data Constrained Regime