Data Selection via Optimal Control for Language Models
Authors: Yuxian Gu
Li Dong
Hongning Wang
Yaru Hao
Qingxiu Dong
Furu Wei
Minlie Huang
Venue: ICLR 2025 Oral (2025)
Keywords: Pre-training Language Models
Data Selection
Optimal Control
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper introduces a framework to select high-quality pre-training data via optimal control.

Abstract: This work investigates the selection of high-quality pre-training data from massive corpora to enhance LMs' capabilities for downstream usage. 
We formulate data selection as a generalized Optimal Control problem, which can be solved theoretically by Pontryagin's Maximum Principle (PMP), yielding a set of necessary conditions that characterize the relationship between optimal data selection and LM training dynamics.
Based on these theoretical results, we introduce **P**MP-based **D**ata **S**election (**PDS**), a framework that approximates optimal data selection by solving the PMP conditions. 
In our experiments, we adopt PDS to select data from CommmonCrawl and show that the PDS-selected corpus accelerates the learning of LMs and constantly boosts their performance on a wide range of downstream tasks across various model sizes.
Moreover, the benefits of PDS extend to ~400B models trained on ~10T tokens, as evidenced by the extrapolation of the test loss curves according to the Scaling Laws.
PDS also improves data utilization when the pre-training data is limited, by reducing the data demand by 1.8 times, which helps mitigate the quick exhaustion of available web-crawled corpora. Our code, model, and data can be found at https://github.com/microsoft/LMOps/tree/main/data_selection.

Published as a conference paper at ICLR 2025
DATA SELECTION VIA OPTIMAL CONTROL FOR LANGUAGE MODELS
Yuxian Gu!:2; Li Dong", Hongning Wang'
Yaru Hao?,
Qingxiu Dong"
Furu Wei?, Minlie Huang
'The CoAl Group, Tsinghua University Microsoff Research
Peking University
ABSTRACT
This work investigates the selection of high-quality pre-training data from mas- sive corpora enhance LMs' capabilities for downstream usage. We formulate data selection generalized Optimal Control problem, which can be solved theoretically by Pontryagin's Maximum Principle (PMP), yielding a set of neces- sary conditions that characterize the relationship between optimal data selection and LM training dynamics. Based on these theoretical results, we introduce PMP-based Data Selection (PDS), framework that approximates optimal data selection by solving the PMP conditions. In our experiments, we adopt PDS to corpus accel-
select data from CommmonCrawi and show that the erates the learning of LMs and constantly boosts their performance on a wide range of downstream tasks across various model sizes. Moreover, the benefits of PDS extend to -400B models trained on -10T tokens, as evidenced by the extrapolation of the loss curves according to the Scaling Laws. PDS also improves data utilization when the pre-training data is limited, by reducing the data demand by 1.8 times, which helps mitigate the quick exhaustion of avail- able web-crawled corpora. Our code, model, and data can be found at ht ttsps:: / //gi tighhuu. com/miimmi/oooooo mi www.oott.ooooooooooooooooooo wwroossfft/MMMpp0 LMOps tree ma n/ datta_aellltttaeeeeeeeeeeeeeleae ion.
Average Accuracy (96)
52
PDS (Ours)
Redpajama
Average Accuracy (%)
52
PDS (Ours)
Radpajama
50
50
2.0x
48
46
44
44
1.07
2.14
3.21
5.35
160M
470M
18
1.78
Computation (x1020 FLOPs)
Model Size
(a) Scaling of Training Computation (1.7B)
(b) Scaling of Model Size
Figure 1: Scaling curves of average accuracy 9 widely-used downstream tasks with respect computation (a) and model sizes (b). We select pre-training corpora from the CommonCraw/ and pre- train LMs on the selected data. PDS compared with the Redpajama data cleaning pipeline (Together, 2023). The computation curves are calculated based on the training of a 1.7B LM.
INTRODUCTION
With the thriving of language models (LMs; Han al., 2021; Bommasani et al., 2021), the role of data selection for pre-training becomes increasingly important, which aims at identifying valuable pre-training instances to accelerate model learning or improve downstream performance (Albalak et al., 2024). This focus enables researchers to explore the limit of LMs in the face of increasing training data demands (Brown al., 2020; OpenAI, 2023; Team et al., 2023). It also helps to reduce the computational costs during pre-training (Sorscher et al., 2022), and addresses the potential
"Contribution during an internship at Microsoft Research. (guyx218 8mai ils tsingnuu..dee en)
'Corrsponding author.

---

