Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures

We sincerely appreciate the reviewer for acknowledging the clear motivation of our work, our strong experimental results, and the insightful ablation studies.

W1. Samples from small sources in mini-batches

As explained in detail in lines 225-237 (also see Line 3 of Algorithm 1 in the appendix), we first randomly draw a mini-batch of samples. Then, the algorithm finds a smaller mini-batch from the examples of the large mini-batch (not from the entire data).

Specifically, we keep all samples of small sources within the larger mini-batch. Then, for every big source, we only select its central examples (medoids) in the gradient space from those within the larger mini-batch. That is, centrality is calculated w.r.t the examples within the mini-batch, not the full data.

This is based on our thm 4.2, which shows that: for every big source, central examples w.r.t the larger mini-batch are also central w.r.t the full source, but this does not hold for small sources.

---

W2. Distinguish between small and large sources

We discussed the application of our method to datasets without source info in Sec 5.4. When the data source is unknown or unavailable, one can perform clustering of the target or a smaller proxy model’s hidden states to group data into different “sources”. Then we identify small sources in the same way as if we have the source information (as discussed in lines 238-239). Specifically, having identified c groups, we consider small sources as those with less than |V|/c examples. Table 7 shows our experiments on the SuperGLUE benchmark, and confirms that CoLM effectively improves the performance without any source information as well.

---

Q1. Sampling or selection for small sources

Please refer to our answer in W1.

---

Q2. Small and large sources

Please refer to our answer in W2.

Q2. Skew distributions

Note that, as written in Line 177, the ratio of the largest to smallest source size in the MathInstruct dataset is 300, which is sufficiently large. Indeed, the number of small sources is a property of the dataset and may vary depending on the dataset used. Nevertheless, we expect that our strategy which considers any data sources whose sizes are below the average account as small sources might still work in other cases.

We sincerely thank the reviewer for recognizing the clarity and thoughtful structure of our paper, the extensive ablation analyses, and the strong empirical results across diverse settings.

W1. CoLM and other memory-efficient methods

Please note that our method can be easily stacked with LoRA and gradient accumulation while lowering the memory and training time compared to training with large batches. Indeed, all our experiments are done with LoRA (see lines 365-366) and gradient accumulation (clarified in our revision - Training details Appendix B). We also note that (i) LoRA reduces the memory but harms the performance; (ii) gradient accumulation reduces the memory but increases the training time; (iii) our method reduces the memory and training time while considerably improving the performance when applied to full training, or stacked with LoRA and gradient accumulation. A more detailed discussion can be found in our general comment.

---

W2. Memory overhead of CoLM

Please see our general comment about actual memory consumption. For Table 1, the activation memory per sample is about 3.9GB while CoLM additional memory is less than 200MB. CoLM effectively reduces the activation memory which scales with the batch size. For example, for fine-tuning Phi-2 with LoRA, CoLM (128->64) requires 20% less memory over bs=128, and CoLM (1024->512) requires 40% memory over bs=1024.

---

W3. Larger vocab size & Llama-3

Our gradient estimation is not affected by the vocab size as we didn’t use the gradient at the lm_embed layer (which is not LoRA fine-tuned) but we used the last V projection matrix (Line 274). In addition, as shown in our general comment, our method is effective for Llama-3 models.

Typos. Thanks for pointing it out. We fixed it in our revision.

---

Q1. Figure 3c

We ran normal fine-tuning with bs=512 and it matched the performance of our method. The average performance over 3 seeds is 56.7% $\pm$ 0.3. Please check the new Fig 3c in our revision.

Q3. Time vs. Accuracy Comparison

We plotted Fig 3a with training time as the x-axis and added it to Figure 7 in revised Appendix D. It still shows that our method converges faster than normal fine-tuning with both smaller and larger batch sizes.

---

Q4. Quantitative Memory Cost Analysis

We reported our actual memory in our general comment. CoLM reduces the memory requirement by up to 2x, compared to fine-tuning with 4x batch size. As discussed in the general comment, The memory overhead of our method is less than 200MB for fine-tuning Phi-2 with LoRA.

Q5. Clarification on Motivation

Memory cost of fine-tuning Phi-2. Consider an LLM with N billion parameters and full fine-tuning with mixed precision. The optimizer maintains a copy of the model in FP32 precision, consuming 4N memory. The gradients, momentum, and second-moment vectors are all stored in FP32 precision with each requiring 4N memory. Consequently, the total memory required is at least 16N = 43.2 GB if N = 2.7 for Phi-2.

