AC-ODM: Актор-критик с онлайн-смешиванием данных для эффективного по выборкам предобучения LLM

Аннотация

Оптимизация состава данных для предобучения является ключевым фактором обобщающей способности больших языковых моделей (LLM). Хотя динамическое смешивание превосходит статические стратегии за счёт учёта эволюции динамики обучения, существующие методы не позволяют одновременно обеспечить вычислительную эффективность, эффективность использования выборок и структурную гибкость для разнообразных конвейеров. Мы представляем метод Actor--Critic Online Data Mixing (AC-ODM), который подходит к смешиванию данных с позиции обучения с подкреплением, используя параметризованную политику. Теоретически мы доказываем, что эта политика действует как динамическая линейная суррогатная функция, максимизирующая конструктивную интерференцию градиентов. Для повышения практической гибкости AC-ODM поддерживает два режима работы: (i) прокси-режим для фиксированных, заранее подготовленных корпусов, когда политика, обученная на малой модели, переносится на более крупную целевую; и (ii) непрокси-режим для сквозного обучения с нуля без априорных данных. Эмпирически AC-ODM значительно превосходит предыдущие методы по скорости сходимости и точности на последующих задачах для различных архитектур. На модели Pythia-1B он достигает оптимальной перплексии на валидации, используя до 66% меньше шагов обучения, чем конкурентоспособные базовые методы, обеспечивая относительное улучшение точности MMLU на 27,5% и показатель pass@1 на HumanEval в 2,23 раза выше, при этом затрачивая практически пренебрежимо малое (0,4%) дополнительное время на шаг и лишь 2% дополнительной памяти. Код доступен по адресу 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.