Improving Pretraining Data Using Perplexity Correlations

## Abstract

Abstract Quality pretraining data is often seen as the key to high-performance language models. However, progress in understanding pretraining data has been slow due to the costly pretraining runs required for data selection experiments. We present a framework that avoids these costs and selects high-quality pretraining data without any LLM training of our own. Our 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 pretraining data selection method. We build a new statistical framework for data selection centered around estimates of perplexity-benchmark correlations and perform data selection using a sample of 90 LLMs taken from the Open LLM Leaderboard on texts from tens of thousands of web domains. In controlled pretraining experiments at the 160M parameter scale on 8 benchmarks, our approach outperforms DSIR on every benchmark, while matching the best data selector found in DataComp-LM, a hand-engineered bigram classifier.

## 1 Introduction

Dataset curation is increasingly crucial for training high-quality large language models (LLMs). As pretraining datasets have grown, from under 200B tokens in 2020 (Raffel et al., ; Gao et al., ) to 240T tokens today (Li et al., ), it has become critical to identify subsets of the available data that will lead to the best LLMs, and a wide range of methods have arisen to meet these needs (Ilyas et al., ; Xie et al., ; ; Engstrom et al., ; Everaert & Potts, ; Liu et al., ; Llama Team, ). However, data-driven approaches to data selection typically involve expensive model retraining steps that limit their effectiveness, and no algorithm has been reported to consistently beat or match hand-crafted classifiers for data selection (Li et al., ).

Is training new LLMs necessary for data selection? Instead of training our own models, can we use the growing collection of publicly available, high-performance LLMs (Wolf et al., ; Beeching et al., ) to perform data valuation and selection? This would have significant benefits: we could leverage the millions of dollars collectively spent on building these LLMs, and we would have coverage over a large, heterogeneous collection of high-performance models varying in size, architectures, and pretraining data distribution.

Despite these advantages, using existing models for pretraining data selection is challenging, as the training data for these models are often unknown and heterogeneous. Our key observation is that data selection can be done using two observable features of *all* public models today: 1) all open-weight models produce a causal language modeling loss for a given text, and 2) all of them can be evaluated on benchmarks. Prior work has found systematic relationships between web corpus loss and benchmark performance (Wei et al., ), which suggests the possibility of using correlations between perplexity and benchmark scores as the basis for a data selection policy.

In the present paper, we pursue this possibility and find a radically simple approach that is also effective: we select data via *perplexity correlations* (Figure 1), where we select data domains (e.g. wikipedia.org, stackoverflow.com, etc.) for which LLM log-probabilities are highly correlated with downstream benchmark performance. To enable our approach, we complement our algorithm with a statistical framework for correlation-based data selection and derive correlation estimators that perform well over our heterogeneous collection of LLMs.

We validate our approach over a large collection of pretrained causal LLMs on the Hugging Face Open LLM Leaderboard (Beeching et al., ) and find that perplexity correlations are often predictive of an LLM’s benchmark performance. More importantly, we find that these relationships are robust enough to enable reliable data selection that targets downstream benchmarks. In controlled pretraining experiments at the 160M parameter scale on eight benchmarks, our approach strongly outperforms DSIR (Xie et al., ) (a commonly used training-free data selection approach based on n-gram statistics) while generally matching the performance of the best overall method validated at scale by (the OH-2.5 +ELI5 fastText classifier (Joulin et al., )) without any parameter tuning or human curation.

Our manuscript focuses on pretraining experiments at the 160M scale. To assess the sensitivity of our experimental results to the choice of pretraining data pool, scale, and experiment design, we use this paper also as a preregistration for further pretraining experiments using different data sources and evaluation benchmarks at the 1.4B model scale. We commit to updating the arXiv version of our manuscript with both positive and negative results on these preregistered validation experiments.

Figure: Figure 1: We want to pretrain on domains where lower loss is generally correlated with higher downstream performance. Our approach does this by taking public, pretrained LLMs and measuring correlations across their log-likelihoods (left, red matrix) and performance on a target benchmark (center, blue vector). We then perform data selection by picking domains with high correlation and training a fastText classifier that distinguishes these domains from others. This approach is on par with the best-known data selection methods in our experiments, despite requiring no human selection of high-quality domains.

## 2 Related Work

To go beyond the status quo of deduplication, perplexity filtering, and hand-curation (Laurençon et al., ; BigScience, ; Abbas et al., ; Groeneveld et al., ; Soldaini et al., ; Penedo et al., ; Llama Team, ), targeted methods have been proposed to filter pretraining data so that the resulting LLM will achieve higher scores on given benchmarks. There are lightweight approaches that use n-gram overlap (Xie et al., ) or embedding similarity (Everaert & Potts, ) to select training data that is similar to data from a given benchmark. There are also less-scalable methods that require training proxy LLMs on different data mixtures (Ilyas et al., ; Xie et al., ; Engstrom et al., ; Liu et al., ; Llama Team, ).

Given the high costs of proxy-based data selection methods, they have primarily been used to select among human-curated pretraining data mixtures (Llama Team, ; Li et al., ) rather than a high dimensional space of mixtures. Our work takes an orthogonal approach and builds upon recent observational studies that have found scaling relationships that hold across collections of uncontrolled and diverse LLMs (Owen, ; Ruan et al., ). While these studies do not examine loss-to-performance relationships or derive useful data selection methods from them, we know that losses and performance are generally highly correlated. Validation losses on samples of text corpora are commonly used as a proxy for downstream performance when comparing LLMs pretrained on the same data distribution (Kaplan et al., ; Hoffmann et al., ; Wei et al., ), even if they have different architectures (Poli et al., ; Peng et al., ; Gu & Dao, ).

According to a recent survey of data selection approaches by Li et al. (), the heavier-weight pretraining data selection methods have not yet shown large gains, and the current state-of-the-art across many tasks is still primitive: a fixed fastText classifier (Joulin et al., ) combined with an English filter as a final layer after extensive deduplication and filtering. Are we missing important information that we can efficiently extract from a diverse collection of already trained models, larger and more diverse than any single organization is likely to produce? We show promising evidence supporting this hypothesis – simple loss-performance correlation coefficients are effective when used for data selection.

## 3 Problem Setting

Our goal is to build predictive models of how pretraining data distributions affect downstream benchmark performance and use them to build better language models. Unfortunately, this task is challenging and computationally expensive. A standard approach adopted in paradigms such as datamodeling (Ilyas et al., ) is to obtain $N$ different pretraining distributions $\{\mathbf{p}_{i}:i\in[N],\mathbf{p}_{i}\in{\mathbb{R}_{0}^{+}}^{D}\}$ over $D\gg N$ domains (e.g. arxiv.org, stackoverflow.com, etc.), pretrain and measure model errors on a target benchmark $y_{i}\in[0,1]$, and fit a model $\mathbf{p}\to y$. This approach requires $N$ LLM training runs, performed at a scale sufficient to obtain non-random performance on $y$. This can cost tens to hundreds of millions of dollars for hard benchmarks such as MMLU, where even the performance of 1B parameter LLMs often do not exceed random chance (Beeching et al., ).

Instead, our work considers the following *observational* setting that requires no training. We obtain $N$ pretrained, high-performance LLMs that vary in pretraining data, tokenizer, architecture, and scale (e.g. models on Huggingface’s OpenLLM leaderboard). Now, if we could train a predictor $\mathbf{p}\to y$ on these $N$ models, we could avoid large scale model training. Unfortunately, this is impossible as the training data for these models is often proprietary, and so we have no knowledge of $\mathbf{p}$.

The key observation of our work is that we can replace $p_{i,j}$ (the unobserved sampling probability of model $i$’s data selection policy on document $j$) with an observable surrogate $x_{i,j}$, which is the negative log-likelihood of document $j$ under model $i$.(^1^11To be precise, we use bits-per-byte, which normalizes the sequence negative log-likelihood with the number of UTF-8 bytes. This is defined in terms of the length of the string in tokens $L_{T}$, the length of the string in UTF-8 bytes $L_{B}$, and the cross entropy loss $\ell$ as $\text{BPB}=\frac{L_{T}\ell}{L_{B}\ln(2)}$) We can then build a regression model that relates negative log-likelihood $\mathbf{x}_{i}$ and benchmark error $y_{i}$. Using this model, we can select pretraining data from domains $j$ for which decreasing the loss $x_{i,j}$ is predicted to rapidly decrease error $y_{i}$.

The perplexity-performance hypothesis. We formulate the task of predicting errors $y_{i}$ from negative log-probabilities $\mathbf{x}_{i}$ as a single-index model (SIM),

$$ $y_{i}=f(\langle\bm{{\bf{\theta}^{*}}},\mathbf{x}_{i}\rangle+\epsilon_{i})$ (1) $$

where $f:\mathbb{R}\mapsto\mathbb{R}$ is some unknown monotonically increasing univariate function, $\epsilon_{i}$ is zero-mean noise which is independent of $\mathbf{x}$, and $\bm{{\bf{\theta}^{*}}}\in\mathbb{R}^{D}$ are unknown weights over $D$ domains.

A single index model is highly flexible (due to the arbitrary, monotone $f$) and has the advantage that we do not need to estimate the nonlinear function $f$ if our goal is to optimize model performance. We can see this directly from the monotonicity of $f$ as

$$ $\langle\bm{{\bf{\theta}^{*}}},\mathbf{x}_{i}\rangle+\epsilon_{i}<\langle\bm{{ \bf{\theta}^{*}}},\mathbf{x}_{j}\rangle+\epsilon_{j}\iff f(\langle\bm{{\bf{ \theta}^{*}}},\mathbf{x}_{i}\rangle+\epsilon_{i})<f(\langle\bm{{\bf{\theta}^{* }}},\mathbf{x}_{j}\rangle+\epsilon_{j}).$ (2) $$

Data selection from perplexity correlations. The weights $\bm{{\bf{\theta}^{*}}}$ tell us which domain perplexities are correlated with downstream performance. However, this isn’t sufficient for data selection. Even if we know how model likelihoods relate to model performance, we do not know how data selection affects likelihoods. Even worse, this data mixture to likelihood relationship *cannot* be learned observationally, as we do not know the data mixture of any of our models.

Despite this, we show that there is a clean approach for optimizing the data mixture. Our core observation is the following: *if we find a nonnegative $\bm{{\bf{\theta}^{*}}}$, sampling proportional to $\bm{{\bf{\theta}^{*}}}$ is always a good choice.* More formally, we see that this sampling distribution defines the pretraining loss such that optimizing the training loss directly optimizes the downstream task via the single index model.

###### Proposition 1

Suppose that $\bm{{\bf{\theta}^{*}}}$ weights are non-negative. Then, for models with associated likelihoods $\mathbf{x}\in\mathcal{X}\subset\mathbb{R}^{D}$, the minimizer of the pretraining loss over the $\bm{{\bf{\theta}^{*}}}$ sampling distribution $\mathbb{E}_{j\sim\bm{{\bf{\theta}^{*}}}}[x_{j}]$ also has the lowest expected downstream error according to the single index model:

$$ $\operatorname*{arg\,min}_{\mathbf{x}\in\mathcal{X}}\mathbb{E}_{j\sim\bm{{\bf{ \theta}^{*}}}}[x_{j}]=\operatorname*{arg\,min}_{\mathbf{x}\in\mathcal{X}} \mathbb{E}[f(\langle\bm{{\bf{\theta}^{*}}},\mathbf{x}\rangle+\epsilon)].$ $$

This observation follows directly from the fact that we can normalize any non-negative $\bm{{\bf{\theta}^{*}}}$ into a distribution (and shift the normalization constant into $f$) which allows us to write the inner product in the single-index model as a monotone function of the expected pretraining loss:

$$ $y=f(\langle\bm{{\bf{\theta}^{*}}},\mathbf{x}\rangle+\epsilon)=f(\mathbb{E}_{j \sim\bm{{\bf{\theta}^{*}}}}[x_{j}]+\epsilon).$ (3) $$

Proposition 1 allows us to entirely avoid the task of finding the optimal data mixture for a target likelihood. Instead, we pick sampling distributions that make the pretraining loss a monotone function of the predicted downstream error. Afterward, we can rely on our ability to optimize the loss to optimize downstream performance.

This view gives us a straightforward roadmap for data selection in the remainder of the paper: estimate a set of domains where loss and downstream benchmark performance is highly correlated, and then constrain our $\bm{{\bf{\theta}^{*}}}$ estimates to be a pretraining data sampling distribution.

###### Proposition 1

## 4 Methods

We now describe the details of our approach, starting by presenting the algorithm itself and the intuitions behind it, followed by a more precise and mathematical justification for the various steps.

### 4.1 Algorithm

##### Estimating 𝜽 ∗ superscript 𝜽 \bm{{\bf{\theta}^{*}}} bold_italic_θ start_POSTSUPERSCRIPT bold_∗ end_POSTSUPERSCRIPT .

The parameter $\theta^{*}_{j}$ measures the relationship between log-likelihoods in domain $j$ and downstream performance. Because of this, we might naturally expect $\theta^{*}_{j}$ to be related to nonlinear correlation coefficients between $x$ and $y$. Our work uses a simple correlation measure,

$$ $\gamma_{j}=\sum_{\begin{subarray}{c}1\leq k,l\leq n\\ k\neq l\end{subarray}}\operatorname{sign}(y_{k}-y_{l})(\operatorname{rank}_{j} (x_{k,j})-\operatorname{rank}_{j}(x_{l,j}))$ $$

