Improving Pretraining Data Using Perplexity Correlations
Under review as a conference paper at ICLR 2025

IMPROVING PRETRAINING DATA USING
PERPLEXITY CORRELATIONS

Anonymous authors

Paper under double-blind review

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 with-
out 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 ex-
periments at the 160M parameter scale on 8 benchmarks, our approach outper-
forms 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., 2020; Gao et al.,
2020) to 240T tokens today (Li et al., 2024), 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., 2022; Xie et al., 2023a;b; Engstrom et al., 2024; Everaert & Potts, 2024; Liu et al.,
2024; Llama Team, 2024). 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., 2024).

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., 2019; Beeching
et al., 2023) 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., 2022; Huang et al., 2024), 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.

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

054
055

056
057
058
059
060
061

062
063
064
065
066
067

068
069
070
071
072

073
074
075
076
077
078

079
080
081
082
083
084

085
086
087
088
089
090

091
092
093
094
095

096
097
098
099
100
101

102
103
104
105
106
107

Domains

Benchmark

bbc arxiv · · · willys-hifi

SciQ

s
M
L
L

Mistral

Llama

Mamba

...

...

...

Pythia

· · ·

· · ·

· · ·
. . .
· · ·

...

...

logprob

accuracy

Correlations

bbc arxiv · · · willys-hifi

· · ·

High Corr (Keep)
arxiv, bbc, · · ·

Low Corr (Discard)
willys-hifi, · · ·

Figure 1: We 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 training a fastText classifier that distinguishes high cor-
relation 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.

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

2 RELATED WORK

To go beyond the status quo of deduplication, perplexity filtering, and hand-curation (Laurençon
et al., 2022; BigScience, 2023; Abbas et al., 2023; Groeneveld et al., 2024; Soldaini et al., 2024;
Penedo et al., 2024; Llama Team, 2024), targeted methods have been proposed to filter pretrain-
ing 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., 2023b) or embedding similarity (Ever-
aert & Potts, 2024) 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.,
2022; Xie et al., 2023a; Engstrom et al., 2024; Liu et al., 2024; Llama Team, 2024).

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, 2024; Li et al., 2024) 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, 2024; Ruan et al., 2024). 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., 2020; Hoffmann et al., 2022; Wei et al., 2022), even if they
have different architectures (Poli et al., 2023; Peng et al., 2023; Gu & Dao, 2024).

According to a recent survey of data selection approaches by Li et al. (2024), the heavier-weight
pretraining data selection methods have not shown large gains, and the current state-of-the-art across
many tasks is primitive: a fixed fastText classifier (Joulin et al., 2016) 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 evidence supporting this hypothesis –
simple loss-performance correlation coefficients are effective when used for data selection.

2

Under review as a conference paper at ICLR 2025

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

3 PROBLEM SETTING

Our goal is to build predictive models of how pretraining data distributions affect downstream bench-
mark performance and use them to build better language models. Unfortunately, this task is challeng-
ing and computationally expensive. A standard approach adopted in paradigms such as datamodel-
ing (Ilyas et al., 2022) is to obtain N different pretraining distributions {pi : i ∈ [N ], pi ∈ R+
}
0
over D ≫ N domains (e.g. arxiv.org, stackoverflow.com, etc.), pretrain and measure model errors
on a target benchmark yi ∈ [0, 1], and fit a model p → 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., 2023).

D

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 p →
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 p.

The key observation of our work is that we can replace pi,j (the unobserved sampling probability of
model i’s data selection policy on document j) with an observable surrogate xi,j, which is the nega-
tive log-likelihood of document j under model i.1 We can then build a regression model that relates
negative log-likelihood xi and benchmark error yi. Using this model, we can select pretraining data
from domains j for which decreasing the loss xi,j is predicted to rapidly decrease error yi.

The perplexity-performance hypothesis. We formulate the task of predicting errors yi from nega-
tive log-probabilities xi as a single-index model (SIM),

yi = f (⟨θ∗, xi⟩ + ϵi)
(1)
where f : R (cid:55)→ R is some unknown monotonically increasing univariate function, ϵi is zero-mean
noise which is independent of x, and θ∗ ∈ RD 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

⟨θ∗, xi⟩ + ϵi < ⟨θ∗, xj⟩ + ϵj ⇐⇒ f (⟨θ∗, xi⟩ + ϵi) < f (⟨θ∗, xj⟩ + ϵj).

(2)

Data selection from perplexity correlations. The weights θ∗ 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 selec-
tion 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 θ∗, sampling proportional to θ∗ 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 θ∗ weights are non-negative. Then, for models with associated like-
lihoods x ∈ X ⊂ RD, the minimizer of the pretraining loss over the θ∗ sampling distribution
Ej∼θ∗ [xj] also has the lowest expected downstream error according to the single index model:

arg min
x∈X

Ej∼θ∗ [xj] = arg min

x∈X

E[f (⟨θ∗, x⟩ + ϵ)].

This observation follows directly from the fact that we can normalize any non-negative θ∗ 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 (⟨θ∗, x⟩ + ϵ) = f (Ej∼θ∗[xj] + ϵ).

(3)

1To 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 LT , the length of the string in
UTF-8 bytes LB, and the cross entropy loss ℓ as BPB = LT ℓ

LB ln(2)

3

Under review as a conference paper at ICLR 2025

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

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 func-
tion 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 θ∗ estimates to be a pretraining data sampling distribution.

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 θ∗. The parameter θ∗
j and downstream performance. Because of this, we might naturally expect θ∗
nonlinear correlation coefficients between x and y. Our work uses a simple correlation measure,

j measures the relationship between log-likelihoods in domain
j to be related to

(cid:88)

γj =

1≤k,l≤n
k̸=l

sign(yk − yl)(rankj(xk,j) − rankj(xl,j))

where rankj(x) is the rank of x among {x1,j . . . xN,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 G), this functional form
has the additional benefit of being a principled estimate of θ∗. In particular, we show in sections
below that in expectation, the ranking of domains in γ exactly matches those of θ∗ (under standard
high-dimensional regression assumptions; see Section 4.2 for a complete discussion).

Selecting pretraining data. Suppose that we have an accurate estimate γj which is nonnegative. In
this case, we could use γj directly as a data selection procedure and Proposition 1 would ensure that
minimizing the population pretraining loss minimizes downstream errors. Unfortunately, γ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 γ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 γ into a pretraining sampling
distribution? Intuitively, if wikipedia.org has a γj = 0.5 and arxiv.org is γk = 0.9, it would be nat-
ural to select tokens in order of γ, preferring tokens from arxiv.org over tokens from wikipedia.org.

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.,
2023) to arrive at a simple selection algorithm: select domains in greatest to least γ, 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 with pseudocode in Algorithm 1 (see Appendix A), and additionally
show an extra step where we train a fastText (Joulin et al., 2016) classifier (using standard settings
and bigram features from Li et al. (2024)) 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 F.

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. (2016). We describe the basic single-index model estimators first, describe our extensions,

4

Under review as a conference paper at ICLR 2025

and then conclude with a discussion on how our estimator and results deviate from the theory. A
discussion of other potential estimation paradigms is provided in Appendix D.

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. (2016)
and Chen & Banerjee (2017). Here, our goal is to estimate the unknown weights θ∗ in a single-index
model yi = f (⟨θ∗, xi⟩ + ϵi), with xi ∼ N (0, I) for ∥θ∗∥2 = 1 (assumed without loss of generality,
as ∥θ∗∥2 can be absorbed by f ).

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

E [ykxk] = cθ∗,
(4)
for some positive constant c and 1 ≤ k ≤ N . Closely related is the result of Chen & Banerjee
(2017) who showed a robust estimator quite similar to ours,

