FastMix: Быстрая оптимизация смеси данных с помощью градиентного спуска

Аннотация

Хотя большие и разнообразные наборы данных способствовали недавним достижениям в области больших моделей, определение оптимального состава данных для предварительного и последующего обучения остается важной открытой проблемой. Мы решаем эту задачу с помощью FASTMIX — новой платформы, которая автоматизирует поиск состава данных, обучая при этом только одну прокси-модель. Вместо того чтобы полагаться на заранее заданные эвристики или ресурсоемкие симуляции, FASTMIX совместно оптимизирует коэффициенты смеси и параметры модели, значительно повышая эффективность и масштабируемость по сравнению с предыдущими подходами. В основе FASTMIX лежит переформулировка выбора смеси как задачи двухуровневой оптимизации. В рамках этой переформулировки мы показываем, что оптимизация пропорций смеси математически эквивалентна присвоению весов потерь для каждого источника при равномерной выборке источников. Это встраивает коэффициенты смеси непосредственно в дифференцируемую итерационную целевую функцию оптимизации, что позволяет эффективно оптимизировать как смесь, так и модель на основе градиентов. Для решения задачи оптимизации FASTMIX реализует приближенную итерационную процедуру, чередующую (i) обновление параметров модели на данных, выбранных в соответствии с текущими пропорциями смеси (внутренний цикл), и (ii) обновление пропорций смеси на основе обратной связи с валидационного набора (внешний цикл). Как на этапе предварительного, так и последующего обучения FASTMIX превосходит базовые подходы, одновременно значительно снижая стоимость поиска. Код (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)