Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining

Thank you for your thorough reviews and constructive comments. We provide our response to your comments below. If our response resolves your concerns, we would greatly appreciate it if you could consider increasing your score.

Q1: The statement “where $L>0$ is the smoothness parameter” under Equation (7) is unclear since $L>0$ does not show up in Equation (7) and could benefit from further explanation.

A1: Thanks for noticing this! This was a typo and we have removed it in our revision.

Q2: Could you elaborate on the selection of LogiQA, LogiQA 2, SciQA, and PiQA as benchmarks? Testing additional benchmarks could further strengthen the few-shot performance evaluation.

A2: We thank the reviewer for the valuable feedback! For the GPT-2 models, we find that only question-answering tasks work well at this scale, which is why we report 5-shot results on LogiQA, LogiQA-2, SciQ, and PiQA. All other benchmarks we explored yielded random results for all compared algorithms, i.e., they were not useful for model evaluation. The reason is the limited capacity of the GPT-2 models.

To address the reviewer’s concern, we have conducted new experiments to train additional 1.4B (billion) and 7B model with Llama architecture on randomly sampled subsets of the FineWeb 15T dataset. For the 1.4B and 7B models, we use 100B and 175B randomly sampled tokens from the FineWeb 15T dataset, respectively. We train the models using the Uniform baseline with uniformly averaged losses and our LinUpper strategy.

For both models, we include benchmark results from 13 language understanding & reasoning (LUR) tasks and 6 question-answering (QA) tasks. Please see the new section 6.4 in our main paper and Appendix A for more details, including hyperparameters and model architectural details. We report the benchmark results on the last checkpoints.

After fully training the 1.4B model on 100B tokens, we observe a performance gain on 6 out of 13 LUR tasks with an average performance gain of 0.66%. For QA benchmarks, we improved on 5 out of 6 tasks with an average of 1.72%. For the 1.4B model, we provide a full overview below and in our updated paper in Appendix A, Tables 8 and 9. The first table shows the language understanding and reasoning benchmarks; the second question-answering tasks.

Update: We have completed the training of the 7B model and now report the final benchmark accuracies after 175B tokens. The training took a total of 9800 GPU hours. Overall, our LinUpper strategy performs better on 9 out of 13 LUR tasks, with an average improvement across all benchmarks of 1.45%. Also, we observe an average performance gain on QA benchmarks of 1.16%. We improve on 5/6 QA tasks. We have updated our paper with these final results. The full overview is available in Appendix A, Tables 10 and 11. Below, the first table shows the language understanding and reasoning benchmarks; the second question-answering tasks.

...continue A2.

These improvements are particularly noteworthy given the scale of these LLama models, highlighting the scalability and effectiveness of our method.

Q3: Including DoGE and DoReMi as baselines in Table 1 would provide useful context on the effectiveness of the proposed approach. Would you consider adding this comparison?

A3: Thanks for the question! Note that the methods DoGE and DoReMi are orthogonal to the uniform domain sampling approaches in Table 1, i.e., all methods in this table can also be run with DoGE or DoReMi. This is why we organized our experiments into two different subsections, one that compares algorithms under uniform domain sampling setting, and a second subsection that deals with the non-uniform domain sampling setup (in which we use DoGE and DoReMi as the optimal domain sampling estimators).

Q4: The theoretical assumption cap on maximum weight at $\frac{2}{M}$ seems restrictive. Could you provide further justification or discuss how this assumption holds in practical applications?

A4: Thanks for the comment! Note that the inequality $w_i \leq \frac{2}{M}$ is not an assumption but rather a finding of our theoretical analysis (Proposition 1) that supports the benefits of some reweighting strategies such as LinUpper. The only assumptions required for Theorem 1 to hold are the widely adopted Assumptions 1, 2, and 3.

Q5: A deeper comparison of computational and implementation efficiency relative to DoGE and DoReMi could underscore the practical advantages of your method. Could you expand on this?