E [sign(yk − yl)(xk − xl)] = βθ∗
(5)
for any 1 ≤ k, l ≤ N (where k ̸= l) and some positive constant β. Both of these results clearly iden-
tify that for the high-dimensional single-index model in the Gaussian setting, generalized correlation
coefficients provide consistent estimates of the true regression coefficient θ∗.

4.2.2 DERIVING OUR ESTIMATOR

Both Plan et al. and Chen & Banerjee provide moment-matching style estimators that consistently
recover θ∗ 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 (2017) to outliers in x.

Recall that our estimate γ is a U-statistic, defined as pairwise sums of

sign(yi − yj)(Φ(xi) − Φ(xj)),
(6)
for any 1 ≤ i, j ≤ N (where i ̸= j), where Φ is the empirical CDF of the x values. This estimate is
significantly less sensitive to outliers than that of Chen & Banerjee (2017), 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 Φ as the CDF of the standard Gaussian. In this case, we can show that this
modified estimator is also consistent in recovering θ∗.

Theorem 1 When ϵ ∼ N (0, σ2), we have:

E[sign(yi − yj)(Φ(xi) − Φ(xj))] =

2
π

sin−1

(cid:18)

θ∗
1 + σ2

√

2

(cid:19)

.

(7)

We provide the proof in Appendix B. Because we assume ||θ∗||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:

ˆθ ∝ sin

(cid:16) π
2

√

(cid:17)
E [sign(yi − yj)(Φ(xi) − Φ(xj))]

(8)

as an estimate for θ∗, where the constant 2

1 + σ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 Chen
& Banerjee estimator: our estimates agree when running the original estimator on rank-transformed
data. More specifically, for two models xi and xj with the estimated model rankings ⟨ ˆθ, xi⟩ >
⟨ ˆθ, xj⟩, the expected ranking under rank-transformation (i.e. Φ(x)) match this ranking.

Corollary 1 Suppose that ˆθ is any vector of fixed weights and x ∼ N (0, I). Then, conditioning on
the event ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩, we have with probability 1 that:

⟨ ˆθ, E[Φ(xi) | ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩]⟩ < ⟨ ˆθ, E[Φ(xj) | ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩]⟩.

(9)

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

5

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

Under review as a conference paper at ICLR 2025

4.2.3 SELECTING DATA FOR PRETRAINING

Recall that our algorithm for data selection is to constrain γ to be a valid sampling distribution
(nonnegative, at the very least) and then sample directly from this estimate. For now, we focus
on constraining ˆθ, and we will see at the end of this section that we can apply the same con-
straint to γ directly to get the same result. The theory of constrained estimation for ˆθ is simple
and well-understood, with both Plan et al. (2016) and Chen & Banerjee (2017) extensively study-
ing the problem of estimating ˆθ under a known convex constraint set C. In particular, Plan et al.
(2016) show that performing a L2 projection via ˆθproj = arg minθ∈C ∥θ − ˆθ∥2 provides improved
convergence rates that depend on the Gaussian mean width of C rather than the ambient dimen-
sion, and Chen & Banerjee (2017) show similar results when maximizing the linear correlation
ˆθproj = arg minθ∈C⊆BD −⟨θ, ˆθ⟩.
We take a similar approach here. We define a convex constraint set C that forces ˆθ 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 ˆθ to lie in the simplex. As we noted above, it is well-known that dupli-
cating data harms performance (Abbas et al., 2023), and so we constrain ˆθ to avoid data duplication
by limiting the maximum weight on domains. Concretely, if want to pretrain on m tokens overall,
we enforce θ∗
i ≤ τi, ∀i ∈ [1, D], where τi is set so τim 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 ˆθproj to 0 and
then iterating through the values in ˆθ from largest to smallest, setting the value at the corresponding
index of ˆθproj to the maximum allowable value, until ˆθproj sums to 1 (see Appendix C for a proof).

Theorem 2 Suppose we want to solve:

subject to:

ˆθproj = arg min
θ∈RD

−⟨θ, ˆθ⟩,

D
(cid:88)

i=1

θi = 1

0 ≤ θi ≤ τi, ∀i ∈ [1, D],

where τi > 0 are fixed values. Then, the solution is:

ˆθproj
k =






τk
1 − (cid:80)
0

j: rj (ˆθj )>rk(ˆθk) τj

if (cid:80)
if (cid:80)
otherwise

j: rj (ˆθj )≥rk(ˆθk) τj ≤ 1
j: rj (ˆθj )≥rk(ˆθk) τj ≥ 1 ∧ (cid:80)

j: rj (ˆθj )>rk(ˆθk) τj ≤ 1

, (10)

where r is some function that breaks all ties between ˆθj and ˆθk for k ̸= 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 (2017), the L2 projection is arguably more natural, and does not require assum-
ing that ∥ ˆθ∥2 = 1 for asymptotic recovery conditions. We derive similar closed-form expressions
for this quadratic case in Appendix C, but do not use this approach for two separate reasons.
First, the L2 projection depends on the L2 norm of ˆθ, unlike the linear program which only depends
on the ranks of the values in ˆθ. 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 ˆθ 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 G).

6

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

Under review as a conference paper at ICLR 2025

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

We conclude by relating our theory to the full algorithm in Section 4.1. The estimation step for γ
is the finite sample, U-estimate of the expectation in Equation (8), dropping the nonlinear transform
sin and π/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 γ rather than an exact estimate for θ∗.

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. 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, X, we collected byte normalized loss values
from a sample of 90 Open LLM Leaderboard (Beeching et al., 2023) LLMs that we could run
LB ln(2) where LT is the token
without errors. Concretely, these values are defined as bits-per-byte
count, LB is the number of UTF-8 bytes, and ℓ is the per-token cross-entropy (Gao et al., 2020).
We collected these values on “sample” subset2 of the RedPajama V2 (RPJv2) dataset (Together
Computer, 2023) for all domains with ≥ 25 pages in the sample. There are 9,841 domains/features.
Specifics are in Appendix E. A detailed principal components analysis of X, which reveals a variety
of salient embedded information in the losses, is in Appendix J.

LT ℓ

Target benchmark performance y. We constructed a target vector, y, for LAMBADA (Paperno
et al., 2016), ARC Easy (Clark et al., 2018), PIQA (Bisk et al., 2020), and SciQ (Welbl et al.,
2017). 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 LAMBADAIT, LAMBADAFR, LAMBADADE, and LAMBADAES, which are subsets of
LAMBADA translated into Italian, French, German, and Spanish by Black (2023). 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 is convenient, as it allows us to easily include language filtering baselines.

5.1 PRETRAINING

We begin by validating our algorithm in the end-to-end task of pretraining data selection with con-
trolled experiments at the 160M parameter, 3.2B token scale. The low compute requirements of
this setting allow us to more extensively study replicates and ablations in Appendix G within the
timeframe of a few days. 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× larger than the
160M scale, and our single index model must extrapolate substantially from ≈7B scale models to
our small-scale validation setting (see Appendix I for a histogram of model sizes).

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., 2022) 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., 2023b): a lightweight training data selection technique based
on n-gram overlaps that Li et al. (2024) found to be competitive with proxy LLM-based techniques

2https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2

7

Under review as a conference paper at ICLR 2025

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. (2024). Our approach significantly beats the default filter in Li et al.
(2024) with the EN filter and loses slightly after additional manual language filtering that depends
on the target task (+ manual Lang Filter).

Method None

Lang
Filt

DSIR
(Xie et al., 2023b)

Handcrafted fastText
+ EN Lang Filter
(Li et al., 2024)

Handcrafted fastText
w/o Lang Filter

Handcrafted fastText
+ manual Lang Filter

Perplexity
Correlations

Avg. Rank

3.750

4.000

4.500

3.750

3.250

1.375

1.750

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.