Published as a conference paper at ICLR 2025
limitations caused by available Internet data (Villalobos al., 2022; Muennighoft et al., 2023).
Without doubt, pre- teeiiiaiing data selection is highly valuable for both research and industry sectors.
Unlike previous works relying primarily manually crafted hcuristies (Tirumala et al., 2023; Xie ct al., 2023), we connect data selection with classical Optimal Control theory (Lewis et al., 2012), where control variables in dynamic system are optimized achieve desired objectives. This mathematical formulation allows fine-graine tne-grained white-box analysis of how the control variables drive dynamic system from one state another. In particular, by treating data selection as the control variables (i.e., whether a data point is included in pre-training), the LM pre-training process as the dynamic system, and the LM's downstream performance as the objective, we leverage Pontryagin's Maximum Principle (PMP: Pontryagin, 2018) to derive the necessary conditions for optimal data selection in theory. These results offer a rigorous, theory-driven alternative to the ad-hoe trial-and- error practices that currently dominate LM pre-training.
Based on our theoretical results, we introduce PMP-based Data Selection (PDS), a framework that selects high-quality pre-training data at scale, by solving the equation system induced by the PMP conditions. Balancing effectiveness and efficiency, PDS first solves the equation system for the optimal data selection on proxy dataset (e.g.. 0.2B tokens), assigning quality score to each instance based on its impact on downstream tasks. After that, a data scorer (c.g.. 125M parameters) is trained to predict the quality scores and then infers scores on the target corpus (e.g., 50B tokens). Finally, the predicted scores guide data selection for pre-training LMs with various sizes (e.g., 1.7B parameters).
Unlike previous pre-training data selection methods based deduplication (Tirumala al., 2023; Abbas et al., 2023), pattern matching (Xie et al., 2023), or single checkpoint performance (Engstrom et al., 2024), which are agnostic the pre-training process of LMs, PDS exploits the highly dynamic nature of LM pre-training through the theoretical optimal control perspective. On the other hand, compared to methods that incorporate signals from the LM training process online (Yu et al., 2024: Wang et al., 2024), PDS operates offline, before the training begins, which avoids additional training-time time computation overhead and allows for training LMs with arbitrary configurations while performing PDS only once. Furthermore, PDS only filters the training corpus, leaving highly optimized pre-training pipelines largely intact. Most importantly, PDS enjoys a strong theoretical basis, opening up the black box of understanding of individual data point impact on LM pre-training.
In our experiments, we select data from the CommonCrawl with PDS using LIMA (Zhou et al., 2024) as the guide for downstream performance and pre-train LMs with 160M, 470M, 1B, and 1.7B parameters from scratch. Then, we test the LM's zero-shot performance on tasks other than LIMA to examine the generalization of PDS. We observe around 2 times speed-up in pre-training on the 1.7B LM and constant improvement in downstream tasks and language modeling performance across all model sizes compared to state-of-the-arr baselines. Extrapolating these results using the Scaling Law (Kaplan et al., 2020; Hoffmann et al., 2022), we show that the benefits remain consistent for -400B LMs trained on 115T tokens. Besides, PDS enhances data utilization in a data-constrained setting, reducing the pre-training data demand by 1.8 times, which is a critical advantage as the LM community is running out of data (Villalobos et al., 2022). We also conduct extensive analysis and ablation studies on the key factors of PDS to facilitate further research on data selection.
2
METHOD
2.1 PROBLEM FORMULATION
We study an LM parameterized with 0 IR pre-trained from scratch on a dataset D = {I, 1 |P] over T' training steps. Data selection (Albalak et al., 2024) aims at identifying subset D' from D. such that LMs trained on D' achieve better performance, measured by a lower downstream loss J(0).
n=
The pre-training process renders (0) as a function of D', which can be fully characterized by data quality score vector 7 = [1. ... ]] T in a [D]-dimensional simplex U, where U =
{ In, 72, 72, I EIP| ?h = 1 and 7n 2 0 for 1 S S [D|} !. A higher quality score in indicates the corresponding instance is more helpful to reduce ((),, and thus the LM should learn
'Only the relative data quality is meaningful for data selection. Therefore, ensure the sum of the quality scores to avoid the impact of their individual scales.
2

---

Published as a conference paper at ICLR 2025
more from the instance. Here, focus the independent importance of each example and provide discussion on the dependence between each data point in Appendix E. This results in the following general pre-training loss, defined as the weighted sum of the per-instance loss by 7:
[P]
L(O. v) Σ 7al 0).
(1)
where [[En: 0) - lo (rr). The goal of data selection is thus to find 7 that reduces the downstream loss J()), and then select instances with the highest scores in Y. For simplicity, we assume that the LM is trained using Gradient Decent (GD) for 0 < < T, with the derivation under the Adam optimizer (Kingma & Ba, 2015) provided in Appendix C:
Oc+1 = 0 1 ().
where 0, represents the model parameters at the time step & during GD and is the learning rate.
Optimization Problem. Motivated by the literature on learned optimizers (Metz et al., 2020), optimize y by minimizing the area under the curve (AUC; Cortes Mohri, 2003) of J(O)), which approximated by the cumulative sum of [(+) over the pre-training process:
min Σ J(O.)
(3)
0.+1 O. -1 L 7), 7 e U.
Intuitively, a lower AUC corresponds to faster convergence of the loss and improved final downstream performance. Unlike evaluating J(,,) at individual time steps, the AUC captures the overall LM training dynamics. As shown in Appendix A, minimizing the AUC essentially enhances the constants in the LM's Scaling Laws (Kaplan al., 2020), leading substantial improvements in LM learning.
2.2 DATA SELECTION AS OPTIMAL CONTROL
We recognize the optimization problem in Eq. (3) is analogous to a discrete optimal control prob- lem (Lewis et al., 2012). where (-) is the cost function, the model parameters 0 are the state variables evolving according to Eq. (2), and the data quality scores y are the control variables to be optimized within U. This perspective makes convenient for theoretically solving the data selection problem.
Theoretically Optimal Solution for Data Selection. Optimal control problems be solved by a powerful tool known as Pontryagin's 's Maximum Principle (PMP; Pontryagin, 2018), which provides a set of necessary conditions for the optimal control variables and their corresponding state variables (See Appendix B for its formal expression). However, standard PMP conditions allow the optimal control to vary over time, whereas in Eq. (3), the control variables r are constrained to be time-invariant due to the offline nature of data selection in our setting. This typically makes the optimization problem more challenging due to the shrinking of feasible region. In the following, we present the PMP conditions for data selection under this constraint:
Theorem 2.1 (PMP Conditions for Data Selection). Let n solve the problem in Eq. (3), and 0; denote the LM parameters trained with y² For 0 < < T. there exists a vector Xi R such that
0: = 0 L Y²), 0 = Ao.
(4)
X; Xi + VJ(()))
2 L(O Nit1, X; =V](O),
(5)
ID
T-
N Vil(nn, (i)
"' arg max
7 EEU,
(6)
where L2 7") denotes the Hessian matrix of L(0, x²) with respect to 0 evaluated at 0 0
We prove Theorem 2.1 in Appendix B using the standard PMP conditions and the Lagrange multiplier method, and provide an illustration for this theorem in Figure 2. Since standard PMP conditions and the Lagrange multiplier method are necessary conditions for the optimality, Theorem 2.1 is also a necessary condition for optimal data selection. By inspecting the PMP conditions for data selection, we can see that Eq. (4) ensures the LM parameters, trained with the optimal data quality scores, continue to evolve via GD. As illustrated in Figure 2 (Left), Eq. (5) defines A; "target vector"
3

---

Published as a conference paper at ICLR 2025
Downstream 1(0.)
Conventional Data Salection
Optimal Selection
Vmm
Vlkt
X
M. Vha M. Pinc Ee Aia1 Vlae
SteiningSteps
Data Quality Soores: Yo
< 7m < 7%
Figure 2: An illustration of Theorem 2.1. Left: Xi+1 E R defines "target vector" aligning with the optimization direction towards optimal data selection, as in Eq. (5). Right: data quality scores are positively correlated with how close the gradient direction of each instance to the target direction, calculated as the dot-product between A and V4, Vi 02) for i = m, k, as in Eq. (6).
suggesting the ideal gradient direction formed only by high-quality data points. In particular, X aggregates information about the downstream loss ()(,)) with respect to current training step and the LM's training dynamics V L y²), from I to t. As a result, X; summarizes the dynamics of LM pre-training (i.e., from future to the current). Since 7 E U, Eq. (6) essentially suggests that Zn with a higher Σ, 1i+1 VI(I,, 07) value should obtain a larger score in y", as shown in Figure (Right). This indicates that the instances whose gradients align closely with the target vector X; should be selected. Note that the PMP conditions for data selection form a complete equation system, where O7 Xj, and 7 are the solutions. In principle, by solving Eqs. (4)-(6) simultaneously, can derive the optimal data quality scores, forming the foundation of our data selection framework PDS.
2.3
PDS: DATA SELECTION BASED ON PMP
PDS selects pre-training data by solving the PMP conditions defined in Eqs. (4)-(6), and consists of three key components, as illustrated in Figure 3. To balance effectiveness and efficiency, PDS first uniformly samples a proxy dataset DD from the pre-training corpus D. Algorithm 1, derived from the PMP conditions, is then applied to D to compute data quality scores for each instance (Section 2.3.1). Then, a data scorer, typically a small LM, is fine-tuned to predict the quality scores based on the instances in Derx The learnt data scorer is subsequently applied to infer quality scores on the entire pre-training corpus D (Section 2.3.2). Finally, the instances with large scores are selected to form high-quality corpus D', which is used to pre-train LMs with any size (Section 2.3.3).
2.3.1 DATA QUALITY SCORES FROM PMP
Algorithm solves the PMP condi- tions for data selection iteratively and returns the data quality scores 7
Algorithm 1 PMP-Solver
Input: LM learning rate 7). Outer loop learning rate Outer loop epochs Tg- Training data hefore selection D. Downstream loss (Ø). Training steps T. Proj[-] that projects point in |!ll to U. Model initialization On.
Overview:
Algorithm 1 solves bi- level optimization problem, consist- ing of an outer loop to compute >-" and two inner loops to compute 0r and A; based on the current y Y is first uniformly initialized and then up- dated for T. epochs, based on 07 and X; obtained in the outer iterations.
Output: Data quality scores "
[n, [n,72, 72.
. -- D
repeat To times
D Outer loop Forward inner loop D Eq. (4)
for = 0, 1... , - do T 0.+1 + 0. V/VI.
end for
A7 < (.(OT)
Inner loops: In each iteration of the outer loop, we run the forward in- ner loop to compute 07 according to
for L = T - 1,T ..... .......11 1 do
4 Reverse inner loop
X Act1 v.)(o)) - ((c. & Eq. (5)
end for
Eq. (4) from = 0 to f T 1, equivalent to training the LM with GD using the current data quality scores to re- we-ghttht the per-instance losses. After that, X; is computed from t - 1 to t = 0 with Eq. (5) in the reverse inner loop, based on 0 obtained from the forward inner loop.
for n = 1, ... DD] do
in Yn + I AT Vi(za, 0.)
D Eq. (6)
end for 7 + Proj [r] end and return Y
4

---

Published as a conference paper at ICLR 2025
Pontryegin's Maximum Principle (Algorithm
Train Data Scorer
Data Scorer
Proxy Dataset ppp
Data Quality Scores y
Inference & Select Data
Uniform Sample
Pre-trein
Language Model
Large-Scale Corpus TD
Selected Corpus
Figure 3: The PDS framework. We compute data quality scores > on a proxy dataset D using Algorithm I, which is derived from the Pontryagin's Maximum Principle (Pontryagin, 2018) (Section 2.3.1). After that, the data scorer learns to predict quality scores from instances, which then infers scores for large corpus D (Section 2.3.2). Finally, high-quality corpus D' is selected based on the inferred scores to pre-train an LM (Section 2.3.3).
Update of y": > is supposed to be updated according Eq. (6) with Of and X computed in the inner loops. Eq. (6) indicates that the new y" should be set one-hot vector, where the element
corresponding the highest Σ X Vl(zn. 07)) value is set to and the others set to O. However, this "hard" update is unstable, as it causes training the LM with only one example in the upcoming outer loop iteration Therefore, Algorithm 1 adopts "soft" update, which increases y AT Vi(a,, 0e) and projects the updated y back into U.
by a value proportional ET- 2
Efficient Implementation. Running Algorithm on Dprx based on the learning of large LM remains computationally intensive, as the inner loops involve training the LM for T steps with GD and computing the Hessian matrix. Therefore, we limit the outer loop to just one epoch and employ stochastic gradient descent (SGD) with small batch size in the inner loops, which is based on small proxy LM with Nprx parameters (Nprx < N) be trained for Tprx steps (T < T). To recover any lost long-range training dynamics, we run Algorithm multiple times by setting Os to the checkpoints at different large-scale pre- vreitraining stages of the proxy LM and then average the obtained data quality scores on Dirx, Specifically, we first train the proxy LM for T steps and save
O(M) in every 1 steps. Then, the quality scores are given by
M checkpoints
Σ PMP-Solver (D- DPIX, , =1=1ph T OD 0 (m)) . - -1)
where PMP-Solver refers to Algorithm 1. We also incorporate several practical optimization techniques to further reduce the computational overhead, as described in Appendix G.I.
2.3.2 DATA SCORER
We fine-tune a small LM as the data scorer to fit the data quality scores 7* on Dptx, Specifically, each instance in Dprx is encoded by averaging the output hidden states of the data scorer. The representation of each instance is then passed through a linear head, outputting scalar. The linear head and the LM are trained together to fit y on D with the Mean Square Error loss:
[DRIA
w. b" arg min
Σ
(w TwTT
h( : ()) + 7:1).
(8)
where & is the parameters of the data scorer and Fe(,.) E Rd is the average output hidden states of an LM along the sequence length, with d representing the hidden state size. E R2, b E IR are the parameters
linear head. After fine-tuning, we infer the data quality scores with the data scorer, where the quality score for v, is given by (xx,) = w Th ( (. x*) + b".
instances D
Thii does not imply that one single example is optimal for data selection due the necessity but not sufficiency of Theorem 2.1. See Appendix FF for detailed discussion
5

---

Published as a conference paper at ICLR 2025
2.3.3 DATA SELECTION
We use the output scores from the data scorer to estimate the value of the instances in D to select the final pre-training corpus D' for the large L.M. Given the importance of data diversity in pre-training LMs, we adopt Gumbel-Top Top- to introduce randomness into the selection process:
D' = Top-K (T(2.) - log( (u.) E D, S [P]},
(9)
where W1 U2, "Ipl are independently sampled from Uniform(0, 1) and 7 is a hyper-parameter control the strength of the Gumbel noise. The size of the selected data is managed by data selection ratio with K = r|D in our experiments.
2.4 DISCUSSION
Effectiveness of PDS. Compared to existing offline data selection methods that curate the pre- training corpus before the LM training starts using pattern information (Xie et al., 2023), deduplica- tion (Tirumala al., 2023; Abbas et al., 2023), or single-step checkpoints (Engstrom al., 2024), PDS incorporates long-range training dynamics into data selection, as reflected by the "target vector" & in Eq. (5). This can be critical for selecting high-quality instances, given the highly dynamic nature of LM pre-training. Although quality scores to the large-scale setting via the data scorer, many previous works (Xie et al., 2024; Yu et al., 2024) have shown that data quality information is learnable and transferable across model sizes. Different LMs also share many common training dynamics (Tirumala et al., 2022).
run Algorithm 1 in proxy environment and transfer the
Efficiency and Flexibility of PDS. Unlike recent online data selection approaches to incorporate LM training dynamics (Yu et al., 2024; Wang et al., 2024), PDS selects the pre-training corpus offiine. and used for pre-training multiple LMs of any sizes, without
This allows PDS to be run only incurring additional computational overhead. The FLOPs needed by PDS in a proxy environment are also negligible compared the demands of large-scale pre-training, analysis in Section 3.3. Besides, the offline nature of PDS makes it flexible to be integrated into optimized pre-training pipelines (Chowdhery al., 2023) by simply replacing the data sources.
shown in complexity
EXPERIMENTS
3.1 EXPERIMENTAL SETUP
Data. We use the CommonCrawl split from Redpajama (Together, 2023) as D to exclude the influence of domain weights (Xie et al., 2024). For the downstream loss ((0), compute the LM's loss on the training split of LIMA (Zhou et al., 2024), a high-quality dataset consisting of 1,030 diverse instruction-responss pairs that cover wide range of downstream scenarios. The results more choices of (0) are shown in Appendix I.5. Our evaluation is conducted on various downstream datasets other than LIMA avoid over-fitting.
Model. We adopt the same model architecture as Mistral (Jiang et al., 2023) and pre-train LMs with 160M, 470M, IB, and 1.7B parameters. Model configurations are detailed in Table 6.
PDS. To compute the data quality scores from PMP, we adopt 160M proxy LM. DeEx consists of 160K instances uniformly sampled from D. We first pre-train the proxy LM on D for 50K steps and then select checkpoints at [10K, 20K, 30K, 1OOK,400 40K, 50K] steps. Initialized from these checkpoints, the proxy LM undergoes inner loops with 77 = 0.008 over T 100 steps with a mini-bater size of 256. Y is updated for one outer loop epoch with Fairseq- -Dnse model (Artetxe et al., 2022) along with the linear head, using the objective in Eq. (8). The training details for the data scorer can be found in Appendix G.2. For Data Selection, we set 0.1, r = 0.4, with further hyper-parameter exploration provided in Appendix I.5.
1. For the data scorer, we fine-tune a 125M
Pre-Training We pre-train all LMs for 100k steps, using batch size of 512 and max input length of 1,024, resulting in roughly 50B tokens. In Section 3.2, we select a 50B token dataset from CC corpus containing 125B tokens to assess how different data selection methods improve LM learning given a sufficiently large D. In Section 3.3 Cont--oont Setting), we also analyze the effectiveness of PDS when D is limited in size. See Appendix G.3 for more pre-training details.

---

Published as a conference paper at ICLR 2025
HS
LAMB
Wino.
OBQA
ARC-e
ARC-c
PIQA
SciQ
BoolQ
Avg.
Model Size
470M
Conventional
36.7
41.4
52.4
30.4
44.8
25.2
61.0
70.6
60.4
47.0
RHO--Lss
42.4
53.0
29.4
43.7
25.2
72.8
59.8
47.0
DSIR
36.4
42.6
51.7
29.8
46.0
24.7
61.0
72.0
55.8
46.7
IF- -core
36.6
41.8
53.4
29.6
44.7
25.1
60.8
68.8
58.7
46.6
PDS
37.9
44.6
52.3
29.8
46.5
25.8
61.8
73.8
61.4
48.2
Model Size = 1B
Conventional
39.9
47.6
52.4
30.6
49.3
26.4
63.1
73.7
60.9
49.3
RHO-Loss
39.8
47.0
53.0
30.8
48.0
26.4
62.9
71.1
61.0
48.9
DSTR
40.8
47.8
53.0
31.2
49.8
26.8
62.7
76.6
58.0
49.6
IF-Score 39.4
47.0
52.6
28.6
49.4
26.4
74.0
60.5
49.0
PDS
42.1
48.8
54.0
33.4
51.3
28.0
64.1
78.5
58.7
51.0
Table 1: Results the downstream evaluation datasets in OLMo (Groeneveld et al., 2024). We report the accuracy scores and the average scores across the datasets. The best scores of each model size are boldfaced. PDS achieves the best performance in cases compared the baselines.
Evaluation. We evaluate the LMs' 0-shot accuracy the downstream test datasets used in OLMo (Groeneveld et al., 2024) and their O-shot performance on MMLU (Hendrycks et al., 2021). I). We also report the LM's language modeling loss on subset of DCLM (Li al.. 2024), high-quality corpus curated with complex pipelines and human heuristics. to verify that models trained on D retain diversity and long-tail knowledge coverage. Further evaluation details are Appendix G.4.
Baselines.
We compare PDS with conventional pre-training and 3 offline data selection methods: Conventional: conventionally pre-training LM on 50B tokens uniformly sampled from also refering "Redpajama" in Figure 1.
RHO-Loss (Mindermann et al., 2022): selecting data with high reducible losses, as in Eq. (55). DSIR (Xie et al., 2023): selecting data with high n-gram overlap with instances LIMA.
IF-Score (Koh & Liang, 2017): selecting data with high influence scores, in Eq. (56).
3.2
MAIN RESULTS
PDS Improves Downstream Performance. We present the evaluation results of the pre-trained LMs on the OLMo evaluation datasets and MMLU in Table and Table 2, respectively. As shown, PDS outperforms the baselines most datasets, achieving the best overall performance across models with 470M and 1B param- eters. See Appendix 1.1 for results 160M LMs. Fig- ure 1(b) shows the scaling curves average accuracy the OLMo evaluation sizes, ranging from 160M to 1.7B, demonstrating that the performance improvement remains consistent as the LM scales These results indicate that the data qual- ity scores obtained in a proxy environment generalize well to various model sizes and downstream tasks.
N
Method
O-shot
PPL
Conventional
27.6
34.8
RHO-Loss
28.4
33.0
470M
DSIR
28.0
34.0
IF-Score
28.4
31.1
PDS
28.9
27.1
with respect the model
Conventional
29.7
26.1
RHO-Loss
30.2
24.9
IB
DSIR
30.0
25.3
IF-Score
30.7
23.6
PDS
31.4
20.5
Table 2: MMLU results. We report the O- shot accuracy and the perplexity (PPL) on the ground truth answers. The best scores of each model size are boldfaced.
PDS Enhances Language Modeling. Besides down- stream tasks, we show that PDS also enhances language modeling on a high-quality pre-training corpus. Fig- ure shows the test losses the DCLM subset for conventionally pre-trained LMs and PDS CoOS-triied LMs with 160M, 470M, 1B, and 1.7B parameters. We can LMs constantly achieve better performance across various model sizes. In see that extrapolate the test losses using the Scaling Law (Hoffmann et al., 2022), showing that Table
7

---

Published as a conference paper at ICLR 2025
Model Size
160M
Model Size 470M
Model Size 1B
Model Size 1.78
Convention PDS (Ours)
Conventional PDS (Ours)
Conventional PDS (Ours)
3.5
3.3
PDS (Ours)
3.2
3.4
3.2
3.1
3.3
25k
50k
75k
100k
25k
50k
75k
100k
25k
50k
75k
100k
25k Pre-Training Steps
50k
75k
100k
Pre-Training Steps
Pre-Training Steps
Pre-Training Steps
Figure 4: Test losses on the DCLM corpus (Li et al., 2024) for 160M, 470M, IB and 1.7B LMs.
N
D
Conventional
PDS
GPT-3 (Brown al., 2020)
175B
300B
2.882
2.872
Llama (Touvron et al., 2023a)
6.7B
1.0T
2.942
2.896
Llama (Touvron et al., 2023b)
70B
2.0T
2.877
2.855
Llama 3.1 (Dubey et al., 2024)
405B
15T
2.851
2.838
Table 3: Test loss extrapolation using the Scaling Law (Hoffmann et al., 2022). We predict the test loss when the LM size N and the trained tokens D meet that of GPT-3 175B, Llama 6.7B, Llama 70B, and Llama 3.1 405B. The improvements of PDS remain consistent for these LMs.
the improvements with PDS persist in pre-training recent large LMs, such GPT-3 (Brown et al., 2020) and Llama family (Touvron et al., 2023a;b; Dubey et al., 2024). Details of the extrapolation are provided in Appendix 1.4. The DCLM corpus is curated using a complex pipeline based on human heuristics and is verified to be diverse and comprehensive in its coverage of human knowledge. Algorithm offers a principled alternative to the complex pipeline for curating pre-training corpus.
PDS Accelerates LM Learning. In Figure (a)) we plot the average accuracy evaluation datasets with respect to the training FLOPPs of the 1.7B model. PDS achieves 2.0 times acceleration in terms of training-time computation compared to conventional pre-training. Similar trends are observed for other model sizes (Figure 8) and DCLM losses (Figure 4).
the OLMo
3.3 ANALYSIS
Setting. In Section 3.2, we assume the original pre-training data is unlimited to ensure that the PDS-selected corpus contains 50B tokens, the same as conventional pre-training. In practice, when the data before selection is limited, the LM has to be trained on the PDS-selectes data for multiple epochs. Figure 5, we restrict D to contain 50B tokens and apply PDS with selection ratios r E [0.125. 0.25, 0.5). corresponding to training the LM on the selected data for (8, 4, 2] epochs, respectively. We can see that selecting 1/4 data with PDS and training for epochs achieves the lowest test losses, consistent with the findings of Muennighoff et al. (2023). This suggests that PDS improves data utilization as the high-quality web-crawled corpora run out (Villalobos et al., 2022). We extrapolate the loss curve of conventional pre-training with the Scaling Law suggested by (Muennighoff et al., 2023), and find that requires another 42B tokens to achieve the performance of PDS (r 0.25, 4 Eps). which means PDS reduces the pre-training data requirement by 1.8 times. See Appendix 1.4 for details of this extrapolation.
Efficient Implementation for Solving Data Quality Scores. Section 2.3.1, we introduce efficient implementations to solve the data quality scores by running Algorithm 1 for single outer epoch, leveraging a small proxy LM trained for a limited number of steps. In Figure 6, we use simulated setting where " can be obtained exactly from Algorithm I within affordable computational budget and compare its performance (measured by J(e,)) and overhead (measured in FLOPs) with our efficient implementation. Details of this simulated setting be found in Appendix G.6. The results show that the efficient implementation significantly reduces computational overhead while preserving most of the performance of the exact solution.
8

---

Published as a conference paper at ICLR 2025
1pp
Conventional
8.5
Conventional
3.40
PDS 0.5, 2 Eps)
PDS (Efficient, -10 FLOPs)
Test Loss
PDS 0.25, 4 Eps)
8.0
PDS (Exact, ~100a FLOPS)
PDS (r= 0.125, 8 Eps)
7.5
3.35
Ep.
Ep
7.0
Ep
0
1000 2000 3000
4000 5000
3.30
4E
Training Steps t
12.5B
25B
37.58
50B
Trained Tokens
Figure 6: Comparison between exact and effi- cient implementation to solve the data quality scores in simulated setting. The effectiveness is measured by J(0,). The efficient implemen- tation saves computation and preserves most of the performance of the exact solution.
Figure 5: Test losses on DCLM corpus (Li et al., 2024) in the data-constrained setting. We select data with PDS for different selection ratios r and train the model for multiple epochs to reach the same token number budgets.
Complexity
FLOPs (x1020)
Actual Time
Proxy 7-soover
O(Npex D + 4M NPrx DPEx)
0.49
15.2 Hours
PDS Data Selection
Data Scorer
(33Nstor D + Nscore D) O(D)
0.063
1.50 Hours
0.0
10.2 Minutes
Pre-Training
O(ND)
5.1
144 Hours
Table Asymptotic complexity, GPU FLOPs, and actually spent time of different PDS steps and 1.7B model pre-training. N and D are the sizes of the proxy LM and proxy data. Nsscor is the size of the data scorer. We claborate on the details of how the complexity and FLOPs are computed in Appendix H. PDS costs little computational overhead compared to pre-training large LMs.
Complexity Analysis. We compare the computational complexity of PDS with pre-training in Table 4. The overhead of running PDS to select data is only about 1/9 of that of pre-training a 1.7B model. More importantly, since PDS is an offline method, the selected corpus can be used pre-train any number of LMs without additional computational cost. In addition, the offline nature of PDS allows it to seamlessly integrate into recent highly optimized pre-traine pre-training pipelines (Chowdhery et al., 2023), requiring only replacement of the data source without altering the pre-training process.
3.4 ABLATION STUDIES ON PMP-SOLVER
Training Dynamics Information. Incorporating the LMs' training dynamics into data selection is key distinction between PDS and other offline data selection approaches. While IFoSe also uses the gradient information of well-trained LMs to estimate data importance, we find that the long-range training dynamics in early training stages are valuable. Table 5 shows the results when different types of learning information are considered. PDS (50K) refers to using the LM checkpoint at 50K as mm)) in Eq. (7). PDS (Tprx 1) means running the inner loops as in Eq. (4) and Eq. (5) for only
one step, excluding long-range information. PDS (50K- 100K) refers to setting 0 (m)) to checkpoints later training stages. Comparing our choice with IF-oee PDS (50K), and PDS (Trex 1), we conclude that the long-range training dynamics is more valuable than single-step gradient, although it may be slightly harder for the data scorer fit. Our choice also outperforms PDS indicating that the carly-stage training dynamics are more beneficial than those from later stages. An explanation is that LMs conduct large-range parameter searching during early training and converge to local optimums in the late stages. Therefore, early-stage information helps the LM choose better local optimums, which can always be achieved by later annealing.
Proxy Model and Proxy Data. In Figure 7, we explore how the sizes of the proxy LM and the proxy dataset affect the performance of the data scorer and the pre- pre-trained LM. As the size of the proxy LM increases, the LM's downstream performance increases, but the data scorer's performance

