Data Selection via Optimal Control for Language Models

We are deeply grateful for the reviewer's highly positive feedback and recognizing the novelty and significance of our contributions. The reviewer's acknowledgment of our theoretical formulation, optimization solution, efficient implementation, and robust experimental validation is really encouraging. We are especially honored by the reviewer's clear and strong support!

The use of our formulation in other scenarios (question #1)

To our best knowledge, there is no existing work applying this formulation in data selection. We believe that it is quite promising to apply this formulation to multi-modal foundation model training, such as vision language models, because we do not make any assumptions of the data forms or its content. In our on-goning project, we apply PDS to select image-text pairs for SFT of vision language models and observed promising results.

The necessary condition (question #2)

Yes, the PMP conditions do not guarantee that the data selection is optimal, since they are necessary conditions. However, according to many empirical results in the field of Optimal Control [1,2], solving the PMP conditions is sufficient to get fairly good solutions. (Much like using Gradient Descent in deep learning even if Gradient Decent also cannot guarantee the global optimality.)

[1] Optimal Control with Engineering Applications. 2007. Springer.

[2] Optimal Control Theory with Economic Applications. 1987. Amsterdam: North-Holland.

The potential negative results (question #3)

A potential scenario that PDS may fail is that the original corpus contains too much duplication. Since PDS selects data individually, duplicated data points may get selected in this setting. Therefore, PDS is better suited for further improving the qualities of data after de-duplication. Since data-deduplication is well-developed in recent years [1,2,3], our method focuses more on the data cleaning stage after deduplication. We have also added a discussion about this scenario in Appendix E and also provide directions for future exploration.

[1] SemDeDup: Data-efficient learning at web-scale through semantic deduplication. 2023. In ICLR Workshop.

[2] D4: Improving LLM Pretraining via Document De-Duplication and Diversification. 2023. In NeurIPS.

[3] Deduplicating Training Data Makes Language Models Better. 2022. In ACL.

The relevance between the resulting algorithm and the theory (weakness #1)

Our algorithm is directly derived from Theorem 2.1.

For Algorithm 1, we use Eq. (4) and Eq. (5) to compute $\theta$ and $\lambda$ in the forward and reverse inner loops. Then, we follow Eq. (6) to obtain $\gamma$ in each outer loop by updating $\gamma$ with the $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values. This update strategy still aligns with Theorem 2.1 because the sample set with the largest $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values still gets the highest $\gamma$. Additionally, this update ensures the samples with the same $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values are assigned the same $\gamma$, avoiding the “single training sample” solution.

The effectiveness of our proposed efficient implementation is empirically verified. In Figure 6, we showed that the efficient implementation preserves most of the performance of the exact algorithm in a simulated setting. In Figure 7, we showed that the data scorer achieves around 0.55 Spearman correlation in fitting the mapping between each data point and its $\gamma$. This is a reasonably high correlation and when we select 40% data using the trained data scorer, we find that 81.6% of the data points are correctly marked as selected/discarded in the validation set.

The scaling law for extrapolating empirical results (weakness #2)

The scaling law still applies to PDS-trained models because only the pre-training corpus is replaced with our selected data, with all other model and optimization configurations unchanged compared to conventional pre-training. Previous works have shown that the power law applies to different data sources [1,2,3,4] when all other configurations remain unchanged. Furthermore, the results in General Response show that the power law fits the observed performance data very well for both conventional pre-training and pre-training with PDS-selected data.

Based on the goodness of fit of our scaling curve and the previous literature, we firmly believe that the fitted scaling curve reveals the performance trend of the PDS method in comparison to the conventional pre-training as model and data sizes scale up.

[1] DeepSeek LLM Scaling Open-Source Language Models with Longtermism. 2024. arxiv pre-print.

[2] Resolving Discrepancies in Compute-Optimal Scaling of Language Models. 2024. In ICML Workshop.

[3] Scaling Laws for Transfer. 2021. arxiv pre-print.

[4] MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. 2024. In COLM.

We thank the reviewer for recognizing contribution of our theoretical formulation and the clarity of our presentation. We are also encouraged to see the reviewer’s willingness to increase the score based on our explanations of our theoretical analysis. The following is the clarification of our theoretical and empirical results:

The intuition behind Eq. (6) (weakness #1 and question #1 (a))

Eq. (6) does not imply that a single training sample is optimal. The reasons are explained below:

• Eq. (6) implies that there could exist several samples having the same largest $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta^*)$ values when the model’s parameter $\theta^*$ is trained under the optimal $\gamma$ scores, i.e. $\exists S=\{n_1,n_2,\cdots,n_{|S|}\}$ such that for any $m \notin S$.

  $$\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_{n_1},\theta^*) = \cdots = \sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_{n_{|S|}},\theta^*)>\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_{m},\theta^*)$$

  Therefore, multiple samples in $S$ can have non-zero $\gamma^*$ scores to satisfy Eq. (6), as long as those samples not in $S$ get zero $\gamma^*$ scores. Accordingly, although "a single training sample" is a solution to Eq. (6), there also exist many more "equally qualified" solutions (such as assigning uniform $\gamma$ values to the data points in $S$) that satisfy Eq. (6).

• As stated in lines 56-60 and 158-169 in the latest revision, the PMP condition for data selection (Theorem 2.1) is a necessary condition of the $\gamma$‘s optimality. This is because the two techniques we used to prove Theorem 2.1 in Appendix B, i.e., (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 of the problem.

Many empirical results, including that in our paper and previous literature in optimal control [1,2], have shown that solving this necessary condition can 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.

The description in Line 235 does not mean that “training an LLM with a single sample is theoretically optimal”. Instead, it means that to solve the PMP conditions, a straightforward implementation can use the max operation to get $\gamma$ during each outer loop iteration. However, since $\theta$ has not reached its optimum yet, the sample set $S$ having the same largest $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values has not emerged, making the max operation unstable in practice. Therefore, we increase $\gamma$ by a value proportional to $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ in each iteration of outer loop as in Algorithm 1. This update strategy still satisfies our theory because it guarantees that the sample set having the largest $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values gets the highest $\gamma$ values. It also ensures that the samples with the same $\sum_{t=0}^T\lambda_{t+1}^\top\nabla l(x_n,\theta)$ values receive the same $\gamma$ scores, which avoids the “single training sample” solution.

[1] Optimal Control with Engineering Applications. 2007. Springer.

[2] Optimal Control Theory with Economic Applications. 1987. Amsterdam: North-Holland.

The Lagrange Multiplier method (question #1 (b))

We omit the $\max_{\mu}$ term in Eq. (22) because this term is equivalent to the time-invariant constraint in Eq. (20). Specifically, by treating $\mu_t$ also as the time-variant control variables and applying the standard PMP conditions (Theorem B.1), the $\max_{\mu}$ term yeilds

$$\mu_t=\arg\max_{\mu_t}\sum_{n=1}^{|D|}\mu_{n,t}(\gamma_{n,t}-\gamma_{n,0}).$$

Since $\mu_t$ does not have any constraint, it exists and is finity ($<+\infty$) if and only if $\gamma_{n,t}=\gamma_{n,0}$ for all $n=1,2,\cdots,|D|$. The same analysis applies to all $1 \le t \le T$. In this way, we have $\mu_t=\arg\max_{\mu_t} 0$, which does not introduce additional constrains to the problem. This means the remaining analysis in our paper still holds.

We highly appreciate the reviewer’s insightful feedback, which clearly helped us improve our paper.

As the reviewer mentioned being willing to raise their rating if the two questions in Question #1 were resolved, we have provided detailed explanations on our theoretical results to address these concerns. Additionally, we have included thorough evidence to support the scaling curve extrapolation, demonstrate our method’s effectiveness on larger models, and highlight the improvements on MMLU.

We sincerely hope the reviewer find our response useful and update the scores if the concerns have been resolved. We are also open to further discussion if there are further questions.

Thank you for your detailed responses.

It makes sense that multiple training samples can have the same highest scores and thus non-zero weights (I played around with a few toy examples analytically and saw that this indeed happened), although the concrete values of the optimal weights cannot be found based on the proposed theory that offers necessary conditions.

I think my major concerns have been resolved, so I've raised my rating as promised. It would be useful to add the clarifications in your responses for this "one-hot vector" issue, as well as the goodness of fit for scaling laws and the additional empirical results, to the revision of your manuscript.

We sincerely thank the reviewer for highlighting the strengths of our work, including the extensive ablations to support our empirical claims, the innovation of our approach, and the promising empirical results. We also appreciate the reviewer's recognition of the computed scaling law as an effort to generalize to larger models.

Regarding the weakness:

The results of the $J(\theta)$ choice :

Thanks for your constructive suggestion. Considering the strict page limit of ICLR, we chose to put the results in the Appendix. A more important reason is that the main contribution of our paper is the proposed theoretical framework and the PDS method, rather than the choice of $J(\theta)$. We would like to elaborate more on the unique key factors of PDS compared to previous methods, including its general empirical results (Table 1-3, Figure 4-5), its relationship to the theoretically optimal results (Figure 6), its computation complexity (Table 4), the effect of different incorporated learning information (Table 5), and the choice of the proxy setting (Figure 7). Furthermore, since the choice of $J(\theta)$ has been discussed in some previous studies [1,2], we do not emphasize it as one of our key findings and choose to put it in the Appendix to keep the full results within the page limit. In Section 3.1 (lines 305-306), we add a cross-reference to the results for the convenience of the readers.

[1] Less: selecting influential data for targeted instruction tuning. 2024. In ICML.

[2] DsDm: Model-Aware Dataset Selection with Datamodels. 2024. In ICML.

Explanation to Figure 1:

• We used the same amount of data for our method and Redpajama to train the models.
• All the points correspond to the training of a 1.7B model in (a).

The way $J(\theta)$ is computed:

In the revision of our paper, we have clarified that $J$ is not computed on the same downstream tasks as those used in evaluation. (lines 80-82)

The independence assumption of each example:

• We have clarified the assumption that we focus on the independent importance of each example in the revision (lines 103-104) and add the following discussion to Appendix E.

• Considering the dependence between each data point is a promising direction for future extension of our method. Some straightforward implementation could be improving the pair-wise difference between the examples to promote the diversity of the selected data. Specifically, we consider the pairwise similarity:

  $$\operatorname{sim}(x_n,x_m)=\frac{h_n^\top h_m}{||h_n||||h_m||},$$

  where $h_n$ and $h_m$ are the representations of each example generated by a model like RoBERTa [1]. Then, we have the following loss to minimize, which encourages the difference between examples:

  $$\mathcal{L}^{\text{diversity}}=\sum_{n,m}\text{sim}(x_n,x_m)\gamma_n \gamma_m = \pmb{\gamma}^\top\pmb{S} \pmb{\gamma},$$

  where $\pmb{S} = \left[\operatorname{sim}(x_n,x_m)\right]_{1\le n \le |\mathcal{D}|, 1\le m \le |\mathcal{D}|}$. We can add this loss to the optimization problem in Eq. (3), using a hyper-parameter $w$ to control its weight:

  $$\min_{\pmb{\gamma}} \sum_{t=1}^{T} J(\theta_t) + w\cdot\pmb{\gamma}^\top\pmb{S}\pmb{\gamma}.$$

  This problem can still be solved by the PMP conditions, and we leave the design for efficient implementation to future work.

  We have included the above discussion to the Appendix E in the revision of the paper.

  [1] RoBERTa: A Robustly Optimized BERT Pre-training Approach. 2019. arxiv pre-print.

The use of Gumbel Sampling

Basically, $\gamma$ is computed based on the distance between the sample gradient and the target vector $\lambda$. Therefore, if the original corpus contains too many duplicated samples, their distances to $\lambda$ will also be similar, which means that if one example gets a high $\gamma$, the others will also get high $\gamma$ and being selected. This would affect the diversity of the selected corpus. Therefore, we adopt Gumbel sampling to introduce some variance to this selection, which means even if a set of examples has similar high scores, it is still possible that not all of them will be selected. Similar techniques are also adopted in previous works [1,2].

[1] MATES : Model-Aware Data Selection for Efficient Pre-training with Data Influence Models. 2024. In NeurIPS. (Section 3.1)

[2] Language Models are Few-Shot Learners. 2020. In NeurIPS. (Appendix A)

We sincerely thank the reviewer for recognizing the good motivation and practical value of our PDS method, the clarity of our writing, and the thoughtful ideas to accelerate its implementation. We deeply appreciate the reviewer's understanding of the computational resource constraints for larger-scale experiments and the reviewer's inclination to support the acceptance of our paper.

The large-scale experiments (the weakness)

As the reviewer has mentioned, we did not conduct larger-scale experiments because of the prohibitively high computational cost since we have to train the LMs from scratch each time. Therefore, we fit the scaling curve and extrapolate the results to show the performance trend of our method compared to the baselines, as training scales up. The results in the General Response show that the scaling law curves fit our observed data well.

Previous work[1] has shown that the scaling curves fitted by performance data obtained based on 40M to 2B models can provide useful guidance for large-scale training. Another work[2] uses performance data obtained under a $3\times 10^{20}$ FLOPs training budget to fit a scaling curve and accurately predicts the performance of 7B and 67B models, while our performance data are obtained under an even larger $5\times 10^{20}$ FLOPs training budget.

Based on the fitting goodness of our curve and the previous literature, we firmly believe that the fitted scaling curve reveals the performance trend of the PDS method compared to conventional pre-training when model and data sizes scale up.

[1] MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. 2024. In COLM.

[2] DeepSeek LLM Scaling Open-Source Language Models with Longtermism. 2024. arxiv pre-print.

Connection between Theorem 2.1 and KKT conditions (the question)

The connection between Theorem 2.1 and KKT conditions is that, they are both built on the method of Lagrange multipliers. $\lambda$ is essentially the Lagrange multiplier in the conventional PMP conditions; and as the proof in Appendix B shows, we use Lagrange multiplier method to adapt the conventional PMP conditions to our data selection setting. One notable difference is that, Theorem 2.1 is a stronger result than the KKT conditions as it is applicable for the cases where the Hamiltonian function $H(\theta, \lambda, \gamma)$, as defined in Eq. (17), is not differentiable with respect to $\gamma$ [1].

[1] Lawrence Craig Evans. An introduction to mathematical optimal control theory. University of California, 2005

We sincerely thank the reviewer for the thoughtful comments and the inclination to support the acceptance of our paper.

Regarding the mentioned weakness, we have provided the goodness of fit for the scaling curves in the General Response. The high correlation ($R^2$) values for both model and data fits demonstrate strong alignment, which strengthens our claims on fitting and extrapolating the scaling curves. Additionally, we reference prior literature showing that scaling curves derived from experiments of a similar scale to ours can accurately predict the performance of larger models (e.g., 7B and 67B).

In response to the question, we have clarified the connection between Theorem 2.1 and the KKT conditions. Both are derived using Lagrange multipliers, but Theorem 2.1 provides a stronger result for non-differentiable dynamic systems compared to the KKT conditions.

As we are approaching the end of the discussion stage, we would appreciate it if you could read our responses and update the scores if your concerns have been addressed. We are more than happy to further discuss any concerns that you find not fully addressed. Thank you.

We thank the reviewer for the thoughtful comments. We are encouraged to find that the reviewer appreciates the novelty of our method and the clear presentation of our paper!

The scaling laws (weakness #1)

As described in Section H.3 in the latest revision, the model of the scaling law consists of 5 parameters, while we have 80 observed data points to fit the scaling law curves. As we explain in the following, this is considered statistically sufficient. And moreover, we did not use more data mainly because of the substantial computational resources required for larger-scale experiments.

From the results in the General Response, we can clearly see that the scaling law curve fits our observed data well, in both the model and data size dimensions.

We also want to point out that there are quite a few existing works fitting the scaling curves and predicting performance for large models with few much less data points. For example,

• [1] fits the model performance curve with 5 points and shows good prediction in its Figure 2.
• [2] uses 36 data points obtained from models ranging from 40M to 2B to fit a scaling curve, which provides useful guidance for large-scale training.
• [3] selects 8 data points obtained under the training budget less than $3\times10^{20}$ FLOPs to fit a scaling curve and accurately predicts the performance of 7B and 67B models, while the maximum training computation in our experiments is larger ($5\times 10^{20}$ FLOPs), as shown in our Figure 1.

Based on the goodness of fit results obtained by our experiments and practices in previous literature, we firmly believe that the fitted scaling curve reveals the trend of our PDS method's performance compared to the conventional pre-training when model and data sizes scale up.

[1] GPT-4 Technical Report. 2023. arxiv pre-print.

[2] MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. 2024. In COLM.

[3] DeepSeek LLM Scaling Open-Source Language Models with Longtermism. 2024. arxiv pre-print.

The empirics (weakness #2)

• We achieve consistent improvement on nearly all 9 benchmarks compared to the baselines. To make the results' significance clearer: conventional pre-training has to spend 2x more computation than PDS to bridge this gap, which can be translated as saving the training time from 7 days to 3.5 days to train a 1.7B model on 8 A100 GPUs. Therefore, our proposed method can strongly benefit the current pre-training practices.
• Regarding the question of “hyper-parameter sweeps on the downstream optimizer”: The results reported in our paper are all zero-shot or few-shot performance of a base pre-trained model, without fine-tuning on downstream tasks. Therefore, the required results about careful hyper-parameter sweeps on the downstream optimizer are not suited for our paper. For the hyper-parameters of the model architectures, optimizers, and pre-training recipes, since we only changed the pre-training data sources of the models, we simply kept these hyper-parameters the same for a fair comparison.

The gain of a single inner step (question #1)

This is because the “target vector” $\lambda_t$ comprises two components of model information: $\nabla J(\theta_t)$ and $\eta \nabla^2L(\theta_t,\gamma)\lambda_{t+1}$, as described in Eq. (5). $\nabla J(\theta_t)$ is the zero-order information and $\eta \nabla^2L(\theta_t,\gamma)\lambda_{t+1}$ is a first-order correction when the learning rate $\eta$ is small, especially for large models. For the “single inner step” ($T=1$) case, we have $\lambda_1=\lambda_T=0$, and thus only $\lambda_0=\lambda_1+\nabla J(\theta_0)-\eta \nabla^2L(\theta_0,\gamma)\lambda_{1}=\nabla J(\theta_0)$ is computed as the “target vector” in Algorithm 1. As a result, “single inner step” can be regarded as a zero-order approximation of our solution, ignoring the first-order correction.

The proxy and data scorer model (question #2)

It is better to scale up the proxy and data scorer model because a large model may need more complex signals for data selection. However, considering the trade-off between computation and final performance, we showed that using a ~100M model to select data for a 10x larger model leads to fairly good results.

Typos

Thanks for pointing this out, and it is indeed a typo in line 883. $B$ and $\beta$ should be $C$ and $c$. We have fixed this together with the typo in line 68 in the latest revision.

Thanks for your appreciation and interest our work!

Algorithm 1

As shown in Eq. (3), $\theta_{t+1}=\theta_t-\eta \nabla L(\theta_t,\gamma)$ serves as the constraint of the trajectory of $\theta_t$, ensuring that the model is still optimized by GD when the optimal $\gamma$ is applied. Executing the outer loop once is only an implementation choice for efficiency and in practice, users can run the Algorithm 1 with multiple outer loop iterations, given sufficient computation budget. In Figure 6, we compared the $J(\theta)$ values of this one iteration implementation and the multi-iteration implementation. We can see that the efficient implementation preserves most of the performance of the exact implementation.

Results in Table 1

The performance of the baseline methods depends on the downstream datasets. In Table 2 and Table 7 in the revision, the baseline methods improve the conventional pre-training on MMLU in most cases. The baseline methods did not consistently outperform conventional pre-training for the following reasons:

• RHO-Loss [1] selects instances based on the loss difference of a fully pre-trained model and a model fine-tuned on the downstream task LIMA, which is originally evaluated in the image classification task on relatively small models. It uses the final model parameters to compute the loss differences, which may be helpful for the final model training stages. However, it ignores the effect of data points on the early and middle stages of LM training.
• DSIR [2] selects instances based on the n-gram overlap between the downstream data (LIMA) and the pre-training documents. This shallow signal works for models with relatively small sizes or continually pre-training models, as verified empirically in [2] and our results for the 160M model. However, for training larger LMs from scratch with more complex training dynamics, the shallow signal does not seem to be much helpful. Instead, it makes the LM to over-fit the common phrases in the downstream data, which can hurt the performance.
• IF-Score [3] is first introduced to evaluate the influence of specific data points. Similar to RHO-Loss, it uses the final model parameters to compute the influence scores, ignoring the early training stages. Furthermore, its assumption $\nabla L(\theta_T)=0$ is hard to satisfy because the gradient norm does not seem to reach 0 when pre-training LMs [4].

[1] Prioritized Training on Points that are Learnable, Worth Learning, and not yet Learnt. 2022. In ICML.

[2] Data Selection for Language Models via Importance Resampling. 2023. In NeurIPS.

[3] Understanding black-box predictions via influence functions. 2017. In ICML.

[4] https://wandb.ai/ai2-llm/OLMo-1B/reports/OLMo-1B--Vmlldzo2NzY1Njk1#optimizer-metrics