FastMix: Schnelle Datenmischungsoptimierung mittels Gradientenabstieg

Zusammenfassung

Während große und vielfältige Datensätze die jüngsten Fortschritte bei großen Modellen ermöglicht haben, bleibt die Identifizierung der optimalen Datenmischung für das Pre-Training und Post-Training ein bedeutendes offenes Problem. Wir begegnen dieser Herausforderung mit FASTMIX, einem neuartigen Rahmenwerk, das die automatisierte Entdeckung von Datenmischungen ermöglicht, während nur ein einziges Proxy-Modell trainiert wird. Anstatt auf vordefinierte Heuristiken oder ressourcenintensive Simulationen angewiesen zu sein, optimiert FASTMIX gleichzeitig die Mischungskoeffizienten und die Modellparameter, was die Effizienz und Skalierbarkeit im Vergleich zu früheren Ansätzen erheblich verbessert. Im Kern von FASTMIX steht eine Neuformulierung der Mischungsauswahl als ein bileveles Optimierungsproblem. Unter dieser Neuformulierung zeigen wir, dass die Optimierung der Mischungsverhältnisse mathematisch äquivalent zur Zuweisung von Quellenverlustgewichten bei gleichmäßiger Quellenabtastung ist. Dies bettet die Mischungskoeffizienten direkt in das differenzierbare iterative Optimierungsziel ein und ermöglicht eine effiziente, gradientenbasierte Optimierung sowohl der Mischung als auch des Modells. Zur Lösung des Optimierungsproblems implementiert FASTMIX ein approximatives iteratives Optimierungsverfahren, das zwischen (i) der Aktualisierung der Modellparameter auf Daten, die gemäß den aktuellen Mischungsverhältnissen abgetastet wurden (innere Schleife), und (ii) der Aktualisierung der Mischungsverhältnisse basierend auf Validierungsfeedback (äußere Schleife) wechselt. Sowohl beim Pre- als auch beim Post-Training übertrifft FASTMIX die Basislinien, während die Suchkosten drastisch reduziert werden. 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)