where $\operatorname{rank}_{j}(x)$ is the rank of $x$ among $\{x_{1,j}\ldots x_{N,j}\}$. This formula is intuitive: when model $k$ does better than model $l$, what percentile is model $k$’s log-likelihood compared to model $l$’s? While this is not the only correlation coefficient that performs well (see Appendix E), this functional form has the additional benefit of being a principled estimate of $\bm{{\bf{\theta}^{*}}}$. In particular, we show in sections below that in expectation, the ranking of domains in $\bm{\gamma}$ exactly matches those of $\bm{{\bf{\theta}^{*}}}$ (under standard high-dimensional regression assumptions; see Section 4.2 for a complete discussion).

##### Selecting pretraining data.

Suppose that we have an accurate estimate $\gamma_{j}$ which is nonnegative. In this case, we could use $\gamma_{j}$ directly as a data selection procedure and 1 would ensure that minimizing the population pretraining loss minimizes downstream errors. Unfortunately, $\gamma_{j}$ can be negative and the finite number of tokens per domain can make it difficult to minimize the population pretraining loss. Thus, we must project $\gamma_{j}$ onto the set of reasonable pretraining data distributions that are nonnegative and account for the per-domain token counts.

What is a good way to project a set of domain rankings estimated via $\bm{\gamma}$ into a sampling distribution? Intuitively, if wikipedia.org has a $\gamma_{j}=0.5$ and arxiv.org is $\gamma_{k}=0.9$, it would be natural to select tokens in order of $\bm{\gamma}$, preferring tokens from arxiv.org over tokens from wikipedia.org when selecting pretraining data.

Having established the ordering of domains, the remaining question is how many tokens we take for each domain. We follow recent observations that repeating data degrades performance (Abbas et al., ) to arrive at a simple selection algorithm: select domains in greatest to least $\bm{\gamma}$, taking all the tokens in each domain once, until we exhaust our total pretraining token budget.

##### Full algorithm.

Together, these steps result in a simple, parameter-free algorithm that calculates our rank correlation coefficient, and selects domains in order from largest to smallest coefficient. We show this process explicitly in Algorithm 1, and additionally show an extra step where we train a fastText (Joulin et al., ) classifier (using standard settings and bigram features from Li et al. ()) which distinguishes our selected documents and domains from the rest of the pool. The fastText classifier allows us to perform data selection at a single-page level, and scale the selection process to larger datasets. We also found the classifier to slightly improve downstream performance over directly selecting the documents. More information on the specifics of the data selection approaches that we tested is given in Appendix D.

Figure: Algorithm 1 Perplexity Correlation Based Data Selection

### 4.2 Theory

We now study the approach closely and show that our choices for the correlation coefficient and projection step are
extensions of the classic, high-dimensional single index model estimator of Plan et al. (). We describe the basic single-index model estimators first, describe our extensions, and then conclude with a discussion on how our estimator and results deviate from the theory.

#### 4.2.1 High-dimensional estimation of single index models

For our theory, we consider the standard high-dimensional regression setting of Plan et al. () and Chen & Banerjee (). Here, our goal is to estimate the unknown weights $\bm{{\bf{\theta}^{*}}}$ in a single-index model $y_{i}=f(\langle\bm{{\bf{\theta}^{*}}},\mathbf{x}_{i}\rangle+\epsilon_{i})$, with $\textbf{x}_{i}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ for $\|\bm{{\bf{\theta}^{*}}}\|_{2}=1$ (assumed without loss of generality, as $\|\bm{{\bf{\theta}^{*}}}\|_{2}$ can be absorbed by $f$).

Our starting point is the classic result of Plan et al. (), who showed

$$ $\mathbb{E}\left[y_{k}\mathbf{x}_{k}\right]=c\bm{{\bf{\theta}^{*}}},$ (4) $$

for some positive constant $c$ and $1\leq k\leq N$. Closely related is the result of Chen & Banerjee () who showed a robust estimator quite similar to ours,

$$ $\mathbb{E}\left[\operatorname{sign}(y_{k}-y_{l})(\mathbf{x}_{k}-\mathbf{x}_{l} )\right]=\beta\bm{{\bf{\theta}^{*}}}$ (5) $$

for any $1\leq k,l\leq N$ (where $k\neq l$) and some positive constant $\beta$. Both of these results clearly identify that for the high-dimensional single-index model in the Gaussian setting, generalized correlation coefficients provide consistent estimates of the true regression coefficient $\bm{{\bf{\theta}^{*}}}$.

#### 4.2.2 Deriving our estimator

Both and provide moment-matching style estimators that consistently recover $\bm{{\bf{\theta}^{*}}}$ in high-dimensional, sparse settings. However, we found that both estimators directly use the values of $x$, and this resulted in brittle estimates due to outliers in language model log-likelihoods. While outlier removal is one possibility, we found that a simpler approach was to robustify the estimator of Chen & Banerjee () to outliers in $x$.

Recall that our estimate $\bm{\gamma}$ is a U-statistic, defined as pairwise sums of

$$ $\operatorname{sign}(y_{i}-y_{j})(\Phi(\mathbf{x}_{i})-\Phi(\mathbf{x}_{j})),$ (6) $$

for any $1\leq i,j\leq N$ (where $i\neq j$), where $\Phi$ is the empirical CDF of the $\mathbf{x}$ values. This estimate is significantly less sensitive to outliers than that of Chen & Banerjee (), as the empirical CDF is bounded between zero and one, and no single model can make the estimator degenerate.

We study this estimate theoretically in the Gaussian setting, where we consider the asymptotically equivalent estimator with $\Phi$ as the CDF of the standard Gaussian. In this case, we can show that this modified estimator is also consistent in recovering $\bm{{\bf{\theta}^{*}}}$.

###### Theorem 1

When $\epsilon\sim\mathcal{N}(0,\sigma^{2})$, we have:

$$ $\mathbb{E}[\operatorname{sign}(y_{i}-y_{j})(\Phi(\mathbf{x}_{i})-\Phi(\mathbf{ x}_{j}))]=\frac{2}{\pi}\sin^{-1}\left(\frac{\bm{{\bf{\theta}^{*}}}}{2\sqrt{1+ \sigma^{2}}}\right).$ (7) $$

We provide the proof in Appendix A. Because we assume $||\bm{{\bf{\theta}^{*}}}||_{2}=1$ and the expected value in Equation 7 must be between $-1$ and $1$, we are always within the domain of $\sin^{-1}$ and able to invert it. After inverting, we get:

$$ $\hat{\bm{\theta}}\propto\sin\left(\frac{\pi}{2}\mathbb{E}\left[\operatorname{ sign}(y_{i}-y_{j})(\Phi(\mathbf{x}_{i})-\Phi(\mathbf{x}_{j}))\right]\right)$ (8) $$

as an estimate for $\bm{{\bf{\theta}^{*}}}$, where the constant $2\sqrt{1+\sigma^{2}}$ term due to noise has been dropped.

Beyond the fact that our estimator is consistent, we can show an even tighter connection to the estimator and show that our estimates agree when running the original estimator on rank-transformed data. More specifically, for two models $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ with the estimated model rankings $\langle\hat{\bm{\theta}},\mathbf{x}_{i}\rangle>\langle\hat{\bm{\theta}},
\mathbf{x}_{j}\rangle$, the expected ranking under rank-transformation (i.e. $\Phi(\mathbf{x})$) match this ranking.

###### Corollary 1

Suppose that $\hat{\bm{\theta}}$ is any vector of fixed weights and $\mathbf{x}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. Then, conditioning on the event $\langle\hat{\bm{\theta}},\mathbf{x}_{i}\rangle<\langle\hat{\bm{\theta}},
\mathbf{x}_{j}\rangle$, we have

$$ $\langle\hat{\bm{\theta}},\mathbb{E}[\Phi(\mathbf{x}_{i})\mid\langle\hat{\bm{ \theta}},\mathbf{x}_{i}\rangle<\langle\hat{\bm{\theta}},\mathbf{x}_{j}\rangle] \rangle<\langle\hat{\bm{\theta}},\mathbb{E}[\Phi(\mathbf{x}_{j})\mid\langle \hat{\bm{\theta}},\mathbf{x}_{i}\rangle<\langle\hat{\bm{\theta}},\mathbf{x}_{j }\rangle]\rangle.$ (9) $$

with probability $1$.

This proof follows from the same calculations as Theorem 1 and is given in Appendix A.

###### Theorem 1

###### Corollary 1

#### 4.2.3 Selecting Data for Pretraining

Recall that our algorithm for data selection is to constrain $\bm{\gamma}$ to be a valid sampling distribution (nonnegative, at the very least) and then sample directly from this estimate. For now, we focus on constraining $\hat{\bm{\theta}}$, and we will see at the end of this section that we can apply the same constraint to $\bm{\gamma}$ directly to get the same result. The theory of constrained estimation for $\hat{\bm{\theta}}$ is simple and well-understood, with both Plan et al. () and Chen & Banerjee () extensively studying the problem of estimating $\hat{\bm{\theta}}$ under a known convex constraint set $C$. In particular, Plan et al. () show that performing a $L_{2}$ projection via $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in
C
}\|\bm{\theta}-\hat{\bm{\theta}}\|_{2}$ provides improved convergence rates that depend on the Gaussian mean width of $C$ rather than the ambient dimension, and Chen & Banerjee () show similar results when maximizing the linear correlation $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in
C
\subseteq B_{D}}-\langle\bm{\theta},\hat{\bm{\theta}}\rangle$.

We take a similar approach here. We define a convex constraint set $C$ that forces $\hat{\bm{\theta}}$ to be a reasonable sampling distribution and find the best sampling distribution via the linear correlation approach.

We define $C$ as the combination of two sets of constraints. First, we must have a valid sampling distribution, so we constrain $\hat{\bm{\theta}}$ to lie in the simplex. As we noted above, it is well-known that duplicating data harms performance (Abbas et al., ), and so we constrain $\hat{\bm{\theta}}$ to avoid data duplication by limiting the maximum weight on domains. Concretely, if want to pretrain on $m$ tokens overall, we enforce $\theta^{*}_{i}\leq\tau_{i},\forall i\in[1,D]$, where $\tau_{i}$ is set so $\tau_{i}m$ is the number of tokens from the $i$-th domain that we can access for training.

The resulting linear program has a simple solution and takes the form of initializing $\hat{\bm{\theta}}^{\textnormal{proj}}$ to $\mathbf{0}$ and then iterating through the values in $\hat{\bm{\theta}}$ from largest to smallest, setting the value at the corresponding index of $\hat{\bm{\theta}}^{\textnormal{proj}}$ to the maximum allowable value, until $\hat{\bm{\theta}}^{\textnormal{proj}}$ sums to $1$ (see Appendix B for a proof).

###### Theorem 2

Suppose we want to solve:

$$ $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in \mathbb{R}^{D}}-\langle\bm{\theta},\hat{\bm{\theta}}\rangle,$ $$

subject to:

$$ $\displaystyle\sum_{i=1}^{D}\theta_{i}=1$ $\displaystyle 0\leq\theta_{i}\leq\tau_{i},\forall i\in[1,D]$ , $$

where $\tau_{i}>0$ are fixed values. Then, the solution is:

$$ $\hat{\theta}^{\textnormal{proj}}_{k}=\begin{cases}\tau_{k}&\text{if }\sum_{j \text{: }r_{j}(\hat{\theta}_{j})\geq r_{k}(\hat{\theta}_{k})}\tau_{j}\leq 1\\ 1-\sum_{j\text{: }r_{j}(\hat{\theta}_{j})>r_{k}(\hat{\theta}_{k})}\tau_{j}& \text{if }\sum_{j\text{: }r_{j}(\hat{\theta}_{j})\geq r_{k}(\hat{\theta}_{k})} \tau_{j}\geq 1\land\sum_{j\text{: }r_{j}(\hat{\theta}_{j})>r_{k}(\hat{\theta}_ {k})}\tau_{j}\leq 1\\ 0&\text{otherwise}\end{cases},$ (10) $$

where $r$ is some function that breaks all ties between $\hat{\theta}_{j}$ and $\hat{\theta}_{k}$ for $k\neq j$, and otherwise leaves the ordinal relationships the same.

We note that while the use of this linear program is in line with the constrained estimators proposed in Chen & Banerjee (), the $L_{2}$ projection is arguably more natural, and does not require assuming that $\|\hat{\bm{\theta}}\|_{2}=1$ for asymptotic recovery conditions. We derive similar closed-form expressions for this quadratic case in Appendix B, but do not use this approach for two separate reasons.

First, the $L_{2}$ projection depends on the $L_{2}$ norm of $\hat{\bm{\theta}}$, unlike the linear program which only depends on the ranks of the values in $\hat{\bm{\theta}}$. The challenge with determining the norm is that the exact recovery result in Equation 7 requires knowledge of the noise level, and the trigonometric functions rely strongly on the Gaussian structure of $x$. Because of this, we are unlikely to be able to estimate the norm of $\hat{\bm{\theta}}$ with any accuracy, and the only way to avoid this would be to treat the norm as a hyperparameter, which adds unnecessary complexity. The second reason is empirical (although possibly a consequence of the first) – we found that the linear projection performed better across a wide range of benchmarks and conditions (See Appendix E).

