# Balanced Data Sampling for Language Model Training with Clustering

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

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 43 references that are in the map.
- First posted on arXiv 22 Feb 2024.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Llemma: An Open Language Model For Mathematics](https://mlsanity.com/paper/2310.10631/) (2023)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [On The Power of Curriculum Learning in Training Deep Networks](https://mlsanity.com/paper/1904.03626/) (2019)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [A Pretrainer's Guide to Training Data: Measuring the Effects of Data Age, Domain Coverage, Quality, & Toxicity](https://mlsanity.com/paper/2305.13169/) (2023)
- [D4: Improving LLM Pretraining via Document De-Duplication and Diversification](https://mlsanity.com/paper/2308.12284/) (2023)
- [Jina Embeddings 2: 8192-Token General-Purpose Text Embeddings for Long Documents](https://mlsanity.com/paper/2310.19923/) (2023)
- [DoGE: Domain Reweighting with Generalization Estimation](https://mlsanity.com/paper/2310.15393/) (2023)
- [The MiniPile Challenge for Data-Efficient Language Models](https://mlsanity.com/paper/2304.08442/) (2023)
- [SlimPajama-DC: Understanding Data Combinations for LLM Training](https://mlsanity.com/paper/2309.10818/) (2023)
- [Efficient Passage Retrieval with Hashing for Open-domain Question Answering](https://mlsanity.com/paper/2106.00882/) (2021)

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

- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Unleashing the Power of Data Tsunami: A Comprehensive Survey on Data Assessment and Selection for Instruction Tuning of Language Models](https://mlsanity.com/paper/2408.02085/) (2024)
- [Maximize Your Data's Potential: Enhancing LLM Accuracy with Two-Phase Pretraining](https://mlsanity.com/paper/2412.15285/) (2024)
- [Measuring Data Diversity for Instruction Tuning: A Systematic Analysis and A Reliable Metric](https://mlsanity.com/paper/2502.17184/) (2025)
- [VersiCode: Towards Version-controllable Code Generation](https://mlsanity.com/paper/2406.07411/) (2024)
- [QuaDMix: Quality-Diversity Balanced Data Selection for Efficient LLM Pretraining](https://mlsanity.com/paper/2504.16511/) (2025)
- [BLAST: Balanced Sampling Time Series Corpus for Universal Forecasting Models](https://mlsanity.com/paper/2505.17871/) (2025)
- [SampleMix: A Sample-wise Pre-training Data Mixing Strategey by Coordinating Data Quality and Diversity](https://mlsanity.com/paper/2503.01506/) (2025)
- [Training an LLM-as-a-Judge Model: Pipeline, Insights, and Practical Lessons](https://mlsanity.com/paper/2502.02988/) (2025)
- [IDEAL: Data Equilibrium Adaptation for Multi-Capability Language Model Alignment](https://mlsanity.com/paper/2505.12762/) (2025)
- [Enhancing Training Data Attribution with Representational Optimization](https://mlsanity.com/paper/2505.18513/) (2025)
- [TiKMiX: Take Data Influence into Dynamic Mixture for Language Model Pre-training](https://mlsanity.com/paper/2508.17677/) (2025)
- [Your Pretrained Model Tells the Difficulty Itself: A Self-Adaptive Curriculum Learning Paradigm for Natural Language Understanding](https://mlsanity.com/paper/2507.09758/) (2025)
- [RefactorAssist: Agentic Refinement for Reliable Code Refactoring](https://mlsanity.com/paper/2608.00924/) (2026)
- [CDT: A Comprehensive Capability Framework for Large Language Models Across Cognition, Domain, and Task](https://mlsanity.com/paper/2509.24422/) (2025)
- [Optimizing the Training Diet: Data Mixture Search for Robust Time Series Forecasting](https://mlsanity.com/paper/2512.11546/) (2025)
- [Dynamic Cluster Data Sampling for Efficient and Long-Tail-Aware Vision-Language Pre-training](https://mlsanity.com/paper/2604.27932/) (2026)
- [MADS: Model-Aware Diverse Core Set Selection for Instruction Tuning](https://mlsanity.com/paper/2605.30857/) (2026)
- [TACTICS: Taxonomy-Aware Intelligent Corpus Sampling for Machine Translation](https://mlsanity.com/paper/2609.17956/) (2026)
- [Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling](https://mlsanity.com/paper/2609.32484/) (2026)
- [Which LLM to pick? Online Active Model Selection for Large Language Models](https://mlsanity.com/paper/2610.01592/) (2026)
- [ML-CLAPPS: A Probabilistic Framework for Imbalanced Multilabel Learning under Class Prioritization](https://mlsanity.com/paper/2610.15816/) (2026)

## Abstract (arXiv)

Data plays a fundamental role in the training of Large Language Models (LLMs). While attention has been paid to the collection and composition of datasets, determining the data sampling strategy in training remains an open question. Most LLMs are trained with a simple strategy, random sampling. However, this sampling strategy ignores the unbalanced nature of training data distribution, which can be sub-optimal. In this paper, we propose ClusterClip Sampling to balance the text distribution of training data for better model training. Specifically, ClusterClip Sampling utilizes data clustering to reflect the data distribution of the training set and balances the common samples and rare samples during training based on the cluster results. A repetition clip operation is introduced to mitigate the overfitting issue led by samples from certain clusters. Extensive experiments validate the effectiveness of ClusterClip Sampling, which outperforms random sampling and other cluster-based sampling variants under various training datasets and large language models.