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

## Abstract

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 1.45 × 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., ; Korbak et al., ), 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., ; Raffel et al., ) or using simple classifiers to determine high-quality samples (Wenzek et al., ).
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. () perform importance resampling where importance scores are calculated based on feature similarity to a target text.
Tirumala et al. () prune datasets by deduplicating and diversifying data based on a pretrained language model’s embeddings of the text samples.
Xie et al. () re-weight domain proportions based on learnability as determined by a smaller proxy model.
Marion et al. () 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., ) 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., ), *hellaswag*(Zellers et al., ), 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^11[https://data.commoncrawl.org/](https://data.commoncrawl.org/))), we consider two datasets with different domain compositions: the Pile (Gao et al., ) and Dolma (Soldaini et al., ).
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\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

Figure: Algorithm 1 Psuedocode for performing 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}$), 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-\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., ; MosaicML, ).
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., ) is composed of 22 different domains that range from general web scrapes to legal text.
Dolma (Soldaini et al., ) is composed of 7 different domains and is derived mainly from general web scrapes.
We tokenize all datasets using the GPT-4 tokenizer (OpenAI, ).

##### 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., ), 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., ) 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, ) 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, ).
Before averaging the accuracy across tasks, we normalize each task by the baseline of random guessing
(^2^22Not to be confused, the random accuracy normalization we use is different from the normalized accuracy reported by the [EleutherAI LM Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness/tree/main), which normalizes based on the Byte-length of the response.).
Specifically, we normalize the accuracy of each individual task as $a_{n}=\frac{a_{m}-a_{r}}{1-a_{r}}$, where $a_{m}$ is the accuracy of the model and $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

**Table 1: Average normalized accuracy grouped by task category for both datasets and both final model sizes. For all datasets and model sizes we find that training on perplexity pruned data outperforms the baseline. Bold results are within one standard error of the highest score.**
| Pruning Method | World Knowledge | Common Sense Reasoning | Language Understanding | Symbolic Problem Solving | Reading Comprehension | Average |
| --- | --- | --- | --- | --- | --- | --- |
| 1B Parameters Trained on Pile |  |  |  |  |  |  |
| No Pruning (Baseline) | 15.51 | 10.31 | 28.11 | 3.53 | 11.16 | 13.73 |
| High Perplexity Selected | 18.18 | 12.75 | 33.2 | 3.36 | 10.63 | 15.62 |
| 3B Parameters Trained on Pile |  |  |  |  |  |  |
| No Pruning (Baseline) | 21.82 | 13.09 | 39.08 | 4.88 | 14.28 | 18.63 |
| High Perplexity Selected | 25.8 | 16.24 | 43.32 | 2.91 | 15.07 | 20.67 |
| 1B Parameters Trained on Dolma |  |  |  |  |  |  |
| No Pruning (Baseline) | 16.48 | 12.32 | 28.86 | 3.58 | 7.95 | 13.84 |
| Medium Perplexity Selected | 17.98 | 13.03 | 31.87 | 3.44 | 10.41 | 15.35 |
| 3B Parameters Trained on Dolma |  |  |  |  |  |  |
| No Pruning (Baseline) | 23.56 | 14.29 | 39.57 | 4.4 | 14.2 | 19.2 |
| Medium Perplexity Selected | 24.19 | 16.48 | 41.8 | 3.3 | 13.19 | 19.79 |

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

Figure: Figure 1: Average normalized task accuracy evaluated intermittently throughout pretraining for each dataset and model size investigated. Perplexity-based data pruning leads to an improvement in performance for all intermediate training steps evaluated.
Refer to caption: x1.png

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\times$ and $1.45\times$ fewer steps for Pile 1B and 3B respectively and in $1.29\times$ and $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

**Table 2: Downstream task performance for Chinchilla Optimal and $5\times$ over-trained data budgets. The “Improvement Over Baseline” column refers to the gain observed from perplexity pruning as compared to the baseline trained in the same setting.**
| Pruning Method | Average | Improvement Over Baseline |
| --- | --- | --- |
| 1B Parameters Trained on High Perplexity Pile |  |  |
| Chinchilla Optimal | $15.62$ | $1.89$ |
| $5\times$ Over-Trained | $18.83$ | $1.74$ |
| 1B Parameters Trained on Medium Perplexity Dolma |  |  |
| Chinchilla Optimal | $15.35$ | $1.51$ |
| $5\times$ Over-Trained | $18.67$ | $0.84$ |

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., ; Gadre et al., ).
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\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

