CLIMB: Clustering-gebaseerde Iteratieve Data Mengsel Bootstrapping voor Vooraf Trainen van Taalmodellen

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

Samenvatting

Pre-trainingsdatasets worden doorgaans verzameld uit webinhoud en missen inherente domeinindelingen. Bijvoorbeeld, veelgebruikte datasets zoals Common Crawl bevatten geen expliciete domeinlabels, terwijl het handmatig samenstellen van gelabelde datasets zoals The Pile arbeidsintensief is. Als gevolg hiervan blijft het identificeren van een optimale pre-trainingsdatamix een uitdagend probleem, ondanks de aanzienlijke voordelen voor de pre-trainingsprestaties. Om deze uitdagingen aan te pakken, stellen wij CLustering-based Iterative Data Mixture Bootstrapping (CLIMB) voor, een geautomatiseerd framework dat datamixen ontdekt, evalueert en verfijnt in een pre-trainingsomgeving. Specifiek embedt en clustert CLIMB grootschalige datasets in een semantische ruimte en zoekt vervolgens iteratief naar optimale mixen met behulp van een kleiner proxy-model en een voorspeller. Wanneer continu getraind op 400B tokens met deze mix, overtreft ons 1B-model de state-of-the-art Llama-3.2-1B met 2,0%. Bovendien observeren we dat optimalisatie voor een specifiek domein (bijvoorbeeld Sociale Wetenschappen) een verbetering van 5% oplevert ten opzichte van willekeurige steekproeven. Ten slotte introduceren we ClimbLab, een gefilterd corpus van 1,2 biljoen tokens met 20 clusters als onderzoeksspeelveld, en ClimbMix, een compact maar krachtige dataset van 400 miljard tokens die is ontworpen voor efficiënte pre-training en superieure prestaties levert binnen een gelijk tokenbudget. We analyseren de uiteindelijke datamix en verduidelijken de kenmerken van een optimale datamix. Onze data is beschikbaar op: 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/