CLIMB

Analysis

Why This Paper Matters

Data quality and composition are increasingly recognized as critical levers for language model performance, yet most pre-training pipelines rely on heuristic filtering or uniform sampling. CLIMB tackles this head-on by treating data mixture optimization as a search problem, using clustering to impose semantic structure on massive unlabeled corpora and iterative bootstrapping to efficiently explore the mixture space. This matters because it moves beyond static, one-size-fits-all datasets toward dynamic, task-aware data selection, which could democratize high-quality pre-training for specialized domains.

How It Works

CLIMB operates in three phases. First, raw documents are embedded using a sentence transformer (stella_en_400M_v5) and clustered via k-means into 1000 initial clusters. Low-quality clusters are pruned using a fasttext classifier trained on annotations from a large teacher model, and remaining clusters are merged by centroid distance to yield about 20 super-clusters. This semantic grouping replaces noisy domain labels with data-driven categories.

Second, an iterative bootstrapping loop searches for optimal mixture weights. Starting with 64 random mixtures, small proxy models (62M or 350M parameters) are trained on 40B tokens each, and their performance on a calibration set (PIQA, ARC_E, HellaSwag validation) is recorded. A LightGBM predictor is fitted to approximate performance as a function of mixture weights. In each subsequent iteration, the predictor scores all unseen mixtures, the top performers are sampled (balancing exploration and exploitation), new proxy models are trained, and the predictor is retrained. After three iterations (64, 32, 16 samples), the best predicted mixture is selected.

Third, the optimal mixture is used to train a target model (62M, 350M, or 1B) on 40B tokens, and performance is compared against baselines. The entire pipeline is applied to a combined Nemotron-CC and smollm-corpus dataset, resulting in ClimbMix (400B tokens) and ClimbLab (1.2T tokens).

Results

CLIMB consistently outperforms random sampling, DoReMi, and RegMix across all model sizes. For the 350M model, average accuracy reaches 54.83% versus 52.17% (random) and 53.78% (RegMix). The 1B model achieves 60.41%, a clear margin. Gains on the calibration tasks generalize to held-out benchmarks like MMLU, GPQA, and BoolQ.

Compared to open models of similar size, CLIMB-trained models lead by 2.0% overall average over Llama-3.2-1B, and outperform AMD-OLMo and other baselines on most reasoning tasks. The ClimbMix dataset also shows better scaling: a 1B model trained on ClimbMix surpasses models trained on the same token budget with existing datasets.

Significance

CLIMB offers a principled, automated alternative to manual data curation. Its clustering step provides interpretable domain structure, while the iterative search reduces the computational cost of exploring mixture spaces from exponential to linear in the number of iterations. The framework is model-agnostic and can be applied to any large corpus. However, the reliance on proxy models and a fixed embedding model may limit generalization to very different tasks or languages. Future work could explore dynamic clustering and multi-task optimization. Overall, CLIMB advances the science of data-centric AI and provides a practical tool for building better pre-training datasets.