---

Published as a conference paper at ICLR 2025
Corr.
Acc.
Average Accuracy (%)
45.00
0.60
Spearman Correlation
Conventional
43.2
44.75
Acouracy Conelation
0.58
IF-Score
0.32
43.0
44.50
[200g] = 160K
PDS (50K)
0,51
44.0
[000]
320K
0.56
PDS (TPxX=1)
0.54
44.6
44.25
PDS (50K-100K)
0.48
43.4
0.54
44.00
PDS (10K-50K, ours)
0.52
45.0
8.7M
20M
66M
160M
Size Proxy Model
Table 5: Data selection based on different kinds of learning information. We report the LM's zero-shot accuracy on the OLMo evalua- tion tasks (Acc.) and the Spearman Correlation of the data scorer (Corr.).
Figure 7: LM performance on the OLMo evalu- ation tasks (Average Accuracy) and data scorer performance (Spearman Correlation) when differ- ent proxy model and proxy data sizes are adopted.
decreases. We notice that when using the 8.7M model, the LM's performance (44.0) is close to that of DSIR (43.8), which selects data based on n-gram matching. This implies that small LMS estimate data quality primarily through shallow patterns that are easy to learn. More complex LM learning information is encoded in the data quality scores computed based on larger proxy LMs, making harder for the data scorer to fit, but this can be mitigated by increasing the size of DPrx
4 RELATED WORK
Data Selection for Language Models. Many works have explored data selection approaches to accelerate LM learning or improve downstream performance (Albalak et al., 2024). Some curate data prior to LM training, which refer offline methods, including domain-mixing (Xie al., 2024; Fan et al., 2023; Gao et al., 2020), data pruning (Marion et al., 2023; Tirumala et al., 2023; Abbas al., 2023), sample-wise data selection (Xia al., 2024b; Du al., 2023; Xie al., 2023; Gu et al., 2023), and data programming (Ratner et al., 2016; Gu et al., 2022a;b). Other works dynamically select data during LM training by adjusting domain -ixing weights (Xia al., 2024a; Chen et al., 2024; Albalak et al., 2023) or more fine-graines reweighting strategies (Fan & Jaggi, 2023; Grangier al., 2023; Gu al., 2024; Thakkar et al., 2023). This work studies data selection from its general principles, theoretically deriving optimal selection and designing scalable algorithms to implement it.
Optimal Control in Deep Learning. The principles of optimal control (Lewis et al., 2012) have been shown be effective in deep learning (Benning al., 2019; Liu & Theodorou, 2019), by treating the forward pass of a multi-layer neural network as a control process where the hidden vectors are state parameters and the model parameters are control variables to optimize. With the help of Pontryagin's 's Maximum Principle (Pontryagin, 2018), some works design optimization algorithms with better convergence rates (Li al., 2017; Zhang et al., 2019) or broader application scenarios (Li & Hao, 2018), and others provide theoretical foundations of neural networks for better interpretation (Han al., 2019; Geshkovski & Zuazua, 2022). Unlike these works, we adopt Optimal Control in a novel and "orthogonal" way, by treating the model's learning pass as the control process.
5 CONCLUSION
In this work, we investigate selecting pre-training data of LMs from both theoretical and practical perspectives. We first formulate data selection as an Optimal Control problem to derive set of necessary conditions for optimal data selection using Pontryagin's 's Maximum Principle (PMP), which establishes theoretical principles for LM pre-training. Then, building on these theoretical results, we introduce PDS, practical framework that solves the PMP conditions in practice based long-range LM training dynamics. Extensive experiments show that PDS selects high-quality pre-training corpora that accelerate LM learning and boost LM performance across various model sizes and downstream tasks. We find that PDS also improves data utilization in cata-constrained settings, which mitigates the pre-training data exhaustion problem.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work supported by the National Science Foundation for Distinguished Young Scholars (with No. 62125604) and Tsinghua University Initiative Scientific Research Program.
REFERENCES
Amro Kamal Mohamed Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari Morcos. SemDeDup: Data-efficient learning at web-scale through semantic deduplication. In ICLR 2023 Workshop, 2023. URL https: / openreveiiie
net/ forum? 1 Gm9gv6c.
Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. JMLR, 2017. URL https: /ww.. jmlr orq 166-911/ 6- 4-1.. pdf.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training. NeurIPS workshop, 2023, URL httpt: https::/openreviiii net/forum?ie 2id 4y
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, et al. A survey data selection for language models. TMLR, 2024. URL htttps://ppereviitt ner/forum? i HWcNTHH..
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, et al. Efficient large scale language modeling with mixtures of experts. In Proceedings EMNLP, 2022. URL https: http:/:///claaahhhoooo.oo
org/2022. emnlp 1p-main main. inn.inn.. 804/.
Heinz H. Bauschke and Patrick L. Combettes. Convex Analysis and Monotone Operator Theory in Hilbert Spaces. CMS Books in Mathematics. Springer, 2011. URL htttps://doi.oott
1007 /9 1119998811-11 9-19-11111 67-77.
Martin Benning. Elena Celledoni. Matthias Ehrhardt, Brynjulf Owren, and Carola-Bibiane Schönlieb. Deep learning as optimal control problems: Models and numerical methods. arXiv preprint 1955051904.05667,, 2019. URL https: /ttp/a/axvvvv.////v//ee////ae.va wrg. abs// 1904. 05657.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek I1m: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.022 2024. URL https::/aaxiitii otgeiiioorrraaaae abss/ 24C1. 02954.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of AAAI, 2020. URL htttss://jss. aaai. orggindex.phiee AAAI AAAI/AAttatt....A /aaaaaicceevviieeeeeeeeeeeeeeeee 6239/ 6095.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, al. On the op- portunities and risks of foundation models. arXiv preprint arXiv:2108.0722 2021. URL https::/arxiss org/ hatp//////// 2208 07258.
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnn and Qiao Zhang. JAX: composable transformations of Python+NumPPy programs, 2018. URL
http:// /ttp:///gtthh////// com/ hoom/googee/////// jax.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, et al. Language models are few-shot learners. In Proceedings of NewrIPS, 2020. URL https: https::/paapeee http:/ppppppss.pp.ee.. cc/ bo/paper/ 2020/ 00000hhah//145cc06555 9674c44499914111
10777488fbbaae wwwbbaclllcc... Sefcc142f4440 Abst html.
Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-it! a data-driven skills framework for understand- ing and training language models. In Proceedings of NeurlPS, 2024. URL httpsss/pproceeeinss...eee http:ppp...pppppppeeeeeee cc/paper http://eeeeeceae.o...../eeeeeeep
11

