Kati8/ICLR_2023_2025 huggingface dataset:
authors: Yuxian Gu
authors: Hongning Wang
authors: Yaru Hao
authors: Qingxiu Dong
authors: Minlie Huang
reviews: 
Comments: 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.
Strengths: This work is sufficiently solid. It covers the fundamental problem formulation, which introduces a new optimization problem. Also, it provides the theoretical optimization solution to the proposed problem, which has not been studied before. From this perspective, this paper has made sufficiently novel contributions.  Moreover, this work designs an efficient implementation to avoid additional computation overhead. These methods are empirically verified with different experiments setting, demonstrating the effectiveness of this method across different settings. Due to its theoretical contributions and solid experiments, I give a clear accept (score 10) to this paper.
Weaknesses: I didn't identify any major weaknesses of this paper.
Recommendation: 10
Reviewer confidence: 2
Soundness/correctness: 3
Clarity: 4
reviews: 
Comments: 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.
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.
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.
Recommendation: 6
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: 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.
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.
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.
Recommendation: 8
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: 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.
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.
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.
Recommendation: 8
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 4
reviews: 
Comments: This work reframes the data selection task as an AUC maximization problem. Notably, the paper establishes the closed-form optimal solution for this data selection formulation and develops a computationally efficient algorithm for implementation. The proposed approach is validated across multiple experimental scenarios.
Strengths: The paper demonstrates strong technical quality with a well-articulated problem formulation that introduces a novel optimization objective. A key contribution is the derivation of the theoretically optimal solution to the newly formulated problem, representing previously unexplored territory in this domain. Additionally, the authors present a practical algorithm design that minimizes computational burden. The comprehensive experimental evaluation across diverse settings effectively demonstrates the robustness and applicability of the proposed method. The combination of strong theoretical insights and thorough empirical validation warrants a clear acceptance (score 10).
Recommendation: 10
Reviewer confidence: 2
Soundness/correctness: 3
Clarity: 4