A5: Thank you for the insightful question! Our reweighting approach is designed with efficiency at its core, minimizing computational overhead associated with the sample weights computation by innovatively leveraging information from sample loss values. In particular, compared to DoGE and DoReMi, which train additional auxiliary models (DoReMi trains 2 additional separate models and DoGE trains 1 additional proxy model), our approach does not require training any additional auxiliary model for computing the instance weights. Hence, in terms of resource consumption, our method requires about 3 times less compute than DoReMi and 2 times less compute than DoGE. Furthermore, in terms of implementation our method enjoys huge simplicity compared to DoGE and DoReMi, which implement additional bilevel/minimax based optimization loops for training the auxiliary models. Instead, our Algorithm 1 can be seamlessly integrated into existing training pipelines for LLMs. To demonstrate how our method can be easily integrated into a traditional multi-GPU training script, we have added a new section in the appendix (Appendix A) that provides a practical implementation in Pytorch. We hope that this further clarifies some details about our method.

Dear Reviewer 9BZv,

Many thanks for your reply. We further clarify the meaning of our theoretical upper bound on the weights as follows.

First, we note that in the case of SGD with a minibatch size of b, the uniformly weighted SGD (i.e., without reweighting) has a weight of 1/b for each sample in the minibatch, since the weights of all the samples in the minibatch sum to one. In the case of reweighting, our theoretical upper bound is that any weight is upper bounded by 2/b. Therefore, this weight upper bound simply means that the weight of any sample (after reweighting) should not be more than twice the uniform weight. Intuitively, the reason of having such an upper bound is that, if the weights of some samples are very large, some other samples would have very small weights, causing an extremely skewed data distribution that is not helpful for training. Instead, if the samples are weighted properly within a reasonable range, emphasizing more useful samples but not causing extreme skewness, it is helpful for training. We also note that Theorem 1 in the paper is a special case of Theorem 2 with b=M, which corresponds to the case of full gradient descent and we use it as a theoretical foundation to motivate our reweighting approach. Theorem 2 is more general and aligns with the practical setting when SGD is used.

We would also like to explain that our LinUpper reweighting method indeed satisfies the upper bound of $w_i < 2/b$ and any strategy that assigns the same weight to half of the data points and smaller weights to the other half (such as LinUpper) will satisfy the upper bound of $w_i < \frac{2}{b}$. The proof is straightforward when noting that the sum of all weights needs to be 1.

To demonstrate that in practice the upper bound of $w_i < 2/b$ is in fact satisfied by our LinUpper method, we provide the distribution of maximum weights (please see Figure 6 in Appendix A.5) during the pretraining of the 7B parameter model with our LinUpper method, which is the new experiment that we added in the revision. In this experiment, the minibatch size is $128$ and hence the upper bound is $\frac{2}{128} = 0.0156$. As you can see, the maximum weight (which is around $0.01$) by our LinUpper method is clearly below the upper bound of $0.0156$. The fact that we obtain better performance when training using our LinUpper method compared to the uniform weighting baseline confirms the usefulness of reweighting within this upper bound.

We also note that we have updated the results of the 7B parameter model training (in both the paper and the previous comments) that we obtained after a full run with 175B tokens of training data.

We hope that this addresses your concern. Please let us know if you have further questions.

Regards,

Authors of Paper 13352

Note: We wanted to let the reviewer know that we have conducted extensive additional experiments in which we trained larger 1.4B (billion) and 7B models. For both models, we include benchmark results from 13 language understanding & reasoning (LUR) tasks and 6 question-answering (QA) tasks.

After fully training the 1.4B model on 100B tokens, we observe a performance gain on 11 out of 19 benchmarks with an average performance gain of 0.99%. For the larger 7B parameter model (trained on 175B tokens), we improved performance on 14 out of 19 tasks with an average gain of 1.35%.

Please see the new section 6.4 in our main paper and Appendix A for more details.

These improvements are particularly noteworthy given the scale of these LLama models, highlighting the scalability and effectiveness of our method.

Thank you for your thorough reviews and constructive comments. We provide our response to your comments below. If our response resolves your concerns, we would greatly appreciate it if you could consider increasing your score.

