Research map / Large language models / Instruction and fine-tuning
Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CoLM, a mini-batch coreset selection method for memory-efficient LLM fine-tuning on imbalanced data mixtures, matching larger-batch gradients under Adam using zeroth-order sparsified V-projection gradients.
Problem
Training LLMs with large mini-batches improves convergence but needs prohibitive GPU memory; existing gradient-matching coreset selection fails for LLMs due to imbalanced language sources, Adam optimizer, and very high gradient dimensionality.
Approach
Keep all examples of small sources, select medoids of big sources using Adam-normalized gradients estimated with MeZO (SPSA) on the last V-projection and sparsified to top-magnitude dimensions, using l1 distance for medoid search.
Key results
- CoLM with bs=64 (selected from 128) reaches 56.6% average accuracy on Phi-2 MathInstruct, outperforming FT bs=256 (55.3%) while using 45% less memory and running 2.7x faster.
- On SuperGLUE with Phi-2, CoLM using clusters from the fine-tuned model reaches 81.0% average vs 78.9% for FT bs=128.
- In Phi-2 MathInstruct ablation, including all small-source examples improves average accuracy by roughly 3 points.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CoLM-selected mini-batches of bs=64 instead of random FT batches of bs=256 | Phi-2 fine-tuning with LoRA · MathInstruct | Avg All accuracy: $55.3_{\pm 0.5}$ -> $56.6_{\pm 0.9}$ | Helped |
| Use clusters from a fine-tuned model instead of updating clusters during fine-tuning | CoLM on Phi-2 · SuperGLUE | Avg accuracy: $78.7_{\pm 1.9}$ -> $81.0_{\pm 2.6}$ | Helped |
| Keep all examples from small sources in the coreset (ablation) | CoLM · MathInstruct | Avg accuracy: $51.5_{\pm 0.5}$ -> $54.6_{\pm 0.6}$ | Helped |
| Select medoids separately for each big source (ablation) | CoLM with small sources retained · MathInstruct | Avg accuracy: $54.6_{\pm 0.6}$ -> $55.1_{\pm 0.5}$ | Helped |
| Normalize gradients for Adam (ablation) | CoLM with separate big-source medoids · MathInstruct | Avg accuracy: $55.1_{\pm 0.5}$ -> $56.6_{\pm 0.9}$ | Helped |
| Use uniform rather than weighted example weights (ablation) | CoLM · MathInstruct | Avg accuracy | Helped |
| Use ℓ1 distance instead of cosine distance (ablation) | Medoid selection · MathInstruct | Avg accuracy: $51.3_{\pm 0.9}$ -> $51.5_{\pm 0.5}$ | No clear effect |
| Retain dimensions with the largest gradient magnitude (ablation) | CoLM sparsification · MathInstruct | Avg accuracy: $55.4_{\pm 1.7}$ -> $56.6_{\pm 0.9}$ | Helped |
| Set sparsity dimension to 2560 rather than 1280 (ablation) | CoLM · MathInstruct | Avg accuracy: $54.1_{\pm 1.3}$ -> $56.6_{\pm 0.9}$ | Helped |
| Use sparsified MeZO gradients instead of projected actual gradients (ablation) | CoLM gradient estimation · MathInstruct | Avg accuracy: $55.2_{\pm 0.4}$ -> $56.6_{\pm 0.9}$ | Helped |
| Use the last V-projection instead of the last FC layer (ablation) | CoLM gradient estimation · MathInstruct | Avg accuracy: $54.5_{\pm 0.1}$ -> $56.6_{\pm 0.9}$ | Helped |
| Use CoLM with bs=128 instead of standard pretraining with bs=256 | Llama-60M · Pile mixture downstream tasks | Avg downstream accuracy: 42.5 -> 43.9 | Helped |
Surprising findings
- Selecting examples with shorter completions than average (not longest) works better than selecting longest-completion examples, contrary to prior work.
- Random sparsification and weight-magnitude sparsification perform comparably, while largest-gradient-magnitude sparsification is best.
Limitations
- Theoretical guarantees rely on strong assumptions (dense neighborhoods, Lipschitz facility location objective) that may not hold for real LLM gradients.
- Experiments are mostly with LoRA; gains on full-parameter training and pretraining are limited to a small Llama-60M study.
- Small-source definition is a simple heuristic (below average source size), and SuperGLUE source discovery requires clustering with a warm-up model.
Experimental setup
Datasets: MathInstruct SuperGLUE (SST-2, CB, MultiRC) The Pile (EuroParl, Github, HackerNews, NIH Exporters, Wikipedia)
Benchmarks: GSM8K MATH NumGLUE SVAMP Mathematics SimulEq SuperGLUE Squad WiC COPA
Metrics: exact match accuracy validation perplexity wall-clock time memory usage
Baselines: FT with bs=64 FT with bs=128 FT with bs=256 Big Loss Grad Norm Middle Perplexity Completion Length Least Confidence Facility Location
Builds on: gradient matching coresets (CRAIG) MeZO GaLore LoRA Big Loss selection
Closest papers by meaning
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.92)
- Benchmarking Optimizers for Large Language Model Pretraining (NeurIPS 2026 · similarity 0.91)
- AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning (ICLR 2025 · similarity 0.91)
- A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models (ICML 2025 · similarity 0.90)
- MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.89)
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (ICML 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.