---

Published as a conference paper at ICLR 2025
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06114, 2016. URL https:t//////ctt http:////xxxxv.oooo.e abss/
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, et al. PaLM: Scaling language modeling with pathways. JMLR, 2023. URL http:
///mlr. wwwprg/ppprs/vv22 2222222222222 1144. html.
Christopher Clark, Kenton Lee, Ming-Wai Weig-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova, BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of NAACL-HLT, 2019. URL https: / anthology org/N1 //1/ 9-030.
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 challenge.
18X3.0800.054455 2018. URL htttts://aaate
org/ abs/ 1803. 1803.05457.
Corinna Cortes and Mehryar Mohri. AUC optimization vs. error rate minimization. In Proceedings of NeurlPS, 2003, URL https: / httpt: htttpp::/ppoceeedingg neurips wttp:/ttttppppppppppeaa/a/eaeeee cctpaper//003/ //aash/ 6ef8 (2880bb2377200 ad / Abstraaat........tt........aaa.. html.
Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How to compute hessian- products? In ICLR Blogposts 2024, 2024. URL https:::/acrr-booopoootttoo gitthuu 1o/
2024 2.024/blee /gg og/ httpbbbnnchhhppp
Qianlong Du, Chengging Zong, and Jiajun Zhang. Mods: Model-oriented data selection tuning. arXiv preprini arXiv:2311.15653, 2023. URL https::/arxii..oott /2311. 15653.
selttion for instruction
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama herd of models. arXiv preprint arXiv:2407.21783, 2024. URL https: www.aaaivv.... oww.iiioorr//ooooooo org/ abs///////a//// 2207 21783.
Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with datamodels. arXiv preprint arXiv:2401.12996 2024. URL httttttts///xiiiiie
Lawrence Craig Evans. An introduction to mathematical optimal control theory. University of California, 2005. URL / www.th.brreeee..eeeeeeeeeee conraansccoorllll.oo contaao.ooooooo..sssee...... course. pdf.
Simin Fan and Martin Jaggi. Irreducible curriculum for language model pretraining. arXiv preprint
arXiv:2310. 15389, 2023. URL https: /ttp///aavvvv//v///////e.///avvv org/ abs//3111 2310.15389.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. DOGE: Domain reweighting with generalization
estimation. In Second Agent Learning in Endedness Workshop, 2023. URL https: //openreview net/ forum?
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0002) 2020. URL https: h/tp:////xxxvoo/o///oo abs/2101. 00027.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman. Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, et al. A framework for few-shot language model evaluation, 07 2024. URL wwt.:/ooooooro/ooooooooooooooeot recoors//26100 http:///enoooooorrrrrrr///////// 602.
Hans P Geering. Optimal control with engineering applications, volume 113. Springer, 2007.
Borjan Geshkovski and Enrique Zuazua. Tumpiik in optimal control of pdes, resnets, and beyond. Acta Numer. 31:135-2633 2022. doi: 10. 10.1017/5099249992220000 10. 1.0.1017//0 0tl.1117750962440 0224499222000042
URL hattps://of org/
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanic Tellex. Openwebtext corpurs 2019. URL http:///kkylon0000 .i thub. o
12

---