Q1: Could you provide more detailed implementation steps for the dynamic reweighting algorithms? Specifically, information on how to normalize the losses, apply the reweighting strategies, and integrate these into existing training pipelines would be beneficial for reproducibility.

A1: Thanks for the question! We provide a detailed explanation of the steps in our dynamic reweighting method (Algorithm 1 in our paper). Additionally, to facilitate reproducibility of our method we have included a practical implementation in Pytorch that demonstrates how our method can be seamlessly integrated into standard multi-GPU training pipelines (please see Appendix A).

Step 1: Loss Normalization [line 5 in Algorithm 1]

The normalization step ensures that the losses are scaled consistently for reweighting. Given a batch of losses $\{f_{i,t}\}, i=1,...,b$, where $b$ is the batchsize, compute:

$h_{i,t} = \frac{2(f_{i,t} - f_{\text{min}})}{f_{\text{max}} - f_{\text{min}}} - 1$.

Here, $f_\text{min}$ and $f_\text{max}$​ are the minimum and maximum loss values in the current batch. This linear transformation maps the losses to the interval [−1,1], ensuring consistent scaling across different training batches. This allows the reweighting functions to operate on a consistent and bounded range of values. The normalization step outputs the normalized losses $h_{i,t}, i=1,...,b$ to be used in the next step.

Step 2: Applying Reweighting Strategies [line 6 in Algorithm 1]

Once the losses are normalized, the chosen reweighting strategy is applied:

LinUpper: This strategy downweights the low-loss samples and caps the weight applied to high-loss samples, ensuring that outliers do not dominate the training process.

Quadratic: This strategy emphasizes the medium-loss samples by penalizing both low- and high-loss values.

Extremes: This strategy assigns higher importance to samples with extreme loss values, whether high or low.

This step outputs the strategic weights $s_{i,t}, i=1,...,b$ to be adjusted by the next curriculum step.

Step 3: Curriculum Adjustment [line 7 in Algorithm 1]

One we obtain the strategic weights $s_{i,t}, i=1,...,b$, apply the softmax-based adjustment controlled by the temperature parameter $r$​:

$w_{i,t} = \frac{\exp(s_{i,t} / r_t)}{\sum_{j=1}^b \exp(s_{j,t} / r_t)}$

Start with a higher $r_t$​ (e.g., $r_0$​), ensuring more uniform weights early in training. Reduce $r$ as training progresses to allow for sharper differentiation in sample importance.

Step 4: Gradient-Based Parameter Update [line 8 in Algorithm 1]

After computing the sample weights $\{w_{i,t}\}, i=1,...,b$​, update the model parameters using the weighted gradient:

$\theta^{t+1} = \theta^t - \eta \sum_{i=1}^b w_{i,t} \nabla f_{i,t}$

This ensures that each sample contributes to the parameter update proportional to its dynamically associated weight.

Please see Appendix A for the Python code for each of these steps and a practical integration into existing training pipelines.

Q2: The curriculum adjustment mechanism using a temperature parameter is an interesting aspect of your approach. Could you provide more details or results on how different values of the temperature parameter affect the performance and convergence of the model? This would help in understanding the sensitivity and effectiveness of this component.

A2: Great question! We have conducted additional experiments using different values of $r$ to understand the sensitivity of our methods to the value of $r$. The perplexity plots and benchmark results can be found in Appendix A.4. Figure 5 in Appendix A.4 shows that when the value of r is large (e.g. r=1), the performance of our method becomes closer to that of the uniform baseline. Decreasing the value of r leads to diminished effect of low-loss samples, but this can also have a negative effect on the performance when r is too small (eg. r=0.2), potentially due to data wastage by overfiltering the low-loss samples.

Q3: Your theoretical analysis mentions that the reweighting schemes avoid overfitting to outliers. Could you provide empirical evidence or additional discussion on how well your methods handle outliers in practice, especially compared to other robust optimization techniques?

