FastMix: Optimización rápida de mezcla de datos mediante descenso de gradiente

Resumen

Aunque los conjuntos de datos grandes y diversos han impulsado los avances recientes en modelos de gran escala, identificar la mezcla óptima de datos para el pre-entrenamiento y el post-entrenamiento sigue siendo un problema abierto significativo. Abordamos este desafío con FASTMIX, un marco novedoso que automatiza el descubrimiento de mezclas de datos mientras entrena solo un único modelo proxy. En lugar de depender de heurísticas predefinidas o simulaciones que consumen muchos recursos, FASTMIX optimiza conjuntamente los coeficientes de mezcla y los parámetros del modelo, mejorando sustancialmente la eficiencia y escalabilidad respecto a enfoques anteriores. El núcleo de FASTMIX es una reformulación de la selección de mezcla como un problema de optimización bilevel. Bajo esta reformulación, mostramos que optimizar las proporciones de mezcla es matemáticamente equivalente a asignar pesos de pérdida por fuente bajo un muestreo uniforme de fuentes. Esto incorpora los coeficientes de mezcla directamente en el objetivo de optimización iterativa diferenciable, permitiendo una optimización eficiente basada en gradientes tanto de la mezcla como del modelo. Para resolver el problema de optimización, FASTMIX implementa un procedimiento de optimización iterativa aproximada, alternando entre (i) la actualización de los parámetros del modelo con datos muestreados según las proporciones de mezcla actuales (bucle interno) y (ii) la actualización de las proporciones de mezcla basándose en la retroalimentación de validación (bucle externo). Tanto en pre-entrenamiento como en post-entrenamiento, FASTMIX supera a las líneas base al tiempo que reduce drásticamente el costo de búsqueda. Código (https://github.com/hrtan/fastmix)

English

While large and diverse datasets have driven recent advances in large models, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. We address this challenge with FASTMIX, a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource-intensive simulations, FASTMIX jointly optimizes mixture coefficients and model parameters, substantially improving efficiency and scalability over prior approaches. At the core of FASTMIX is a reformulation of mixture selection as a bilevel optimization problem. Under this reformulation, we show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling. This embeds the mixture coefficients directly into the differentiable iterative optimization objective, enabling efficient, gradient-based optimization of both mixture and model. To solve the optimization problem, FASTMIX implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop). Across pre- and post-training, FASTMIX outperforms baselines while drastically reducing search cost. Code (https://github.com/hrtan/fastmix)