AC-ODM: Mistura de Dados Online com Actor-Critic para Pré-treinamento Eficiente de LLMs

Resumo

Otimizar a composição dos dados de pré-treinamento é fundamental para a generalização de LLMs. Embora a mistura dinâmica supere as estratégias estáticas ao capturar dinâmicas de treinamento em evolução, os métodos atuais não conseguem conciliar eficiência computacional com eficiência amostral e flexibilidade estrutural para pipelines diversos. Apresentamos o Actor--Critic Online Data Mixing (AC-ODM), que aborda a mistura de dados a partir de uma perspectiva de aprendizado por reforço com uma política parametrizada que provamos teoricamente atuar como um substituto linear dinâmico que maximiza a interferência construtiva dos gradientes. Para aumentar a flexibilidade prática, o AC-ODM suporta dois modos operacionais: (i) um modo proxy para corpora fixos e pré-preparados, onde uma política aprendida em um modelo pequeno é transferida para um alvo maior; e (ii) um modo não proxy para treinamento direto ponta a ponta do zero sem conhecimentos prévios. Empiricamente, o AC-ODM supera significativamente os métodos anteriores em velocidade de convergência e precisão downstream em várias arquiteturas. No Pythia-1B, ele atinge a perplexidade de validação ideal usando até 66% menos passos de treinamento do que as linhas de base competitivas, proporcionando uma melhoria relativa de 27,5% na precisão do MMLU e um pass@1 2,23x maior no HumanEval, tudo isso com um aumento praticamente insignificante (0,4%) no tempo de execução por passo e apenas 2% de sobrecarga adicional de memória. O código está disponível em 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.