Improving Pretraining Data Using Perplexity Correlations

This paper proposes a novel approach to data selection for pretraining large language models (LLMs) using perplexity correlations. This work is based on a simple observation: LLM losses on many pretraining texts are correlated with downstream benchmark performance, and selecting high-correlation documents is an effective pertaining data selection method. This approach is compute efficient and has competitive performance.

This paper proposes a new framework for pre-training data selection by leveraging large amounts of open-weight language models. Concretely, a data domain is selected for pre-training when open-weight LLMs’ log probabilities on this domain highly correlate with performance on a target downstream task. The higher correlation a data domain demonstrates, the more weight the domain occupies the pre-training data. The proposed data selection algorithm is evaluated with small models of 160M parameters trained on a data scale of 3.2B tokens. Experimental results show that the proposed data selection algorithm, without relying on hand-crafted heuristics and expensive model training runs, approximately matches the handcrafted fastText baseline.

This paper proposes a new LLM pre-training data selection framework where no LLM training is needed. Their framework leverages publicly accessible pre-trained LLMs and is based on the correlation between downstream benchmark performance and LLMs' losses. In the paper, the authors provide theoretical background for the proposed framework and present experiments where they use 160M parameters models with a 3.2B token budget, comparing DSIR and DataComp-LM data selection. Concretely, they selected data from RedPajamas V2, and used LAMBADA, ARC-E, PIQA, and SciQ as target downstream benchmarks. In their experiments, the proposed method outperforms DSIR and shows competitive performance with DataComp-LM best data selection method.

The authors build on a large number of existing open LMs and the observation that perplexity of an LM on a sample of text is predictive of performance on some (presumably related) benchmarks to produce a method for selecting pretraining data, aiming to optmize downstream performance. The method doesn’t require any training of models. Specifically they greedily sample all available data, without repetition, from each domain in order of correlation until a desired token budget is met. They also filter at the webpage level using a FastText classifier trained to discriminate the selected and non-selected domains from the previous step. They apply this set up over:

• 90 models from 33M to 9B parameters from the Open LLM leaderboard, including both intermediate and final checkpoints.
• 4 downstream benchmarks appropriate for small models with an additional 4 evaluations created from one of the benchmarks translated into 4 different languages
• Perplexity evaluations on 9841 internet domains (e.g. wikipedia.org) represented by a sample of 25 web pages each, derived from Common Crawl.
• Sampling training data from a 100B token corpus, derived from Common Crawl
• Pretraining experiments at the 160M parameter 3.2B token (“chinchilla optimal”) comparing their proposed data recipe to baselines including uniform sampling, language filtering, an n-gram overlap method (DSIR; Xie et al., 2023), and a SOTA handcrafted FastText classifier (Li et al., 2024)

With this set up they find that:

• Their proposed approach performs comparably to the state of the art handcrafted baseline (which was constructed with expensive pretraining experiments not required by this paper’s approach) and outperforms all other baselines.
• Using Spearman rank correlation as opposed to the authors’ proposed correlation measure performs comparably (figure 6, appendix)
• Ablating the webpage-level filtering using the FastText classifier trained on their selected domains, does lower performance (though the difference is smaller when using Spearman rank correlation).

Their method can also be used to predict the ranking of the 90 public models by downstream performance using only domain perplexities. They find:

• A baseline approach that predicts the downstream rank by a model’s rank w.r.t. to mean loss over domains already gets R^2 values from 75 to 92 (always within 5 points of the proposed approach)
• Testing over all 8 downstream benchmarks, the proposed approach gets a statistically significantly better R^2 with p=0.035, though no individual benchmark is significantly improved.

They also provide an extensive theoretical statistical framework for using a single-index model (SIM) to represent the relationship between domain perplexity and downstream performance. They explain the use of the empirical CDF rather than raw perplexity values to better handle outliers. They also provide an approach to constrain the correlations to be a valid distribution over domains and to respect the limited number of tokens available per domain. However their empirical results “suggest that the performance of [their] approach is not dependent on the precise form of the estimator that is coupled to [their] theory results, but holds broadly across perplexity-correlation relationships” (line 470).

This paper analyzes how to select high-quality pretraining data for large language models (LLMs) without training new LLMs. The authors find that LLM losses on certain pretraining texts correlate with downstream benchmark performance. The key idea is to select data domains where lower loss generally corresponds to higher downstream performance. To do this, the authors use a collection of pre-trained LLMs and measure the log-likelihoods for texts from various domains and the performance of these LLMs on downstream benchmarks. They then calculate the correlation between these two measures for each domain. Domains with high correlations are selected for pretraining.

We appreciate this thoughtful and detailed set of reviews. In this response to all reviewers, we address their shared concerns. Separately, we address specific reviewer concerns in individual messages below the reviews.

