QuaDMix: Seleção de Dados Equilibrada entre Qualidade e Diversidade para Pré-treinamento Eficiente de LLMs

QuaDMix: Quality-Diversity Balanced Data Selection for Efficient LLM Pretraining

April 23, 2025

Autores: Fengze Liu, Weidong Zhou, Binbin Liu, Zhimiao Yu, Yifan Zhang, Haobin Lin, Yifeng Yu, Xiaohuan Zhou, Taifeng Wang, Yong Cao

cs.AI

Resumo

Qualidade e diversidade são duas métricas críticas para os dados de treinamento de modelos de linguagem de grande escala (LLMs), impactando positivamente o desempenho. Estudos existentes frequentemente otimizam essas métricas separadamente, geralmente aplicando primeiro filtros de qualidade e depois ajustando as proporções dos dados. No entanto, essas abordagens ignoram a relação intrínseca de trade-off entre qualidade e diversidade, exigindo sua consideração conjunta. Dada uma cota fixa de treinamento, é essencial avaliar tanto a qualidade de cada ponto de dados quanto seu efeito complementar no conjunto de dados geral. Neste artigo, introduzimos um framework unificado de seleção de dados chamado QuaDMix, que otimiza automaticamente a distribuição de dados para o pré-treinamento de LLMs, equilibrando qualidade e diversidade. Especificamente, propomos primeiro múltiplos critérios para medir a qualidade dos dados e empregamos classificação de domínio para distinguir os pontos de dados, medindo assim a diversidade geral. O QuaDMix então utiliza uma função de amostragem de dados parametrizada e unificada que determina a probabilidade de amostragem de cada ponto de dados com base nessas etiquetas relacionadas à qualidade e diversidade. Para acelerar a busca pelos parâmetros ótimos envolvidos no framework QuaDMix, realizamos experimentos simulados em modelos menores e usamos o LightGBM para a busca de parâmetros, inspirados no método RegMix. Nossos experimentos em diversos modelos e conjuntos de dados demonstram que o QuaDMix alcança uma melhoria média de desempenho de 7,2% em múltiplos benchmarks. Esses resultados superam as estratégias independentes para qualidade e diversidade, destacando a necessidade e a capacidade de equilibrar qualidade e diversidade dos dados.

English

Quality and diversity are two critical metrics for the training data of large language models (LLMs), positively impacting performance. Existing studies often optimize these metrics separately, typically by first applying quality filtering and then adjusting data proportions. However, these approaches overlook the inherent trade-off between quality and diversity, necessitating their joint consideration. Given a fixed training quota, it is essential to evaluate both the quality of each data point and its complementary effect on the overall dataset. In this paper, we introduce a unified data selection framework called QuaDMix, which automatically optimizes the data distribution for LLM pretraining while balancing both quality and diversity. Specifically, we first propose multiple criteria to measure data quality and employ domain classification to distinguish data points, thereby measuring overall diversity. QuaDMix then employs a unified parameterized data sampling function that determines the sampling probability of each data point based on these quality and diversity related labels. To accelerate the search for the optimal parameters involved in the QuaDMix framework, we conduct simulated experiments on smaller models and use LightGBM for parameters searching, inspired by the RegMix method. Our experiments across diverse models and datasets demonstrate that QuaDMix achieves an average performance improvement of 7.2% across multiple benchmarks. These results outperform the independent strategies for quality and diversity, highlighting the necessity and ability to balance data quality and diversity.