Published as a conference paper at ICLR 2025
David Grangier, Pierre Ablin, and Awni Hannun. Adaptive training distributions with scalable online bilevel optimization. arXiv preprint arXiv:2311.1193,, 2023. URL
Abs/2331.11973.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. OLMo: Ac- celerating the science of language models. arXiv preprint arXiv:2402.00838, 2024. URL
bttps://arxer wwwpxivv.orr/eeooo orgo als////////////////////// 2202. 00838.
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. arXiv preprint (rxxi:2338..3229 2023. URL https::////// /ttp:////axxvvo//////o abs/2308 03296.
Yuxian Gu, Xu Han, Zhiyuan Liu, and Minlie Huang. PPT: Pre-trained prompt tuning for few-shor learning. In Proceedings of ACL, 2022a. URL https://axxxxi wwwp:/iwwwiiiooiiioooooooooooore abs///sss/111//// 04332.
Yuxian Gu, Pei Ke, Xiaoyan Zhu, and Minlie Huang. Learning instructions with unlabeled data for zero-shot cross-task generalization. In Proceedings of EMNLP, 2022b. URL https://// /tcttp:///rxiv org/abs 2210. 09175.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Pre-training to learn in context. In Proceedings of AC,, 2023. URL https: /aacl aoaatoo org/ acl-long.266 267/.
Yuxian Gu, Li Dong, Yaru Hao, Qingxiu Dong, Minlie Huang, and Furu Wei. Towards optimal learming of language models. arXiv preprint arXiv:2402.1775, 2024. URL https: org/abs /2402 17759.
Jiequn Han, Qianxiao Li, al. A mncln-field optimal control formulation of deep learning. Research in the Mathematical Sciences, 2019. URL https: ://ink/ springes com/ article/10.
(007/440687-011- 012-yy.
Xu Han, Zhengyan Zhang, Ning Ding, Yuxian Gu, et al. Pre-trained models: Past, present and future. AI Open, 2:225-250, 2021. doi: ATT.://...PE........... 10. 10..116/j.aa. aiwpen..........ooe.oo.o 2021.0222222222 08 002.
URL https://doi. org/
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In Proceedings of ICLR, 2021.
URL https:t//////. htap::://pppeee hwtp//eveeeeeeeeeeee/ee net forrmmormmm id=d7kBBBBB jmI 3GmQ.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, et al. Training compute-optimal large language models. In Proceedings of NeurlPS, 2022. URL https::/prcceednns.s.eeo neurips cc/paper fneer_fflese files//app::/2////i. clezaafff68887733511eeee08338 Pasissee----pe----
oorpraaeeeee--ppp-----oeeeeeeeeo Conference.
Peter J. Huber. Robust Estimation of Location Parameter. Springer New York, 1992. doi: 10.1007/ 111118-----11---331 1111111-------
978-1-4612-4380-9J33 35.
URL https: http:///o...//111100//////1////0 /
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b, arXiv preprint ArXiv:2310.068// 2023. URL https: pettps://arxtt //tp/aaaivv/////////////////aall wwe.xiivorr///oo orbbas//////////oooooooo/../ss /2310. 06825.
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 preprini arXiv:2001.088: 000361, 2020. URL
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of ICLR, 2015. URL httttpss//pppp/e///// nwt. foot/ffoum?? r d-8gmWwjFytt.j
13

---

Published as a conference paper at ICLR 2025
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In Proceedings of ICML, 2017. URL
koh17a. 1.n.17..pdl pdf.
Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Proceed- ings of KR, 2012. URL https: /ttp//dl..// acm. wwtp:///goooooooooo//o//oo//ottt 5555/ 30555555/311550 3031909,
Frank I. Lewis, Draguna Vrabie, and Vassilis I. Syrmos. Optimal control. John Wiley & Sons, 2012.
Jeffrey Li. Alex Fang, Georgios Smymiis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, et al. DataComp-LM: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406. 1794, 2024. URL https: //arxiv. orgg abs/ 2406 11794.
Qianxiao Li and Shuji Hao. An optimal control approach to deep learning and applications to discrete-weight neural networks. In Proceedings of ICML, 2018. URL httL:tttp///prcceeiiiei
mlr. wtl.spress/vvl v80/
1118b. html.
Qianxiao Li, Long Chen, Cheng Tai, and Weinan E. Maximum principle based algorithms for deep learning. JMLR, 2017. URL http: fmmrr. org/ http//ppprr//////e vvt//v///////// 17-67555 .tml.
Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang. Jian Jiao, Nan Duan, et al. Rho-1: Not all tokens what you need, arXiv preprint arXiv:2404.07965, 2024. URL https: / aawaaiiivvvoo.m wwtp:/ii.oorraeooooo wttppbg/g/.a//a//////////// 07965.
Dong C Liu and Jorge Nocedal. the limited memory bfgs method for large scale optimization. Mathematical programming, 1989. URL https: :ttpps://lnnn springer. com/ hatm///rrieee/0 1.007//FF F1 0199111811 6.
Guan-Horny Liu and Evangelos A Theodorou. Deep learning theory review: An optimal control and dynamical systems perspective. arXiv preprint arXiv: 1908. 10920, 2019. URL https:
//arxiv. org/ abs//bbs///99 10920.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprint 1rXiv:1907.1169) 2019. URL https: /aapaaxxvvv.rgg. http://////axv..orrr///////1111 11692.
llya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of ICLR, 2019. URL https: http:///ppeeeeeeeeeeeeeeeee http://wwww....eeeeeoooooooeoeaa
Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang, Marzieh Fadace, and Sara Hooker. When less is more: Investigating data pruning for pretraining LLMs at scale. In NeurIPS Workshop on Attributing Model Behavior Scale, 2023. URL https:tp///..mc http:p:://opeeeeeeee wpwpeevieeeeeeeeeeeeeeeeeeeeeeee net forum?id-
SUIYn3jo511
Luke Metz, Niru Maheswaranathan, C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture. and compute: Training more effecttive learned optimizers, and using them to train themselves. arXiv preprint arXiv:2009.11244 2020. URL https: /arxiv.iiiv////i///aaaaaaaaaavva org/ abs/ 2009. 11243.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of EMNLP, 2018. URL https:ttp///.mm //pi. www.etiisshhhootoooooooo
org/ Cttp//orpussDD 5118:7527777
Sören Mindermann, Jan Markus Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N. Gomez, Adrien Morisot, Sebastian Farquhar, and Yarin Gal. Prioritized training on points that are learnable, worth learning, and not yet learnt. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), Proceedings of ICML, 2022. URL ptttps://proccedinng mlr. mi html.
14

---

Published as a conference paper at ICLR 2025
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling contttcod language models. In Proceedings NeurlPS, 2023. URL
j33UTrR
OpenAl. GPT-4 technical report. 2023. URL
0877
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of ACL, 2016. URL https: 11tp//16-1111 pdf.
//aclaachologyy
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, al. PyTorch: An imperative style, high-performance deep learning library. In Proceedings of NeurlPS, 2019. URL https: httptthtt:///eeee /ttp///eeoceeeeeeeeee/eeoo neurips cc /eeppere/////////// 2019 2ttp://aaaaccccc2222eee000022220
-bsttt0tbbbaatttt..aa..aaaa html.
Lev Semenovich Pontryagin. Mathermatical theory of optimal processes. Routledge, 2018. URL
https://doi.e. /ttp//0..111111100751523023
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In Proceedings of SC, 2020.
Alexander Ratner, Christopher De Sa, Sen Wu, Daniel Selsam, and Christopher Ré. Data programming: Creating large training sets, quickly. In Proceedigns of NeurlPS, 2016. URL https::/proceedingg neurips httpiipp.ccppppppeeeeee htcppaper/2016/hast html.
67C9e3d64a5E4722 0rr.:662247722255eeeee22222-----
Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
URL https:tp///.m. /tsttps://ppeee http//ppeeeeeeeeeeeeeeeeeeeeeeee net. wtwpiiewww.ee//ounttt..mm forum? id- IdwIIIWIdIIWDD
Atle Seierstad and Knut Sydsaeter. Optimal control theory with economic applications. Elsevier North-Holland, Inc., 1986.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Be- yond neural scaling laws: beating power law scaling via data pruning. In Proceedings of NeurlPS, 2022. URL http: /tatt:::/pppprss.t nips oe wwwppppprrraeeeeee lesilee/eeee httppeer//pppp/22//// / 7b75da9b61eda
html..
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv: /8005, 2023. URL htttts:///aaeetee org/abs/
Megh Thakkar, Tolga Bolukbasi, Sriram Ganapathy, Shikhar Vashishth, Sarath Chandar, and Partha Talukdar. Self-influence guided data reweighting for language model pre-training. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of EMNLP, 2023. URL https:
//aclanthor loy.ooggy.... www.y.org//000 orgg/222. mmmlp-main.. 125/.
Kushal Tirumala, Aram Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. Memorization without overfitting: Analyzing the training dynamics of large language models. In Proceedings of NeurlPS, 2022. URL https:////// /ttpes//oppeeeveeeeeeeee net forum? d=u3vEuR
Kushal Tirumala, Daniel Simig, Armen Aghajanyan, and Ari Morcos. D4: improving LLM pretraining via document de-duplication and diversification. In Alice Oh, Tristan Naumann. Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Proceedings of httpapprrssllse/pppp/////0
NeurIPS, 2023. URL http:// http://ppppsssssssee hash/a 8 (efffccd77. o Absttaatt-aattsss__mm Benchmarks html.
Together. Redpajama: an open dataset for training large language models, October 2023. URL http://gt..h...cooo//oggeeccroom www.m/tttteeeecetoteeeerer jama-aaaaaa-aaa
https: petpss:/githhe thu.. com
15

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix. Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a. URL https: /ttpaaaiiiv///...//aaa bttettps//aaxii.oorr/aabeaee/e //b 2300.. 13971.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092220 2023b. URL https: wtt.tttss///axiiv //axiv.. org/ has//bb/// 2307. 09288.
Pablo Villalobos, Jaime Sevilla, Lennart Heim, Tamay Besiroglu, Marius Hobbhahn, and Anson Ho. Will we run out of data? an analysis of the limits of scaling datasets in machine learning. arXiv preprint arXiv:22/ 04325, 2022. URL https: // //wrxvvvvvvv////////////////vvvv wrg/abs/////oooooo.oo.o.ooee...r 2211 22211.044 04325.
Jiachen T Wang, Prateek Mittal, Dawn Song, and Ruoxi Jia. Data shapley in one training run. arXiv org. http:/bb//24/00 11011.
preprint arXiv:2406. (tt:::2206.111114 1001, 2024. URL
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User-generated Text (ACL 2017), 2017. URL
https: /ttp:///allaaaatcc wog.l www.orgg/oo....oo //17/1111 4413.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre-training via structured pruning. In Proceedings of ICLR, 2024a. URL https://
openreview.o wwtp:/eedtppddddd///dde
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS: Selecting influential data for targeted instruction tuning. In Forty-first International Conference on Machine Learning, 2024b. URL https: http///pppeeeeeeeeeeeeeeeeeeeeee net/ forum? id= FV5 0maR.
Sang Michael Xic, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Proceedings of NeurlPS, 2023. URL http:: 3S0008241-BB22200050500202222-22
cttpp/srccppppeeeeeeeeeeeeeeee firerfiielss/eeee http://pppppe////eeeee html.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. DoReMi: Optimizing data mixtures speeds up language model pretraining. In Proceedings of NeurlPS, 2024. URL net/forum? 1edidddxx0 wdddxdxyyUHHhhh..
Zichun Yu, Spandan Das, and Chenyan Xiong. Mates: Model-aware data selection for efficient pretraining with data influence models. arXiv preprint arXiv:2406.0600 2024. URL https:
//arxiv.oort abs/ 2406. 06046.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of ACL, 2019. URL
org/P19- 1472/.
Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong. You only propagate once: Accelerating adversarial training via maximal principle. In Proceedings of NeurlPS, 2019, URL https: //ppc//ed/ioc/s//sseoo/ooe./ neurips www..iss...paaper//e..e C12btba2877Seebbc9dde wwweeeeedeeebdebbeebeeeee
wwtppppee//11//aaaeeeeeeaeeeeeee 2220099/hahh Abstract htmll.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer. Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, al. Lima: Less more for alignment. In Proceedings of NewrIPS, 2024.
URL https::////. bettps://pperreee /ttp/////p//////e/////pp// net forum? id KBMOKmX2he.
16

---

Published as a conference paper at ICLR 2025
A CONNECTION BETWEEN AUC AND SCALING LAW CONSTANTS
We show that the area under the loss curve (AUC) is directly connected to the scaling law (Kaplan al., 2020) constants, and minimizing AUC essentially improves the scaling properties of LMs. As suggested by existing literature (Kaplan et al., 2020; Hoffmann et al., 2022), the test losses of LMs follow power-law with respect to the training steps after warmup stage:
C L(t) + L irre t> To.
(10)
where and c are scaling law constants, Lirre the irreducible loss related to the noise in the test set, and To is the end steps of the warmup stage. Lime is invariant to optimizing pre-tran pre-training data selection strategies. Therefore, we care about the reducible loss, whose constants depend on the data quality scores
L (t) = L(t) - - C(r)
t> To.
(11)
We then consider the AUC of the reducible loss for sufficiently large T:
C(7) dt C(r) (Tl-c(7) - To T))d))).
AUC(7)= T= (<(7) 1 1-y())
(12)
For c(7) < 1, when 7' is sufficiently large, AUC(n) ss C(r) to decrease and c(>) to increase improving the scaling properties LMs. For c(7) > AUC(7) ~ C(a)
I- 1
C(r) T)-(()) Minimizing AUC causes
<[>]- (o)--1m60)-1 T- which also exhibit a trend of decreasing C'(y) and increasing c(t) when AUC is minimized.
B PROOF OF THEOREM 2.1
To prove Theorem 2.1, we first describe the standard discrete-time Pontryagin's Maximum Principle (PMP; Pontryagin, 2018) in Optimal Control (Lewis et al., 2012) for time-variant control variables: Theorem B.1 (PMP), Consider the following optimization problem in discrete dynamical system:
T-1
min 74
Σ
J(e,, 71) + (10 J(0r).
(13)
1=
S.t. 01+1 = f(0.. Yt EU
where the state variable 0, E R the control variable Tt RD, and J. RND + R, f NxD + R. are continuous XxD Let 7i be the solution to this problem, and 0 denote the corresponding state variable. For 0 < t < T, there exists co-state vector X; E R. such that
01 V (O;; X-11. 0% = Oo:
(14)
A; V H (.-177), Ni = V
(15)
7 arg min H X411 Yt E U,
(16)
7:
where R.V RNN R + R is the Hamiltonian function defined by
H(0,A.7) = J(0,)) X f(0.r).
(17)
Proof of Theorem B.1 can be found in various texthooks in Optimal Control (Lewis et al., 2012: Evans, 2005). In our context, we interpret O, as the model parameters, J(-) as the downstream loss function, as the [D)-dimensional simplex as defined in Section 2.1 and f as the GD operation where the data weights 7 changes with respect to the training steps t:
D
0 = 0 v
(n, 0.)
(18)
n=1
0, 1E--VLCC L 7c).
*f(c) = LTTl- is increasing with respect to e when - c < In T., which is easily satisfied for sufficiently large T.
17

---

Published as a conference paper at ICLR 2025
In this way, the Hamilton function is
H(0,X. J(0m)) + A 7/VL(0, -
(19)
The key challenge to prove Theorem 1 that the control variables are constrained be invariant of the training step t in data selection, as discussed in Section 2.1 and 2.2, and introducing more constraint usually makes an optimization problem more complex. Formally, the requirement of invariant data weights can be expressed by T equations:
71 = 70, = Yo. Yr-1 = Yo-
(20)
Therefore, the optimization of data selection, as in Eq. (3), is equivalent to the following problem:
T
min
J(O.),
71
1=1
(21)
s.t. 0,+1 0 - V), 7t); E U,
Yo 7/1 -... Y--1-
We adopt the method of Lagrange multipliers solve Eq. (21) by considering the following opti- mization problem:
T-1
T-1 D
min
Σ J(O,) + Σ Σ f=0
Aanelmm--haai - + J(Or),
(22)
==1
s.t. 0++1 = 0, - 1) L(,,, 7), Tt E
where (Jn..) are Lagrange multipliers. Note that we split J(Or) out and add J(0p) to the sum of J(0,), which does not affect the minimization. In this way, when J(0,, Ye) takes the following form:
J(0O)
Nn 0n
if
#=1 n=1
(23)
P
J(O.)
if Sttt--1
can apply Theorem B.1 to Eq. (21). By substituting Eq. (19), Eq. (20), and Eq. (23) into Eq. (14) and Eq. (15), we have:
0; = 0 - VL( 76), 0i =
(24)
A; = Xit1 V)(0,) - ((2 N1 vJ(0r),
(25)
which prove Eq. (4) and Eq. (5) in Theorem 2.1 when we set 7 7" By substituting Eq. (19) and Eq. (23) into Eq. (16), have
D
arg max Yo
nn.0
A
Vllnn 06) + Σ Uny
i
n'1l
(26)
D
arg max
Tnat
Nit1
Vi(xn,
([)
Hn,e
iflstsT-1
Considering the time-invariant constraint in Eq. (20), we set 70 = Y" =... 7ー = " and get
D
arg max
I'm
ai
Vi(xn, 03) + Σ Hayt
ift=0
t'1l
(27)
D
arg max
miiti
Vi(In, 0,) - Mnat
iflsIsT-1
18

---

Published as a conference paper at ICLR 2025
which forms complete equation system containing T equations and T unknown variables (T number pt = [M1,,, M2,1,' , plus one y²), which has the solution:
1
Sn
Pn. = Mitt1 Vi(xn: 0,) -1 "T T
is 0
(28)
1
Sn T
7 = arg max
(29)
1=1
where S, Ei= 141 Vi(inn 07). Note that Eq. (29) equivalent to Eq. (6). far, Theorem 2.1 proved by combining Eq. (24), Eq. (25), and Eq. (29).
C DERIVATION FOR ADAM
We provide our formulation and derivation for Adam (Kingma Ba, 2015) in this section. For 0 s t < - L, the parameter update rules of Adam is given by
m4+ m, + (1 - B1 )))
Ut+1 8202 + (1 B2)9
m m4+1/ (1 3/+11;
(30)
6 0t+1/(1 - 1/-1111
0.11 = 0, -1jm++1/ 1 (ov Tz+1 + 6)2
where 31, $2, e, 1) are hyper-parameters and g: = VI(O,, 7t. D). We set = 0 and vo = 0. To formulate the LM training with Adam as an Optimal Control (Lewis al., 2012) problem, we treat
the vector θ = [6., m., v.]] E RaN as the state variable. Let denote the sate-transition function from + 1
Oct1 = F(O,,))
(31)
and thus F represents the following relations:
Bim: + (1 B1)g:
0.+1 =0, - 8(+
(32)
V
(B20r + (1 - 32)92)/(1 -8ht1) +
8im, (1 B1)g:
m+41
1 - 187+1
(33)
820, + (1 - B2)g
U1+1 1 g5+1
(34)
Similar to GD, we can still define the Hamiltonian function of the problem and obtain the following theorem with Pontryagin's 's Maximum Principle (PMP; Pontryagin, 2018):
Theorem C.1 (PMP Data Selection for Adam). Let Y" solve the problem Eq. (3), and 0 denote the state variable corresponding to ry" Then, there exists co-state vector A E RR3N such that
Ø VAH(O,, 7)),),,, 0 = =0o. 0.00
(35)
A; VoH(O,, A ATr))"), AG [VJ(0r).0,0'
(36)
7 arg min Σ H(O; Ain & E U.
(37)
t=
where H: R3N R3N
RDD
R is the Hamiltonian function defined by
H X, J(0) + A (TF F(O,7).
(38)
Similar to the derivation for GD, Eq. (35) controls the state transition, equivalent to Eq. (31). To dutt1
simplify the further derivation of Eq. (36) and Eq. (37), assume
0, which is reasonable dg.
19

---

Published as a conference paper at ICLR 2025
Algorithm 2 PMP Solver for Adam
Input:
LM learning rate 77. Outer loop learning rate a. Outer loop epochs T, Training data D. Downstream loss function J-)). Training steps T'. Function Proj[:] that projects a point in Rº to the D-simplex.
Output: Data quality score y
( 0.0
Y [7. 72. .
Ou
DI
D
[D]
repeat To times
Outer loop
for / 0, ,
D Forward inner loop Eq. (35), expanded to Eq. (32-33)
0 -- VAH(OO:,A+1:)))
end for
Ar
[ 0)
for / = 7 1 ... do
Reverse inner loop Eq. (36), expanded Eq. (39-42)
A:
end for for
1,2, ; do
In E Yn + aaV 7/ 1.(O, 7
> Eq. (37), expanded to Eq. (43)
end for
4- Proj [y] end and return y
because 01+1 an exponential moving average of of and the weight 33 usually much smaller than practice. Therefore, we have der+1 0ri11 = 0 and thus Eq. (36) can be written to J0,
0,
a%..
A; = A (2) A 1
(39)
(1) 0 (VJ())) + AA0) (1 (1-31)m (700),)""
i+1
1-
V
04+1 +
+ (1 (1- B1)
( AA)
(40)
8741
A
A (2)
n1
31
(2)
1 8441
Vōtt +e
1-
A
i+1
(41)
A (8)
T13
(1)
B2
(3)
(42)
8
V Dtt1
(Votti+e)
3 2++
H+1
To achieve the minimum in Eq. (37), need compute the gradient of H with respect to 7:
71 81)
V(( 0.)
1- 81
V.,,
1- - Att1 0 +
(1)
+
1-
1-+
A
A41
(2)
Vi(zn,
0,)
(43)
In this way, we can use Algorithm 2 to solve for the data quality scores on the proxy dataset Drrx just like Algorithm in Section 2. Then, we train a data scorer with the solved data weights y as in Section 2.3.2 and conduct data selection based on the scores inferred by the data scorer on D as in Section 2.3.3. In pilot experiments, we find that computing data quality scores based on Adam does not show substantial improvement over that based GD. Given that Algorithm 2 requires twice as much computation and times as much GPU memory as Algorithm 1, we adopt PMP condition for data selection based GD (Theorem 2.1) to build PDS in our main paper.
20

---

Published as a conference paper at ICLR 2025
D ALGORITHM 1 AS PROXIMAL GRADIENT DECENT
We provide another view of Algorithm using Proximal Gradient Decent method (Bauschke & Combettes, 2011). Specifically, we can optimize Eq. (3) with the following rules:
A(r) Σ J(Or),
(44)
7 Proj by a A
(45)
where A(7) denotes the cost function in Eq. (3), a is the learning rate and Proj[-] projects a point in E. to the D-simplex. V,A=
is the gradient of A(y) with respect the data weights Now, compute each element of V 7 A with the chain rule:
ByD
0A
0J(Or)
8%
0m
==1
:0. (O,)
8%
00, 00,
(((O..)
00, On
F=1
(46)
T
00,
V((02)
T
Vl(an, 0c-1)
00,
00,
VJ(Ov)
Vi(InZ O₂)
00,+1
00,
= (O(0o))
Vi(In, 0c).
00.-1
where 20,- 30.41
satisfies
00,
00 00.
80,
[I - L0- 7:)
(47)
00,
00,
00,+1
00
according to Eq. (2). In the following, we show that:
T 00, VJ((.))
A411
Σ
(48)
00
where Aet1 is the co-state vector in Algorithm 1. Let X Et
80,,
V((O,), we then have
=1==
00
N/- vJ(Or) and the following difference equation for X;:
X
J(,).)
a0,
T
00. (J().)
-V)(0p) + Σ
00,
(49)
r'=++1
00, (J(;))
=VJ(0+; Σ
I L
T) 00.41
7)]
t'1+1
-v - X L(0:, 7
21

---

Published as a conference paper at ICLR 2025
Given that 1/41 satisfies the same difference equation X++1 in Algorithm I and has the same value at t = T, we have X,+1 = A Therefore, the gradient in Eq. (46) can be written as
DJ - - Vi 0c). 8%
(50)
Combining Eq. (45) and Eq. (50), we recover the update rules of Tn, in Algorithm by setting a' =a/n.
E EXTENSION TO CONSIDERING DATA DEPENDENCE
In Section 2.1, we formulate our problem by focusing the individual importance of each data point to the downstream loss. This potentially affects the performance PDS when the original corpus to select from contain too much duplication. The reason is that, similar examples tend to have similar distances to A, which means that if one example individually get high y, the others will also get high y and be selected, affecting the diversity of the selected corpus. In this section, provide discussion on how the dependence between the selection data points can be considered to improve the diversity of the selected data, improve the pair-wise difference between the examples to promote the diversity of the selected data by considering the following pair-wise similarity:
extension
theoretical framework. Specifically,
h2 hm
sim Tm)
(51)
h2
where fe and hm are the representations of each example generated by model like RoBERTa (Liu et al., 2019). Then, we have the following loss to minimize, which encourages the difference between examples:
diversity
sim(Zn; - S7,
(52)
where S (1n; F'm) (ml) ( We can add this loss the optimization problem Eq. (3), using a hyper-parameter w to control its weight. Additionally, also useful to add an 4o-norm constraint y to constrain the number of non-zero elements the selected data number K. This is because when pairwise interactions between examples are considered, adding or removing one data point would affect the contribution of the others. Therefore, the optimization problem becomes:
min
J(O.) w -Sy
(53)
N Σ a=1
S.L. O111 = 0: - (1-VVL((..))) 7 U,
th > 0] = K.
However, the problem with the no-ormm constraint is computationally intractable to solve because computing the lu-momm involves determining the exact sparsity pattern, which is equivalent to solving a combinatorial optimization problem. For digh--immnsional 7, this becomes computationally prohibitive as the complexity grows exponentially with the size of the problem. Furthermore, is unclear how the constraint would affect the physical property of the dynamics and resulting in the optimization problem. For example, whether it will introduce any singularity during the dynamic process or the inaccuracy from the approximated optimization procedure (given fo-norm is not differentiable) would lead any unexpected outcome. We leave the design of effective and efficient methods to solve the problem in Eq. (53) to future work.
F DISCUSSION ON EQ. (6) OF THEOREM 2.1
The Optimality. In Theorem 2.1, Eq. (6) takes a maximum operation to get y", which in- creases the scores with the highest Σin XT,1 Vi(z,, 0*) values and reduce those with lower Ei= AT, Vi(x,, 8*) values to However, this does not imply that a single training example is the optimal for data selection for the following reasons:
22

---

Published as a conference paper at ICLR 2025
First, there could exist several samples having the same largest Σ A V2 8*) values when the model's parameter 8" is trained under the optimal y scores, i.c. SS = {nl, 12,, m1s)) such that for any m $ S.
T Σ
T A41 (( 0))
T
A Vi(
Σ
8}>
Xht VI(Tm, 0*)
(54)
Therefore, multiple samples in S can have non-zero 7 scores to satisty Eq. (6), as long as those samples not in S get zero scores. Accordingly, although "a single training sample" is solution to Eq. (6), there also exist many more "equally qualified" solutions (such as assigning uniform y values to the data points in S) that satisfy Eq. (6).
Second, Theorem 2.1 is a necessary condition of the y optimality. This is because the two techniques we used to prove Theorem 2.1 in Appendix B, i.c., (1) time-variant PMP for discrete dynamical system (Theorem B.1) and (2) The Lagrange Multiplier method, are all necessary conditions. Therefore, the unreasonable "a single training sample" is not necessarily the optimal solution to the problem.
Many empirical results (Geering, 2007; Seierstad & Sydsaeter, 1986) have shown that solving this necessary condition successfully lead to fairly good solutions, which is much like using gradient descent in deep learning even if it does not guarantee the optimality of the solution.
Intuition Behind 7 Updates. During the update of y in Algorithm 1, since has not reached its optimum yet, the sample set S having the same largest E A Vl(zn, 0) values has not emerged, making the max operation unstable in practice. Therefore, we increase ? by value proportional to Et= > Vi(zn, 0) in each iteration of the outer loop as in Algorithm updating strategy still satisfies our theory because it guarantees that the sample set having the largest EL AT. VI(In, (f) values gets the highest y values. It also ensures that the samples with the same Ei=o X__ VI(zn, 0) values receive the same scores, which avoids the "single training sample" solution.
This
G IMPLEMENTATION DETAILS
G.1 SOLVING DATA QUALITY SCORES
Algorithm 1 needs to compute the Hessian matrix, in Eq. (5), and per-instance vector-gradien product, as in Eq. (6). We adopt the Jacobian- Vector-Product" in PyTorch (Paszke et al., 2019) to efficiently implement these operations. There is still large room for this implementation, such as adopting other deep learning frameworks (Dagréou et al., 2024; Bradbury et al., 2018) or using lower-complcxity algorithms (Wang et al., 2024). Algorithm I requires storing all the LM checkpoints O, from i = 0 to f = T - 1 in the forward inner loop for computing A in the reverse inner loop. To reduce the single-device GPU memory use, we implement an activation partition algorithm inspired by ZeRO-2 (Rajbhandari et al., 2020), where O, in one inner-loop pass are stored in different GPU devices. We also adopt strategy inspired by activation checkpointing (Chen al., 2016).
G.2 TRAINING DATA SCORER
We fine-tune the Fairseq-Densed 100--115MM 125M model (Artetxe et al., 2022) on the solved data weights y* obtain the data scorer. as in Eq. (8), we adopt a linear transformation on the mean pooling of the instance's representations along the sequence length. The size of the hidden state is 768. We optimize Eq. (8) with the AdamW (Loshchilov
Hutter, 2019) optimizer for 5 epochs, using a learning rate 1 10-1 and batch size 512. We split 10% samples from Dprx for validation and select the checkpoint achieving the highest Spearman correlation score on the validation set to infer data quality scores in D.
hettp:://pyttrchhorr//dbo.
func. api html
23