and was also validated at scale (Parmar et al., 2024). Finally, we run the state-of-the-art method
for pretraining data quality filtering found by Li et al., which is a fastText classifier that beats all of
the heavier-weight proxy-LLM methods tested. The classifier was trained on a benchmark-agnostic
and handcrafted objective, which is to classify data as Common Crawl3 (low quality) or OH2.5
(Teknium, 2023) and Reddit ELI5 (Fan et al., 2019) (high quality). It is combined with an English
filter in Li et al.; we present results for this fastText filter with and without the English filter.

Model and hyperparameters. We use the Pythia 160M LLM configuration from Biderman et al.
(2023) 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 F.

3https://commoncrawl.org

8

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

Under review as a conference paper at ICLR 2025

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

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).

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. (2024), 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 G.

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 bench-
mark: within each block, the benchmark column matching each row typically performs best. The
pattern is much less obvious for DSIR – the heatmap looks 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 do-
mains 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 τ constraint to prevent duplication has a large impact on the weights when the benchmarks are
non-English. We provide the same figure when the τ values are made 5× as large in Appendix H.

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 G that varying the projection method
(linear, L2) and even using Spearman rank correlations (Spearman, 1904) often work better than the
baselines. 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. Additionally, our approach performs better with the optional fastText classifier that
our algorithm trains, possibly because it operates at the page-level instead of the domain-level

5.2 PERFORMANCE RANK PREDICTIONS

We have shown that our approach succeeds at selecting useful pretraining data, but how good are the
single index model’s predictions? A good map of loss to benchmarks would be helpful in selecting
among candidate pretraining data mixtures generally, even without using our specific 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 LAMBADAFR
with the rank predictions given by ⟨ ˆθproj, Φ(x)⟩. Every point in the plot is a held-out point: we
estimated θ∗ 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 R2 scores for all tasks in Figure 5. However, we complement these strong

9

Under review as a conference paper at ICLR 2025

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

Figure 4: Rank predictions given by ⟨ ˆθproj, Φ(x)⟩ for PIQA and LAMBADA FR. A standard devia-
tion (σ) from the ideal fit is shown in red. 2σ is shown in orange. Many models outside 2σ (shown
in blue) are trained on atypical data such as multilingual data, code, or GPT-4 (Brown et al., 2020)
outputs. Models with atypical architectures (i.e. Mamba (Gu & Dao, 2024)) are shown in black.
Generally, our estimate tightly predicts ordinal benchmark performance from web corpus losses.

Figure 5: Held-out R2 score of our raw correlation estimate ˆθ, our projected estimate ˆθproj, and
the average loss baseline. The 95% bootstrapped confidence intervals are wide enough that no
individual comparison is significant. Across benchmarks, ˆθproj has statistically significant gains
over the baseline (p=0.035) as it is unlikely that ˆθproj beats mean loss 7 times out of 8 by chance.

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, 2024; Wei et al., 2022;
Kaplan et al., 2020) 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, 2024)), 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., 2023)). 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.

6 CONCLUSION

Does high-performance data selection require careful hand-crafted heuristics or prohibitively ex-
pensive model training runs? Our work demonstrates an alternative, viable approach – leveraging
existing, public models as a source of information for data selection. Pretraining experiments sug-
gest 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.

10

Under review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

REFERENCES

Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S. Morcos. Semdedup: Data-

efficient learning at web-scale through semantic deduplication. arXiv, 2023.

Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky,
Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will
Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael
Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael La-
zos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yunjie Pan,
Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael
Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren
Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2:
Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Com-
pilation. ACM International Conference on Architectural Support for Programming Languages
and Operating Systems, 2024.

Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge,
Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu,
Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren, Xuancheng Ren, Chuanqi
Tan, Sinan Tan, Jianhong Tu, Peng Wang, Shijie Wang, Wei Wang, Shengguang Wu, Benfeng
Xu, Jin Xu, An Yang, Hao Yang, Jian Yang, Shusheng Yang, Yang Yao, Bowen Yu, Hongyi Yuan,
Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou,
Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report. arXiv, 2023.

Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Ra-
jani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf, 2023. URL https://huggingface.
co/spaces/HuggingFaceH4/open_llm_leaderboard. Open LLM Leaderboard.

Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O’Brien, Eric Hal-
lahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya
Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language
models across training and scaling. arXiv, 2023.

BigScience. BLOOM: A 176b-parameter open-access multilingual language model. arXiv, 2023.

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about

physical commonsense in natural language. AAAI, 2020.

Sid Black, 2023. URL https://huggingface.co/datasets/EleutherAI/lambada_openai.

Multilingual LAMBADA.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari-
wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal,
Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M.
Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin,
Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford,
Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv, 2020.

Sheng Chen and Arindam Banerjee. Robust structured estimation with single-index models. ICML,

2017.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and
Oyvind Tafjord. Think you have solved question answering? Try ARC, the AI2 reasoning chal-
lenge. arXiv, 2018.

John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto

the l1-ball for learning in high dimensions. ICML, 2008.

Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection

with datamodels. arXiv, 2024.

11

Under review as a conference paper at ICLR 2025

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

Dante Everaert and Christopher Potts. Gio: Gradient information optimization for training dataset

selection. ICLR, 2024.

Angela Fan, Yacine Jernite, Ethan Perez, David Grangier, Jason Weston, and Michael Auli. ELI5:

long form question answering. arXiv, 2019.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason
Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile:
An 800gb dataset of diverse text for language modeling. arXiv, 2020.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos-
ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muen-
nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang
Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for
few-shot language model evaluation. Zenodo, 2023.

Xinyang Geng and Hao Liu. Openllama: An open reproduction of llama, 2023. URL https:

//github.com/openlm-research/open_llama.

Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord,
Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, Shane Arora, David Atkin-
son, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar,
Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoff,
Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander,
Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Worts-
man, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle
Lo, Luca Soldaini, Noah A. Smith, and Hannaneh Hajishirzi. Olmo: Accelerating the science of
language models. arXiv, 2024.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv,

2024.

Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth
Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital
Shah, Harkirat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai,
Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need. arXiv, 2023.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza
Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen-
nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy,
Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre.
Training compute-optimal large language models. arXiv, 2022.

Yuzhen Huang, Jinghan Zhang, Zifei Shan, and Junxian He. Compression represents intelligence

linearly. COLM, 2024.

Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Data-

models: Predicting predictions from training data. ICML, 2022.

Armand Joulin, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov. Bag of tricks for efficient

text classification. arXiv, 2016.

Adam Tauman Kalai and Ravi Sastry. The isotron algorithm: High-dimensional isotonic regression.

COLT, 2009.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child,
Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language
models. arXiv, 2020.

Hugo Laurençon, Lucile Saulnier, Thomas Wang, Christopher Akiki, Albert Villanova del Moral,
Teven Le Scao, Leandro Von Werra, Chenghao Mou, Eduardo González Ponferrada, Huu Nguyen,
Jörg Frohberg, Mario Šaško, Quentin Lhoest, Angelina McMillan-Major, Gerard Dupont, Stella
Biderman, Anna Rogers, Loubna Ben allal, Francesco De Toni, Giada Pistilli, Olivier Nguyen,

12

Under review as a conference paper at ICLR 2025

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

Somaieh Nikpoor, Maraim Masoud, Pierre Colombo, Javier de la Rosa, Paulo Villegas, Tristan
Thrush, Shayne Longpre, Sebastian Nagel, Leon Weber, Manuel Muñoz, Jian Zhu, Daniel Van
Strien, Zaid Alyafeai, Khalid Almubarak, Minh Chien Vu, Itziar Gonzalez-Dios, Aitor Soroa,
Kyle Lo, Manan Dey, Pedro Ortiz Suarez, Aaron Gokaslan, Shamik Bose, David Adelani, Long
Phan, Hieu Tran, Ian Yu, Suhas Pai, Jenny Chim, Violette Lepercq, Suzana Ilic, Margaret
Mitchell, Sasha Alexandra Luccioni, and Yacine Jernite. The bigscience roots corpus: A 1.6tb
composite multilingual dataset. NeurIPS Datasets and Benchmarks, 2022.

Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash
Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel,
Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bit-
ton, Marianna Nezhurina, Amro Abbas, Cheng-Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej
Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras,
Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic,
Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer,
Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groen-
eveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair
Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-lm: In search of the
next generation of training sets for language models. arXiv, 2024.

Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing
Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training. arXiv,
2024.

