Agendador Holístico de Dados para Pré-treinamento de LLM via Aprendizado por Reforço Multiobjetivo

Resumo

A composição dos dados de treinamento, governada pela diversidade das fontes e sua estratégia de mistura, é um pilar fundamental do pré-treinamento de Modelos de Linguagem de Grande Porte (LLM). A Mistura Online de Dados (ODM), técnica de ajustar adaptativamente as misturas de dados durante o treinamento, surgiu como uma direção promissora para melhorar a eficiência. No entanto, os métodos existentes são limitados por sua dependência de uma perspectiva de otimização singular, o que desconsidera fundamentalmente a necessidade de que o pré-treinamento complexo de LLMs considere a composição dinâmica dos dados a partir de múltiplas dimensões. Para superar essa limitação, apresentamos o Agendador Holístico de Dados (HDS), uma nova estrutura de mistura online de dados. O HDS formula o desafio do agendamento de dados como um problema de aprendizado por reforço em um espaço de controle contínuo e utiliza o algoritmo Ator-Crítico Suave (SAC) por sua estabilidade e eficiência amostral na exploração do espaço de políticas de alta dimensão. No núcleo do HDS está uma nova função de recompensa holística multi-objetivo que integra três perspectivas críticas: uma recompensa orientada por dados para qualidade, uma recompensa orientada por perda capturando a influência entre domínios e uma recompensa orientada pelo modelo com base em normas de pesos. Para validar nosso projeto e determinar sua configuração ideal, conduzimos experimentos sistemáticos em LLMs de vários tamanhos. No benchmark The Pile, o HDS atinge a perplexidade de validação final do segundo melhor método com 44% menos iterações de treinamento. Além disso, obtém uma melhoria de 7,2% na tarefa MMLU 0-shot, juntamente com ganhos consistentes em outros benchmarks, demonstrando sua capacidade de aprimorar tanto a eficiência do treinamento quanto a capacidade final do modelo.

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.