AC-ODM: Actor-Critic Online Datamenging voor Sample-Efficiënte LLM Pre-training

Samenvatting

Het optimaliseren van de samenstelling van pretrainingdata is cruciaal voor de generalisatie van grote taalmodellen (LLM's). Hoewel dynamische menging beter presteert dan statische strategieën door zich aan te passen aan de evoluerende trainingsdynamiek, slagen huidige methoden er niet in om computationele efficiëntie te verenigen met steekproefefficiëntie en structurele flexibiliteit voor diverse pipelines. We introduceren Actor-Critic Online Data Mixing (AC-ODM), dat datamenging benadert vanuit een reinforcement learning-perspectief met een geparametriseerd beleid waarvan we theoretisch aantonen dat het fungeert als een dynamische lineaire vervanger die de constructieve interferentie van gradiënten maximaliseert. Om de praktische flexibiliteit te vergroten, ondersteunt AC-ODM twee operationele modi: (i) een proxy-modus voor vaste, voorbereide corpora, waarbij een op een klein model geleerd beleid wordt overgedragen naar een groter doelmodel; en (ii) een niet-proxy-modus voor directe end-to-end training vanaf nul zonder voorkennis. Empirisch gezien presteert AC-ODM aanzienlijk beter dan eerdere methoden wat betreft convergentiesnelheid en downstream-nauwkeurigheid over verschillende architecturen. Op Pythia-1B bereikt het de optimale validatieperplexiteit met tot 66% minder trainingsstappen dan competitieve baselines, met een relatieve verbetering van 27,5% in MMLU-nauwkeurigheid en een 2,23x hogere pass@1 op HumanEval, terwijl het een vrijwel verwaarloosbare (0,4%) toename in wandkloktijd per stap en slechts 2% extra geheugenoverhead met zich meebrengt. Code is beschikbaar op 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.