Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures
Authors: Dang Nguyen
Wenhan Yang
Rathul Anand
Yu Yang
Baharan Mirzasoleiman
Venue: ICLR 2025 Poster (2025)
Keywords: Data selection
Memory efficient
Large Language Models
Data Mixtures
Primary area: foundation or frontier models, including LLMs
TL;DR: We select small and high-quality mini-batches to reduce the memory consumption while boosting the performance of training LLMs.

Abstract: Training with larger mini-batches improves the convergence rate and can yield superior performance. However, training with large mini-batches becomes prohibitive for Large Language Models (LLMs), due to the large GPU memory requirement. To address this problem, an effective approach is finding small mini-batch coresets that closely match the gradient of larger mini-batches. However, this approach becomes infeasible and ineffective for LLMs, due to the highly imbalanced mixture of sources in language data, use of the Adam optimizer, and the very large gradient dimensionality of LLMs. In this work, we address the above challenges by proposing *Coresets for Training LLMs* (CoLM). First, we show that mini-batch coresets found by gradient matching do not contain representative examples of the small sources w.h.p., and thus including all examples of the small sources in the mini-batch coresets is crucial for optimal performance. Second, we normalize the gradients by their historical exponential to find mini-batch coresets for training with Adam. Finally, we leverage zeroth-order methods to find smooth gradient of the last *V*-projection matrix and sparsify it to keep the dimensions with the largest normalized gradient magnitude. We apply CoLM to fine-tuning Phi-2, Phi-3, Zephyr, and Llama-3 models with LoRA on MathInstruct and SuperGLUE benchmark. Remarkably, CoLM reduces the memory requirement of fine-tuning by 2x and even outperforms training with 4x larger mini-batches. Moreover, CoLM seamlessly integrates with existing memory-efficient training methods like LoRA, further reducing the memory requirements of training LLMs.