The scale of the pretraining experiments is small (all reviewers):

• We have duplicated some of our original pretraining experiments at the 410M scale instead of the 160M scale, and are currently running them. Due to the time constraints for our rebuttal, we are only running the perplexity-correlations models and the uniform-sampling baseline model. To make the task harder for perplexity-correlations (and in the interest of time), we are using the same 3.2B datasets computed by perplexity correlations for the 160M parameter models, and duplicating them to be 8.2B tokens, which is chinchilla optimal at the 410M scale. The uniform-sampling model selects 8.2B tokens from the pre-filtered pool, meaning that it sees many more unique tokens than the perplexity-correlations models. If perplexity-correlations didn’t work at this scale, we would expect the target models to perform worse (not even the same) as the uniform-sampling model. We already have results for the SciQ target, the Arc Easy target, the PIQA target, and the uniform baseline. In these cases so far, our target models score well over a point higher on the respective benchmarks than the uniform baseline, except for PIQA which gets about the same (0.2% difference). We will add the 410M suite of experiments to the next version of our paper.
• EDIT: all 410M models have finished for this experiment, and perplexity-correlations does better than uniform sampling at all of our 8 tasks by a decent margin (~one percent to a few percent), except for PIQA where it is about the same (0.2% difference).

Our experiments only target individual benchmarks, not an aggregate (reviewers 54hW, 2JKP, 9vdE):

• Although we would hope that increasing performance for one benchmark would increase performance on many others (https://x.com/gblazex/status/1746295870792847562), we see no theoretical issue with using perplexity correlations to target an aggregate of benchmarks. This aggregation would itself be another benchmark. We have rented an H100 node and are currently running pretraining experiments where we target the DataComp-LM “core” benchmark, which is an aggregate of 22 benchmarks. This is particularly challenging because many of these 22 benchmarks yield random performance at the 160M and even 410M scales. We are currently training 160M, 410M, and 1B models with chinchilla optimal token counts. Although we may not be able to provide final numbers before the rebuttal deadline, we will put DataComp-LM numbers in the next version of our paper, regardless of whether they are good or bad.
• By choosing DataComp-LM core, we also hope to address a concern raised by reviewer 54hW about the limited number of benchmarks we tested so far.

The theory did not explain why another estimator (Spearman) also does well (reviewers 2JKP, 9vdE):

• We have now proven that the Spearman estimator is also monotonic with respect to the optimal weights in expectation, so it fits into our theory. We will update the next version of the paper with this proof. We can see from the common formula for Spearman rank correlation (where there are no ties between points) that we can reduce the problem to finding $E[(\Phi(y_i) - \Phi_j(x_{i,j}))^2]$. We can just multiply this out and use linearity of expectation to get $E[\Phi(y_i)^2] - 2E[\Phi(y_i) \cdot \Phi_j(x_{i,j})] + E[\Phi_j(x_{i,j})^2]$. Because $x$ and $y$ are both gaussian under the high dimensional regression assumptions, we can compute most of the terms easily. The middle term $2E[\Phi(y_i) \cdot \Phi_j(x_{i,j})]$ takes a bit more work, but we can use the law of total expectation and prove an analogous lemma to lemma 1. Then we can continue to proceed in a way analogous to the original estimator proof in our paper. Eventually we get to an integral that we found in a table.

We do not provide experiments where we train our estimator on smaller models and extrapolate to larger models (reviewers 9vdE, TyHg):

• We agree that such an experiment would be ideal to run. It would be exciting to use our approach to gain information about which sampling distributions are useful for next-gen models which would likely be larger than current-gen models. We plan on another experiment in a future version of our paper where we use only models on the order of 100M parameters for our estimate and train a model on the order of 1B parameters. Due to compute constraints, we may have to choose only a few small scale benchmarks out of the 8 that we used in our paper. Alternatively, if any of the reviewers have suggestions for any benchmarks on which ~100M parameter models get significantly above random, we would be open to considering them too.

The paper proposed a novel framework for selecting pretraining data based on correlations between LLM's perplexity and downstream benchmark performance. The authors leverage publicly available LLMs, their perplexity on different domain, their performance on downstream tasks, to estimate these correlations (within in a statistical framework), thereby avoiding costly pretraining ablations. Experiments at the 160M parameter scale with a 3.2B token budget, shows that the framework is superior to existing baselines.

The proposed approach is theoretically grounded, computationally efficient, and inspiring. It achieves promising results in data selection, but the paper’s contribution is relatively limited by the small scale of experiments, especially considering 3.2B token budget is much smaller than real training of LLMs.

The paper is a well-designed academic work and also has the potential to become a practical framework for pretraining data selection. Based on the reviewers' scores, the positive trajectory of the rebuttal, the submission is recommended for acceptance.

Accept (Poster)