One-sentence Summary

NVIDIA and Georgia Institute of Technology propose Nemotron-CLIMB, an automated framework that semantically embeds and clusters large-scale pre-training data, then iteratively searches for optimal data mixtures via a smaller proxy model and a predictor to enable effective domain adaptation without curated labels, achieving a $2.0\%$2.0% improvement over Llama-3.2-1B when continuously trained on a $400\text{B-token}$400B-token mixture and releasing NEMOTRON-CLIMBLAB (a filtered $1.2\text{T-token}$1.2T-token corpus with $20\ \text{clusters}$20 clusters) and NEMOTRON-CLIMBMIX (a 400B-token compact dataset) for efficient pre-training.

Key Contributions

• The paper introduces Nemotron-CLIMB, an automated framework that embeds and clusters large-scale unlabeled web data in semantic space, then iteratively searches for optimal data mixtures using a smaller proxy model and a predictor, removing the need for predefined domain labels or manual curation required by prior methods such as DoReMi and DoGE.
• The paper releases NEMOTRON-CLIMBLAB, a filtered 1.2T-token corpus with 20 clusters for research, and NEMOTRON-CLIMBMIX, a 400B-token compact dataset designed for efficient pre-training that delivers superior performance under an equal token budget.
• A 1B model trained on 400B tokens with the discovered mixture outperforms Llama-3.2-1B by 2.0%, domain-specific optimization (e.g., Social Sciences) yields a 5% improvement over random sampling, and 350M and 1B models achieve state-of-the-art performance across 12 reasoning tasks.

Introduction

Pre-training datasets for large language models now scale to trillions of tokens, typically combining large-scale web crawls with smaller, high-quality domain-specific corpora. This scale creates a difficult balance between general knowledge and specialized expertise, and high-value data is often underused for targeted capabilities. Recent work shows that a final pre-training stage, called mid-training, can sharpen domain performance on benchmarks such as GSM8K, MMLU, and HumanEval, but optimizing the data mixture for that stage remains an open problem.

The technical challenge is twofold. First, massive corpora like Common Crawl lack explicit domain labels, so filtering relies on general heuristics such as perplexity, which do not capture domain-specific quality. Second, even when curated datasets like The Pile provide annotations, the relationship between mixture composition and model performance is complex and nonlinear, making manual or static mixing strategies insufficient.

The authors introduce CLIMB, a clustering-based iterative data mixture bootstrapping framework. They approach mixture building as a search problem: embedding and clustering large-scale datasets, constructing mixture-performance pairs with sampled and pruned mixtures using lightweight proxy models, and fitting a regression predictor. This enables iterative refinement of data mixtures without predefined domain labels, dynamically adapting throughout training. CLIMB reduces computational cost through progressive pruning and proxy model evaluation. The authors show that 350M and 1B models trained on 40B tokens surpass prior mixing methods like Doremi and RegMix, and a 1B model trained on 400B tokens exceeds Llama-3.2-1B. They also release a filtered 1.2-trillion-token corpus with 20 clusters and a 400-billion-token dataset for efficient pre-training.

Dataset

Dataset Description

The authors construct and release two new pre-training datasets, NEMOTRON-CLIMBLAB and NEMOTRON-CLIMBMIX, based on a clustering-based approach to semantic reorganization.

• Data Sources: The authors use two existing large-scale corpora, Nemotron-CC and smollm-corpus, as the raw source material. These are combined into a single pool for processing.
• Dataset Composition via Clustering: The combined corpus is processed using a clustering pipeline to semantically reorganize the data. Documents are first mapped to embeddings using an embedding model. These embeddings are then clustered into a large number of initial clusters (e.g., 1000), followed by a pruning step to remove low-quality clusters and a merging step to combine similar fine-grained clusters.
• NEMOTRON-CLIMBLAB: After applying the clustering pipeline to the combined Nemotron-CC and smollm-corpus, the authors obtain a 1.2-trillion-token high-quality corpus. This corpus is organized into 20 distinct semantic clusters and is released as NEMOTRON-CLIMBLAB.
• NEMOTRON-CLIMBMIX: Using the CLIMB search method, the authors identify an optimal data mixture from the 20 clusters. They extract a 400-billion-token subset based on this optimal mixture, which is released as NEMOTRON-CLIMBMIX.
• Data Statistics and Schema: The final datasets are organized by semantic cluster. The source corpus is reduced to 20 clusters for the released 1.2T dataset. In experimental settings, a similar clustering process applied to other source data yielded 21 super-clusters containing 800B tokens for training.
• How the Data Is Used: The authors train a 1B model from scratch using the NEMOTRON-CLIMBMIX mixture. This model is compared against models pre-trained on other datasets, including Nemotron-CC, SmolLM, DCLM-baseline, and FineWeb-Edu, under the same token budget. The authors note that a more balanced cluster distribution is needed for pre-training from scratch compared to continuous pre-training.

Method

The authors propose CLIMB (CLustering-based Iterative Data Mixture Bootstrapping), an automated framework designed to curate high-quality pre-training data by optimizing the mixture of different data domains. The goal is to maximize downstream task performance while maintaining general language capabilities, eliminating the need for manual data curation. The overall framework consists of two primary phases: data preprocessing to create semantic clusters and an iterative bootstrapping process to search for optimal mixture weights.

