Data Selection via Optimal Control for Language Models

## Abstract

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 mitigates the quick exhaustion of available web-crawled corpora. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/data_selection .

## 1 Introduction

With the thriving of language models (LMs; , ), 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 [].
This focus enables researchers to explore the limit of LMs in the face of increasing training data demands [, , ]. It also helps to reduce the computational costs during pre-training [], and addresses the potential limitations caused by available Internet data [, ].
Without doubt, pre-training data selection is highly valuable for both research and industry sectors.

Unlike previous works relying primarily on manually crafted heuristics [, ], we connect data selection with classical Optimal Control theory [], where control variables in a dynamic system are optimized to achieve desired objectives.
This mathematical formulation allows fine-grained white-box analysis of how the control variables drive a dynamic system from one state to 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; ) to derive the necessary conditions for optimal data selection in theory.
These results offer a rigorous, theory-driven alternative to the ad-hoc 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 a proxy dataset (e.g., 0.2B tokens), assigning a quality score to each instance based on its impact on downstream tasks.
After that, a data scorer (e.g., 125M) 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) from scratch.

Unlike previous pre-training data selection methods based on deduplication [, ], pattern matching [], or single checkpoint performance [, ], which are agnostic to 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 in an online fashion [, ], PDS operates offline, before the training begins, which avoids additional training-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 the impact of individual data points on LM pre-training.

In our experiments, we select data from the CommonCrawl(^1^11[https://commoncrawl.org/](https://commoncrawl.org/)) with PDS and pre-train LMs with 160M, 470M, 1B, and 1.7B parameters from scratch. 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-art baselines. Extrapolating these results using the Scaling Law [, ], we show that the benefits remain consistent for ~400B LMs trained on ~15T 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 []. 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 ${\bm{\theta}}\in{\mathbb{R}}^{N}$, pre-trained from scratch on a dataset $\mathcal{D}=\{x_{n}\}_{n=1}^{|\mathcal{D}|}$, over a total $T$ steps.
Data selection [, ] aims at identifying a subset $\mathcal{D}^{\prime}$ from $\mathcal{D}$, such that LMs trained on $\mathcal{D}^{\prime}$ achieve better downstream performance, measured by a lower downstream loss $J({\bm{\theta}})$.

The pre-training process renders $J({\bm{\theta}})$ as a function of $\mathcal{D}^{\prime}$, which can be fully characterized by a data quality score vector ${\bm{\gamma}}=\left[\gamma_{1},\gamma_{2},\cdots,\gamma_{|\mathcal{D}|}\right]
^{\top}$ in a $|\mathcal{D}|$-dimensional simplex $U$, where $U=\left\{\left.[\gamma_{1},\gamma_{2},\cdots,\gamma_{|\mathcal{D}|}]^{\top}
\right|\sum_{n=1}^{|\mathcal{D}|}\gamma_{n}=1\text{ and }\gamma_{n}\geq 0\text
{ for }1\leq n\leq{|\mathcal{D}|}\right\}$
(^2^22Only the relative data quality is meaningful for data selection. Therefore, we ensure the sum of the quality scores to 1 to avoid the impact of their individual scales.).
A higher quality score in ${\bm{\gamma}}$ indicates the corresponding instance is more helpful to reduce $J({\bm{\theta}})$, and thus the LM should learn more from the instance.
This results in the following general pre-training loss, defined as the weighted
sum of the per-instance loss by ${\bm{\gamma}}$:

$$ $L({\bm{\theta}},{\bm{\gamma}})=\sum^{|\mathcal{D}|}_{n=1}\gamma_{n}l(x_{n},{ \bm{\theta}}),$ (1) $$

where $l(x_{n},{\bm{\theta}})=-\log p_{{\bm{\theta}}}(x_{n})$.
The goal of data selection is thus to find ${\bm{\gamma}}$ that reduces the downstream loss $J({\bm{\theta}})$, and then select instances with the highest scores in ${\bm{\gamma}}$.
For simplicity, we assume that the LM is trained using Gradient Decent (GD) for $0\leq t<T$, with the derivation under the Adam optimizer [] provided in Appendix C:

$$ $\displaystyle{\bm{\theta}}_{t+1}$ $\displaystyle={\bm{\theta}}_{t}-\eta\nabla L({\bm{\theta}}_{t},{\bm{\gamma}}),$ (2) $$

where ${\bm{\theta}}_{t}$ represents the model parameters at the time step $t$ during GD and $\eta$ is the learning rate.

##### Optimization Problem.

Motivated by the literature on learned optimizers [], we optimize ${\bm{\gamma}}$ by minimizing the area under the curve (AUC; ) of $J({\bm{\theta}}_{t})$, which is approximated by the cumulative sum of $J({\bm{\theta}}_{t})$ over the pre-training process:

$$ $\displaystyle\min_{{\bm{\gamma}}}$ $\displaystyle\sum_{t=1}^{T}J({\bm{\theta}}_{t}),$ (3) s.t. $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta\nabla L({\bm{\theta}}_ {t},{\bm{\gamma}}),\ \ {\bm{\gamma}}\in U.$ $$

Intuitively, a lower AUC corresponds to faster convergence of the loss and improved final downstream performance. Unlike evaluating $J({\bm{\theta}}_{t})$ 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 [], leading to 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 time optimal control problem [], where $J(\cdot)$ is the cost function, the model parameters ${\bm{\theta}}$ are the state variables evolving according to Eq. (2), and the data quality scores ${\bm{\gamma}}$ are the control variables to be optimized within $U$. This perspective provides a convenient framework to formulate the pre-training data selection problem from a more theoretical perspective.

##### Theoretically Optimal Solution for Data Selection.

Optimal control problems can be solved by a powerful tool known as Pontryagin’s Maximum Principle (PMP; ), 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 ${\bm{\gamma}}$ 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 ${\bm{\gamma}}^{*}$ solve the problem in Eq. (3), and ${\bm{\theta}}^{*}_{t}$ denote the LM parameters trained with ${\bm{\gamma}}^{*}$. For $0\leq t<T$, there exists a vector ${\bm{\lambda}}^{*}_{t}\in{\mathbb{R}}^{N}$ such that

$$ $\displaystyle{\bm{\theta}}_{t+1}^{*}$ $\displaystyle={\bm{\theta}}^{*}_{t}-\eta\nabla L({\bm{\theta}}^{*}_{t},{\bm{ \gamma}}^{*}),\ \ {\bm{\theta}}_{0}^{*}={\bm{\theta}}_{0},$ (4) $\displaystyle{\bm{\lambda}}^{*}_{t}$ $\displaystyle={\bm{\lambda}}^{*}_{t+1}+\nabla J({\bm{\theta}}_{t}^{*})-\eta \nabla^{2}L({\bm{\theta}}^{*}_{t},{\bm{\gamma}}^{*}){\bm{\lambda}}^{*}_{t+1}, \ \ {\bm{\lambda}}^{*}_{T}=\nabla J({\bm{\theta}}^{*}_{T}),$ (5) $\displaystyle{\bm{\gamma}}^{*}$ $\displaystyle=\arg\max_{{\bm{\gamma}}}\sum^{{|\mathcal{D}|}}_{n=1}\gamma_{n} \left[\sum_{t=0}^{T-1}{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{ \theta}}^{*}_{t})\right],\ \ {\bm{\gamma}}\in U,$ (6) $$

where $\nabla^{2}L({\bm{\theta}}^{*}_{t},{\bm{\gamma}}^{*})$ denotes the Hessian matrix of $L({\bm{\theta}},{\bm{\gamma}}^{*})$ with respect to ${\bm{\theta}}$ evaluated at ${\bm{\theta}}={\bm{\theta}}^{*}_{t}$.

Figure: Figure 2: An illustration of Theorem 2.1. Left: ${\bm{\lambda}}^{*}_{t+1}\in\mathbb{R}^{N}$ defines a “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 is to the target direction, calculated as the dot-product between ${\bm{\lambda}}^{*}_{t+1}$ and $\nabla l_{i,t}=\nabla l(x_{i},{\bm{\theta}}^{*}_{t})$ for $i=n,m,k$, as in Eq. (6).
Refer to caption: x3.png

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. 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 ${\bm{\lambda}}^{*}_{t}$, a “target vector” suggesting the ideal gradient direction formed only by high-quality data points. In particular, ${\bm{\lambda}}^{*}_{t}$ aggregates information about the downstream loss $\nabla J({\bm{\theta}}_{t})$ with respect to current training step and the LM’s training dynamics $\nabla^{2}L({\bm{\theta}}^{*}_{t},{\bm{\gamma}}^{*})$, from $T$ to $t$. As a result, ${\bm{\lambda}}^{*}_{t}$ summarizes the dynamics of LM pre-training (i.e., from future to the current).
Since ${\bm{\gamma}}\in U$, Eq. (6) essentially suggests that $x_{n}$ with a higher $\sum_{t}{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{\theta}}^{*}_{t})$ value should obtain a larger score in ${\bm{\gamma}}^{*}$, as shown in Figure 2 (Right). This indicates that the instances whose gradients align closely with the target vector ${\bm{\lambda}}^{*}_{t}$ should be selected.
Note that the PMP conditions for data selection form a complete equation system, where ${\bm{\theta}}^{*}_{t}$, ${\bm{\lambda}}^{*}_{t}$, and ${\bm{\gamma}}^{*}$ are the solutions.
In principle, by solving Eqs. (4)-(6) simultaneously, we can derive the optimal data quality scores, forming the foundation of our data selection framework PDS.

###### Theorem 2.1 (PMP Conditions for Data Selection) .