Our experiments so far were conducted in the setting where there exists a sufficient abundance of data such that even after pruning with the desired selection rate there are enough data points to fill the desired token budget without requiring any data to be repeated.
However, there are many training settings that do not fall under this data-abundant regime.
Consequently, we evaluate how perplexity-based data pruning performs when the number of tokens is constrained, and pruning induces a greater number of repetitions of the data.
For each dataset, we vary the available data such that training for a Chinchilla optimal number of tokens requires a different number of repetitions.
Specifically, we investigate data budgets that require {0.5,1,2,4,8} repetitions to reach the Chinchilla optimal(^3^33Repeat=0.5 means that the available number of tokens is twice the training budget, i.e. the data-abundant setting).
As each number of repeats refers to the total number of tokens available, for all pruning experiments the number of repetitions after pruning is actually greater by a factor of $\frac{1}{r_{s}}$ since we prune the available tokens according to $r_{s}$, the selection rate.
Since all models use a selection rate of $0.5$, the models trained on the pruned data see the data for $2\times$ more repetitions.

Figure: Figure 2: Downstream task performance as a function of available dataset size. The number of repeats denotes the number of repeats over the raw dataset necessary to achieve the Chinchilla optimal number of tokens. Training on perplexity pruned data leads to an improvement for up to two repeats on both the Pile Dolma.
Refer to caption: x2.png

We plot the average downstream performance as a function of the number of repetitions in Figure 2.
On both the Pile and Dolma, we find that training on perplexity pruned data yields an improvement for up to two repetitions.
These results suggest that perplexity-based data pruning can still provide performance gains for some degree of data constraint.
Furthermore, our results replicate the findings of Muennighoff et al. () that more than four repetitions yields negligible gains.
Specifically, the baseline model without pruning maintains commensurate performance for up to four repetitions.
Similarly, models trained on perplexity-pruned data maintain commensurate performance for up to two repetitions through the base data, which corresponds to four repetitions after pruning.
That training on repeated perplexity-pruned data leads to diminishing gains after four repetitions post-pruning suggests that the higher quality data resulting from pruning does not change the point for which repeating data yields diminishing improvements in performance.

### 3.6 Upstream Perplexity is not a Reliable Evaluation Metric for Data Pruning

**Table 3: Performance as evaluated by perplexity on a test split of the original dataset as well as average normalized task accuracy for 1 billion parameter final models trained on the Pile. The model trained on pruned data has worse pretraining test split perplexity even though it significantly improves average downstream normalized accuracy.**
| Pruning Method | Test Set Pplx. ($\downarrow$) | Downstream Task Avg. ($\uparrow$) |
| --- | --- | --- |
| 1B Parameters Trained on Pile |  |  |
| No Pruning (Baseline) | 7.83 | $13.73$ |
| High Perplexity Selected | 8.51 | 15.62 |
| 1B Parameters Trained on Dolma |  |  |
| No Pruning (Baseline) | 13.53 | $13.84$ |
| Medium Perplexity Selected | 14.33 | 15.35 |

As previous works have used the perplexity of the model on a test split of the pretraining dataset as an approximation to downstream performance, we wanted to explore how well such perplexity-based evaluations agree with downstream performance for data intervention techniques.
Pruning performs an intervention on the dataset, making models trained on the pruned dataset biased estimators of the original data distribution.
Therefore, it is unlikely that the performance on the original data distribution is a fair evaluation of model quality.
We compare the test set perplexity and average downstream performance for 1 billion parameter models trained on the original and pruned version of the Pile and Dolma in Table 3.
For both the Pile and Dolma, training on perplexity pruned data significantly worsens perplexity on a test split of the pretraining data, while the average downstream performance is significantly improved.
This result suggests that test set perplexity may not always be a sound metric for data pruning work and that researchers should instead directly evaluate on downstream benchmarks.

## 4 Understanding the Effects of Perplexity-Based Pruning

In this section, we investigate how data pruning works by exploring some of the properties of perplexity-based pruning.

### 4.1 How Are Reference Perplexities Distributed

Figure: Figure 3: Distribution of sample perplexities as evaluated by the reference model for the Pile and Dolma. We show both the original distribution over the full dataset without pruning as well as the distribution after applying the optimal perplexity-based data pruning technique for a given dataset.
Refer to caption: x3.png