Data Preprocessing The first phase transforms a massive raw dataset into a set of semantically aligned clusters. This involves three steps:

1. Text Embedding: The raw documents are mapped into an embedding space using a pre-trained embedding model. This captures deeper semantic alignments compared to word-space representations.
2. Embedding Clustering: The embedding vectors are grouped using a clustering algorithm, such as K-means. To ensure fine-grained granularity, an initial large number of clusters (e.g., $K_{\text{init}} = 1000$Kinit​=1000) is generated.
3. Cluster Merging and Pruning: To improve quality and reduce redundancy, low-quality clusters are pruned based on model-based classifiers (e.g., fasttext models evaluating educational and informational value). Subsequently, similar clusters are merged based on the Euclidean distance between their centroids. This step reduces the initial clusters to a smaller, high-quality set of enhanced clusters ($K_{\text{enhanced}}$Kenhanced​), which form the basis for the subsequent mixture search.

Iterative Bootstrapping: Mixture Weight Search The second phase optimizes the sampling mixture weights to maximize performance on a calibration set. This is formulated as a bi-level optimization problem where the objective is to find mixture weights $\alpha$α that maximize the task performance $\ell(\alpha, \omega)$ℓ(α,ω), subject to $\sum \alpha_i = 1$∑αi​=1 and $\alpha_i \ge 0$αi​≥0. Since training a model for every possible mixture configuration is computationally prohibitive, the framework employs a predictor $f_\theta(\alpha)$fθ​(α) to approximate the performance.

The optimization is solved via an iterative coordinate descent method that alternates between sampling configurations and fitting the predictor. In each iteration $k$k:

• Configuration Sampling: The predictor evaluates untried mixture configurations. The top $N$N predicted configurations are selected, and $M$M new configurations are randomly sampled from this top set to balance exploitation of high-performing areas and exploration. These new samples are combined with previously evaluated configurations to form the training set for the next step.
• Predictor Fitting: Proxy models are trained on the newly sampled mixture configurations to obtain their actual performance. The predictor (e.g., a LightGBM regression model) is then updated by minimizing the loss between the predicted and actual performance on this expanded dataset.

This iterative process progressively refines the predictor's accuracy and guides the sampling strategy toward subspaces with higher-quality mixture weights. After a predefined number of iterations, the configuration with the highest predicted performance is selected as the optimal data mixture.

Experiment

The experiments evaluate CLIMB, a data mixing method, using proxy models (350M) to search for optimal data mixtures and then training target models (350M and 1B) on 40B tokens (or 400B for scaling) for reasoning benchmarks. CLIMB consistently outperforms random selection and existing baselines like DoReMi and RegMix, and also surpasses SOTA models such as Llama-3.2 at similar scales, with gains generalizing across all tested tasks. Further analysis on domain-specific optimization (MMLU) and compute allocation shows that iterative search with balanced depth and breadth (4:2:1 ratio) yields the best results, and that both cluster relevance and diversity are crucial for effective mixtures.

CLIMB consistently outperforms all baseline data mixture methods for both 350M and 1B models after continuous training on 40B tokens. The gains observed on validation tasks used for optimization also carry over to other benchmarks, indicating robust generalization of the learned data mixture. CLIMB achieves the highest average accuracy for both model sizes compared to random mixing and other mixture methods like Doremi and RegMix. Performance improvements from CLIMB extend beyond the tasks used for optimization, covering a broader set of reasoning benchmarks. The relative gains over the base model and random mixing are larger for the 1B model than for the 350M model, showing scalability with model size.

CLIMB models trained with the optimal data mixture outperform state-of-the-art language models of similar or larger scale on general reasoning benchmarks. The 950M CLIMB achieves the highest overall average score, surpassing the next best model by a noticeable margin, while the 350M variant also leads among sub-500M models. CLIMB (950M) attains the top average score across all compared models, exceeding Llama-3.2 (1.2B) by 2.0%. Among sub-500M models, CLIMB (350M) achieves the highest average, outperforming Qwen2.5 and SmolLM. The performance advantage of CLIMB is consistent across diverse benchmarks, including MMLU, GPQA, OBQA, BoolQ, and RACE.

The ablation study examines how search compute budget, compute allocation across iterations, and proxy model size affect the CLIMB data mixing method. Increasing the total search compute from 100% to 200% yields steady accuracy improvements, while a 4:2:1 allocation across three iterations outperforms fewer or more split allocations. Using a larger proxy model also boosts performance, confirming that these design choices are important for robust data mixture optimization. Raising search compute from 100% to 150% and 200% consistently improves average accuracy across all evaluated tasks. The 4:2:1 compute allocation across iterations achieves the best average performance compared with 6:1 and 2:2:1:1 allocations. A 350M proxy model leads to higher average accuracy than a 62M proxy under the same compute budget.

The evaluations show that CLIMB consistently outperforms baseline data mixture methods across model sizes, with performance gains generalizing to broader reasoning benchmarks beyond the optimization tasks. CLIMB models also surpass state-of-the-art models of similar or larger scale on general reasoning tasks. An ablation study confirms the importance of search compute budget, compute allocation across iterations, and proxy model size, where larger budgets, balanced allocations, and larger proxies yield better results.