### 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 $\mathcal{D}^{\mathrm{prx}}$ from the pre-training corpus $\mathcal{D}$.
Algorithm 1, derived from the PMP conditions, is then applied to $\mathcal{D}^{\mathrm{prx}}$ 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 $\mathcal{D}^{\mathrm{prx}}$. The learnt data scorer is subsequently applied to infer quality scores on the entire pre-training corpus $\mathcal{D}$ (Section 2.3.2). Finally, the instances with large scores are selected to form a high-quality corpus $\mathcal{D}^{\prime}$, which is used to pre-train LMs with any size (Section 2.3.3).

#### 2.3.1 Data Quality Scores from PMP

Algorithm 1 solves the PMP conditions for data selection iteratively and returns the data quality scores ${\bm{\gamma}}^{*}$. We elaborate on the details of the algorithm and provide an efficient implementation.

Overview: Algorithm 1 solves a bi-level optimization problem, consisting of an outer loop to compute ${\bm{\gamma}}^{*}$ and two inner loops to compute ${\bm{\theta}}^{*}_{t}$ and ${\bm{\lambda}}^{*}_{t}$ based on the current ${\bm{\gamma}}^{*}$. ${\bm{\gamma}}^{*}$ is first uniformly initialized and then updated for $T_{o}$ epochs, based on ${\bm{\theta}}^{*}_{t}$ and ${\bm{\lambda}}^{*}_{t}$ obtained in the outer iterations.

Inner loops: In each iteration of the outer loop, we run the forward inner loop to compute ${\bm{\theta}}^{*}_{t}$ according to Eq. (4) from $t=0$ to $t=T-1$, equivalent to training the LM with GD using the current data quality scores to re-weight the per-instance losses.
After that, ${\bm{\lambda}}^{*}_{t}$ is computed from $t=T-1$ to $t=0$ with Eq. (5) in the reverse inner loop, based on ${\bm{\theta}}^{*}_{t}$ obtained from the forward inner loop.

Update of ${\bm{\gamma}}^{*}$: ${\bm{\gamma}}^{*}$ is supposed to be updated according to Eq. (6) with ${\bm{\theta}}^{*}_{t}$ and ${\bm{\lambda}}^{*}_{t}$ computed in the inner loops.
Eq. (6) indicates that the new ${\bm{\gamma}}^{*}$ should be set as a one-hot vector, where the element corresponding to the highest $\sum_{t=0}^{T-1}{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{\theta}}^
{*}_{t})$ value is set to 1 and the others are set to 0.
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 a “soft” update, which increases ${\bm{\gamma}}^{*}$ by a value proportional to $\sum_{t=0}^{T-1}{\bm{\lambda}}_{t+1}^{\top}\nabla l(x_{n},{\bm{\theta}}_{t})$ and projects the updated ${\bm{\gamma}}^{*}$ back into $U$.

Figure: Figure 3: The PDS framework. We compute data quality scores ${\bm{\gamma}}^{*}$ on a proxy dataset $\mathcal{D}^{\mathrm{prx}}$ using Algorithm 1, which is derived from the Pontryagin’s Maximum Principle [] (Section 2.3.1). After that, the data scorer learns to predict quality scores from instances, which then infers scores for a large corpus $\mathcal{D}$ (Section 2.3.2). Finally, a high-quality corpus $\mathcal{D}^{\prime}$ is selected based on the inferred scores to pre-train an LM (Section 2.3.3).
Refer to caption: x4.png

Figure: Algorithm 1 PMP-Solver

##### Efficient Implementation.

Running Algorithm 1 on $\mathcal{D}^{\mathrm{prx}}$ based on the learning of a large LM remains computationally intensive, as the inner loops involve training the LM for all $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 a small batch size in the inner loops, which is based on a small proxy LM with $N^{\mathrm{prx}}$ parameters ($N^{\mathrm{prx}}\ll N$) to be trained for $T^{\mathrm{prx}}$ steps ($T^{\mathrm{prx}}\ll T$).
To recover any lost long-range training dynamics, we run Algorithm 1 multiple times by setting ${\bm{\theta}}_{0}$ to the checkpoints at different large-scale
pre-training stages of the proxy LM and then average the obtained data quality scores on $\mathcal{D}^{\mathrm{prx}}$. Specifically, we first train the proxy LM for $T$ steps and save $M$ checkpoints $\left[{\bm{\theta}}_{0}^{(1)},{\bm{\theta}}_{0}^{(2)},\cdots,{\bm{\theta}}_{0}
^{(M)}\right]$ in every $\lfloor\frac{T}{M}\rfloor$ steps. Then, the quality scores are given by

$$ ${\bm{\gamma}}^{*}=\frac{1}{M}\sum_{m=1}^{M}\operatorname{PMP-Solver}\left( \mathcal{D}=\mathcal{D}^{\mathrm{prx}},T=T^{\mathrm{prx}},{\bm{\theta}}_{0}={ \bm{\theta}}_{0}^{(m)},T_{o}=1\right),$ (7) $$

where $\operatorname{PMP-Solver}$ refers to Algorithm 1.
We also incorporate several practical optimization techniques to further reduce the computational overhead, as described in Appendix E.1.

#### 2.3.2 Data Scorer

We fine-tune a small LM as the data scorer to fit the data quality scores ${\bm{\gamma}}^{*}$ on $\mathcal{D}^{\mathrm{prx}}$. Specifically, each instance in $\mathcal{D}^{\mathrm{prx}}$ 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 a scalar. The linear head and the LM are trained together to fit ${\bm{\gamma}}^{*}$ on $\mathcal{D}^{\mathrm{prx}}$ with the Mean Square Error loss:

$$ $\bm{\phi}^{*},{\bm{w}}^{*},b^{*}=\arg\min_{\bm{\phi},{\bm{w}},b}\sum_{n=1}^{| \mathcal{D}^{\mathrm{prx}}|}({\bm{w}}^{\top}\overline{{\bm{h}}}(x_{n}^{\mathrm {prx}},\bm{\phi})+b-\gamma^{*}_{n})^{2},$ (8) $$

where $x_{n}^{\mathrm{prx}}$ is the $n^{\text{th}}$ instance in $\mathcal{D}^{\mathrm{prx}}$, $\bm{\phi}$ is the parameters of the data scorer, and $\overline{\bm{h}}(\cdot,\cdot)\in\mathbb{R}^{d}$ is the average output hidden states of an LM along the sequence length, with $d$ representing the hidden state size. ${\bm{w}}\in\mathbb{R}^{d},b\in\mathbb{R}$ are the parameters of the linear head. After fine-tuning, we infer the data quality scores of the instances in $\mathcal{D}$ with the data scorer, where the quality score for $x_{n}$ is given by
$\gamma(x_{n})={{\bm{w}}^{*}}^{\top}\overline{{\bm{h}}}(x_{n},\bm{\phi^{*}})+b^
{*}$.

#### 2.3.3 Data Selection

We use the output scores from the data scorer to estimate the value of the instances in $\mathcal{D}$ to select the final pre-training corpus $\mathcal{D}^{\prime}$ for the large LM. Given the importance of data diversity in pre-training LMs, we adopt Gumbel-Top-$K$ to introduce randomness into the selection process:

$$ $\mathcal{D}^{\prime}={\text{Top-}K\left\{\gamma(x_{n})-\tau\log(-\log(u_{n})) \mid x_{n}\in\mathcal{D},1\leq n\leq|\mathcal{D}|\right\}},$ (9) $$