Llama Team. The llama 3 herd of models. arXiv, 2024.

Edward W. Ng and Murray Geller. A table of integrals of the error functions. Journal of Research

of the Natianal Bureau of Standards, Section B: Mathematical Sciences, 1968.

David Owen. How predictable is language model benchmark performance? arXiv, 2024.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi,
Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset:
Word prediction requiring a broad discourse context. ACL, 2016.

Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Bo Liu, Aastha Jhunjhunwala, Zhilin
Wang, Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Data, data everywhere:
A guide for pretraining dataset construction. arXiv, 2024.

Karl Pearson. On lines and planes of closest fit to systems of points in space. Philosophical Maga-

zine, 1901.

Guilherme Penedo, Hynek Kydlíˇcek, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin
Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the
finest text data at scale. arXiv, 2024.

Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman,
Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen
Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden
Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang,
Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang
Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing rnns for the
transformer era. arXiv, 2023.

Yaniv Plan, Roman Vershynin, and Elena Yudovina. High-dimensional estimation with geometric

constraints. Information and Inference: A Journal of the IMA, 2016.

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua
Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional
language models. arXiv, 2023.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi
Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text
transformer. Journal of Machine Learning Research, 1–67, 2020.

13

Under review as a conference paper at ICLR 2025

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

Yangjun Ruan, Chris J. Maddison, and Tatsunori Hashimoto. Observational scaling laws and the

predictability of language model performance. arXiv, 2024.

Shai Shalev-Shwartz and Yoram Singer. Efficient learning of label ranking by soft projections onto

polyhedra. JMLR, 2006.

Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur,
Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Harsh
Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas
Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Abhilasha Ravichander, Kyle
Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Pete Walsh, Luke
Zettlemoyer, Noah A. Smith, Hannaneh Hajishirzi, Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and
Kyle Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research.
arXiv, 2024.

Charles Spearman. The Proof and Measurement of Association between Two Things. The American

Journal of Psychology, 1904.

Teknium. Openhermes 2.5: An open dataset of synthetic data for generalist llm assistants, 2023.

URL https://huggingface.co/datasets/teknium/OpenHermes-2.5.

Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical

Society Series B: Statistical Methodology, 58(1):267–288, 1996.

Together Computer, 2023. URL https://github.com/togethercomputer/RedPajama-Data.

RedPajama: an Open Dataset for Training Large Language Models.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko-
lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher,
Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy
Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn,
Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel
Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee,
Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra,
Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi,
Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh
Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen
Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic,
Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models.
arXiv, 2023.

Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. JMLR, 2008.

Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo-
gatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol
Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models.
TMLR, 2022.

Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions.

W-NUT, 2017.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi,
Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface’s trans-
formers: State-of-the-art natural language processing. arXiv, 2019.

Jeffrey M. Wooldridge. Econometric Analysis of Cross Section and Panel Data. MIT Press, 2010.

Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang,
Quoc V. Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up
language model pretraining. NeurIPS, 2023a.

Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language

models via importance resampling. NeurIPS, 2023b.

14

Under review as a conference paper at ICLR 2025

A MAIN ALGORITHM

Algorithm 1 Perplexity Correlation Based Data Selection

N ×D

, available tokens per domain a ∈ ND, and pretraining token target b ∈ N.

Input: Benchmark error vector y ∈ [0, 1]N , log-loss matrix normalized as bits-per-byte X ∈
R+
0
Output: Target token counts per domain t ∈ ND
Initialize: γ ← 0 ∈ RD, t ← [0 . . .] ∈ ND
r0, r1, . . . , rN ← rank(x0, x1, . . . , xN )
for i, j ∈ 0 to N do

0 , a fastText classifier to filter pretraining data.

▷ 1. Compute the γ correlation coefficient

0 , counter ← 0.

γ ← γ + sign(yi − yj) · (ri − rj)
for i ∈ArgSort(γ, descending=True) do

ti ← min(ai, b − counter)
counter ← counter + ai
if counter ≥ b then

▷ 2. Select most to least correlated domains

Break

classifier = trainFastText(positive = 1t>0, negative = 1t=0)
Return t, classifier

B ESTIMATOR SOLUTION

B.1 LEMMA 1

Statement of Lemma 1 Define the PDF of HalfNormal as f (x; σ) =
otherwise. Now, suppose:

• β is a vector with ||β||2 = 1

• Z1, Z2 are vectors ∼ N (0, I)

• ϵ ∼ N (0, σ2)

• Z ′ ∼ N (0, 1)

• Z+ ∼ HalfNormal(1).

Then we have:

√
2
√

π e− x2

σ

2σ2 for x > 0 and 0

Z1j|⟨Z1 − Z2, β⟩ + ϵ > 0 d= Z ′

1 −

(cid:115)

β2
j
2 + σ2 +

βj√

2 + σ2

Z+,

where Z1j is the j-th entry of Z1.

Proof: First, note:

Z1j|⟨Z1−Z2, β⟩+ϵ > 0 d= Z1j|

(cid:42)













,



β





−β

σ

(cid:43)







> 0 d= Z1j|

Z1

Z2

ϵ/σ

(cid:42)













,



β





−β

σ







(cid:112)
/

(cid:43)

2 + σ2

> 0,

Z1

Z2

ϵ/σ

denotes the vector-valued result of concatenating vectors and scalars. For readability, we

where



·


·




·



set Zc =





Z1

Z2

ϵ/σ







and βc =



β





−β

σ







√

/

2 + σ2.

15

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

Under review as a conference paper at ICLR 2025

Given that βc is unit-norm (by supposition, β is unit-norm), and every element of Zc is ∼ N (0, 1)
(even ϵ/σ), we can easily split a conditional random vector containing Z1j into a conditionally
dependent component and independent component:

Zc|⟨Zc, βc⟩ > 0 d= (I − βcβ⊤

c )Z′′ + βcZ+.

The first term is orthogonal to βc and so it is the part of Zc that is not subject to the condition. In the
unconditional case, Zc ∼ N (0, I) and so Z′′ ∼ N (0, I). The second term is the part of Zc that is
in the direction of βc. Z+ ∼ HalfNormal(I) because our dot product condition is satisfied for half
of the possible non-orthogonal Zc values. Now, we focus on finding Zc|⟨Zc, βc⟩ > 0 for a single
index j. We have (for C defined to be the dimensionality of βc):

((I − βcβ⊤

c )Z′′)j + (βcZ+)j = Z ′′

j (1 − βc

(cid:88)

2
j ) −

Z ′′

i βcjβci + βjZ+j

1≤i≤C
i̸=j

= Z ′′

j −

C
(cid:88)

i=1

Z ′′

i βcjβci + βjZ+j.

j −(cid:80)C
2
j βc

Now, note that Z ′′
given by 1 and βc
also use the fact that (cid:80)C
have that the conditional Z1j is given by:

i=1 Z ′′
i , so it itself is a zero-mean Gaussian Y ∼ N (0, 1 − (cid:80)C

i βcjβci is the sum of independent zero-mean Gaussians with variances
2
i ). We can
2
j ). So we

