FastMix : Optimisation rapide du mélange de données par descente de gradient

Résumé

Alors que des jeux de données vastes et diversifiés ont favorisé les récentes avancées des grands modèles, identifier le mélange de données optimal pour le pré-entraînement et le post-entraînement reste un problème ouvert majeur. Nous relevons ce défi avec FASTMIX, un nouveau cadre qui automatise la découverte du mélange de données tout en n’entraînant qu’un seul modèle proxy. Au lieu de s’appuyer sur des heuristiques prédéfinies ou des simulations coûteuses en ressources, FASTMIX optimise conjointement les coefficients de mélange et les paramètres du modèle, améliorant ainsi considérablement l’efficacité et la passage à l’échelle par rapport aux approches antérieures. Au cœur de FASTMIX se trouve une reformulation de la sélection du mélange comme un problème d’optimisation bi-niveau. Sous cette reformulation, nous montrons qu’optimiser les ratios de mélange est mathématiquement équivalent à attribuer des poids de perte par source sous un échantillonnage uniforme des sources. Cela intègre directement les coefficients de mélange dans l’objectif itératif différentiable de l’optimisation, permettant une optimisation efficace, basée sur les gradients, à la fois du mélange et du modèle. Pour résoudre le problème d’optimisation, FASTMIX implémente une procédure d’optimisation itérative approchée, alternant entre (i) la mise à jour des paramètres du modèle sur des données échantillonnées selon les ratios de mélange courants (boucle interne) et (ii) la mise à jour des ratios de mélange à partir de retours de validation (boucle externe). Tant en pré-entraînement qu’en post-entraînement, FASTMIX surpasse les méthodes de référence tout en réduisant drastiquement le coût de recherche. Code (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)