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

arXiv: https://arxiv.org/abs/2407.19580 · HTML: https://mlsanity.com/paper/2407.19580/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 53 references that are in the map.
- First posted on arXiv 28 Jul 2024.

## Builds on (most cited first, 30 of 37)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [LIMA: Less Is More for Alignment](https://mlsanity.com/paper/2305.11206/) (2023)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [DoRA: Weight-Decomposed Low-Rank Adaptation](https://mlsanity.com/paper/2402.09353/) (2024)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Textbooks Are All You Need II: phi-1.5 technical report](https://mlsanity.com/paper/2309.05463/) (2023)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning](https://mlsanity.com/paper/2309.05653/) (2023)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [TinyStories: How Small Can Language Models Be and Still Speak Coherent English?](https://mlsanity.com/paper/2305.07759/) (2023)
- [WiC: the Word-in-Context Dataset for Evaluating Context-Sensitive Meaning Representations](https://mlsanity.com/paper/1808.09121/) (2018)
- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [Not All Samples Are Created Equal: Deep Learning with Importance Sampling](https://mlsanity.com/paper/1803.00942/) (2018)
- [From Quantity to Quality: Boosting LLM Performance with Self-Guided Data Selection for Instruction Tuning](https://mlsanity.com/paper/2308.12032/) (2023)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [D4: Improving LLM Pretraining via Document De-Duplication and Diversification](https://mlsanity.com/paper/2308.12284/) (2023)
- [NumGLUE: A Suite of Fundamental yet Challenging Mathematical Reasoning Tasks](https://mlsanity.com/paper/2204.05660/) (2022)
- [MoDS: Model-oriented Data Selection for Instruction Tuning](https://mlsanity.com/paper/2311.15653/) (2023)
- [Accelerating Deep Learning by Focusing on the Biggest Losers](https://mlsanity.com/paper/1910.00762/) (2019)
- [Long Is More for Alignment: A Simple but Tough-to-Beat Baseline for Instruction Fine-Tuning](https://mlsanity.com/paper/2402.04833/) (2024)
- [Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models](https://mlsanity.com/paper/2210.14199/) (2022)
- [Chain of LoRA: Efficient Fine-tuning of Language Models via Residual Learning](https://mlsanity.com/paper/2401.04151/) (2024)

## Built on by (most cited first, 12 of 12)

- [Guiding Data Collection via Factored Scaling Curves](https://mlsanity.com/paper/2505.07728/) (2025)
- [Learning as Reasoning Unfolds: Progressive Rollout Allocation for Efficient Reinforcement Learning](https://mlsanity.com/paper/2607.22002/) (2026)
- [Fine-tuning Quantized Neural Networks with Zeroth-order Optimization](https://mlsanity.com/paper/2505.13430/) (2025)
- [Understanding the Role of Training Data in Test-Time Scaling](https://mlsanity.com/paper/2510.03605/) (2025)
- [Theoretical Perspectives on Data Quality and Synergistic Effects in Pre- and Post-Training Reasoning Models](https://mlsanity.com/paper/2603.01293/) (2026)
- [GUARD: Guided Unlearning and Retention via Data Attribution for Large Language Models](https://mlsanity.com/paper/2506.10946/) (2025)
- [UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees](https://mlsanity.com/paper/2604.10952/) (2026)
- [Efficient Content-based Recommendation Model Training via Noise-aware Coreset Selection](https://mlsanity.com/paper/2601.10067/) (2026)
- [Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs](https://mlsanity.com/paper/2602.00576/) (2026)
- [Minibatch Selection for Language Models via Partition Matroid Constrained Gradient Matching](https://mlsanity.com/paper/2606.07954/) (2026)
- [Unifying Data, Memory, and Compute Efficiency in LLM training: A Survey](https://mlsanity.com/paper/2606.10706/) (2026)
- [Every Batch Is Its Own Validation Set: Leave-One-Out Gradient Matching for Online Data Selection in LLM Fine-Tuning](https://mlsanity.com/paper/2610.00436/) (2026)

## Abstract (arXiv)

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. Our code is available at https://github.com/BigML-CS-UCLA/CoLM.