---

Published as a conference paper at ICLR 2025
Model Size
dmodel
deFN
Flaayrrs
Theaa
dhead
learning rate
160M
768
3,072
12
12
64
6 x 10-1
470M
1,024
4,096
24
16
64
3 x 10-4 44
IB
1,536
6,144
24
16
96
2.5 x 10
1.7B
2,048
8,192
24
16
128
2 x 10-1
Table 6: Model configurations and corresponding learning rates.
G.3 PRE-TRAINING
We pre-train all our models for about 50B tokens with a batch size of 512 and max input length of 1,024. We use the AdamW (Loshchilov Hutter, 2019) optimizer and cosine learning rate scheduler, which warmups the learning rates for 2K steps and then decays it to 10% of the maximal value. We follow (Brown al., 2020) to set the maximal learning rates listed in Table 6, together with the model configurations.
G.4 EVALUATION
We evaluate our trained models on MMLU (Hendrycks et al., 2021) and the evaluation sets used in OLMo (Groeneveld al., 2024), including Hellaswag (HS; Zellers et al., 2019), Winograde (Wino.; Levesque et al., 2012), LAMBADA (LAMB; Paperno et al., 2016), OpenbookQA (OBQA: Mi- haylov al., 2018), ARC-easy/challenge (ARC-elc; Clark al., 2018), PIQA (Bisk al., 2020), SeiQ (Welbl et al., 2017), and BoolQ (Clark et al., 2019). We adopt the L.M-valuationn-arnees library (Gao et al., 2024)* to conduct the zero-shot evaluation, which formulates the tasks as the multiple-choice problems and chooses the answer by comparing the answer- mnswer-length-nomee across candidates (a (ac__nor)). We also compute the test loss of our trained models on the DCLM corpus (Li et al., 2024), 10K subset uniformly sampled from the high-quality pre-training corpus Li al. (2024).
loss
G.5 BASELINES
Reducible Holdout Loss Selection (RHO-Loss; Mindermann et al., 2022; Lin et al., 2024). RHO-Loss selects data with high "reducible holdout loss" defined as follows:
RHO-Loss, l(z, 0T) [( @downn,
(55)
where (down is the model parameters trained on the downstream tasks, which is LIMA (Zhou et al., 2024) in our experiments. We first split 10% of LIMA for validation and choose @down with the lowest validation loss. Then, we train 0, using In with the top a x 100% RHO-Loss,, We tried (0.1, 0.3, 0.5, 0.7) and find that 0.5 performs the best.
Data Selection via Importance Resampling (DSIR; Xie et al., 2023). DSIR selects pre- training data based the n-gram feature overlap between the instances in the downstream dataset (LIMA (Zhou et al., 2024) in our experiments) and the large-scale corpus. Pre-training instances whose features have high probabilities under the feature distribution of the downstream dataset will obtain higher sampling weights. We use the official implementation of D
Influence Score (IF-Score Koh & Liang, 2017; Grosse al., 2023). We adopt the influence function (Koh & Liang, 2017; Grosse et al., 2023) to measure the quality of x using the downstream loss J(8) as the target:
IF-Score(z) = Vi(a, 0T)
(J(()).
(56)
where O is the LM parameters trained for steps and L(0r) = + Exep 1x, 0r). We adopt a linear-time iterative algorithm to compute the
(Agarwal et al., 2017).
// harnessss/ httpps: cttp:tthoooooooaaaa/aaaaa
24

---

Published as a conference paper at ICLR 2025
HS
LAMB
Wino.
OBQA
ARC-e
ARC-c
PIQA
SciQ
BoolQ
Avg.
Conventional
32.2
34.9
51.4
25.6
40.9
22.5
58.3
65.5
57.6
43.2
RIOO-Loss
32.2
35.3
53.2
28.1
40.5
24.1
58.5
53.0
43.1
DSIR
32.8
35.7
52.5
26.6
41.2
23.8
57.8
68.7
54.7
43,8
IF-Score
32.2
35.7
51.1
27.4
40.8
22.6
57.6
64.1
55.8
43,0
PDS
33.5
38.2
51.4
28.4
42.3
24.1
59.2
68.8
58.7
45.0
Table 7: Downstream evaluation results on 160M models. We report the accuracy scores and the average scores across the datasets. The best scores of each model size boldfaced. PDS achieves the best performance most cases compared to the baselines.
To reduce the computational cost, we adopt similar efficient implementation PDS by computing the IF-Scores on a small proxy data based on small proxy LM and then transferring these scores to large pre-training corpus with fine-tuned data scorer. We select examples with the top 40% scores inferred by the data scorer on the large pre-training corpus.
G.6 SIMULATED SETTING FOR EXPERIMENTS TABLE 6
To exactly run Algorithm with a feasible computational overhead, we adopt a 12M LM from the Mistral (Jiang et al., 2023) family, with a small vocabulary. We uniformly sample 4,096 instances as D, with a max sequence length of 256, and construct a 16K vocabulary from and LIMA. We run the inner loops for 5K steps and the outer loop for 5 epochs to get the PDS (exact) line in Figure 6. For PDS (Efficient), we adopt an 8.7M model as the proxy LM and set M = 5, Tprx 100. We run the inner loop using SGD with batch size of 128. The outer loop epoch number set to 1.
H COMPLEXITY ANALYSIS
Following Hoffmann al. (2022), for an LM with N parameters trained on D tokens, we assume the computational FLOPs of forward and a backward pass are 2N D and 4ND, respectively. We compute the FLOPs and asymptotic complexities of different stages PDS as follows:
Solving Data Quality Scores: According to Section 2.3.1, we first pre-train a proxy LM on D which consumes 6.NPEx D) FLOPs. Then, we perform Algorithm M times on Drxx based on the proxy LM. The forward inner loop in Algorithm 1 consumes 6.NPXX FLOPs. The reverse inner loop can be treated as the "backward" propagation of the forward inner loop as discussed in Appendix D, which consumes x 6NPrx Deix FLOPs. The update of results in one forward and backward pass of the proxy LM on DPxx which consumes 6 D FLOPs. In summary, the asymptotic complexity of solving data quality scores is O(NPEX D + NeXx
Data Scorer: The data scorer with Nscore is trained Derx and used to infer data quality scores on D. Therefore, the computation overhead is A.yssoreDDrss + +2N N. D and the asymptotic complexity is O(3N Dorr + N'wore D).
Data Selection:
Selecting pre-training corpus requires iterating D, whose asymptotic plexity is OD). This process can be done on CPUs and does not require GPU FLOPs.
Pre- Fre-Training: fre-training an LM with N parameters requires 6N D FLOPs, whose asymptotic complexity O(ND).
I MORE RESULTS
I.1 RESULTS ON 160M MODELS
We present the results on the OLMo (Groeneveld et al., 2024) evaluation sets based on the 160M models in Table 7. PDS achieves the best performance in most compared the baselines.
25

---

Published as a conference paper at ICLR 2025
Average Accuracy (%)
Average Accuracy (8)
Average Accuracy
64-
43-
46
48
45
42-
46-
41
44
Conventional
8
0.1
033
0.5
0.3
0.84
1.18
1.40
0,63
1.28
1.89
2.52
315
Computation
Ix10 FLOPs)
Computation Ix10
Computation (x100 FLOPS)
(a) Model Size 160M
(b) Model Size 470M
(c) Mode Size IB
Figure 8: Scaling curves of average accuracy on the OLMo (Groeneveld et al., 2024) evaluation datasets with respect to computation for 160M, 470M, and IB sizes.
Method
PPL (160M)
PPL (470M)
Conventional
44.7
34.8
RHO-Loss
42.1
33.0
DSIR
40.5
34.0
IF-Score
41.5
31.1
PDS
36.6
27.1
Table 8: The LMs' perplexity (PPL) the ground truth for zero-shot evaluation
MMLU (Hendrycks et al., 2021). The best scores of each model size are boldfaced.
1.2 SCALING CURVES OF COMPUTATION FOR OTHER MODEL SIZES
We plot scaling curves of computation for 160M, 470M, and 1B models in Figure 8. PDS-selected data accelerates the model learning across model sizes.
1.3 RESULTS ON MMLU MEASURED BY PERPLEXITY ON GROUND TRUTH
As suggest by Schaeffer al. (2023), using smooth and continuous metrics like the LM's perplexity on the ground truth better reveals the gap between different base models. Therefore, we include this results in Table supplement to Table 2, which shows substantial improvement of PDS over the baselines.
I.4 TEST Loss EXTRAPOLATION WITH THE SCALING LAW
We extrapolate the test losses on the DCLM corpus (Li et al., 2024) of the conventionally trained and PDS- -saitrined LMs with the Scaling Law (Hoffmann et al., 2022; Kaplan et al., 2020). Following Hoff- mann al. (2022), we consider the scaling law with the following form:
B
D)
(57)
No
D
where N is the model parameters, D is the trained tokens, and A, B, F,, a,, are constants. We obtain these constants by minimizing the Huber loss (Huber, 1992):
min a,
Σ
Huber, (LSE(a a log N,, - 3og D;, e) log L,);
(58)
(N.L).
where LSE(-) is the log-sum-exp operation. The loss is summed over all (N;, Di, Li), which is obtained by the test losses 160M, 470M, IB, and 1.7B LM during training from OB to 50B tokens. We record the losses every 2.5B tokens, resulting in total 4 x 50B/2.5B 80 tuples like (N,, D., L,). After solving a, b, and e from Eq. (58), we have A exp(a), = exp(b). and E exp(e).
26

---

Published as a conference paper at ICLR 2025
A
B
Ct
8
E
Conventional
8.09 x 102
7.50
10
0.397
0.651
2.829
PDS
6.21 x 10a
1.76
10a
0,518
0.585
2.829
Figure 9: Scaling law constants by fitting the test losses on the DCLM corpus.
160M
=470M
N = IB
N =1.7B
50
10
Conventional
0.990 0.992
0,998 0.995
0.993 0.997
0.996 0.993
0,999
PDS
0.998
Table 9: Correlations (R²) of fitting the scaling curves.
Hoffmann et al. (2022) optimizes Eq. (58) using the LBEGS algorithm (Liu & Nocedal, 1989). However, we find this algorithm sensitive the initialization of the parameters be opti- mized. Therefore, we apply a two-stage optimization. Specifically, we first fit the following data scaling curves for N =160M, 470M, IB, and 1.7B with non-linear least squares from scipy. owt.opiimiiii curve_ fit?, which is much more robust to the initialization:
Bo(N) D
L(D) = E'N) +
(59)
where EN(N), (o(N) and Bo(N) are the fitted parameters. Then, we fit the following model size scaling curve:
Ao
E'
Eo +
(60)
Noo
We use the constants from Eq. (60) and the average constants from Eq. (59) to compute the initializa- tion for the LBFGS algorithm:
4g log Ao:
Bo + (4700) + B1B) + Bo(7.B
bo log
4
0g Co
(61)
(1600) + (o(40) + (o((B)) + (1.7B))
Bu =
4
Co log Eu,
where dg, bo, (rg, Bo eo are the parameter initialization for the LFBGS algorithm to optimize Eq. (58). We set 8 = 1 x 10-3 and learning rate to 0.05 when running LFBGS and obtain the constants in Table 9. We use these constants and Eq. (57) to compute the predicted loss in Table 3.
In Section 3.3 (Data-Constrained Setting), compute the expected token demand of conventional pre-training to match the performance of PDS (r = 0.25, 4 Eps.), we solve for D using the constants Table 3 and use as the token demand, indicating the LM can be trained for epochs as suggested by Muennighoff et al. (2023).
Goodness of Fit. We evaluate the goodness of of the scaling curves with respect to the training token size D and model size N respectively, by computing the correlation coefficient It 1
-.-- 2.0 = where is the ground truth value and ÿ. is the prediction.
Regarding the training token size, to make the original problem a linear regression problem, convert Eq. (57) into
log L(N,D)
A Na
= log B log D.
(62)
httpps://dcss scwsyccppo...o. http://wwwwt..ooooo////////////e curve_ html
optimi ze.
27

---

Published as a conference paper at ICLR 2025
Average Accuracy (96)
45.0
Average Accuracy
44.5
44.5
Conventional (Ours)
Conventional Pre-Training
44.0
44.0
43.5
43.5
0.4
0.6
0
1.0
0.0
0.1
0.3
0.5
Data Se ection Ratio r
Figure 10: Effect of the data selection ra- tio a. We report the average accuracy on the OLMo evaluation datasets for [0.3, 0.4, 0.5, 0.0,
Figure 11: Effect of the strength in Gumble- Top-K. We report the average accuracy on the OLMo evaluation datasets for [0.0, 0.1, 0.3, 0.5].
Then, consider log (1: AA) as the ground truth value for regression, where 4 is the observed loss, and log B log D, the prediction. For each N [160M, 470M, 1B, 1.7B) compute an R² respectively. Similarly, regarding the model size, we convert Eq. (57) into
log (LNN
B D3
= log A - a log N,
(63)
to compute the corresponding R² that measures its goodness of fit. For simplicity, we only consider the models at the end of training, where D 50 10 The results in Table I.4 show that the correlations are sufficiently high, suggesting the scaling curve fits the impact from both the data and model sizes very well.
1.5 ABLATION STUDIES
Data Selection Ratio. In Figure 10, investigate how the data selection ratio affects the perfor- mance of PDS when the original training corpus is sufficiently large (in Section 3.3, we explore the data selection ratio in the data-constrained setting.). A lower data selection ratio results in better final model performance. However, to ensure that the selected data contains enough training tokens, larger original corpus needed for lower a. Therefore, keep a 0.4 our main experiments to balance effectiveness and data demand.
Gumbel Noise in Data Selection In Figure 11, we explore the effect the strength in Gumble- Top-K used for data selection in Eq. (9). We can see that = 0.1 achieves the best performance, verifying the benefits of increasing the diversity of the pre-training corpus. Too large a 7 value makes PDS degenerate random selection, causing the performance to decrease.
Choice of (0). The desired data plays an important role in determining the quality of the selected data, We test different downstream datasets to compute ((0) in PDS and report the model performance in Table 10. The comparison between the results of using LAMBADA and LIMA shows the impor- tance of the downstream data diversity. Instances in LAM- BADA mostly come from stories, while LIMA is composed of instructiin-responn pairs that cover various tasks, which yields better overall performance. When comparing LIMA. Open WebText Gokaslan et al. (2019), and CC, we conclude that data quality is another major concern. Although Open- WebText has been shown to have better scaling factors (Bi al., 2024) and used the target (Brown et al., 2020), replacing it with higher quality LIMA further improves perfor-
J(0)
Acc.
Conventional
43.2
LAMB.
43.7
CC
43.0
PDS
OWT
44.1
LIMA
45.0
Table 10: Effect of using differ-
ent downstream datasets compute J()). We report average accuracy on OLMo evaluation datasets.
mance. Compared diversity and quality, large sizes of downstream datasets seem less important, because LIMA performs the best with the least instance number.
28


Reviews:
[{"review_id": "RNLRdBc3cB", "reviewer": "Reviewer_tANF", "summary": "The paper introduces a new data selection methodology inspired by control theory. The authors design a control problem that corresponds to training dynamics of a data-reweighted objective where the control action is the reweighting, the dynamics are gradient descent or Adam dynamics on the parameters, and the objective is minimal training error of the trajectory induced by such dynamics.  This is similar to some works in the learned optimizer literature but the optimizer parameters being optimized correspond to data selection.  \n\nOn a small proxy LM, the algorithm applies a fixed-point iteration (which is shown to be projected gradient descent) in order to induce Pontryagin's Maximum Principle for this control problem and to produce a dataset of quality scores (example weights).  A model is then learned to label examples, which can be used in a data filtering process for a much larger training run.  The authors evaluate this methodology on 160M to 1.7B parameter models with some initial promising results in comparison to baseline data selection approaches.", "questions": "Seems A single step of the inner loop gets most of the gains, why do you think this is?  Given this is much cheaper is this a more feasible algorithm.\n\nDoes the proxy model and the scorer models complexity need to scale as we scale up the final pretrained model we want to use this with?\n\nLine 883, what is B and what is beta? \n\nTypo:\nLine 68: \"(e.g., 125M paramaters)\"", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "As far as i can tell, this approach is novel, in that similar methods have not been applied to the data selection setting. This is a theoretically motivated approach that attempts to optimize the training trajectory directly.  The idea to learn a labeler to translate the ideas from leaned optimizers into a data filtering mechanism is very clever.\n\nThe paper is very well written in my opinion--the general algorithm pipeline is quite complicated, but the main technical/algorithmic ideas were well motivated and clear.\n\nThe algorithm seems to perform well against baselines, although the compute requirements are likely higher (though maybe not prohibitively so as can be seen in Table 4).", "weaknesses": "I am very skeptical of the scaling laws presented in this work.  There does not appear to be enough data to fit such a complicated model, and I do not think these results should be included nor trusted.\n\nWhile the empirics seem to be better than baselines, there seems to be fairly marginal improvements across the board and without very careful hyperparameter sweeps on the downstream optimizer, there is reason to be skeptical of the empirics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JvOmbF31NL", "reviewer": "Reviewer_KX1L", "summary": "This paper proposes a data selection method (called PDS) for language model pretraining based on optimal control principles. Specifically, the PDS method aims to compute a data score for re-weighting data, such that the trained model (via a fixed procedure) achieves the best downstream test error. In practice, this is done using small-scale proxies. Additionally, a model is fitted to predict scores for general data. One important benefit of the proposed method is its offline nature, enabling an effortless plug-in to existing language model training pipelines. The effectiveness of the proposed method is verified across multiple benchmarks.", "questions": "Theorem 2.1 seems related to the KKT conditions of the optimization problem in equation (3). Could you briefly discuss the connection between Theorem 2.1 and the KKT conditions? Of course, this is not a weakness of the paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written, and the experimental setup is clearly explained.\n- The PDS method is well-motivated by control theory and demonstrated to be effective in practice.\n- The offline nature of the proposed method makes it very easy to use in existing language model training pipeline.  \n- Multiple thoughtful ideas are proposed to accelerate the implementation of PDS.", "weaknesses": "I didn’t observe significant weaknesses in the work. If I had to mention one, it would be the lack of large-scale verification of the method's effectiveness. This limits the strength of some claims (e.g., fitting scaling laws and extrapolating results). However, this is understandable given the substantial computational resources required for such large-scale experiments.\n\nOverall, I think the paper demonstrates the promise of the proposed method, and I am inclined to support its acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CAmTwonSBg", "reviewer": "Reviewer_snbg", "summary": "The paper suggests a method for data selection based on applying optimal control theory to compute the importance of individual data points. Crucially, this importance is computed in the context of end-to-end training, to optimize the AUC of downstream performance - this is in contrast to previous works that don’t look at training trajectory. Several optimizations are used to approximate the PMP solution. Moreover, the data selection is done offline so there is just a constant cost.", "questions": "- Why did you decide at line 259 to use the average of the output hidden states rather than the last one? Intuitively this gives more weight to how the datapoint starts (since the first activations at the first few positions affect both the hidden states at these positions as well as later on).\n- Is it computationally feasible to constrain the $\\ell_0$-norm of $\\gamma$? If so, do you expect to get the same data points as you would if you took the biggest ones by the current $\\gamma$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Extensive ablations (optimizing Algorithm 1, choice of downstream task) are performed to support the empirical claims and many of the choices made.\n- The tackled problem is very relevant and the approach is innovative, using optimal control to account for the complete training trajectory.\n- The empirical results are very promising on the small-to-medium scale experiments that they are run on. A scaling law is also computed to try to support generalization to larger models.", "weaknesses": "- It would be good if Appendix G3, choice of $J(\\theta)$ was covered in the main body (regarding the influence of the downstream data on performance) since that is an important driver of the performance.\n- Better explain Figure 1: in (b), do you have the same amount of data being used by your method and Redpajama’s to train? Do all points correspond to the same model size in (a)? If not, how do you vary these?\n- It would be useful to mention at line 89 that J is not computed on the same downstream tasks that you measure one (as you mentioned at line 307) since this distinction is crucial and not made clear until line 307.\n- It seems like line 100, defining $\\gamma$’s makes an implicit assumption about points having a notion of independent importance, whereas one could envision cases where a point’s importance is contingent on other points being in the dataset. It would be useful to add a discussion on this - as stated, it looks like this is a fact, not a choice.\n- It would be useful to add a discussion on why it makes sense to need Gumbel sampling - after all, one could expect $\\gamma$’s to already have accounted for the need of diversity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "shYHXaxado", "reviewer": "Reviewer_3Rqd", "summary": "This work proposes a novel method of data selection for training machine learning models, or LLMs in particular.\nIt starts by theoretically formulating data selection as an optimal control problem, \nwhere the weights of training samples are optimized in order to minimize the total area under the downstream loss curve during training with gradient descent on the weighted data.\nBased on this, a practical and efficient implementation is developed: \nfirst solve the optimal control problem on a small proxy dataset (approximately, via bi-level optimization),\nand then train a data scorer (by fitting the obtained scores from the first step),\nwhich can then be applied for training any LLM with other datasets.\nExperiments show improvements over prior data selection algorithms,\nin terms of test losses and scores on some commonly used benchmarks.", "questions": "- Two questions about the theory:\n(a) Could you provide some intuition about why theoretically a single training sample is optimal, according to Eq. (6)? \nWhat's missing in this theoretical analysis?\n(b) In Eq. (22) where the method of Lagrange multiplier is applied, it seems that it misses a $\\max_{\\mu}$ term. Is this correct? If so, does the remaining analysis still hold, or need corrections as well?\nI'm willing to raise my rating if these two questions can be resolved.\n\n\n- Line 883, \"B\" and \"\\beta\" seems like typos; should be \"C\" and \"c\" respectively?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The formulation of data selection as optimal control, and the theoretical connection between two seemingly unrelated fields, are interesting.\n\n- The presentation is good overall. Ideas are clearly explained.", "weaknesses": "- According to the theory of this work, in particular Eq. (6), the optimal value of $\\gamma$ (the weights for training samples) is an one-hot vector, \nwhich leads to the surprising conclusion that training an LLM with a single sample is theoretically optimal.\nThe authors also mention this in Line 235. \nMy major concern is that this seems to reveal a fundamental caveat of the proposed methodology.\nAlso, with many approximations in deriving the practical implementation (Section 2.3), \nit becomes less and less clear how much the empirical advantages of the resulting algorithm are actually relevant to the theory in Sections 2.1 and 2.2.\n\n\n- Using the \"scaling law\" for extrapolating empirical results with small-scale models to 175B or 405B models (Table 3) doesn't feel like science. \nIt is not obvious why scaling laws for conventionally trained LLMs should be applicable to PDS-trained LLMs.\nIn addition, the goodness of fit for calculating the scaling law's coefficients is not reported, which further makes the extrapolation less convincing.\n\n\n- Regarding the empirical results, the proposed method is compared with prior data selection algorithms only in Tables 1 and 2, for 160M and 470M models that are relatively small-scale.\nMoreover, many scores in these results are close to random choice (e.g., 25% accuracy for MMLU in Table 2), \nhence it is not clear what the minor differences between the scores of different methods really imply.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "etcdljQ6MX", "reviewer": "Reviewer_NWbj", "summary": "This paper presents a formulation of turning the data selection to the AUC optimization problem. More importantly, this paper derives the optimal solution for the data selection problem and proposes an efficient implementation of the proposed method. The proposed methods are empirically verified in diverse settings.", "questions": "It is really interesting to see the formulation of optimization problem (3) in the data selection problem. Has this method been used in other tasks such as computer vision, diffusion model, or simple MNIST dataset classification? It seems to be a very general approach that can be used in many senarios. \n\nIn Line 60, the PMP derives the necessary condition for optimal data selection. Does it mean that it is not sufficient to guarantee the data selection is optimal? \n\nIt seems that the PDS method always achieves better performance. Is there any negative results where the PDS method fails? Do we have any senarios where we should avoid using this method for the data selection?", "limitations": "", "rating": 10, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This work is sufficiently solid. It covers the fundamental problem formulation, which introduces a new optimization problem. Also, it provides the theoretical optimization solution to the proposed problem, which has not been studied before. From this perspective, this paper has made sufficiently novel contributions.  Moreover, this work designs an efficient implementation to avoid additional computation overhead. These methods are empirically verified with different experiments setting, demonstrating the effectiveness of this method across different settings. Due to its theoretical contributions and solid experiments, I give a clear accept (score 10) to this paper.", "weaknesses": "I didn't identify any major weaknesses of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)