CLIMB: Amostragem Iterativa de Mistura de Dados Baseada em Agrupamento para Pré-treinamento de Modelos de Linguagem

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

April 17, 2025

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

Resumo

Os conjuntos de dados de pré-treinamento são tipicamente coletados a partir de conteúdo da web e carecem de divisões de domínio inerentes. Por exemplo, conjuntos de dados amplamente utilizados, como o Common Crawl, não incluem rótulos de domínio explícitos, enquanto a curadoria manual de conjuntos de dados rotulados, como The Pile, é trabalhosa. Consequentemente, identificar uma mistura ideal de dados para pré-treinamento continua sendo um problema desafiador, apesar de seus benefícios significativos para o desempenho do pré-treinamento. Para enfrentar esses desafios, propomos o CLustering-based Iterative Data Mixture Bootstrapping (CLIMB), uma estrutura automatizada que descobre, avalia e refina misturas de dados em um cenário de pré-treinamento. Especificamente, o CLIMB incorpora e agrupa conjuntos de dados em larga escala em um espaço semântico e, em seguida, busca iterativamente por misturas ideais usando um modelo proxy menor e um preditor. Quando treinado continuamente em 400 bilhões de tokens com essa mistura, nosso modelo de 1 bilhão supera o estado da arte Llama-3.2-1B em 2,0%. Além disso, observamos que a otimização para um domínio específico (por exemplo, Ciências Sociais) resulta em uma melhoria de 5% em relação à amostragem aleatória. Por fim, apresentamos o ClimbLab, um corpus filtrado de 1,2 trilhão de tokens com 20 clusters como um ambiente de pesquisa, e o ClimbMix, um conjunto de dados compacto, porém poderoso, de 400 bilhões de tokens projetado para um pré-treinamento eficiente que oferece desempenho superior sob um orçamento igual de tokens. Analisamos a mistura final de dados, elucidando as características de uma mistura ideal de dados. Nossos dados estão disponíveis em: 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/