2
j βc
2
i = 1 (recall that βc is unit norm) to get: Y ∼ N (0, 1 − βc

i=1 βc

i=1 βc

2

Z ′(cid:113)

1 − βc

(cid:115)

2
j + βcjZ+ = Z ′

1 −

β2
j
2 + σ2 +

βj√

2 + σ2

Z+,

for Z ′ ∼ N (0, 1). As a corollary, we can see that Z2j under the same condition is given by:

(cid:115)

Z ′

1 −

β2
j
2 + σ2 +

−βj√

2 + σ2

Z+.

B.2 LEMMA 2

Statement of Lemma 2 Suppose that Φ is the CDF of a standard Gaussian, a and c are constants,
and Z ∼ N (0, 1). Then we have:

E[Φ(aZ + c)] = Φ

(cid:18)

√

c
1 + a2

(cid:19)

.

Proof: By the definition of the CDF of a standard Gaussian, we have:
E[Φ(aZ + c)] = E[P (X ≤ aZ + c)],

where X ∼ N (0, 1). Continuing, we have:

= E[P (X − aZ − c ≤ 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 ∼ N (−c, a2 + 1). To find P (X − aZ − c ≤ 0),
we can evaluate its CDF at 0:

(cid:20)

(cid:18)

Φ

= E

√

c
a2 + 1

(cid:19)(cid:21)

(cid:18)

= Φ

√

c
a2 + 1

(cid:19)

.

B.3 LEMMA 3

Statement of Lemma 3 Suppose Φ is the standard Gaussian CDF, Z+ ∼ HalfNormal(1), and b
and a are constants. Then we have:

(cid:20)

Φ

E

(cid:18) Z+b
√

a2 + 1

(cid:19)(cid:21)

=

1
2

+

1
π

tan−1

(cid:18)

√

b
a2 + 1

(cid:19)

.

16

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

Under review as a conference paper at ICLR 2025

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

Proof: By the definition of expected value, we can take the following integral where fZ+ is the PDF
of Z+. We integrate from 0 instead of −∞ because the PDF of the Standard Half Normal is 0 in the
domain below 0:
(cid:20)

(cid:19)(cid:21)

(cid:19)

(cid:90) ∞

(cid:18) Z+b
√

E

Φ

a2 + 1

=

=

=

0
(cid:90) ∞

0
1
√
2π

(cid:18) zb
√

(cid:18) zb
√

Φ

Φ

a2 + 1

(cid:18)(cid:90) ∞

a2 + 1
−z2
2 dz +

e

0

0

fZ+(z)dz

(cid:19) √
2
√
π
(cid:90) ∞

e

−z2
2 dz

(cid:18)

erf

√

zb
a2 + 1

√
2

(cid:19)

(cid:19)

−z2
2 dz

e

(*).

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 (1968), which states:
(cid:19)
(cid:18) d
c

erf(cx)e−d2x2

tan−1

π
2d

dx =

(cid:90) ∞

1
√

√

−

π

d

0

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 erf(·) is positive (and so we can use the integral table). In every case, we find that the solution is
2 + 1
1

π tan−1 (cid:16) b√

(cid:17)

.

a2+1

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

(cid:32) √
√

√
√

−

+

(*) =

1
√
2π

π
2

π
2

tan−1

(cid:32) √

(cid:33)(cid:33)

a2 + 1
b

2
π

Then, using the identity:

we find the following:

=

1
2

+

1
2

−

1
π

tan−1

(cid:32) √

(cid:33)

.

a2 + 1
b

tan−1 x + tan−1 1
x

=

π
2

if x > 0,

=

1
2

+

1
π

tan−1

(cid:18)

√

b
a2 + 1

(cid:19)

.

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

(cid:18) √
√

π
2

(cid:19)

+ 0

(*) =

1
√
2π

=

1
2

.

Because tan−1(0) = 0, we have:

=

1
2

+

1
π

tan−1

(cid:18)

√

b
a2 + 1

(cid:19)

.

Case 3: b < 0. Because erf(·) is an odd function, we can pull the negative out:

1
√
2π
Now we can use the integral table as in the b > 0 case:

−z2
2 dz −

(*) =

erf

e

0

0

(cid:90) ∞

(cid:18)(cid:90) ∞

(cid:18)

√

z|b|
√
a2 + 1

2

(cid:19)

−z2
2 dz

e

(cid:19)

.

1
√
2π

(cid:32) √
√

π
2

−

√
√

π
2

√
√

2
π

+

tan−1

(cid:32) √

(cid:33)(cid:33)

a2 + 1
|b|

1
2

+

1
2

−

1
π

tan−1

(cid:32) √

(cid:33)

.

a2 + 1
|b|

=

=

17

Under review as a conference paper at ICLR 2025

We can then use the same identity again:

to get:

tan−1 x + tan−1 1
x

=

π
2

if x > 0

=

1
2

−

1
π

tan−1

(cid:18) |b|
√

a2 + 1

(cid:19)

.

Because tan−1 is an odd function, we can put the negative inside of it:

=

1
2

+

1
π

tan−1

(cid:18)

√

b
a2 + 1

(cid:19)

.

B.4 FULL PROOF

Here, we prove:

E[sign(y1 − y2)(Φ(x1) − Φ(x2))] =

(cid:32)

2
π

sin−1

θ∗
(cid:112)4 + 2σ2

1 + 2σ2
2

(cid:33)

with y1, y2, Φ(x1), Φ(x2), and θ∗ defined in the main text, for the case where ϵ1 and ϵ2 are zero-
mean Gaussian noise ∼ N (0, σ2

1) and ∼ N (0, σ2

2), respectively.

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

Theorem 1 When ϵ ∼ N (0, σ2), we have:

E[sign(yi − yj)(Φ(xi) − Φ(xj))] =

2
π

sin−1

(cid:18)

θ∗
1 + σ2

√

2

(cid:19)

.

(7)

Proof: By symmetry, we have:

E[sign(y1 − y2)(Φ(x1) − Φ(x2))]

=

1
2

E[Φ(x1) − Φ(x2)| sign(y1 − y2) > 0] +

1
2

E[−(Φ(x1) − Φ(x2))| sign(y1 − y2) < 0].

By increasing monotonicity of f , we have sign(y1 − y2) > 0 ⇐⇒ ⟨x1 − x2, θ∗⟩ + ϵ∆ > 0, for
ϵ∆ = ϵ1 − ϵ2 ∼ N (0, σ2

2). So:

1 + σ2
1
2

=

+

E[−(Φ(x1) − Φ(x2))|⟨x1 − x2, θ∗⟩ + ϵ∆ < 0].

E[Φ(x1) − Φ(x2)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0]
1
2
d= −ϵ∆, the two expected values above are the same:
= E[Φ(x1) − Φ(x2)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0].

Because x1

d= x2 and ϵ∆

By linearity of expectation:

= E[Φ(x1)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0] − E[Φ(x2)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0].

Now, we focus on finding the overall estimate for a single index j. By Lemma 1, we have, for
Z ∼ N (0, 1) and Z+ ∼ HalfNormal(1):

Φ(x1j)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0 d= Φ(Za + Z+b1).

(cid:114)

Here, a =

1 −

(θ∗
2+σ2

j )2
1 +σ2
2

and b1 =

θ∗
j√
2+σ2

1 +σ2
2

. As a corollary of Lemma 1, we can see:

Φ(x2j)|⟨x1 − x2, θ∗⟩ + ϵ∆ > 0 d= Φ(Za + Z+b2).

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

Where b2 = −

θ∗
j√
2+σ2

1 +σ2
2

. So for the index j, our estimate is:

E[Φ(Za + Z+b1)] − E[Φ(Za + Z+b2)]
= E[E[Φ(Za + c)|c = Z+b1]] − E[E[Φ(Za + c)|c = Z+b2]].

Using Lemma 2, we have:

= E

(cid:20)

Φ

(cid:18) Z+b1√
a2 + 1

(cid:19)(cid:21)

− E

(cid:20)

Φ

(cid:18) Z+b2√
a2 + 1

(cid:19)(cid:21)

.

Then, using Lemma 3, we have:

=

=

1
2
1
π

+

1
π

tan−1

(cid:18) b1√
a2 + 1
(cid:19)

tan−1
(cid:18) b1√

(cid:19)

−

1
2

1
tan−1
−
π
(cid:18) b2√

(cid:19)

(cid:18) b2√
(cid:19)

a2 + 1

.

−

tan−1

1
π

a2 + 1
Using the fact that tan−1 is an odd function and b2 = −b1, we get:

a2 + 1

=

2
π

tan−1

(cid:18) b1√

a2 + 1

(cid:19)

.

Now, we write a and b1 in terms of θ∗
j :

=

2
π

tan−1

=

2
π

tan−1















(cid:114)

(cid:115)

θ∗
j√
2+σ2

1 +σ2
2

2 −

(θ∗
2+σ2

j )2
1 +σ2
2







θ∗
j√
4+2σ2

1 +2σ2
2

(cid:18)

1 −

θ∗
j√
4+2σ2

1 +2σ2
2

(cid:19)2









.

Using the identity sin−1 x = tan−1 (cid:16) x√

1−x2

(cid:17)

, we have:

=

2
π

sin−1

(cid:32)

θ∗
j
(cid:112)4 + 2σ2
1 + 2σ2
2

(cid:33)

.

B.5 COROLLARY 1

Corollary 1 Suppose that ˆθ is any vector of fixed weights and x ∼ N (0, I). Then, conditioning on
the event ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩, we have with probability 1 that:

⟨ ˆθ, E[Φ(xi) | ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩]⟩ < ⟨ ˆθ, E[Φ(xj) | ⟨ ˆθ, xi⟩ < ⟨ ˆθ, xj⟩]⟩.

(9)

To see this, we can find:

E[Φ(x1) − Φ(x2)|⟨ ˆθ, x1⟩ + ϵ1 > ⟨ ˆθ, x2⟩ + ϵ2] = E[Φ(x1) − Φ(x2)|⟨ ˆθ, x1 − x2⟩ + ϵ∆ > 0]
Note that we have already computed this expected value in the proof above; for an index j, it is:

(cid:32)

2
π

sin−1

ˆθj
(cid:112)4 + 2σ2

1 + 2σ2
2

(cid:33)

.

Because sin−1 is an odd function, the above expression has the same sign as ˆθj. Because the
values at every index of E[Φ(x1) − Φ(x2)] under our condition and ˆθ are the same sign, we have
⟨E[Φ(x1) − Φ(x2)], ˆθ⟩ > 0, so ⟨ ˆθ, E[Φ(x1)]⟩ > ⟨ ˆθ, E[Φ(x2)]⟩.

19

972

973
974
975
976
977

978
979
980
981
982
983

984
985
986
987
988
989

990
991
992
993
994
995

996
997
998
999
1000

1001
1002
1003
1004
1005
1006

1007
1008
1009
1010
1011
1012

1013
1014
1015
1016
1017
1018

1019
1020
1021
1022
1023

1024
1025

Under review as a conference paper at ICLR 2025

C OPTIMAL PROJECTED WEIGHTS SOLUTIONS

C.1 LINEAR PROJECTION

Theorem 2 Suppose we want to solve:

subject to:

ˆθproj = arg min
θ∈RD

−⟨θ, ˆθ⟩,

D
(cid:88)

i=1

θi = 1

0 ≤ θi ≤ τi, ∀i ∈ [1, D],

where τi > 0 are fixed values. Then, the solution is:

ˆθproj
k =






τk
1 − (cid:80)
0

j: rj (ˆθj )>rk(ˆθk) τj

if (cid:80)
if (cid:80)
otherwise

j: rj (ˆθj )≥rk(ˆθk) τj ≤ 1
j: rj (ˆθj )≥rk(ˆθk) τj ≥ 1 ∧ (cid:80)

j: rj (ˆθj )>rk(ˆθk) τj ≤ 1

, (10)

where r is some function that breaks all ties between ˆθj and ˆθk for k ̸= j, and otherwise leaves the
ordinal relationships the same.

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 ˆθproj and yet ˆθproj
for some ˆθproj
satisfying the projection constraints that is the same as ˆθproj except in these places:

k < τk
falling under the first case of Equation 10. Now suppose that we construct a θ′ also

k

k + ∆ = τk
p − δ1 ≥ 0

k = ˆθproj
θ′
p = ˆθproj
θ′
...
q = ˆθproj
θ′

q − δn ≥ 0

for some ∆ = (cid:80)n
i=1 δi > 0 where ˆθp ≥ · · · ≥ ˆθq are all of the ˆθ values which do not fall under the
first condition and where the corresponding ˆθproj values are nonzero. We know that there must be
some ˆθproj
from which we can subtract δ1, · · · , δn (and so from which we can take the ∆)
because (cid:80)

j: rj (ˆθj )≥rk(ˆθk) τj ≤ 1. Now, we have:

p , · · · , ˆθproj

q

q − ˆθk
ˆθproj

⟨ ˆθ, ˆθproj⟩ − ⟨ ˆθ, θ′⟩
= ˆθk
p + · · · + ˆθq
ˆθproj
k + ˆθp
ˆθproj
= −ˆθk∆ + ˆθpδ1 + · · · + ˆθqδn
≤ ˆθp(δ1 + · · · + δn) − ˆθk∆
= ˆθp∆ − ˆθk∆
≤ 0.

k − ˆθk∆ − ˆθp
ˆθproj

p − · · · − ˆθq
ˆθproj

q + ˆθpδ1 + · · · + ˆθqδn
ˆθproj

At this point, the only way to avoid the contradiction result would be if ˆθk = ˆθp = · · · = ˆθq.
Otherwise, the above non-strict inequality would be a strict inequality. If ˆθk = ˆθp = · · · = ˆθq, then
we know that ˆθk is the smallest ˆθ value satisfying condition 1 and all of the other greater ˆθ values
satisfying condition 1 must be projected to their τ threshold value (otherwise we would get the
contradiction result). In this edge case can see above that rearranging the remaining weight among

20

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

Under review as a conference paper at ICLR 2025

equal ˆθ 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 ˆθproj and yet ˆθproj
falling under the third case of Equation 10. Now suppose that
we construct a θ′ also satisfying the projection constraints that is the same as ˆθproj except in these
places:

k > 0 for some ˆθproj

k

k − ∆ = 0
p + δ1 ≤ τp

k = ˆθproj
θ′
p = ˆθproj
θ′
...
q = ˆθproj
θ′

q + δn ≤ τq

q

p , · · · , ˆθproj

to which we can add δ1, · · · , δn. Now, we have:

for some ∆ = (cid:80)n
i=1 δi > 0 where ˆθp ≥ · · · ≥ ˆθq are all of the ˆθ values which do not fall under the
third condition and where the corresponding ˆθproj values are not at their thresholds. By construction
we know that there must be some ˆθproj
⟨ ˆθ, ˆθproj⟩ − ⟨ ˆθ, θ′⟩
= ˆθk
ˆθproj
p + · · · + ˆθq
ˆθproj
k + ˆθp
= ˆθk∆ − ˆθpδ1 − · · · − ˆθqδn
≤ −ˆθq(δ1 + · · · + δn) + ˆθk∆
= −ˆθq∆ + ˆθk∆
≤ 0.

q − ˆθpδ1 − · · · − ˆθqδn
ˆθproj

ˆθproj
k + ˆθk∆ − ˆθp

p − · · · − ˆθq
ˆθproj

q − ˆθk
ˆθproj

At this point, the only way to avoid the contradiction result would be if ˆθk = ˆθp = · · · = ˆθq.
Otherwise, the above non-strict inequality would be a strict inequality. If ˆθk = ˆθp = · · · = ˆθq, then
we know that ˆθk is the largest ˆθ value satisfying condition 3 and all of the other smaller ˆθ 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 ˆθ 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 ˆθproj not covered by either case.

C.2 QUADRATIC PROJECTION

C.2.1 LEMMA 4

Statement of Lemma 4 Suppose that ˆθproj is the optimal solution to:

subject to:

ˆθproj = arg min
θ∈RD

|| ˆθ − θ||2
2,

D
(cid:88)

i=1

θi = 1

where τi > 0 are fixed values. Then, ˆθproj

0 ≤ θi ≤ τi, ∀i ∈ [1, D],
s = 0 implies that any j with ˆθs > ˆθj must have ˆθproj

j = 0.

Proof: This is similar to Lemma 2 from Shalev-Shwartz & Singer (2006). Assume for the sake of
contradiction ˆθproj

s = 0 and ˆθs > ˆθj, yet we have ˆθproj

j > 0.

21

1080
1081

1082
1083
1084
1085
1086
1087

1088
1089
1090
1091
1092
1093

1094
1095
1096
1097
1098

1099
1100
1101
1102
1103
1104

1105
1106
1107
1108
1109
1110

1111
1112
1113
1114
1115
1116

1117
1118
1119
1120
1121

1122
1123
1124
1125
1126
1127

1128
1129
1130
1131
1132
1133

Under review as a conference paper at ICLR 2025

Now we can construct another vector θ′ that is the same as ˆθproj, except in two places:
s = ˆθproj
θ′
j = ˆθproj
θ′
for some ∆ satisfying 0 < ∆ < min(ˆθproj
, τs − ˆθproj
within the thresholds. We know that ∆ can exist because min(ˆθproj
τs − ˆθproj
j > 0).
Now we can compute:
|| ˆθ − ˆθproj||2

s = τs − 0 > 0 and ˆθproj

s + ∆
j − ∆,

)2 − (ˆθs − (ˆθproj

)2 + (ˆθj − ˆθproj

, τs − ˆθproj

2 − || ˆθ − θ′||2

s

s

s

j

j

). This bound on ∆ ensures that θ′ is still
) > 0 (by supposition,

s + ∆))2 − (ˆθj − (ˆθproj

j − ∆))2

s

2 = (ˆθs − ˆθproj
= 2∆((ˆθs − ˆθproj
> 2∆((ˆθs − ˆθproj
≥ 2∆((ˆθs − ˆθproj
s
= 2∆(ˆθs − ˆθj)
> 0.

s

j

j
) − (ˆθj − ˆθproj
) − (ˆθj − ˆθproj
) − (ˆθj − ˆθproj

j

j

) − ∆)
) − min(ˆθproj
) − ˆθproj
j

)

j

, τs − ˆθproj

s

))

So ˆθproj cannot be the optimal solution.

C.2.2 LEMMA 5

Statement of Lemma 5 Suppose that ˆθproj is the optimal solution to:

subject to:

ˆθproj = arg min
θ∈RD

|| ˆθ − θ||2
2,

D
(cid:88)

i=1

θi = 1

where τi > 0 are fixed values. Then, ˆθproj

j = τj for any ˆθj − τj > ˆθs − τs.

0 ≤ θi ≤ τi, ∀i ∈ [1, D],
s = τs implies ˆθproj

Proof: Again, this is similar to Lemma 2 from Shalev-Shwartz & Singer (2006). Assume for the
sake of contradiction ˆθproj

s = τs and ˆθj − τj > ˆθs − τs, yet we have ˆθproj

j < τj.

Now we can construct another vector θ′ that is the same as ˆθproj, except in two places:
s = ˆθproj
θ′
j = ˆθproj
θ′
for some ∆ satisfying 0 < ∆ < min(ˆθproj
, τj − ˆθproj
within the thresholds. We know that ∆ can exist because min(ˆθproj
j > 0 and ˆθproj
τj − ˆθproj
Now we can compute:
|| ˆθ − ˆθproj||2

s − ∆
j + ∆,

s = τs > 0).

)2 − (ˆθs − (ˆθproj

, τj − ˆθproj

2 − || ˆθ − θ′||2

s

s

s

j

j

). This bound on ∆ ensures that θ′ is still
) > 0 (by supposition,

s − ∆))2 − (ˆθj − (ˆθproj

j + ∆))2

j

s

)2 + (ˆθj − ˆθproj

j
) − (ˆθs − ˆθproj
) − (ˆθs − ˆθproj
) − (ˆθs − ˆθproj

2 = (ˆθs − ˆθproj
= 2∆((ˆθj − ˆθproj
> 2∆((ˆθj − ˆθproj
≥ 2∆((ˆθj − ˆθproj
s
j
= 2∆((ˆθj − τj) − (ˆθs − ˆθproj
= 2∆((ˆθj − τj) − (ˆθs − τs))
> 0.

))

s

s

j

) − ∆)
) − min(ˆθproj
) − (τj − ˆθproj

s

j

))

, τj − ˆθproj

j

))

So ˆθproj cannot be the optimal solution.

22

1134

1135
1136
1137
1138
1139

1140
1141
1142
1143
1144
1145

1146
1147
1148
1149
1150
1151

1152
1153
1154
1155
1156
1157

1158
1159
1160
1161
1162

1163
1164
1165
1166
1167
1168

1169
1170
1171
1172
1173
1174

1175
1176
1177
1178
1179
1180

1181
1182
1183
1184
1185

1186
1187

Under review as a conference paper at ICLR 2025

C.2.3 FULL PROOF

Theorem 3 Suppose we want to solve:

subject to:

ˆθproj = arg min
θ∈RD

|| ˆθ − θ||2
2,

D
(cid:88)

i=1

θi = 1

0 ≤ θi ≤ τi, ∀i ∈ [1, D],

where τi > 0 are fixed values. Then the solution is:

where λ is found (through e.g. bisection search) to satisfy:

k = min(max(ˆθk − λ, 0), τk),
ˆθproj

D
(cid:88)

i=1

min(max(ˆθi − λ, 0), τi) = 1.

Proof: Note that this problem is the same as the simplex projection problem from Shalev-Shwartz
& Singer (2006) and Duchi et al. (2008), except here we have additional θi ≤ τi constraints. The
Lagrangian for this problem is4:

L(θ, µ, ζ, λ) =

(cid:32)

|| ˆθ − θ||2

2 + λ

−1 +

1
2

N
(cid:88)

i=1

(cid:33)

θi

− ⟨µ, θ⟩ + ⟨ζ, θ − τ ⟩.

To find the optimality condition with respect to a single index of θ, we set the derivative to zero:

dL
dθi

= θi − ˆθi + λ − µi + ζi = 0.

The complimentary slackness KKT condition gives us that ζi = µi = 0 when 0 < θi < τi, so for θi
not at the boundary of our constraints, we get:

θi = ˆθi − λ.
So, we have that for all θi ∈ (0, τi), there is a shared value λ which we subtract from ˆθi to get the
value of θi. How do we know which θi are 0 and which θi are τi, though?

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

ˆθproj
k = max(ˆθk − λ, 0),

for ˆθproj
k

for ˆθproj
k

̸= τk. By Lemma 5, we can characterize the optimal solution as:
k = min(ˆθk − λ, τk),
ˆθproj

̸= 0. So, we can combine these two forms to get:

k = min(max(ˆθk − λ, 0), τk).
ˆθproj

Now recall that we have the following constraint:

D
(cid:88)

i=1

min(max(ˆθi − λ, 0), τi) = 1.

Given this constraint, we can find λ through search (moving the value up or down). We can see this
by noticing that (cid:80)D
i=1 min(max(ˆθi − λ, 0), τi) is a strictly decreasing function of λ between the
setting of λ that makes ˆθi − λ > 0 for at least one i, and the setting of λ that makes ˆθi − λ < τi for
at least one i. So in this range, there is only one setting of λ that satisfies this equation. We can only
choose a λ outside of this range when (cid:80)D
i = τi
for all i.

i=1 τi = 1, and in this case the solution is trivial: ˆθproj

4Note that multiplying || ˆθproj − θ||2

2 by 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.

23

1188
1189

1190
1191
1192
1193
1194
1195

1196
1197
1198
1199
1200
1201

1202
1203
1204
1205
1206

1207
1208
1209
1210
1211
1212

1213
1214
1215
1216
1217
1218

1219
1220
1221
1222
1223
1224

1225
1226
1227
1228
1229

1230
1231
1232
1233
1234
1235

1236
1237
1238
1239
1240
1241

Under review as a conference paper at ICLR 2025

1242
1243

1244
1245
1246
1247
1248
1249

1250
1251
1252
1253
1254
1255

1256
1257
1258
1259
1260

1261
1262
1263
1264
1265
1266

1267
1268
1269
1270
1271
1272

1273
1274
1275
1276
1277
1278

1279
1280
1281
1282
1283

1284
1285
1286
1287
1288
1289

1290
1291
1292
1293
1294
1295

D 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, 2010) and the isotron algorithm (Kalai & Sastry, 2009).
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., 2020; Llama Team, 2024; Ruan
et al., 2024). 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
(2017), and even the standard Spearman correlation (Spearman, 1904) (see Appendix G) performed
well. We believe that in general, robust per-feature correlations are likely to perform well as D ≫ N ,
and extreme levels of regularization are needed to obtain reasonable models. Sparse methods such
as the Lasso (Tibshirani, 1996) are one classic answer, but we cannot necessarily assume that the
underlying correlations θ∗ are sparse, and we did not find these techniques to perform well.

E 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. 2023).
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.

