Planificateur de données holistique pour le pré-entraînement de LLM via l'apprentissage par renforcement multi-objectif

Résumé

La composition des données d'entraînement, régie par la diversité des sources et leur stratégie de mélange, constitue un pilier fondamental du pré-entraînement des grands modèles de langage (LLM). Le mélange en ligne des données (ODM), technique consistant à ajuster de manière adaptative les mixtures de données durant l'entraînement, s'est imposé comme une voie prometteuse pour en améliorer l'efficacité. Cependant, les méthodes existantes sont limitées par leur dépendance à une perspective d'optimisation unique, ce qui néglige fondamentalement la nécessité, pour le pré-entraînement complexe des LLM, de considérer la composition dynamique des données selon de multiples dimensions. Pour surmonter cette limitation, nous introduisons le Holistic Data Scheduler (HDS), un nouveau cadre de mélange de données en ligne. HDS formalise le problème d'ordonnancement des données comme un problème d'apprentissage par renforcement dans un espace de contrôle continu et exploite l'algorithme Soft Actor-Critic (SAC) pour sa stabilité et son efficacité d'échantillonnage dans l'exploration d'un espace de politiques de haute dimension. Au cœur de HDS se trouve une fonction de récompense holistique multi-objectifs innovante qui intègre trois perspectives critiques : une récompense guidée par les données pour la qualité, une récompense guidée par la perte capturant l'influence inter-domaines, et une récompense guidée par le modèle basée sur les normes des poids. Pour valider notre conception et déterminer sa configuration optimale, nous avons mené des expériences systématiques sur des LLM de différentes tailles. Sur le benchmark The Pile, HDS atteint la perplexité de validation finale de la meilleure méthode suivante avec 44 % d'itérations d'entraînement en moins. De plus, il obtient une amélioration de 7,2 % sur la tâche MMLU 0-shot, accompagnée de gains constants sur d'autres benchmarks, démontrant sa capacité à améliorer à la fois l'efficacité de l'entraînement et la capacité finale du modèle.

English

The composition of training data, governed by the diversity of sources and their mixing strategy, is a cornerstone of Large Language Model (LLM) pre-training. Online Data Mixing (ODM), the technique of adaptively adjusting data mixtures during training, has emerged as a promising direction to improve efficiency. However, existing methods are constrained by their reliance on a singular optimization perspective, which fundamentally overlooks the need for complex LLM pre-training to consider the dynamic data composition from multiple dimensions. To overcome this limitation, we introduce the Holistic Data Scheduler (HDS), a novel online data mixing framework. HDS formulates the data scheduling challenge as a reinforcement learning problem in a continuous control space and leverages the Soft Actor-Critic (SAC) algorithm for its stability and sample efficiency in exploring the high-dimensional policy space. At the core of HDS lies a novel multi-objective, holistic reward function that integrates three critical perspectives: a data-driven reward for quality, a loss-driven reward capturing inter-domain influence, and a model-driven reward based on weight norms. To validate our design and determine its optimal configuration, we conducted systematic experiments on LLMs of various sizes. On The Pile benchmark, HDS reaches the final validation perplexity of the next best method with 44% fewer training iterations. Furthermore, it achieves a 7.2% improvement on the MMLU 0-shot task along with consistent gains on other benchmarks, showcasing its ability to enhance both training efficiency and final model capability.