AC-ODM: Mezcla de Datos en Línea Actor-Critic para el Preentrenamiento de LLM Eficiente en Muestras

Resumen

Optimizar la composición de los datos de preentrenamiento es fundamental para la generalización de los LLM. Si bien la mezcla dinámica supera a las estrategias estáticas al capturar dinámicas de entrenamiento en evolución, los métodos actuales no logran conciliar la eficiencia computacional con la eficiencia de muestreo y la flexibilidad estructural para canalizaciones diversas. Presentamos Actor--Critic Online Data Mixing (AC-ODM), que aborda la mezcla de datos desde una perspectiva de aprendizaje por refuerzo con una política parametrizada que demostramos teóricamente que actúa como un sustituto lineal dinámico que maximiza la interferencia constructiva de los gradientes. Para mejorar la flexibilidad práctica, AC-ODM admite dos modos operativos: (i) un modo proxy para corpus fijos y previamente preparados, donde una política aprendida en un modelo pequeño se transfiere a un objetivo más grande; y (ii) un modo no proxy para entrenamiento directo de extremo a extremo desde cero sin conocimientos previos. Empíricamente, AC-ODM supera significativamente a métodos anteriores en velocidad de convergencia y precisión en tareas posteriores en diversas arquitecturas. En Pythia-1B, alcanza la perplejidad de validación óptima utilizando hasta un 66% menos de pasos de entrenamiento que las líneas base competitivas, ofreciendo una mejora relativa del 27.5% en precisión en MMLU y un pass@1 2.23 veces mayor en HumanEval, todo ello con un aumento de tiempo real por paso prácticamente insignificante (0.4%) y solo un 2% adicional de sobrecarga de memoria. El código está disponible en 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.