Published as a conference paper at ICLR 2025
MINI-BATCH CORESETS FOR MEMORY-EFFICIENT LANGUAGE MODEL TRAINING ON DATA MIXTURES
Dang Nguyen
Wenhan Yang Rathul Anand
Yu Yang
Babaran Mirzasoleiman
{dangnth, hangeryang 18, rathul, yuyang, baharan @cs.ucla.edu Computer Science Department, UCLA
ABSTRACT
Training with larger mini-batches improves the convergence rate and can yield superior performance. However, training with large mini-batches becomes pro- hibitive for Large Language Models (LLMs), due the large GPU memory re- quirement. To address this problem, an effective approach is finding small mini- batch coresets that closely match the gradient of larger mini-batches. However, this approach becomes infeasible and ineffective for LLMs, due to the highly im- balanced mixture of sources in language data, use of the Adam optimizer, and the very large gradient dimensionality of LLMs. In this work. we address the above challenges by proposing Coresets for Training L.LMs (CoLM). First, we show that coresets found by gradient matching do not contain representative ex- amples of the small sources w.h.p., and thus including all examples of the small sources in the mini-batch coresets is crucial for optimal performance. Second, we normalize the gradients by their historical exponential to find mini-batch core- sets for training with Adam. Finally, we leverage zeroth-order methods to find smooth gradient of the last V-projection matrix and sparsify it to keep the dimen- sions with the largest normalized gradient magnitude. We apply CoLM fine- tuning Phi-2, Phi-3, Zephyr, and Llama-3 models with LoRA on MathInstruct and SuperGLUE benchmark. Remarkably, CoLM reduces the memory requirement of fine-tuning by 2x and even outperforms training with 4x larger mini-batches. Moreover, CoLM seamlessly integrates with existing memory-efficient training methods like LoRA, further reducing the memory requirements of training LLMs.
INTRODUCTION
Large Language Models (LLMs) have achieved remarkable success in variety of tasks, ranging from machine translation to conversational AL. However, pretraining and fine-tuning LLMs with billions of parameters requires a large amount of compute and GPU memory, not only store the parameters but also to compute gradients and optimizer states (e.g., momentum and historical gradients in Adam). For example, full finetuning a relatively small LLM, such Phi-2 with 2.7B parameters, using a batch size of 128 requires at least 44 GB of GPU memory. The large memory requirement makes prohibitive to train such models with larger batch sizes, which effectively improves the convergence and can improve performance. This raises a key question: can we train LLMS with smaller mini-batches and get the benefits of training with larger batch sizes?
To address this problem, many memory efficiien techniques have been recently proposed, mainly enable efficient fine-tuning of pretrained language models. a high level, such methods aim to find smaller set of parameters (Adelman et al., 2021), or find low-rank (Hu et al., 2021; Zhao et al., 2024b) or quantized (Dettmers et al., 2022) weights or optimizer states to train the model in memory-efficice manner. There have also been efforts to adapt gradient-free optimization for training LLMs (Malladi et al., 2023). Yet, most memory-efficient techniques cannot achieve a com- parable performance to training the full model parameters, considerably increase the training time.
In this work, we address the above problem from the data perspective. Specifically, we target find- ing smaller batches of examples that simulate outperform training with larger mini-bre batches. If this can be done, it directly improves the convergence rate of training or tne-tunnng with mini- batch stochastic gradient methods, and can yield superior performance. To achieve this, an effective

---

Published as a conference paper at ICLR 2025
approach is to find smaller mini-batches that closely capture the gradient of large random batches (Yang et al., 2023). Smaller mini batches (coresets) selected by this approach are medoids (cen- troids) of the data in gradient space, weighted by their corresponding cluster size (Mirzasoleiman et al., 2020). Despite its promise on classification tasks, this approach is infeasible and ineffective for pre-training fine-tuning L.MMs, we discuss below.
Firstly, language data is often a mixture of highly imbalanced sources (e.g. categories or types of instructions). In this case, we show that smaller mini-batch coresets do not contain representative examples from the small sources, and obtain poor performance. Secondly, Adam is the standard op- timizer for training LLMs, and small are not optimal for training with Adam. Finally, the very large dimensionality of the LLM gradients makes pairwise distances vacuous and the medoids cannot be found accurately. These challenges make finding mini-batch coresets for training LLMSs inherently much more challenging to address.
that capture the vanilla gradient of larger batches
In this work, we propose Coresets for Training LLMs(CoLMM by making the following contributions:
First, we show that w.h.p. mini-batch coresets only contain medoids of the big sources with a large-enough number of samples. Thus, examples selected via gradient matching from small sources are not representative and do not benefit learning other examples in their sources. There- fore. it crucial to include all examples of the small sources in the bati- coresets. Besides, to enhance learning small sources, we weight all examples in the mini-batch coresets uniformly.
Next, to find mini-batch coresets for training with Adam, we normalize gradients by their histor- ical exponential average, where the historical terms are only calculated for examples in the big sources. We find medoids of the big sources based on their normalized gradients.
Finally, to enable finding medoids in the very high-dimensional gradient space, we use zeroth- order methods find smooth gradient of the last -projecction matrix in a memory-efficient way, and sparsify it via a source-wiss mask, which keeps dimensions with the largest normalized gradient magnitude. We use & distance between sparse gradients find medoids of big sources.
We evaluate CoLM on the challenging task of mathematical problem-sslving, by fine-tuning Phi-2, Phi-3 (Li et al., 2023b), Zephyr (Tunstall et al., 2023), and Llama-3 models (Dubey et al., 2024) using LoRA on the MathInstruct dataset (Yue et al., 2023) containing 14 highly imbal- anced sources. Additionally, we apply CoLM to datasets in SuperGLUE benchmark (Wang et al., 2019), where we find sources by clustering the model's hidden states during the training. Remarkably, CoLM reduces the memory requirement of fine-tuning by 2x and even outperforms training with 4x larger random mini-batches. Compare to mini-batch of the same size, CoLM outperforms by up to 7.1 % and 20% on several in- and out-domain tasks.
Notably, our approach can be easily stacked with LoRA, and other memory-efficient training meth- ods to further reduce the memory requirements of training LLMs, as we confirm in our experiments.
2 RELATED WORK
Memory-efficint training of LLMss To address the large memory requirements of training LLMs, several methods have been recently proposed. LoRA (Hu et al., 2021) freezes the pre-trained model weights and trains two low-rank adaptor weight matrices adapt the weights of each layer. How- ever, LoRA suffers from a performance drop compared to training with full-rank matrices. To im- prove upon this, several variations of LoRA (Liu al., 2024; Renduchintala et al., 2023; Xia al., 2024b) have been proposed. Besides, GaLore (Zhao et al., 2024b) proposed to reduce the mem- ory cost of optimizer states by calculating the gradients and projecting them into a low-rank space. However, the above approaches also lead to increased computational costs.
Another line of methods approximate backpropagation by sparsifying gradients (Frantar and Alis- tarh. 2023), subsampling the computational graph (Adelman et al., 2021), gradient check-pointing (Chen et al., 2016), and quantization of weights and optimizer states (Dettmers et al., 2022). However, these approaches can incur large approximation errors and cause performance drops. Zeroth-order gradient approximation has also been used for memory-efficient training (Malladi et al., 2023). However, this approach cannot reach comparable performance normal training.
Our method can be easily stacked with existing memory-efficient methods to improve convergence and further reduce memory requirements.
2

---

Published as a conference paper at ICLR 2025
Data selection for training LLMs. Data selection for training LLMs has garnered significant at- tention due to its potential to enhance model performance while reducing computational costs. For pre-training, examples with middle perplexity rankings are shown beneficial (Marion et al., 2023). Clustering based on embeddings of a pretrained model and sampling from the clusters to drop re- dundancies has been also investigated (Tirumala al., 2024).
For fine-tuning, training on manually crafted high-quality instruction/response pairs has shown highly effective (Zhou al., 2023a). Building on this observation, data selection using LLMs such as ChatGPT or training on textbooks is proposed (Eldan and Li, 2023; Li et al., 2023c; Chen et al., 2024: Li et al., 2023a), and metrics such as diversity (Bukharin and Zhao, 2023; Du et al., 2023), difficulty (Bhatt et al., 2024; Marion et al., 2023; Zhou et al.. 2023b), and completion length (Zhao et al., 2024a) are shown relevant. Given a high-quality validation set, using influence functioon to select the most beneficial subsets of fine-tuning data has been also explored (Xia et al., 2024a). How- ever, a high-quality validation set not always available (e.g. for MathInstruct). Existing methods select data in a one-shot manner before fine-tuning, and either require access to another open LLM or a large preprocessing time to fine-tune the original or a proxy LLM on the target data.
We study data selection from different perspective, i.e. by selecting small mini-batches that match the performance of training with larger mini- bathee As baselines, we adapt several one-shot methods based on loss (Jiang et al., 2019), gradient norm (Katharopoulos and Fleuret, 2018), middle perplexity (Marion et al., 2023), completion length (Zhao et al., 2024a), confidence, and hidden-state centrality (Bhatt al., 2024) to iteratively select small mini-batches from larger random batches.
PRELIMINARY: MATCHING GRADIENT OF LARGE BATCHES
Consider training a machine learning model on a dataset indexed by V, by minimizing the loss function L(0) E,ev E,ev[C,(0)) [O). Mini-batch SGD with learning rate n iteratively updates the model parameters as 0.+1 = 0, g.M.,I. where g.M.,' = E,e,M, [g...] is the gradient of a mini-batch M₁ of random examples, and g (..). As long as the mini-batch size b = [M&ll is not too large, the convergence rate of mini-batch SGD directly scales with a factor of 1/b. Formally, for a non- convex L-gradient Lipschitz loss, mini-batch SGD with small enough 17 will visit an e-stationary point w.h.p. at least once in the following number of iterations (Ghadimi and Lan, 2013):
Ô ( [T(C(O)) -()) ( e²
be
where lg.. - g < o² is the variance of the individual gradients.
To improve the convergence of mini-batch SGD with mini-batch b, one can iteratively find weighted subsets (mini-batch coresets) of size that closely match the gradients of large random batches M of size >> b (Yang et al., 2023). As the mini- coresets have a similar gradient to large batches, they have a smaller variance of o² /r. Thus, they improve the convergence of mini-batch SGD by r/b. The mini- batch coresets found by gradient matching medoids (centroids) of the larger hatch in the gradient space, weighted by their corresponding cluster size, and can be found by maximizing monotone submodular' function via the greedy algorithm (Mirzasoleiman et al., 2020):
S E arg max
Σ max .ES iEM!
||ga, gr,elll,
SCMH,1S1Sb
where C' is big constant. To find subsets efficiently, gradient of the loss w.r.t the input to the last layer of the model (which best captures the variation of gradient norm (Katharopoulos and Fleuret, 2018)) is commonly used (Mirzasoleiman al., 2020; Pooladzandi et al., 2022; Yang et al., 2023).
4 COLM: MINI-BATCH CORESETS FOR TRAINING LLMS
Despite its success of image classification tasks (Yang et al., 2023), the above gradient matching formulation performs poorly for training LLMs, due to the following reasons:
A set function F 2 + RR is submodular if F(c(S)= F(S (()) F(S) 2 U - F(T), for any V and E V VT.. F monotone for all S T, F(S) < F(T).
3

---

Published as a conference paper at ICLR 2025
V3
V3
V3 V2
V2
V2
V1
V1
V1
Figure 1: A toy imbalance data. (Left) Full data V with two big (blue, green) and one small sourees (purple). & = 3 medoids of the data are shown in red. (Middle & right) Two random samples of the data, with their corresponding k = 3 medoids. The a -nighborroood of big sources are dense and thus medoids of random samples contain central examples of the big sources. However, the medoids of random sample do not necessarily contain central examples of the small source.
Highly Imbalanced Language Data. Language data often contain highly imbalanced sources (c.g. categories or types of instructions). For example, the ratio of the largest to smallest source size in the MathInstruct data is 300. Here, subsets found by gradient matching do not contain representative examples (medoids) of the small sources and yield suboptimal performance.
Adam optimizer. Adam (Kingma and Ba, 2014) is the commonly used optimizer for language tasks. Adam scales each gradient dimension to speed up learning along flatter dimensions and slow down learning along sharper ones. Matching the vanilla gradient of large batches yields suboptimal performance when training with Adam.
Very Large Gradient Dimensionality. Finding medoids via Eq. (2) requires calculating pair- wise distances between per-example gradients. But, in the very high-dimensional gradient space of LLMs, distances become vacuous. Even the last layer is too high-dimensional to find medoids accurately. For instance, the dimensionality of the last V-projection matrix of Phi-2 is 6.5M when training the full parameters and 327K (matrix B) when using LoRA with rank 128.
Next, we will discuss how we address each of the above challenges.
4.1 DEALING WITH THE IMBALANCED LANGUAGE DATA
First, we address the challenge of dealing with language data containing highly imbalanced mixture of sources. Our key observation is that larger random batches contain many examples from the big sources. Thus, mini-batch coresets selected via gradient matching contain central examples of the big source. In this case, training on the coresets benefits learning other examples in the big sources. On the other hand, small sources have a few examples in the large batches. Hence, the mini-batch coresets do not contain central examples of the small sources. In this case, training on the coresets does not benefit learning other examples from small sources. This implies that one cannot reliably select representative examples of small sources from the larger batches. Indeed, for optimal performance, it is crucial to train on all examples of small sources in the larger batches.
Next, we formalize this problem. Consider a dataset V containing Q sources, i.e., V = [Vi U ... U Vo Suppose gradients of examples in source V,, at iteration are drawn from an underlying infinite set, according to an unknown probability distribution. Let A be the set of A medoids of the infinite set, such that around each i E A there is a neighborhood of radius at least a', where the probability density is least B at all points, for some constants a", B. This implies that the medoids are from reasonably dense and therefore representative regions of the gradient space. Let us consider g : IR + R. to be the volume of ball of radius a centered at a point in the metric space. The following theorem shows that for a large enough source that is randomly partitioned into m parts, there are many examples from the dense neighborhoods in every partition.
Theorem 4.1. Let examples in Vg be partitioned into m parts. A number of examples Vgl 2 2km (k)) where a at, is suffice (1) have at least km kmlog(km//) elements in the Bgla)
a-neighborhood of each i E At and (2) have each partition contain elements from all k a- neighborhoods with probability at least (1 for small 8 > 0.
4

---

Published as a conference paper at ICLR 2025
Next, show that the medoids of every partition central examples from the dense neighbor- hoods of the (infinite) data, with high probability. That is, they are in neighhboodooo of Ay
Theorem 4.2, Let ò, 0 and let ny IVa and no be an integer such that for ny 2 no have
Rg mk If ni 2 max (no: vn log(2m/8)) with probability of at least 1 - 6, medoids found In[m,) c
by the greedy algorithm from every purtition are in a-neighborhoods of each i E A/-
Note that in every large random batch, there is part of every source Vq with expected number of examples |Vl/LV Hence, the above theorems imply that for small sources without a large-enough sample size, coresets found by gradient matching do not necessarily contain their central examples. Hence, training on them yields poor performance on small sources. Fig shows an illustration.
Coresets for Imbalanced Data. Consider a data V = (Vi ..... U V, U Vp+1 U... V2). with P small and Q large sources. For a dataset with e sources, we regard small sources as those with less than /Vi/c examples. To learn the small sources, we include all of their examples from the large batch in the small mini-batch coreset. That is, S = {u E I E UsElo) V.). But, for every big source V, where E {p + 1,..,,,.) we apply the greedy algorithm to its examples in the larger batch Vi (u E Mf E V2) and add medoids St E V to the small mini- batch coreset, where by = | is proportional to the number of examples from Vg in MM, i.e., by - (b-1s)) (/9 -( The mini-batch coreset at step / is S. = ( St 5p+1 U..) We note that sources are mostly separable based on their gradients. Thus, one subset can found from all examples of big sources in the larger batch. However, selecting subsets separately yields slightly better performance, as we show in our experiments. Finally, to ensure learning various big and small groups For datasets such as MathInstruct, the sources are labeled in the training data. For datasets without specified sources, we cluster the hidden state of the model find sources during fine-tuning.
4
uniform speed,
assign uniform weights to all the selected examples.
The following theorem shows that the small mini-batch coresets have a smaller variance compared to random mini-batches of the same size. Therefore, they guarantee superior convergence rate.
Theorem 4.3 (Variance reduction). Let the number of outliers which do not belong to any k dense areas be k. Let Ou >> a* be the largest distance from an outlier to any centroids. Assume that all the selected samples St belong to the dense areas. The variance of the mini-batch coresets of size b is smaller than the variance of the random subset of size b by up
E (a, a*) a2a(""a + (a. a a))). m
4.2 FINDING CORESETS FOR TRAINING WITH ADAM OPTIMIZER
Next, we address finding mini-batch coresets for training with Adam optimizer. Adam adapts the learning rate across dimensions by scaling the gradient updates by square roots of exponential mov- ing averages of squared past gradients. In doing so, reduces the learning rate across sharp dimen- sions and increases the learning rate across flatter dimensions improve convergence. Formally,
m. 8im,- + (1 31)ge
U2 = 32v;- + (1 32)g;
0, 00-1111
m.
(3)
1- 81
54
c Ut
For selecting mini-batch coresets for training with Adam, matching the vanilla gradient is not enough. To do so, we normalize every gradient dimension by the exponential average of its his- torical values. Additionally, terms m, u only based on the big groups' gradients, which we denote by m, ò. This allows a more precise selection of the subsets, as also confirm in our experiments. Formally, we select the medoids of the normalized gradients of big sources, by solving the following submodular facility location function:
we only select medoids of big sources, we calculate the historical
ST E arg max
Σ max SES
- é + V e
My.i
Mr,s
(4)
SCV S
The very high dimensional gradient of LLMs makes solving Eq. (4) prohibitively expensive. Be- sides, in such a high dimensional space, pair-wise distances become vacuous. Next, we discuss lowering the gradient dimensionality to find medoids of big sources more accurately.
5

---

Published as a conference paper at ICLR 2025
4.3 FINDING
GRADIENT ESTIMATES
The very high-dimensional gradients of LLMs are very noisy. To find smoother lower dimensional gradients in a memory efficient manner, we use zeroth-order methods to calculate the gradient of the last V- -proojection matrix, and then sparsify it to lower its dimensionality. The V-projections matrix allows finding higher-quality subsets, as we will confirm in our experiments. Notably, this approach stacks well with memory-efficitee training methods such as LoRA, as we will confirm empirically.
Simultaneous Perturbation Stochastic Approximation (SPSA) (Spall, 1992) zeroth-order tech- nique that estimates the gradient as:
g [( + cz) C(0 cz) 2 zz g.
(5)
2e
where 2 E IRd is a random vector with (O, Ia). and d is the number of model parameters and e is the perturbation scale. As e + O, the SPSA estimate provides a rank-1 reconstruction of the gradient, and this is smoother than the actual gradient calculated with backpropation. SPSA requires two forward passes through the model to compute the gradient estimate.
Estimating the Last V V-Projection Gradient. To get the gradient of the last V -projecciion matrix for example i, instead of perturbing all the parameters, we sample random perturbations for parameters corresponding to the last (LoRA) V-projection in the perturbation vector z, and use zero for the other entries:
(., + (2up) L.t
(.: vZvp)
Zup:
(6)
2
where uup E Ru with (d : Eup] and upp the flattened last V -projectioon matrix. Eq. (6) can be calculated very efficiently in just one forward pass. To do so, first make forward pass to get the activations X of the penultimate layer of the model. Then, we perturb the last-layer parameters twice to calculate g based on the pre- calculated X The time of getting the lower dimensional last-layer gradients will be dominated by the time of computing X and the cost of the second step is negligible. To minimize the memory requirement, one can follow (Malladi et al., 2023) to use a fix seed to generate the same perturbation 2 multiple times. Hence, the memory overhead is also negligible.
I ). and dup is the dimensionality of
We use the zeroth-order gradient estimates in Eq. (6) to calculate normalized gradients m., û, for the big sources. Nevertheless, these gradients are still too high-dimensional to find medoids accurately.
Sparsifying the Last-layer Normalized Gradient Estimates for Adam. To further reduce the gradient dimensionality, we sparsify the normalized V-rooeection gradients. The subsets are selected for each big source separately. Thus, for every big source we find dimensions that best preserve the normalized gradient norm of its examples Vt C .M in the larger random batch. The normalized gradient norm to the first order indicates how much each gradient update achieves a loss reduction:
L( E m/(c vv.)) (()) = c
m
m C
AC = lim c00
Dimensions that best preserve the normalized gradient norm are those with the largest magnitude. Therefore, for every big source, we sparsify the normalized zeroth-order gradient mi/((s + Vor) of the last (LoRA) V V-projection by mask vector Mt. which has for the top h parameters with the largest magnitudes and O elsewhere.
Using l, distance in high dimensions. We calculate the pair-wise normalized gradient dissimilarity between sparsified gradients using & distance, which is preferable to Euclidean distance in high dimensions (Aggarwal al., 2001). The medoids for each big source are found by solving:
$
arg max SCV leV;
Σ max
miss E Vor.
M
Mi.
M/
(8)
e+ Vur
In experiments, we show that selecting as small as 0.7% of the (LoRA) last layer gradient dimensions enables finding high-quality subsets efficiently. Additionally, we show that compared to random projection (Johnson, 1984), this approach enables finding higher-quality subsets and is orders of magnitude faster in practice.
The pseudo-code of CoLM is illustrated in Alg. in Appendix B.

---

Published as a conference paper at ICLR 2025
Table Accuracies (f) on in-domain and out-of-domain datasets when fine-tuning Phi-2 with LoRA
on the MathInstruct for IK iterations. One-shot selection techniques (CL, GN, LC, FL, MP) are adapted to select small mini-batches on the fly. CoLM with batch size (bs 64) outperforms all the baselines. Notably, CoLM even outperforms fine-tuning with bs 256, while using 45% less memory, and achieves similar performance to fine-tuning for 2K iterations with bs = 128.
Method
In domain
Out domain
Avg All
GSM8K
NumGL.UE
Aug
SVAMP
Math.
SimulFq
Avg
Pretrained
$2.9
67.9
31.9
28.8
42.9
FT (bs=64)
66.31a.
28.4g0.3
50.2g0.9
48.3_0.2
79. 2.0.
52.4
24.121.
51.9
50.1
CL
56.1+3.2
26.4
32.9=3.2
38.5-2.
33.3+44
16.9+4.9
11.9+3.2
30.7+2.8
31.6+22
BL
58.0+0.s
21.1+a.5
43.7=22
40.9
77.1=0.7
38.0+4.4
18.4+0.
44.5+1.5
42.7=1.0
GN
65.0+1.2
24.9
45.5+112
45.1-1.1
76.7=1a
42.9+2.7
16.8+2.a
45.5+2.0
45.3+1.6
LC
59.3+n.9
24.0+a.7
48.0+u.5
43.8-c.4
79.5=o.a
45.9+0.3
23.6+2.
49.7+1.1
46.7+0.7
FL
68.011.1
29.21um
51.411.
49.5_..7
80.410.3
55.611.2
30.5
55.5
52.511.0
MP
65.3gn.2
28.4g0.a
64.6g1.6
40.4_o.5
70.8
53.611.0
36.613.0
58.7
53.0go.7
CoLM (Ours)
68.4g..a
29. 20..4
57.3
51.9.0.5
80.2#1.
59.8=1.1
41.1=a.8
61.4211
56. 66.6g0.0
FT (bx=128)
67.4+n.
28.8+u.
53.2+1.2
49.8-o.a
80.4+1a
55.6.0.4
29.9+2.4
55.3+1.o
52.640.6
FT (hc=256)
67.51a.1
29.61u.2
58.3
51.8_c.4
79.
56.310.7
40.5
58.911.2
55.31o.
FT (bs= (28) 2K
67.7
30.3
58.4gu.
52.1ge.
79.5gm.
67.9
45.5gu.
60.9
56.5go.s
5 EXPERIMENTS
In this section, evaluate the performance of CoLM, for fine-tuning LLMs, by comparing the performance, memory requirement, and wall-clock training time of training with small and large random mini-batches, with that of our method, In addition, conduct an ablation study showing the effectiveness of different design choices of CoL.M We further demonstrate the effectiveness of CoLM in the pre-training setting Appendix F.
5.1 SETTINGS
Training datasets. We the MathInstruct (Yue al., 2023) dataset which consists of about 260K instruction tuning examples, curated from 14 highly imbalanced open-source math datasets, with broad coverage of mathematical fields and wide range of difficulty levels. The ratio of the largest to smallest source in MathInstruct is almost 300, and the distribution sources can be found in Fig 4a in the Appendix. For datasets without specific sources, we use three datasets from the SuperGLUE benchmark (Wang et al., 2019) for the classification task: SST-2, CB, and MultiRC.
Models. We utilize the Phi-2, Phi-3 (Li et al., 2023b), Zephyr (Tunstall et al., 2023), and Llama-3 models (Dubey et al., 2024).
Training details. We use LoRA with rank of 128, alpha of 512, and dropout rate of 0.05. For Phi models, we apply LoRA to all attention matrices (i.e. QKV -roj)) and two fully connected layers while for Zephyr, we apply LoRA to all attention matrices (i.e. QKVO-proj). All experiments run on 4 NVIDIA A40 GPUs We repeat each experiment three times.
Baselines. We compare CoLM with normal fine-tuning (FT) using small and large batch sizes and an online selection method called Big Loss (BL) (Jiang et al., 2019). In addition, we adapt one- shot selection techniques, including Grad Norm (GN) (Katharopoulos and Fleuret, 2018), Middle Perplexity (MP) (Marion et al., 2023), Completion Length (CL) (Zhao et al., 2024a), Least Confi- dence (LC), and selecting centroids of the model's hidden state by maximizing submodular facility location (FL) (Bhatt et al., 2024), to select small mini-batches from larger ones during fine-tuning.
Evaluation datasets. Following (Yue et al., 2023), we use a variety of popular datasets across both in-domain and out-of- datasets. The in-domain datasets include GSM8K (Cobbe et al., 2021), MATH (Hendrycks et al., 202 I), and NumGLUE (Mishra et al., 2022). For the out-of-domain datasets, we include SVAMP (Patel et al., 2021), Mathematics (Davies et al., 2021), and SimulEq et al., 2016).
Additional training details and evaluation metrics are specified in Appendix C.
5.2
MAIN RESULTS: MATHINSTRUCT
CoLM achieves a superior performance. Table shows the in-distribttion and dustof-distt accuracies of fine-tuning Phi-2 with LoRA the MathInstruct dataset for 1K iterations. First,
7

---

Published as a conference paper at ICLR 2025
Sad
FT
Coult Ibeeed)
Accuracyy (%6
Memory (G8)
202
CoLM (half bs)
155
302
Patt ta
CoLM
Pat
Liastd
Lamid
Mini-batch
Minl-batch
(a) Different LLMs
(b) Changing mini-batch size
(c) Me)oory usage batch size
Figure 2: (a) CoLM with bs = 64 (from 128) outperforms fine-tuning different models with bs 64 and bs = 128 by a large margin; (b) CoLM improves the performance training with different batch sizes. The size of each circle is proportional to the training time of the corresponding method. (c) CoLM reduces memory consumption, with reduction increasing as the batch size grows.
126
Timett Actrnscy
1.
19c
Mervay
Varance
16
50.0
(b5-64)
(b5-64)
(bs=128)
425
FT (bs=128)
Prodes 1003
CoLM (bs 664
CaLM (bs=64)
Training iteration
of training Iterations
(a) Time vs Accuracy vs Memory
(b) Variance of mini-batch grads
(c) Convergence rate
(a) Wall-clock time (including the time for CoI.M's selection), memory consumption, and performance of fine-tuning. CoLM outperforms normal fine- tuning for IK iterations with bs 128 (256), while being 1.3x (2.7x) faster and consuming 20% (45%) less memory, respectively; (b) CoLM has a smaller variance than random mini-batches the same size; (c) CoLM converges much faster than normal fine-tuning (FT).
Figure 3: Fine-tuning Phi-2 on Mathinstruct.
see that using a larger mini-batch size improves the performance, which is consistent with the the- oretical results in Eq. (1). Besides, we note that training for IK steps with mini-batch size of 64, 128, 256 corresponds to training on 25%, 50% and 100% of the data, respectively. Remarkably, training on only 25% of the data with CoLM with bs 64 (selected from 128) outperforms all the baselines and achieves a similar performance to fine inn-tuning for 2K iterations with bs 128. Interest- ingly, training for IK iterations with CoLM using bs 64 even outperforms training with bs 256.
CoLM improves the performance of different models and batch sizes. Figure 2a shows that CoLM significantly outperforms normal fine-tuning across 6 different model architectures. Specifi- cally, fine-tuning Phi-3 for 4K iterations using CoLM with bs = 64 outperforms normal fine-tuning for 4K iterations with bs 64 and bs 128 by 5% and 4.2%. Notably, for better-performing models, CoLM provides more performance improvement. This confirms its applicability to state-of-the-an architectures. Fig. 2b shows that CoLM improves the performance of different batch sizes, including bs = 32, bs = 64, and bs = 128, without significantly increasing the training time.
CoLM effectively reduces the Activation Memory. The memory required for training an LLM can be decomposed into three parts: activation memory + weight memory + optimizer state mem- ory. Memory efficient methods often reduce the weight or optimizer-state memory. For example, LoRA reduces the optimizer state memory but slightly increases the weight and activation memory by adding low-rank matrices. Orthogonal such methods, CoLM effectively reduces the activation memory by reducing the batch size. Hence, stacked with and gradient accumulation, can further reduce the memory, particularly when batch size is large. In that scenario, the activation memory dominates the optimizer state and weight memory. Figure 2c shows the memory usage normal fine-tuning and CoLM with half the batch size. Notably, for total bs 2048, CoLM (bs 1024) requires almost 2x less memory than fine-tuning with bs=2048. A larger batch size is useful in particular for pre-training. Our experiments in Appendix F, confirm the benefits of CoLM to pre-training. Furthermore, while memory efficient methods harm the perfor- mance, CoLM effectively improves the performance over training with larger batches. A detailed analysis of memory overhead of our method can be found in Appendix E.
methods such as LoRA
CoLM speeds up training and improves convergence. Fig. 3a compares the wall-clock time and average performance of CoLM and normal fine-tuning Phi-2, using LoRA. For CoLM, the
8

---

Published as a conference paper at ICLR 2025
Table 2: Effect of different components in CoLM.
Method
In-domain
Out-domain
Avg
Weighted medoids
48.5.1.1
51.1=0.
Medoids (using cosine distance)
48.7
54.021.
51.32a
Medoids (using & distamce)
48.6 48.620.
54.420s
51.5205
Medoids of big sources keep small sources
5..9
58.420.s
54.62008
Medoids of big sources selected separately keep small sources
50.6.0.
59.622.
55.1208
CoLM: Modoids of big sources selected separately for Adam & keep small sources
51.9:0.5
61.4:1.
56.640.9
wall-clock time includes the time for selecting the mini-batches, Remarkably, CoLM with bs 64 (selected from 128) speeds up training for 2K iterations with bs = 128 and IK iterations with bs = 256 by 2.7x, while having 20% and 45% less memory requirements and superior performance. Figure 3b shows that throughout training, the variance of CoLM (bs 64) gradients is smaller than normal fine-tuning with bs = 64, which confirms our theoretical results in Sec. 4.1. This yields a faster convergence compared to random mini-batches of the same size, as shown in Fig 3c. At the same time, although random bs = 128 has lower variance than CoLM with bs = 64, the more uniform speed of learning sources by CoLM enables it to obtain a superior performance. Furthermore, show that CoLM yields smaller loss compared to baselines throughout training and achieves the optimal performance in less training time Appendix D.
5.3 ABLATION STUDIES
The importance of different components. Tab 2 highlights the importance of different components in CoLM. Notably, including all examples of the small data sources improves the accuracy signif- ieantly by around 3% on average. This finding well aligns with our analysis in Sec. 4.1. Selecting subsets separately per source and normalizing gradients for Adam further boost the performance by 0.5% and 1.5%, respectively, justifying our methods in Sections 4.2 and 4.3. Using uniform weights for selected example slightly improves the performance by 0.4%.
Sparsification criteria for V-projection. Table 3 compares the performance of CoLM for different sparsification criteria. Keeping parameters with the largest gradient magnitude achieves the best performance. This result is consistent with that of (Guo et al., 2024), which show that parameters with the largest gradient magnitude are the most salient. Weight magnitude, common approach in network pruning (Han et al., 2015), yields a slightly lower performance than random sparsification.
Sparsity level. Table 4 illustrates the performance of CoLM when changing the dimensionality (h) of the sparsified gradients. The accuracy peaked at h = 2560, which equals the dimensionality of the hidden state of Phi-2, and gradually decreases for larger values of h. This result expected as gradients in high dimension suffer from the curse of dimensionality, yielding a sub-optimal solution.
Choices of oowddimeniional gradient approximations. We compare the usage of sparsified MeZO gradient in our method with the sparsified actual gradient (via backprop), projected actual gradient, and low-rank MeZO gradient. For sparsified actual gradient, we apply the same sparsification technique as our CoLM. For projected actual gradient, we apply a random projection to the actual gradient and leverage the memory-efficien implementation introduced by Park et al. (2023). For low-rank Me70 gradient, we use the low-rank projection technique with SVD in Gal.ore Zhao et al. (2024b) and adopt their best setting with rank = and subspace change frequency T = 200. Table 5 shows that the sparsified McZO gradient has clear margin over the other choices of low-dimensional gradient estimates, highlight the effectiveness of zero-th order gradient and our sparsification strategy in selecting high-quality subsets.
Choices of layers. We replace the last V- V-projection layer with the last FC layer and the combination of the last Q, K, V projections in CoLM. As can be seen in Table 6, using the MeZO gradient of the last V projection matrix yields gap of almost 2% compared to other choices of layers.
Completion length. (Zhao al., 2024a) found that fine-tuning on examples with the longest com- pletion length improves the performance. Figure 4b in the Appendix shows that CoLM, in contrast, selects examples with shorter answers (avg length 120) than the average completion length of the data, which is about 130. Nevertheless, CoLM significantly improves over selecting examples with the longest completion length (avg length z 210) in random mini-batch as indicated in Table 1.

---

Published as a conference paper at ICLR 2025
Table Effect of the sparsification criteria.
Table 4: Effect of the sparsity level.
Criteria
In-domain
Out-domain
Avg
Dim
In-domain
Out-domain
Avg
random weight
51.120.9
59.6g2.
55.42177
1280
50.420.8
57. 87.8
54.1
51.110.7
59.6 69.6
55.3103
2560
51.9=o.s
61.421a
56.6+0.9
weight grad
51.420.3
58.8=1.5
15.1
5120
51. 51.2.o.1
20.1
55.7
grad
51.9=0.3
61.4:16
56.6+0.9
10240
51.620.5
59.420-7
55.520.4
Table 5: Comparison between different low- dimensional gradient approximations.
Table 6: Comparison between different
choices of layers.
Approx
In nn-domain
Out-domain
Avg
Layer(s)
In-domain
Out-domain
Avg
Sparsified actual grad
51,0go3
58.3:03
51.7:03
FC
50.620.8
58.420.6
54.3
Projected nctual gmad
50.9
50.41o.
55.2104
V proj
51.9
61.4+16
56.6.0.9
Low-rank McZ0 grad
51.0z
58.1
64.62204
Sparsified Mc20 grad
51.9+0.a
61. 4+1.
56 66.6.0.0
QKV projs
51.321.
58.221-1
54.711.0
Table Accuracies (1) when fine-tuning Phi-2 with LoRA on three datasets from the SuperGI.UE benchmark for 80 iterations. CoL.M with bs = 64 (from 128) effectively improves the performance of normal fine-tuning with bs 64. When using clusters found by the fine-tuned model, CoLM outperforms fine-tuning with bs = 128.
SST-2
CB
MultiRC
Avg
Pretrained
56.6
45,5
46.3
49.5
FT (hs=64)
91.440.2
69.1
62.0aa.2
74.2
CoLM (clustering during fine-tuning)
92.3+0.4
73.3+1.7
70.5
78.7+119
CoLM (clustering of the fine-tuned model)
92.410.7
77.6 77.613.7
73. 73.02a.4
81.02z
FT (bs=128)
92.1+1.0
72.1
72.6g5.2
78.9
5.4 DATASETS WITHOUT SPECIFIC SOURCES: SUPERGLUE BENCIMARK
We apply COLM to fine-tuning Phi-2 with bs 64 (selected from 128) for 80 iterations on three classification datasets (SST-2, CB, MultiRC) in the SuperGLUE benchmark. Note that the Super- GLUE datasets do not have any source information. To find sources, we warm up the model for 20 iterations with bs = 64, and then cluster the model's hidden states. We consider each cluster as a source and define small sources as those with less than /V/c examples, where c is the number of clusters. We update the clustering four times during fine-tuning. As shown in Table 7, CoLM outperforms normal fine-tuning with bs = 64 by 4.5% on average. This shows CoLM's applicability datasets without specified sources. Additionally, model can significantly enhance the results. Compared to updating the clusters during training, us- ing the clusters by fine-tuned model improved the performance by 6.8% on average. Compared to standard fine-tuning with bs 128, CoLM also improved the performance by 2.1 To leverage this, one can fine-tune tune smaller proxy model on smaller random subset of the data and cluster its hidden states to find sources more accurately, without large overhead, For SST-2 and MultiRC, we trained the model on a randomly selected subset of 3,000 examples to find the clusters.
find that clusters found by the fine-tuned
6 CONCLUSION
To simulate training with larger mini-batch sizes with limited memory, an effective approach is to find small mini- batch coresets that match the gradient of larger random batches. We showed that for language data with highly imbalanced sources, mini-batch coresets found by gradient matching do not contain representative examples of the small sources. Thus, one should keep all examples of the small sources and augment them with examples that match the gradient of big sources in the larger batch. To enable solving the gradient matching problem effectively, we used techniques from zeroth- order optimization and model pruning to find siverrdimensional gradient estimates. We also showed that matching the normalized gradient of larger batches provides superior performance for training with Adam. Our method, CoLM, outperforms fine-tuning Phi models on MathInstruct with 4x larger batch size, while being 2.7x faster, and also improves fine-tuning the SuperGI GUE benchmark.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research was partially supported by the National Science Foundation CAREER Award 21 2146492, the NSF-Simons Al Institute for Cosmic Origins, and Okawa Research Award.
REFERENCES
Menachem Adelman, Kfrr Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations. Advances in Neural Information Processing Systems, 34; 27877-27889, 2021.
Charu C Aggarwal, Alexander Hinneburg, and Daniel A Keim. On the surprising behavior of dis- tance metrics in high dimensional space. In Database Theory-ICDT 2001: 8th International Conference London, UK, January 6-6, 2001 Proceedings pages 420-4344444 Springer, 2001.
Gantavya Bhatt, Yifang Chen, Arnav M Das, Jifan Zhang, Sang T Truong, Stephen Mussmann, Yinglun Zhu, Jeffrey Bilmes, Simon S Du, Kevin Jamieson, et al. An experimental design framework for label-efficient supervised finetuning of large language models. arXiv preprint arXiv:2401.06692, 2024.
Alexander Bukharin and Tuo Zhao. Data diversity matters for robust instruction tuning. arXiv preprint arXiv:2311.14736, 2023.
Lichang Chen, Shiyang Li, Jun Yan, Hai Wang, Kalpa Gunaratna, Vikas Yadav, Zheng Tang, Vijay Srinivasan, Tianyi Zhou, Heng Huang, et al. Alpagasus: Training better alpaca with fewer data. In The Twelfth International Conference on Learning Representations, 2024. URL https: / ///ppe //ooperreiie ew. net/ fort///rrm? FV
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 2016.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 10.1122100.14168, 2021.
Alex Davies, Petar Veličković, Lars Buesing, Sam Blackwell, Daniel Zheng, Nenad Tomašev, Richard Tanburn, Peter Battaglia, Charles Blundell, András Juhász, et al. Advancing mathematics by guiding human intuition with ai. Nature, 600(7887):70-74 2021.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 O: 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems, 35: 2022.
Qianlong Du, Chengqing Zong, and Jiajun Zhang. Mods: Model-oriented data selection for instruc- tion tuning. arXiv preprint arXiv:2311.1565: 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprini arXi::2305.0775) 2023.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pages 10323-10337. PMLR, 2023.
Lco 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:210J. 00002,, 2020.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochas- tic programming. SIAM journal on optimization, 23(4)22341-2368, 2013.

---

Published as a conference paper at ICLR 2025
Wentao Guo, Jikai Long, Yimeng Zeng, Zirui Liu, Xinyu Yang, Yide Ran, Jacob R Gardner, Osbert Bastani, Christopher De Sa, Xiaodong Yu, et al. Zeroth-order fine-tuning of Ilms with extreme sparsity. arXiv preprint arXiv:2406.00 2xXh::2406.02911, 2024.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Dan Hendryeks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL httpt:ttpp/.. aattps://oppnne hwtp//eeepereeeeeeeeeeeeeeeeeeee net./ forum? (ooorrrmdd7Byyeee Bys=77ywtttmsCee
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 2021.
Angela H Jiang, Daniel L-K Wong, Giulio Zhou, David G Andersen, Jeffrey Dean, Gregory R Ganger, Gauri Joshi, Michael Kaminksy, Michael Kozuch, Zachary C. Lipton, et al. Accelerating deep learning by focusing the biggest losers. arXiv preprint arXiv: 1910.0072,, 2019.
William B Johnson. Extensions of lipshitz mapping into hilbert space. In Conference modern analysis and probability, 1984, pages 189-206, 1984.
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pages 2525-2534. PMLR, 2018.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Rik Mencel-Kedziorski Subhro Roy, Aida Amini, Nate Kushman, and Hannanch Hajishirzi. MAWPS: A math word problem repository. In Kevin Knight, Ani Nenkova, and Owen Rambow, editors, Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1152-1157, San Diego, California, June 2016. Association for Computational Linguistics. doi: 10.18653/vI/N16- 1136. 1136.
URL https:tt/// www.llannhooooyorrrgoooo
Ming Li, Yong Zhang, Zhitao Li, Jiuhai Chen, Lichang Chen, Ning Cheng, Jianzong Wang, Tianyi Zhou, and Jing Xiao. From quantity to quality: Boosting I1m performance with self-guided data selection for instruction tuning. arXiv preprint arXiv:2308. /2032, 2023a.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.05450 05463, 2023b.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.05463 2023c.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In International Conference Machine Learning, pages 22188-22214. PMLR, 2023.
Shih- Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Ting Cheng, and Min-Hung Chen. DoRA: Weight-decomposed low-rank adaptation. 24Xi::2402.09333.. 2024. URL arxiv. org/ abss//22.2... 09353.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36:53038-53077 2023.
Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang, Marzieh Fadace, and Sara Hooker. When less is more: Investigating data pruning for pretraining 1lms at scale. arXiv preprint arXiv:2300.0444 045644, 2023.
12

---

Published as a conference paper at ICLR 2025
Baharan Mirzasoleiman, Jeff Bilmes, and Jure Leskovec. Coresets for data-efficient training of machine learning models. In International Conference on Machine Learning, pages 6950-6960. PMLR. 2020.
Swaroop Mishra, Arindam Mitra, Neeraj Varshney, Bhavdeep Sachdeva, Peter Clark, Chitta Baral, and Ashwin Kalyan. NumGLUE: A suite of fundamental yet challenging mathematical reason- ing tasks. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), pages 1505-3522 Dublin, Ireland, May 2022. Association for Computational Linguis- tics. doi: 10. 18663///1/2022.acc-ooog acl-long: 2.
URL https: https////aaaaaataaaaaaaee prg/2022.
Sung Min Park. Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior scale. arXiv preprint arXiv:2303. 14/86, 2023.
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are NLP models really able to solve simple math word problems? In Proceedings of the 2021 Conference of the North Americun Chupter of the Association for Computational Linguistics: Human Lunguage Technologies, pages 2080- 2094, Online, June 2021. Association for Computational Linguistics. doi: 00.18633/1/202/ naacl-main. 168. URL https: httppp/s///aaaaoooooo...oo org/ 2rtptt/hoooooooo/22221111111 naacl monel-mainng www.. 168.
Mohammad Taher Pilehvar and Jose Camacho-Collados. Wic: the word-in-context dataset for eval- uating context-sensitive meaning representations. arXiv preprint arXiv:1808.0912) 2018.
Omead Pooladzandi, David Davini, and Baharan Mirzasoleiman. Adaptive second order coresets for data-efficient machine learning. In International Conference on Machine Learning, pages 17848-17869. PMLR, 2022.
Rajpurkar. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.00220, 2016.
Adithya Renduchintala, Tugrul Konuk, and Oleksii Kuchaiev. Tied-lora: Enhacing parameter effi- ciency of lora with weight tying. arXiv preprint arXiv:2311.09577 2023.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 20/1 AAAI spring symposium series, 2011.
James Spall. Multivariate stochastic approximation using simultaneous perturbation gradient approximation. IEEE transactions automatic control, (733):33-341, 1992.
Kushal Tirumala, Daniel Simig, Armen Aghajanyan, and Ari Morcos. D4: Improving I1m pretrain- ing via document de-duplication and diversification. Advances in Neural Information Processing Systems, 36, 2024.
Lewis Tunstall, Edward Beeching. Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, et al. Zephyr: Direct distillation of Im alignment. arXiv preprint arXiv:2310.1694 2023.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Se- lecting influential data for targeted instruction tuning. arXiv preprint 240.02300. 2024a.
Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of language models via residual learning. arXiv preprint arXiv:2401.044 0411111 2024b.
Yu Yang. Hao Kang, and Baharan Mirzasoleiman. Towards sustainable learning: Coresets for data- efficient deep learning. In International Conference on Machine Learning, pages 39314-39330. PMLR, 2023.
13

---

Published as a conference paper at ICLR 2025
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653, 2023.
Hao Zhao, Maksym Andriushchenko, Francesco Croce, and Nicolas Flammarion. Long more for alignment: A simple but tough-to-beat baseline for instruction fine-tuning. arXiv preprint
arXiv:2402.04833, 2024a.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024b.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, LILI YU, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. LIMA: Less is more for alignment. In Thirty-seventh Conference on Neural Information Process- ing Systems, 2023a. URL https: / wpwp//peeeeewweeeeeeeeeeeo net fo
Haotian Zhou, Tingkai Liu, Qianli Ma, Jianbo Yuan, Pengfei Liu, Yang You, and Hongxia Yang. Lobass: Gauging learnability in supervised fine-tuning data. arXiv preprint 1it::2300.13000 2023b.
14

---

Published as a conference paper at ICLR 2025
A VARIANCE REDUCTION BY FACILITY LOCATION
Outline. In this section, we present the proofs for our theoretical results Section 4.1. Firstly, we introduce the notations, problem formulation, and all assumptions. Secondly, Theorem 4.1 is the result of Lemma A.4 and the first part of Lemma A.6. In addition, we provide a bound for the local optimal solution in Corollary A.7. Thirdly, prove Theorem 4.2 using Lemma A.8. Finally, Theorem 4.3 is subsequent to Lemma A.9.
Notations.
d(:;-) V x V R is the distance between two elements. (v) = {w: d(v, w) S {a) is the set of elements within a distance a from v, called o-neighborhood. g(a) + Ih is the volume of a ball radius centered
point in the metric space. In TR, we have g(a) O(aP).
Facility Location. For a set V. we solve the k-medoid problem by finding a subset S such that [S] = k and S minimizes L(S) T Vvev minees d(v, e). We can turn into monotone submodular function by using an auxiliary element up: f(S) = [((vo)) - L(S u (vo}).
Settings. We have dataset V with examples in IRD which is drawn from an underlying infinite set, according to some unknown probability distribution. Let A such that |A| = 1 be the global optimal solution of the facility location problem in the infinite set.
Assumption A.1 (Data structure). For each e₂ E A, there is a neighborhood of radius at least a², where the probability is at least at all points, for some constant a* and [3.
is known that f is decomposable. In other words, f can be written as sum of (non-negative) monotone submodular functions follows: f(S) = [V] Evev fe(S). We define the evaluation of
f restricted to D E V as follows: fD(S) = [ ΣieD f((). Assume that the objective function f have the following two properties.
1
Assumption A.2 (Lipschitz property). f 2V - R X Lipschiit. In other words, for equal and S" {ví, 12 v2.,..,, vk,, and for any matching of elements M
sized sets S = {v), 02,... {(V1, v, (v2, r2),........ ... (V), r,)), the difference between ((S) and f(S") is bounded by the total of distances between respective elements (SS) - f(S') S X, d(u,, v²).
y ,
Assumption A.3 (Bound property). fü is bounded, and without loss of generality 0 S f((S) S for 1 < < VV,S V.
The dataset randomly partition into m large the local optimal solution for each M, as A, where |A,1 = k. We are showing in the next Lemma that when the size of the training set is large enough, it has many examples from all the dense areas.
{M of size r a We denote
m
2km log(km//)) where S a suffices to have at least Sgla)
Lemma A.A. A number of elements 2
each e, E A with probability at least (1 - for km log( (og(k//) (5) elements in the a neiighboorhood eacch small values of 6.
Proof. The probability of a random element being in No (es) is at least 3g(a). Thus, the expected number of a- neighbors of 6 A is ElIN, (e+)1] 2km log(km/5).
From the Chernoff bound, we have for every < O, (e:) < km log(km/6)) S Elexp(t * INo (ex)))eep(( km log(km/6))
exp(t (E]INN - km log(km//)))
S exp(( km log(km/5)).
Let t = - km
in the above equation, we have
Pl]N(ee:)l S km log(km/00) S
(10)
km
Therefore, the probability that some e; E A does not have large enough neighborhood is
P PU (e:)l < km log(km/00) <
(e)) $ km log(km/
im1
6
k
< 8.
(II)
km
15

---

Published as a conference paper at ICLR 2025
Therefore, with probability least 1 - 8, the a-neighborhood of each element 6 E A contains least km log(km/0 elements.
Next, we prove that sampling with replacement guarantees that each mini-batch has elements from all the dense areas.
Lemma A.5 (Sampling with replacement). for each e; E A, [aa(e]]] m (k/), and if Mj is a mini-batee batch of size n/m sampling with replacement, then Mj contains elements from all is dense areas with probability at least (1 8).
Proof. The number of mini-batches Mj does not contain elements from (es) is (n m log(k/8)) (m/m)))) The total number of mini-batches of size n/m is n (n/m), Thus, the probability of My does not contain elements from (e;e;) is (2-m (==mlor(k/b)) loo(k/0) ((n/m) (1
)(//m)
(kg()7))
crr((-lo(k/// = k/k. Therefore, the probability that M, does not contain elements from all dense areas is
PL M,
Σ N(MM) =
(12)
The above guarantee also holds for sampling without replacement as shown in the following lemma.
Lemma A.6 (Sampling without replacement). for each e; A. [Na(e:)l 2 km log (logkm//0) ( and V is partitioned into m mini-batch M1.M2,.... MMn... then each Mj contains elements from all the dense areas and ((A f(A,) I xak with probability at least (1 8).
Proof. Because (e:) 2 km log(km/5), we can construct i mutually disjoint subsets {S;) ;=1 such that S; E No (e:) and |S|| = m log(km/ö) ö). Each element in S; goes into a particular M with probability of 1/m. The probability that particular M, does not contain element in S is P(IM, 0 S.
(1 1/m)m log(km/6)) /)
km The last equality hold because lim
m-++00
1/m)m (-p(-1). The probability that M, does not intersect with at least one S;
p PIUM, =
Σ PIMM St
8
(13)
m
5 Therefore, the probability that Mj contains elements from every is least probability that every Mj contains elements from every S; is
Thus, the
m
h P PNN M > 0)] = II P M, ns > 0] =
ym m 1 1-6. 6.
(14)
Thus, with high probability 8, every M, has subset S such that are |S [S, N. (e.e)) for e; A. Therefore, f(A) f(A)) s f(A) f(S,) I dak.
|A| = 1: and
From Lemmas A.4 and A.6, we have the following corollary
Corollary A.7 (Bound for local optimal solution). For n 2 2km log(Akm//)
Bgl it) where As I a², Vf is
partitioned into mini-batches M1,M2, M.. then for sufficiently small values of 6, we have (f(A) - f(A)) < E with a probability of at least 1
Lemma A.8 (Bound for local evaluation). Let no be an integer such that for n 2 no we have with a probability of at least 1 8, we can evaluate on each mini-batch M, with small error of e, i.e. 1SM f(S)
Tn(n] mkk n 2 max (a.. ne mllog(2m66)
16

---

Published as a conference paper at ICLR 2025
Proof. Note that each mini-batch has exactly |Mjl n/mm elements. Let define E, (S) the event that [fM, (S) - f(S) < c, for some fixed e < and a fixed S with |S] - k. Note that E, (S) denotes the event that the empirical mean fM, (S) is close to the true mean. Because f is decomposable,
we have fM, (5) M, ΣiEM, {S) = LieM, [M,T Also remember that 0 s
f+(8)
f,(8)
M, = M
Based the Hoeffding inequality (without replacement) have
P|-S(S)) P( ((S)
(S)v
E(sv 2 c)
26a
S 2xx
[Vil(TVl - -0)3
-2exp(-22
= 22xxp(-22 -anc /m))
(15)
Let &i be an event that (S) [(S) < for any S such that |S| S k. Note that there are at most n sets of size at most (because sampling samples with replacement results in subset of size most k). Hence,
PI-E.] 2n exp(-2ne" /m)
(16)
There are m mini-batches, by the union bound we can conclude that
m
PL
G
Σ Pi-E- < 2mn exp(- /m)
(17)
The above calculation implics that we need to choose 2 2mn exp(-2ne" (imee/mm) /m) so that w.h.p
1 8 we evaluate locally each mini-batch. For large n, the function In(n) is increasing
function, thus, there exists no such that for n 2 1op. In(n]
2 mk 2 Then, we choose n as follows
m log(2m/5))
n max
fl:
(18)
2
hold for all subsets S such that |S] S k. the data distribution of Vi should be similar to that of V. Hence, k-medoids of V; are in close neighborhood of k-medoids of V.
For -
Lemma A.9 (Upper bound for the variance). Let the number of outliers which do not belong to any k dense area be k. Let au >> a be the largest distance from an outlier to any centroids. Assume that all the selected samples A belong to the dense areas. The upper bound of the variance of the local optimal solution A is smaller than that of the random subset of size k.
Proof. For each subset S, we use the notation S to denote the centroid of this subset. Let ej be the distance between the centroid of subset S3 of size k to the centroid of A. The variance of the subset S3 has an upper bound as follow.
Var(S;) = Var(A +ez)
Var(e,)
(19) For each local optimal solution Aj. we know that ej S a= For a random subset Sj size k, there k/m outliers and & (k/m) examples from the dense areas in the subset average. Thus, the distance e is bounded as éj < (1 & )aa + _a, 2 a*. Therefore, the upper bound of a random subset 5; larger than that of the select subset Aj.
I Eleg).
From the above lemma, we can conclude that
Theorem A.10 (Variance reduction). The variance of the mini-batch coresets of size b is smaller than the variance of the random subset of size b by up (Or, a² (2a" + z (a., a²)).
17

---

4 NVIDIA A40 Published as conference paper at ICLR 2025
Algorithm Coresets for Training LLMs (CoLM) on Imbalanced Language Data
1: Input: 0 E R.d, loss IRd - I,, step budget T, batch size learning rate schedule (ne).
small sources (V1,-.. ... Vpt. large sources (Vp+1,'* V&}
2: for 1,... ,T' do
3: Sample batch M, C
4:
5% - Molo Uselp Vpb
I/ Keep all samples from small sources in the batch
5:
for (p+1, ......... ,4] do
6:
t ("E Vi) (b b
// Find all samples from each big source
7:
(s; VV1/(M - (s;
// Calculate the number selected samples
8:
Get the zeroth-ores oooer gradient I1.t of the last (LoRA) V-projection using Eq 6.
9:
Calculate the normalized gradient
from historical terms and zeroth-order gradient Vv.,
m.s
Gi.t
uisng Eq
10:
Create mask vector My for top h parameters with largest magnitude. mi.
11:
5 <
arg maxsc şb; Eiey max,es|C - -1 0r,
mi. 2 M2 et /t.
// Solve the submodular facility location optimization problem
end for
13: S. - (5t US U... U...U U.S
14: 0 Ls, (0)
15: end for
B PSEUDO-CODE
Algorithm illustrates the pseudo-code of our CoLM.
C FINE-TUNING SETTINGS
Training datasets. We use the MathInstruct (Yue et al., 2023) dataset for the challenging task of mathematical reasoning. MathInstruct consists of about 260K instruction tuning examples, curated from 14 highly imbalanced open-source math datasets, with broad coverage of mathematical fields and a wide range of difficulty levels. The ratio of the largest to smallest source MathInstruct is almost 300, and the distribution of sources can be found in Fig 4a in the Appendix. Fine-tuning MathInstruct has shown state-of-the-art performance on a variety of standard math evaluation bench- marks. For datasets without specific sources, use three datasets from the SuperGLUE benchmark (Wang et al., 2019) for the classification task: SST-2, CB, and MultiRC. For CB, we use the full training dataset, which consists of 250 examples. For SST-2 and MultiRC, 3K examples for fine-tuning. Notably, compared to the 1K example setting used in (Malladi et al., 2023), we sample more data for datasets with larger sizes because we use much larger batch size.
randomly sample
Training details. Following the setup used in (Yue et al., 2023), we adopt a training regime with learning rate of 2e-5 and cosine scheduler with 3% warm-up period, i.e. the learning rate linearly increases from 0 to 2e-5 over the first 3% of training steps, then follows a cosine decay to 0 at the end of training. We set maximum sequence length of 512. For experiments on MathInstruct, we standardize the number of gradient steps to correspond to 1K, unless explicitly specified. To simulate larger batch size, we have also used a gradient accumulation step of 8 in our experiments. We use LoRA with rank of 128, alpha of 512, and dropout rate of 0.05. For Phi models, apply LoRA to all attention matrices (i.c. QKV-proj) and two fully connected layers while for Zephyr, we apply LoRA all attention matrices (i.e. QKVO -roj). All experiments GPUS. We repeat each experiment three times.
Evaluation datasets. Following (Yue al., 2023), both in-domain and out-of-domain datasets. The in in-domain datasets include GSMSK (Cobbe et al., 2021), MATH (Hendrycks et al., 2021), and NumGLUE (Mishra al., 2022). For the out-of-domain
variety of popular datasets across
18

---

Published as a conference paper at ICLR 2025
Average length
Average court 18717
Count
10a
200
Random
CoLM
200
Longest per
10a
120
Data source
Training iteration
Figure 4: (a) Data distribution of different data sources in MathInstruct. (b) The average completion length examples selected by MvL. random examples and longest examples random batches.
(bs=64)
55.0
1.4
(bs=128)
CoLM (bs=64)
LOSS
47.5
FT (bs=64)
0.8
FT (bs=128)
42.5
CoLM (bs=64)
40.0
100 200 300 400 500 600 700 800 900 1000
2000
4000
6000
8000 10000 12000
Iteration
Training time (s)
(a) Training loss trajectory
(b) Convergence plot
Figure 5: Fine-tuning Phi-2 on MathInstruct. (a) CoLM yields smallest loss throughout the whole training process; (b) CoLM reaches the optimal performance in less training time.
datasets, we include SVAMP (Patel et al., 2021). Mathematics (Davies et al., 2021), and SimulEq "Koncel-Kedziorskt et al., 2016). These datasets collectively cover a wide range of mathemati- cal areas such as algebra, probability, number theory, calculus, and geometry. Furthermore, some questions in these datasets require the application of commonsense, reading comprehension, and multi-step reasoning. All questions are formatted as open-ended.
Evaluation metric. We use the standard evaluation metric for form-formed questions, exact match, which measures the model's accuracy by comparing its generated answers against the correct solu- tions. For an answer to be considered correct, it must match the reference solution precisely. We evaluate under the 0-shot setting with maximum sequence length of 2048 tokens for decoding. The ooggram-of-Thoughh (PoT), falling back Chain-of-Thouss (CoT) prompting the former does not work (Yue et al., 2023).
default prompt
Small sources. Figure 4a visualizes the data distribution of the Mathunstruct dataset. It can be that the dataset is highly imbalanced with most of the samples belonging to large sources (10-1 13). Therefore, use a simple heuristic to consider any data sources whose sizes below the average count as small sources.
D ADDITIONAL FINE-TUNING RESULTS
D. TRAINING LOSS
For fine-tuning LLMs, downstream performance is generally better metric than training loss or perplexity. In addition, perplexity does not always correlate with actual task performance on diverse downstream tasks Liu et al. (2023). Furthermore, the MathInstruct dataset doesn't 't have held-out validation set to calculate the perplexity. For completeness, we added the training loss in Figure 5a. C consistently yields smallest loss throughout the whole training process.
19

---

Published as a conference paper at ICLR 2025
Table 8: Statistics of pre-training mixture.
Table 9: Downstream acc of pre-trained Llama-60M.
Dataset
Weight (%)
Small/Large
Dataset
PT (bs=128)
PT (bs=256)
CoLM (bs= (bs=128)
EuroParl
2.8
Small
Squad
25.4
24.8
25.9
Github
28.2
Large
WiC
50.2
50.7
51.6
HackerNews
5.0
Small
COPA
49.0
52.0
54.0
Exporters
Small
Wikipedia (en)
60.6
T.arge
Avg
41.5
42.5
43,9
PT (bs= (bs=
Val Perplexity
PT (bs= 256)
45
CoLM (bs=128)
40
20
15.
1000
2000
3000
4000
5000
6000
7000
8000
9000
10000
Iteration
Figure 6: Validation perplexity when pre-training Llama-60M on a mixture of the Pile dataset.
D.2 CONVERGENCE
TERMS TRAINING TIME
To highlight the faster convergence rate of CoLM, plotted Figure 3c with training time as the x-axis. Figure 5b still shows that our method converges faster than normal fine-tuning with both smaller and larger batch sizes.
E MEMORY CONSUMPTION
In all our experiments, we used LoRA and a gradient accumulation step of 8, as 4xA40 GPUs did have enough memory to hold bs=
Memory overhead for CoLM. The memory overhead of CoLM stems from three main sources the last layer zeroth-order gradient 0 in Eq 6, the historical terms of Adam in Eq 8, and the pairwise dissimilarity matrix when solving Eq 8. Because with LoRA, the last layer dimension 327K and 2560 before and after sparsification and the batch size is 128, the memory overhead is less than 200MB. When fine-tuning Phi-2 on MathInstruct with 4xA40 GPUs each with 45G memory (except for bs=256) with LoRA, each GPU can have a maximum device batch-size 5, trained all models with gradient accumulation step of with LoRA 4 GPUs. The total batch size = num GPUs x device batch size x gradient accumulation step. Figure 3a illustrates that CoLM (bs=64) outperforms fine-tuning (FT) with bs=256, while requiring 1.8x less memory and being 2.7x faster. Compared to FT bs= CoLM requires 20% less memory, while being 30% faster, and obtains 4% higher accuracy.
F PRE-TRAINING EXPERIMENTS
Settings. We used Llama-60M, which is also used in Zhao et al. (2024b), on mixture of datasets from the Pile dataset Gao et al. (2020). We selected 5 different datasets without copyright in- fringement including EuroParl, Github, HackerNews, NIH Exporter, and Wikipedia (en). For pre- processing the dataset, we divide each dataset into 1024-token chunks and the statistics are given in Table 8 in which weight means the percentage of 1024-token chunks of each dataset in the mixture. Following Zhao et al. (2024b), we pretrained the model for 10K iterations with a max sequence length of 1024 on 4 GPUs. We also used gradient accumulation step of 8 similar to fine-tuning
20

---

Published as a conference paper at ICLR 2025
experiments. For evaluation, calculated the perplexity held-out validation set. In addition, we calculated the down-stream accuracy on 3 datasets Squad Rajpurkar (2016), WiC Pilehvar and Camacho-Collados (2018), and COPA Roemmele al. (2011).
Validation perplexity. Figure 6 illustrates the validation perplexity of pre-trained models at differ- checkpoints during training. CoLM achieves almost the same perplexity pre-training with 2x batch size.
Down-stream performance. Table 9 demonstrates that CoLM improves the downstream accuracy on all 3 evaluation datasets, yielding an improvement of at least 1.4% on average.
21


Reviews:
[{"review_id": "R6vFmP4KLT", "reviewer": "Reviewer_AKmp", "summary": "The paper introduces CoLM (Coresets for Training LLMs), a method for memory-efficient training of large language models (LLMs) using mini-batch coresets that approximate the gradients of larger batches. CoLM addresses specific challenges in LLM training, such as data imbalance, high gradient dimensionality, and compatibility with the Adam optimizer. CoLM addresses them by including all examples from smaller data sources, normalizing gradients for Adam, and using zeroth-order gradient estimates. It reduces memory requirements by 2x while maintaining or surpassing performance compared to larger mini-batches.", "questions": "1. Is zeroth-order necessarily needed? How about estimating using standard first-order methods?\n2. What hypeperparametrs involved in CoLM? Like how many small sources and large sources needed? How much is the cost of tuning those hyper parameters? Are they sensitive to different data structures and types of data?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The approach takes a unique and innovative perspective on memory-efficient training. Unlike typical methods in field, it focuses on data selection and reduces batch size to achieve memory-efficient training.\n2. The paper is clear to understand and well motivated. \n2. The paper presents comprehensive experiments on diverse datasets and models, substantiating the benefits of CoLM over various baselines.", "weaknesses": "1. Computational overhead: Figure 2(b) suggests that the computational overhead for CoLM is notably high, doubling the training time in some cases, whereas typical memory-efficient methods usually incur 5-20% overhead. A comparison of \"memory vs. computational overhead vs. performance\" with other methods would be beneficial.\n2. The paper mainly evaluates CoLM in a fine-tuning context, where model and optimizer state memory are the main bottlenecks. It would be valuable to see CoLM tested in pre-training or continual pre-training settings, where activation memory dominates due to large batch sizes. This could pose additional challenges given complex data structures and extensive amount of training data.\n3. Not enough details of algorithm and discussion on hyper-parameters.", "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": "N67lzgrICj", "reviewer": "Reviewer_dq66", "summary": "This paper proposes a method called Coresets for Training Large Language Models (CoLM) to enable memory-efficient training of large language models (LLMs) by constructing mini-batch \"coresets.\" These coresets are small, representative subsets of mini-batches that approximate the gradient of larger batches, allowing effective model training with reduced memory usage. CoLM addresses challenges such as the high memory demands of large mini-batches, the imbalanced distribution in language data, and the complexities of using the Adam optimizer in high-dimensional gradient spaces.", "questions": "- **Computational Complexity**: CoLM introduces additional computational costs to find mini-batches. Could the authors discuss the complexity and time requirements for this estimation algorithm and expand on it in the main paper?\n\n- **Convergence Rate Comparison**: Figure 3 compares CoLM and fine-tuning (FT) by plotting accuracy against training iterations to highlight CoLM’s faster convergence. Could the authors include a comparison of CoLM with LoRA in Figure 3? It would also be valuable to add perplexity or training loss versus iterations to provide further insight into convergence rates.\n\n- **Time vs. Accuracy Comparison**: Given that CoLM involves additional computation to determine mini-batches, comparing accuracy against time would offer a fairer evaluation alongside FT and LoRA in Figure 3(a).\n\n- **Quantitative Memory Cost Analysis**: While the paper claims CoLM uses 4x less memory, it lacks a quantitative breakdown of memory usage. Could the authors add specific memory costs for LoRA, other low-rank adaptation methods, and CoLM in a table? This would provide stronger evidence of CoLM’s memory efficiency.\n\n- **Clarification on Motivation**: 1) In the introduction, the authors state that fine-tuning a model like Phi-2 with 2.7B parameters and a batch size of 128 requires at least 44 GB of GPU memory. However, fine-tuning LLMs primarily incurs memory costs due to model weights and optimizer states, which together cost approximately three times the model weight. Since Phi-2’s model size is under 7 GB, the total memory cost should be around 22 GB rather than 44 GB. Furthermore, batch size typically only affects activation memory and does not significantly impact memory for optimizer states or model weights, and the activation memory cost is not an issue for PEFT methods such as LoRA. 2) **Comparison to LoRA**: LoRA also reduce optimization memory significantly. Besides, LoRA also achieve significantly speedup since it has less trainable parameters. CoLM achieve similar wall clock time with full fine-tuning, but when it applies to full fine-tuning, it will be much slower than LoRA. Could the authors clarify the motivation of the paper from these aspects? \n\n\n\n- **Title Consistency**: Please align the submission title, \"Memory-efficient Training of Large Language Models with Larger Mini-batches,\" with the manuscript title, \"MINI-BATCH CORESETS FOR MEMORY-EFFICIENT TRAINING OF LARGE LANGUAGE MODELS.\"\n\nI would like to discuss the questions I raised regarding the weaknesses and concerns with the authors. If my concerns are adequately addressed, I would be willing to reconsider my rating.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The presentation is clear, and the paper is easy to follow, with only a few minor typos.\n- The proposed method, CoLM, is straightforward and demonstrates strong empirical performance.", "weaknesses": "- **Limited Base Models**: While the authors mention using Phi-2, Phi-3, and Zephyr, the main results in Tables 1 and 7 only report Phi-2. To more fully demonstrate CoLM’s performance across various models, I recommend including Phi-3 in these tables and consider adding state-of-the-art models, such as LLaMA-3 8B and LLaMA-3.1.\n\n- **Insufficient Discussion of Related Work and Novelty**: One of CoLM’s contributions is gradient normalization for the Adam optimizer, which improves mini-batch approximation by using historical averages. However, previous studies, such as [1], have explored renormalization techniques with mini-batches. It would strengthen the paper if the authors could discuss distinctions between [1] and CoLM to clarify CoLM’s unique contributions.\n\n\n\nReference:\n\n[1] Batch Renormalization: Towards Reducing Minibatch Dependence in Batch-Normalized Models", "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": "kEeETlDEj7", "reviewer": "Reviewer_5bEa", "summary": "The paper identifies three key difficulties when applying previous coreset findings methods (methods for finding small mini-batches that have gradients whose mean closely approximates the full gradient) to fine-tuning LLMs: imbalanced data, Adam's update, and large gradient dimensionality. The authors systematically formalize the difficulties and design methods to address them: by treating small-source examples and large-source examples differently, by comparing gradients in the space of Adam updates, and by using zeroth order gradient estimations methods, sparsifying the gradient, and using an appropriate distance metric. The combination of these improvements makes up their method,  Coresets for Training LLMs (CoLM). In their empirical study, authors demonstrate that their method leads to stronger performance compared to a number of baselines (including strong ones) and uses less memory than comparable methods while adding a reasonable overhead in training time.", "questions": "(Q1) For Figure 3 (c), how large does the batch size need to be to match your method?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well written. It provides clear motivations of the problem and in thoughtfully structured around the three key difficulties of applying coreset finding methods to LLM fine-tuning. This structure makes the paper clear and very easy to follow. \nquality, \n- The authors provide extensive ablation analyses about the great majority (if not all) of the improvements they propose, which is greatly appreciated and helps show they are worthwhile.\n- The authors empirically demonstrate the effectiveness of CoLM for different pre-trained language models and across many datasets (Math Instruct, SST-2, CB, MultiRC).\n- The method often outperforms fine-tuning with twice the batch size.\n- The authors demonstrate that the method achieves a smaller memory overhead than a number of representative baselines.", "weaknesses": "- Using gradient accumulation + LoRA trivially reduces the memory overhead. The paper would benefit from experiments making an explicit comparison to this case. It seems from Fig 2 (a) that you should be able to claim CoLM is still faster in this case, but not showing it and claiming memory efficiency makes the paper seem weak.\n\n- While the memory overhead is reduced due to the smaller batch size, there is no figure showing how much the total memory decreased by making the batch size smaller when fine-tuning with LoRAs using CoLM. Does the batch size reduction account for a substantial portion of the memory consumption or is the memory consumption dominated by the model parameters given the smaller optimizer states required for one LoRA? Providing a figure with the memory consumption of CoLM compared to fine-tuning at different batch sizes is warranted given the author's claims about memory efficiency.\n\n- Do the methods to estimate gradient distances scale to larger vocabulary sizes? Recent SOTA models (Llama 3 has 128k vocab size whereas phi2,phi3,and zephyr have 50k at most) have been using much larger vocabulary sizes and recent work \"Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies\" shows this is an important axis along which to scale LLMs. Simply including one experiment with a larger tokenizer 100k vocab+ would strengthen this aspect.\n\n- *I believe the above concerns are easily addressable with small additional experiments. I would be happy to raise my score if these concerns are addressed. *\n\n*List of Typos*:\n \"to iteratively selecting small mini-batches from larger random batches.\" selecting --> select", "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": "OTAiEED30Z", "reviewer": "Reviewer_TWbY", "summary": "The paper introduces a method in using smaller mini-batches to simulate the performance of training with larger mini-batches, in the context of Large Language Models (LLMs). In particular, the authors propose coresets for training LLMs (CoLM) which involves the following aspects:\n\n1. Including all examples from small sources in the mini-batch coresets\n2. Adapt the mini-batch coreset selection process with ADAM, selecting medoids of big sources.\n3. Use zero-order optimization technique based on MeZO in combination with LoRA to estimate the gradients.\n\nExperiments show that the proposed method achieves comparable performance to regular fine-tuning that uses a much larger batch size, while reducing the memory requirements.", "questions": "Q1: Related to W1 - Is there some sampling or representative subset selection method for the small sources going on behind the scenes?\n\nQ2: Related to W2 - How to determine the boundary between \"small\" and \"large\" partitions, especially with skewed distributions where average count may not serve as a good cutoff?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1: Well written paper with nice motivation to enable both memory efficiency and performance gains by looking at a data perspective.\n\nS2: Strong experimental results that demonstrate the benefits of their approach compared to baselines. The increase in training time due to mini-batch selection also justified given significant improvement in accuracy.\n\nS3: Ablation studies give insight to the relative importance to each of the three components proposed, with emphasis on gradient approximation, which helps to better understand the benefits of the method.", "weaknesses": "W1: There should be some further clarification on how \"all\" examples of small sources can be included in the mini-batch coreset. Based on experiments, figure 4 shows that the data count for even small data sources can be over 1000 examples. Including all the sources in a single mini-batch does not seem to be plausible, but implied from the methodology and Algorithm 1.\n\nW2: Dependency on having access to the true source data distribution may not always be possible. In these cases, it is difficult to distinguish between small and large sources. In addition, how to determine the boundary between \"small\" and \"large\" partitions could be discussed in more detail with experiments.", "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)