where $u_{1},u_{2},\cdots,u_{|\mathcal{D}|}$ are independently sampled from $\operatorname{Uniform}(0,1)$ and $\tau$ is a hyper-parameter to control the strength of the Gumbel noise. The size of the selected data is managed by a data selection ratio $r$, with $K=r|\mathcal{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 [], deduplication [, ], or single-step checkpoints [], PDS incorporates long-range training dynamics into data selection, as reflected by the “target vector” ${\bm{\lambda}}^{*}_{t}$ in Eq. (5). This can be critical for selecting high-quality instances, given the highly dynamic nature of LM pre-training. Although we run Algorithm 1 in a proxy environment and transfer the quality scores to the large-scale setting via the data scorer, many previous works [, ] have shown that data quality information is learnable and transferable across model sizes. Different LMs also share many common training dynamics [].

##### Efficiency and Flexibility of PDS .

Unlike recent online data selection approaches to incorporate LM training dynamics [, ], PDS selects the pre-training corpus offline. This allows PDS to be run only once and used for pre-training multiple LMs of any sizes, without incurring additional computational overhead. The FLOPs needed by PDS in a proxy environment are also negligible compared to the demands of large-scale pre-training, as shown in a complexity analysis in Section 3.3. Besides, the offline nature of PDS makes it flexible to be integrated into optimized pre-training pipelines [] by simply replacing the data sources.

## 3 Experiments

### 3.1 Experimental Setup

##### Data.

We use the CommonCrawl split from Redpajama [] as $\mathcal{D}$ to exclude the influence of domain weights []. During pre-processing, multiple documents may be merged into a single pre-training instance with 1,024 tokens. For the downstream loss $J({\bm{\theta}})$, we compute the LM’s loss on the training split of LIMA [], a high-quality dataset consisting of 1,030 diverse instruction-response pairs that cover a wide range of downstream scenarios. Our evaluation is conducted on various downstream datasets other than LIMA to avoid over-fitting.

##### Model.

We adopt the same model architecture as Mistral [] and pre-train LMs with 160M, 470M, 1B, and 1.7B parameters. Model configurations are detailed in Table 6.

##### PDS .

To compute the data quality scores from PMP, we adopt a 160M proxy LM. $\mathcal{D}^{\mathrm{prx}}$ consists of 160K instances uniformly sampled from $\mathcal{D}$. We first pre-train the proxy LM on $\mathcal{D}$ for 50K steps and then select checkpoints at $\left[10K,20K,30K,40K,50K\right]$ steps. Initialized from these checkpoints, the proxy LM undergoes inner loops with $\eta=0.008$ over $T^{\mathrm{prx}}=100$ steps with a mini-batch size of 256. ${\bm{\gamma}}^{*}$ is updated for one outer loop epoch with $\alpha=1$. For the data scorer, we fine-tune a 125M Fairseq-Dense model [] along with the linear head, using the objective in Eq. (8). The training details for the data scorer can be found in Appendix E.2.
For Data Selection, we set $\delta=0.1$, $r=0.4$, with further hyper-parameter exploration provided in Appendix G.3.

##### Pre-Training.

We pre-train all LMs for 100k steps, using a batch size of 512 and a max input length of 1,024, resulting in roughly 50B tokens. In Section 3.2, we select a 50B-token dataset from a CC corpus containing 125B tokens to assess how different data selection methods improve LM learning given a sufficiently large $\mathcal{D}$. In Section 3.3 (Data-Constrained Setting), we also analyze the effectiveness of PDS when $\mathcal{D}$ is limited in size. See Appendix E.3 for more pre-training details.

##### Evaluation.

We evaluate the LMs’ 0-shot accuracy on the downstream test datasets used in OLMo [] and their 0/5-shot accuracy on MMLU []. We also report the LM’s language modeling loss on a subset of DCLM [], a high-quality corpus curated with complex pipelines and human heuristics, to verify that models trained on $\mathcal{D}^{\prime}$ retain diversity and long-tail knowledge coverage. Further evaluation details are in Appendix E.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 $\mathcal{D}$, also referring to “Redpajama” in Figure 1.
- •
RHO-Loss []: selecting data with high reducible losses, as in Eq. (51).
- •
DSIR []: selecting data with high n-gram overlap with instances in LIMA.
- •
IF-Score []: selecting data with high influence scores, as in Eq. (52).

**Table 1: Results on the downstream evaluation datasets in OLMo []. 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 most cases compared to the baselines.**
|  | HS | LAMB | Wino. | OBQA | ARC-e | ARC-c | PIQA | SciQ | BoolQ | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model Size = 160M |  |  |  |  |  |  |  |  |  |  |
| Conventional | 32.2 | 34.9 | 51.4 | 25.6 | 40.9 | 22.5 | 58.3 | 65.5 | 57.6 | 43.2 |
| RHO-Loss | 32.2 | 35.3 | 53.2 | 28.1 | 40.5 | 24.1 | 58.5 | 63.1 | 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 |
| 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-Loss | 36.6 | 42.4 | 53.0 | 29.4 | 43.7 | 25.2 | 60.4 | 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-Score | 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 |

### 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 1 and Table 2, respectively. As shown, PDS outperforms the baselines on most datasets, achieving the best overall performance across models with 160M and 470M parameters. Figure 1(b) shows the scaling curves of average accuracy on the OLMo evaluation sets with respect to the model sizes, which range from 160M to 1.7B, demonstrating that the performance improvement remains consistent as the model scales up.
These results indicate that the data quality scores obtained in a proxy environment generalize well to various model sizes and downstream tasks.

**Table 2: The 0/5-shot accuracy scores on MMLU. The best scores of each model size are boldfaced.**
| $N$ | Method | 0-shot | 5-shot |
| --- | --- | --- | --- |
| 160M | Conventional | 27.2 | 26.7 |
| RHO-Loss | 27.3 | 27.1 |  |
| DSIR | 27.4 | 27.3 |  |
| IF-Score | 27.4 | 27.1 |  |
| PDS | 27.4 | 27.7 |  |
| 470M | Conventional | 27.6 | 28.5 |
| RHO-Loss | 28.4 | 28.5 |  |
| DSIR | 28.0 | 28.2 |  |
| IF-Score | 28.4 | 28.8 |  |
| PDS | 28.9 | 29.0 |  |

##### PDS Enhances Language Modeling.

Besides downstream tasks, we show that PDS also enhances language modeling on a high-quality pre-training corpus.
Figure 4 shows the test losses on the DCLM subset for conventionally pre-trained LMs and PDS-trained LMs with 160M, 470M, 1B, and 1.7B parameters.
We can see that PDS-trained LMs constantly achieve better performance across various model sizes.
In Table 3, we extrapolate the test losses using the Scaling Law [], showing that the improvements with PDS persist in pre-training recent large LMs, such as GPT-3 [] and Llama family [, , ].
Details of the extrapolation are provided in Appendix 7.
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 1 offers a principled alternative to the complex pipeline for curating pre-training corpus.

##### PDS Accelerates LM Learning.

In Figure 1(a), we plot the average accuracy scores on the OLMo evaluation datasets with respect to the training FLOPs 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).

Figure: Figure 4: Test losses on the DCLM corpus [] for 160M, 470M, 1B and 1.7B LMs.
Refer to caption: x5.png

**Table 3: Test loss extrapolation using the Scaling Law []. 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 2 70B, and Llama 3.1 405B. The improvements of PDS remain consistent for these LMs.**
|  | $N$ | $D$ | Conventional | PDS |
| --- | --- | --- | --- | --- |
| GPT-3 [] | 175B | 300B | 2.882 | 2.872 |
| Llama [] | 6.7B | 1.0T | 2.942 | 2.896 |
| Llama 2 [] | 70B | 2.0T | 2.877 | 2.855 |
| Llama 3.1 [] | 405B | 15T | 2.851 | 2.838 |

### 3.3 Analysis

Figure: Figure 5: Test losses on DCLM corpus [] 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.
Refer to caption: x6.png

##### Data-Constrained 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 in conventional pre-training.
In practice, when the data before selection is limited, the LM has to be trained on the PDS-selected data for multiple epochs.
In Figure 6, we restrict $\mathcal{D}$ to contain 50B tokens and apply PDS with selection ratios $r\in\left[0.125,0.25,0.5\right]$, corresponding to training the LM on the selected data for $\left[8,4,2\right]$ epochs, respectively.
We can see that selecting 1/4 data with PDS and training for 4 epochs achieves the lowest test losses, consistent with the findings of Muennighoff et al. [].
This suggests that PDS improves data utilization as the high-quality web-crawled corpora run out [].
We extrapolate the loss curve of conventional pre-training with the Scaling Law suggested by [], and find that it 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.

##### Efficient Implementation for Solving Data Quality Scores.

In Section 2.3.1, we introduce efficient implementations to solve the data quality scores by running Algorithm 1 for a single outer epoch, leveraging a small proxy LM trained for a limited number of steps.
In Figure 6, we use a simulated setting where ${\bm{\gamma}}^{*}$ can be obtained exactly from Algorithm 1 within an affordable computational budget and compare its performance (measured by $J({\bm{\theta}}_{t})$) and overhead (measured in FLOPs) with our efficient implementation.
Details of this simulated setting can be found in Appendix E.6.
The results show that the efficient implementation significantly reduces computational overhead while preserving most of the performance of the exact solution.

##### 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 to 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-training pipelines [], requiring only a replacement of the data source without altering the pre-training process.

**Table 4: Asymptotic complexity, GPU FLOPs, and actually spent time of different PDS steps and 1.7B model pre-training. $N^{\mathrm{prx}}$ and $D^{\mathrm{prx}}$ are the sizes of the proxy LM and proxy data. $N^{\mathrm{score}}$ is the size of the data scorer. We elaborate on the details of how the complexity and FLOPs are computed in Appendix F. PDS costs little computational overhead compared to pre-training large LMs.**
|  |  | Complexity | FLOPs ($\times 10^{20}$) | Actual Time |
| --- | --- | --- | --- | --- |
| PDS | Proxy $\gamma$-solver | $O(N^{\mathrm{prx}}D+4MN^{\mathrm{prx}}D^{\mathrm{prx}})$ | 0.49 | 15.2 Hours |
| Data Scorer | $O(3N^{\mathrm{score}}D^{\mathrm{prx}}+N^{\mathrm{score}}D)$ | 0.063 | 1.50 Hours |  |
| Data Selection | $O(D)$ | 0.0 | 10.2 Minutes |  |
| Pre-Training | $O(ND)$ | 5.1 | 144 Hours |  |

### 3.4 Ablation Studies on PMP-Solver

Figure: Table 5: Data selection based on different kinds of learning information. We report the LM’s zero-shot accuracy on the OLMo evaluation tasks (Acc.) and the Spearman Correlation of the data scorer (Corr.).
Refer to caption: x8.png

##### Training Dynamics Information.