We conclude by relating our theory to the full algorithm in Section 4.1. The estimation step for $\bm{\gamma}$ is the finite sample, U-estimate of the expectation in Equation 8, dropping the nonlinear transform $\sin$ and $\pi/2$ as these two terms do not change the rankings of the domains. The data selection step directly applies our projection in Equation 10, and we make use of the fact that this projection only relies on rankings among the domains to use $\bm{\gamma}$ rather than an exact estimate for $\bm{{\bf{\theta}^{*}}}$.

###### Theorem 2

### 4.3 Alternative Methods

Our estimator is far from the only reasonable high-dimensional, single-index model estimator. We briefly discuss some alternatives and the tradeoffs involved before moving to experimental results.

We could use classic low-dimensional methods regularized for the high-dimensional setting. This includes ordinal regression (Wooldridge, ) and the isotron algorithm (Kalai & Sastry, ). We found these methods to underperform correlation-based estimators, and tuning hyperparameters added additional complexity that was not needed in the correlation-based approaches.

Another class of methods involve scaling laws (Kaplan et al., ; Llama Team, ; Ruan et al., ). We could transform the $y$ values via an inverse sigmoid or power law, and fit high-dimensional linear regression methods (e.g. ridge, partial least squares, or Lasso). We initially found this approach promising, but the inverse transforms were unstable, and the combination of fitting the nonlinear transform and regularization required significant amounts of tuning.

Rank-correlation methods, including our robustified version of the estimator from Chen & Banerjee (), and even the standard Spearman correlation (Spearman, ) (see Appendix E) performed well. We believe that in general, robust per-feature correlations are likely to perform well as $D\gg N$, and extreme levels of regularization are needed to obtain reasonable models. Sparse methods such as the Lasso (Tibshirani, ) are one classic answer, but we cannot necessarily assume that the underlying correlations $\bm{{\bf{\theta}^{*}}}$ are sparse, and we did not find these techniques to perform well.

## 5 Results

We empirically validate our approach to predicting downstream performance and data selection. Our validation consists of three sets of experiments: we first pretrain 160M-parameter LLMs from scratch to study our primary goal of selecting pretraining data to improve downstream performance, followed by analyzing the ability of losses to predict downstream performance, and conclude with an analysis of the loss matrix $\mathbf{X}$. Throughout our experiments, we use the same single-index model that we train using Algorithm 1. As shown in the algorithm, we train the fastText classifier on selected vs unselected domains and use the classifier to filter the pretraining data at the page-level.

Input data matrix X. To build the input data matrix, $\mathbf{X}$, we collected byte normalized loss values from a sample of 90 Open LLM Leaderboard (Beeching et al., ) LLMs that we could run without errors. Concretely, these values are defined as bits-per-byte $\frac{L_{T}\ell}{L_{B}\ln(2)}$ where $L_{T}$ is the token count, $L_{B}$ is the number of UTF-8 bytes, and $\ell$ is the per-token cross-entropy (Gao et al., ). We collected these values on the smallest sample provided by the RedPajama V2 (RPJv2) dataset (Together Computer, ) for all domains with $\geq$ 25 pages in the sample. Specifically, we used the “sample” subset(^2^22[https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2)) of RPJv2 resulting in 9,841 domains/features. Specifics of how we computed losses are in Appendix C.

Target benchmark performance $y$.
We constructed a target vector, $\mathbf{y}$, for LAMBADA (Paperno et al., ), ARC Easy (Clark et al., ), PIQA (Bisk et al., ), and SciQ (Welbl et al., ). These are all of the tasks reported in the Pythia scaling experiments for which a model in the 160M parameter range could meaningfully perform above chance. We also constructed target vectors for LAMBADA${}_{\text{IT}}$, LAMBADA${}_{\text{FR}}$, LAMBADA${}_{\text{DE}}$, and LAMBADA${}_{\text{ES}}$, which are subsets of LAMBADA translated into Italian, French, German, and Spanish by Black (). These languages match those in RPJv2 where each page is conveniently tagged as one of five languages: English, Spanish, French, German, and Italian. The correspondence between our target benchmark languages and the RPJv2 metadata will be convenient for us, as it will allow us to include language filtering baselines at no additional cost.

### 5.1 Pretraining

We begin by validating our algorithm in the end-to-end task of pretraining data selection with controlled experiments at the 160M parameter, 3.2B token scale. The
low compute requirements of this setting allowed us to more extensively study replicates and ablations in Appendix E within the timeframe of a few days. We also note that while 160M models are small, this is far from an easy setting for our data selection algorithm. Most of the Open LLM Leaderboard models are 10 to 100$\times$ larger than the 160M scale, and our single index model must extrapolate substantially from $\approx$7B scale models to our small-scale validation setting (see Appendix G for a histogram of model sizes). Finally, the focus on small-scale validation in this manuscript allows us to commit to a clear preregistration-based strategy toward scaling up (Section 6) rather than cherry-picking experiments that succeed at scale.

Pretraining data and setting.
For pretraining, we used the “sample-100B” subset of RPJv2. This is larger than the sample that we used to compute our estimate. We filtered this data so it contains only the domains used for our estimate, and then tokenized the data with the Pythia tokenizer. The vast majority of the domains from our BPB matrix were present in this larger sample of text. However, 42 (out of 9,841) were not, and so we removed them from our estimate. For every data selection method that we tested, the task was to further select 3.2B tokens for pretraining, which is Chinchilla-optimal (Hoffmann et al., ) for the 160M-parameter LLM used in our tests.

Baselines. We compare against several baseline data-selection methods. First, we present the results of uniformly sampling from the available pretraining data. Then we use the language tags present in RPJv2 to filter only for the language matching the target task. In addition to these commonsense baselines, we also run DSIR (Xie et al., ): a lightweight training data selection technique based on n-gram overlaps that Li et al. () found to be competitive with proxy LLM-based techniques and was also validated at scale (Parmar et al., ). Finally, we run the state-of-the-art method for pretraining data quality filtering found by , which is a fastText classifier that beats all of the heavier-weight proxy-LLM methods that were tested. The classifier was trained on a benchmark-agnostic and handcrafted objective, which is to classify data as coming from Common Crawl(^3^33[https://commoncrawl.org](https://commoncrawl.org)) (low quality) or OH2.5 (Teknium, ) and Reddit ELI5 (Fan et al., ) (high quality). It is combined with an English filter in , and so we present results for this fastText filter with and without the additional English filter.

Model and hyperparameters. We use the Pythia 160M LLM configuration from Biderman et al. () and optimize the hyperparameters including learning rate, weight decay, and warmup to minimize loss on the uniform sampling (no selection algorithm) baseline. Training hyperparameters were fixed across all methods. We provide
additional training and evaluation details in Appendix D.

**Table 1: Average rankings of each data selection method (lower is better) across 8 benchmarks shows that correlation-based filtering beats baselines by a wide margin, and matches the current best open data filter from Li et al. (). Our approach significantly beats the default filter in Li et al. () with the EN filter and loses slightly after additional manual language filtering that depends on the target task ($+$ manual Lang Filter).**
| Method | None | Lang<br>Filt | DSIR<br>(Xie et al., ) | Handcrafted fastText<br>$+$ EN Lang Filter<br>(Li et al., ) | Handcrafted fastText<br>w/o Lang Filter | Handcrafted fastText<br>$+$ manual Lang Filter | Perplexity<br>Correlations |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Avg. Rank | 3.750 | 4.000 | 4.500 | 3.750 | 3.250 | 1.375 | 1.750 |

Figure: Figure 2: Pretraining results with different data selection methods. Each row is an LLM, and each column is a task. The number in the upper left indicates the ranking of the method when targeting that benchmark compared to other methods (lower is better). Numbers within the heatmap denote accuracy for all benchmarks except the LAMBADA tasks for which the values are log perplexities (where lower scores are better). We find that our approach appropriately optimizes data mixes for the target language and benchmark, and matches the fastText baseline across most benchmarks.
Refer to caption: extracted/5843524/images/llm-perf-heatmap-fasttext.png

Results. We report average rankings over all benchmarks in Table 1, and we find that our approach significantly outperforms the basic baselines of random sampling, language filtering, and DSIR. Compared to the existing state of the art from Li et al. (), our approach beats the performance of the default, English-filtered fastText classifier, but loses slightly once we add in a manual language filtering step to enable better performance on the multilingual LAMBADA datasets. For the maintext comparisons, we use the optional fastText classifier from our algorithm to select pretraining data at the page levels, but we show ablations without the classifier in Appendix E.

Figure 2 shows how each data selection method affects benchmark performance in more detail. Each block of rows represents a data selection method, while an individual row represents an LLM within a method that targets a particular benchmark or set of benchmarks. Columns represent benchmarks. We see that language filtering and perplexity correlations both clearly optimize for the target benchmark – within each block, the benchmark column matching each row typically performs best. Surprisingly, the pattern is much less obvious for DSIR – the heatmap looks much more uniform across LLMs with different task targets. We also see that while language filtering has significant impacts on model performance, our performance significantly exceeds the impact of language filtering across all tested benchmarks.

Figure 3 shows the distribution of languages in pretraining data selected by our method, targeting each benchmark. Our algorithm provides significant enrichment of the corresponding languages for the multilingual benchmarks (LAMBADA_*), but we also find that it does not *exclusively* select domains in one language. In contrast, for English benchmarks our approach selects nearly exclusively English data, likely due to the large quantity of high-quality English data in our pretraining data pool. There are significantly fewer tokens in non-English languages in the pretraining data pool and our $\tau$ constraint to prevent duplication has a large impact on the weights when the benchmarks are non-English. We provide the same figure when the $\tau$ values are made $5\times$ as large in Appendix F.

Figure: Figure 3: Language distributions of pretraining data selected by perplexity correlations. The default RPJv2 distribution is given in the left column for reference. The English benchmark targets often exclusively select English but the reverse is not the case. In every case, our approach selects more data than the default from the benchmark-matched language (shown as a green box in each column).
Refer to caption: extracted/5843524/images/token-dist-heatmap.png

Finally, we note that our results are somewhat insensitive to the specifics of the perplexity-correlation procedure we present in Algorithm 1. We show in Appendix E that varying the projection method (linear, $L_{2}$) and even using Spearman rank correlations (Spearman, ) often work better than the baselines. These results suggest that the performance of our approach is not tightly dependent on the precise form of the estimator that is coupled to our theory results but holds more broadly across general perplexity-correlation relationships. Additionally, our approach performs better with the optional fastText classifier that our algorithm trains, possibly because it can operate at the page-level instead of the domain-level(^4^44One might ask if the fastText classifier alone can be an effective data selector, since it appears in both high-performance selection methods. This is not possible – the classifier can only amplify the effectiveness of an existing selector rather than replace it – as training a classifier requires supervision, which is provided by either the perplexity correlation algorithm or human curation (in the case of Li et al. ()).).

### 5.2 Performance Rank Predictions

Figure: Figure 4: Rank predictions given by $\langle\hat{\bm{\theta}}^{\textnormal{proj}},\Phi(\mathbf{x})\rangle$ for PIQA and LAMBADA FR. A standard deviation ($\sigma$) from the ideal fit is shown in red. $2\sigma$ is shown in orange. Many models outside $2\sigma$ (named and shown in blue) are trained on atypical data such as heavily multilingual data, code, or GPT-4 (Brown et al., ) outputs. Models with atypical architectures (i.e. Mamba (Gu & Dao, )) are named and shown in black. Generally, our estimate tightly predicts ordinal benchmark performance from web corpus losses.
Refer to caption: extracted/5843524/images/piqa-5-fold-leave-out-rank.png

Figure: Figure 5: Held-out $R^{2}$ score of our raw correlation estimate $\hat{\bm{\theta}}$, our projected estimate $\hat{\bm{\theta}}^{\textnormal{proj}}$, and the average loss baseline. The $95\%$ bootstrapped confidence intervals are wide enough that no individual comparison is significant. Across benchmarks, $\hat{\bm{\theta}}^{\textnormal{proj}}$ has statistically significant gains over the baseline (p=0.035) as it is unlikely that $\hat{\bm{\theta}}^{\textnormal{proj}}$ beats mean loss 7 times out of 8 by chance.
Refer to caption: extracted/5843524/images/rank-r2-heatmap.png

We have shown that our approach succeeds at the goal of selecting useful pretraining data, but how good are single index model’s predictions? An accurate map between loss to benchmarks would be helpful in selecting among candidate pretraining data mixtures generally, even when not using our pretraining data selection algorithm.

Comparing model performance rankings predicted by our regression to the ground truth, we find generally accurate predictions. Figure 4 shows 5-fold leave-out plots for PIQA, and LAMBADA${}_{\text{FR}}$ with the rank predictions given by $\langle\hat{\bm{\theta}}^{\textnormal{proj}},\Phi(\mathbf{x})\rangle$. Every point in the plot is a held-out point: we estimated $\bm{{\bf{\theta}^{*}}}$ five times, holding out a different $20\%$ of the data each time, and plotted the prediction for every point when it was held out.

We find that our estimator achieves high ordinal prediction performance across all target tasks. We include 5-fold leave-out $R^{2}$ scores for all tasks in Figure 5. However, we complement these strong results with the additional observation that simply taking the *mean* loss across all domains is a strong predictor of model performance (bottom row). The surprising effectiveness of average loss over uniformly sampled documents has been discussed extensively (Owen, ; Wei et al., ; Kaplan et al., ) and our results further suggest that regressions with correlations only slightly above the mean loss baseline still can result in effective data selection methods.

Finally, we discuss outliers in our prediction of model performance. Our predictions are accurate for LLMs with usual architectures (e.g. Mamba (Gu & Dao, )), the smallest/largest vocabulary sizes, context sizes, and parameter sizes. However, we also see that LLMs that were trained on unusual data are not as well predicted by our approach (e.g. Phi (Gunasekar et al., )). We may simply require a bigger or more diverse pretraining data pool and set of models to find estimates that work well for models that expect different styles of text.

### 5.3 Analysis of the model-loss matrix 𝐗 𝐗 \mathbf{X} bold_X

What information is contained in the matrix of model losses $\mathbf{X}$? Clearly, it must contain semantically meaningful information about the data, such as the language that a piece of text is in. We performed PCA (Pearson, ) and t-SNE (van der Maaten & Hinton, ) on $\mathbf{X}$ and plotted the first two components for each of our 9,841 domains. As shown in the first row of Figure 6, we found two components with relatively high singular values. The first component clearly corresponds with the language of a domain. The second component corresponds with the average bits-per-byte or entropy of a domain. The t-SNE components show the same general pattern as well as showing that the language clusters are very well separated. As shown in our plots, there are several salient clusters within the language clusters. Within the English cluster, we found a subcluster for luxury goods, another for legal services and information, another for academic research, and even a cluster for funeral homes.

Figure: (a)
Refer to caption: extracted/5843524/images/eigenvalue-decay.png

The second row of Figure 6 shows plots for the loss matrix when we take the principal components of the other dimension, where points correspond to the 90 LLMs. For PCA, PC1 corresponds to entropy. For both cases, it is less clear what the other PCs are, but when we color the three largest families of models in our data (Pythia (Biderman et al., ), Qwen (Bai et al., ), and OpenLlama (Geng & Liu, )), we see that model families are clustered together in the PC graphs.

## 6 Robustness Checks with Pre-Registration

In small-scale experiments, our approach is competitive with the leading approach from ’s survey: a fixed fastText model (Joulin et al., ), manually augmented with the best language filtering. This leading approach is heuristic and hand-crafted, requiring appropriate language filtering matched to the target benchmark and assumptions about what good pretraining data looks like. Our approach does not make these assumptions and could potentially improve as more public models are released and we have better data to estimate $\bm{{\bf{\theta}^{*}}}$.

While our results are generally positive, many past data selection methods have reported initially positive results, only to later break: they may fail to scale to larger models or rely on specific details of their experimental setting. Our 160M-scale experiments may also raise such concerns.

We design a pre-registered scaling experiment that addresses both the concerns of scale and external validity. We use the permanence of arXiv preprints as a mechanism to preregister a series of scaling experiments within the DataComp-LM framework (Li et al., ), which is a testbed for data-selection techniques released with the recent survey. Pre-registering held-out scaling experiments commits us to reporting negative results, and avoid overfitting to our chosen experimental settings.

DataComp-LM is ideal for this preregistered scaling experiment, as it standardizes the setting by providing a pool of 240 trillion tokens, pretraining code for 412M to 7B parameter models, and evaluation code for 53 benchmarks, 22 of which are labelled as “core” benchmarks that scale predictably. Importantly, we have not trained *any* models on DataComp-LM using our methods or baselines, making this a true held-out experiment with known high-performance baselines.

##### Preregistered experiment.

We will run the best-performing approach from our paper: a fastText filter trained on our correlation estimator. We will define the target benchmark for our estimator as the average of the “core” DataComp-LM benchmarks and run our estimator with perplexities from our set of 90 OpenLM Leaderboard LLMs on a uniform subsample of the DataComp-LM pool of data. We will use the provided DataComp-LM code for training LLMs for the “Filtering 1B-1x” track, where a 1.4B parameter LLM is trained on 28.8B tokens chosen from a 1.64T sample of the DataComp-LM pool of data. In the DataComp-LM paper, they apply their fixed fastText filter as a final step after several complicated deduplication and filtering steps. We will report results where our fastText classifier is used as a direct substitute for this last step alone, as well as another test in which we replace the entire pipeline with one classifier. We will report results where our estimator is trained at the domain-level (following this paper) and where our estimator is trained at the page-level (which we have not tried yet). Finally, we will report analogous results where we replace the “core” evaluation score with the average score on all of the non-English LAMBADA translations, and compare the raw fastText classifier from Li et al. () to our approach, using both of these approaches in place of the full filtering pipeline from 1.64T tokens. We preregister this additional multilingual task because “core” does not include multilingual evaluations.

We chose the “Filtering 1B-1x” track because it is the most compute-intensive track that we can perform within a couple weeks, given our resources. For these experiments, the compute needs for our data selection procedure are negligible compared to LLM pretraining. Per Li et al. (), the pretraining for these experiments combined is estimated to take 240 hours on an 8-GPU H100 node.

## 7 Conclusion

Does high-performance data selection require careful hand-crafted heuristics or prohibitively expensive model training runs? Our work demonstrates an alternative, viable approach – leveraging existing, public models as a source of information for data selection. Pretraining experiments suggest that a simple, correlation-based approach to selecting data can be effective, but more broadly, we show how to 1) use single-index models as a surrogate for downstream performance and 2) build models that relate *losses* to downstream performance and use these surrogates effectively in data selection. Finally, we propose pre-registered scaling experiments on held-out data to test external validity of reported results. We hope this type of experiment will help improve the trustworthiness of experimental results for data selection.

