Gradient-based Optimization of Dataset Mixtures

First, the gradient approximation assumes a small, constant learning rate ($\eta$), yet the paper lacks evidence on its performance with varying strategies like adaptive or scheduled rates. This raises concerns about its use in dynamic environments. Second, while the softmax function normalizes mixture coefficients for interpretability, its impact on optimization and performance compared to unconstrained coefficients is not explored. An empirical comparison would clarify this. Additionally, the paper omits potential theoretical limitations of gradient-based mixture optimization, leaving unclear if specific data distributions or tasks could hinder its performance. Please see questions below.

1. The experiments lack comparisons to several baseline methods on data mixture optimization: (i) DoReMi [1] (ii) DOGE [2] (iii) The clustering-based method [3], where the data mixture coefficient $\alpha$ can be obtained by assigning all validation samples into training datasets. [2] should be considered closely relevant since it also applies gradient-based method in data mixture weights optimization with similar first-order approximation. Also, the uniform distribution baseline is not considered on the language task (dyck grammar). In the vision experiments, the baseline methods included in the paper need a more clear clarification (e.g. the Exact method in Table2 and Base in Table5).
2. The results from the language experiment (dyck grammar) is not clearly demonstrated. e.g. what is the test perplexity compared to uniform baseline (train with uniform $\alpha$ distribution)? and what does the "uniform" line stands for in Figure 4,5 (a,b) given there are only three training datasets considered?
3. The improvement on the vision tasks seem to be trivial. In Table 2, the proposed method is not able to outperform the Exact baseline in most of cases, and all listed results in Table 3.
4. The authors did not give detailed analysis on the computation overheads by applying the methods.

[1] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining [2] DoGE: Domain Reweighting with Generalization Estimation [3] Automatic Data Curation for Self-Supervised Learning: A Clustering-Based Approach

• Writing and Notation: The paper’s writing could be improved for clarity. For example, notation consistency should be addressed:
  • In Equation 5 and 7, the gradient term should use boldface as $\mathbf{G}$ instead of $G$ to maintain consistency with earlier notation.
  • Caption in Table 1, $\theta$-optimization should be boldface as well.
• Concept Explanation: Certain concepts could be explained more clearly to improve readability:
  • Lines 120-121 mention that $\mathbf{\theta}_T$​ implicitly depends on $\mathbf{\alpha}$, but it lacks an explanation. It would be helpful to elaborate on this a little bit.
  • Line 194 references a change of variables that could benefit from additional explanation, perhaps in the appendix, as it may not be immediately clear to readers.
• Limited Checkpointing Experiment: The checkpointing experiment only includes 5 and 10 checkpoints. Expanding this to additional checkpoints would better illustrate the effect of this extension.
• Efficiency Comparison: While the method can be viewed as a more efficient fine-tuning approach where dataset weights are optimized, no direct comparison with conventional fine-tuning in terms of runtime or computational cost is provided.
• Lack of Real-World NLP Task: Although the motivation of the paper stems from natural language processing (NLP) applications, only a toy example (Dyck Grammar) is presented. Including a real-world NLP task would better demonstrate the method’s effectiveness and relevance for NLP applications.

Typos:

• line 307, paris -> pairs
• line 402, Qunatized -> Quantized

The idea of this paper is simple and obvious. The working of the proposed approach depends on how similar the testing distribution is to the validation distribution and the size of the validation set. Clearly, if the testing distribution is far from the data distribution from which the validation set is drawn, then optimization of the validation loss (on the validation set) would not lead to good testing performance. On the other hand, if the validation distribution is identical to the testing distribution, but the validation set is too small, this approach will overfit to the validation set, and also fail to generalize to the testing set. To this reviewer, theoretically understanding the interplay between distribution shifts (from training, to validation, to testing) and the sizes of the training set and testing set is an important aspect of this research, but is ignored by the authors.

The proposed approach makes use the validation set. Then a natural question to be answered is this: among all approaches having access to both the training set and the validation set, is the proposed approach the best?

Overall, I feel the research results presented so far are shallow and lack nutrient.

• My main concern is that the experimental evaluation is mostly made of toy/invented problems. Given the generality of the broader problem, I would have expected to see at least one real instance where the approach shines. The only real case in the paper is MS-COCO, but I am having a hard time understanding why it is a study case in the first case. What are the different data sources in MS-COCO? The introduction states “The Pile” as a case, being made of data from ArXiv, GitHub, etc. Can an experiment be made on this one instead?
• Lack of real baselines: the paper lists 6 previous works that investigated the optimization process of data mixture coefficients, yet does not compare with any of these. Having virtually no baseline tackling the same problem, it is hard to assess the virtues of the approach.
• The constant $\mathbf{G}$ approximation is not empirically validated. It would be helpful to see some quantitative measure assessing how much it actually holds, and whether it correlates with the benefits of the approach.