FastMix: Snelle optimalisatie van datamengsels via gradiëntafdaling

Samenvatting

Hoewel grote en diverse datasets recente vooruitgang in grote modellen hebben gestimuleerd, blijft het identificeren van het optimale datamengsel voor voortraining en natraining een belangrijk open probleem. We pakken deze uitdaging aan met FASTMIX, een nieuw raamwerk dat de ontdekking van datamengsels automatiseert terwijl er slechts één enkel proxymodel wordt getraind. In plaats van te vertrouwen op vooraf gedefinieerde heuristieken of resource-intensieve simulaties, optimaliseert FASTMIX gezamenlijk mengselcoëfficiënten en modelparameters, wat de efficiëntie en schaalbaarheid aanzienlijk verbetert ten opzichte van eerdere benaderingen. De kern van FASTMIX is een herformulering van mengselselectie als een tweeledig optimalisatieprobleem. Onder deze herformulering laten we zien dat het optimaliseren van mengselverhoudingen wiskundig equivalent is aan het toekennen van per-bron verliesgewichten onder uniforme bronbemonstering. Dit plaatst de mengselcoëfficiënten direct in de differentieerbare iteratieve optimalisatiedoelstelling, waardoor efficiënte, gradiëntgebaseerde optimalisatie van zowel mengsel als model mogelijk wordt. Om het optimalisatieprobleem op te lossen, implementeert FASTMIX een benaderende iteratieve optimalisatieprocedure, die afwisselt tussen (i) het bijwerken van modelparameters op data die zijn bemonsterd volgens de huidige mengselverhoudingen (binnenste lus) en (ii) het bijwerken van mengselverhoudingen op basis van validatiefeedback (buitenste lus). Zowel bij voortraining als natraining presteert FASTMIX beter dan basislijnen, terwijl de zoekkosten drastisch worden verlaagd. 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)