In order to better understand how perplexity-based data pruning works, we investigate the distribution of the computed reference model perplexities for each dataset.
For each dataset, we randomly sample 10% of the calculated perplexities and perform kernel density estimation to estimate the distribution of log perplexities for a given dataset.
We repeat this procedure for the optimal pruned version of the dataset.
We plot the resulting estimates of the log perplexity distribution in Figure 3.
We find that the log perplexity distribution for the Pile is multimodal and asymmetric, while for Dolma and it is unimodal and symmetric.

### 4.2 How Pruning Affects Domain Composition

Figure: Figure 4: Proportion of the total dataset each domain makes up before and after pruning. For all datasets, pruning tends to select more samples from general web domains while leaving out samples from highly specific domains.
Refer to caption: x4.png

We can also interpret the effect that perplexity-based data pruning has on a dataset by examining how pruning affects each domain’s proportion of the total dataset.
We plot the pre and post-pruning domain compositions for the Pile and Dolma in Figure 4.
Interestingly, for all datasets pruning increases the proportion of data coming from web-scraped domains while decreasing the proportion of data coming from highly specific technical domains such as code or scientific papers.
This trend is more pronounced in the Pile, where the proportions of Pile-CC and OpenWebText2 nearly double, while the proportions of domains such as Pubmed Central, ArXiv, and Github are all reduced by at least a factor of three.
Future work should investigate how perplexity-based pruning affects a model’s performance on downstream tasks that are in the same category as the highly pruned domains.

## 5 Related Work

##### Classical methods for pruning text data.

In order to improve the quality of raw web scrapes, which often contain very noisy samples, pruning via quality filtering has become a common practice.
Simple rules-based methods have been employed to prune datasets by filtering out low-quality samples according to some hand-crafted heuristic such as whether the text contains prohibited words, is predominantly English, etc. (Bane et al., ; Raffel et al., ; Rae et al., ; Penedo et al., ).
N-gram perplexity-based methods, in which an n-gram model is first trained on a high quality, curated corpus and then used to score another corpus, have also been applied to filter text data (Moore & Lewis, ; Axelrod, ; Gao, ; Laurençon et al., ; Muennighoff et al., ).
Although our method also uses perplexity to prune data, it does so in a very different manner.
In n-gram perplexity pruning, perplexity is used to estimate whether new text is in distribution as compared to the currated text the n-gram was trained on, while in our model-based perplexity pruning, the reference model is trained on the same distribution of text and the perplexity is more akin to an estimate of the difficulty of an example.
In this work, the datasets we leverage already have some basic rules-based pruning applied, and as such, the method we investigate is largely complementary to these existing techniques.

##### Neural network based methods for pruning text data.

Recently, there has been much interest in using neural networks to compute metrics that can be used to intelligently prune datasets.
A common technique in this family of methods is using a model to sample high-quality data from large datasets based on the sample’s similarity to a curated high-quality corpus that serves as a target distribution (Feng et al., ; Xie et al., ).
Xie et al. () also consider how to use a small reference model to prune pretraining data for a much larger model, by using a small reference model to learn the optimal weighting of domain proportions to maximize the "learnability" of the resulting dataset.
Pruning based on the difficulty or loss of a sample has previously been explored for text data, but the majority of such work focuses on curating data for finetuning (Swayamdipta et al., ; Attendu & Corbeil, ; Coleman et al., ; Mindermann et al., ; Mekala et al., ).
Marion et al. (), however, investigate multiple model-based sample difficulty heuristics for pruning pretraining text datasets.
Although we use the same method for pruning text pretraining datasets, our analysis differs substantially as we evaluate model quality based on downstream metrics and extend our analysis to multiple different dataset compositions which enables us to conclude that the reference model can be smaller than the final model.

##### Data pruning on vision tasks.

While data pruning is becoming more and more relevant with large amounts of text data, it has also been extensively applied in the vision domain (Paul et al., ; Toneva et al., ; Park et al., ). These works often prune data points based on their loss or gradients during training (Killamsetty et al., ; Mirzasoleiman et al., ). Model-based methods have also been leveraged for image data pruning (Fang et al., ; Schuhmann et al., ). Note that in the literature, data pruning is also sometimes referred to as coreset selection (Guo et al., ). More recently, Park et al. () show that, somewhat surprisingly, active learning (Castro & Nowak, ) based algorithms tend to outperform most data subset selection algorithms. In the context of contrastive learning, hard-negative mining has been effective as a data pruning method (Kalantidis et al., ; Robinson et al., ; Zhang & Stratos, ).
Recently, Goyal et al. () investigated scaling laws for training on pruned data in the context of vision models.

