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

We would like to thank Reviewer 13P9 for the time they spent reviewing our paper. Our responses to the reviewer’s feedback are listed below.

“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 baseline results presented in Table 1 are implemented as described in the reviewer’s comment as our pruning results are constructed such that the number of tokens post-pruning equals the training duration. Namely, if our experiment has a selection rate of $r_s$ and we have a desired training duration of $D$ examples, then we compute the perplexity and sample from a pool of $\frac{1}{r_s}D$ examples from the original dataset. As such, post-pruning we are left with $D$ examples. The baseline results in Table 1 are constructed by just training on a random set of $D$ examples from the original dataset. As such, the baseline is a random selection that is the same size as the perplexity-pruned data. We understand that this is not stated clearly in the paper, and we will clarify these details in our experimental setup.

The results in Table 1 are conducted under the assumption that the dataset is sufficiently large such that we can prune without requiring multiple passes through the pruned data. We relax this assumption in Section 3.5 where we investigate the data-constrained regime, and the baseline results in Figure 2 are randomly sub-sampled such that they are of the same size as the perplexity pruned data.

“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.”

We agree that evaluating the impact of reference model size on final performance is an interesting direction for future work. Based on the findings of Marion et al. [1] we expect that increasing the reference model size would lead to increases in pruning performance, but it would indeed be valuable to reconfirm this for our experimental setup. We will add this result for the final paper.

The results we present in the paper do factor in the sensitivity to the random split used to train the reference model. Each experiment in the paper is conducted across two seeds, and the same random seed used for training the final model is also used to split the dataset and train the reference model. Thus the gain’s from perplexity pruning are robust to the random split used as each result presented in the paper is the performance across two different random splits of the training data. We will clarify this detail in our experimental setup for the final paper.

“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.”

Developing a better understanding for why different domain compositions admit different optimal pruning strategies is an interesting direction for future work. One possible explanation can be seen from the fact that the post-pruning domain compositions become more similar (Figure 4). Namely, for both the Pile and Dolma the web crawled domains (cc, openwebtext2, etc.) are upweighted while specialized domains (pubmed-central, the stack) are downweighted suggesting that in both cases it is optimal to have a larger portion of web-data. In order to prune to more web data on the Pile, high perplexity samples must be selected as there is a large proportion of specialized domains which typically have samples with lower perplexity. For Dolma on the otherhand, it is primarily composed of web-data, and as such the medium perplexity samples are retained. We will make this discussion more explicit in the final version of the paper.

We would like to thank Reviewer wYdV for their review and the questions that they asked. We also appreciate that the reviewer finds the method to be “well motivated”. Our responses to the reviewer’s feedback are listed below.

“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.”

While training to convergence would be an interesting data point, we don’t believe this is computational feasible with any reasonable budget. Consider the experiments in which we trained for 5x the Chinchilla training duration for the 1B models. We have plotted the performance at intermediate checkpoints throughout training as done in Figure 1 for the 1B parameter model trained 5x chinchilla on the Pile, and the resulting plot can be accessed here. As can be seen, even training for 5x the Chinchilla training duration does not saturate the performance.

We would like to make the further meta point that training to convergence is not the standard practice when training LLMs and the standard is training a model for the compute optimal duration [1]. Additionally, in the age of scaling when one benchmark saturates we evaluate on harder benchmarks instead, and as such the performance people care about is not in the convergened regime.

“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.”

We are only aware of one paper that examines model-based perplexity pruning for pretraining LLMs before ours [2] and as such we believe that it is a false characterization to say this setting is “well-studied”. With regard to this earlier paper, the differences we outline are very significant. Without evaluating based on downstream evaluations, one would conclude that the technique does not work unless the reference model is significantly larger than the final model. This conclusion would severely limit the applicability of perplexity-based data pruning as it would only be a useful technique for training smaller language models.

We also strongly believe that evaluating multiple domain compositions and non-standard training regimes have significant implications. By evaluating multiple domain compositions, our research is actually applicable for practitioners as they can choose the proper pruning setting based on their dataset composition. As stated, we find that the optimal settings for one domain composition may actually lead to worse performance than no pruning on another composition. Furthermore, by evaluating the over-trained and data-constrained settings, we provide the first guidance on when the technique should be expected to work in non-standard settings.

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

