CLIMB: Bootstrapping Iterativo di Miscelazione Dati basato su Clustering per il Pre-addestramento di Modelli Linguistici

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

April 17, 2025

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

Abstract

I dataset di pre-training sono tipicamente raccolti da contenuti web e mancano di divisioni intrinseche per dominio. Ad esempio, dataset ampiamente utilizzati come Common Crawl non includono etichette esplicite per dominio, mentre la creazione manuale di dataset etichettati come The Pile è un processo laborioso. Di conseguenza, identificare una miscela ottimale di dati per il pre-training rimane un problema complesso, nonostante i suoi significativi benefici per le prestazioni del pre-training. Per affrontare queste sfide, proponiamo CLustering-based Iterative Data Mixture Bootstrapping (CLIMB), un framework automatizzato che scopre, valuta e affina le miscele di dati in un contesto di pre-training. Nello specifico, CLIMB incorpora e raggruppa dataset su larga scala in uno spazio semantico, per poi cercare iterativamente miscele ottimali utilizzando un modello proxy più piccolo e un predittore. Quando addestrato continuamente su 400 miliardi di token con questa miscela, il nostro modello da 1 miliardo supera lo stato dell'arte Llama-3.2-1B del 2,0%. Inoltre, osserviamo che l'ottimizzazione per un dominio specifico (ad esempio, Scienze Sociali) produce un miglioramento del 5% rispetto al campionamento casuale. Infine, introduciamo ClimbLab, un corpus filtrato da 1,2 trilioni di token con 20 cluster come ambiente di ricerca, e ClimbMix, un dataset compatto ma potente da 400 miliardi di token progettato per un pre-training efficiente che offre prestazioni superiori con un budget di token equivalente. Analizziamo la miscela finale di dati, chiarendo le caratteristiche di una miscela ottimale. I nostri dati sono disponibili all'indirizzo: 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/