#### Acknowledgments

We thank Jack Spilecki for conversations on the mathematical aspects of the work. We also thank Zitong Yang, Yangjun Ruan, and Lisa Li for their helpful feedback throughout the project, Ludwig Schmidt and Samir Gadre for discussions on scaling laws involving benchmark perplexity, Rohan Pandey for conversations about scaling laws, Sung Min Park for discussions on drafts of this work, and William Held for conversations about data selection.
This work is supported in part by a grant from Sandia National Laboratories, and gifts from Open Philanthropy, Meta, Panasonic Research, and the Tianqiao and Chrissy Chen Institute. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of Sandia National Laboratories. Tristan Thrush is supported in part by the Stanford Graduate Fellowship.

## Appendix A Estimator solution

### A.1 Lemma 1

###### Statement of Lemma 1

Define the PDF of HalfNormal as $f(x;\sigma)=\frac{\sqrt{2}}{\sigma\sqrt{\pi}}e^{-\frac{x^{2}}{2\sigma^{2}}}$ for $x>0$ and 0 otherwise. Now, suppose:

- •
$\bm{\beta}$ is a vector with $||\bm{\beta}||_{2}=1$
- •
$\mathbf{Z}_{1},\mathbf{Z}_{2}$ are vectors $\sim\mathcal{N}(\mathbf{0},\mathbf{I})$
- •
$\epsilon\sim\mathcal{N}(0,\sigma^{2})$
- •
$Z^{\prime}\sim\mathcal{N}(0,1)$
- •
$Z_{+}\sim\textnormal{HalfNormal}(1)$.

Then we have:

$$ ${Z_{1}}_{j}|\langle\mathbf{Z}_{1}-\mathbf{Z}_{2},\bm{\beta}\rangle+\epsilon>0 \overset{d}{=}Z^{\prime}\sqrt{1-\frac{\beta_{j}^{2}}{2+\sigma^{2}}}+\frac{ \beta_{j}}{\sqrt{2+\sigma^{2}}}Z_{+},$ $$

where ${Z_{1}}_{j}$ is the $j$-th entry of $\mathbf{Z}_{1}$.

Proof: First, note:

$$ ${Z_{1}}_{j}|\langle\mathbf{Z}_{1}-\mathbf{Z}_{2},\bm{\beta}\rangle+\epsilon>0 \overset{d}{=}{Z_{1}}_{j}|\left\langle\begin{bmatrix}\mathbf{Z}_{1}\\ \mathbf{Z}_{2}\\ \epsilon/\sigma\end{bmatrix},\begin{bmatrix}\bm{\beta}\\ -\bm{\beta}\\ \sigma\end{bmatrix}\right\rangle>0\overset{d}{=}{Z_{1}}_{j}|\left\langle\begin {bmatrix}\mathbf{Z}_{1}\\ \mathbf{Z}_{2}\\ \epsilon/\sigma\end{bmatrix},\begin{bmatrix}\bm{\beta}\\ -\bm{\beta}\\ \sigma\end{bmatrix}/\sqrt{2+\sigma^{2}}\right\rangle>0,$ $$

where $\begin{bmatrix}\cdot\\
\cdot\\
\cdot\end{bmatrix}$ denotes the vector-valued result of concatenating vectors and scalars. For readability, we set $\mathbf{Z}_{c}=\begin{bmatrix}\mathbf{Z}_{1}\\
\mathbf{Z}_{2}\\
\epsilon/\sigma\end{bmatrix}$ and $\bm{\beta}_{c}=\begin{bmatrix}\bm{\beta}\\
-\bm{\beta}\\
\sigma\end{bmatrix}/\sqrt{2+\sigma^{2}}$.

Given that $\bm{\beta}_{c}$ is unit-norm (by supposition, $\bm{\beta}$ is unit-norm), and every element of $\mathbf{Z}_{c}$ is $\sim\mathcal{N}(0,1)$ (even $\epsilon/\sigma$), we can easily split a conditional random vector containing ${Z_{1}}_{j}$ into a conditionally dependent component and independent component:

$$ $\mathbf{Z}_{c}|\langle\mathbf{Z}_{c},\bm{\beta}_{c}\rangle>0\overset{d}{=}( \mathbf{I}-\bm{\beta}_{c}\bm{\beta}_{c}^{\top})\mathbf{Z}^{\prime\prime}+\bm{ \beta}_{c}\mathbf{Z}_{+}.$ $$

The first term is orthogonal to $\bm{\beta}_{c}$ and so it is the part of $\mathbf{Z}_{c}$ that is not subject to the condition. In the unconditional case, $\mathbf{Z}_{c}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ and so $\mathbf{Z}^{\prime\prime}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. The second term is the part of $\mathbf{Z}_{c}$ that is in the direction of $\bm{\beta}_{c}$. $\mathbf{Z}_{+}\sim\text{HalfNormal}(\mathbf{I})$ because our dot product condition is satisfied for half of the possible non-orthogonal $\mathbf{Z}_{c}$ values. Now, we focus on finding $\mathbf{Z}_{c}|\langle\mathbf{Z}_{c},\bm{\beta}_{c}\rangle>0$ for a single index $j$. We have (for $C$ defined to be the dimensionality of $\bm{\beta}_{c}$):

$$ $\displaystyle((\mathbf{I}-\bm{\beta}_{c}\bm{\beta}_{c}^{\top})\mathbf{Z}^{ \prime\prime})_{j}+(\bm{\beta}_{c}\mathbf{Z}_{+})_{j}$ $\displaystyle=Z^{\prime\prime}_{j}(1-{\beta_{c}}_{j}^{2})-\sum_{\begin{ subarray}{c}1\leq i\leq C\\ i\neq j\end{subarray}}Z^{\prime\prime}_{i}{\beta_{c}}_{j}{\beta_{c}}_{i}+\beta _{j}{Z_{+}}_{j}$ $\displaystyle=Z^{\prime\prime}_{j}-\sum_{i=1}^{C}Z^{\prime\prime}_{i}{\beta_{c }}_{j}{\beta_{c}}_{i}+\beta_{j}{Z_{+}}_{j}.$ $$

Now, note that $Z^{\prime\prime}_{j}-\sum_{i=1}^{C}Z^{\prime\prime}_{i}{\beta_{c}}_{j}{\beta_{
c}}_{i}$ is the sum of independent zero-mean Gaussians with variances given by $1$ and ${\beta_{c}}_{j}^{2}{\beta_{c}}_{i}^{2}$, so it itself is a zero-mean Gaussian $Y\sim\mathcal{N}(0,1-\sum_{i=1}^{C}{\beta_{c}}_{j}^{2}{\beta_{c}}_{i}^{2})$. We can also use the fact that $\sum_{i=1}^{C}{\beta_{c}}_{i}^{2}=1$ (recall that $\bm{\beta_{c}}$ is unit norm) to get: $Y\sim\mathcal{N}(0,1-{\beta_{c}}_{j}^{2})$. So we have that the conditional ${Z_{1}}_{j}$ is given by:

$$ $Z^{\prime}\sqrt{1-{\beta_{c}}_{j}^{2}}+{\beta_{c}}_{j}Z_{+}=Z^{\prime}\sqrt{1- \frac{\beta_{j}^{2}}{2+\sigma^{2}}}+\frac{\beta_{j}}{\sqrt{2+\sigma^{2}}}Z_{+},$ $$

for $Z^{\prime}\sim\mathcal{N}(0,1)$. As a corollary, we can see that ${Z_{2}}_{j}$ under the same condition is given by:

$$ $Z^{\prime}\sqrt{1-\frac{\beta_{j}^{2}}{2+\sigma^{2}}}+\frac{-\beta_{j}}{\sqrt{ 2+\sigma^{2}}}Z_{+}.$ $$

###### Statement of Lemma 1

### A.2 Lemma 2

###### Statement of Lemma 2

Suppose that $\Phi$ is the CDF of a standard Gaussian, $a$ and $c$ are constants, and $Z\sim\mathcal{N}(0,1)$. Then we have:

$$ $\mathbb{E}[\Phi(aZ+c)]=\Phi\left(\frac{c}{\sqrt{1+a^{2}}}\right).$ $$

Proof:
By the definition of the CDF of a standard Gaussian, we have:

$$ $\displaystyle\mathbb{E}[\Phi(aZ+c)]$ $\displaystyle=\mathbb{E}[P(X\leq aZ+c)],$ $$