F 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.

F.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., 2019) and appropriate PyTorch (Ansel et al., 2024)
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., 2023) for more information; we initialized the model from scratch using their 160M model
configuration at https://huggingface.co/EleutherAI/pythia-160m. Other hyperparameters
can be assumed to be Hugging Face Trainer defaults at the time of this writing.

F.2 LLM EVALUATION

At the end of the pretraining script, we used the Eleuther AI Eval Harness (Gao et al., 2023). 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.

F.3 DSIR

DSIR (Xie et al., 2023b), 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

24

Under review as a conference paper at ICLR 2025

Table 2: LLM Pretraining Hyperparameters

Parameter

Per-device Batch Size

Learning Rate

Warmup Ratio

Adam β1

Adam β2

Adam ϵ

Weight Decay

LR Scheduler

Max Grad Norm

BF 16

Distributed Backend

Gradient Accumulation Steps

Value

128
5 × 10−3

0.1

0.9

0.95
1 × 10−8

0.1

cosine

1.0

True

nccl

1

Table 3: Unique pretraining tokens selected per benchmark, from DSIR.

Benchmark

Tokens

ARC Easy

2,905,206,499

PIQA

SCIQ

2,910,486,295

2,920,734,042

LAMBADA

3,022,219,424

LAMBADADE