Incorporating the LMs’ training dynamics into data selection is a key distinction between PDS and other offline data selection approaches.
While IF-Score 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 more valuable.
Table 7 shows the results when different types of learning information are considered. PDS (50K) refers to using the LM checkpoint at 50K as ${\bm{\theta}}_{0}^{(m)}$ in Eq. (7).
PDS ($T^{\mathrm{prx}}=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 ${\bm{\theta}}_{0}^{(m)}$ to checkpoints at later training stages.
Comparing our choice with IF-Score, PDS (50K), and PDS ($T^{\mathrm{prx}}=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 to fit. Our choice also outperforms PDS (50K-100K), indicating that the early-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-trained LM. As the size of the proxy LM increases, the LM’s downstream performance increases, but the data scorer’s performance 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 it harder for the data scorer to fit, but this can be mitigated by increasing the size of $\mathcal{D}^{\mathrm{prx}}$.

## 4 Related Work

##### Data Selection for Language Models.

Many works have explored data selection approaches to accelerate LM learning or improve downstream performance [].
Some curate data prior to LM training, which we refer to offline methods, including domain-mixing [, , ], data pruning [, , ], sample-wise data selection [, , , ], and data programming [, , ].
Other works dynamically select data during LM training by adjusting domain-mixing weights [, , ] or more fine-grained reweighting strategies [, , , ].
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 [] have been shown to be effective in deep learning [, ], 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 Maximum Principle [], some works design optimization algorithms with better convergence rates [, ] or broader application scenarios [],
and others provide theoretical foundations of neural networks for better interpretation [, ].
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 a set of necessary conditions for optimal data selection using Pontryagin’s Maximum Principle (PMP), which establishes theoretical principles for LM pre-training.
Then, building on these theoretical results, we introduce PDS, a practical framework that solves the PMP conditions in practice based on 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 data-constrained settings, which mitigates the pre-training data exhaustion problem.

## Appendix 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 [] constants, and minimizing AUC essentially improves the scaling properties of LMs. As suggested by existing literature [, ], the test losses of LMs follow a power-law with respect to the training steps $t$ after a warmup stage:

$$ $L(t)=\frac{C}{t^{c}}+L^{\mathrm{irre}},\ \ t>T_{0},$ (10) $$

where $C$ and $c$ are scaling law constants, $L^{\mathrm{irre}}$ is the irreducible loss related to the noise in the test set, and $T_{0}$ is the end steps of the warmup stage. $L^{\mathrm{irre}}$ is invariant to optimizing pre-training data selection strategies. Therefore, we care about the reducible loss, whose constants depend on the data quality scores ${\bm{\gamma}}$:

$$ $L^{\mathrm{re}}(t)=L(t)-L^{\mathrm{irre}}=\frac{C({\bm{\gamma}})}{t^{c({\bm{ \gamma}})}},\ \ t>T_{0}.$ (11) $$

We then consider the AUC of the reducible loss for sufficiently large $T$:

$$ $\text{AUC}({\bm{\gamma}})=\int_{t=T_{0}}^{T}\frac{C({\bm{\gamma}})}{t^{c({\bm{ \gamma}})}}\mathrm{d}t=\frac{C({\bm{\gamma}})}{1-c({\bm{\gamma}})}(T^{1-c({\bm {\gamma}})}-T_{0}^{1-c({\bm{\gamma}})}).$ (12) $$

For $c({\bm{\gamma}})<1$, when $T$ is sufficiently large, $\text{AUC}({\bm{\gamma}})\approx\frac{C({\bm{\gamma}})}{1-c({\bm{\gamma}})}T^{
1-c({\bm{\gamma}})}$. Minimizing AUC causes $B({\bm{\gamma}})$ to decrease and $\beta({\bm{\gamma}})$ to increase(^3^33$f(c)=\frac{1}{1-c}T^{1-c}$ is increasing with respect to $c$ when $1-c<\ln T$, which is easily satisfied for sufficiently large $T$.), improving the scaling properties of LMs. For $c({\bm{\gamma}})>1$, $\text{AUC}({\bm{\gamma}})\approx\frac{C({\bm{\gamma}})}{c({\bm{\gamma}})-1}
\frac{1}{T_{0}^{c({\bm{\gamma}})-1}}$, which also exhibit a trend of decreasing $C({\bm{\gamma}})$ and increasing $c({\bm{\gamma}})$ when AUC is minimized.

## Appendix B Proof of Theorem 2.1

To prove Theorem 2.1, we first describe the standard discrete-time Pontryagin’s Maximum Principle (PMP; ) in Optimal Control [] for time-variant control variables:

###### Theorem B.1 (PMP) .

Consider the following optimization problem in a discrete dynamical system:

$$ $\displaystyle\min_{{\bm{\gamma}}_{t}}\sum_{t=0}^{T-1}{\mathcal{J}}({\bm{\theta }}_{t},{\bm{\gamma}}_{t})+J({\bm{\theta}}_{T}),$ (13) s.t. $\displaystyle{\bm{\theta}}_{t+1}=f({\bm{\theta}}_{t},{\bm{\gamma}}_{t}),\ \ { \bm{\gamma}}_{t}\in U,$ $$

where the state variable ${\bm{\theta}}_{t}\in{\mathbb{R}}^{N}$, the control variable ${\bm{\gamma}}_{t}\in{\mathbb{R}}^{D}$, and ${\mathcal{J}}:{\mathbb{R}}^{N\times D}\mapsto{\mathbb{R}}$, $f:{\mathbb{R}}^{N\times D}\mapsto{\mathbb{R}}^{N}$ are continuous in ${\mathbb{R}}^{N\times D}$. Let ${\bm{\gamma}}^{*}_{t}$ be the solution to this problem, and ${\bm{\theta}}^{*}_{t}$ denote the corresponding state variable. For $0\leq t<T$, there exists a co-state vector ${\bm{\lambda}}^{*}_{t}\in{\mathbb{R}}^{N}$ such that

$$ $\displaystyle{\bm{\theta}}^{*}_{t+1}$ $\displaystyle=\nabla_{{\bm{\lambda}}}H({\bm{\theta}}^{*}_{t},{\bm{\lambda}}_{t +1}^{*},{\bm{\gamma}}^{*}_{t}),\ \ {\bm{\theta}}^{*}_{0}={\bm{\theta}}_{0},$ (14) $\displaystyle{\bm{\lambda}}^{*}_{t}$ $\displaystyle=\nabla_{{\bm{\theta}}}H({\bm{\theta}}^{*}_{t},{\bm{\lambda}}_{t+ 1}^{*},{\bm{\gamma}}^{*}_{t}),\ \ {\bm{\lambda}}^{*}_{T}=\nabla J({\bm{\theta} }_{T}),$ (15) $\displaystyle{\bm{\gamma}}^{*}_{t}$ $\displaystyle=\arg\min_{{\bm{\gamma}}_{t}}H({\bm{\theta}}^{*}_{t},{\bm{\lambda }}_{t+1}^{*},{\bm{\gamma}}_{t}),\ \ {\bm{\gamma}}_{t}\in U,$ (16) $$

where $H:{\mathbb{R}}^{N}\times{\mathbb{R}}^{N}\times{\mathbb{R}}^{D}\mapsto{\mathbb{
R}}$ is the Hamiltonian function defined by

$$ $H({\bm{\theta}},{\bm{\lambda}},{\bm{\gamma}})={\mathcal{J}}({\bm{\theta}},{\bm {\gamma}})+{{\bm{\lambda}}}^{\top}f({\bm{\theta}},{\bm{\gamma}}).$ (17) $$

Proof of Theorem B.1 can be found in various textbooks in Optimal Control [, ].
In our context, we interpret ${\bm{\theta}}_{t}$ as the model parameters, $J(\cdot)$ as the downstream loss function, $U$ as the $|\mathcal{D}|$-dimensional simplex as defined in Section 2.1 and $f$ as the GD operation where the data weights ${\bm{\gamma}}_{t}$ changes with respect to the training steps $t$:

$$ $\displaystyle{\bm{\theta}}_{t+1}$ $\displaystyle={\bm{\theta}}_{t}-\eta\nabla\sum_{n=1}^{|\mathcal{D}|}\gamma_{n, t}l(x_{n},{\bm{\theta}}_{t})$ (18) $\displaystyle={\bm{\theta}}_{t}-\eta\nabla L({\bm{\theta}}_{t},{\bm{\gamma}}_{ t}).$ $$

In this way, the Hamilton function is

$$ $H({\bm{\theta}},{\bm{\lambda}},{\bm{\gamma}})={\mathcal{J}}({\bm{\theta}},{\bm {\gamma}})+{\bm{\lambda}}^{\top}\left[{\bm{\theta}}-\eta\nabla L({\bm{\theta}} ,{\bm{\gamma}})\right].$ (19) $$

The key challenge to prove Theorem 2.1 is that the control variables are constrained to 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-1$ equations:

$$ ${\bm{\gamma}}_{1}={\bm{\gamma}}_{0},\ {\bm{\gamma}}_{2}={\bm{\gamma}}_{0}, \cdots,{\bm{\gamma}}_{T-1}={\bm{\gamma}}_{0}.$ (20) $$

Therefore, the optimization of data selection, as in Eq. (3), is equivalent to the following problem:

$$ $\displaystyle\min_{{\bm{\gamma}}_{t}}\sum_{t=1}^{T}J({\bm{\theta}}_{t}),$ (21) s.t. $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta\nabla L({\bm{\theta}}_ {t},{\bm{\gamma}}_{t}),\ \ {\bm{\gamma}}_{t}\in U,$ $\displaystyle{\bm{\gamma}}_{0}={\bm{\gamma}}_{1}=\cdots={\bm{\gamma}}_{T-1}.$ $$

We adopt the method of Lagrange multipliers to solve Eq. (21) by considering the following optimization problem:

$$ $\displaystyle\min_{{\bm{\gamma}}_{t}}\sum_{t=0}^{T-1}J({\bm{\theta}}_{t})+\sum _{t=1}^{T-1}\sum_{n=1}^{|\mathcal{D}|}\mu_{n,t}(\gamma_{n,t}-\gamma_{n,0})+J({ \bm{\theta}}_{T}),$ (22) s.t. $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta\nabla L({\bm{\theta}}_ {t},{\bm{\gamma}}_{t}),\ \ {\bm{\gamma}}_{t}\in U,$ $$

where $(\mu_{n,t})_{1\leq n\leq D,0\leq t\leq T-1}$ are Lagrange multipliers. Note that we split $J({\bm{\theta}}_{T})$ out and add $J({\bm{\theta}}_{0})$ to the sum of $J({\bm{\theta}}_{t})$, which does not affect the minimization. In this way, when ${\mathcal{J}}({\bm{\theta}}_{t},{\bm{\gamma}}_{t})$ takes the following form:

$$ ${\mathcal{J}}({\bm{\theta}}_{t},{\bm{\gamma}}_{t})=\begin{dcases}J({\bm{\theta }}_{0})-\sum_{t^{\prime}=1}^{T-1}\sum_{n=1}^{|\mathcal{D}|}\mu_{n,t^{\prime}} \gamma_{n,0},&\text{if }t=0\\ J({\bm{\theta}}_{t})+\sum_{n=1}^{|\mathcal{D}|}\mu_{n,t}\gamma_{n,t},&\text{if }1\leq t\leq T-1\end{dcases},$ (23) $$

we 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:

$$ $\displaystyle{\bm{\theta}}_{t}^{*}$ $\displaystyle={\bm{\theta}}^{*}_{t}-\eta\nabla L({\bm{\theta}}^{*}_{t},{\bm{ \gamma}}^{*}_{0}),\ \ \ {\bm{\theta}}_{0}^{*}={\bm{\theta}}_{0},$ (24) $\displaystyle{\bm{\lambda}}^{*}_{t}$ $\displaystyle={\bm{\lambda}}^{*}_{t+1}+\nabla J({\bm{\theta}}_{t}^{*})-\eta \nabla^{2}L({\bm{\theta}}^{*}_{t},{\bm{\gamma}}^{*}_{0}){\bm{\lambda}}^{*}_{t+ 1},\ \ \ {\bm{\lambda}}^{*}_{T}=\nabla J({\bm{\theta}}_{T}),$ (25) $$

which prove Eq. (4) and Eq. (5) in Theorem 2.1 when we set ${\bm{\gamma}}^{*}_{0}={\bm{\gamma}}^{*}$. By substituting Eq. (19) and Eq. (23) into Eq. (16), we have

$$ ${\bm{\gamma}}^{*}_{t}=\begin{dcases}\arg\max_{{\bm{\gamma}}_{0}}\sum_{n=1}^{| \mathcal{D}|}\gamma_{n,0}\left[{{\bm{\lambda}}_{1}^{*}}^{\top}\nabla l(x_{n},{ \bm{\theta}}_{0}^{*})+\sum_{t^{\prime}=1}^{T-1}\mu_{n,t^{\prime}}\right],& \text{if }t=0\\ \arg\max_{{\bm{\gamma}}_{t}}\sum_{n=1}^{|\mathcal{D}|}\gamma_{n,t}\left[{{\bm{ \lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{\theta}}^{*}_{t})-\mu_{n,t} \right],&\text{if }1\leq t\leq T-1\end{dcases}$ (26) $$

Considering the time-invariant constraint in Eq. (20), we set ${\bm{\gamma}}^{*}_{0}={\bm{\gamma}}^{*}_{1}=\cdots={\bm{\gamma}}^{*}_{T-1}={
\bm{\gamma}}^{*}$ and get

$$ $\begin{dcases}{\bm{\gamma}}^{*}=\arg\max_{{\bm{\gamma}}}\sum_{n=1}^{|\mathcal{ D}|}\gamma_{n}\left[\eta{{\bm{\lambda}}_{1}^{*}}^{\top}\nabla l(x_{n},{\bm{ \theta}}_{0}^{*})+\sum_{t^{\prime}=1}^{T-1}\mu_{n,t^{\prime}}\right],&\text{if }t=0\\ {\bm{\gamma}}^{*}=\arg\max_{{\bm{\gamma}}}\sum_{n=1}^{|\mathcal{D}|}\gamma_{n} \left[\eta{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{\theta}}^{*}_{t })-\mu_{n,t}\right],&\text{if }1\leq t\leq T-1\end{dcases},$ (27) $$

which forms a complete equation system containing $T$ equations and $T$ unknown variables ($T-1$ number of $\bm{\mu}_{t}=\left[\mu_{1,t},\mu_{2,t},\cdots,\mu_{{|\mathcal{D}|},t}\right]$ plus one ${\bm{\gamma}}^{*}$), which has the solution:

$$ $\displaystyle\mu_{n,t}$ $\displaystyle=\eta{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{\theta} }^{*}_{t})-\eta\frac{S_{n}}{T},\ \ 1\leq n\leq{|\mathcal{D}|},\ \ 0\leq t\leq T -1,$ (28) $\displaystyle{\bm{\gamma}}^{*}$ $\displaystyle=\arg\max_{{\bm{\gamma}}}\sum_{n=1}^{|\mathcal{D}|}\gamma_{n}\eta \frac{S_{n}}{T},$ (29) $$

where $S_{n}=\sum_{t=0}^{T-1}{{\bm{\lambda}}^{*}_{t+1}}^{\top}\nabla l(x_{n},{\bm{
\theta}}^{*}_{t})$. Note that Eq. (29) is equivalent to Eq. (6). So far, Theorem 2.1 is proved by combining Eq. (24), Eq. (25), and Eq. (29).

###### Theorem B.1 (PMP) .

## Appendix C Derivation for Adam

We provide our formulation and derivation for Adam [] in this section. For $0\leq t\leq T-1$, the parameter update rules of Adam is given by

$$ $\displaystyle{\bm{m}}_{t+1}$ $\displaystyle=\beta_{1}{\bm{m}}_{t}+(1-\beta_{1}){\bm{g}}_{t},$ (30) $\displaystyle{\bm{v}}_{t+1}$ $\displaystyle=\beta_{2}{\bm{v}}_{t}+(1-\beta_{2}){\bm{g}}_{t}^{2},$ $\displaystyle\widehat{{\bm{m}}}_{t+1}$ $\displaystyle={\bm{m}}_{t+1}/(1-\beta_{1}^{t+1}),$ $\displaystyle\widehat{{\bm{v}}}_{t+1}$ $\displaystyle={\bm{v}}_{t+1}/(1-\beta_{2}^{t+1}),$ $\displaystyle{\bm{\theta}}_{t+1}$ $\displaystyle={\bm{\theta}}_{t}-\eta\widehat{{\bm{m}}}_{t+1}/(\sqrt{\widehat{{ \bm{v}}}_{t+1}}+\epsilon),$ $$

where $\beta_{1}$, $\beta_{2}$, $\epsilon$, $\eta$ are hyper-parameters and ${\bm{g}}_{t}=\nabla L({\bm{\theta}}_{t},{\bm{\gamma}}_{t},\mathcal{D})$.
We set ${\bm{m}}_{0}=\bm{0}$ and ${\bm{v}}_{0}=\bm{0}$.
To formulate the LM training with Adam as an Optimal Control [] problem, we treat the vector $\bm{\Theta}_{t}=\left[{\bm{\theta}}_{t},{\bm{m}}_{t},{\bm{v}}_{t}\right]^{\top
}\in\mathbb{R}^{3N}$ as the state variable. Let $F$ denote the state-transition function from $\bm{\Theta}_{t}$ to $\bm{\Theta}_{t+1}$:

$$ ${\bm{\theta}}_{t+1}=F(\bm{\Theta}_{t},{\bm{\gamma}}),$ (31) $$

and thus $F$ represents the following relations:

$$ $\displaystyle{\bm{\theta}}_{t+1}=$ $\displaystyle{\bm{\theta}}_{t}-\frac{\eta}{1-\beta_{1}^{t+1}}\frac{\beta_{1}{ \bm{m}}_{t}+(1-\beta_{1}){\bm{g}}_{t}}{\sqrt{(\beta_{2}{\bm{v}}_{t}+(1-\beta_{ 2}){\bm{g}}_{t}^{2})/(1-\beta_{2}^{t+1})}+\epsilon},$ (32) $\displaystyle{\bm{m}}_{t+1}=$ $\displaystyle\frac{\beta_{1}{\bm{m}}_{t}+(1-\beta_{1}){\bm{g}}_{t}}{1-\beta_{1 }^{t+1}},$ (33) $\displaystyle{\bm{v}}_{t+1}=$ $\displaystyle\frac{\beta_{2}{\bm{v}}_{t}+(1-\beta_{2}){\bm{g}}_{t}^{2}}{1- \beta_{2}^{t+1}}.$ (34) $$

Similar to GD, we can still define the Hamiltonian function of the problem and obtain the following theorem with Pontryagin’s Maximum Principle (PMP; ):

###### Theorem C.1 (PMP Data Selection for Adam) .

Let ${\bm{\gamma}}^{*}$ solve the problem in Eq. (3), and $\bm{\Theta}^{*}_{t}$ denote the state variable corresponding to ${\bm{\gamma}}^{*}$. Then, there exists a co-state vector $\bm{\Lambda}^{*}_{t}\in{\mathbb{R}}^{3N}$ such that

$$ $\displaystyle\bm{\Theta}^{*}_{t+1}$ $\displaystyle=\nabla_{\bm{\Lambda}}{\mathcal{H}}(\bm{\Theta}^{*}_{t},\bm{ \Lambda}^{*}_{t+1},{\bm{\gamma}}^{*}),\ \ \bm{\Theta}^{*}_{0}=\left[{\bm{ \theta}}_{0},\bm{0},\bm{0}\right]^{\top},$ (35) $\displaystyle\bm{\Lambda}^{*}_{t}$ $\displaystyle=\nabla_{\bm{\Theta}}{\mathcal{H}}(\bm{\Theta}^{*}_{t},\bm{ \Lambda}^{*}_{t+1},{\bm{\gamma}}^{*}),\ \ \bm{\Lambda}^{*}_{T}=\left[\nabla J( {\bm{\theta}}_{T}),\bm{0},\bm{0}\right]^{\top},$ (36) $\displaystyle{\bm{\gamma}}^{*}$ $\displaystyle=\arg\min_{{\bm{\gamma}}}\sum_{t=0}^{T-1}{\mathcal{H}}(\bm{\Theta }^{*}_{t},\bm{\Lambda}^{*}_{t+1},{\bm{\gamma}}),\ \ {\bm{\gamma}}\in U,$ (37) $$

where ${\mathcal{H}}:{\mathbb{R}}^{3N}\times{\mathbb{R}}^{3N}\times{\mathbb{R}}^{D}
\mapsto{\mathbb{R}}$ is the Hamiltonian function defined by

$$ ${\mathcal{H}}(\bm{\Theta},{\bm{\lambda}},{\bm{\gamma}})=J({\bm{\theta}})+{\bm{ \Lambda}}^{\top}F(\bm{\Theta},{\bm{\gamma}}).$ (38) $$

Similar to the derivation for GD, Eq. (35) controls the state transition, equivalent to Eq. (31).
To simplify the further derivation of Eq. (36) and Eq. (37), we assume $\frac{\partial{\bm{v}}_{t+1}}{\partial{\bm{g}}_{t}}\approx\bm{0}$, which is reasonable because ${\bm{v}}_{t+1}$ is an exponential moving average of ${\bm{g}}_{t}^{2}$ and the weight $1-\beta_{2}$ is usually much smaller than 1 in practice.
Therefore, we have $\frac{\partial v_{t+1}}{\partial{\bm{\theta}}_{t}}\approx\bm{0}$, $\frac{\partial v_{t+1}}{\partial{\bm{\gamma}}_{t}}\approx\bm{0}$ and thus Eq. (36) can be written to

$$ $\displaystyle\bm{\Lambda}^{*}_{t}=$ $\displaystyle\left[\bm{\Lambda}^{(1)}_{t},\bm{\Lambda}^{(2)}_{t},\bm{\Lambda}^ {(3)}_{t}\right]^{\top},$ (39) $\displaystyle\bm{\Lambda}^{(1)}_{t}\approx$ $\displaystyle\nabla J({\bm{\theta}}^{*}_{t})+\bm{\Lambda}_{t+1}^{(1)}-\frac{(1 -\beta_{1})\eta}{1-\beta_{1}^{t+1}}\nabla^{2}({\bm{\theta}}^{*}_{t},{\bm{ \gamma}}^{*})\frac{\bm{\Lambda}_{t+1}^{(1)}}{\sqrt{\widehat{{\bm{v}}}_{t+1}+ \epsilon}}$ $\displaystyle+\frac{(1-\beta_{1})}{1-\beta_{1}^{t+1}}\nabla^{2}({\bm{\theta}}^ {*}_{t},{\bm{\gamma}}^{*})\bm{\Lambda}_{t+1}^{(2)},$ (40) $\displaystyle\bm{\Lambda}^{(2)}_{t}=$ $\displaystyle-\frac{\eta\beta_{1}}{1-\beta_{1}^{t+1}}\frac{\bm{\Lambda}_{t+1}^ {(1)}}{\sqrt{\widehat{{\bm{v}}}_{t+1}}+\epsilon}+\frac{\beta_{1}}{1-\beta_{1}^ {t+1}}\bm{\Lambda}_{t+1}^{(2)},$ (41) $\displaystyle\bm{\Lambda}^{(3)}_{t}=$ $\displaystyle\frac{\eta\beta_{2}}{1-\beta_{2}^{t+1}}\frac{\widehat{{\bm{m}}}_{ t+1}\bm{\Lambda}_{t+1}^{(1)}}{\sqrt{\widehat{{\bm{v}}}_{t+1}}\left(\sqrt{ \widehat{{\bm{v}}}_{t+1}+\epsilon}\right)^{2}}+\frac{\beta_{2}}{1-\beta_{2}^{t +1}}\bm{\Lambda}_{t+1}^{(3)}.$ (42) $$

To achieve the minimum in Eq. (37), we need to compute the gradient of ${\mathcal{H}}$ with respect to ${\bm{\gamma}}_{t}$:

$$ $\nabla_{\gamma_{n}}{\mathcal{H}}(\bm{\Theta}^{*}_{t},\bm{\Lambda}^{*}_{t+1},{ \bm{\gamma}})=\frac{\eta(1-\beta_{1})}{1-\beta_{1}^{t+1}}{\bm{\Lambda}^{(1)}_{ t+1}}^{\top}\frac{\nabla l(x_{n},{\bm{\theta}}_{t})}{\sqrt{\widehat{{\bm{v}}}_ {t+1}}+\epsilon}+\frac{1-\beta_{1}}{1-\beta_{1}^{t+1}}{\bm{\Lambda}^{(2)}_{t+1 }}^{\top}\nabla l(x_{n},{\bm{\theta}}_{t})$ (43) $$

In this way, we can use Algorithm 2 to solve for the data quality scores on the proxy dataset $\mathcal{D}^{\mathrm{prx}}$, just like Algorithm 1 in Section 2.3.1. Then, we train a data scorer with the solved data weights ${\bm{\gamma}}^{*}$ as in Section 2.3.2 and conduct data selection based on the scores inferred by the data scorer on $\mathcal{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 on GD. Given that Algorithm 2 requires twice as much computation and 3 times as much GPU memory as Algorithm 1, we adopt PMP condition for data selection based on GD (Theorem 2.1) to build PDS in our main paper.

Figure: Algorithm 2 PMP Solver for Adam

###### Theorem C.1 (PMP Data Selection for Adam) .

## Appendix D Algorithm 1 as Proximal Gradient Decent

We provide another view of Algorithm 1 as using Proximal Gradient Decent method []. Specifically, we can optimize Eq. (3) with the following rules:

$$ $\displaystyle A({\bm{\gamma}})$ $\displaystyle=\sum_{t^{\prime}=1}^{T}J({\bm{\theta}}_{t^{\prime}}),$ (44) $\displaystyle{\bm{\gamma}}$ $\displaystyle\leftarrow\operatorname{Proj}\left[{\bm{\gamma}}-\alpha^{\prime} \nabla_{{\bm{\gamma}}}A\right],$ (45) $$

where $A({\bm{\gamma}})$ denotes the cost function in Eq. (3), $\alpha^{\prime}$ is the learning rate and $\operatorname{Proj}[\cdot]$ projects a point in ${\mathbb{R}}^{D}$ to the $D$-simplex.
$\nabla_{{\bm{\gamma}}}A=\left[\frac{\partial A}{\partial\gamma_{1}},\frac{
\partial A}{\mathrm{d}\gamma_{2}},\cdots,\frac{\partial A}{\partial\gamma_{D}}\right]$ is the gradient of $A({\bm{\gamma}})$ with respect to the data weights ${\bm{\gamma}}$. Now, we compute each element of $\nabla_{{\bm{\gamma}}}A$ with the chain rule:

$$ $\displaystyle\frac{\partial A}{\partial\gamma_{n}}$ $\displaystyle=\sum_{t^{\prime}=1}^{T}\frac{\partial J({\bm{\theta}}_{t^{\prime }})}{\partial\gamma_{n}}$ (46) $\displaystyle=\sum_{t^{\prime}=1}^{T}\nabla J({\bm{\theta}}_{t^{\prime}})^{ \top}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial\gamma_{n}}$ $\displaystyle=\sum_{t^{\prime}=1}^{T}\nabla J({\bm{\theta}}_{t^{\prime}})^{ \top}\sum_{t=1}^{t^{\prime}}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial {\bm{\theta}}_{t}}\frac{\partial{\bm{\theta}}_{t}}{\partial\gamma_{n}}$ $\displaystyle=-\eta\sum_{t^{\prime}=1}^{T}\nabla J({\bm{\theta}}_{t^{\prime}}) ^{\top}\sum_{t=1}^{t^{\prime}}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{ \partial{\bm{\theta}}_{t}}\nabla l(x_{n},{\bm{\theta}}_{t-1})$ $\displaystyle=-\eta\sum_{t^{\prime}=1}^{T}\sum_{t=0}^{t^{\prime}-1}\nabla J({ \bm{\theta}}_{t^{\prime}})^{\top}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{ \partial{\bm{\theta}}_{t+1}}\nabla l(x_{n},{\bm{\theta}}_{t})$ $\displaystyle=-\eta{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb} {1,0,0}\sum_{t=0}^{T-1}}\left[{\color[rgb]{1,0,0}\definecolor[named]{ pgfstrokecolor}{rgb}{1,0,0}\sum_{t^{\prime}=t+1}^{T}}\nabla J({\bm{\theta}}_{t ^{\prime}})^{\top}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{ \theta}}_{t+1}}\right]\nabla l(x_{n},{\bm{\theta}}_{t}),$ $$