## 6 Conclusion

In this work, we conduct an empirical investigation of the impact that perplexity-based data pruning has on model performance.
We demonstrate that small reference models can be used to prune the data of models with up to *30$\times$* more parameters, leading to both significant downstream performance improvements and increased training efficiency.
We then investigate perplexity-based data pruning in two non-standard settings: the over-trained and data-constrained regimes.
We find that for both settings, training on perplexity pruned data can outperform training on unpruned data, demonstrating that perplexity-based data pruning is a widely applicable and extensible technique.
We also investigate upstream metrics for evaluating data pruning techniques and provide an example where evaluating models based on their perplexity on the test split of the pretraining dataset does not align with evaluating based on downstream model performance.
Additionally, we demonstrate that optimal pruning techniques can vary greatly for different dataset compositions.
Although we do not present a predictive theory for how pruning parameters should be selected for different datasets, we demonstrate that the optimal pruning parameters for a 1 billion parameter model can successfully transfer to 3 billion parameter models, potentially suggesting that empirically determining the optimal pruning parameters can be done cheaply.
Our work takes a key step towards establishing perplexity-based data pruning as a primary technique in the modern data researcher’s toolkit.

#### Acknowledgments

There are a few people who we would like to express our deepest gratitude for the assistance they provided.
Sean Owen helped us with his encyclopedic knowledge of PySpark.
Sam Havens and Daniel King both helped advise the early stages of this work.
Brett Larsen provided feedback on the presentation of our results.

## 7 Full Data Pruning Settings Sweep

In this section, we report the results of sweeping over different perplexity-based pruning setting configurations.
In particular, for each dataset, we first sweep over the selection criteria to determine where from the distribution of perplexities samples should be selected.
Then, using the best selection criteria, we sweep the selection rate to determine how much we should prune.

##### Setup.

We use the same training and evaluation setup as detailed in Section 3.1.
We only perform the sweep over pruning settings for 1 billion parameter final models for computational budget reasons; however, we find that the best selection criteria at the 1 billion parameter scale also confers a performance improvement at the 3 billion parameter scale, as detailed in 3.2.

### 7.1 Finding the Best Selection Criteria

**Table 4: Results from sweeping different selection criteria. We report the average normalized accuracy for each task grouping as well as across all tasks. While high perplexity selection is optimal for the Pile, medium perplexity selection is optimal for Dolma. Bold results are within one standard error of the highest normalized accuracy.**
| Pruning Method | World Knowledge | Common Sense Reasoning | Language Understanding | Symbolic Problem Solving | Reading Comprehension | Average |
| --- | --- | --- | --- | --- | --- | --- |
| 1B Parameters Trained on Pile |  |  |  |  |  |  |
| No Pruning (Baseline) | 15.51 | 10.31 | 28.11 | 3.53 | 11.16 | 13.73 |
| Low Perplexity Selected | 11.14 | 5.76 | 18.66 | 3.54 | 8.72 | 9.56 |
| Medium Perplexity Selected | 16.12 | 9.01 | 28.1 | 3.41 | 10.86 | 13.5 |
| High Perplexity Selected | 18.18 | 12.75 | 33.2 | 3.36 | 10.63 | 15.62 |
| 1B Parameters Trained on Dolma |  |  |  |  |  |  |
| No Pruning (Baseline) | 16.48 | 12.32 | 28.86 | 3.58 | 7.95 | 13.84 |
| Low Perplexity Selected | 16.13 | 10.1 | 27.28 | 3.45 | 7.85 | 12.96 |
| Medium Perplexity Selected | 17.98 | 13.03 | 31.87 | 3.44 | 10.41 | 15.35 |
| High Perplexity Selected | 16.65 | 13.12 | 31.14 | 3.15 | 8.55 | 14.52 |

For each dataset, we first sweep the selection criteria while keeping the selection rate fixed at 50%.
We report the performance of each selection criteria in Table 4.
We find that on the Pile high perplexity selection works the best and on Dolma medium perplexity selection works the best, improving the average downstream performance by 1.89 and 1.51 respectively.
An important observation from the sweep is that the best selection criteria from one dataset does not transfer to another dataset and may actually degrade performance compared to the baseline.
Although medium-perplexity selection is the best method on Dolma, selecting medium-perplexity samples on the Pile leads to a decrease in the average downstream performance of 0.23 as compared to not performing pruning.
These results inform us that high and medium perplexity selection are the optimal selection criteria for the Pile and Dolma respectively, and that the optimal selection criteria does not necessarily transfer between datasets with different domain compositions.

