FastMix: Otimização Rápida de Mistura de Dados via Gradiente Descendente

Resumo

Embora conjuntos de dados grandes e diversos tenham impulsionado avanços recentes em modelos de grande escala, identificar a mistura de dados ideal para pré-treinamento e pós-treinamento continua sendo um problema em aberto significativo. Abordamos esse desafio com o FASTMIX, uma nova estrutura que automatiza a descoberta de mistura de dados enquanto treina apenas um único modelo proxy. Em vez de depender de heurísticas predefinidas ou simulações com uso intensivo de recursos, o FASTMIX otimiza conjuntamente os coeficientes de mistura e os parâmetros do modelo, melhorando substancialmente a eficiência e a escalabilidade em relação às abordagens anteriores. No cerne do FASTMIX está uma reformulação da seleção de mistura como um problema de otimização bilevel. Sob essa reformulação, mostramos que otimizar as proporções da mistura é matematicamente equivalente a atribuir pesos de perda por fonte sob amostragem uniforme de fontes. Isso incorpora os coeficientes de mistura diretamente no objetivo iterativo diferenciável de otimização, permitindo uma otimização eficiente baseada em gradiente tanto da mistura quanto do modelo. Para resolver o problema de otimização, o FASTMIX implementa um procedimento de otimização iterativa aproximada, alternando entre (i) atualizar os parâmetros do modelo em dados amostrados de acordo com as proporções atuais da mistura (loop interno) e (ii) atualizar as proporções da mistura com base no feedback de validação (loop externo). Tanto no pré-treinamento quanto no pós-treinamento, o FASTMIX supera as referências enquanto reduz drasticamente o custo de busca. 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)