3,210,986,137

LAMBADAES

3,396,528,704

LAMBADAFR

3,413,930,081

LAMBADAIT

3,384,854,845

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 LAMBADAFR. 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.

F.4

FASTTEXT

The “SOTA” fastText model from Li et al. (2024) is available here: 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

25

1296
1297

1298
1299
1300
1301
1302
1303

1304
1305
1306
1307
1308
1309

1310
1311
1312
1313
1314

1315
1316
1317
1318
1319
1320

1321
1322
1323
1324
1325
1326

1327
1328
1329
1330
1331
1332

1333
1334
1335
1336
1337

1338
1339
1340
1341
1342
1343

1344
1345
1346
1347
1348
1349

Under review as a conference paper at ICLR 2025

1350
1351

1352
1353
1354
1355
1356
1357

1358
1359
1360
1361
1362
1363

1364
1365
1366
1367
1368

1369
1370
1371
1372
1373
1374

1375
1376
1377
1378
1379
1380

1381
1382
1383
1384
1385
1386

1387
1388
1389
1390
1391

1392
1393
1394
1395
1396
1397

1398
1399
1400
1401
1402
1403

(a) Estimate with linear projection. This is our algo-
rithm from the main text without training the addi-
tional fastText filter.

(b) Estimate with quadratic projection. Same as (a) ex-
cept the linear projection is replaced with the quadratic
projection.