A3: Thank you for your insightful question! Unlike traditional robust optimization methods, such as those based on distributionally robust optimization (DRO), our methods impose a cap on sample weights ($w_i \leq 2/M$), as justified by our theoretical analysis (Proposition 1). This cap prevents over-reliance on a small subset of high-loss samples, which can lead to overfitting to outliers. In contrast, DRO methods often focus heavily on worst-case scenarios, potentially leading to suboptimal generalization when the training data contains a substantial amount of noisy samples. To demonstrate this benefit of our method in practice, we conducted new experiments to compare with the traditional KL-divergence regularized distributionally robust optimization (DRO-KL). First, we created a synthetic regression problem (similar to the regression problem in Appendix B) with 25% of the data being outliers to compare the robustness of the different methods against outliers.

The experiment's plots can be found in Figure 7 of Appendix B.3. We note the following: The DRO-KL approach because it focuses heavily on the hard samples, it diverges when we use the same learning rate as the other 2 methods (our LinUpper & the Uniform baseline). The DRO-KL method required a smaller learning rate to converge, in which case it converged slower than our method. Furthermore, we have also included the DRO-KL method among the compared baselines for the GPT-2-medium experiments. On average, our LinUpper method outperforms the DRO-KL method, which we attribute to its ability to handle outliers (see Table 12 in Appendix B.3 for the detailed results).

Q4: Your results indicate that the benefits of your reweighting strategies are more pronounced in larger models. Could you provide more insights or hypotheses on why this might be the case? Additionally, discussing any observed limitations or challenges when applying these strategies to smaller models would be helpful.

A4: Thanks for the question! Indeed, our results show that the benefits of our reweighting strategies become more pronounced as model size increases. Below, we provide further insights and hypotheses to explain this observation, as well as discuss potential limitations and challenges when applying our methods to smaller models.

(a) Insights for Larger Models

Larger models have a higher capacity to learn from both the clean (informative) and noisy data. As a result, they are more sensitive to the training dynamics and can better leverage reweighting strategies that emphasize certain subsets of the data. Specifically:

Overfitting Control: Larger models are inherently more prone to overfitting due to their increased parameter space. Our reweighting schemes, particularly those that down-weight low-loss samples (e.g. LinUpper), help mitigate this risk by preventing the model from overemphasizing redundant or easy examples.

Representation Capacity: Larger models can represent complex patterns in the data more effectively. This allows them to benefit more from focusing on medium- or high-loss samples, which might correspond to harder but useful examples that smaller models struggle to learn.

(b) Hypotheses for Smaller Models

In contrast, smaller models often lack the capacity to learn complex patterns effectively and tend to underfit. Hence, since these limited-capacity models are less capable of exploiting harder examples, the benefits of emphasizing medium- or high-loss samples are diminished. These models may even rely more on easy examples to make incremental improvements, which reduces the effectiveness of reweighting strategies like LinUpper or Quadratic.

(c) Potential limitations for Smaller Models

Intuitively, in some cases reweighting strategies like LinUpper could even introduce instability for smaller models by shifting focus away from easy-to-learn samples, which are crucial for their limited capacity. However, in practice, as validated by our various experiments, our reweighting method LinUpper provides a consistent performance boost for larger models, while still achieving competitive performance for smaller models.

Dear Reviewer AvfZ,

Thank you once again for your valuable time and effort in reviewing our paper. We have responded to your comments in the threads above and also added new experiments to answer your questions. All edits and added parts in our revised paper are in blue.

As the reviewer-author discussion period ends soon, we would greatly appreciate your feedback on whether our responses and updates address your concerns. If you have any further questions, please feel free to let us know, and we will be happy to provide more clarifications. Otherwise, if you find that your concerns have been resolved, we kindly ask you to consider raising your score.

Thank you for your thoughtful review and feedback!

Regards,

Authors of Paper 13352

Thanks for the detailed response! Most of my concerns have been addressed, and I appreciate the effort you put into clarifying the points I raised. I've increased the score to reflect these improvements. It would be great if the authors could include these improvements in the final version.

Thank you for your thorough reviews and constructive comments. We provide our response to your comments below. If our response resolves your concerns, we would greatly appreciate it if you could consider increasing your score.

Q1: Although the authors state that the proposed method is designed for LLM pretraining, the evaluation is conducted solely on QA tasks. Generally, QA tasks are simpler than generation or reasoning tasks. Therefore, the benefits of pretraining LLMs may not be adequately assessed under this setting.