where $\frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t+1}}$ satisfies

$$ $\displaystyle\frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{ t}}=\frac{\partial{\bm{\theta}}_{t+1}}{\partial{\bm{\theta}}_{t}}\frac{ \partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t+1}}=\left[{\bm{I} }-\eta\nabla^{2}L({\bm{\theta}}_{t},{\bm{\gamma}}_{t})\right]\frac{\partial{ \bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t+1}},$ (47) $$

according to Eq. (2). In the following, we show that:

$$ ${\bm{\lambda}}_{t+1}=\sum_{t^{\prime}=t+1}^{T}\frac{\partial{\bm{\theta}}_{t^{ \prime}}}{\partial{\bm{\theta}}_{t+1}}\nabla J({\bm{\theta}}_{t^{\prime}}),$ (48) $$

where ${\bm{\lambda}}_{t+1}$ is the co-state vector in Algorithm 1. Let ${\bm{\lambda}}_{t+1}^{\prime}=\sum_{t^{\prime}=t+1}^{T}\frac{\partial{\bm{
\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t+1}}\nabla J({\bm{\theta}}_{t^{
\prime}})$, we then have ${\bm{\lambda}}^{\prime}_{T}=\nabla J({\bm{\theta}}_{T})$ and the following difference equation for ${\bm{\lambda}}_{t}^{\prime}$:

$$ $\displaystyle{\bm{\lambda}}^{\prime}_{t}=$ $\displaystyle\sum_{t^{\prime}=t}^{T}\frac{\partial{\bm{\theta}}_{t^{\prime}}}{ \partial{\bm{\theta}}_{t}}\nabla J({\bm{\theta}}_{t^{\prime}})$ (49) $\displaystyle=$ $\displaystyle\nabla J({\bm{\theta}}_{t^{\prime}})+\sum_{t^{\prime}=t+1}^{T} \frac{\partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t}}\nabla J({ \bm{\theta}}_{t^{\prime}})$ $\displaystyle=$ $\displaystyle\nabla J({\bm{\theta}}_{t^{\prime}})+\sum_{t^{\prime}=t+1}^{T} \left[{\bm{I}}-\eta\nabla^{2}L({\bm{\theta}}_{t},{\bm{\gamma}})\right]\frac{ \partial{\bm{\theta}}_{t^{\prime}}}{\partial{\bm{\theta}}_{t+1}}\nabla J({\bm{ \theta}}_{t^{\prime}})$ $\displaystyle=$ $\displaystyle\nabla J({\bm{\theta}}_{t^{\prime}})+{\bm{\lambda}}^{\prime}_{t+1 }-\eta\nabla^{2}L({\bm{\theta}}_{t},{\bm{\gamma}}){\bm{\lambda}}^{\prime}_{t+1}.$ $$

