CLIMB: Mezcla Iterativa de Datos Basada en Agrupamiento para el Preentrenamiento de Modelos de Lenguaje

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

Resumen

Los conjuntos de datos de preentrenamiento suelen recopilarse de contenido web y carecen de divisiones de dominio inherentes. Por ejemplo, conjuntos de datos ampliamente utilizados como Common Crawl no incluyen etiquetas de dominio explícitas, mientras que la curación manual de conjuntos de datos etiquetados como The Pile es intensiva en mano de obra. En consecuencia, identificar una mezcla óptima de datos para el preentrenamiento sigue siendo un problema desafiante, a pesar de sus beneficios significativos para el rendimiento del preentrenamiento. Para abordar estos desafíos, proponemos CLustering-based Iterative Data Mixture Bootstrapping (CLIMB), un marco automatizado que descubre, evalúa y refina mezclas de datos en un entorno de preentrenamiento. Específicamente, CLIMB incrusta y agrupa conjuntos de datos a gran escala en un espacio semántico y luego busca iterativamente mezclas óptimas utilizando un modelo proxy más pequeño y un predictor. Cuando se entrena continuamente con 400 mil millones de tokens utilizando esta mezcla, nuestro modelo de 1 mil millones supera al estado del arte Llama-3.2-1B en un 2.0%. Además, observamos que la optimización para un dominio específico (por ejemplo, Ciencias Sociales) produce una mejora del 5% sobre el muestreo aleatorio. Finalmente, presentamos ClimbLab, un corpus filtrado de 1.2 billones de tokens con 20 clusters como un espacio de investigación, y ClimbMix, un conjunto de datos compacto pero potente de 400 mil millones de tokens diseñado para un preentrenamiento eficiente que ofrece un rendimiento superior bajo un presupuesto de tokens igual. Analizamos la mezcla final de datos, elucidando las características de una mezcla óptima de datos. Nuestros datos están disponibles en: 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/