where $X\sim\mathcal{N}(0,1)$. Continuing, we have:

$$ $\displaystyle=\mathbb{E}[P(X-aZ-c\leq 0)].$ $$

Now, note that $X-aZ-c$ is the sum of independent Gaussian random variables with given mean and variance; it itself is a Gaussian random variable $\sim\mathcal{N}(-c,a^{2}+1)$. To find $P(X-aZ-c\leq 0)$, we can evaluate its CDF at 0 0:

$$ $\displaystyle=\mathbb{E}\left[\Phi\left(\frac{c}{\sqrt{a^{2}+1}}\right)\right] =\Phi\left(\frac{c}{\sqrt{a^{2}+1}}\right).$ $$

###### Statement of Lemma 2

### A.3 Lemma 3

###### Statement of Lemma 3

Suppose $\Phi$ is the standard Gaussian CDF, $Z_{+}\sim\textnormal{HalfNormal}(1)$, and $b$ and $a$ are constants. Then we have:

$$ $\mathbb{E}\left[\Phi\left(\frac{Z_{+}b}{\sqrt{a^{2}+1}}\right)\right]=\frac{1} {2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b}{\sqrt{a^{2}+1}}\right).$ $$

Proof: By the definition of expected value, we can take the following integral where $f_{Z_{+}}$ is the PDF of $Z_{+}$. We integrate from 0 0 instead of $-\infty$ because the PDF of the Standard Half Normal is 0 0 in the domain below 0 0:

$$ $\displaystyle\mathbb{E}\left[\Phi\left(\frac{Z_{+}b}{\sqrt{a^{2}+1}}\right)\right]$ $\displaystyle=\int_{0}^{\infty}\Phi\left(\frac{zb}{\sqrt{a^{2}+1}}\right)f_{Z_ {+}}(z)dz$ $\displaystyle=\int_{0}^{\infty}\Phi\left(\frac{zb}{\sqrt{a^{2}+1}}\right)\frac {\sqrt{2}}{\sqrt{\pi}}e^{\frac{-z^{2}}{2}}dz$ $\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\int_{0}^{\infty}e^{\frac{-z^{2}}{2}} dz+\int_{0}^{\infty}\text{erf}\left(\frac{zb}{\sqrt{2}\sqrt{a^{2}+1}}\right)e^ {\frac{-z^{2}}{2}}dz\right)\textbf{ (\textasteriskcentered)}.$ $$

The second integral is generally non-trivial to solve, but luckily we can solve it by using Equation 2 in Section 4.3 of the integral table from Ng & Geller (), which states:

$$ $\int_{0}^{\infty}\text{erf}(cx)e^{-d^{2}x^{2}}dx=\frac{\sqrt{\pi}}{2d}-\frac{1 }{d\sqrt{\pi}}\tan^{-1}\left(\frac{d}{c}\right)$ $$

Where $c$ and $d$ are real and positive. We split the solution by cases: $b>0$, $b=0$, and $b<0$. We find that in every case, we can manipulate our integral so that the solution is trivial or the constant inside the $\text{erf}(\cdot)$ is positive (and so we can use the integral table). In every case, we find that the solution is $\frac{1}{2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b}{\sqrt{a^{2}+1}}\right)$.

Case 1: $b>0$. We can use the integral table directly:

$$ (*) $\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\frac{\sqrt{\pi}}{\sqrt{2}}+\frac{ \sqrt{\pi}}{\sqrt{2}}-\frac{\sqrt{2}}{\sqrt{\pi}}\tan^{-1}\left(\frac{\sqrt{a^ {2}+1}}{b}\right)\right)$ $\displaystyle=\frac{1}{2}+\frac{1}{2}-\frac{1}{\pi}\tan^{-1}\left(\frac{\sqrt{ a^{2}+1}}{b}\right).$ $$

Then, using the identity:

$$ $\tan^{-1}x+\tan^{-1}\frac{1}{x}=\frac{\pi}{2}\text{ if }x>0,\\$ $$

we find the following:

$$ $=\frac{1}{2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b}{\sqrt{a^{2}+1}}\right).$ $$

Case 2: $b=0$. Note that $\text{erf}(0)=0$; we do not have to use the integral table:

$$ (*) $\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\frac{\sqrt{\pi}}{\sqrt{2}}+0\right)$ $\displaystyle=\frac{1}{2}.$ $$

Because $\tan^{-1}(0)=0$, we have:

$$ $=\frac{1}{2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b}{\sqrt{a^{2}+1}}\right).$ $$

Case 3: $b<0$. Because $\text{erf}(\cdot)$ is an odd function, we can pull the negative out:

$$ (*) $\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\int_{0}^{\infty}e^{\frac{-z^{2}}{2}} dz-\int_{0}^{\infty}\text{erf}\left(\frac{z|b|}{\sqrt{2}\sqrt{a^{2}+1}}\right) e^{\frac{-z^{2}}{2}}dz\right).$ $$

Now we can use the integral table as in the $b>0$ case:

$$ $\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\frac{\sqrt{\pi}}{\sqrt{2}}-\frac{ \sqrt{\pi}}{\sqrt{2}}+\frac{\sqrt{2}}{\sqrt{\pi}}\tan^{-1}\left(\frac{\sqrt{a^ {2}+1}}{|b|}\right)\right)$ $\displaystyle=\frac{1}{2}+\frac{1}{2}-\frac{1}{\pi}\tan^{-1}\left(\frac{\sqrt{ a^{2}+1}}{|b|}\right).$ $$

We can then use the same identity again:

$$ $\tan^{-1}x+\tan^{-1}\frac{1}{x}=\frac{\pi}{2}\text{ if }x>0\\$ $$

to get:

$$ $=\frac{1}{2}-\frac{1}{\pi}\tan^{-1}\left(\frac{|b|}{\sqrt{a^{2}+1}}\right).$ $$

Because $\tan^{-1}$ is an odd function, we can put the negative inside of it:

$$ $=\frac{1}{2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b}{\sqrt{a^{2}+1}}\right).$ $$

###### Statement of Lemma 3

### A.4 Full Proof

Here, we prove:

$$ $\mathbb{E}[\operatorname{sign}(y_{1}-y_{2})(\Phi(\mathbf{x}_{1})-\Phi(\mathbf{ x}_{2}))]=\frac{2}{\pi}\sin^{-1}\left(\frac{\bm{{\bf{\theta}^{*}}}}{\sqrt{4+2 \sigma_{1}^{2}+2\sigma_{2}^{2}}}\right)$ $$

with $y_{1},y_{2},\Phi(\mathbf{x}_{1}),\Phi(\mathbf{x}_{2})$, and $\bm{{\bf{\theta}^{*}}}$ defined in the main text, for the case where $\epsilon_{1}$ and $\epsilon_{2}$ are zero-mean Gaussian noise $\sim\mathcal{N}(0,\sigma^{2}_{1})$ and $\sim\mathcal{N}(0,\sigma^{2}_{2})$, respectively.

It is easy to see that this is a more general version of the following theorem.

See 1

Proof: By symmetry, we have:

$$ $\displaystyle\mathbb{E}[\operatorname{sign}(y_{1}-y_{2})(\Phi(\mathbf{x}_{1})- \Phi(\mathbf{x}_{2}))]$ $\displaystyle=\frac{1}{2}\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})| \operatorname{sign}(y_{1}-y_{2})>0]+\frac{1}{2}\mathbb{E}[-(\Phi(\mathbf{x}_{1 })-\Phi(\mathbf{x}_{2}))|\operatorname{sign}(y_{1}-y_{2})<0].$ $$

By increasing monotonicity of $f$, we have $\operatorname{sign}(y_{1}-y_{2})>0\iff\langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm
{{\bf{\theta}^{*}}}\rangle+\epsilon_{\Delta}>0$, for $\epsilon_{\Delta}=\epsilon_{1}-\epsilon_{2}\sim\mathcal{N}(0,\sigma^{2}_{1}+
\sigma^{2}_{2})$. So:

$$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})| \langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm{{\bf{\theta}^{*}}}\rangle+\epsilon_{ \Delta}>0]$ $\displaystyle+\frac{1}{2}\mathbb{E}[-(\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2} ))|\langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm{{\bf{\theta}^{*}}}\rangle+\epsilon _{\Delta}<0].$ $$

Because $\mathbf{x}_{1}\overset{d}{=}\mathbf{x}_{2}$ and $\epsilon_{\Delta}\overset{d}{=}-\epsilon_{\Delta}$, the two expected values above are the same:

$$ $=\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})|\langle\mathbf{x}_{1}- \mathbf{x}_{2},\bm{{\bf{\theta}^{*}}}\rangle+\epsilon_{\Delta}>0].$ $$

By linearity of expectation:

$$ $\displaystyle=\mathbb{E}[\Phi(\mathbf{x}_{1})|\langle\mathbf{x}_{1}-\mathbf{x} _{2},\bm{{\bf{\theta}^{*}}}\rangle+\epsilon_{\Delta}>0]-\mathbb{E}[\Phi( \mathbf{x}_{2})|\langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm{{\bf{\theta}^{*}}} \rangle+\epsilon_{\Delta}>0].$ $$

Now, we focus on finding the overall estimate for a single index $j$. By Lemma 1, we have, for $Z\sim\mathcal{N}(0,1)$ and $Z_{+}\sim\text{HalfNormal}(1)$:

$$ $\displaystyle\Phi({x_{1}}_{j})|\langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm{{\bf{ \theta}^{*}}}\rangle+\epsilon_{\Delta}>0\overset{d}{=}\Phi(Za+Z_{+}b_{1}).$ $$

Here, $a=\sqrt{1-\frac{(\theta_{j}^{*})^{2}}{2+\sigma^{2}_{1}+\sigma^{2}_{2}}}$ and $b_{1}=\frac{\theta_{j}^{*}}{\sqrt{2+\sigma^{2}_{1}+\sigma^{2}_{2}}}$. As a corollary of Lemma 1, we can see:

$$ $\displaystyle\Phi({x_{2}}_{j})|\langle\mathbf{x}_{1}-\mathbf{x}_{2},\bm{{\bf{ \theta}^{*}}}\rangle+\epsilon_{\Delta}>0\overset{d}{=}\Phi(Za+Z_{+}b_{2}).$ $$

Where $b_{2}=-\frac{\theta_{j}^{*}}{\sqrt{2+\sigma^{2}_{1}+\sigma^{2}_{2}}}$. So for the index $j$, our estimate is:

$$ $\displaystyle\mathbb{E}[\Phi(Za+Z_{+}b_{1})]-\mathbb{E}[\Phi(Za+Z_{+}b_{2})]$ $\displaystyle=\mathbb{E}[\mathbb{E}[\Phi(Za+c)|c=Z_{+}b_{1}]]-\mathbb{E}[ \mathbb{E}[\Phi(Za+c)|c=Z_{+}b_{2}]].$ $$

Using Lemma 2, we have:

$$ $\displaystyle=\mathbb{E}\left[\Phi\left(\frac{Z_{+}b_{1}}{\sqrt{a^{2}+1}} \right)\right]-\mathbb{E}\left[\Phi\left(\frac{Z_{+}b_{2}}{\sqrt{a^{2}+1}} \right)\right].$ $$

Then, using Lemma 3, we have:

$$ $\displaystyle=\frac{1}{2}+\frac{1}{\pi}\tan^{-1}\left(\frac{b_{1}}{\sqrt{a^{2} +1}}\right)-\frac{1}{2}-\frac{1}{\pi}\tan^{-1}\left(\frac{b_{2}}{\sqrt{a^{2}+1 }}\right)$ $\displaystyle=\frac{1}{\pi}\tan^{-1}\left(\frac{b_{1}}{\sqrt{a^{2}+1}}\right)- \frac{1}{\pi}\tan^{-1}\left(\frac{b_{2}}{\sqrt{a^{2}+1}}\right).$ $$

Using the fact that $\tan^{-1}$ is an odd function and $b_{2}=-b_{1}$, we get:

$$ $\displaystyle=\frac{2}{\pi}\tan^{-1}\left(\frac{b_{1}}{\sqrt{a^{2}+1}}\right).$ $$

Now, we write $a$ and $b_{1}$ in terms of $\theta_{j}^{*}$:

$$ $\displaystyle=\frac{2}{\pi}\tan^{-1}\left(\frac{\frac{\theta_{j}^{*}}{\sqrt{2+ \sigma^{2}_{1}+\sigma^{2}_{2}}}}{\sqrt{2-\frac{(\theta_{j}^{*})^{2}}{2+\sigma^ {2}_{1}+\sigma^{2}_{2}}}}\right)$ $\displaystyle=\frac{2}{\pi}\tan^{-1}\left(\frac{\frac{\theta_{j}^{*}}{\sqrt{4+ 2\sigma^{2}_{1}+2\sigma^{2}_{2}}}}{\sqrt{1-\left(\frac{\theta_{j}^{*}}{\sqrt{4 +2\sigma^{2}_{1}+2\sigma^{2}_{2}}}\right)^{2}}}\right).$ $$

Using the identity $\sin^{-1}x=\tan^{-1}\left(\frac{x}{\sqrt{1-x^{2}}}\right)$, we have:

$$ $\displaystyle=\frac{2}{\pi}\sin^{-1}\left(\frac{\theta_{j}^{*}}{\sqrt{4+2 \sigma^{2}_{1}+2\sigma^{2}_{2}}}\right).$ $$

### A.5 Corollary 1

See 1

To see this, we can find:

$$ $\displaystyle\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})|\langle\hat{ \bm{\theta}},\mathbf{x}_{1}\rangle+\epsilon_{1}>\langle\hat{\bm{\theta}}, \mathbf{x}_{2}\rangle+\epsilon_{2}]=\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi( \mathbf{x}_{2})|\langle\hat{\bm{\theta}},\mathbf{x}_{1}-\mathbf{x}_{2}\rangle+ \epsilon_{\Delta}>0]$ $$

Note that we have already computed this expected value in the proof above; for an index $j$, it is:

$$ $\displaystyle\frac{2}{\pi}\sin^{-1}\left(\frac{\hat{\theta}_{j}}{\sqrt{4+2 \sigma^{2}_{1}+2\sigma^{2}_{2}}}\right).$ $$

Because $\sin^{-1}$ is an odd function, the above expression has the same sign as $\hat{\theta}_{j}$. Because the values at every index of $\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})]$ under our condition and $\hat{\bm{\theta}}$ are the same sign, we have $\langle\mathbb{E}[\Phi(\mathbf{x}_{1})-\Phi(\mathbf{x}_{2})],\hat{\bm{\theta}}
\rangle>0$, so $\langle\hat{\bm{\theta}},\mathbb{E}[\Phi(\mathbf{x}_{1})]\rangle>\langle\hat{
\bm{\theta}},\mathbb{E}[\Phi(\mathbf{x}_{2})]\rangle$.

## Appendix B Optimal projected weights solutions

### B.1 Linear Projection

See 2

Proof: We proceed by considering each of the three cases from Equation 10.

Case 1.
Suppose for the sake of contradiction that the optimal solution is $\hat{\bm{\theta}}^{\textnormal{proj}}$ and yet $\hat{\theta}^{\textnormal{proj}}_{k}<\tau_{k}$ for some $\hat{\theta}^{\textnormal{proj}}_{k}$ falling under the first case of Equation 10. Now suppose that we construct a $\bm{\theta}^{\prime}$ also satisfying the projection constraints that is the same as $\hat{\bm{\theta}}^{\textnormal{proj}}$ except in these places:

$$ $\displaystyle\theta^{\prime}_{k}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{k}+\Delta=\tau_{k}$ $\displaystyle\theta^{\prime}_{p}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{p}-\delta_{1}\geq 0$ $\displaystyle\vdots$ $\displaystyle\theta^{\prime}_{q}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{q}-\delta_{n}\geq 0$ $$

for some $\Delta=\sum_{i=1}^{n}\delta_{i}>0$ where $\hat{\theta}_{p}\geq\cdots\geq\hat{\theta}_{q}$ are all of the $\hat{\theta}$ values which do not fall under the first condition and where the corresponding $\hat{\theta}^{\textnormal{proj}}$ values are nonzero. We know that there must be some $\hat{\theta}^{\textnormal{proj}}_{p},\cdots,\hat{\theta}^{\textnormal{proj}}_{q}$ from which we can subtract $\delta_{1},\cdots,\delta_{n}$ (and so from which we can take the $\Delta$) because $\sum_{j\text{: }r_{j}(\hat{\theta}_{j})\geq r_{k}(\hat{\theta}_{k})}\tau_{j}\leq
1$. Now, we have:

$$ $\displaystyle\langle\hat{\bm{\theta}},\hat{\bm{\theta}}^{\textnormal{proj}} \rangle-\langle\hat{\bm{\theta}},\bm{\theta}^{\prime}\rangle$ $\displaystyle=\hat{\theta}_{k}\hat{\theta}^{\textnormal{proj}}_{k}+\hat{\theta }_{p}\hat{\theta}^{\textnormal{proj}}_{p}+\cdots+\hat{\theta}_{q}\hat{\theta}^ {\textnormal{proj}}_{q}-\hat{\theta}_{k}\hat{\theta}^{\textnormal{proj}}_{k}- \hat{\theta}_{k}\Delta-\hat{\theta}_{p}\hat{\theta}^{\textnormal{proj}}_{p}- \cdots-\hat{\theta}_{q}\hat{\theta}^{\textnormal{proj}}_{q}+\hat{\theta}_{p} \delta_{1}+\cdots+\hat{\theta}_{q}\delta_{n}$ $\displaystyle=-\hat{\theta}_{k}\Delta+\hat{\theta}_{p}\delta_{1}+\cdots+\hat{ \theta}_{q}\delta_{n}$ $\displaystyle\leq\hat{\theta}_{p}(\delta_{1}+\cdots+\delta_{n})-\hat{\theta}_{ k}\Delta$ $\displaystyle=\hat{\theta}_{p}\Delta-\hat{\theta}_{k}\Delta$ $\displaystyle\leq 0.$ $$

At this point, the only way to avoid the contradiction result would be if $\hat{\theta}_{k}=\hat{\theta}_{p}=\cdots=\hat{\theta}_{q}$. Otherwise, the above non-strict inequality would be a strict inequality. If $\hat{\theta}_{k}=\hat{\theta}_{p}=\cdots=\hat{\theta}_{q}$, then we know that $\hat{\theta}_{k}$ is the smallest $\hat{\theta}$ value satisfying condition 1 and all of the other greater $\hat{\theta}$ values satisfying condition 1 must be projected to their $\tau$ threshold value (otherwise we would get the contradiction result). In this edge case can see above that rearranging the remaining weight among equal $\hat{\theta}$ values does not change the dot product, so all of the solutions that we can get without the contradiction result are equivalently optimal (including the solution from Equation 10).

Case 3.
This is analogous to case 1. Suppose for the sake of contradiction that the optimal solution is $\hat{\bm{\theta}}^{\textnormal{proj}}$ and yet $\hat{\theta}^{\textnormal{proj}}_{k}>0$ for some $\hat{\theta}^{\textnormal{proj}}_{k}$ falling under the third case of Equation 10. Now suppose that we construct a $\bm{\theta}^{\prime}$ also satisfying the projection constraints that is the same as $\hat{\bm{\theta}}^{\textnormal{proj}}$ except in these places:

$$ $\displaystyle\theta^{\prime}_{k}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{k}-\Delta=0$ $\displaystyle\theta^{\prime}_{p}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{p}+\delta_{1}\leq\tau_{p}$ $\displaystyle\vdots$ $\displaystyle\theta^{\prime}_{q}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{q}+\delta_{n}\leq\tau_{q}$ $$

for some $\Delta=\sum_{i=1}^{n}\delta_{i}>0$ where $\hat{\theta}_{p}\geq\cdots\geq\hat{\theta}_{q}$ are all of the $\hat{\theta}$ values which do not fall under the third condition and where the corresponding $\hat{\theta}^{\textnormal{proj}}$ values are not at their thresholds. By construction we know that there must be some $\hat{\theta}^{\textnormal{proj}}_{p},\cdots,\hat{\theta}^{\textnormal{proj}}_{q}$ to which we can add $\delta_{1},\cdots,\delta_{n}$. Now, we have:

$$ $\displaystyle\langle\hat{\bm{\theta}},\hat{\bm{\theta}}^{\textnormal{proj}} \rangle-\langle\hat{\bm{\theta}},\bm{\theta}^{\prime}\rangle$ $\displaystyle=\hat{\theta}_{k}\hat{\theta}^{\textnormal{proj}}_{k}+\hat{\theta }_{p}\hat{\theta}^{\textnormal{proj}}_{p}+\cdots+\hat{\theta}_{q}\hat{\theta}^ {\textnormal{proj}}_{q}-\hat{\theta}_{k}\hat{\theta}^{\textnormal{proj}}_{k}+ \hat{\theta}_{k}\Delta-\hat{\theta}_{p}\hat{\theta}^{\textnormal{proj}}_{p}- \cdots-\hat{\theta}_{q}\hat{\theta}^{\textnormal{proj}}_{q}-\hat{\theta}_{p} \delta_{1}-\cdots-\hat{\theta}_{q}\delta_{n}$ $\displaystyle=\hat{\theta}_{k}\Delta-\hat{\theta}_{p}\delta_{1}-\cdots-\hat{ \theta}_{q}\delta_{n}$ $\displaystyle\leq-\hat{\theta}_{q}(\delta_{1}+\cdots+\delta_{n})+\hat{\theta}_ {k}\Delta$ $\displaystyle=-\hat{\theta}_{q}\Delta+\hat{\theta}_{k}\Delta$ $\displaystyle\leq 0.$ $$

At this point, the only way to avoid the contradiction result would be if $\hat{\theta}_{k}=\hat{\theta}_{p}=\cdots=\hat{\theta}_{q}$. Otherwise, the above non-strict inequality would be a strict inequality. If $\hat{\theta}_{k}=\hat{\theta}_{p}=\cdots=\hat{\theta}_{q}$, then we know that $\hat{\theta}_{k}$ is the largest $\hat{\theta}$ value satisfying condition 3 and all of the other smaller $\hat{\theta}$ values satisfying condition 3 must be projected to 0 (otherwise we would get the contradiction result). In this edge case, we can see above that rearranging the remaining weight among equal $\hat{\theta}$ values does not change the dot product, so all of the solutions that we can get without the contradiction result are equivalently optimal (including the solution from Equation 10).

Case 2.
Above, we show that both Case 1 and Case 3 are true. So, the remaining weight must be given to the single value of $\hat{\bm{\theta}}^{\textnormal{proj}}$ not covered by either case.

### B.2 Quadratic Projection

#### B.2.1 Lemma 4

###### Statement of Lemma 4

Suppose that $\hat{\bm{\theta}}^{\textnormal{proj}}$ is the optimal solution to:

$$ $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in \mathbb{R}^{D}}||\hat{\bm{\theta}}-\bm{\theta}||_{2}^{2},$ $$

subject to:

$$ $\displaystyle\sum_{i=1}^{D}\theta_{i}=1$ $\displaystyle 0\leq\theta_{i}\leq\tau_{i},\forall i\in[1,D],$ $$

where $\tau_{i}>0$ are fixed values. Then, $\hat{\theta}^{\textnormal{proj}}_{s}=0$ implies that any $j$ with $\hat{\theta}_{s}>\hat{\theta}_{j}$ must have $\hat{\theta}^{\textnormal{proj}}_{j}=0$.

Proof: This is similar to Lemma 2 from Shalev-Shwartz & Singer (). Assume for the sake of contradiction $\hat{\theta}^{\textnormal{proj}}_{s}=0$ and $\hat{\theta}_{s}>\hat{\theta}_{j}$, yet we have $\hat{\theta}^{\textnormal{proj}}_{j}>0$.

Now we can construct another vector $\bm{\theta}^{\prime}$ that is the same as $\hat{\bm{\theta}}^{\textnormal{proj}}$, except in two places:

$$ $\displaystyle\theta^{\prime}_{s}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{s}+\Delta$ $\displaystyle\theta^{\prime}_{j}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{j}-\Delta,$ $$

for some $\Delta$ satisfying $0<\Delta<\min(\hat{\theta}^{\textnormal{proj}}_{j},\tau_{s}-\hat{\theta}^{
\textnormal{proj}}_{s})$. This bound on $\Delta$ ensures that $\theta^{\prime}$ is still within the thresholds. We know that $\Delta$ can exist because $\min(\hat{\theta}^{\textnormal{proj}}_{j},\tau_{s}-\hat{\theta}^{\textnormal{
proj}}_{s})>0$ (by supposition, $\tau_{s}-\hat{\theta}^{\textnormal{proj}}_{s}=\tau_{s}-0>0$ and $\hat{\theta}^{\textnormal{proj}}_{j}>0$).

Now we can compute:

$$ $\displaystyle||\hat{\bm{\theta}}-\hat{\bm{\theta}}^{\textnormal{proj}}||^{2}_{ 2}-||\hat{\bm{\theta}}-\bm{\theta}^{\prime}||^{2}_{2}$ $\displaystyle=(\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})^{2}+( \hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})^{2}-(\hat{\theta}_{s}-( \hat{\theta}^{\textnormal{proj}}_{s}+\Delta))^{2}-(\hat{\theta}_{j}-(\hat{ \theta}^{\textnormal{proj}}_{j}-\Delta))^{2}$ $\displaystyle=2\Delta((\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})- (\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})-\Delta)$ $\displaystyle>2\Delta((\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})- (\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})-\min(\hat{\theta}^{ \textnormal{proj}}_{j},\tau_{s}-\hat{\theta}^{\textnormal{proj}}_{s}))$ $\displaystyle\geq 2\Delta((\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{ s})-(\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})-\hat{\theta}^{ \textnormal{proj}}_{j})$ $\displaystyle=2\Delta(\hat{\theta}_{s}-\hat{\theta}_{j})$ $\displaystyle>0.$ $$

So $\hat{\bm{\theta}}^{\textnormal{proj}}$ cannot be the optimal solution.

###### Statement of Lemma 4

#### B.2.2 Lemma 5

###### Statement of Lemma 5

Suppose that $\hat{\bm{\theta}}^{\textnormal{proj}}$ is the optimal solution to:

$$ $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in \mathbb{R}^{D}}||\hat{\bm{\theta}}-\bm{\theta}||_{2}^{2},$ $$

subject to:

$$ $\displaystyle\sum_{i=1}^{D}\theta_{i}=1$ $\displaystyle 0\leq\theta_{i}\leq\tau_{i},\forall i\in[1,D],$ $$