Activation cost for LoRA. LoRA only reduces the optimizer state memory but increases the model weights and activation memory, by adding additional low-rank matrices. The activation memory still scales with larger device batch sizes.

Stack CoLM with other memory-efficient methods. We note that our method is not a replacement for LoRA and other memory-efficient methods and can be easily stacked with them to further reduce the memory requirements. Indeed, as mentioned in lines 365-366 and the caption of table 1, all our experiments are done with LoRA. Fig 2b shows that when stacked with LoRA, CoLM is indeed much faster than training with the larger batch. Fig 2b shows that our method (CoLM 128->64) is actually 30% faster than normal fine-tuning with a batch size of 128. Notably, unlike existing memory-efficient training methods that harm accuracy, our method outperforms training with bs=128 by around 4%!

---

Q6. Title Consistency

Thanks for pointing this out. We will change the submission title in the final version.

We sincerely appreciate the reviewer for acknowledging the clarity of our paper, as well as the straightforward nature of our proposed method, and its strong empirical performance.

W1. Limited Base Models

Figure 2a shows that CoLM outperforms fine-tuning Phi-3 and Zephyr with larger batches on MathInstruct. We report the exact numbers here and have added them to the appendix of our revised version.

To further show the effectiveness of our method across different state-of-the-art models, we ran our method for the LLaMa-3 series and reported the results in our general comment. We see that CoLM consistently outperforms training with small and larger batch sizes, by up to 5.7% and 4% respectively.

---

W2. Insufficient Discussion of Related Work and Novelty

Please note that we use our normalized gradient only to select smaller mini-batches, not to train with the normalized gradients. Indeed, our method does not interfere with the optimizer. The optimizer can apply its own normalization, including the gradient normalization introduced by [1] to train on our selected smaller mini-batches. Hence, our contribution is orthogonal to [1].

[1] Batch Renormalization: Towards Reducing Minibatch Dependence in Batch-Normalized Models

---

Q1. Computational Complexity

The additional time for our method stems from three sources: the time to get the zeroth-order gradient of the last layer (Eq 6), the normalized gradient, the sparsified gradient, and solving the submodular facility location optimization problem (Eq 8). The most time-consuming step is getting the zeroth-order gradients, which can be calculated efficiently in only one forward pass as we discussed in Lines 292-296. Figure 2b reports the wall-clock time including the additional time of the above steps, and confirms that CoLM is 30% faster than training with bs=128 while using less memory, and outperforms it by 4%. In detail, one iteration of bs=64, bs=128, and CoLM (128->64) in Table 1 cost 3.15, 6.3, and 4.7 (s), respectively.

---

Q2. Convergence Rate Comparison

As mentioned in our experiments section (see lines 365-366 and caption of table 1), in all our experiments including Fig 3, we have already used LoRA for all methods. Please note that LoRA doesn’t make convergence faster, and often harms the final accuracy of the model [2, 3]. In contrast, our method improves the convergence and performance of the model. For fine-tuning, downstream performance is often a better indicator of the model performance than training loss or perplexity. Hence, we reported the downstream accuracy in Fig 3. In our revision, we also added the training loss to Figure 6 in revised Appendix D. At the end of fine-tuning, the training loss of FT (bs=64), FT (bs=128), and CoLM (bs=64) are 0.74, 0.57, and 0.47, respectively. Because the MathInstruct dataset doesn’t have a held-out validation set to calculate the perplexity, we can only calculate the training perplexity. For a fixed-length model, the train perplexity can be approximated by taking the exponential of the training loss. Therefore, the same conclusion holds for training perplexity.

[2] Evans, Paul. "How fast do economics converge?." Review of Economics and Statistics 79.2 (1997): 219-225.

[3] Liu, Shih-Yang, et al. "Dora: Weight-decomposed low-rank adaptation." arXiv preprint arXiv:2402.09353 (2024).

Thank you for your response. We believe your concerns are already addressed in our original submission and rebuttal. We will iterate over these points below:

Regarding our method, point (2) is not accurate, as we “do not revise the Adam optimizer”. As discussed in Sec 4.2 (e.g. see line 257: “For selecting mini-batch coresets for training with Adam …”), and our response to your review (see W2), our method does not change or revise the Adam optimizer. We only use the normalized gradient to select the subset, not to train with them. In fact, (2) and (3) in addition to (1) allow “finding the subsets”.

