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는 대규모 데이터셋을 의미 공간에 임베딩하고 클러스터링한 후, 더 작은 프록시 모델과 예측기를 사용하여 반복적으로 최적의 혼합물을 탐색합니다. 이 혼합물로 400B 토큰을 지속적으로 학습했을 때, 우리의 1B 모델은 최신 기술인 Llama-3.2-1B를 2.0% 능가했습니다. 또한, 특정 도메인(예: 사회과학)에 최적화하면 무작위 샘플링보다 5%의 성능 향상을 관찰했습니다. 마지막으로, 우리는 연구 플레이그라운드로 20개의 클러스터로 구성된 필터링된 1.2조 토큰 코퍼스인 ClimbLab과, 동일한 토큰 예산 하에서 우수한 성능을 제공하도록 설계된 간결하면서도 강력한 4000억 토큰 데이터셋인 ClimbMix를 소개합니다. 우리는 최종 데이터 혼합물을 분석하여 최적의 데이터 혼합물의 특성을 설명합니다. 우리의 데이터는 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/