where $\tau_{i}>0$ are fixed values.
Then, $\hat{\theta}^{\textnormal{proj}}_{s}=\tau_{s}$ implies $\hat{\theta}^{\textnormal{proj}}_{j}=\tau_{j}$ for any $\hat{\theta}_{j}-\tau_{j}>\hat{\theta}_{s}-\tau_{s}$.

Proof: Again, this is similar to Lemma 2 from Shalev-Shwartz & Singer (). Assume for the sake of contradiction $\hat{\theta}^{\textnormal{proj}}_{s}=\tau_{s}$ and $\hat{\theta}_{j}-\tau_{j}>\hat{\theta}_{s}-\tau_{s}$, yet we have $\hat{\theta}^{\textnormal{proj}}_{j}<\tau_{j}$.

Now we can construct another vector $\bm{\theta}^{\prime}$ that is the same as $\hat{\bm{\theta}}^{\textnormal{proj}}$, except in two places:

$$ $\displaystyle\theta^{\prime}_{s}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{s}-\Delta$ $\displaystyle\theta^{\prime}_{j}$ $\displaystyle=\hat{\theta}^{\textnormal{proj}}_{j}+\Delta,$ $$

for some $\Delta$ satisfying $0<\Delta<\min(\hat{\theta}^{\textnormal{proj}}_{s},\tau_{j}-\hat{\theta}^{
\textnormal{proj}}_{j})$. This bound on $\Delta$ ensures that $\theta^{\prime}$ is still within the thresholds. We know that $\Delta$ can exist because $\min(\hat{\theta}^{\textnormal{proj}}_{s},\tau_{j}-\hat{\theta}^{\textnormal{
proj}}_{j})>0$ (by supposition, $\tau_{j}-\hat{\theta}^{\textnormal{proj}}_{j}>0$ and $\hat{\theta}^{\textnormal{proj}}_{s}=\tau_{s}>0$).

Now we can compute:

$$ $\displaystyle||\hat{\bm{\theta}}-\hat{\bm{\theta}}^{\textnormal{proj}}||^{2}_{ 2}-||\hat{\bm{\theta}}-\bm{\theta}^{\prime}||^{2}_{2}$ $\displaystyle=(\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})^{2}+( \hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})^{2}-(\hat{\theta}_{s}-( \hat{\theta}^{\textnormal{proj}}_{s}-\Delta))^{2}-(\hat{\theta}_{j}-(\hat{ \theta}^{\textnormal{proj}}_{j}+\Delta))^{2}$ $\displaystyle=2\Delta((\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})- (\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})-\Delta)$ $\displaystyle>2\Delta((\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{j})- (\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})-\min(\hat{\theta}^{ \textnormal{proj}}_{s},\tau_{j}-\hat{\theta}^{\textnormal{proj}}_{j}))$ $\displaystyle\geq 2\Delta((\hat{\theta}_{j}-\hat{\theta}^{\textnormal{proj}}_{ j})-(\hat{\theta}_{s}-\hat{\theta}^{\textnormal{proj}}_{s})-(\tau_{j}-\hat{ \theta}^{\textnormal{proj}}_{j}))$ $\displaystyle=2\Delta((\hat{\theta}_{j}-\tau_{j})-(\hat{\theta}_{s}-\hat{ \theta}^{\textnormal{proj}}_{s}))$ $\displaystyle=2\Delta((\hat{\theta}_{j}-\tau_{j})-(\hat{\theta}_{s}-\tau_{s}))$ $\displaystyle>0.$ $$

So $\hat{\bm{\theta}}^{\textnormal{proj}}$ cannot be the optimal solution.

###### Statement of Lemma 5

#### B.2.3 Full Proof

###### Theorem 3

Suppose we want to solve:

$$ $\hat{\bm{\theta}}^{\textnormal{proj}}=\operatorname*{arg\,min}_{\bm{\theta}\in \mathbb{R}^{D}}||\hat{\bm{\theta}}-\bm{\theta}||_{2}^{2},$ $$

subject to:

$$ $\displaystyle\sum_{i=1}^{D}\theta_{i}=1$ $\displaystyle 0\leq\theta_{i}\leq\tau_{i},\forall i\in[1,D],$ $$

where $\tau_{i}>0$ are fixed values.
Then the solution is:

$$ $\hat{\theta}^{\textnormal{proj}}_{k}=\min(\max(\hat{\theta}_{k}-\lambda,0), \tau_{k}),$ $$

where $\lambda$ is found (through e.g. bisection search) to satisfy:

$$ $\sum_{i=1}^{D}\min(\max(\hat{\theta}_{i}-\lambda,0),\tau_{i})=1.$ $$

Proof: Note that this problem is the same as the simplex projection problem from Shalev-Shwartz & Singer () and Duchi et al. (), except here we have additional $\theta_{i}\leq\tau_{i}$ constraints. The Lagrangian for this problem is(^5^55Note that multiplying $||\hat{\bm{\theta}}^{\textnormal{proj}}-\bm{\theta}||_{2}^{2}$ by $\frac{1}{2}$ does not change the minimization problem and enables us to get rid of a factor of $2$ after taking the derivative of the Lagrangian.):

$$ $\mathcal{L}(\bm{\theta},\mathbf{\mu},\mathbf{\zeta},\lambda)=\frac{1}{2}||\hat {\bm{\theta}}-\bm{\theta}||_{2}^{2}+\lambda\left(-1+\sum_{i=1}^{N}\theta_{i} \right)-\langle\mathbf{\mu},\bm{\theta}\rangle+\langle\mathbf{\zeta},\bm{ \theta}-\mathbf{\tau}\rangle.$ $$

To find the optimality condition with respect single index of $\bm{\theta}$, we set the derivative to zero:

$$ $\frac{d\mathcal{L}}{d\theta_{i}}=\theta_{i}-\hat{\theta}_{i}+\lambda-\mu_{i}+ \zeta_{i}=0.$ $$

The complimentary slackness KKT condition gives us that $\zeta_{i}=\mu_{i}=0$ when $0<\theta_{i}<\tau_{i}$, so for $\theta_{i}$ not at the boundary of our constraints, we get:

$$ $\theta_{i}=\hat{\theta}_{i}-\lambda.$ $$

So, we have that for all $\theta_{i}\in(0,\tau_{i})$, there is a shared value $\lambda$ which we subtract from $\hat{\theta}_{i}$ to get the value of $\theta_{i}$. How do we know which $\theta_{i}$ are 0 0 and which $\theta_{i}$ are $\tau_{i}$, though?

Assume that we know $\lambda$. By Lemma 4, we can characterize the optimal solution as:

$$ $\hat{\theta}^{\textnormal{proj}}_{k}=\max(\hat{\theta}_{k}-\lambda,0),$ $$

for $\hat{\theta}^{\textnormal{proj}}_{k}\neq\tau_{k}$. By Lemma 5, we can characterize the optimal solution as:

$$ $\hat{\theta}^{\textnormal{proj}}_{k}=\min(\hat{\theta}_{k}-\lambda,\tau_{k}),$ $$

for $\hat{\theta}^{\textnormal{proj}}_{k}\neq 0$. So, we can combine these two forms to get:

$$ $\hat{\theta}^{\textnormal{proj}}_{k}=\min(\max(\hat{\theta}_{k}-\lambda,0), \tau_{k}).$ $$

Now recall that we have the following constraint:

$$ $\sum_{i=1}^{D}\min(\max(\hat{\theta}_{i}-\lambda,0),\tau_{i})=1.$ $$

Given this constraint, we can find $\lambda$ through search (moving the value up or down). We can see this by noticing that $\sum_{i=1}^{D}\min(\max(\hat{\theta}_{i}-\lambda,0),\tau_{i})$ is a strictly decreasing function of $\lambda$ between the setting of $\lambda$ that makes $\hat{\theta}_{i}-\lambda>0$ for at least one $i$, and the setting of $\lambda$ that makes $\hat{\theta}_{i}-\lambda<\tau_{i}$ for at least one $i$. So in this range, there is only one setting of $\lambda$ that satisfies this equation. We can only choose a $\lambda$ outside of this range when $\sum_{i=1}^{D}\tau_{i}=1$, and in this case the solution is trivial: $\hat{\theta}^{\textnormal{proj}}_{i}=\tau_{i}$ for all $i$.

###### Theorem 3

## Appendix C Loss Matrix Computation Specifics

For all of our experiments, we computed the loss matrix as follows. For efficiency purposes, we sampled only 25 pages for a domain’s bits-per-byte (BPB) computation even if a domain had more than 25 pages. To get an LLM’s BPB on a page, we split the page into chunks of text that were 512 tokens according to a reference tokenizer (we used the Llama 2 7B tokenizer; Touvron et al. ). These text chunks turned out to be small enough to fit in the context of every LLM we tested. We then averaged BPB across chunks for each page and then across pages for each domain.

## Appendix D Additional Details for Pretraining Experiments

In this section, we specify hyperparameters and methods used for LLM pretraining and evaluation for our LLM pretraining experiments. We also specify settings used for the data-selection methods.

### D.1 LLM Pretraining

