AC-ODM : Mixage de Données en Ligne par Acteur-Critique pour un Pré-entraînement de LLM Économe en Échantillons

Résumé

Optimisation de la composition des données de pré-entraînement est cruciale pour la généralisation des grands modèles de langage. Bien que le mélange dynamique surpasse les stratégies statiques en capturant la dynamique évolutive de l'entraînement, les méthodes actuelles ne parviennent pas à concilier efficacité computationnelle, efficacité d'échantillonnage et flexibilité structurelle pour des pipelines diversifiés. Nous introduisons le Mélange de Données en Ligne Acteur-Critique (AC-ODM), qui aborde le mélange de données sous l'angle de l'apprentissage par renforcement avec une politique paramétrée dont nous prouvons théoriquement qu'elle agit comme un substitut linéaire dynamique maximisant l'interférence constructive des gradients. Pour améliorer la flexibilité pratique, AC-ODM prend en charge deux modes opérationnels : (i) un mode proxy pour des corpus fixes et préétablis, où une politique apprise sur un petit modèle est transférée à une cible plus grande ; et (ii) un mode non-proxy pour un entraînement direct de bout en bout à partir de zéro sans a priori. Empiriquement, AC-ODM surpasse significativement les méthodes antérieures en termes de vitesse de convergence et de précision en aval sur diverses architectures. Sur Pythia-1B, il atteint la perplexité de validation optimale en utilisant jusqu'à 66 % d'étapes d'entraînement en moins que les références concurrentes, offrant une amélioration relative de 27,5 % de la précision MMLU et un pass@1 2,23 fois plus élevé sur HumanEval, le tout avec une augmentation quasi négligeable du temps par étape (0,4 %) et seulement 2 % de surcoût mémoire supplémentaire. Le code est disponible sur 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.