### 7.2 Finding the Best Selection Rate

**Table 5: Results from sweeping different selection rates. We report the average normalized accuracy for each task grouping as well as across all tasks. Bold results are within one standard error of the highest normalized accuracy.**
| Pruning Method | World Knowledge | Common Sense Reasoning | Language Understanding | Symbolic Problem Solving | Reading Comprehension | Average |
| --- | --- | --- | --- | --- | --- | --- |
| 1B Parameters Trained on Pile |  |  |  |  |  |  |
| 25% Selection Rate | 18.21 | 12.88 | 34.44 | 3.73 | 9.44 | 15.74 |
| 50% Selection Rate | 18.18 | 12.75 | 33.2 | 3.36 | 10.63 | 15.62 |
| 75% Selection Rate | 17.08 | 10.11 | 31.37 | 3.81 | 9.02 | 14.28 |
| 1B Parameters Trained on Dolma |  |  |  |  |  |  |
| 25% Selection Rate | 17.94 | 12.16 | 31.63 | 3.58 | 8.91 | 14.85 |
| 50% Selection Rate | 17.98 | 13.03 | 31.87 | 3.44 | 10.41 | 15.35 |
| 75% Selection Rate | 18.2 | 11.78 | 29.96 | 3.32 | 10.82 | 14.82 |

Using the optimal selection criteria that we found for each dataset, we next investigate the best selection rate for each dataset.
We investigate three different selection rates: 25%, 50%, and 75%.
We present the results for each selection rate in Table 5.
On the Pile, we find that there is no significant difference in downstream performance for selection rates of 25% and 50%; on Dolma we find that a selection rate of 50% achieves the best average downstream performance.
For simplicity, we chose to conduct the rest of the experiments in the paper using a selection rate of 50% on both datasets.
Furthermore, we find that all the selection rates tested outperform the baseline of no data pruning as measured by average downstream performance.
This suggests that the selection criteria has a greater impact on the performance of a pruning configuration than the selection rate.

## 8 Detailed Evaluation Setup

Jha et al. () also use the MosaicML evaluation gauntlet to perform evaluations in their work.
As such, with explicit permission from the authors, we exactly reproduce their text describing the tasks and tasks categories in the evaluation gauntlet. The following is from Section D of their paper:

The World Knowledge category includes the following datasets:

- •
Jeopardy (2,117 questions that are a custom subset of the dataset originally obtained from Wolfe et al. ())
- •
MMLU (14,042 four-choice multiple choice questions distributed across 57 categories Hendrycks et al. ()
- •
BIG-bench wikidata (20,321 questions regarding factual information pulled from wikipedia) Srivastava et al. ()
- •
ARC easy (2,376 easy multiple choice middle school science questions) Clark et al. ()
- •
ARC challenge (1,172 hard multiple choice science questions) Clark et al. ()
- •
BIG-bench: misconceptions (219 true or false questions regarding common misconceptions) Srivastava et al. ()

The Commonsense Reasoning category loosely assesses a model’s ability to do basic reasoning tasks that require commonsense knowledge of objects, their properties, and their behavior. It includes the following datasets:

- •
BIG-bench Strategy QA (2,289 very eclectic yes/no questions on a wide range of commonsense subjects e.g “Can fish get Tonsilitis?”)Srivastava et al. ()
- •
BIG-bench Strange Stories (174 short stories followed by questions about the characters)Srivastava et al. ()
- •
BIG-bench Novel Concepts (32 find-the-common-concept problems)Srivastava et al. ()
- •
COPA (100 cause/effect multiple choice questions) Roemmele et al. ()
- •
PIQA (1,838 commonsense physical intuition 2-choice questions) Bisk et al. ()
- •
OpenBook QA (500 questions that rely on basic physical and scientific intuition about common objects and entities) Mihaylov et al. ().

Language Understanding tasks evaluate the model’s ability to understand the structure and properties of languages, and include the following datasets:

- •
LAMBADA (6,153 passages take from books - we use the formatting adopted by OpenAI’s version)Paperno et al. ()
- •
HellaSwag (10,042 multiple choice scenarios in which the model is prompted with a scenario and choose the most likely conclusion to the scenario from four possible options)Zellers et al. ()
- •
Winograd Schema Challenge (273 scenarios in which the model must use semantics to correctly resolve the anaphora in a sentence. The Eval Gauntlet uses the partial evaluation technique technique introduced in Trinh & Le ()) Levesque et al. ()
- •
Winogrande (1,267 scenarios in which two possible beginnings of a sentence are presented along with a single ending) Sakaguchi et al. ()
- •
BIG-bench language identification (10,000 questions on multiple choice language identification) Srivastava et al. ()
- •
BIG-bench conceptual combinations (103 questions using made up words) Srivastava et al. ()
- •
BIG-bench conlang translation (164 example problems in which the model is given translations of simple sentences between English and some fake constructed language) Srivastava et al. ()

Symbolic problem solving tasks test the model’s ability to solve a diverse range of symbolic tasks including arithmetic, logical reasoning, algorithms, and algebra. These datasets include:

- •
BIG-bench elementary math QA (38,160 four-choice multiple choice arithmetic word problems) Srivastava et al. ()
- •
BIG-bench dyck languages (1000 complete-the-sequence questions) Srivastava et al. ()
- •
BIG-bench algorithms (1,320 questions) Srivastava et al. ()
- •
BIG-bench logical deduction (1500 four-choice multiple choice questions relating to relative ordering of objects) Srivastava et al. ()
- •
BIG-bench operators (210 questions involving mathematical operators) Srivastava et al. ()
- •
BIG-bench repeat copy logic (32 samples in which the model is required to follow some instructions for copying words/symbols)
- •
Simple arithmetic with spaces (1000 arithmetic problems consisting of up to 3 operations and using numbers of up to 3 digits, developed by MosaicML)
- •
Simple arithmetic without spaces (1000 arithmetic problems consisting of up to 3 operations and using numbers of up to 3 digits, developed by MosaicML)
- •
Math QA (2,983 four-choice multiple choice math word problems) Amini et al. ()
- •
LogiQA (651 four-logical word problems) Liu et al. ()

The Reading comprehension benchmarks test a model’s ability to answer questions based on the information in a passage of text. The datasets include:

- •
BIG-bench Understanding fables (189 short stories) Srivastava et al. ()
- •
Pubmed QA Labeled (1000 hand-labeled medical documents followed by a related question for which the model must respond yes/no/maybe) Jin et al. ()
- •
SQuAD (10,570 short documents followed by a related question. The model is expected to output the exact correct answer) Rajpurkar et al. ()
- •
BoolQ (3,270 short passages on a diverse range of subjects followed by a yes/no questions) Clark et al. ()

### 8.1 Evaluation Procedure

To compute model performance on the above datasets, the Eval Gauntlet uses one of the following three ICL metrics for each dataset (from MosaicML’s composer library).

- 1.
[InContextLearningQAAccuracy](https://docs.mosaicml.com/projects/composer/en/latest/api_reference/generated/composer.metrics.InContextLearningQAAccuracy.html): This metric uses the query, the corresponding correct answer and a list of alternative answers to measure a model’s prediction. If the model’s response conditioned on the query starts with either the correct answer or with one of the alternative answers, it is considered correct. This is used for question-answering tasks such as TriviaQA.
- 2.
[InContextLearningLMAccuracy](https://docs.mosaicml.com/projects/composer/en/latest/api_reference/generated/composer.metrics.InContextLearningLMAccuracy.html): This metric tests a model’s ability to output a precise set of tokens. A model’s output conditioned on a given query is judged to be correct only if the model’s highest probability tokens match the correct sequence of tokens. This is used for language modeling tasks such as LAMBADA.
- 3.
[InContextLearningMultipleChoiceAccuracy](https://docs.mosaicml.com/projects/composer/en/latest/api_reference/generated/composer.metrics.InContextLearningMultipleChoiceAccuracy.html): This metric is used for testing a model’s ability to answer multiple choice questions accurately. It compares the respective perplexity of the query prepended to each of the possible choices, according to the model. If the query-choice pair with the lowest per token perplexity is indeed the correct choice, then the model’s output is judged to be correct. This is used for multiple choice tasks such as HellaSwag, Winograd etc.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: How do you expect the results to scale on models larger than 3B parameters? 

How does models' performance change on domains which are pruned the most?
soundness: 3
strengths: The goal, and the process, and algorithm are defined and presented very clearly. Experiments cover multiple settings, with different model sizes and training algorithms. 
The proposed method is super useful for researchers who investigate practical techniques for data curation, with insightful empirical results. 
Experiments include two very different dataset distributions, the Pile dataset and Dolma. The work shows thorough experiments for various selection rates and perplexity criteria, presenting strong evidence about settings in which perplexity pruning does and does not work.
summary: The paper proposes that smaller language models effectively prune large datasets in a way that benefits the training of much larger model. Applying perplexity-based pruning techniques, they explore using a small model to filter high-quality subsets of data for training larger models. This approach is interesting because it’s a cost-effective alternative to using large models for pruning, and is applicable in real settings. The findings indicate benefits for downstream accuracy and training efficiency.

The paper demonstrates that a 125m parameter model can successfully prune data for large models and improve downstream task performance. The paper shows empirical results testing on The Pile and Dolma, two datasets with very different domain structures.
They also study the two settings of over-training and data-constrained setups and provide additional insights.
weaknesses: Authors claim that datasets pruning increases the proportion of general domain data from web-scraped domains, and decreases the proportion of specific and technical domains. But it is unclear and counter intuitive why training on general domain data improves performance of models on benchmarks. I think the paper lacks analysis to explain this observation.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - A quantized model may lead to better inference efficiency while calculating the perplexity. Was this considered while running the experiments?
- High perplexity selection will also inevitably lead to the inclusion of a significant portion of the noisier examples in the overall dataset. How can we determine the proportion of such examples in the final dataset and exclude them reliably?
- Minor typo (line 66): perplexity-basd -> perplexity-based
- It would be useful to include the following closely related data pruning works in the related work section:
    - https://arxiv.org/abs/2403.07384
    - https://arxiv.org/abs/2402.09668
soundness: 3
strengths: - The paper addresses an important problem of pruning the pre-training datasets to enable efficient training of LLMs.
- The experiments are thorough and cover different dimensions of perplexity-based pruning. 
- The paper is well-written and the results are presented clearly. 
-  The findings are significant, as they show that perplexity-based data filtering can not only reduce the size of the pre-training datasets, it also leads to better performance on certain downstream tasks.
summary: This paper presents a perplexity-based pruning method for reducing the size of pre-training datasets. The effect of pruning is evaluated through the performance on downstream tasks as well. Two datasets are used for evaluation: Pile and Dogma. The pruning efficacy is determined for over-trained and data-constrained regimes as well.
weaknesses: - The paper does not currently cover the computational complexity of the proposed pruning procedure. A few important questions that need to be considered in this regard:
    - How do the computational requirements for perplexity-based pruning increase with the size of the dataset to be pruned?
    - How does the cost of computing perplexity (before pruning) amortize over the efficiency improvements achieved while pretraining the model on the pruned datasets? 
- A discussion for choosing the right perplexity pruning method (low, medium, high) for the dataset should be included for the practitioners. From the experimental results, we can see that high perplexity selection performs better on Pile while medium perplexity selection is better for dolma. Can we extract any patterns from these results and other experiments that can be generalized to other datasets? 
    - For example, prior theory on data pruning for vision tasks shows that the optimal pruning strategy changes depending on the amount of initial data. When data is abundant, the better pruning strategy is to keep harder examples. In contrast, for smaller datasets, keeping the easier examples leads to better performance. [1] 
- The results show that test set perplexity may not always be a sound metric for evaluating a pruning strategy and that downstream evaluation is necessary. What should be the cheapest way of conducting the downstream evaluation of the correct perplexity pruning method, i.e., the one that can yield reliable results at a minimal cost? For example, could there be a small set of representative downstream tasks or metrics that could serve as efficient proxies for full downstream evaluation?

References:

[1] https://arxiv.org/abs/2206.14486

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - Fig.4 is interesting, but I'm not sure how Fig. 3 is relevant in practice - could you clarify?
soundness: 3
strengths: The method is well motivated. Except for some uncommon terminology that is explained in later sections like "non-standard training regime", "over-training" (which is not over-fitting) the paper is clearly written.
summary: The authors filter LLM pre-training data by using the perplexity of a smaller language model. They demonstrate that dataset filtering improves the [initial] learning curve of LLM pre-training.
weaknesses: L186 suggests that the final models are (pre-)trained for a fixed number of steps, no matter the dataset size. This sets the stage for dataset filtering, since training on the full dataset may go through fewer epochs. It would be interesting to train for long enough to show convergence in the plots in Fig. 1.  The story would be more convincing if there is an offset between the blue and red curves even after convergence. In fact, the "over-training" experiment in Sec. 3.4 shows diminishing gains, so I can imagine that they disappear fully at some point. The method would still have merits (steeper pre-training curve), just not the ones claimed in the paper.

Novelty. Perplexity-based pruning and countless variations of it are well-studied. The authors set their work apart from prior work in L058, but neither of the arguments (i)-(iii) (evaluation on downstream task, exploration of domain compositions, "non-standard" evaluation regimes) strike me as particularly strong.

I don't think that Algorithm 1 is really helping clarity. 1-2 normal equations would be just as expressive and more concise.

Edit: my point about novelty was unjustified - I have increased my scores after the rebuttal

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: * L290: "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." It would be good to understand why this is the case since there are no repeats. 
* L314: "That training on repeated perplexity-pruned data leads to diminishing gains after four repetitions post- pruning suggests that the higher quality data resulting from pruning does not change the point for which repeating data yields diminishing improvements in performance." This sentence is confusing and should be reworded.
* In section 4.2, the paper presents results showing that the pruning affects data composition such that some domains (e.g. web) are oversampled compared to others (e.g. pubmed). It would be useful to perform additional analysis to understand why this is the case e.g. is it possible that the training split (L113) resulted in a smaller proportion of these domains for the reference dataset?
soundness: 3
strengths: * Describes a simple approach to improve the performance of large language models using perplexity based data filtration using a smaller reference model.
* Presents useful results e.g. 1) filtration criteria varies by dataset type and 2) test set perplexity is not a good indicator of the downstream task performance.
summary: This paper investigates whether a small model can be used to perform perplexity based data selection for a larger model. The key findings are that 1) a reference model with 30x fewer parameters compared to the larger model can be used to identify a subset of the training data which can improve the performance of the larger model relative to no pruning. 2) the filtered data subset can speed up training of the larger model, 2) the improvements carry over to some extent to over training and data constrained regimes, 3) ideal pruning criteria can vary by dataset e.g. for Pile, a high perplexity subset performs better while for Dolma, a medium perplexity subset works the best. The paper shows that test data perplexity is not a good indicator of the downstream task performance when using perplexity based pruning.
weaknesses: * The main results (Table 1) do not include a random baseline i.e. what is the performance of a model trained on a subset of the data which has a similar size as the perplexity filtered buckets but is selected randomly?
* The paper does not contain ablations on the size of the reference model and sensitivity of the results to the random split (L113) used for training the reference model. Though exploring this space is computationally expensive, it may be useful to present 1-2 additional data points.
* It would be good to see some additional analysis to understand why a high perplexity set works better for one domain while a medium perplexity set works better for others.

