Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining
Authors: Daouda Sow
Herbert Woisetschläger
Saikiran Bulusu
Shiqiang Wang
Hans Arno Jacobsen
Yingbin Liang
Venue: ICLR 2025 Poster (2025)
Keywords: sample reweighing
large language models
pretraining
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose new loss-based sample reweighing techniques for improved efficiency and effectiveness of LLMs pretraining.

Abstract: Pretraining large language models (LLMs) on vast and heterogeneous datasets is crucial for achieving state-of-the-art performance across diverse downstream tasks. However, current training paradigms treat all samples equally, overlooking the importance or relevance of individual samples throughout the training process. Existing reweighting strategies, which primarily focus on group-level data importance, fail to leverage fine-grained instance-level information and do not adapt dynamically to individual sample importance as training progresses. In this paper, we introduce novel algorithms for dynamic, instance-level data reweighting aimed at improving both the efficiency and effectiveness of LLM pretraining. Our methods adjust the weight of each training sample based on its loss value in an online fashion, allowing the model to dynamically focus on more informative or important samples at the current training stage. In particular, our framework allows us to systematically devise reweighting strategies deprioritizing redundant or uninformative data, which we find tend to work best. 
Furthermore, we develop a new theoretical framework for analyzing the impact of loss-based reweighting on the convergence of gradient-based optimization, providing the first formal characterization of how these strategies affect convergence bounds. We empirically validate our approach across a spectrum of tasks, from pretraining 7B and 1.4B parameter LLMs to smaller-scale language models and linear regression problems, demonstrating that our loss-based reweighting approach can lead to faster convergence and significantly improved performance.

Published as a conference paper at ICLR 2025
DYNAMIC LOSS-BASED SAMPLE REWEIGHTING FOR IMPROVED LARGE LANGUAGE MODEL PRETRAINING
Daouda A. Sow" The Ohio State University SOw. . 530csu. edu
Herbert Woisetschläger"
Saikiran Bulusu
Technical University of Munich h. woi setschl saegee utum de
The Ohio State University bu uulus.. 1 @osu. edu
Shiqiang Wang
Hans-Arno Jacobsen
Yingbin Liang
IBM Research
University of Toronto
The Ohio State University liang889 edu
wangshiqeus. ibm. com
jaccbsent eecg torento.
ABSTRACT
Pretraining large language models (LLMs) on vast and heterogeneous datasets is crucial for achieving state-of-the-art performance across diverse downstream tasks. However, current training paradigms treat all samples equally, overlooking the importance or relevance of individual samples throughout the training pro- cess. Existing reweighting strategies, which primarily focus on group-level data importance, fail leverage fine-grained instance-level information and do not adapt dynamically to individual sample importance as training progresses. In this paper, we introduce novel algorithms for dynamic, instance-level data reweighting aimed at improving both the efficiency and effectiveness of LLM pretraining. Our methods adjust the weight of each training sample based on its loss value an online fashion, allowing the model dynamically focus more informative important samples at the current training stage. In particular, our framework allows to systematically devise reweighting strategies deprioritizing redundant unin- formative data, which we find tend to work best. Furthermore, we develop a new theoretical framework for analyzing the impact of loss-based reweighting on the convergence of gradient-based optimization, providing the first formal characteri- zation of how these strategies affect convergence bounds. We empirically validate our approach across spectrum of tasks, from pretraining 7B and 1.4B parameter LLMs smaller-scale language models and linear regression problems, demon- strating that our loss-based reweighting approach can lead to faster convergence and significantly improved performance.
INTRODUCTION
The rapid advancement of large language models (LLMs) (Brown et al. 2020 Raffel et al. 2020
Touvron et al.) 2023, Chowdhery et al.) 2023) Achiam al. 2023] Dubey et al. 2024) has revolution- ized natural language processing and artificial intelligence (Al) capabilities. These models, trained on billions or trillions of tokens, exhibit remarkable generalization capabilities across wide range of downstream tasks. As datasets grow web-scale proportions and models become increasingly large, the need for more efficient training techniques has become paramount.
Existing approaches to LLM pretraining predominantly involve two phases; heavy data curation (Longpre et al.) 2023 Dubey et al., 2024; Wettig et al. 2024 Penedo et al. 2024] and training with uniform sampling the constructed corpus. Data curation for LLM pretraining typically involves a combination of automated filtering techniques and manual quality checks. For instance, heuristic-based filters often employed remove low-quality content and deduplicate data. Some approaches use perplexity-basee filtering or auxiliary classifiers (Gao et al.) 2020) |Pneedo et al. (023) to identify high-quality samples. Manual curation then applied to refine these filtered datasets, often involving human evaluation of subsets of the data to ensure quality and relevance. However,
Daouda Sow and Herbert Woisetschläger were interns IBM Research when most of this work was done.

---

Published as a conference paper at ICLR 2025
these approaches face significant limitations due scalability issues and the static nature of data selection. First, as the pretraining corpora grow to hundreds of billions of tokens, manual curation becomes increasingly infeasible. The sheer volume of data makes it impractical for humans to review even a small fraction of the dataset. Second, data curation methods cannot adapt to the changing importance of samples during the training process, e.g., the model improves over time certain samples that were useful early in training can become irrelevant in later training stages.
In addition, conventional gradient-basee based training pipelines often treat all data points as equally informative regardless of their individual importance at current training stage. This uniform sampling strategy, while simple, overlooks the nuanced diversity within large-scale datasets and potentially wastes computational resources on less informative samples. Recent work (Xie et al. 2023 Fan et al.) 2023) has explored group-level reweighting strategies, which adjust the importance of entire domains or groups of data. While these methods have shown promise, they operate at coarse level, do not fully leverage the fine-grained information available the instance level, and still lack dynamicity in importance assignment during the training phase.
Given these limitations, this paper seeks to address the following fundamental question:
How can we dynamically leverage instance-level information to accelerate training and improve model performance without incurring significant computational overhead while potentially reducing the need for extensive data curation?
Answering this question presents several intertwined technical challenges, particularly in the context of LLM pretraining. First, unlike in computer vision, where samples are seen repeatedly over the course of training (often hundreds of times), in LLM pretraining, individual samples are typically en- countered only once due to the vast size of the datasets. This renders dynamic reweighting approaches primarily devised for multi-epoch training or relying on historical statistics, such as previous loss or gradient norms (Loshchilov & Hutter] 2015; Jiang et al. 2019),, inadequate within the context of LLM pretraining, Second, storing previous statistics for individual samples becomes computationally infeasible when working with web-scale datasets, and methods requiring the persistent storage of sample-level data would significantly inflate resource requirements. Furthermore, for general- purpose LLM, leveraging small hold-out validation sets to compute importance weights through, e.g.. bilevel optimization as in Grangier et al. (2023), is ineffective.
Addressing these challenges, our work makes the following significant contributions:
Instance-Level Loss-Based Reweighting Strategies. We introduce and systematically study a variety of instance-level, loss-based reweighting strategies for improving both the efficiency and effectiveness of ML training, especially within the context of LLM pretraining. Each strategy is meticulously designed to achieve specific goals, such as focusing the learning dynamics on different parts of the loss distribution. This builds on recent work that emphasizes the importance of data diversity and sample importance in pretraining and extends previous works on group-level reweighting for LLM pretraining, such as DoReMi Xie et al. (2023). DoGE Fan et al. (2023). It also adds new dimension by incorporating more fine-grained per-sample dynamics rather than donain-level adjustments. In fact, combining our reweighting schemes with DoGF/DoReMi achieves better or comparable performance on various few-shot reasoning benchmarks when we train on the SlimPajama dataset. Moreover, our study reveals that, in general, strategies down-weighting the importance of low-loss samples tend to consistently yield performance improvements across various scenarios.
New Theoretical Framework. We develop a new theoretical framework for analyzing the effects of loss-based reweighting on training acceleration. To the best of our knowledge, this represents the first explicit characterization of loss reweighting effects within the convergence bounds of gradient methods under widely adopted loss geometries. Our derived convergence bounds provide theoretical justification for the empirical success of downweighting low-loss samples, demonstrating faster convergence under this strategy.
We conduct extensive experiments that corroborate our claims and support our theoretical findings. Notably, we demonstrate that the advantages of downweighting low-loss samples are observed across a spectrum of problem scales and complexities: (i) complex, large- scale problems such LLM pretraining with number of parameters ranging from 124M to 7B, approach leads to notably improved performance and faster convergence; (ii) in simple, small-scale problems like linear regression, the results highlight the fundamental nature of our findings.
Empirical Validation.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Training Data Re- Ne-weighting/Selectne for LLMs. Several recent studies (Xie et al. 2023 Chen
al. 2023 an et al. 2023 Thakkar al. 2023 have explored various reweighting techniques
to enhance the generalization and efficiency of language models pretraining. For instance, Xie al. (203) and Fan et al. (023) optimize the composition of pretraining corpora to achieve better performance across pretraining domains or for out-of-domain generalization. Chen et al. (2023) introduce a framework for ordered skill learning. optimizing data selection based on how effectively it teaches interdependent skills for continual pretraining and fine-tuning regimes. Although effective, these techniques operate at the group level, whereas our work explores reweighting at the instance level, offering finer control over how individual samples are treated based on their loss values. Furthermore, we demonstrate that combining domain-level methods such as DoReMi (Xie et al.] 2023) or DoGE (Fan et al., 2023) with our instance-level reweighting methods results in improved performance across multiple domains. Instance-level reweighting has been used in post-training settings of LLMs (Chen al. 2024; Jiang al.) 20244. Jiang et al. (2024) boost the self-improvement abilities of LLMs by employing sample reweighting to filter out self-generated data that have correct answers but exhibit high distribution shifts. Chen et (2024) reweight individual to focus on medium- loss samples. In contrast,
samples during continual shining/instructiontuuum our work systematically studies the effects of various sample-level, loss-based reweighting strategies on the efficiency and effectiveness of LLMs pretraining- The approach in |Fan & Jaggi (2023) offers a curriculum learning framework that prioritizes samples with a higher learnability score, which is precomputed using another auxiliary model similar to DoReMi and DoGE. While we do not explicitly address curriculum learning in this work, our re-weighting mechanisms naturally allow for implementing form of loss- based curriculum learning algorithms without the need to train and store additional proxy models as in Fan & Jaggi (2023).
Sample-level Re-weighting generic ML solution. Sample-level reweighting has been extensively explored in other machine learning areas. For image classification, Loshchilov Ulassfication,Loohhllloo & Hutter (2015), Jiang
et al. 2019), and Katharopoulos & Fleuret (2018) introduce pioneering approaches that prioritize samples based on loss values or gradient norms with the goal of accelerating training speed. Although these methods emphasize the importance of selecting high-loss samples during training, they typically require additional forward/backwarn passes on each training sample and are primarily designed for multi-epoch training. limiting their applicability with the vast pretraining corpus used for LLMs. Using bilevel optimization, Grangier et al.) (2023) adapt the data distribution during training to focus more on relevant samples for target data distribution. Similarly, Ren et al.) (2018) employ a meta-learning approach to adjust sample weights based on validation set performance, which excels in noisy and imbalanced datasets. this work, we introduce and investigate various lightweight, loss-based reweighting techniques that add little to computational overhead compared to uniform sampling and do not require any nested optimization routines, often arising with bilevel optimization and/or meta-learning. Sample-level reweighting has also been explored in the context of adversarial
machine learning (Zhang et al. 2020 Liu et al. 2021 ZZng et al.) 2021 Sow et al. 2023). ), domain
adaptation (Jiag & Zhai [0077] Fang et al. 2020), data augmentation (Yi al.] 2021), ) and imbalanced
classification (Qi et al. [2021; Ren et al. 20118). Our reweighting mechanisms also have the potential
to be studied under these contexts, which we leave future work.
PRELEMINARIES: AUTOREGRESSIVE LANGUAGE MODELING & GOAL
Large language models (LLMs) are typically trained using autoregressive language modeling, where the objective is to predict the next token in a sequence given the previous tokens. Formally, given a sample sequence of tokens x (21,, 22, 2T), where o; represents the 1-th token in the sequence, the L.L model parameterized by E RRd is trained to maximize the likelihood of the sequence:
T
(IIM(x; 0)
II P( ((..11)..........:0))
where P( (,-1;,; denotes the conditional probability of generating token I, given 0 and all previously seen tokens æ1,... ... It-1- In practice, autoregressive LLMs typically compute the loss f (x;; 8) for sample x using the negative likelioiddood (NLL) of the predicted tokens.
T
(xx:0) = - llg PLIM(x;)) 0) = log P(ze|
Σ
3

---

Published as a conference paper at ICLR 2025
Based on these sample losses, conventional Sik-like algorithms will then update the model parame- ters with the average gradient over all samples in a batch B at each training step t, i.e.
0 = 0²
B Σ [x((x;)))
(I)
which results in each sample contributing equally the overall training update. While this approach enjoys simplicity, it inherently treats all data points as equally informative, regardless of their individual relevance difficulty. This uniform treatment overlooks the nuanced diversity within the data, which is particularly critical for current LLMs trained on vast and heterogeneous web-scale datasets. In such scenarios, the ability dynamically differentiate between high-value and lesser- value data samples at different training stages is paramount for optimizing training efficiency and model performance.
In contrast, our work introduces and benchmarks a variety of simple, lightweight reweighting strate- gies that can be seamlessly integrated into existing training pipelines with little to no computational overhead. These strategies aim to address the inherent limitations of the standard uniform averaging method by assigning dynamic importance to individual data samples based on their loss values, thus providing more refined control over the training process. Ultimately, our objective is twofold: (i) to improve the efficiency and effectiveness of LLM training by allowing the model to dynamically focus on more relevant or informative data at different training stages and (ii) to potentially pro- vide a scalable, automated approach to online data selection during pretraining. This ambitiously seeks to reduce the need for extensive manual data curation and selection efforts, which are increas- ingly becoming infeasible given the growing size and diversity of pretraining corpora. We envision these reweighting strategies foundation for more efficient and adaptive data utilization in LLM pretraining, ultimately reducing the time, cost, and complexity associated with model pretraining.
4 APPROACH
4.1
TECHNICAL CHALLENGE: DATA REWEIGHTING FOR LLMS MUST BE DYNAMIC AND FULLY ONLINE
The core of our framework is to dynamically reweight individual training samples based on their importance at different training stages. Although sample importance seems static-intuitively, "garbage data is garbage data" regardless of model or training stage-data reweighting should be dynamic. For instance, similar or duplicated samples in training corpora may be useful early in training, but should be deprioritized once the model captures their patterns. Similarly, hard but useful sample could be assigned larger weight, however, as the model learns it could be beneficial to reduce its weight. Based on these observations, update the model at each step using
g = 17p Σ w V fx;;;))
E
where Eieg w(x wxx;; 0¹) and the weight w 0¹) is dynamically assigned based on sample and current model 0'. However, designing effective dynamic reweighting methods for LLM pretraining involves several inherent complexities. L.LM pretraining, each sample typically seen only once due to the vast size of the datasets, which contrasts with problems where repeated exposure to data is common, such as in computer vision problems. This makes methods that depend on accumulating historical data, such previous loss gradient norms (Loshchllov & Hutter 2015 Jiang al.) 2019), unsuitable. Furthermore, the sheer scale of these large corpus makes it computationally prohibitive to track and store sample-specific statistics, as it would significantly increase resource requirements. Additionally, approaches that rely on hold-out validation sets to compute sample weights, as in Grangier al. (2023), are also irrelevant for pretraining a general-purposs LLM. These unique technical challenges necessitate the development of a fully online, lightweight reweighting approach that scales efficiently with the size of a modern LLM training corpus, which we discuss next.
4.2 LOSS- REWEIGHTING STRATEGIES
To address the aforementioned challenges, we propose a novel framework of lass-based sample reweighting strategies motivated by the following theorem. Theorem characterizes the effects
4

