AC-ODM: Actor-Critic Online-Datenmischung für probeneffizientes LLM-Vortraining

Zusammenfassung

Die Optimierung der Zusammensetzung von Vortrainingsdaten ist entscheidend für die Generalisierung großer Sprachmodelle (LLMs). Während dynamisches Mischen statische Strategien übertrifft, indem es die sich entwickelnden Trainingsdynamiken erfasst, gelingt es aktuellen Methoden nicht, Recheneffizienz mit Stichprobeneffizienz und struktureller Flexibilität für vielfältige Pipelines in Einklang zu bringen. Wir stellen Actor-Critic Online Data Mixing (AC-ODM) vor, das Datenmischen aus einer Reinforcement-Learning-Perspektive mit einer parametrisierten Policy angeht, von der wir theoretisch beweisen, dass sie als dynamischer linearer Surrogat fungiert, der die konstruktive Interferenz von Gradienten maximiert. Um die praktische Flexibilität zu erhöhen, unterstützt AC-ODM zwei Betriebsmodi: (i) einen Proxy-Modus für feste, vorbereitete Korpora, bei dem eine an einem kleinen Modell gelernte Policy auf ein größeres Zielmodell übertragen wird; und (ii) einen Nicht-Proxy-Modus für direktes End-to-End-Training von Grund auf ohne Vorwissen. Empirisch übertrifft AC-ODM frühere Methoden in Bezug auf Konvergenzgeschwindigkeit und nachgelagerte Genauigkeit bei verschiedenen Architekturen deutlich. Auf Pythia-1B erreicht es die optimale Validierungs-Perplexität mit bis zu 66% weniger Trainingsschritten als wettbewerbsfähige Baselines, erzielt eine relative Verbesserung von 27,5% bei der MMLU-Genauigkeit und einen 2,23-fach höheren pass@1 auf HumanEval, bei einem praktisch vernachlässigbaren (0,4%) Anstieg der Wanduhrzeit pro Schritt und einem zusätzlichen Speicheraufwand von nur 2%. Der Code ist verfügbar unter https://github.com/DANG-ai/AC-ODM.

English

Optimizing pretraining data composition is pivotal for LLM generalization. While dynamic mixing outperforms static strategies by capturing evolving training dynamics, current methods fail to reconcile computational efficiency with sample efficiency and structural flexibility for diverse pipelines.We introduce Actor--Critic Online Data Mixing (AC-ODM), which approaches data mixing from a reinforcement learning perspective with a parameterized policy that we theoretically prove to act as a dynamic linear surrogate maximizing the constructive interference of gradients. To enhance practical flexibility, AC-ODM supports two operational modes: (i) a proxy mode for fixed, pre-prepared corpora, where a policy learned on a small model is transferred to a larger target; and (ii) a non-proxy mode for direct end-to-end training from scratch without priors. Empirically, AC-ODM significantly outperforms prior methods in convergence speed and downstream accuracy across various architectures. On Pythia-1B, it reaches optimal validation perplexity using up to 66% fewer training steps than competitive baselines, delivering a 27.5% relative improvement in MMLU accuracy and a 2.23 x higher pass@1 on HumanEval, all while incurring a virtually negligible (0.4%) per-step wall-clock increase and only 2% additional memory overhead. Code is available at https://github.com/DANG-ai/AC-ODM.