Усечённый метод оптимизации проксимальной политики

Аннотация

В последнее время крупные языковые модели (LLM), масштабируемые на этапе тестирования, продемонстрировали исключительные способности к рассуждению в научных и профессиональных задачах, генерируя длинные цепочки рассуждений (CoT). Важным компонентом для разработки таких моделей рассуждений является обучение с подкреплением (RL), примером которого служит метод Proximal Policy Optimization (PPO) и его варианты, позволяющие моделям обучаться методом проб и ошибок. Однако PPO может быть затратным по времени из-за своей внутренней природы on-policy, что усугубляется увеличением длины ответов. В данной работе мы предлагаем Truncated Proximal Policy Optimization (T-PPO), новое расширение PPO, которое повышает эффективность обучения за счет оптимизации обновления политики и генерации ответов с ограничением длины. T-PPO устраняет проблему низкой загрузки оборудования, присущую полностью синхронизированным процедурам длинной генерации, где ресурсы часто простаивают в ожидании завершения полных прогонов. Наш вклад заключается в двух аспектах. Во-первых, мы предлагаем Extended Generalized Advantage Estimation (EGAE) для оценки преимуществ, полученных из неполных ответов, при сохранении целостности обучения политики. Во-вторых, мы разрабатываем вычислительно оптимизированный механизм, позволяющий независимо оптимизировать модели политики и ценности. Благодаря избирательной фильтрации промптов и усеченных токенов этот механизм сокращает избыточные вычисления и ускоряет процесс обучения без ущерба для производительности сходимости. Мы демонстрируем эффективность T-PPO на AIME 2024 с базовой моделью на 32 млрд параметров. Результаты экспериментов показывают, что T-PPO повышает эффективность обучения моделей рассуждений до 2,5 раз и превосходит существующих конкурентов.

English

Recently, test-time scaling Large Language Models (LLMs) have demonstrated exceptional reasoning capabilities across scientific and professional tasks by generating long chains-of-thought (CoT). As a crucial component for developing these reasoning models, reinforcement learning (RL), exemplified by Proximal Policy Optimization (PPO) and its variants, allows models to learn through trial and error. However, PPO can be time-consuming due to its inherent on-policy nature, which is further exacerbated by increasing response lengths. In this work, we propose Truncated Proximal Policy Optimization (T-PPO), a novel extension to PPO that improves training efficiency by streamlining policy update and length-restricted response generation. T-PPO mitigates the issue of low hardware utilization, an inherent drawback of fully synchronized long-generation procedures, where resources often sit idle during the waiting periods for complete rollouts. Our contributions are two-folds. First, we propose Extended Generalized Advantage Estimation (EGAE) for advantage estimation derived from incomplete responses while maintaining the integrity of policy learning. Second, we devise a computationally optimized mechanism that allows for the independent optimization of the policy and value models. By selectively filtering prompt and truncated tokens, this mechanism reduces redundant computations and accelerates the training process without sacrificing convergence performance. We demonstrate the effectiveness and efficacy of T-PPO on AIME 2024 with a 32B base model. The experimental results show that T-PPO improves the training efficiency of reasoning LLMs by up to 2.5x and outperforms its existing competitors.