Holistische Datascheduler voor LLM Pre-training via Multi-Objectief Reinforcement Learning

Samenvatting

De samenstelling van trainingsgegevens, beheerst door de diversiteit van bronnen en hun mengstrategie, is een hoeksteen van de pre-training van grote taalmodellen (LLM's). Online Data Mixing (ODM), de techniek om datamengsels tijdens de training adaptief aan te passen, is naar voren gekomen als een veelbelovende richting om de efficiëntie te verbeteren. Bestaande methoden worden echter beperkt door hun afhankelijkheid van een enkelvoudig optimalisatieperspectief, dat fundamenteel voorbijgaat aan de noodzaak voor complexe LLM-pre-training om de dynamische datasamenstelling vanuit meerdere dimensies te beschouwen. Om deze beperking te overwinnen, introduceren wij de Holistic Data Scheduler (HDS), een nieuw online data-mixing raamwerk. HDS formuleert de uitdaging van datascheduling als een reinforcement learning-probleem in een continue controleruimte en maakt gebruik van het Soft Actor-Critic (SAC)-algoritme vanwege zijn stabiliteit en sample-efficiëntie bij het verkennen van de hoogdimensionale beleidsruimte. De kern van HDS is een nieuwe multi-objectieve, holistische beloningsfunctie die drie kritische perspectieven integreert: een data-gedreven beloning voor kwaliteit, een verliesgedreven beloning die inter-domein invloed vastlegt, en een modelgedreven beloning op basis van gewichtsnormen. Om ons ontwerp te valideren en de optimale configuratie ervan te bepalen, hebben we systematische experimenten uitgevoerd met LLM's van verschillende groottes. Op de The Pile benchmark bereikt HDS de uiteindelijke validatie-perplexiteit van de op een na beste methode met 44% minder trainingsiteraties. Bovendien behaalt het een verbetering van 7,2% op de MMLU 0-shot-taak, samen met consistente verbeteringen op andere benchmarks, wat zijn vermogen aantoont om zowel de trainingsefficiëntie als de uiteindelijke modelcapaciteit te verbeteren.

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.