Thank you for this feedback. While we do believe that there are some important details conveyed in the algorithm that would be harder to communicate in text, we agree that it is more complex than necessary. We will simplify the algorithm in the final paper.

“Fig.4 is interesting, but I'm not sure how Fig. 3 is relevant in practice - could you clarify?”

The purpose of Figure 3 is to provide readers with intuition for both the differences in the distribution of text between the Pile and Dolma and to demonstrate what effect pruning has on the distributions. Namely, the Pile is composed of many domains and as such its distribution has multiple modes while Dolma is predominantly a single domain and correspondingly unimodal. Additionally, we make the point that the perplexity distribution of both datasets has a very similar shape post-pruning. This suggests a potential reason why different pruning strategies are superior on different domain compositions, as different pruning strategies are needed to achieve similar perplexity distributions post-pruning.

References

[1] Hoffmann, Jordan, et al. "Training compute-optimal large language models." arXiv preprint arXiv:2203.15556 (2022).

[2] ​​Marion, Max, et al. "When less is more: Investigating data pruning for pretraining llms at scale." arXiv preprint arXiv:2309.04564 (2023).

We would like to thank Reviewer wjeV for the time that they spent reviewing our paper. Our responses to the reviewer’s feedback are listed below.

“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?”

We will add a discussion of the compute requirements to the paper as we agree it is important. Assuming that the reference model is fixed (i.e., the same size reference model trained on the same number of tokens) as we do, then the only compute required is for performing inference from the reference model over the dataset to be pruned. So the pruning compute grows as $O(N_{ref} \frac{D}{r_s})$ where $N_{ref}$ is the number of reference parameters, $r_s$ is the selection rate, and $D$ is the number of training tokens. As $N_{ref} << N_{final}$, this operation is relatively cheap compared to the overall cost.

We now turn towards the total computational complexity of perplexity pruning. As our method does not change the computation performed on the forward or backward pass, and only affects which tokens are trained on, we can analyze the compute requirements in terms of the total number of operations. We will approximate the cost of training as $6ND$ where $N$ is the number of parameters and $D$ is the number of training tokens, and we will approximate the cost of computing a sequence’s perplexity in inference mode as $2ND$ [1][2]. Assuming a reference model of size $N_{ref}$, number of reference training tokens $D_{ref}$, final model size $N_{final}$, number of final tokens $D_{final}$, selection rate $r_s$, and fraction of tokens for the pruned data to achieve the same performance as the baseline $F$, the relative cost of perplexity pruning compared to the baseline is:

All our reference models are 125M parameters, trained for 26B tokens, and we use $r_s = 0.5$ throughout. The 1B models are trained on 26B tokens giving us $\frac{N_{ref}}{N_{final}}(\frac{1}{3r_s} + \frac{D_{ref}}{D_final}) = \frac{125 \times 10^6}{1.3 \times 10^9}(\frac{2}{3} + \frac{26 \times 10^9}{26 \times 10^9}) = 0.16$. For the Pile, $F=0.76$ giving us a relative cost of $0.16 + 0.76 = 0.92$ and for Dolma $F=0.78$ giving us a relative cost of $0.16 + 0.78 = 0.94$. The 3B models are trained on 54B tokens giving us $\frac{N_{ref}}{N_{final}}(\frac{1}{3r_s} + \frac{D_{ref}}{D_final}) = \frac{125 \times 10^6}{2.7 \times 10^9}(\frac{2}{3} + \frac{26 \times 10^9}{54 \times 10^9}) = 0.05$. For the Pile $F=0.69$ giving us a relative cost of $0.69 + 0.05 = 0.74$ and for Dolma $F=0.88$ giving us a relative cost of $0.88 + 0.05 = 0.93$. As can be seen, perplexity pruning leads to a total reduction in cost all experiments. We would also like to emphasize that the cost of reference model training is typically amortized as it's only performed once per dataset. I.e. we used the same reference models to prune the data for both the 1B and 3B models.

