CLIMB: Кластеризация на основе итеративного бутстрэппинга смеси данных для предварительного обучения языковых моделей

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

April 17, 2025

Авторы: 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

Аннотация

Наборы данных для предварительного обучения обычно собираются из веб-контента и не имеют явного разделения на домены. Например, широко используемые наборы данных, такие как Common Crawl, не содержат явных меток доменов, в то время как ручное создание размеченных наборов данных, таких как The Pile, требует значительных усилий. В результате определение оптимальной смеси данных для предварительного обучения остается сложной задачей, несмотря на ее значительное влияние на производительность предварительного обучения. Для решения этих проблем мы предлагаем CLustering-based Iterative Data Mixture Bootstrapping (CLIMB) — автоматизированную структуру, которая обнаруживает, оценивает и улучшает смеси данных в контексте предварительного обучения. В частности, CLIMB встраивает и кластеризует крупномасштабные наборы данных в семантическом пространстве, а затем итеративно ищет оптимальные смеси с использованием меньшей модели-заместителя и предсказателя. При непрерывном обучении на 400 миллиардах токенов с такой смесью наша модель объемом 1 миллиард параметров превосходит современную модель Llama-3.2-1B на 2,0%. Более того, мы наблюдаем, что оптимизация для конкретного домена (например, социальных наук) дает улучшение на 5% по сравнению со случайной выборкой. Наконец, мы представляем ClimbLab — отфильтрованный корпус объемом 1,2 триллиона токенов с 20 кластерами в качестве исследовательской площадки, а также ClimbMix — компактный, но мощный набор данных объемом 400 миллиардов токенов, предназначенный для эффективного предварительного обучения, который демонстрирует превосходную производительность при равном бюджете токенов. Мы анализируем итоговую смесь данных, раскрывая характеристики оптимальной смеси. Наши данные доступны по адресу: 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/