CLIMB : Amorçage itératif de mélange de données basé sur le clustering pour le pré-entraînement de modèles de langage

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

April 17, 2025

Auteurs: 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

cs.AI

Résumé

Les ensembles de données de pré-entraînement sont généralement collectés à partir de contenu web et ne présentent pas de divisions de domaine inhérentes. Par exemple, des ensembles de données largement utilisés comme Common Crawl ne comportent pas d'étiquettes de domaine explicites, tandis que la curation manuelle d'ensembles de données étiquetés tels que The Pile est laborieuse. Par conséquent, identifier un mélange optimal de données de pré-entraînement reste un problème complexe, malgré ses avantages significatifs pour les performances de pré-entraînement. Pour relever ces défis, nous proposons CLustering-based Iterative Data Mixture Bootstrapping (CLIMB), un cadre automatisé qui découvre, évalue et affine les mélanges de données dans un contexte de pré-entraînement. Plus précisément, CLIMB intègre et regroupe des ensembles de données à grande échelle dans un espace sémantique, puis recherche itérativement des mélanges optimaux à l'aide d'un modèle proxy plus petit et d'un prédicteur. Lorsqu'il est continuellement entraîné sur 400 milliards de tokens avec ce mélange, notre modèle de 1 milliard de paramètres dépasse le modèle de pointe Llama-3.2-1B de 2,0 %. De plus, nous observons que l'optimisation pour un domaine spécifique (par exemple, les sciences sociales) améliore les performances de 5 % par rapport à un échantillonnage aléatoire. Enfin, nous présentons ClimbLab, un corpus filtré de 1,2 trillion de tokens regroupés en 20 clusters comme terrain de recherche, et ClimbMix, un ensemble de données compact mais puissant de 400 milliards de tokens conçu pour un pré-entraînement efficace, offrant des performances supérieures avec un budget de tokens équivalent. Nous analysons le mélange de données final, en clarifiant les caractéristiques d'un mélange de données optimal. Nos données sont disponibles à l'adresse : https://research.nvidia.com/labs/lpr/climb/

English

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 (CLIMB), an automated framework that discovers, evaluates, and refines data mixtures in a pre-training setting. Specifically, 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 ClimbLab, a filtered 1.2-trillion-token corpus with 20 clusters as a research playground, and 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/