Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training

TL;DR

CLIMB (NVIDIA, Nemotron family) automates the search for an LLM pretraining data mixture by (1) embedding and clustering raw web data into ~21 semantic clusters, then (2) iteratively bootstrapping mixture weights via a proxy-model + regression-predictor loop that optimizes directly against a held-out reasoning-benchmark validation set. With the discovered mixture, a 1B model continuously pretrained on 400B tokens beats Llama-3.2-1B by 2.0% on a broad reasoning suite, and targeted-domain mixtures (e.g. MMLU Social Sciences) gain 5% over random sampling. The framework also yields two released artifacts — ClimbLab (1.2T tokens, 20 clusters, research playground) and ClimbMix (400B-token efficient pretraining corpus).

Key claims

Pre-training data mixture is framed as a bi-level optimization problem: outer loop over mixture weights, inner loop trains proxy models, with a regression predictor approximating the mixture→performance map [§3.2, Eq. 2].
A coordinate-descent iterative bootstrap alternates between sampling (top-K from the current predictor, with random exploration) and predictor fitting (LightGBM regression on observed mixture-performance pairs), which the paper argues generalizes single-iteration methods like RegMix [§3.2, Subroutines 1–2].
Three search iterations of 64 / 32 / 16 mixtures (112 proxy runs total) suffice; a Dirichlet prior over per-cluster token counts initializes the search; LightGBM with L1/L2 reg + early stopping is used as the predictor [§4.1].
At 40B continuous-pretraining tokens, CLIMB beats Random, DoReMi and RegMix on average reasoning accuracy at both 350M (54.83 vs 53.78 best baseline) and 1B (60.41 vs 59.37) [Table 1].
Continuing pretraining a 1B model on 400B tokens with the CLIMB mixture exceeds Llama-3.2-1B by 2.0% average across 12 reasoning benchmarks (53.54 vs 51.56) and beats AMD-OLMo-1.2B and TinyLlama-1.1B [Table 2].
Domain-targeting works: optimizing CLIMB toward MMLU Social Sciences yields +5% over random sampling; ablations show iterative search (CLIMB-iter3) beating both Random and a search-budget-matched CLIMB-Best@N baseline on STEM, Humanities, and Social Sciences [§6, Fig. 5].
Search-compute scaling: increasing the total search budget from 100% (112 runs) to 150% / 200% monotonically improves headline average from 60.41 → 60.72 → 61.12 [Table 3, Abl.comp].
Proxy-model size is forgiving: 62M and 132M proxies recover 60.11 / 60.19 vs 60.41 for the 350M proxy — i.e. cheaper proxies cost only ~0.3 points of average accuracy [Table 3, Abl.proxy].
Cluster-count and initial-K choices are also forgiving in the studied range (initial K ∈ {48, 64, 100, 1000, 2000}, final cluster counts ∈ {15, 21, 30}); 1000→21 is the chosen setting [Table 3, Abl.clus].
Two released datasets: ClimbLab (1.2T tokens, 20 clusters, derived from Nemotron-CC + smollm-corpus) and ClimbMix (400B tokens, designed for efficient pretraining under equal token budget) [Abstract, §1].

Method

CLIMB has two phases. Data preprocessing: documents from the source corpus (Nemotron-CC + smollm-corpus) are embedded with stella_en_400M_v5 and k-means-clustered (FAISS, K=1000); fastText quality classifiers (trained on 1M Nemotron-340B-annotated examples across four axes: overall quality, educational value, informational value, advertisement score) prune low-quality clusters at threshold 3.0, yielding ~240 clusters; finally, clusters are merged by Euclidean distance between centroids (threshold 1.5) down to 16–21 super-clusters that act as the mixture-weight axes.

Iterative bootstrapping: the search space is the simplex over the 21 cluster weights. Phase-1 pretraining establishes a stable WSD-schedule checkpoint at 10T tokens for proxy and target models (62M / 350M / 1B); the data-mixing search happens in the decay stage, where each candidate mixture is evaluated by training a proxy model on 40B tokens and measuring validation accuracy on PIQA / ARC_E / HellaSwag. A LightGBM regressor is fit to (weight, performance) pairs after each iteration; the next iteration samples new candidates from the top-K predicted region (with random exploration) and adds them to the pool. After three iterations, the predicted best mixture is used to train the target model on 40B (or 400B) tokens.

Results

At 40B target-training tokens (Table 1), CLIMB beats Random, DoReMi, and RegMix at both 350M (54.83 avg vs 53.78 RegMix) and 1B (60.41 vs 59.37 RegMix). Per-benchmark, the gains are concentrated on ARC_C, ARC_E, HellaSwag, and WinoGrande — the same families used as the search optimization targets, plus generalization to held-out benchmarks (SIQA, MMLU, etc.).

At 400B target-training tokens (Table 2), CLIMB-950M reaches 53.54 average across 12 benchmarks, beating Llama-3.2-1.2B (51.56), AMD-OLMo-1.2B (49.93), and TinyLlama-1.1B (48.42). CLIMB-350M (48.93) beats SmolLM-360M (47.78) and Qwen2.5-490M (48.14). The headline 2.0% margin over Llama-3.2-1B is the abstract’s flagship number.

Domain targeting: optimizing toward MMLU Social Sciences, CLIMB-iter3 at 350M hits 39.36% vs Random’s baseline and CLIMB-Best@N’s intermediate level; at 1B, CLIMB-iter3 hits 41.79% Social Sciences accuracy, +1.13% over the strong CLIMB-Best@N baseline [§6, Fig. 5]. Cross-domain results: STEM 28.67%, Humanities 29.56% at 350M.

Ablations are systematic: increasing search compute (×1.5, ×2) buys 0.3–0.7 points; shrinking the proxy from 350M to 62M costs ~0.3 points; iteration allocations (6:1 vs 4:2:1 vs 2:2:1:1) prefer 4:2:1 (60.41 vs 60.14–60.05); Dirichlet initialization beats uniform random (60.41 vs 60.21).

Why it’s interesting

CLIMB makes the team’s implicit thesis from the Smol Training Playbook (GPU MODE talk on SmolLM3) talk operational: data-mix curation is the dominant lever in modern pretraining, and the loop that finds the right mix can itself be automated against benchmark validation performance — i.e. evals become the explicit objective function, not a downstream measurement. That framing matters for any Luma pretraining-scale project (video, multimodal) where “what data to mix in what ratio” is currently chosen by intuition; CLIMB’s recipe is roughly: embed, cluster, fit a cheap predictor on proxy runs, search. The bi-level + LightGBM-predictor decomposition is domain-agnostic — the only domain-specific pieces are the embedding model and the validation benchmark. The reported insensitivity to proxy size (62M ≈ 350M) is the most encouraging operational result: the search itself is cheap.

The user’s note flags the cost: “evals is super important because that is the explicit metric they are trying to optimize for.” This is the standard hazard of Bayesian-optimization-style pipelines — once the eval is the loss, Goodhart applies. CLIMB partially defends against this by optimizing on PIQA/ARC_E/HellaSwag validation splits and reporting transfer to ~12 held-out benchmarks (Table 2 generalization is consistent), but the failure mode for a generative-models team would be more severe: video / multimodal “evals” are noisier and more gameable than text reasoning suites.