Note: The authors have addressed some of these concerns (random baseline/sensitivity to random split) in  the rebuttal.

Meta-review: This paper presents a study on using smaller language models to prune large datasets effectively, aiming to improve the training efficiency of larger models. The proposed perplexity-based pruning technique is evaluated on two distinct datasets, The Pile and Dolma, which vary in domain structure. The study encompasses different experimental settings, namely over-training and data-constrained setups. The paper concludes that smaller models can successfully select high-value data subsets, which enhance training efficiency and downstream accuracy for larger models.

This paper tackles a significant challenge in natural language processing: the efficient training of LLMs through dataset pruning. By leveraging smaller models for perplexity-based data pruning, the proposed method is a cost-effective alternative to using large models. The extensive empirical results across varied datasets and experimental settings provide persuasive evidence of the effectiveness and practical applicability of this approach. The paper's insights into data composition effects and pruning strategies enrich the understanding of dataset management in model pre-training, making it a valuable contribution to the field.

Suggestions:
1. A deeper analysis is recommended to clarify why training on general domain data enhances model performance on benchmarks, despite domain-specific prunings which seem counterintuitive. Although the authors note this is outside their current scope, addressing this point would enhance the paper's comprehensiveness.
2. The paper should include a discussion of the computational requirements regarding perplexity pruning complexity. While responses address this, embedding it in the paper would guide practitioners better.
3. A clearer presentation of the random baseline setup is vital, as well as an explicit clarification in the experimental setup, to make these comparisons transparent. Some further ablation studies or sensitivity analysis related to the reference model size and pruned datasets would benefit the study, albeit acknowledged as resource-intensive.
4. Simplifying some parts of the paper, such as Algorithm 1, will assist in enhancing the accessibility of the methodologies delineated.

Decision: Accept (Poster)