---

Published as a conference paper at ICLR 2025
X10-
rewraightho aratbogg
0.0
weignt
0.8
0.6
LinUpper, 1-0.5
0.6
LitUpper. r=1
Quadratic
0.4
UnUpperr /m2
02
Lirkppger
Linlpper, 1-10
Extremes
0.
Unitormitlight
0.0
0.0
1.0
normalized
normalized
Geometric curves of the different reweighting functions. Right: Shape of the Li Lipuppe strategy after applying Eq. 4 top of it for different values of r. These plots obtained for a batch of 128 uniformly drawn losses. As increases, LinUpper converges to the uniform averaging method.
Figure 1: Left:
of loss-based reweighting on the convergence bound of reweighted full gradient descent when the sample losses f(xx;; (²)) are convex. The extension to the stochastic setting and analysis for other types of loss geometries is deferred to Section [ Our full statement of Theorem I and detailed proofs for all theoretical statements can be found the appendix.
Theorem 1. Consider M data points and let each loss f(x;;) be convex. Further, assume the
interpolation regime holds, i.e,, 30" E IRd. such that 6* E arg mingera f(x;;) 0) vi. Then, for reweighting scheme that satisfies max;e (1 w(x;; 0*) 2/M, we have
1
Σ
M
f(x;00)
1 M M
Σ 1=1
50 f(x;;0) 50
(1000
8
Σ
T-1
8
(3)
M
T
T
1=1
1=0
where 8¹
Ei
(x:x;;00 0²) (f(x;; 05) [(x;; 0*)) and 0th = + - 0..
Theorem illustrates how setting the importance weight wxx;; 0*) based on the loss gap f(x;; 0*) f(x;; 8*)) influences the convergence bound. Specifically, make the following key observations: (i) fixing w(x; 0*) = M Vi, leads to 6, and we recover the traditional convergence bound of gradient descent for convex functions; (ii) assigning larger weights to the smaller loss gaps leads to 6₂ 2 0, which yields a looser convergence bound; (iii) to achieve a better bound (6, < 0), one should deigh the importance of samples with low loss gaps. Further, note that Eq. ( critically holds only for reweighting schemes satisfying w(,;;; 0') I 2/M, which puts an upper bound on the maximum weight that can be assigned to single data point, and hence, eliminates worst- case importance assignment strategies-such as those based on distributionally robust optimization (Qian et al., 2019 Qi al., 202T] et al., 2023) which are prone to overfitting to outliers, particularly problematic issue when dealing with the large-scale, noisy corpora used in LLM wx,;; 0*) S 2/M essentially means that the weight of any sample (after reweighting) should not be more than twice the uniform weight. This upper bound on the sample weights is directly captured by our convergence analysis. Please see Appendix A.5 for discussions on how it be satisfied by our proposed reweighting methods.
pretraining. Finally, we also note that max;e
Based on the insights from Theorem T on the relationship between loss and importance weight, we study series of reweighting strategies, each carefully designed focus the learning dynamics different parts of the loss distribution during training. Let wi be the importance weight for sample x; where we drop the dependency on O' for clarity. We first normalize the individual sample losses into a bounded range [-a, or). This step crucially ensures the loss values are scaled consistently and amenable to the subsequent weighting functions. We then apply the following analytical functions on the normalized loss h, Figure [ shows the geometric curves of these different functions.
Linear Upper-Bound Strategy (LinUpper). In this strategy, the weight is proportional to the normalized loss but is capped at predefined a value, ensuring that outliers do not dominate the training process. The functional form is 8; == min (h: + a..
Quadratic Strategy (Quadratic). This strategy prioritizes samples with moderate loss values while down-weightin both low-loss (easy or repetitive) and high-loss (potential
outliers) samples by applying a quadratic function 8; a x²
5

---

Published as a conference paper at ICLR 2025
Extremes-Based Strategy (Extremes). This scheme emphasizes both the hardest (high- loss) and the easiest (low-loss) samples by applying s; := h4s, ensuring that samples at the extremes of the loss distribution receive higher importance weights.
To further enhance training stability and effectiveness, we add curriculum-baa based adjustment mechanism top of the strategic sample weight s; by controlling the sharpness of the importance weighting using a temperature parameter r. The final weight w; is computed as:
e///r
E/e5/F
where r annealed during training. Early in training, when the model is still learning fundamental patterns and losses less informative, we employ more uniform weighting scheme ensure diverse feature learning by adopting a large r (see Figure [.). As training progresses, we decrease the value of r to enact full use of reweighting strategies. Algorithm| [ depicts our approach. It is important to note that our reweighting approach adds nearly zero computational overhead. During the forward pass, we already compute the sample losses, which are then used to derive the importance weights through simple analytical transformations. Since no extra loss or gradient computations are required, our reweighting strategies can be seamlessly integrated into standard training pipelines without incurring significant additional costs in terms of time or resources.
Algorithm 1 Fully Online Instance Reweighting
I: Input: stepsize n, initial model parameters 0, IRd temperature {rx (ri)( fmin and fmax-
for 0, ...,
3: Draw minibatch of data samples 8 {x: )x:
4: Run forward pass compute losses (fre) for all samples in B
Normalize losses into interval [-a, al using linear transform: lu,e 2(Se, - fmin)
Jmax Jann
6: Apply reweighting strategy normalized losses obtain {st,eff_ (LinUpper, Quadratic, Extremes]
7: Apply curriculum adjustment using Eq. (4) obtain sample weights {wi,e)i=1
8: Update 0²=1 = EiEB Wi,e V fi,
9: end for
Optimal strategy minimizing 8*. We next demonstrate that applying Eq. (4 on top of our strategy
nippp,rr, fact, corresponds the optimal strategy that minimizes a KL. Ki--ivergence regularized version of 8 in Theorem
Proposition 1. The optimal strategy that minimizes the KI-divergence regularized 8', i.e., 8 r +E w; log(M (Mws), is given by
2
C min
exp
(5)
M
where C. is normalizing constant that ensures Σx Wt 1.
Proof. Proof relegated to Appendix
As shown in Theorem D and Proposition T our experiments confirm that the Linnppper reweighting scheme indeed achieves faster convergence and improved performance across spectrum of LLMs pretraining scales. The incorporation of KL-divergence regularization in the optimization of 8 prevents extreme solutions that focus only the highest-loss samples. Without this regularization term, the optimal solution becomes trivial: it assigns a weight of 2 M to the 2 samples with the highest losses and weight to the remaining samples. It is important discards half of the dataset and performs poorly in practice. However, our L LinUpper method offers a balance between focusing on high-loss samples and maintaining the diversity of data utilized during training, leading to faster convergence and performance improvements.
M
note that such strategy
5 ANALYSIS OF
REWEIGHTING ON CONVERGENCE BOUND
In this section, we analyze the effects of reweighting on the convergence bound in the more generic setting of stochastic gradient methods possibly with momentum. Consider the problem

---

Published as a conference paper at ICLR 2025
of minimizing finite sum of M objective functions:
M ((0)
1 M
min
((
BER4
where each fi Rd + R corresponds to sample loss function f(x;:)) We make the following widely adopted assumptions:
Assumption (Convexity). Each function fi is convex, i.e, for all 0,0' E Rd and i 1...., M:
[=(0') 2 f:(0) + (VI,(0),''' -)..
Assumption 2 (L-smoothness). Each function Ja is L--smooth. Formally, for all 0,0 l Rdd and i e ,,
[Vi(0') - v)(00)| s L: O0]
(8)
and I. [aX T4-
Assumption (Interpolation condition). There exists a global optimum 0- E IRd such that: 6" arg min f((0), E (1....,,M
BERd
Assumption is widely adopted in modern ML settings (Dar al. 2021 Radhakrishnan et al. 2020
Jhunjhunwala et al., 2023), particularly when analyzing overparameterized deep neural networks. This ensures that all individual sample objectives can be minimized simultaneously. effectively eliminating adversarial trade-offss where minimizing one objective could worsen another.
Next, we characterize the effects loss-based reweighting on the convergence bound of reweighted minibatch SGD when the sample losses Ji are convex. For ease of notation, we denote wx(;;; 0*) by Wi, in this section and the corresponding proof's in the Appendix.
Theorem 2 (Minibatch SGD with momentum). Let Assumptions 7 and 3 hold. Consider minibatch of size [B] = b with reweighting scheme satisfying maXçeB WA, S 2/b. Then, for 8LVT41 d > O, we have
E[f(0") - {(0'))
8L 80
1-11
2
T-1
6
+
2
T-1
Σ
(10)
VT+1 +1
T 1
T 1
t=0
1
where 6N = E2 Σ iEB
(1: wa,) (fi(0') - f((())) )),
= E
ΣiEB (b: -wr,)
and 0" = + T-11000 0º.
Proof. Proof provided in Appendix
We note similar observations as in Theorem I and again down woiihhiig the importance of samples with low loss gaps leads to a better bound. Typically, the minibatch gradient with uniform weights is unbiased, E 1>ie (0')] functions of the sample functions, hence, leading to biased gradients and the analysis of Theorem [ being more involved. Further, we provide the analysis for minibatch SGD when the sample functions are non-convex in Appendix C.4
(0')). However, in our setting, the importance weights are
EXPERIMENTAL EVALUATION
6.1 TRAINING AND EVALUATION SETUP
We conduct our first set experiments using decoder-only transformer models /assanni 2017 Rad- ford et al. 2019) with parameter sizes of 1 120M, 210M, and 300M, which we refer to respectively, GPT2-mini, GPT2-small, and GPT2-medium. We train on the SlimPajama (Soboleva et al. 2023) corpus that includes seven diverse domains: Common Crawl (CC), C4, GidHub, StackExchange, Book, Arxiv, and Wikipedia. In the first step, we pretrain the three GPT2 models on all seven domains where compare
sample-level reweighting methods (IinUpper, Quadratii icc. Ext remes) s) against the uniform averaging bascline in which each sample contributes equally. In a second step, we assess the impact of combining our reweighting techniques with existing domain-level reweighting
7

---

Published as conference paper at ICLR 2025
4.5
Pupplotty (00g
4..0
Unioom
Purplastnn (lag)
UnUpper (Ders)
2.5
Quadiatic (ours)
5.0
4.0
2.5
4.0
4.5
15020
15020
12030 Stap
15000
Step
Step
Step
(a) Arxiv Dataset
(b) Book Dataset
(c) C4 Dataset
(d) CC Dataset
4.0
Perplssty (oo)
A.6
Perpodit (lo)
4.5
4.6
500
6.0
5.0
255
3.0
5000
12000 Stop
15000
5000
10000
15000
5000
10000
15000
Slop
GitHub Dataset
StackExchange Dataset
(g) Wikipedia Dataset
Per-domain perplexities on hold-out validation sets under the uniform domain sampling setting for the GPT2-medium model. Our reweighting strategy LinUpper strategy achieves better or at least comparable perplexity on 5 out of 7 domains.
Figure 2:
methods, specifically DoGE and DoRcMi, to demonstrate the potential for further performance gains when fine-graine lnn-grained sample reweighting is added on top of domain-level adjustments. Then, we conduct experiments of pretraining larger models with 1.4B and 7B parameters (Section (.4). We also provide results for regression problem on synthetic data in Appendix [ to validate the fundamental advantage of reweighting. Model architecture details can be found in Appendix A.2 Our codebase for the GPT2 experiments is publicly available
We measure performance using two types of metrics:
Perplexity on Training Distribution. We evaluate per-domain perplexity and average perplexity on all training domains using hold-out validation sets. By doing we aim to generate insights on the sample efficiency during pretraining, i.e., how well a model can learn from a given dataset and adapt to the data distribution.
Few-Shot Reasoning Ability. assess generalization beyond the pretraining distribution, we evaluate models series of diverse reasoning tasks in 5-shot setting, including LogiQA Liu et al.) (2020), LogiQA Liu et al. (2023), SciQ/Well et al. (2017). and PiQA Bisk et al. (2020).
6.2 COMPARISONS UNDER UNIFORM DOMAIN SAMPLING
Comparisons on Pretrainining Distribution. Figure shows the per-domain perplexity for the GPT2-medium model on all seven domains. We provide the final per-domain perplexity results for the GPT2-mini and GPT2-small models in Table I The perplexity plots for GPT2-mini and GPT2-small can be found in Appendix A.33 Figure 2 shows that our proposed LinUpper strategy outperforms matches the other baseline in 5 out of 7 domains, with particularly notable improvements in CC, C4, and Book domains. upper achieves significantly lower perplexity in these domains, demonstrating its effectiveness at handling noisier data sources. This suggests that down-well weighting low-loss samples -central to LinUpper- helps reduce the effects of redundant or similar data, and thus accelerates convergence as supported by theoretical findings. We also note competitive results for LinUpper method at the smaller scale models in Table I but the advantage is more significant in the larger GPT2-medium model.
Few-Shot Evaluations. Table provides the 5-shot evaluations on different tasks for the GPT2- medium model. The LinUpper reweighting strategy consistently yields competitive results, out- performing the other compared baselines on out of 4 tasks. These results further validate the general effectiveness of LinUpper across different types of reasoning challenges, including logi- cal inference (LogiQA), physical commonsense reasoning (PiQA), and scientific reasoning (SciQ). indicating that reducing the effects of Notably, Quadratic performs best LogiQA L 2 (28.6%).
Rewei ght in- CCIR-2225

---

Published as a conference paper at ICLR 2025
Table Per-domain perplexities on hold-out validation datasets under the uniform domain sampling setting for the GPT2 models. Across a wide variety of domains, our Linüpper method outperforms the uniform baseline.
Benchmark GitHluh
Model
Methed
Arxiv
Book
C4
CC
Stacknxchange
Wikipedia
Meun
Uniform
2.48
4.21
4,47
4.44
1.84
2,66
3.14
3.32
LinUpper tours)
2.49
4.41
4.38
1.89
3.12
3,30
GPT2-mini
Quadratic (ours)
2.48
4.26
4.55
4.53
1.87
2.67
3.17
3.36
Extremes (ours
2.54
4.23
4.47
4.44
1,90
2.74
3.19
3.36
Uniform
2.36
4.04
4.29
4.26
1.67
2.51
2.90
3.15
GPT2-small
LinUpper
2.36
3.99
4.23
4.19
1.73
2.53
2,88
3.13
Quadratic (ours)
2.38
4.14
4.43
4.41
1.72
2.54
2.99
3.23
Extremes (ours)
Uniform
2.34
4.02
4.27
4.22
1,66
2.49
2.88
3.13
GPT2-medium
Lipupper tours)
2.34
3,96
4,20
4.15
1.72
2.50
2.85
3,10
Quadratic (ours)
Extremes (ours)
2.45
4.08
4.31
4.27
1,76
2.62
3.00
3.21
Table 2: 5-shot Performance evaluations under uniform domain sampling for the GPT2-medium model. We report the accuracy for every benchmark. The highest accuracy per benchmark is highlighted bold. The LinUpper reweighting strategy consistently yields competitive results, outperforming the other compared baselines 3 of tasks.
Method
Benchmark Accuracy
Uniform
LinUpper (ours)
Quadratic (ours)
(ours)
LogiQA
25.7
27.9
LogiQA 2
27.5
27.6
28.6
27.7
SciQ
49.0
51.8
48.7
49.6
PIQA
55.5
56.2
54.6
55.2
Mean
39.4
40.9
39.3
39.8
both high-loss (potential outliers) and low-loss (potentially redundant data) samples sometimes improve generalization performance even though it generally does not achieve best efficiency for learning the pretraining distribution. This suggests that the Quadratic scheme could be useful when training under noisier and less curated datasets. Also note that the remes strategy consis- tently lags behind the other methods, which further confirms the advantage of reducing the importance of low-loss samples.
6.3 COMPARISONS UNDER NON-UNIFORM DOMAIN SAMPLING
The performance of the compared methods under non-uniform domain sampling setting are presented in the Table The results clearly demonstrate the value of incorporating our LinUpper strategy on top of existing domain-level reweighting methods such as DoGE and DoReMi. Across all tasks, LinUpper consistently improves the performance of both domain reweighting baselines, reinforcing the effectiveness of dynamically adjusting sample weights at the instance level in conjunction with Nomain-leved reweighting. In particular, our reweighting strategy LinUpper provides notable improvements for LogiQA task, boosting accuracy from 27.2% to 28.6% with DoGE and from 27.2% to 27.6% with DoReMi. The gains in LogiQA are similar, where our reweighting method LinUpper consistently improves the performance of both methods. For scientific reasoning (SciQ) and physical commonsense (PiQA) tasks, LinUpper strategy also enhances both baseline methods, with the most substantial improvement observed in SciQ (from 52.8% to 53.2% with DoGE and from 53.3% to 54.5% with DoReMi). These results indicate that LinUpper is particularly effective in tasks requiring nuanced reasoning and factual knowledge.
6.4 EXPERIMENTS ON 1.4B AND 7B LLAMA MODELS
We conduct experiments to train 1.4B (billion) and 7B parameter models with Llama architecture randomly sampled subsets of the Fine 15T dataset. For the 1.4B and 7B parameter models, we use approximately 100B and 175B randomly sampled tokens from the Feb 15T dataset, respectively. We report the benchmark results the last checkpoints.

---

Published as a conference paper at ICLR 2025
Table 5-shot performance evaluations under non-uniform domain sampling for the GPT2-medium model. Our Li Linupper instance reweighting approach creates synergies with state-of- of-te-tf domain reweighting techniques.
Method
Benchmark Accuracy
DoGE
DoGE LinUpper
DoRcMi
DoRcMi + LinUpper
LogiQA
27.2
28.6
27.2
27.6
LogiQA
27.5
28.0
27.7
27.9
SciQ
52.8
PiQA
55.8
56.3
55.7
56.1
40.8
41.5
Table 4: Average performance of the 1.4B pa-
Table 5: Results of the 7B parameter model.
rameter model across LUR and QA benchmarks.
Overall, our LinUpper method improves re- sults for 14 out of 19 tasks with an average per- formance gain of 1.44% among LUR tasks and 1.16% among QA tasks.
Overall, our LinUpper method improves re- sults for 11 out of 19 tasks with an average per-
formance gain of 0.66% among LUR tasks and 1.72% among QA tasks.
Avg. Performance
Uniform
Lint LinUpper pper (ours)
Difference
Avg. Performance
Uniform
LinUpper (ours)
Difference
LUR (13 tasks)
48.51
49,16
0.66
LUR (13 tasks)
50.54
$1.98
1.44
QA (6 tusks)
42.35
44.07
1.72
QA (6 tasks)
48.15
49.31
1.16
Across tasks
46.56
47.55
0.99
Acmos 19 tnoks
49.79
51.14
1.35
For these larger models, we use the -anwwtinnannwwrring (QA) benchmarks and additionally employ language understanding and reasoning (LUR) tasks since their capacity allows for more complex tasks. Please see Appendix A for more details, including hyperparameters, model architectural details, and performance the full list of benchmarks.
Tables and show the average performance among LUR and QA tasks for the 1.4B and 7B parameter models, respectively. For the 1.4B parameter model, we that our Lin strategy improves the average model performance by 0.66% and 1.72% for LUR and QA tasks, respectively. In particular, we note that our L,i nUppe method improves performance for 11 out of 19 tasks in total (see Tables 8 and 9 in Appendix A.3). For the 7B parameter model, our LinUpper method improves results for 14 out of 19 tasks (see Tables 10 and II in Appendix (A.3) with an average performance gain of 1.44% among LUR tasks and 1.16% among QA tasks. These improvements are particularly noteworthy given the scale of these models, highlighting the scalability and effectiveness of our method.
7 CONCLUSION
In this paper, introduce novel sample-level reweighting framework aimed improving the efficiency and effectiveness of large language model (LLM) pretraining. By dynamically adjusting the importance of individual samples based their loss values, approach overcomes the limita- tions of traditional uniform averaging methods and adds a new dimension to existing domain-level reweighting methods by incorporating theoretical and empirical validation, we demonstrate that down-weightte ates convergence and improves performance. The experiments show that our proposed L LinUpper strategy consistently outperforms uniform sampling on common LLM reasoning and commonsense benchmarks. Similarly, our approach creates synergies with existing domain reweighting techniques, which underpins the importance of sample-level dynamics. We observe that the benefits of our method are more pronounced larger models, while smaller models show less significant improvements on benchmarks. This could be due to the fact that smaller models may have limited capacity to fully exploit the advantages of our approach. However, our overall findings highlight the potential of loss-based reweighting strategies to optimize LLM pretraining both in training efficiency and in model evaluation performance.
fine-grained sample-level dynamics. Through extensive
low-loss samples acceler-
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank the IBM Data Engineering Team for their support in running experiments with 1.4B and 7B parameter models. The work of D. Sow and Y. Liang was supported in part by the U.S. National Science Foundation with the grants ECCS-21 13860, ECCS-241352S and CNS-2112471, The work of H. Woisetschläger was partially supported by the Bavarian State Ministry of Economic Affairs, Regional Development and Energy with grant DIK0446/01.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprint 1xXiv:2333.0877: 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp. 1432-7439, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1 2020.
Mayee F Chen, Nicholas Roberts, Kush Bhatia, Jue WANG, Ce Zhang, Frederic Sala, and Christopher Re. Skill-it! data-driven skills framework for understanding and training language models. In shirth- Conference on Neural Information Processing Systems, 2023. URL https //openreview.oo net/ Dot.:e/oooommi---oloooo000.
Xuxi Chen, Zhendong Wang, Daouda Sow, Junjie Yang, Tianlong Chen, Yingbin Liang, Mingyuan Zhou, and Zhangyang Wang. Take the bull by the horns: Hard sample- reweighte continual training improves Ilm generalization. arXiv preprint arXiv:2402. 14270, 2024.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 141240):1-133, 2023.
Ychuda Dar, Vidya Muthukumar, and Richard G Baraniuk. A farewell to the bias-variance tradeoff? overview of the theory of overparameterized machine learning. arXiv preprint arXiv:2009.02355 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint 2rXiv:2407.21783 2024.
Simin Fan and Martin Jaggi. Irreducible curriculum for language model pretraining. arXiv preprint arXiv:2310.1533 2023.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. DOGE: Domain reweighting with generalization estimation. In Second Agent Learning in -eddenndssssss Workshop, 2023. URL https
//openreview net http//orrmmmiicee wtw.et//oormmm.ddi/ieesrr/y
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. Advances in neural information processing systems, 33: 11996-12007, 2020.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00 000027 2020.
David Grangier, Pierre Ablin, and Awni Hannun. Adaptive training distributions with scalable online bilevel optimization. arXiv preprint 1rXiv:2311.11997. 2023.
11

---

Published as a conference paper at ICLR 2025
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp:
Speeding up federated averaging via extrapolation. In The Eleventh International Conference on Learning Representations, 2023.
Angela Jiang, Daniel L-K Wong. Giulio Zhou, David G Andersen, Jeffrey Dean, Gregory R Ganger, Gauri Joshi, Michael Kaminksy, Michael Kozuch, Zachary C Lipton, al. Accelerating deep learning by focusing on the biggest losers. arXiv preprint arXiv: 2019.
Chunyang Jiang, Chi-min Chan, Wei Xue, Qifeng Liu, and Yike Guo. Importance weighting can help large language models self-improve. arXiv preprint (xXx:22408.09849 2024.
Jing Jiang and ChengXiang Zhai. Instance weighting for domain adaptation NLP. Annie Zaenen and Antal van den Bosch (eds.). Proceedings of the 45th Annual Meeting of the Association of Computational Linguistics, pp. 264 271, June 2007.
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pp. 2525-2534 PMLR, 2018.
Ramnath Kumar, Kushal Majmundar, Dheeraj Nagaraj, and Arun Sai Suggala. Stochastic re-weighted gradient descent via distributionally robust optimization. arXiv preprint (rXiv:2306.0922, 2023.
Feng Liu, Bo Han, Tongliang Liu, Chen Gong, Gang Niu, Mingyuan Zhou, Masashi Sugiyama, et al. Probabilistic margins for instance reweighting in adversarial training. Advances in Neural Information Processing Systems, 34:23258-22 23269,, 2021.
Hanmeng Liu, Jian Liu, Leyang Cui, Zhiyang Teng, Nan Duan, Ming Zhou, and Yue Zhang. Logiqa 2.0-an an improved dataset for logical reasoning in natural language understanding. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 2023.
Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08 /24, 2020.
Shayne Longpre, Gregory Yauney, Emily Reif, Katherine Lee, Adam Roberts, Barret Zoph. Denny Zhou, Jason Wei, Kevin Robinson, David Mimno, et al. A pretrainer's guide to training data: Measuring the effects of data age, domain coverage, quality, toxicity. arXiv preprint arXiv:2305.13169, 2023.
Ilya Loshchilov and Frank Hutter. Online batch selection for faster training of neural networks. arXiv preprint arXiv:1511.0634. 2015.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon 1l:: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01 2023.
Guilherme Penedo, Hynek Kydliček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data scale, 2024. URL https: / /atpivavivv.////////aa////aaavv/ org/ a2tp///s////2// 17557
Qi Qi. Zhishuai Guo, Yi Xu, Rong Jin, and Tianbao Yang. An online method for a class of distributionally robust optimization with non-convex objectives. Advances in Neural Information Processing Systems, -000077--000,, 2021.
Qi Qian, Shenghuo Zhu, Jiasheng Tang, Rong Jin, Baigui Sun, and Hao Li. Robust optimization over multiple domains. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 4739-4746, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019,
Adityanarayanan Radhakrishnan, Mikhail Belkin, and Caroline Uhler. Overparameterized neural networks implement associative memory. Proceedings of the National Academy of Sciences, 117 (44):27162-271700 2020.
12

---

Published as a conference paper at ICLR 2025
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5551, 2020.
Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In International conference on machine learning, pp. 4334-4343. PMLR, 2018.
Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, pp. 3935-3971. PMLR, 2021.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and
Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of Red-
Pajama. httpps//cerrerra..a//lll/// and
ama -a -0---a--22660 1o-627b-tooee
of
June 2023.
URL https://
Sttps/ssimaa... ama- 627B
Daouda Sow, Sen Lin, Zhangyang Wang, and Yingbin Liang. Doubly robust instance-
adversarial training. arXiv preprint arXiv:2308 00311, 2023.
Megh Thakkar, Tolga Bolukbasi, Sriram Ganapathy, Shikhar Vashishth, Sarath Chandar, and Partha Talukdar. Self-influence guided data reweighting for language model pre-training. arXiv preprint arXiv:2311.00913 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy Usser-generad Text, pp. 94-106, Copenhagen, Denmark, September 2017. Association for Computational Linguisties.
Alexander Wettig, Aatmik Gupta, Saumya Malik, and Danqi Chen. QuRating: Selecting high-quality data for training language models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver. Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 022915-52971 PMLR, 21-27 Jul 2024. URL hnttts::/pprceet ngs press/v2333
www.wettti.e http:////tttgg22e..... h
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. arXiv preprint arXiv:2305. /0429, 2023.
Mingyang Yi, Lu Hou, Lifeng Shang, Xin Jiang, Qun Liu, and Zhi-Ming Ma. Reweighting augmented samples by minimizing the maximal expected loss. arXiv preprint arXiv:2103.08933. 2021.
Huimin Zeng, Chen Zhu, Tom Goldstein, and Furong Huang. Are adversarial examples created equal? learnable weighted minimax risk for robustness under non-uniform attacks. In Proceedings of the AAAI Conference Artificial Inielligence, volume 35, pp. 10815- 15-808 2021.
Jingfeng Zhang, Jianing Zhu, Gang Niu, Bo Han, Masashi Sugiyama, and Mohan Kankanhalli. Geometry-aware reweigheee adversarial training. arXiv preprint arXiv:2010.01736. 2020.
13

---

Published as a conference paper at ICLR 2025
Appendix
Experimental Details
15
A.1
Reference Implementation & Reproducibility
15
Training Evaluation Details
16
A.3
Additional Experimental Results
17
A.4
Sensitivity the value of
19
A.5
Further Discussions About our Reweighting Upper Bound
19
Additional Discussions
B.1
Toy Regression Problem Setup
21
Comparison Between Different Reweighting Methods
................................
B.3
Comparison with Robust Optimization Methods
22
Proofs
24
Full Gradient
24
Optimal Weight Strategy
26
Convex Minibatch SGD
27
Non-convex Minibatch SGD
30
14

---

Published as a conference paper at ICLR 2025
A
EXPERIMENTAL DETAILS
For full reproducibility of our work we provide insights on extending existing training routines with minimal-invasioo loss reweighting scheme.
A.1 REFERENCE IMPLEMENTATION REPRODUCIBILITY
To integrate our loss reweighting method with PyTorcc training routine, we have to add two plug-and-play components. First, we need to define four utility functions to ensure the proper functionality of our approach with reference implementations provided in Listing
1.
(...): Implements different reweighting strategies such as focusing on medium losses (quadratic), capping high losses (linupper), or emphasizing extreme values (extremes).
2. scale_ losses scale_losses ....): Applies exponential scaling to adjust the importance of losses based on a temperature parameter
:
3. normalize losses (. .) Normalizes losses bounded interval [-6, 6] for consis- tency in subsequent reweighting.
4. get_batcch_osee fross_fromm_o wro__looitttoooooooo .)..): Computes per-sample losses using cross- entropy, ensuring padding tokens do not affect the results.
These four functions then employed inside the training step. We need to capture the unprocessed raw losses and apply our reweighting scheme. We provide reference implementation for Py Torch FSDP training in Listing 2 to enable multi-GPU training.
For a full reference, we provide a ready-to-use implementation on GitHub for the GPT2 ex- periments: https://ttttt:tt http://ww....aaaocetoeeeeeeeaaaa ICLR-2025
Listing I: Utility functions
import torch
Exxonential scaling function def scale_lossed (lssses, r): return torch.exy (losses / r)
Normalization function for losses def normalize losses (losses, delta=l., min=0., 1_max=1.): return 2. delta losses / max(1 max 1 min, 1e-6) delta x (
m_x + 1_min) / max (1_mx max min, 1e-6)
Reweighting strategies def
(losses, delta=1.0, stratogy="linuppor"
if strategy "linupper";
14
return torch.minimun (losses + delta, delta * torch.ones like losses)
elif strategy "uniform":
We do not reweight here. This is our baseline.
17
return losses
18
elif strategy == "quadratic";
return losses**2 / deltax+2
20
elif strategy "extremes":
21
return torch.abs (losses)
22
else:
raise Not Jemenmneeeetterrrrr
25
Compute batch losses from logits def lot_batch_loss__ lts__osss__so_ogitteeeeee (logits, labels): ignore_ = -100 shi shift_logits logits[..., :-1, :] contiguous
26
27
label Tabels[..., 1:] cont iuous
15

---

Published as a conference paper at ICLR 2025
active (shift labels != ignore index) sum (din=1) 31 loss_fct crrchh..m.crrssaatrrre..... (reduction= onn'none)) none')
32 loss =
logits.size (-1)),
view (-1) .ong
return loss.view (logits. -1) sum (dim=1) / hum_active, nun_active
Listing 2: Integration reweighting strategies into a multi-GPU training loop.
import torch
2. import distribued as dist
Define reweighting function mapping STRATEGY "Linupper"
7 This has to placed inside a training step.
Assume model outputs (logits) and labels already computed device_losses, norms
11
trossatto_ossssttoooooooooooo (outputs, labels)
12
13
Initialize placeholder to store the losses from all GPUs
gathered_losses
(
dist.get_ sorrewwwwrdd.. sized)sszee))
len (device (ossess,
17
device=device lasieesseviie__osseesssoso wtw.sl..sss................ letpeeddevice_osssss.ooo
20
21
Gather lesses across all GPUs into tenor
22
losses,, loases_ detach ()
rsshhedulerrstepp (step, ofg.initial_2
25
Compute sample weights
26
with torch -ra.nooggrrd())
min loss = Lossess .item()
max_loss gatheree_lossee..mxx item
normalized_losses g (ataosssssssttteeeeeoosssssss
ottpr//ooooooooooooooooooooooooo view delta
-1., l_mio-min_loss, 1_max=max_loss)
30
rexeighted_losses- orgp_//sssss.pyyytttttteoooooooe losses, delta=l.,
strategy=STRAATEGG)
scaled losses scale scale_losses (reweighted_ reweighted losses.
scaled_losses
max ())m()))) 1=r)
32
weights = scased / scaled_
33
for instance local_ rank can be obtained with jast.getrrank device_wedt we__eeghtttt weights. view :]
-1) [local_rank,
36
Reweight losses and scale appropriately loss torch. sum (dweideeeeeeeeeeeeeeeeeeeeeeeeee device losses) loss. backward
dist. wllst.getwworld.ooo world size ()
A.2 TRAINING & EVALUATION DETAILS
In the following, we outline the hyper-parameter configuration and hardware requirements to fully
reproduce experiments.
Models. We employ three different GPT2 models to study the effectiveness of our reweighting approach on various model sizes. We report the architectural details in Table 6 Furthermore, to demonstrate the performance our reweighting approach on state-of of-the-r-t models, we employ a 1.4B parameter model with Laama architecture.
16

---

Published as a conference paper at ICLR 2025
Table 6: Model Architecture
Architecture
Nume
Parameters
Attenion Heads
Embedding Dimensions
Hödöen Dimensions
Mix. Sequence Length
124M
GPT-2
Small
210M
Medium
300M
36
24
785
3072
512
LLAMA
7B
1.4
7006M
1400M
24
32
32
16
4006
2018
2018 4006
8192
8192
Datasets. As a dataset for training our GPT2 models, we employ the SlimPajama-6i dataset, consisting of 7 domain partitions (ArXiv, Book, CC, C4, GitHub, StackExchange, and Wikipedia). The exact number of documents seen is reported in Table 7 We use the FineWeb 15T dataset to train the larger Llama- 1.4B and Llama-7B models. For the 1.4B and 7B models, 100B and 175B randomly sampled tokens from the Web 15T dataset, respectively.
approximately
Pretraining procedure. We pretrain all models with masked language modeling and do not apply any instruction tuning routine for subsequent evaluations. We set r to 100 for a specific number of warmup steps and then reduce it for the remaining steps. The GPT2 models are trained for 20,000 steps in total, and the Llama models are trained for 100,000 and 170,000 steps, respectively for the 1.4B and 7B model. More hyperparameters are provided in Table 7 For the GPT2 models, we use the Huggingface t trainer interface and the acce lerate library for distributed training. For the Llama models, we use fully sharded data parallel (FSDP).
Table 7: Training Hyperparameters for our benchmark evaluations
Architecture
Name
Minibatch Size
Learning Rate (1.R)
Weight Docsy
LR Schedule
LR
Warmup Steps
Max, Grad Nort
Training Steps
Total
End
Documents Scon
Mini
32
0.0005
0.01
Lincur Warmup Cosine
0.0001
500
1.0
20,000
640,000
Linear Warmup Cusine
0.0055
Linear Warmup Cosine
0.001
20,000
960,000
ILLAMA
7B
1:4
128
128
0.0033
0.0033
0.01
0.01
Linear Warmup Cusine
Linear Warmup Cosine
0.0003
0.0003
2,000
2,000
1.0
1.0
2.0
100,000
170,000
21,760.000
12,800.000
Benchmarks. We the LM Eval Harness generate our benchmark results. We evaluate the GPT2 models every 2000 steps and the Llama model every 5000 steps. For the GPT2 models, we find that only question-answentt results on LogiQA, LogiQA-2, SciQ, and PiQA. All other benchmarks yielded random results at this scale, i.e., they were not useful for model evaluation. The reason is the limited capacity of the GPT2 models compared to larger models like the Llama-1.4p or 7B models. For the Llama architecture models, we employ 13 language understanding and reasoning (LUR) and QA tasks to provide holistic picture.
(QA) tasks work well at this scale, which is why we report 5-shot
A.3 ADDITIONAL EXPERIMENTAL RESULTS
Benchmark results. Using our LinUpper method, the experiments with the 1.4B parameter Llama model show average performance gain of 0.66%, improving on 6 out of 13 benchmarks and tied on one benchmark (Table K). For QA benchmarks, we improved on 5/6 tasks with an average of 1.72% (Table 9..
We observe that our Linupper strategy also improves the performance of the 7B model. Our strategy
outperforms the baseline in 9/13 LUR (Table I0 and 5/6 QA tasks (Table [T). respectively. On average, across all tasks, our LinUpper strategy improves performance by 1.44% for LUR tasks and 1.16% for QA tasks.
Perplexity. We present additional perplexity visualizations of our GPT2 pretraining runs. Figure shows the results for GPT2-mini and Figure shows the results for GPT2-small.
17

---

Published as a conference paper at ICLR 2025
Table 8: Results of 1.4B parameter model on lan- guage understanding and reasoning benchmarks. Overall, our Lim Linupper method improves the results for 6 out of 13 tasks, and there is tie for one task, with an average performance gain of 0.66%.
Table 9: Results of 1.4B parameter model lan- guage question answering benchmarks. Overall, our LinUpper method improves results for out of 6 tasks with an average performance gain of 1.72%.
Benchmark Name
Unilomm
Lint pper (ouns)
Diflerence
Benchmark Name
Uniform
LinU/pper (ours)
Difference
Challenge
ARC Easy
63.80
63.38
-0.42
LogiQA
20.28
23.66
3.38
COPA
77.0
77.0
0.00
LogiQA2
26.02
27.48
1.46
Lambada (OpenA))
49.60
SciQ
91.30
91,90
0.60
Lambada (Stundard)
SocialIQA
MMLU
26.04
25.49
-0.54
TriviaQA (Exact Matchi
11.05
10.73
-0.52
MNLII
MNLI
31.91
36,92
5.01
Meon
42.35
44.07
1.72
RTE
51
52.35
1,08
SST-2
66.86
59.40
-7.45
TRn ARC
40.46
40.85
0.39
Thny WinnGrande
WinoGrande
56.75
59.12
2.37
48.51
49.16
0.66
Table 10: Results of the 7B parameter model on language understanding and reasoning bench- mark tasks. Overall, our LinUpper method improves results for 9 out of 13 tasks with an average performance gain of 1.44%.
Table 11: Results the 7B parameter model on question answering benchmark tasks. Overall, our Li Lipupper method improves results for out of 6 tasks with an average performance gain of 1.16%.
Benchmark k Name
Unilomm
Lint pper (ours)
Diflerence
Benchmarks Name
Uniform
LinUpper (ours)
Difference
ARC Challenge
41.81
39.42
-2.39
Bool()
73.52
74.40
0.89
Easy
LogiQA
COPA
83.00
84.00
1.00
LogiQA2
26.72
27,86
1.15
Lambuda (OpenAI)
61.28
62.43
1.14
SciO
95.20
95.00
-0.20
Lambada (Standurd)
58.08
80
0.72
SacialIQA
47,80
48.11
0.31
TriviaQA (Exact Match)
MNLI (mismatch)
MNII
32.25
32.35
34.24
34,01
1.97
1.66
Mear
48.15
49.31
1.16
69.38
Tiny TiyAARC ARC
36.17
52,80
-3.37
Tiny WinGGrande
49.09
44,20
4.88
Winofirande
65.43
67.32
Meon
50.54
51.98
1.44
5.0
Metnod
8.5
6.0
2up-000yy (bog)
Unifom
Purpoaity (bgg)
Pupp aaty (log)
8.0
4.0
Lintloper (surs))
8.0
6.6
0.5
(purs)
3.0
4.5
10000
(a) Arxiv Dataset
(b) Book Dataset
(c) C4 Dataset
(d) CC Datasel
5.0
66
Perpiatte
Perplodt
4.6
5.
5.8
2.8
2.0
6000
10000
16000
5300
10000
16000
6000
10003
19000
9top
Stsp
Stop
(e) GitHub Dataset
StackExchange Dataset
(g) Wikipedia Dataset
Figure 3: Per-domain perplexities on hold-out validation sets under the uniform domain sampling setting for the GPT2-mini model. Our reweighting strategy LinUpper strategy achieves better or at least comparable perplexity out of 7 domains.
18

---

Published as a conference paper at ICLR 2025
0.5
Method
Puppletty fool.
Undomm
Parploaty (log)
Purphelty (oa))
Parpletly (es))
4.0
Linloper (sars)
Quadruti
100
(ours)
3.0
4.5
4.0
12000
(a) Arxiv Dataset
(b) Book Dataset
(c) C4 Dataset
(d) CC Datasel
4.0
5.0
Perplotty (000
Perpletty
4.0
2.5
8.5
k.0
35
2.0
2.5
3.0
8000
10000 Stap
19000
6000
10000 Step
16000
6000
10003
19000
Step
(e) GitHub Dataset
(f) StackExchange Dataset
g Wikipedia Dataset
Figure 4: Per-domain perplexities on hold-out validation sets under the uniform domain sampling setting for the GPT2-small model. Our reweighting strategy LinUpper strategy achieves better or least comparable perplexity 6 out of domains.
A.4 SENSITIVITY TO THE VALUE OF
We conduct experiments using different values of r to understand the sensitivity of our methods to the value of The perplexity plots for method LinUpper with different values of provided in Figure which shows that when the value of r is large (e.g. r= 1) then, as expected, the performance of our method becomes closer that of the uniform baseline. Decreasing the value of leads to the diminished effect of low-loss samples, but this can also have negative effect on the performance when r small (e.g. 0.2), potentially due data wastage by over filtering the low-loss samples.
anxivlog perplexity
perplesity
obilog
ccilog, perplexity
Nothos
5:32
Rethos
Mothes
5.25
Irapeetr rw
Inusia_
Inuspa_jecs
3.54
Irupper
Irupper
Irupper r-0.8
F4.75
5.00
Irupost -0
Iriporr 044
trupost.m00 m0
Inupser
Irupser
Inupser
Inusser
Tnusser
Inusser
230
busch::
4.5
3000
10030
15000
5002
10030
15550
3003
10003
15000
5003
10300
13500
olthublkog perplenity
becchhhhangggoggg perplexity
wikipedia lng perplexity
Methad
Nethod Inuspe
Nethid
inuter_ro03
Knupper
berplosky
3.0
Inuster
Inusper
inuserr
inuspe 1w2.6
Frupper
inuster
Snusper_rwa 1w0.4
Trupperrea
tanelie
taselne
S3.3
baselin
3.0
Staps
Stops
Shaps
Figure 5: Sensitivity of our method LinUpper to the value of r. When the is large (e.g. r 1) the performance of our method becomes closer to that of the uniform baseline. Decreasing the value of leads to diminished effect of low-loss samples, but can also have a negative effect on the performance when is too small (eg. 0.2)
A.5 FURTHER DISCUSSIONS ABOUT OUR REWEIGHTING UPPER BOUND
To demonstrate that in practice the upper bound of w; < 2 is, in fact, satisfied by our LinUpper method, we provide the distribution of the maximum weights per step (see Figure (6) during the pretraining the 7B parameter model with our LinUpper method. this experiment, the minibatch size is 128 and hence the upper bound is 128 0.0156. As depicted in Figure 6 the maximum
19

---

Published as a conference paper at ICLR 2025
weight per step (which is around 0.01) by our LinUpper method is clearly below the upper bound of 0.0156.
2000
Density
1500
Upper Bound per Our Definition:
2
0.0156
128
1000
500
0
0.008
0.010
0.012
0.014 Max. Weight Value W)
0.016
0.018
Figure 6: Distribution of maximum weight values when training 7B parameter model with Llama architecture. The maximum weight values always stay well below the theoretical upper bound defined in Theorem 2
20

---

Published as a conference paper at ICLR 2025
B ADDITIONAL DISCUSSIONS
B.1 TOY REGRESSION PROBLEM SETUP
To further illustrate the effectiveness of our reweighting strategies simple ML problems, we apply them to regression problem where the goal is to learn a linear mapping from input features to output labels, while dealing with noisy data and outliers. We generate synthetic dataset follows.
The clean imput features Xelemn E Rxx are sampled from the standard normal distribution, where the dimensionality of the data, n is the number of clean data points. The target output y E RR" for the clean data is generated using the following ground truth linear model with added noise:
Velcan = Xolean W + 1 c.
where W E R and b" R are the true model parameters, c e~ No I,) represents Gaussian noise, and c is small constant controlling the noise level. For the m outlier data, we generate each input features of X mxxx as 0.1 (O, 1) + with the corresponding target outputs yood drawn randomly from standard normal distribution, i.e., Yood (0,, 1).
The full training set consists of both the clean data (Xalean Yelean) and the outlier data (Xood Y Yood); concatenated together:
X all Xeleam Xmnd]; Yall
Yelean; y Yood
By incorporating outliers into the training set, we test the robustness of our reweighting strategies, which should down-welt weight the influence of noisy or irrelevant samples, allowing the model to learn the underlying clean data distribution more effectively.
We define the objective function for this regression task the mean squared error between the model predictions and the target outputs. Given model parameters W E R2 and b E R, the loss for a single data point (x,, y:) computed as:
1 [ W + - bb));)² b) 2
(W.
(11)
To evaluate model performance, compute the mean squared error a held-out test set (X Yrea ) which was generated using the same process as the clean training data. We use p = 64, n = 3200, and 800.
B.2 COMPARISON BETWEEN DIFFERENT REWEIGHTING METHODS
The results from Figure cleerrly show that Li LinUpper outperforms all other reweighting strategies and the uniform baseline, achieving the fastest convergence. This supports our theoretical findings that down- weighting low-loss samples allows the model to focus more informative examples, thus accelerating the optimization. Interestingly, the Quadratir ic strategy also performs better than the uniform baseline, indicating that selectively up-wet ppeieighting the moderate loss samples improve performance. While Quadratie converges more slowly than LinUpper, it still highlights that more balanced reweighting provide benefits uniform sampling when dealing with noisy data.
21

---

Published as a conference paper at ICLR 2025
x10
x10
LinUpper
12
Quadratic
Extremes
10
Uniform
eva 1055
0.8
0.6
1
qualiakk [m)
otteres
0.2
0.5
1.0
1.5
2.0
2.5
3.0
-1.0
step
x10
Figure 7: Left: Test loss training steps for the toy regression problem. Our LinUpper strategy achieves faster convergence compared to all the other methods. Notably, the Quadrat ic strategy also outperforms the uniform baseline, further highlighting the benefit of selective reweighting. Right: Shape of the different reweighting methods used with r 1.
B.3 COMPARISON WITH ROBUST OPTIMIZATION METHODS
Comparing with traditional robust optimization methods, such as those based on distributionally
robust optimization (DRO), our methods impose cap sample weights (w, < 2/M), justified by our theoretical analysis (Proposition []]. This cap prevents over-reliance on a small subset of high-loss samples, which lead to overfitting outliers. In contrast, DRO methods often focus heavily worst-case scenarios, potentially leading to suboptimal generalization when the training data contains substantial amount of noisy samples. To demonstrate this benefit of our method in practice, conducted experiments to compare with the traditional KL-divergence regularized distributionally robust optimization (DRO-KL).
First, we created synthetic regression problem (similar to the regression problem described above) with 25% of the data being outliers compare the robustness of the different methods against outliers. As shown in Figure g DRC-KL diverges when we use the same learning rate as the other two methods (our sUnupeer method & the Un form baseline). The DRO-KL method requires smaller learning rate to converge, which case it converges slower than our method. Furthermore, we include the DRO KL method among the compared baselines for the GPT2-medium experiments, as shown in Table On average, our Li Linüpper strategy outperforms other DRO-K KL configurations, which we attribute to its ability to handle outliers.
we note the following: because it focuses heavily on the hard samples,
x101
LinUpper (ours)
5
DRO-KL (same LR)
DRO-KL (reduced LR)
Uniform
loss
3
eval
2
0.5 step
1.0
x10²
Figure 8: Regression problem with DRO-KL.
22

---

Published as a conference paper at ICLR 2025
Table 12: Comparison of our LinUpper loss reweighting strategy with various DRO-KL configura- tions on the GPT2- experimental setup. On average, our LinUpper strategy outperforms the other DRO-KL configurations.
Uniform
LinUpper
LogiQA
25.7
27.9
23.2
25.5
26.8
25.9
26.3
LogiQA2
27.5
27.6
24.4
28.7
27.5
28.2
27.6
SciQ
PIQA
55.3
56.2
54.0
56.2
55.7
56.5
53.4
Menn
39.4
40.9
35.7
39,8
29.9
40.3
39.7
23

---

Published as a conference paper at ICLR 2025
C
PROOFS
Lemma 1, Let Assumptions
1 and
hold. Then we have
ST ( [V(0)=
-(0)) 07))
Proof. Consider the following
f(01) - f(02) = f(0,) f(0) + f(0) f(02).
(12)
From the L-smoothness of f. we have for all 82, 0 E R
L
f(0) f(02) - (Vf( 0 02)
2 110 0211
(13)
From the the convexity of we have for all 0,, 0 E RRd
f(0:) (0) < (V/( - 0)-
(14)
Substituting Eq.
13
Eq.
14
in Eq.
we get
L
f(0)) - f(0z) & (Vf(0:). 01 - 0) + (Vf((2), -02) + 2
lle - 0211²
(15)
Let 0 = 82 (V.(())) V f(t))) Substituting the definition of 0 in Eq.
15 we obtain
1
f(0)) - ((02) & (Vf(0)). 9, 0g + -
(7((0), ()())))
+ 2T VS((2) (
(VI(0)), 01 02) 1
T, [ V/(001 + 1 VS(02) <S(01)]
1
2T.
(VS(0)), 01 - 02) 2L
Vf(02) 3.
(16)
From lemmal) 1
we obtain the optimum l"
|V - V 22.((00) 221000 S(0"))
C.I FULL GRADIENT
We have
0012 ||0ct1 0-112= = =110- -0'12 -1100-1 11100-1-0011 2 + 2(0011 00,00 8*)
M
M
= =110- 0¹1l + 1l Σ
Wis ( -2 11t
Wi,t f. ).00 -0"
1=1
i=1
As
M
< 1100 0-11 + ni M Σ
w
2np Σ
Wi, (VI((0"), 0² - -00 (17)
Using the convexity of function fú, we have
V.f(("", 88)) 2 f(((t) f((8)).
(18)
Therefore combining, we obtain
|0++1 - 0-112= < 110- 0-112 + 17 M
2 2 2n Σ wie (((00)){(00)) (19)
24

---

Published as a conference paper at ICLR 2025
Next, we upper-bound the quantity
||
< |
2 <
-V + +
<L (fi(0') fi(0)) +20
(20)
where Eq.
20 follows from lemma I
Now, combining Eq. 19 and Eq.
20
we obtain
M
M
||0+| - 0011 - 0 112
+ i i
Σ
Wo (1:(0) - Si(0')) + 2
(21)
2T
Wiçt (Fil0') fc(00))
i=1
M
= =100- 1-1
(4M Lm - 2) newee (fi(0") f((0))
M
+ 2M 2 0
whtt:
(22)
Selecting vi such that
-2 2 - Vi, Vi, i.e. n1 S LM yields
M
1 - 001 S 100 1000-0'1' TT Σ
whe - + 2M
(23)
where we also used the fact that EMe w², S due to
1i=1 Wift and waft 2 0. Hence,
1gt+1
1-1
< 0- 1-11
Th
(f((0)- - f((0))
Σ
1
wis)) (f((0)-f((0')) - f((0)) + 2M
7/6
M 1=1
M
i=1
M
= =110* 8r
Tie (f(t*) f(0r)) + me
wo.1 (f((0) - f((01) 2M
M
(24)
Telescoping Eq. (2 from t = t T 1 yields
T-1
M
100 - -011 s 1000 - 0-1/
- Te (
+
7t
Σ i=1
wi,t)
- (1((0)-1(00")
M
t=
T-1
+ 2M
n
(25)
For a weighting scheme such that U <
2
so that can be fixed to 87. we obtain
1.
37
T-1
(S(0') - r(0)) < 10 - 0-1 +1
1
wss))
(1:(00)):(0)) -i(00)) + 2M..n?
11
M
1=0
t=
f(0T)-f(07 < 8T 10 1-1
M wist (1:10)) (009)-f(00 f((00) + Mo?
1
T
4L
1=0 ==1
where ØT = Σt= 0r Define 8 = M
-M
(1 wr.1) (fi(0)) - -f((0)). Hence, we obtain
80
8+ ||
T-1
Mo?
,(0") 1(00) <
+
Σ
8t +
(26)
T
T
4L
z=
25

---

Published as a conference paper at ICLR 2025
C.2 OPTIMAL WEIGHT STRATEGY
Proposition 1. The optimal strategy thut minimizes the KL-divergence regularized 8', i.e, 8 + M
Σ i== W/ M (og(Mss), is given by
h;
2 M f
W: min
exp
(5)
where C' is a normalizing constant that ensures Σ; Wig 1.
Proof. For the sake of clarity, we drop the index t in 8' in the following proof.
Let A, fxx;0)) f(x;; 07) be the loss gap. After dropping all terms that do not depends on the weights wi 8, we have the following constrained optimization problem
M
M
min
W A
+
Wc log
(27)
1=1
2
subject to
Wy S
1 [M]
(28)
M
M
W = 1.
(29)
As we have inequality constraints in the above problem, we introduce the slack variable and consider the Lagrangian function to solve the above problem as follows:
M
M
M -
2
L(WI,... WM, p) =
A
W log
V/
Σ
wi
+
M
1=1
1=1
1=1
(30)
and we have the gradients
0L.
A + r
log W/
P
Ow,
0L
Σ Wi
01
aL
2 M
Wi
+ 12
(31)
0p
We consider the complementary slackness condition. We have one slack variable i, and the corre- sponding Lagrange multiplier is v. We now consider whether a slack variable is zero (which the corresponding incquality constraint is active) or the Lagrange multiplier is zero (the constraint is inactive), we obtain
/ 0,12 > 0: Using at
0, O, we get A + 1 log Wi 0. Hence, we have w; exp (u+A Therefore, obtain
de 0, we get AM w; Using BL Jhc.
-r (MM M exp (557) yielding w; exp
- In (IM exp( (00 ()
Cexp (مه where C exp
In EM exp (Arrr) 1).
n + 1== Using 3y 0, 1 = 0, we get w; = 2 and μ + v A, + + 3.
Hence, obtain w; min (Cexp (~:), 2
26

---

Published as a conference paper at ICLR 2025
C.3 CONVEX MINIBATCH SGD
Theorem 3 (Re-statement of The Let Assumptions 7 size [8] b with reweighting scheme satisfying maxies Wi,t 2/b. Then,
and 3 hold. Consider a minibatch of
minibatch SGD: for rj 8I.TT. - we have
8L 8
1-
T-1
B[ - s(6))]
6,
(32)
VT
T
t=0
minibatch SGD with momentum: for 8LVT41
1
de > 0, we have
Bif(0") f(0))) 8L 00 1-1
T-1
2
8,
Aeples
(33)
VT+I
T+1
T+1
where 6, = 12 and HTh 7 ET-1 gr.
icB G ws.)) (Si(0') - 14= [(0))))= [[EE G wis) (1((09)-:(00-))))
Proof. (i) Minibatch SGD: For the minibatch SGD based on Eq. 2 we have
||0tt1 -011 = |10º 0-112 + | 0² + 20 - 0',00' 0²)
= =1000= 0- 11 + 1m Σ wii (10)0') -2 1ht iEB
Σ
Wiit V vi((0', 0" (34)
iEB
Applying expectation both sides, get
E 0++
11
||0r
8"
1|2
+ n7
Wi, (0
-
Σ
Witt V)((0)) 8*
SEB
WEB
(4)
$ 100 0r 1|2 + n
E
Wre f(('))11
2 Σ
Wie (f:(0') - f((0)))
iEB
iEB
< 1100 0-1
+ n2bB
Σ 1 V
27 L
Σ
Wie (Ji(0) - S&(0%))
EE8
110-
0-1 + 4Lm bE
Σ
W (F)(0)) 5.(005))
27,B
Σ
Wiit (f:(0) fi(00))
Lie
iEB
$1100 1- +E
wie (LLm (n))(f -(((0)))
(35)
LiEB
where (a) follows from Assumption
I and (b) follows from lemma
Selecting ne such that
4bLapwi. 2 < -1 v i, i.e., the abLwaas. yields
Bl|0ct1 0-112 - 1100 0-112 T/
- -
LREB
1100
-0²11 1/.1E Σ 1 - -1(00))
1
(1:(0)) -1((0)))
+ n.E
b
Wiet
LiB
27

---

Published as a conference paper at ICLR 2025
110 0-1 11([(00)) - -((0)) + 1/.de: where 5, [ Ligs ( wi.t) (fi(0') - [(())))] and Eies wf, S Hence, telescoping Eq. (36)
(36)
from -- to T 1 yields
T-1 (f(0)) f(00)) + Σ
E E| 0 0-
112 < 80 0-1
Σ t=
mod
(37)
1=
For weighting scheme such that Wmax,t
z so that 7/ can be fixed to 11
we obtain
8LVT
8L 8
101
T-1
-(0")- 5(0") 5
6,,
(38)
VT
T
t=0
where 0- = - Σ 0¹
The traditional convergence bound of minibatch SGD for convex functions can be recovered when ó, = which occurs when the samples in the batch have uniform weights, Wi,t = $ Vi E B. We obtain a looser convergence bound when 6, 2 O resulting from the assignment of larger weights to smaller loss gaps. Moreover, deighteighting the importance of samples with low loss gaps leads to better bound (6, S 0). Note that the bound in Eq. (32) only holds for reweighting schemes that satisfy Wi,t I 2/b. Thus, enforcing an upper bound on the maximum weight that can be assigned to single data point. Typically, the minibatch gradient with uniform weights is unbiased, [Eie [0') = ((('). However, in our setting, the importance weights are functions of the sample functions. Hence, E [E>EB Wi,s Vi +f VS((') leading to the analysis of Theorem being more involved. Based on our analysis, we note that the step size The is a function of the maximum value of the weights, Wmax,t F which leads to the choice of v) SLVT
1 for convergence.
(ii) Minibatch SGD with Momentum: Given the loss-based sample reweighting strategies, Eq. 33 characterizes the effects of loss-based reweighting on the convergence bound of reweighted minibatch SGD with momentum when the sample losses ft are convex. Given the loss-based sample reweighting strategies, we incorporate heavy ball momentum method to improve the convergence result and to characterize the effects of loss- reweighting on the convergence bound. Hence, we have the following update equations:
2+11
2t-
17: Σ WE f(())..
(39)
Ae+1
0,11
02
2+11
(40)
I+ Aet1
1 A+t1
where 7) is the step size and me A > O. The above iterates 0, are equal to the iterates of the following stochastic minibatch heavy ball momentum method (Sebbouh et al. (202T).
0 = 0:
1/K
wiS (.,,) + I Ae+1 - 00,1).
A
1+A+1
I+
TEB
Eq. (33) characterizes the effects loss-based reweighting on the convergence bound of reweighted minibatch SGD with momentum when the sample losses Ji are convex.
For the minibatch SGD with momentum based on Eq. 440). we have
|| 1-1 ||ad 8²1|2 + +||22t1 2012 (2
zt 2t 8t)
= ||2 8¹1 + 11 Σ Wie (0
-2 1t Σ
Wi, (i("),, -8" ) (41)
EEB
iEB
Applying expectation both sides, we get
E
1-+++--11 0-1
Wi.t V.((()),) 8- }
|| 8-1 T E Σ Wi.t (
2
Σ
VEB
28

---

Published as a conference paper at ICLR 2025
-
- A
Σ NEB
Wint Vi((0"),
(0)
|| 0-1 112 + T
Σ
wi V.:(('0]]]
21
Σ
Wi,e (J:10") - 11(03))
BE
.E8
2
Σ
win (J:(0") 11(0-11))
LIEB
(6)
-
|| 0-1 + E
Σ
(AL. ( -2n)(((()=)(000) -1((0))
LiEB
- A,
Σ Whi (fi(0") - f((0-11)
LiEB
= =I1= 0r1 1-1 +E Σ
wi wi,e(4Ln),buee - 21 ()))))((((0'>))))))))))))))))))) - -s((0))
LiEB
+ ArE
Σ
win (1i(0--1)-1) 1i(00)) -
LiEB
(0)
S || 0r
241 13
Σ
wi, -((0))))))((1))
21%
Σ
Wie
(1:(0-1)) (((03))
LiEB
LiEB
21, 1 = =||2* 1²1
Σ (1i00)) (
+ +2
Σ -w7,)) (f((0') -
- 5.(or))
LiBB
LiES
b (fi(0'-)) - 11(07))
Wiit - b)) (fi(0-1)) -0--1)--()))) 5((07))
+ 2m
+ +2MAAE
LiE
LiEB
= =||= 0-112- 2 /(0')) + 2 Ae
Σ
② -wss) (1:00) - ((0O7)) -1(0))))00)))))
1
LiEB
1oe - - (fi(0--1) - (((0-))-(((00))
+ 2n2 (f(0--1)- - -((0)) + ArE
LieB
-21, = =||2- - -0-12-- - A f(0))) + Me
b -w01s) (f((0))-00) S((07))
LiEB
+ An (f(0"-1)- f((*)) + ArE
Σ LiB
wis
(f((0)) - -
(42)
5
where (a) follows from applying Assumption to to the third and fourth terms, (b) follows from applying
Lemma to the second term, and (c) follows from setting 7t as 1/: S aLbroma, and A = A + 3
Let & Eies ( w...) (0') f(()))) pe = Eies (t w..) (F,(0)) - -f(((-
and Eies whe S 1. Hence, telescoping Eq. 42 from t to T and noting that 2g To yields
T-1
E E||=T 0-112 < 1100 0-11 2n E (f(O)) f(0))
Σ
2np5
Σ
2mpaj44
(43)
With A = t/2, for a weighting scheme such that Wmax, < we obtain
2 so that 7/ can be fixed to 8LVT+1'
BIS(0") f(0') 8T | 00 0-113
T-1
T 1
2
2
Aaple
(44)
VT+I
T+1
29

---

Published as a conference paper at ICLR 2025
C.4 NON-CONVEX MINIBATCH SGD
Unlike the convex case, the importance weight Wi, is set based on the gradient norm instead of the loss gap f((0") f((0*) when the sample functions fi are non-convex. Theorem underlines the influence of these importance weights on the convergence bound for the non-convex case. Note that computing the gradient norm requires extra computations, whereas the loss values are readily available in practice.
Non-convex:
Assumption 4. The norm of the difference between the sample gradient V and full gradient f0,) is bounded by ((..) f s V2, v0,, E [M]
Theorem 4. [Non-convex Minibatch SGD] Given that the objective functions fr, E [m] satisfy
Assumption 2 and Assumption 4 Consider minibatch of data points |B| b with reweighting
scheme Wits if E B and 0 otherwise where maxigB Wi, S 2/b. Then, for 1| SLVT we have
16L(100)) M
6V2
E {
R
Wit
1VS(0))1
+ 4VT
T
VT
b
1=
(45)
Proof. Applying L-smoothness condition
s(0(t))
L -10 -011
$((0) + (V 10+- - -0) + + 2
Ln?
-(00) - (S(())) Σ
Wie Vi((O))
Σ
V((00)
2
tEB
1EB
Lap
-(0'' (((0)) Σ b
Vii(0))
+1 VS(P',
Σ b wse V.i((9)
+ 2
Σ Wie Fi(A
NEB
= -011 -
(VI(0', Vf((09) + Lny²
((()') Σ b 5.(00)
+1
T - 1im
Σ
Whe V.((09)
2
iEB
iEB
BEB
s s((0)
V((0'),
v v((0')
+ 2n
b Wist
|V + 2n Σ
b Who
|
EEB
EB
EE
La7
Σ iEB
Wa (0t)
2
+ Tnp
= =((0
(((At),
b S(((0))
+ 2m
Σ
Wit
|
Σ Wa
b
2
E
iEB
S (011 7 V (O')),
Lbnp
b
V (0')
2n Σ
b
Wit
|
+ 2
W 2
1EB
EBB
EEB
5110') 1
( Σ VI(0%)
2n
b
Wit
|
SB
L.bb Σ wh (1'00 ||
Lbn²
w
\
VEB
5S(00) - 1/ ( V5(0"),
Σ 1 VI((0))
wit)
|
2n
iEB
SEB
+ L2 + Lbn2" Σ 3 |
15B
30

---

Published as a conference paper at ICLR 2025
(46)
where we have used the bound ((*) tion on both sides, we get
v = 12 and Eies wf, S 1. Applying expecta-
Ef((++1) - 7/
[ (2n) 18
Σ T wit |V
+ Lbn) Lbmj"
2
S $1(0") - Ba V (2n) E
Σ
wis
|
b
LEB
+ Lbn Lb 17 w mux. E V
?
S s((0') E |
2ny
Σ
1
wa
||
+
L2
(47)
LEB
where 1) S 21.6 wiss
in the last inequality.
Telescoping and rearranging the terms yields
2(f(6)) 5
E
wit
2Lmb
T t=0
1/T
t=U
LiB
(48)
For a weighting scheme such that Wmax,t S so that Nt can be fixed to to
BLVT
1
we get
T-1
16L()(0)) - f*)
6V2
Σ
wio
+
VT
4VT
1=0
LiB8
(49)
From Theorem 4 we can obtain the convergence bound of traditional minibatch SGD for non- convex functions when wi, = F Vi B. Moreover, looser convergence bound is obtained when E |Eies (t= wa) V (1 0, for Wiy S t.
31


Reviews:
[{"review_id": "AgbAYSuPnB", "reviewer": "Reviewer_u8Tc", "summary": "This paper introduces a dynamic, instance-level data reweighting approach aimed at enhancing the efficiency and effectiveness of large language model (LLM) pretraining. Traditional LLM training paradigms apply uniform sampling to datasets, ignoring individual sample importance. Current reweighting strategies are mostly group-level and lack the adaptability required for instance-specific optimization during training. The authors propose a novel approach that reweights samples dynamically based on their loss values, helping the model to focus on more informative samples. They also provide a new theoretical framework for analyzing how loss-based reweighting impacts convergence rates, demonstrating that deprioritizing low-loss samples leads to faster convergence and improved performance. Empirical results confirm the effectiveness of this approach across various tasks, from large-scale LLM pretraining to smaller problems like linear regression.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Innovative Instance-Level Reweighting Mechanism: The paper introduces a fine-grained, instance-level reweighting strategy, which addresses the limitations of traditional group-level reweighting methods. \n\n2. Theoretical Framework for Convergence Analysis: The authors provide a novel theoretical framework that explicitly characterizes the effect of loss-based reweighting on convergence bounds. This contribution not only supports the empirical findings but also adds a theoretical foundation that enhances the credibility of the proposed approach.", "weaknesses": "1. 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.\n\n2. 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.\n\n3. 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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jhf8ng5wh5", "reviewer": "Reviewer_AvfZ", "summary": "The paper introduces novel algorithms for dynamic, instance-level data reweighting aimed at improving the efficiency and effectiveness of large language model (LLM) pretraining. The proposed methods adjust the weight of each training sample based on its loss value in an online fashion, allowing the model to dynamically focus on more informative or important samples during different training stages. The paper also provides a theoretical framework for analyzing the impact of these reweighting strategies on the convergence of gradient-based optimization.", "questions": "1. 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.\n﻿\n2. 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.\n\n3. 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?\n﻿\n4. 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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "(1) The paper introduces an innovative and practical methodology for dynamically reweighting training samples based on their loss values, which could significantly enhance the efficiency and effectiveness of LLM pretraining.\n\n(2) The empirical findings are robust, demonstrating that the proposed reweighting strategies lead to faster convergence and improved performance across a variety of tasks and model scales.\n\n(3) The proposed algorithms can be seamlessly integrated into existing training pipelines with little to no computational overhead, making them highly practical for real-world applications.", "weaknesses": "(1) The paper might benefit from additional implementation details to enable reproducibility of the study, especially for the proposed reweighting algorithms.\n\n(2) While the evaluation is comprehensive, further ablation studies could strengthen the paper by isolating the impact of different components of the proposed methods.\n\n(3) It would be helpful to have more detailed discussions comparing the proposed methods with other widely-known reweighting or data selection techniques beyond those already mentioned.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DKDbhOqcOW", "reviewer": "Reviewer_9BZv", "summary": "The paper presents a framework for dynamically reweighting training samples in large language model (LLM) pretraining based on loss values. This approach contrasts with traditional uniform sampling by adjusting sample importance during training, allowing the model to focus on more informative data and deprioritize redundant samples. The authors introduce three reweighting strategies, each designed to emphasize specific portions of the loss distribution, and establish a theoretical foundation for their effectiveness on model convergence. Experimental results show that these methods lead to faster convergence and improved performance on various benchmarks, especially when combined with existing domain-level reweighting methods like DoGE and DoReMi.", "questions": "1. **Benchmark Justification**: Could you elaborate on the selection of LogiQA, LogiQA 2, SciQA, and PiQA as benchmarks? Testing additional benchmarks, such as MMLU, ARC, hellaswag, boolq, and TruthfulQA, could further strengthen the few-shot performance evaluation.\n\n2. **Baseline Comparisons**: 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?\n\n3. **Justification for $2/M$**: The theoretical assumption cap on maximum weight at $2/M$ seems restrictive. Could you provide further justification or discuss how this assumption holds in practical applications?\n\n4. **Efficiency Discussion**: 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Originality: The paper introduces a novel approach to LLM pretraining by leveraging instance-level, dynamic reweighting based on loss values. Also, the theoretical results justify the methods.\n\n- Quality: The paper combines theoretical analysis with empirical results across multiple benchmarks. The authors provide convergence bounds and use diverse model scales to demonstrate the generalization and scalability.\n\n- Clarity: The methodology and notations are well-defined, with clear mathematical formulation and supporting visualizations.\n\n- Significance: The paper’s instance-level reweighting method enhances training efficiency and model performance. Its compatibility with existing domain-level techniques (e.g., DoGE, DoReMi) highlights its practical value", "weaknesses": "1. **Clarity**: The statement “where $L>0$ is the smoothness parameter” under Equation (7) is unclear since $L$ does not show up in Equation (7) and could benefit from further explanation.\n2. **Benchmark Choice**: The choice of benchmarks (LogiQA, LogiQA 2, SciQA, PiQA) is not well-justified. Testing on additional benchmarks like MMLU, ARC, hellaswag, boolq and TruthfulQA could strengthen the few-shot evaluation.\n3. **Baseline Comparisons**: Table 1 lacks a direct comparison with baseline methods (DoGE and DoReMi), which would provide valuable context for the effectiveness of the proposed approach.\n4. **Theoretical Assumptions**: The assumption that the maximum weight is capped at $2/M$ may be too strict, particularly in data selection scenarios with fewer than half the dataset selected. Justification for this assumption's practicality is needed.\n5. **Efficiency Discussion**: The paper could better address the efficiency benefits of the proposed methods compared to DoGE and DoReMi, especially regarding computation and implementation costs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)