(c) Spearman rank correlation with linear projection.
Same as (a) except we replaced our estimator with the
Spearman rank correlation.

(d) fastText filter trained on data selected in (c). This
is the same as our algorithm in the main text, replacing
our estimator with the Spearman rank correlation.

Figure 6: Pretraining results for different methods within our paradigm. Overall, we see that many
rank-correlation pretraining data selection approaches perform well.

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.

G ADDITIONAL PRETRAINING RESULTS

In Figure 6, we present additional pretraining results for methods in our loss-performance correlation
data selection paradigm. We find that using Spearman rank correlation (Spearman, 1904) 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.

26

Under review as a conference paper at ICLR 2025

Figure 7: This figure is analogous to Figure 3, except the τ thresholds have been multiplied by 5.
We see that our approach selects even more relevant data when the selection pool is larger.

Figure 8: The parameter-count histogram of the 90 models from the Open LLM Leaderboard
(Beeching et al., 2023) that we used to compute our estimate for pretraining data selection. Bar
widths are 160M. The smallest model in the sample has ≈33M parameters and the largest has ≈9B.
The spike around 6.7B parameters is due to a large number of partially trained Pythia (Biderman
et al., 2023) 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.

H PRETRAINING TOKEN DISTRIBUTION WITH 5 × τ

Figure 7 shows what the projected estimate in our pretraining experiments would be if we had a
pretraining data pool 5× 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.

I PARAMETER COUNT DISTRIBUTION FOR ESTIMATOR LLMS

In Figure 8, we present the parameter-count histogram of the 90 models from the Open LLM Leader-
board (Beeching et al., 2023) 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 effec-
tively pretrain 160M parameter LLMs.

J ANALYSIS OF THE MODEL-LOSS MATRIX X

What information is contained in the matrix of model losses 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, 1901) and t-SNE (van der Maaten & Hinton, 2008) on X and plotted the first two
components for each of our 9,841 domains. As shown in the first row of Figure 9, 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.

27

1404
1405

1406
1407
1408
1409
1410
1411

1412
1413
1414
1415
1416
1417

1418
1419
1420
1421
1422

1423
1424
1425
1426
1427
1428

1429
1430
1431
1432
1433
1434

1435
1436
1437
1438
1439
1440

1441
1442
1443
1444
1445

1446
1447
1448
1449
1450
1451

1452
1453
1454
1455
1456
1457

Under review as a conference paper at ICLR 2025

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

(i)

(j)

Figure 9: Analysis of the loss matrix. The first row treats domains as examples to be projected via
PCA, while the second row treats models as examples. Panels (a): eigenvalue decay for the eigende-
composition of the D×D covariance matrix resulting from the loss matrix; a few dominant PCs are
seen. (b) and (c): domains plotted by the first two PCA components showing separation of language
in b and entropy in c. (d,e) show analogous plots in t-SNE with a clearer separation of language. (f):
eigenvalue decay analogous to (a). (g,h): models plotted by the first two PCA components showing
clustering by model family (clusters show Pythia (Biderman et al., 2023), Qwen (Bai et al., 2023),
and OpenLlama (Geng & Liu, 2023) derivatives – the three largest clusters in our data), and average
model loss. (i,j) show analogous results under t-SNE where (i) is normalized to remove per-model
entropy differences.

The second row of Figure 9 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., 2023), Qwen (Bai et al., 2023), and OpenL-
lama (Geng & Liu, 2023)), we see that model families are clustered together in the PC graphs.

1458
1459

1460
1461
1462
1463
1464
1465

1466
1467
1468
1469
1470
1471

1472
1473
1474
1475
1476

1477
1478
1479
1480
1481
1482

1483
1484
1485
1486
1487
1488

1489
1490
1491
1492
1493
1494

1495
1496
1497
1498
1499

1500
1501
1502
1503
1504
1505

1506
1507
1508
1509
1510
1511

28