1. Please see Lines 365-366 in our experiments under the “training details header”, where the first sentence mentions that “our experiments use LoRA”. This is also clearly mentioned in the caption of Table 1, where the first sentence mentions that “Accuracies on in-domain and out-of-domain datasets when fine-tuning Phi-2 with LoRA”. We’ll highlight this in the caption of the figures too.

2. The one-shot baselines are already briefly discussed in our related work section (lines 108-110), where we mentioned middle perplexity as the most successful one-shot method. However, as discussed there, such baselines are most useful for pretraining, not fine-tuning LLMs. The one-shot baselines are further listed and cited under the “baselines header” in our experiments section (lines 369-374). In fact, the one-shot methods mentioned by the reviewer are already the baselines we compared to in our Table 1, and confirmed the clear advantage of our method! It is also mentioned in the caption of Table 1, that “we compare to one-shot methods”: CL (completion length), BL (big loss), GN (gradient norm), LC (low confidence), FL or SO (submodular facility location), MP (middle perplexity). Consistent with the discussion in our related work section, MP is the best one-shot baseline and CoLM outperforms MP by 3%.

3. Please see line 365 of our experiments section, the first line of the “training details header” mentions: “We use LoRA with a rank of 128, alpha of 512”. This is even mentioned in our introduction (see line 79 which mentions that we use LoRA in our experiments.). Please note that alpha in LoRA doesn’t affect the memory consumption of LoRA. For precision, we used the default precision of each pre-trained model, i.e. float16 for Phi-2, bfloat16 for Phi-3, Zephyr-3B, and Llama-3. We will add this to our version.

4. Please note that there is no memory saving corresponding to different parts of our method. As discussed in our general comment, our method reduces the activation memory by selecting smaller mini-batches. Without part (3) which reduces the gradient dimensionality, one cannot select a subset (this is discussed in lines 183-187 before we discuss our method, and is iterated over in lines 268-269 before we start Sec 4.3): “This is because in such a high dimensional space, distances become vacuous”. Our Table 5 shows other ways of finding low-dimensional gradient estimates to enable subset selection. Again, note that (1-3) are all to select a subset not to change the optimizer or reduce the weight or optimizer memory, etc. The total memory overhead of our method is less than 200MB, from which 82.5MB is due to using historical terms in part (2) and the rest is due to part (3). Thus, one can use CoLM without (2) to slightly reduce the memory overhead but trading off 1% of the performance (see Table 2). (1) doesn’t introduce any overhead. While the total memory overhead of our method is negligible, per reviewer’s suggestion, we’ll add the above numbers to our revision.

5. As discussed in our general comment, our method reduces the batch size and thus reduces the activation memory, not the optimizer memory. Hence, it is complementary to LoRA and other parameter efficient methods. In our general comment cont., we discussed in detail how much memory CoLM saves over LoRA, under different settings (this is also added to our revised version). Note that all the experiments in our rebuttal are all done with LoRA and show the memory saving over LoRA. Particularly, for larger batches, our method reduces the (activation) memory requirement by around 2x, and this includes the 200MB overhead of our method. Notably, our method improves the performance, while LoRA and other parameter-efficient methods harm the performance.

Given that most of your concerns were already addressed in our original submission or rebuttal, we do hope that you consider giving our paper stronger support. We believe our work is complete, and our novel contributions enable training LLMs with larger batch size and superior performance. Thus, our work deserves more than a borderline score.

Q1. Zeroth-order vs First-order gradient

As discussed at the beginning of Sec 4.3 (line 272), the high-dimensional gradients of LLMs are very noisy. The zeroth-order gradients have two benefits: (1) they are smoother than actual gradients calculated with backprop, hence allowing finding higher-quality subsets (see our ablation study in Table 5 that confirms that using zeroth-order gradients outperform standard first-order gradients); and (2) as discussed in lines 293-298, they can be calculated quickly using just one forward pass in a memory efficient manner (which is our objective).

---

Q2. Hyper-parameters

Our only hyper-parameter is the sparsity level. Our ablation study in Table 4 showed the dimensionality of the last hidden state works well, and hence in practice one does not need to tune this parameter.

Q2. Data source & data structure

NLP datasets are usually a mixture of data from different sources. For example, the MathInstruct data contains 14 sources, including TheoremQA, Camel-Math, College-Math, etc. Thus, the number of small sources is a property of the dataset and is not a parameter to tune by our method. As discussed in lines 238-239, we simply consider any data sources whose sizes are below the average account as small sources. We expect our method to work on datasets with various numbers of sources and also on datasets without any specific sources, as we confirmed in Sec 5.4 (see Table 7), on the SuperGLUE benchmark without specific sources.