A1: We thank the reviewer for the valuable suggestions! We have evaluated on more benchmarks, including reasoning and language understanding tasks.

For the GPT-2 models, we find that only question-answering tasks work well at this scale, which is why we report 5-shot results on LogiQA, LogiQA-2, SciQ, and PiQA. All other benchmarks we explored yielded random results for all compared algorithms, i.e., they were not useful for model evaluation. The reason is the limited capacity of the GPT-2 models.

To address the reviewer’s concern, we have conducted new experiments to train additional 1.4B (billion) and 7B model with Llama architecture on randomly sampled subsets of the FineWeb 15T dataset. For the 1.4B and 7B models, we use 100B and 175B randomly sampled tokens from the FineWeb 15T dataset, respectively. We train the models using the Uniform baseline with uniformly averaged losses and our LinUpper strategy.

For both models, we include benchmark results from 13 language understanding & reasoning (LUR) tasks and 6 question-answering (QA) tasks. Please see the new section 6.4 in our main paper and Appendix A for more details, including hyperparameters and model architectural details. We report the benchmark results on the last checkpoints.

After fully training the 1.4B model on 100B tokens, we observe a performance gain on 6 out of 13 LUR tasks with an average performance gain of 0.66%. For QA benchmarks, we improved on 5 out of 6 tasks with an average of 1.72%. For the 1.4B model, we provide a full overview below and in our updated paper in Appendix A, Tables 8 and 9. The first table shows the language understanding and reasoning benchmarks; the second question-answering tasks.

Update: We have completed the training of the 7B model and now report the final benchmark accuracies after 175B tokens. The training took a total of 9800 GPU hours. Overall, our LinUpper strategy performs better on 9 out of 13 LUR tasks, with an average improvement across all benchmarks of 1.45%. Also, we observe an average performance gain on QA benchmarks of 1.16%. We improve on 5/6 QA tasks. We have updated our paper with these final results. The full overview is available in Appendix A, Tables 10 and 11. Below, the first table shows the language understanding and reasoning benchmarks; the second question-answering tasks.

...continue A1.

These improvements are particularly noteworthy given the scale of these LLama models, highlighting the scalability and effectiveness of our method.

Q2: Experimental settings. In the experimental section, the authors evaluate only on smaller language models, specifically GPT-2 models with fewer than 1 billion parameters. Typically, effective LLM pretraining methods should demonstrate scalability to models with at least 10 billion parameters to ensure practical applicability. Consequently, the practicality of the proposed method is not fully substantiated in this study.

A2: Thanks for the suggestions! We understand 300M parameter models may not yield a full picture of our reweighting strategy’s effectiveness on significantly larger foundation models. Therefore, We have conducted new experiments on larger 1.4B and 7B models with Llama architecture. Please see our response to Q1 for the experiments’ details and comparisons. We also added these experiments in section 6.4 and in Appendix A of our revision.

We appreciate the reviewer’s suggestion of even larger 10+ billion parameter models, however pretraining from scratch models at this scale requires a significant amount of resources and time (which is typically not available to academic groups). We are exploring possibilities to pretrain at this scale, which could indeed further justify the practicability of our reweighting framework. Additionally, note that the size of 1 billion parameters is widely adopted for conducting ablation studies for LLMs [1]. As the benchmark results in our response to Question 1 reinforce the results we present in the original paper, we hope the increased model size is satisfactory for the reviewer.

Q3: There are several typos in this paper. For example, "LANGGUAGE" should be 'Language". "fined-grained" should be "fine-grained". "deprioritize redundant or uninformative data" should be "deprioritizing redundant or uninformative data".

A3: We thank the reviewer for noticing these typos! We have corrected these in our revision.

Reference:

[1] Cody Blakeney, et. al. Does your data spark joy? Performance gains from domain upsampling at the end of training, 2024.

Dear Reviewer u8Tc,

We thank the reviewer very much for checking our response and increasing the rating. Many thanks for your effort into this review process!

Regards,
Authors of Paper 13352