QuaDMix: Selección de Datos Equilibrada en Calidad y Diversidad para el Pretrenamiento Eficiente de Modelos de Lenguaje de Gran Escala

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

Resumen

La calidad y la diversidad son dos métricas críticas para los datos de entrenamiento de los modelos de lenguaje de gran escala (LLMs), impactando positivamente en su rendimiento. Los estudios existentes suelen optimizar estas métricas por separado, generalmente aplicando primero un filtrado de calidad y luego ajustando las proporciones de los datos. Sin embargo, estos enfoques pasan por alto la inherente compensación entre calidad y diversidad, lo que hace necesario considerarlas de manera conjunta. Dada una cuota fija de entrenamiento, es esencial evaluar tanto la calidad de cada punto de datos como su efecto complementario en el conjunto de datos general. En este artículo, presentamos un marco unificado de selección de datos llamado QuaDMix, que optimiza automáticamente la distribución de datos para el preentrenamiento de LLMs mientras equilibra calidad y diversidad. Específicamente, primero proponemos múltiples criterios para medir la calidad de los datos y empleamos clasificación de dominio para distinguir los puntos de datos, midiendo así la diversidad general. QuaDMix luego utiliza una función de muestreo de datos parametrizada y unificada que determina la probabilidad de muestreo de cada punto de datos basándose en estas etiquetas relacionadas con calidad y diversidad. Para acelerar la búsqueda de los parámetros óptimos involucrados en el marco de QuaDMix, realizamos experimentos simulados en modelos más pequeños y utilizamos LightGBM para la búsqueda de parámetros, inspirados en el método RegMix. Nuestros experimentos en diversos modelos y conjuntos de datos demuestran que QuaDMix logra una mejora promedio del rendimiento del 7.2% en múltiples benchmarks. Estos resultados superan a las estrategias independientes para calidad y diversidad, destacando la necesidad y capacidad de equilibrar la calidad y diversidad de los datos.

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.