SampleMix: A Sample-wise Pre-training Data Mixing Strategey by Coordinating Data Quality and Diversity

2025-03-03

AI Summary

Current LLM pretraining data mixing predominantly employs coarse-grained domain-level weighting, neglecting inter-domain overlaps and sample-level heterogeneity—thus failing to ensure balanced global diversity and quality. To address this, we propose the first bottom-up, sample-level mixing paradigm: it jointly models per-sample quality and cross-domain diversity, and derives an optimal sampling distribution via global optimization. Unlike conventional within-domain uniform sampling, our approach enables fine-grained, cross-domain collaborative data selection. Experiments demonstrate that our method significantly outperforms state-of-the-art domain-weighting baselines across multiple downstream tasks and perplexity metrics. Notably, it achieves baseline performance in only 1.4–2.1× fewer training steps, substantiating substantial improvements in both pretraining efficiency and effectiveness through sample-level mixing.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Sampling/Simulation-based Search
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Existing pretraining data mixing methods for large language models (LLMs) typically follow a domain-wise methodology, a top-down process that first determines domain weights and then performs uniform data sampling across each domain. However, these approaches neglect significant inter-domain overlaps and commonalities, failing to control the global diversity of the constructed training dataset. Further, uniform sampling within domains ignores fine-grained sample-specific features, potentially leading to suboptimal data distribution. To address these shortcomings, we propose a novel sample-wise data mixture approach based on a bottom-up paradigm. This method performs global cross-domain sampling by systematically evaluating the quality and diversity of each sample, thereby dynamically determining the optimal domain distribution. Comprehensive experiments across multiple downstream tasks and perplexity assessments demonstrate that SampleMix surpasses existing domain-based methods. Meanwhile, SampleMix requires 1.4x to 2.1x training steps to achieves the baselines' performance, highlighting the substantial potential of SampleMix to optimize pre-training data.

Problem

• Addresses limitations in domain-wise pretraining data mixing for LLMs.
• Proposes sample-wise mixing to optimize data quality and diversity.
• Demonstrates improved performance across tasks with SampleMix strategy.

Innovation

• Sample-wise data mixing for LLMs
• Global cross-domain sampling strategy
• Dynamic domain distribution optimization