The appropriate group to review this topic would be Senior Machine Learning Research Scientists and Lead Data Architects specializing in Large Language Model (LLM) pre-training, data engineering, and optimization.
Abstract:
This research introduces CLIMB (CLustering-based Iterative Data Mixture Bootstrapping), an automated framework designed to discover, evaluate, and refine optimal data mixtures for LLM pre-training without relying on manual domain labels. The methodology utilizes embedding-based clustering to categorize massive web-scale datasets into semantic groups, followed by an iterative bi-level optimization process. By training small, computationally efficient proxy models on sampled mixtures, CLIMB fits a performance predictor (LightGBM) to prune suboptimal candidates and converge on high-performing configurations.
Experimental results demonstrate that a 1B-parameter model mid-trained on 400B tokens using the CLIMB mixture exceeds the performance of Llama-3.2-1B by 2.0% across 12 general reasoning benchmarks. The framework is shown to be highly effective for both general-purpose reasoning and domain-specific adaptation, yielding a 5% improvement in specialized areas like Social Sciences compared to random sampling. Alongside the framework, the authors release ClimbLab, a 1.2-trillion-token clustered corpus, and ClimbMix, a 400-billion-token optimized dataset for high-efficiency pre-training.
CLIMB: Automated Data Mixture Optimization for Large Language Models
[Context] Data Mixture Challenges: Identifying optimal proportions of diverse data sources (e.g., Common Crawl vs. curated STEM) is critical for performance but typically relies on labor-intensive manual labeling or static heuristics.
[Method] Semantic Data Preprocessing: CLIMB maps raw documents into an embedding space using models like stella_en_400M_v5, followed by K-means clustering (initially 1,000 clusters) and hierarchical merging to create semantically distinct "super-clusters."
[Method] Iterative Search Framework: The system employs a "coordinate descent" approach to mixture weights, alternating between sampling new configurations and fitting a LightGBM predictor to estimate downstream task performance.
[Method] Computational Efficiency: To minimize costs, the framework uses lightweight proxy models (62M to 350M parameters) to evaluate mixtures. These proxies effectively estimate the performance gains applicable to larger target models (1B+).
[Results] Performance Benchmarks: A 1B model trained with CLIMB-optimized mixtures achieves a 60.41% average across benchmarks, outperforming state-of-the-art baselines like DoReMi and RegMix.
[Results] Efficient Scaling: When pre-training from scratch on 400B tokens, the ClimbMix dataset demonstrates a superior scaling trend compared to Nemotron-CC, SmolLM, and FineWeb-Edu.
[Analysis] Domain-Specific Optimization: CLIMB allows for targeted "mid-training," where optimizing for specific MMLU domains (STEM, Humanities, Social Sciences) consistently outperforms baseline random sampling and single-iteration searches.
[Analysis] Compute Allocation (4:2:1 Ratio): Researchers found that allocating search compute across three iterations in a 4:2:1 ratio (64, 32, then 16 searches) balances exploration and exploitation better than "fat" (single-iteration) or "tall" (many-iteration) search trees.
[Artifacts] ClimbLab & ClimbMix: The study introduces a 1.2T token research playground (ClimbLab) organized into 20 semantic clusters (e.g., Mathematics, Biotechnology, Legal Content) and a compact 400B token high-quality training set (ClimbMix).
[Takeaway] Interaction of Relevance and Diversity: Optimal performance is driven not only by the relevance of individual clusters to the target task but also by the semantic diversity among the selected clusters in the final mixture.