Source-linked AI summary
Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training
Shizhe Diao, Yu Yang, Yonggan Fu, Xin Dong, Dan Su, Markus Kliegl, Zijia Chen, Peter Belcak, Yoshi Suhara, Hongxu Yin, Mostofa Patwary, Yingyan Lin, Jan Kautz, Pavlo Molchanov
TL;DR
Pre-training data mixtures are difficult to optimize because large web corpora lack domain labels and curated mixtures are labor-intensive to construct. Nemotron-CLIMB clusters data semantically and iteratively searches mixtures with proxy models and a predictor, achieving stronger reported performance while producing research and pre-training datasets. Its domain-specific evaluation uses coarse MMLU categories and does not yet establish effectiveness in high-stakes domains.
Problem
Large web datasets lack explicit domain labels, while curated labeled datasets are labor-intensive and optimal mixture selection remains non-trivial.
Method
CLIMB embeds and clusters datasets, evaluates sampled mixtures with proxy models, and trains a predictor to iteratively refine mixtures without predefined domain labels.
Results
CLIMB-trained 350M and 1B models achieve state-of-the-art performance across 12 reasoning tasks, while a 1B model trained on 400B tokens exceeds Llama-3.2-1B by 2.0%.
Takeaways & Limitations
CLIMB supports automated discovery of domain-aware mixtures and provides a 1.2-trillion-token, 20-cluster corpus plus a compact 400-billion-token pre-training dataset.
Takeaways & Limitations
Domain-specific benefits are evaluated using coarse MMLU categories, and effectiveness in high-stakes domains such as finance or healthcare remains unevaluated.
Abstract
from arXiv · showhide
Pre-training datasets are typically collected from web content and lack inherent domain divisions. For instance, widely used datasets like Common Crawl do not include explicit domain labels, while manually curating labeled datasets such as The Pile is labor-intensive. Consequently, identifying an optimal pre-training data mixture remains a challenging problem, despite its significant benefits for pre-training performance. To address these challenges, we propose CLustering-based Iterative Data Mixture Bootstrapping (Nemotron-CLIMB), an automated framework that discovers, evaluates, and refines data mixtures in a pre-training setting. Specifically, Nemotron-CLIMB embeds and clusters large-scale datasets in a semantic space and then iteratively searches for optimal mixtures using a smaller proxy model and a predictor. When continuously trained on 400B tokens with this mixture, our 1B model exceeds the state-of-the-art Llama-3.2-1B by 2.0%. Moreover, we observe that optimizing for a specific domain (e.g., Social Sciences) yields a 5% improvement over random sampling. Finally, we introduce Nemotron-ClimbLab, a filtered 1.2-trillion-token corpus with 20 clusters as a research playground, and Nemotron-ClimbMix, a compact yet powerful 400-billion-token dataset designed for efficient pre-training that delivers superior performance under an equal token budget. We analyze the final data mixture, elucidating the characteristics of an optimal data mixture. Our data is available at: https://research.nvidia.com/labs/lpr/climb/
1 Introduction
CLIMB addresses the difficulty of optimizing heterogeneous pre-training mixtures by clustering data semantically and iteratively searching for effective combinations with proxy models and predictors. The framework targets efficient, domain-relevant training and reports stronger scaling and downstream performance with its discovered mixtures.
- Method: CLIMB embeds and clusters large-scale datasets, then iteratively searches data mixtures using proxy models and a performance predictor.The predictor maps mixture features to performance metrics, enabling refinement without predefined domain labels.
- Method: The framework dynamically refines mixtures to balance diversity and domain relevance while avoiding manually predefined domain labels.Its bootstrapping strategy proposes, prunes, and refines candidate mixtures throughout training.
- Search behavior: The search process narrows from broad configurations in early iterations toward more optimal mixtures in later iterations.The progression is visualized with t-SNE across CLIMB-Iter1, CLIMB-Iter2, and CLIMB-Iter3.
- Efficiency: Lightweight proxy models and progressive pruning reduce the cost of evaluating mixture quality compared with brute-force hyperparameter sweeps.The approach allocates computation across iterative verification and generation rather than relying on a single large search.
- Results: A 1B model trained on 400B tokens with the discovered mixture exceeds Llama-3.2-1B by 2.0%.The same experiments report a 5% improvement over random sampling when optimizing for Social Sciences.
- Resources: CLIMB introduces a filtered 1.2-trillion-token corpus with 20 clusters and a separate 400-billion-token dataset for efficient pre-training.These datasets support research on data mixing and efficient pre-training.
2 CLIMB: CLustering-based Iterative Data Mixture Bootstrapping
CLIMB converts large raw corpora into semantically organized clusters, then searches their mixture weights through iterative proxy evaluation and predictor-guided refinement. This replaces exhaustive mixture evaluation with a progressively narrowed search over promising configurations.
- 2.1 Data Preprocessing: CLIMB embeds raw documents and clusters their vectors to organize semantically aligned data into groups for mixture construction.The preprocessing pipeline maps documents into embedding vectors, applies clustering, prunes low-quality clusters, and merges nearby clusters.
- 2.2 Iterative Bootstrapping: Mixture Weight Search: A data mixture is represented by weights over clusters, and its optimization is formulated as a bi-level search for task performance.The inner model-training problem produces performance for a given mixture, while the outer problem searches mixture weights.
- 2.2 Iterative Bootstrapping: Mixture Weight Search: At each iteration, CLIMB samples promising configurations, trains proxy models on them, updates the predictor, and selects the final configuration predicted to perform best.The implementation combines previously sampled configurations with newly selected candidates, while promising samples improve prediction accuracy for high-performing configurations.
- 2.2 Iterative Bootstrapping: Mixture Weight Search: CLIMB fits a predictor from sampled mixture-performance pairs to approximate the expensive objective without training a proxy model for every configuration.The predictor uses a subset of evaluated configurations, reducing the training cost of searching the full design space.
- 2.2 Iterative Bootstrapping: Mixture Weight Search: Instead of uniform sampling alone, CLIMB alternates configuration sampling and predictor fitting to focus later evaluations on higher-quality mixture subspaces.Uniform initialization can waste a fixed budget on low-quality weights; iterative coordinate descent uses predicted performance to guide new samples.
3 Experimental Settings
The experiments use clustered Nemotron-CC and smollm-corpus data, reasoning benchmarks, proxy models, and 40B-token continuous pre-training against established mixture baselines. The setup also specifies the clustering, iterative search, and predictor regularization choices used in evaluation.
- Data: CLIMB-clustering produces 21 super-clusters containing 800B tokens from Nemotron-CC and smollm-corpus for training.Evaluation uses PIQA, ARC_C, ARC_E, HellaSwag, WinoGrande, and SIQA, with optimization on PIQA, ARC_E, and HellaSwag validation data.
- Model: The study trains 62M, 350M, and 1B decoder-only models after 10T-token phase-1 pre-training, using 62M and 350M models as proxies.Target models are continuously trained on 40B tokens with the discovered mixture unless otherwise stated.
- Baselines: The baselines are Random selection, DoReMi, and RegMix.These methods provide the comparison set for evaluating CLIMB's data-mixture search.
- Data Preprocessing: The preprocessing uses stella_en_400M_v5 embeddings and FAISS K-means initialized with Kinit = 1000 clusters.Cluster quality is assessed across overall quality, educational value, informational value, and advertisement score before pruning and merging.
- Iterative Bootstrapping: The mixture search runs for three iterations with 64, 32, and 16 searches, using a LightGBM predictor trained on current and past data.Regularization, early stopping, depth limits, and minimum leaf size are used to reduce predictor overfitting.
4 Experimental Results
CLIMB improves data-mixture performance over established baselines and remains competitive after 400B-token training, including a 2.0% margin over Llama-3.2 at similar scale.
- Comparison with Data Mixture Baselines: CLIMB outperforms all baseline data-mixture methods in the main comparison.With 350M targets, it reaches 54.83% average accuracy versus 52.17% for Random and 53.78% for RegMix; the 1B model reaches 60.41%.
- Comparison with SOTA LMs: 2.0%: CLIMB surpasses the next-best Llama-3.2 model on overall average score after 400B-token training.At approximately 1B parameters, CLIMB also outperforms Llama-3.2 and AMD-OLMo across most general-reasoning benchmarks.
- Comparison with SOTA LMs: CLIMB remains consistently better than baseline models on additional benchmarks including MMLU, GPQA, OBQA, BoolQ, and RACE.These results are reported as evidence of generalization performance.
- Comparison with SOTA LMs: Iteratively refined data mixtures lead to better pre-training performance.
5 Analysis
The analysis examines domain-specific optimization, search compute, compute allocation, and final mixture weights. It finds that balanced iterative search and both relevance and diversity of clusters matter for robust performance.
- Optimizing towards Specific Domains: CLIMB optimizes mixtures separately for STEM, humanities, and social-sciences objectives using MMLU validation performance.The analysis compares this setup with CLIMB-Best@N, which directly searches for the best parameters.
- Effects of Compute Allocation: 4:2:1 compute allocation achieves the best overall average performance at 60.41%.Allocations that are too shallow or spread compute across too many iterations perform worse.
- Effects of Search Compute Budget: A 150%-200% compute increase yields noticeable gains, while more search iterations improve performance.The search procedure explores more candidates as total search budget increases.
- Analysis of Final Weights: C8, C9, C18, and C19 account for most of the final mixture weight in the general-reasoning setting.C8, C9, and C19 correlate strongly with general reasoning, while the four clusters collectively provide diverse topics.
- Analysis of Final Weights: Effective mixtures require both downstream-task relevance and diversity among clusters.
6 NEMOTRON-CLIMBLAB and NEMOTRON-CLIMBMIX: New SOTA Pre-training Data
CLIMB reorganizes and filters Nemotron-CC and smollm-corpus into a 1.2-trillion-token corpus, then selects a 400-billion-token mixture that outperforms existing datasets under the same token budget.
- NEMOTRON-CLIMBLAB: CLIMB-clustering reorganizes and filters Nemotron-CC and smollm-corpus into 20 clusters forming the 1.2-trillion-token NEMOTRON-CLIMBLAB corpus.
- NEMOTRON-CLIMBMIX: CLIMB-search extracts NEMOTRON-CLIMBMIX, a 400-billion-token dataset based on the identified optimal mixture.
- Performance: Models trained on NEMOTRON-CLIMBMIX significantly outperform models trained on Nemotron-CC, SmolLM, DCLM-baseline, and FineWeb-Edu under the same token budget.The comparison uses 1B models trained from scratch.
7 Related Work
Prior data-mixture methods often depend on datasets with clear domain distinctions. CLIMB instead identifies approximate domains from web data through clustering and automatically refines the mixture.
- Data Mixture for LLM Pre-training: Manually crafted mixtures such as those in The Pile, GLaM, and ROOTS rely on heuristics lacking standardization and transferability.
- Data Mixture for LLM Pre-training: DoReMi and DoGE iteratively optimize domain proportions using reference and proxy models.
- Data Mixture for LLM Pre-training: Existing optimization methods generally require datasets to possess clear and natural domain distinctions.
- Data Mixture for LLM Pre-training: CLIMB automatically identifies approximate domains from large web datasets and refines their mixture through clustering and iterative optimization.The approach is presented as straightforward and scalable relative to classifier-based annotation.
- Data Selection for Specific Domains: Domain-specific selection methods include DSIR’s hashed-n-gram relevance estimation and CRISP’s cluster-based sampling.
8 Conclusion
CLIMB automates domain-aware data-mixture optimization through clustering, proxy models, and iterative search, achieving state-of-the-art performance across 12 reasoning tasks under fixed computational budgets.
- CLIMB automates discovery, evaluation, and refinement of data mixtures for LLM pre-training.It combines unsupervised clustering, proxy model training, and prediction to navigate data-composition search without predefined domain labels.
- 350M and 1B models trained with CLIMB-selected mixtures achieve state-of-the-art performance across 12 reasoning tasks.
- Iterative mixture refinement balances unstructured corpora with targeted domain data under fixed computational budgets.
A Limitations
The paper identifies computational and evaluation limitations, including non-negligible proxy-model costs and coarse domain categories that have not been validated in high-stakes applications.
- Computational limitations: Proxy-model training remains a non-negligible computational cost despite using lightweight models.The authors identify parameter-efficient tuning, distillation, and zero-shot evaluation as possible directions for reducing this burden.
- Evaluation scope: Domain-specific evaluation relies on MMLU’s coarse-grained categories, which may not represent real-world applications.
- Evaluation scope: CLIMB has not yet been evaluated in high-stakes domains such as finance or healthcare.The authors leave validation in these settings for future research because their data characteristics and requirements may differ substantially.
- Societal impacts: Automated data-mixture optimization may overrepresent some domains while underrepresenting others, potentially reinforcing knowledge disparities.The paper calls for safeguards addressing fairness, unintended clustering bias, transparency, and interpretability.
C.2 Data
The experiments use high-quality Nemotron-CC data organized into clusters, evaluate general reasoning across multiple benchmarks, and characterize clusters with sampled-document topic summaries.
- Data source: Nemotron-CC’s highest-quality bucket contains approximately 800 billion tokens distributed across 21 hierarchical clusters.
- Downstream evaluation: General reasoning evaluation covers PIQA, ARC_C, ARC_E, HellaSwag, WinoGrande, TruthfulQA, and SIQA.Optimization uses validation data from PIQA, ARC_E, and HellaSwag, while evaluation uses test sets from all listed datasets.
- Pre-training setup: Phase-1 pre-training trains 62M, 350M, and 1B decoder-only Transformer models with next-token language modeling loss.All models are trained on 10 trillion tokens, although the authors note this over-training does not strictly align with scaling laws.
- Pre-training setup: Training uses AdamW, a 5e-5 learning rate annealed to 1e-5, a 2M-token batch size, and 256 NVIDIA H100 GPUs.A lightweight proxy model takes approximately 45 GPU hours, compared with around 6,400 GPU hours for the large target model.
- Cluster analysis: Cluster topics are characterized by asking GPT-4o to summarize four to seven representative topics from 100 randomly sampled documents per cluster.The authors describe this analysis as auxiliary rather than a complete account of cluster content.
- Cluster analysis: The dataset section includes a table listing cluster topics.
D.2 Relationship between Clusters and Downstream Tasks
Cluster similarity, evolving mixture weights, proxy size, clustering choices, initialization, and compute allocation are examined as factors in downstream performance and search robustness.
- Relationship between clusters and downstream tasks: In-domain similarity can help, but out-of-domain clusters may also become important during iterative search.For general reasoning, C8 and C19 resemble ARC-E, while C8 gains importance despite initially appearing out of domain and C21 declines despite high similarity.
- Relationship between clusters and downstream tasks: Similarity is an indicator rather than a decisive rule because optimal cluster contributions depend on the broader mixture.The analysis concludes that selection requires more than similarity-based filtering.
- Proxy models: Larger proxy models improve average score from 60.11 to 60.41 as size increases from 62M to 350M.The gains are modest but consistently favor the largest feasible proxy model.
- Clustering sensitivity: The study tests Kinit values of 48, 64, 100, 1000, and 2000 alongside Kenhanced values of 15, 21, and 30 to assess clustering sensitivity.
- Initialization: Dirichlet initialization reaches 60.41% average score versus 60.21% for random initialization.The comparable performances suggest robustness to initialization choice.
- Compute allocation: The 4:2:1 compute allocation achieves the best overall average performance at 60.41%.The allocation distributes total compute across three iterations in a 64:32:16 ratio.
- Evolution of cluster weights: General-reasoning weights shift across iterations, with C8 and C9 rising to 0.13 and 0.18 in Iteration 3.C19 and C21 decrease from initially high weights, indicating adaptive changes in cluster importance.
D.8 Analysis of Final Weights
The final mixture concentrates weight on clusters associated with target capabilities, while proxy models and predictors preserve useful structure for efficient search.
- C8, C9, C18, and C19 account for most general-reasoning mixture weight and collectively provide a diverse distribution.
- C7, C11, and C19 are particularly important for humanities, while C7 and C8 are highly influential for STEM.
- A 62M proxy model can retain key cluster-weight patterns while reducing computational costs.
- 94% Spearman rank correlation measures agreement between predicted and true accuracy for the 350M proxy models.
2. Limitations
The paper reports limitations concerning responsible data release and the statistical reporting of experimental results, while reviewer materials also emphasize documenting scope and reproducibility.
- The work states that reproducibility information, including framework, training, proxy-model, clustering, and evaluation details, is provided in the paper and appendix.
- The paper reports no error bars or statistical significance tests, stating that large-scale training makes task performance stable across runs.
- Optimizing data mixtures may overrepresent some domains and underrepresent others, potentially reinforcing knowledge disparities.
- The authors report no additional safeguards for released data because it is based on publicly available open datasets.