Given that ${\bm{\lambda}}_{t+1}^{\prime}$ satisfies the same difference equation as ${\bm{\lambda}}_{t+1}$ in Algorithm 1 and has the same value at $t=T$, we have ${\bm{\lambda}}_{t+1}^{\prime}={\bm{\lambda}}_{t+1}$. Therefore, the gradient in Eq. (46) can be written as

$$ $\displaystyle\frac{\partial J}{\partial\gamma_{n}}=-\eta{\bm{\lambda}}_{t+1}^{ \top}\nabla l(x_{n},{\bm{\theta}}_{t}).$ (50) $$

Combining Eq. (45) and Eq. (50), we recover the update rules of $\gamma_{n,t}$ in Algorithm 1 by setting $\alpha^{\prime}=\alpha/\eta$.

## Appendix E Implementation Details

### E.1 Solving Data Quality Scores

Algorithm 1 needs to compute the Hessian matrix, as in Eq. (5), and per-instance vector-gradient product, as in Eq. (6). We adopt the Jacobian-Vector-Product(^4^44[https://pytorch.org/docs/stable/func.api.html](https://pytorch.org/docs/stable/func.api.html)) in PyTorch [] to efficiently implement these operations. There is still a large room for this implementation, such as adopting other deep learning frameworks [, ] or using lower-complexity algorithms [].
Algorithm 1 requires storing all the LM checkpoints ${\bm{\theta}}_{t}$ from $t=0$ to $t=T-1$ in the forward inner loop for computing ${\bm{\lambda}}_{t}$ in the reverse inner loop.
To reduce the single-device GPU memory use, we implement an activation partition algorithm inspired by ZeRO-2 [], where ${\bm{\theta}}_{t}$ in one inner-loop pass are stored in different GPU devices. We also adopt a strategy inspired by activation checkpointing [].

### E.2 Training Data Scorer

We fine-tune the Fairseq-Dense-125M model [] on the solved data weights ${\bm{\gamma}}^{*}$ to 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 [] optimizer for 5 epochs, using a learning rate $1\times 10^{-4}$ and a batch size 512. We split 10% samples from $\mathcal{D}^{\mathrm{prx}}$ for validation and select the checkpoint achieving the highest Spearman correlation score on the validation set to infer data quality scores in $\mathcal{D}$.

### E.3 Pre-Training

We pre-train all our models for about 50B tokens with a batch size of 512 and a max input length of 1,024. We use the AdamW [] 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 [] to set the maximal learning rates as listed in Table 6, together with the model configurations.

**Table 6: Model configurations and corresponding learning rates.**
| Model Size | $d_{\text{model}}$ | $d_{\text{FFN}}$ | $n_{\text{layers}}$ | $n_{\text{head}}$ | $d_{\text{head}}$ | learning rate |
| --- | --- | --- | --- | --- | --- | --- |
| 160M | 768 | 3,072 | 12 | 12 | 64 | $6\times 10^{-4}$ |
| 470M | 1,024 | 4,096 | 24 | 16 | 64 | $3\times 10^{-4}$ |
| 1B | 1,536 | 6,144 | 24 | 16 | 96 | $2.5\times 10^{-4}$ |
| 1.7B | 2,048 | 8,192 | 24 | 16 | 128 | $2\times 10^{-4}$ |

### E.4 Evaluation

We evaluate our trained models on MMLU [] and the evaluation sets used in OLMo [], including Hellaswag (HS; ), Winograde (Wino.; ), LAMBADA (LAMB; ), OpenbookQA (OBQA; ), ARC-easy/challenge (ARC-e/c; ), PIQA [], SciQ [], and BoolQ []. We adopt the LM-evaluation-harness library [](^5^55[https://github.com/t1101675/lm-harness/](https://github.com/t1101675/lm-harness/)) to conduct the zero-shot evaluation, which formulates the tasks as the multiple-choice problems and chooses the answer by comparing the answer-length-normed loss across candidates (acc_norm). We also compute the test loss of our trained models on the DCLM corpus [], a 10K subset uniformly sampled from the high-quality pre-training corpus in Li et al. [].

### E.5 Baselines

##### Reducible Holdout Loss Selection (RHO-Loss; 58 , 51 ).

RHO-Loss selects data with high “reducible holdout loss” defined as follows:

$$ $\text{RHO-Loss}_{n,t}=l(x_{n},{\bm{\theta}}_{t})-l(x_{n},{\bm{\theta}}^{\text{ down}}),$ (51) $$

where ${\bm{\theta}}^{\text{down}}$ is the model parameters trained on the downstream tasks, which is LIMA [] in our experiments. We first split 10% of LIMA for validation and choose ${\bm{\theta}}^{\text{down}}$ with the lowest validation loss. Then, we train ${\bm{\theta}}_{t}$ while selecting $x_{n}$ with top $\alpha\times 100\%$ $\text{RHO-Loss}_{n,t}$ to compute gradients. We tried $\alpha\in\{0.1,0.3,0.5,0.7\}$ and find that $\alpha=0.5$ performs the best.

##### Data Selection via Importance Resampling (DSIR; 79 ).

DSIR selects pre-training data based on the n-gram feature overlap between the instances in the downstream dataset (LIMA [] 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 DSIR(^6^66[https://github.com/p-lambda/dsir](https://github.com/p-lambda/dsir)).

##### Influence Score (IF-Score 43 , 32 ).

We adopt the influence function [, ] to measure the quality of $x$ using the downstream loss $J({\bm{\theta}})$ as the target:

$$ $\text{IF-Score}(x)=\nabla l(x,{\bm{\theta}}_{T})^{\top}\left[\nabla^{2}L({\bm{ \theta}}_{T})\right]^{-1}\nabla J({\bm{\theta}}_{T}),$ (52) $$

where ${\bm{\theta}}_{T}$ is the LM parameters trained for $T$ steps and $L({\bm{\theta}}_{T})=\frac{1}{|\mathcal{D}|}\sum_{x\in\mathcal{D}}l(x,{\bm{
\theta}}_{T})$. We adopt a linear-time iterative algorithm to compute the inverse-Hessian-vector-product [].
To reduce the computational cost, we adopt a similar efficient implementation as PDS by computing the IF-Scores on a small proxy data based on a small proxy LM and then transferring these scores to a large pre-training corpus with a fine-tuned data scorer. We select examples with the top 40% scores inferred by the data scorer on the large pre-training corpus.

### E.6 Simulated Setting for Experiments in Table 6

To exactly run Algorithm 1 with a feasible computational overhead, we adopt a 12M LM from the Mistral [] family, with a small vocabulary. We uniformly sample 4,096 instances as $\mathcal{D}$, with a max sequence length of 256, and construct a 16K vocabulary from $\mathcal{D}$ 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$, $T^{\mathrm{prx}}=100$. We run the inner loop using SGD with a batch size of 128. The outer loop epoch number is set to 1.

## Appendix F Complexity Analysis

Following Hoffmann et al. [], for an LM with $N$ parameters to be trained on $D$ tokens, we assume the computational FLOPs of a forward and a backward pass are $2ND$ and $4ND$, respectively. We compute the FLOPs and asymptotic complexities of different stages in PDS as follows:

- •
Solving Data Quality Scores: According to Section 2.3.1, we first pre-train a proxy LM on $\mathcal{D}$ which consumes $6N^{\mathrm{prx}}D$ FLOPs. Then, we perform Algorithm 1 $M$ times on $\mathcal{D}^{\mathrm{prx}}$ based on the proxy LM. The forward inner loop in Algorithm 1 consumes $6N^{\mathrm{prx}}D^{\mathrm{prx}}$ FLOPs. The reverse inner loop can be treated as the “backward” propagation of the forward inner loop as discussed in Appendix D, which consumes $2\times 6N^{\mathrm{prx}}D^{\mathrm{prx}}$ FLOPs. The update of ${\bm{\gamma}}$ results in one forward and backward pass of the proxy LM on $\mathcal{D}^{\mathrm{prx}}$, which consumes $6N^{\mathrm{prx}}D^{\mathrm{prx}}$ FLOPs. In summary, the asymptotic complexity of solving data quality scores is $O(N^{\mathrm{prx}}D+4MN^{\mathrm{prx}}D^{\mathrm{prx}})$.
- •
Data Scorer: The data scorer with $N^{\text{score}}$ is trained on $\mathcal{D}^{\mathrm{prx}}$ and used to infer data quality scores on $\mathcal{D}$. Therefore, the computation overhead is $6N^{\text{score}}\mathcal{D}^{\mathrm{prx}}+2N^{\text{score}}\mathcal{D}$ and the asymptotic complexity is $O(3N^{\text{score}}\mathcal{D}^{\mathrm{prx}}+N^{\text{score}}\mathcal{D})$.
- •
Data Selection: Selecting pre-training corpus requires iterating over $\mathcal{D}$, whose asymptotic complexity is $O(D)$. This process can be done on CPUs and does not require GPU FLOPs.
- •
Pre-Training: Pre-training an LM with $N$ parameters requires $6ND$ FLOPs, whose asymptotic complexity is $O(ND)$.

## Appendix G More Results

### G.1 Scaling Curves of Computation for Other Model Sizes

Figure: (a) Model Size = 160M
Refer to caption: x9.png

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.

### G.2 Test Loss Extrapolation with the Scaling Law

**Table 7: Scaling law constants by fitting the test losses on the DCLM corpus.**
|  | $A$ | $B$ | $\alpha$ | $\beta$ | $E$ |
| --- | --- | --- | --- | --- | --- |
| Conventional | $8.09\times 10^{2}$ | $7.50\times 10^{5}$ | 0.397 | 0.651 | 2.829 |
| PDS | $6.21\times 10^{3}$ | $1.76\times 10^{5}$ | 0.518 | 0.585 | 2.829 |

We extrapolate the test losses on the DCLM corpus [] of the conventionally trained and PDS-trained LMs with the Scaling Law [, ]. Following Hoffmann et al. [], we consider the scaling law with the following form:

$$ $L(N,D)=E+\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}},$ (53) $$

where $N$ is the model parameters, $D$ is the trained tokens, and $A$, $B$, $E$, $\alpha$, $\beta$ are constants. We obtain these constants by minimizing the Huber loss []:

$$ $\min_{a,b,e,\alpha,\beta}\sum_{(N_{i},D_{i},L_{i})}\text{Huber}_{\delta}(\text {LSE}(a-\alpha\log N_{i},b-\beta\log D_{i},e)-\log L_{i}),$ (54) $$

where $\text{LSE}(\cdot)$ is the log-sum-exp operation. The loss is summed over all $(N_{i},D_{i},L_{i})$, which is obtained by the test losses of 160M, 470M, 1B, and 1.7B LM during training from 0B to 50B tokens. We record the losses every 2.5B tokens, resulting in a total $4\times 50B/2.5B=80$ tuples like $(N_{i},D_{i},L_{i})$. After solving $a$, $b$, and $e$ from Eq. (54) , we have $A=\exp(a)$, $B=\exp(b)$, and $E=\exp(e)$.

Hoffmann et al. [] optimizes Eq. (54) using the LBFGS algorithm []. However, we find this algorithm sensitive to the initialization of the parameters to be optimized. Therefore, we apply a two-stage optimization. Specifically, we first fit the following data scaling curves for $N=$160M, 470M, 1B, and $1.7B$ with non-linear least squares from scipy.optimize.curve_fit(^7^77[https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html)), which is much more robust to the initialization:

$$ $L(D)=E^{\prime}(N)+\frac{B_{0}(N)}{D^{\beta_{0}(N)}},$ (55) $$

where $E^{\prime}(N)$, $B_{0}(N)$ and $\beta_{0}(N)$ are the fitted parameters. Then, we fit the following model size scaling curve:

$$ $E^{\prime}=E_{0}+\frac{A_{0}}{N^{\alpha_{0}}}.$ (56) $$

We use the constants from Eq. (56) and the average constants from Eq. (55) to compute the initialization for the LBFGS algorithm:

$$ $\displaystyle a_{0}$ $\displaystyle=\log A_{0},$ (57) $\displaystyle b_{0}$ $\displaystyle=\log\frac{B_{0}(\text{160M})+B_{0}(\text{470M})+B_{0}(\text{1B}) +B_{0}(\text{1.7B})}{4},$ $\displaystyle\alpha_{0}$ $\displaystyle=\alpha_{0},$ $\displaystyle\beta_{0}$ $\displaystyle=\frac{\beta_{0}(\text{160M})+\beta_{0}(\text{470M})+\beta_{0}( \text{1B})+\beta_{0}(\text{1.7B})}{4},$ $\displaystyle e_{0}$ $\displaystyle=\log E_{0},$ $$

where $a_{0},b_{0},\alpha_{0},\beta_{0},e_{0}$ are the parameter initialization for the LFBGS algorithm to optimize Eq. (54). We set $\delta=1\times 10^{-3}$ and learning rate to $0.05$ when running LFBGS and obtain the constants in Table 7. We use these constants and Eq. (53) to compute the predicted loss in Table 3.

In Section 3.3 (Data-Constrained Setting), to 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 in Table 3 and use $\frac{D}{4}$ as the token demand, indicating the LM can be trained for 4 epochs as suggested by Muennighoff et al. [].

### G.3 Ablation Studies

Figure: Figure 9: Effect of the data selection ratio $\alpha$. We report the average accuracy on the OLMo evaluation datasets for $\alpha\in\left[0.3,0.4,0.5,0.6,1.0\right]$.
Refer to caption: x12.png

##### Data Selection Ratio.

In Figure 10, we investigate how the data selection ratio affects the performance 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, a larger original corpus is needed for lower $\alpha$. Therefore, we keep $\alpha=0.4$ in our main experiments to balance effectiveness and data demand.

##### Gumbel Noise in Data Selection

In Figure 10, we explore the effect of the strength $\tau$ in Gumble-Top-$K$ used for data selection in Eq. (9). We can see that $\tau=0.1$ achieves the best performance, verifying the benefits of increasing the diversity of the pre-training corpus. Too large a $\tau$ value makes PDS degenerate to random selection, causing the performance to decrease.

**Table 8: Effect of using different downstream datasets to compute $J({\bm{\theta}})$. We report average accuracy on the OLMo evaluation datasets.**
|  | $J({\bm{\theta}})$ | Acc. |
| --- | --- | --- |
| Conventional | - | 43.2 |
| PDS | LAMB. | 43.7 |
| CC | 43.0 |  |
| OWT | 44.1 |  |
| LIMA | 45.0 |  |

##### Choice of J ⁢ ( 𝜽 ) 𝐽 𝜽 J({\bm{\theta}}) italic_J ( bold_italic_θ ) .

The desired data plays an important role in determining the quality of the selected data. We test different downstream datasets to compute $J({\bm{\theta}})$ in PDS and report the model performance in Table 8. The comparison between the results of using LAMBADA and LIMA shows the importance of the downstream data diversity. Instances in LAMBADA mostly come from stories, while LIMA is composed of instruction-response pairs that cover various tasks, which yields better overall performance. When comparing LIMA, OpenWebText [], and CC, we conclude that data quality is another major concern. Although OpenWebText has been shown to have better scaling factors [] and used as the target set [], replacing it with higher quality LIMA further improves performance. Compared to diversity and quality, large sizes of downstream datasets seem less important, because LIMA performs the best with the least instance number.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
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.

Does 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?

Line 883, what is B and what is beta? 

Typo:
Line 68: "(e.g., 125M paramaters)"
soundness: 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.

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

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

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

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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
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.
soundness: 3
strengths: - The paper is well-written, and the experimental setup is clearly explained.
- The PDS method is well-motivated by control theory and demonstrated to be effective in practice.
- The offline nature of the proposed method makes it very easy to use in existing language model training pipeline.  
- Multiple thoughtful ideas are proposed to accelerate the implementation of PDS.
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.
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.

Overall, I think the paper demonstrates the promise of the proposed method, and I am inclined to support its acceptance.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
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).
- 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$?
soundness: 3
strengths: - Extensive ablations (optimizing Algorithm 1, choice of downstream task) are performed to support the empirical claims and many of the choices made.
- The tackled problem is very relevant and the approach is innovative, using optimal control to account for the complete training trajectory.
- 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.
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.
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.
- 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?
- 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.
- 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.
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Two questions about the theory:
(a) Could you provide some intuition about why theoretically a single training sample is optimal, according to Eq. (6)? 
What's missing in this theoretical analysis?
(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?
I'm willing to raise my rating if these two questions can be resolved.


- Line 883, "B" and "\beta" seems like typos; should be "C" and "c" respectively?
soundness: 3
strengths: - The formulation of data selection as optimal control, and the theoretical connection between two seemingly unrelated fields, are interesting.

- The presentation is good overall. Ideas are clearly explained.
summary: This work proposes a novel method of data selection for training machine learning models, or LLMs in particular.
It starts by theoretically formulating data selection as an optimal control problem, 
where 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.
Based on this, a practical and efficient implementation is developed: 
first solve the optimal control problem on a small proxy dataset (approximately, via bi-level optimization),
and then train a data scorer (by fitting the obtained scores from the first step),
which can then be applied for training any LLM with other datasets.
Experiments show improvements over prior data selection algorithms,
in terms of test losses and scores on some commonly used benchmarks.
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, 
which leads to the surprising conclusion that training an LLM with a single sample is theoretically optimal.
The authors also mention this in Line 235. 
My major concern is that this seems to reveal a fundamental caveat of the proposed methodology.
Also, with many approximations in deriving the practical implementation (Section 2.3), 
it 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.


- Using the "scaling law" for extrapolating empirical results with small-scale models to 175B or 405B models (Table 3) doesn't feel like science. 
It is not obvious why scaling laws for conventionally trained LLMs should be applicable to PDS-trained LLMs.
In addition, the goodness of fit for calculating the scaling law's coefficients is not reported, which further makes the extrapolation less convincing.


- 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.
Moreover, many scores in these results are close to random choice (e.g., 25% accuracy for MMLU in Table 2), 
hence it is not clear what the minor differences between the scores of different methods really imply.

---

confidence: 2
contribution: 3
initial_rating: 10
presentation: 4
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. 

In 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? 

It 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?
soundness: 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.
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.
weaknesses: I didn't identify any major weaknesses of this paper.

Meta-review: The paper introduces a novel framework, PMP-based Data Selection (PDS), to enhance language model pre-training by leveraging optimal control principles. Reviewers enjoyed the paper’s theoretical foundation, particularly its use of Pontryagin's Maximum Principle (PMP), and its strong empirical results demonstrating improved data usage and scaling efficiency across various model sizes. The authors effectively addressed reviewer feedback during the rebuttal, clarifying technical details and strengthening their claims. Given its significance for large-scale data selection and LM training, I recommend acceptance.

Decision: Accept (Oral)