We trained each LLM on 4 NVIDIA A100 GPUs. At 3.2B tokens, each training run took under 3 hours with the Hugging Face Trainer (Wolf et al., ) and appropriate PyTorch (Ansel et al., ) compile flags. We provide pretraining hyperparameters in Table 2. Given our per-device batch size, we found the learning rate by increasing it by a factor of 2 until we saw instability and then using the highest learning rate where no instability was observed. Refer to the Pythia paper (Biderman et al., ) for more information; we initialized the model from scratch using their 160M model configuration at [https://huggingface.co/EleutherAI/pythia-160m](https://huggingface.co/EleutherAI/pythia-160m). Other hyperparameters can be assumed to be Hugging Face Trainer defaults at the time of this writing.

**Table 2: LLM Pretraining Hyperparameters**
| Parameter | Value |
| --- | --- |
| Per-device Batch Size | 128 |
| Learning Rate | $5\times 10^{-3}$ |
| Warmup Ratio | 0.1 |
| Adam $\beta_{1}$ | 0.9 |
| Adam $\beta_{2}$ | 0.95 |
| Adam $\epsilon$ | $1\times 10^{-8}$ |
| Weight Decay | 0.1 |
| LR Scheduler | cosine |
| Max Grad Norm | 1.0 |
| BF 16 | True |
| Distributed Backend | nccl |
| Gradient Accumulation Steps | 1 |

### D.2 LLM Evaluation

At the end of the pretraining script, we used the Eleuther AI Eval Harness (Gao et al., ). For efficiency, we set the sample limit to 5000 examples per benchmark. Elsewhere, we used the default settings. On 4 NVIDIA A100s, it took only a few minutes per LLM to compute evaluation results for SciQ, ARC Easy, PIQA, LAMBADA, and all of the translations of LAMBADA.

### D.3 DSIR

DSIR (Xie et al., ), despite its simplicity, requires some tuning. A decision must be made about how to format the bemchmark data into a single piece of text per example so that it can be compared with potential pretraining data in terms of n-gram overlap. The LAMBADA tasks only have one text column per example, so the decision here is trivial. Examples from the other tasks each have a question, possibly a context, and a set of multiple choice answers to choose from. We chose to concatenate all of these columns together with spaces to form one piece of text per example, duplicating the same question as a prefix for each different answer.

DSIR does not allow the user to specify the exact number of unique tokens desired for pretraining. It only allows the specification of the number of unique pages, which can have wildly varying token counts. For every DSIR job, we set the desired number of pages to 3325589, which we found through binary search to produce slightly more than 3.2B unique tokens for LAMBADA${}_{\text{FR}}$. It was expensive to find this number for even one bechmark, because for each iteration of the binary search, we had to run DSIR and then the Pythia tokenizer to know how many tokens resulted from the input page number parameter. We provide the number of unique tokens from DSIR for each task in Table 3. We pretrained on 3.2B tokens for every LLM regardless of whether all of them were unique.

**Table 3: Unique pretraining tokens selected per benchmark, from DSIR.**
| Benchmark | Tokens |
| --- | --- |
| ARC Easy | 2,905,206,499 |
| PIQA | 2,910,486,295 |
| SCIQ | 2,920,734,042 |
| LAMBADA | 3,022,219,424 |
| LAMBADA${}_{\text{DE}}$ | 3,210,986,137 |
| LAMBADA${}_{\text{ES}}$ | 3,396,528,704 |
| LAMBADA${}_{\text{FR}}$ | 3,413,930,081 |
| LAMBADA${}_{\text{IT}}$ | 3,384,854,845 |

### D.4 fastText

The “SOTA” fastText model from Li et al. () is available here:
[https://huggingface.co/mlfoundations/fasttext-oh-eli5](https://huggingface.co/mlfoundations/fasttext-oh-eli5). We used this model to filter data by sorting pages by the model’s “high quality” score, including the top pages in order until we had either reached or gone slightly over 3.2B unique tokens. This aligns with the data-selection procedure in the original paper, and is also essentially the same as running the linear projection (Equation 10) at the page-level. We also applied this method when selecting data using our own fastText filter trained by our algorithm.

## Appendix E Additional Pretraining Results

In Figure 7, we present additional pretraining results for methods in our loss-performance correlation data selection paradigm. We find that using Spearman rank correlation (Spearman, ) in place of our estimator achieves comparable performance. On some tests, it performs even better than our estimator. We also find that using the quadratic projection, while perhaps more intuitive, leads to worse performance than the linear projection.

Figure: (a) Estimate with linear projection. This is our algorithm from the main text without training the additional fastText filter.
Refer to caption: extracted/5843524/images/llm-perf-heatmap-lin-proj.png

## Appendix F Pretraining Token Distribution with 5 × τ 5 𝜏 5\times\tau 5 × italic_τ

Figure 8 shows what the projected estimate in our pretraining experiments would be if we had a pretraining data pool $5\times$ as large. We see here that the estimate does an even better job at selecting pretraining data with the language that matches the target task.

Figure: Figure 8: This figure is analogous to Figure 3, except the $\tau$ thresholds have been multiplied by 5. We see that our approach selects even more relevant data when the selection pool is larger.
Refer to caption: extracted/5843524/images/token-dist-heatmap-5-tau.png

## Appendix G Parameter Count Distribution for Estimator LLMs

In Figure 9, we present the parameter-count histogram of the 90 models from the Open LLM Leaderboard (Beeching et al., ) that we used to compute our estimate for pretraining data selection. Only 8 models here are less than 160M parameters. Despite this, our estimate can be used to effectively pretrain 160M parameter LLMs.

Figure: Figure 9: The parameter-count histogram of the 90 models from the Open LLM Leaderboard (Beeching et al., ) that we used to compute our estimate for pretraining data selection. Bar widths are 160M. The smallest model in the sample has $\approx$33M parameters and the largest has $\approx$9B. The spike around $6.7$B parameters is due to a large number of partially trained Pythia (Biderman et al., ) checkpoints from the same training run at that scale. Our algorithm has the hard task of selecting pretraining data for 160M parameter models, which is abnormally small in the set of models used to compute the estimate.
Refer to caption: extracted/5843524/images/param-count-dist.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Clarification for Line 201: `select domain greatest to least \gamma, taking all the tokens in each domain once, until we exhaust our total pretraining token budget.`
Does this mean that if a domain with the greatest \gamma includes more tokens than the token budget, only tokens from that domain will be selected?

2. For the proposed perplexity correlation, have you tested aggregating benchmark scores in some way to use as an overall benchmark metric? From Figure 2, I only see LLMs where data selection (ppl correlations) is optimized for each target task.
soundness: 3
strengths: 1. This paper focuses on avoiding costly pre-training ablations for LLM pre-training data selection and proposes a method that leverages public, pre-trained LLMs based on the correlation between downstream performance and perplexity. This aim is quite important for today's LLM practitioners and the proposed framework is very original.  

2. Although limited (see below "weaknesses"), the proposed framework demonstrates promising results compared to baselines. 

3. Although the presentation can be improved, the paper includes a theoretical background for the proposed framework.
summary: 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.
weaknesses: 1. The main weakness of this paper is the scale of the experiments. The paper only includes pre-training experiments using 160M parameters LLM with a 3.2B token budget. This is very small and creates questions about the validity of the results. For example, the smallest scale for DataComp-LM paper is 400M and 8.2B and they experimented with 5 different scales which goes up to 7B parameters model trained on 276B tokens (as DCLM competition scale).

2.  Experiments only include 8 evaluation benchmarks (5 of them are Lambada from different languages). This is a main limitation for both showcasing the validity of results but also limit the understanding how the proposed framework. work when there are many more benchmarks to consider for perplexity validation. Again for comparison, DataComp-LM includes an extended evaluation suite that consists of up to 53 downstream tasks. 

3. If the framework only optimizes data selection per task, is it not expected to have good scores on optimized tasks? Are there any experiments where authors tested aggregating benchmark scores as one overall metric?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - Relevant works that are not cited:

[1] Niklas Muennighoff, Alexander M. Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Nouamane
Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models, 2023

[2] Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang, Marzieh Fadaee, Sara Hooker. When Less is More: Investigating Data Pruning for Pretraining LLMs at Scale, 2023
soundness: 3
strengths: The paper presents a new and remarkably simple approach to data selection for LLM pretraining. It leverages readily available resources (publicly available LLMs) and a straightforward correlation-based method.  The authors ground their approach in a statistical framework, providing a theoretical basis for their correlation-based data selection. The authors demonstrate the robustness of their approach by showing its effectiveness across various benchmarks and conditions.
summary: 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.
weaknesses: While the paper presents promising results at the 160M parameter scale, further validation at larger scales is necessary to assess the scalability and effectiveness of the approach for slightly more massive LLMs. The method inherits the biases present in the publicly available LLMs used for estimating correlations. This raises potential concerns about bias amplification, especially if these LLMs were trained on data with inherent biases.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. It seems that factors such as benchmark contamination or formatting overlap between pretraining data and perplexity validation data could alter the relationship between domain perplexity and downstream performance for some observed models. This would introduce more structured noise than simple run-to-run variance due to differing training seeds. In the observed data do correlations vary dramatically for any subsets of models (i.e. are there subsets of models perhaps with systematically biased perplexity to downstream correlations)? Does your approach mitigate this possibility in any way other than by accumulating enough diverse observations that it’s unlikely a substantial number of models will be affected by the same biases?
2. Is the projection method used in algorithm 1 when selecting data or only when predicting downstream performance ranking? Figure 6 implies that the data selection is different when using different projection methods but algorithm 1 doesn’t mention the projection step. 
3.  What does “features” refer to on line 344?
4. Line 460: what is meant by “​​significant enrichment of the corresponding languages”
5. Since the domain perplexities are measured over samples of 25 webpages, do the number of tokens used to estimate perplexity vary considerably from domain to domain? This may add noise to your perplexity estimates proportional to how severely some domains are subsetted with respect to tokens as observed by https://arxiv.org/abs/2312.10523 in their subsampling guideline.
6. How are the standard deviations computed for the figure 4 ranking plots?

Other comments (not questions and no response needed; wasn't sure where to put them):

The paper grounds the method in some theory by relating it to estimators for single index models.  This is not extremely important from my point of view, because the method is fairly intuitive, but it may give further intuition or justification to some readers.  The value may be in revealing the assumptions made by the model (through the lens of statistical estimation), so that they can be revised in future explorations.  I thought the discussion of the projections (4.2.3) was mostly unnecessary; I would have rather seen the experimental comparison of the different variants reported in the main paper and the derivations in the appendix, rather than the reverse.

Suggestions (no response needed):

It would be interesting to see more analysis of specifically what domains are selected by the proposed approach contrasted with other baselines. Likewise how do the domains selected differ for different target downstream benchmarks (other than in terms of language)? Also are the most correlated domains also domains with few or many tokens? That is, does the final data selection only include a small subset of the positive correlation domains before satisfying the token budget or are many smaller domains selected?

Line 93: There is some work which uses pretraining experiments to explore high dimensional mixtures of pretraining data (https://arxiv.org/abs/2410.15661).

Line 97: There is work that directly shows the relationship between loss to downstream performance (https://arxiv.org/abs/2403.08540), though this work does not propose a data selection method.

In line 125 and 128 j switches from indexing documents and domains. Maybe use different indices to make this clearer.

On line 208 and 337 “single-page level” was confusing to me until I recalled that you are working with purely webtext and thus documents are webpages. It would be best to define this term before use to avoid possible confusion with other meanings of the word page.

Line 458: avoid use of “significantly” when not discussing statistical significance.

Line 526: “usual” should be “unusual” here, right?

In 4.2.1, I think the epsilons, not x, should be drawn from a standard normal.

Figure 2 doesn’t clearly label methods the way they are referred to in the text.  It would be helpful to clearly mark, for example, the state-of-the-art methods from the DCLM paper as such and “our” methods (those introduced in this paper).
soundness: 4
strengths: This approach provides a promising way to make use of “the millions of dollars collectively spent” (line 39) on the experimentation represented by all open weight models. The approach extracts information about performant pretraining data selection even when information about the pretraining data of these open-weight models is not available.

Their approach is intuitive and the core of it is quite simple (appearing to be insensitive to the trickier details suggested by theory). This means it should be straightforward for others to apply it.

The authors find that their approach achieves comparable results to a SOTA quality classifier that utilized extensive, costly pretraining experiments, while crucially the authors’ approach requires no new pretraining experiments, unlike other methods..
The paper provides a sound empirical validation of their approach at the 160M parameter scale. They use downstream evaluations that are appropriate for the parameter scale of their experiments (even including some basic multilingual evaluations). When optimizing for a single benchmark, the approach outperforms the state-of-the-art data selection method from DCLM if the latter is handicapped by removing the “manual filtering” step, but not if that manual filtering is included. Their set up demonstrates the utility of their approach for data with a very large number of domains (far greater than can be optimized with brute force pretraining experiments). These results should be more than enough to motivate further exploration of this technique at larger scales.

I appreciated the additional analysis of outliers where the approach doesn’t do so well at predicting relative performance.  The observation that models with “unusual” data mixtures is interesting and will be useful in future work in this space.
summary: 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).
weaknesses: I think the biggest weakness is that the paper’s method targets a single benchmark.  Are there obvious ways to extend the work to target some aggregation of benchmark scores?  Sure, but the paper offers no theory to suggest how we might do that, and no experiments to show whether the obvious ways of doing that would work.
Lambada, which accounts for 5 out of 8 of downstream evaluations, is a language modeling task (i.e. next word prediction) just like the perplexity measures being used as proxies in the proposed approach. This may somewhat obscure whether the proposed approach works as well for all kinds of downstream tasks which typically are not language modeling tasks. That said, the other 3 non-lambada tasks do also benefit from the proposed approach.
It is unclear in the empirical results of the paper what parts of the theoretical results are actually necessary for the evident success of the authors’ approach. The results in Appendix G do not conclusively favor the author’s proposed correlation measure nor their projection method. While the authors do say so themselves (line 467), it would be kinder to the reader to let the reader know this before they have waded through the dense theoretical sections of the paper. 
The presentation of the paper also makes it unclear what theoretical results are necessary as the discussion of theory and methods are quite entangled. The paper would be much more readable if there was a simple discussion of what the proposed approach does and then only after that provide theoretical arguments for why this approach is used. I believe Algorithm 1 in the appendix provides the clearest end-to-end description of what exactly the proposed approach is. But even there I am left uncertain of details not mentioned there such as the theta projection or the use of empirical CDF. Is theta projection always used for the greedy data selection, or can the raw gamma values be used without projection since all you need is the order of the domains by the correlation?
I believe the experiments work only at the internet domain granularity for selection.  The paper actually says very little about what a “domain” is – it seems to follow from variables defined in the redpajama data, and there are around 9000 “domains,” so the problem is of interesting size.  But I didn’t come away with a strong sense that – beyond empirical reports – there really should be a relationship between these particular “domains” and model performance on the benchmarks.  The paper would be stronger with some discussion of the qualitative assumptions that lead us to believe perplexity on specific slices of pretraining data (whether defined by internet domains or some other attributes) should be predictive of task performance. Are those assumptions grounded in any prior work?  Are they pervasive?  Are they novel?

The largest model on which the method is tested is only 160M parameters; it’s not clear whether the method will scale.  (I have no reason to believe it won’t, or that it will.)  This scale of model also implies a choice of target benchmarks that deviates pretty far from metrics most builders of language models would care about; I understand that the experiments needed to use benchmarks that would show meaningful differences at the 160M scale, but this creates even more distance between what the paper shows and what is needed in practical settings.

Apart from the most obvious case (based on language ID), the paper really doesn’t delve into what the predictive model itself tells us about the relationship between pretraining data composition and the benchmarks.  This is in part due to the unexciting set of benchmarks, I think.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - How can the proposed method be extended to selecting data for multiple downstream tasks, including unseen and/or dynamic benchmarks?
- Will the data domains selected using models of a certain scale (e.g., < 8B) be helpful for training larger models (e.g., 70B)?  
- How would varying the token ratio of pre-training domains affect the downstream performance? That is, mixing domains beyond "selecting all the tokens in each domain once" until token budget is exhausted.
soundness: 3
strengths: - This paper proposes a new framework for pre-training data selection without extensive manual labor by leveraging large amounts of pre-trained open-weight models.
- The results validated on small 160M parameter models look promising.
- The studied topic is important and relevant to many language model practitioners.
summary: 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.
weaknesses: - It is unclear how well the proposed framework works for larger scale pre-training where multiple static benchmarks should be considered, and when the model size can be larger than the maximum model size (~8B) used for computing perplexity-benchmark correlation.
- The theory section is a bit difficult to follow. In the paragraph above section 5.2 "*This suggests that the performance of our approach is not dependent on the precise form of the estimator that is coupled to our theory results, but holds broadly across perplexity-correlation relationships.*" If the theoretical result is less universal than empirical observation, it can be interesting to add more discussion on the more general relationship.
- The proposed model provides the ordering of pre-training domains, however, from what I understand, the sampling ratio across domains is not determined by the benchmark-perplexity correlation, but rather “*select domains in greatest to least γ, taking all the tokens in each domain once, until we exhaust our total pretraining token budget.*” The discussion about the exact ratio to mix various domains is a little limited.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: N/A
soundness: 2
strengths: 1. Using perplexity correlations for data selection is well-motivated and easy to implement.

2. The paper provides a solid theoretical foundation for the proposed method, including a detailed discussion of the single-index model and the estimation of correlation coefficients. The mathematical derivations and proofs are rigorous and well-presented.

3. In experiments, authors demonstrate that their method outperforms existing data selection techniques in terms of downstream benchmark performance.
summary: 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.
weaknesses: 1. While the paper demonstrates the effectiveness of the approach at the 160M parameter scale (though the data selection parameters are generated with LLM of 7B),  how the method would scale to larger models' pertaining remains unclear. The extrapolation from smaller models to larger ones could be a potential limitation.

Meta-review: 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.

Decision: Accept (Poster)