“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?”

We agree that such a discussion would be useful and will add it to the final paper. Our results are still useful for practitioners without needing to extrapolate any trends as the Pile and Dolma cover the two primary types of domain compositions for pretrainnig datasets. Namely, datasets are either composed of lots of specialized, skilled domains or predominantly general web scrapes. As we test both these settings, we believe practitioners can use the same settings we find depending on which of the two types of dataset compositions they have.

“A quantized model may lead to better inference efficiency while calculating the perplexity. Was this considered while running the experiments?”

We did not consider such a method for reducing the perplexity calculation costs at the time of writing the paper but we suspect that it would work and should be used in future 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?”

We agree that while selecting higher perplexity samples leads to an improvement in performance, it may bias the dataset to contain noisier samples. One potential method for removing such examples is combining our perplexity-based pruning method with other data filtering methods specifically targeted at noisy examples. One example would be training a simple classifier on curated examples of noisy text [4] and then using that classifier to further prune the perplexity-pruned dataset.

“Minor typo (line 66): perplexity-basd -> perplexity-based”

Thank you for catching this. It is now fixed.

“It would be useful to include the following closely related data pruning works in the related work section.”

Thank you for bringing these related works to our attention. We will include the updated related works in the final version of our paper.

References

[1] Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

[2] Anthony, Quentin, et al. "Transformer Math 101." EleutherAI, 18 Apr. 2023, blog.eleuther.ai/transformer-math/.

[3] Barton, Tessa. "Calibrating the Mosaic Evaluation Gauntlet." Databricks, 30 Apr. 2024, https://www.databricks.com/blog/mosaic-research/calibrating-mosaic-evaluation-gauntlet.

[4] Gao, Leo. "An empirical exploration in quality filtering of text data." arXiv preprint arXiv:2109.00698 (2021).

We would like to thank Reviewer qYWL for their detailed review of the paper and for finding the proposed method “super useful”. Our responses to the reviewer's questions are listed below.

“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.”

As the upweighting of web domains is observed across both the Pile and Dolma for two different pruning strategies, we are inclined to believe that the increased performance of web-domains is a more general phenomenon of llm pretraining. While we agree that developing a deeper understanding of why web domain data leads to better LLM performance, we believe it is outside of the scope of our research.

“How do you expect the results to scale on models larger than 3B parameters?”

While we believe that using a 125M parameter reference model to prune data for the Pile would continue to yield gains past 3B parameters, there is too much uncertainty to make any claims for Dolma. On the Pile, the gap between the average pruned performance and the average baseline performance actually grows from 1.89 to 2.04 for the 1B and 3B parameter final models respectively. However, on Dolma the gap between the average pruned performance and the average baseline performance shrinks from 1.51 to 0.59 for the 1B and 3B parameter final models respectively.

We would like to emphasize that this is not to say that we don’t believe the method would for models larger than 3B on Dolma. Even if using a 125M parameter reference model stopped showing gains past 3B final parameters, we believe that through scaling the size of the reference model proportionality to the final model size we would continue to see gains on Dolma.

“How does models' performance change on domains which are pruned the most?”

The model sizes and token budgets we evaluate aren’t sufficiently large to get signal on the domains that are primarily pruned (i.e., code tasks). The closest category from the evaluation gauntlet that we use is the “Symbolic Problem Solving” category. As seen in Table 1, while there is no statistically significant difference for the 1B models, the baseline models outperform the models trained on pruned data for the symbolic problem-solving category at the 3B model scale. We hope to fully investigate how pruning affects coding performance at scale in future work. If performance does degrade on coding tasks, we believe that there would be promising strategies to mitigate this such as domain upsampling at the end of training [1] or applying pruning on a per-domain basis.

References

[1] Blakeney, Cody, et al. "Does your data spark joy? Performance gains from domain upsampling at the end of training." arXiv preprint arXiv:2406.03476 (2024).

Based on your response, since you do not have the capacity to empirically test this for the paper, I think it is critical to mention the unknown effect of pruning on performance in the pruned domains in your paper.