CLIMB: Clustering-basierte iterative Datenmischung zur Bootstrapping-Vorverarbeitung von Sprachmodellen

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

April 17, 2025

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

Zusammenfassung

Vorverarbeitungsdatensätze werden typischerweise aus Webinhalten gesammelt und weisen keine inhärenten Domänenunterteilungen auf. Beispielsweise enthalten weit verbreitete Datensätze wie Common Crawl keine expliziten Domänenlabels, während die manuelle Kuratierung von gelabelten Datensätzen wie The Pile arbeitsintensiv ist. Folglich bleibt die Identifizierung einer optimalen Vorverarbeitungsdatenmischung ein herausforderndes Problem, obwohl sie erhebliche Vorteile für die Vorverarbeitungsleistung bietet. Um diese Herausforderungen zu bewältigen, schlagen wir CLustering-based Iterative Data Mixture Bootstrapping (CLIMB) vor, ein automatisiertes Framework, das Datenmischungen in einem Vorverarbeitungskontext entdeckt, bewertet und verfeinert. Konkret bettet CLIMB groß angelegte Datensätze in einen semantischen Raum ein, clustert sie und sucht dann iterativ nach optimalen Mischungen unter Verwendung eines kleineren Proxy-Modells und eines Prädiktors. Wenn unser 1B-Modell kontinuierlich mit dieser Mischung auf 400B Tokens trainiert wird, übertrifft es den State-of-the-Art Llama-3.2-1B um 2,0%. Darüber hinaus beobachten wir, dass die Optimierung für eine spezifische Domäne (z. B. Sozialwissenschaften) eine Verbesserung von 5 % gegenüber der zufälligen Stichprobenziehung erzielt. Schließlich stellen wir ClimbLab vor, ein gefiltertes 1,2-Billionen-Token-Korpus mit 20 Clustern als Forschungsumgebung, und ClimbMix, einen kompakten, aber leistungsstarken 400-Milliarden-Token-Datensatz, der für eine effiziente Vorverarbeitung konzipiert ist und unter einem gleichen Token-Budget eine überlegene Leistung liefert. Wir analysieren die endgültige Datenmischung und erläutern die Merkmale einer optimalen Datenmischung. Unsere Daten sind verfügbar unter: 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/