We sincerely thank the reviewer for recognizing our innovative approach, the strong motivation of our work, as well as our comprehensive experiments across diverse datasets and models.

W1. Computational overhead

There might be a misread of our Fig 2b. Our method (CoLM) with bs=64 simulates and outperforms training with bs=128, using lower memory requirements. Hence, our method should be compared with bs=128. Fig 2b shows that (CoLM) is actually 30% faster than normal fine-tuning with bs=128. Notably, unlike existing memory-efficient training methods that harm accuracy, our method outperforms training with bs=128 by around 4%! Besides, our method can be easily stacked with existing memory-efficient methods. Indeed, as mentioned in the Experiments section (line 365 and Table 2), all our experiments were done with LoRA and improve the memory efficiency over LoRA.

---

W2. Pre-training experiments

Our method mimics training with larger batch sizes (by matching their gradient), and thus can deal with complex and big datasets, as normal training with larger batch sizes does.

Indeed, for pretraining when the batch size is larger, the activation memory dominates the optimizer state and model weight memory. Thus, as discussed in our general comment, CoLM reduces the memory requirement even further when batch size is larger.

Pre-training experiments. To demonstrate the effectiveness of our method in the pre-training setting, we pre-trained Llama-60M (used in [1]) on a mixture of datasets from the Pile dataset [2]. We selected 5 different datasets without copyright infringement including EuroParl, Github, HackerNews, NIH Exporter, and Wikipedia (en). For preprocessing the dataset, we divide each dataset into 1024-token chunks, and the statistics are given in the table below.

Following [1], we pretrained the model for 10K iterations with a max sequence length of 1024 on 4 GPUs. For evaluation, we calculated the perplexity on a held-out validation set. In addition, we calculated the downstream accuracy on 3 datasets Squad, WiC, and COPA.

In terms of validation perplexity (cf. Figure 9 in revised Appendix F), CoLM achieves a similar value as FT (bs=256), which is lower than FT (bs=128) at the end of training.

The following table demonstrates that CoLM improves the downstream accuracy

on all 3 evaluation datasets, yielding an improvement of at least 1.4% on average.

We believe that this is a great addition to our work, and we will add larger pre-training experiments to our final version. Thank you for the suggestion!

[1] Zhao, Jiawei, et al. "Galore: Memory-efficient llm training by gradient low-rank projection." arXiv preprint arXiv:2403.03507 (2024).

[2] Gao, Leo, et al. "The pile: An 800gb dataset of diverse text for language modeling." arXiv preprint arXiv:2101.00027 (2020).

---

W3. Algorithm & Hyper-parameters

Our only hyper-parameter is the sparsity level. As we showed in Table 4, using the same number of dimensions as that of the model’s last hidden state yields the best results. Thus, in practice one does not need to tune this hyper-parameter. The pseudo-code of our algorithm is illustrated in revised Appendix C.

CoLM Reduces the Memory Requirement and Obtains a Higher Accuracy, when Applied with (i) Smaller Gradient Accumulation Steps, or (ii) Larger Batch Sizes

(i) Smaller Gradient Accumulation Steps. The following table shows 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=128, CoLM requires 20% less memory, while being 30% faster, and obtains 4% higher accuracy.

Gradient accumulation linearly increases the training time. With larger GPU memory (H100 94GB), one can use no gradient accumulation (step = 1) to significantly speed up training. In this case, the activation memory dominates the model weight and optimizer state memory (even when training with LoRA). Therefore, CoLM can reduce the memory consumption over bs=128 by around 2x.

(ii) Larger Batch Sizes. When the batch size is larger, the activation memory dominates the optimizer state and weight memory. The following table shows how the estimated memory requirement (per GPU) of CoLM for fine-tuning Phi-2 increases when using larger batch sizes (also shown in Figure 8 in our revised Appendix E). Notably, with 4xA40 and using gradient accumulation step=8 with LoRA, for bs=2048, CoLM (bs=1024) will require almost 2x less memory than fine-tuning with bs=2048. A larger batch size is useful in particular for pretraining. Our new experiments for pre-training Llama on Pile, confirm the benefits of CoLM to pretraining.

(iii) Llama-3 experiments. To demonstrate the effectiveness of our method to the recent state-of-the-art models with larger vocabulary size, we fine-tuned Llama-3 models with LoRA using the same setting as Phi-2 on MathInstruct. We report the average accuracy on in- and out-of-domains below. We see that CoLM with bs=64 outperforms normal fine-tuning with bs=128, with smaller memory requirements.