Optimización de Política Proximal Truncada

Resumen

Recientemente, los modelos de lenguaje de gran escala (LLMs) ajustados en tiempo de prueba han demostrado capacidades excepcionales de razonamiento en tareas científicas y profesionales mediante la generación de largas cadenas de pensamiento (CoT). Como componente crucial para desarrollar estos modelos de razonamiento, el aprendizaje por refuerzo (RL), ejemplificado por la Optimización de Política Proximal (PPO) y sus variantes, permite que los modelos aprendan mediante prueba y error. Sin embargo, PPO puede ser lento debido a su naturaleza inherentemente on-policy, lo que se ve agravado por el aumento en la longitud de las respuestas. En este trabajo, proponemos la Optimización de Política Proximal Truncada (T-PPO), una extensión novedosa de PPO que mejora la eficiencia del entrenamiento optimizando la actualización de políticas y la generación de respuestas con longitud restringida. T-PPO mitiga el problema de la baja utilización del hardware, una desventaja inherente de los procedimientos de generación larga completamente sincronizados, donde los recursos suelen permanecer inactivos durante los períodos de espera para completar las ejecuciones. Nuestras contribuciones son dobles. Primero, proponemos la Estimación de Ventaja Generalizada Extendida (EGAE) para la estimación de ventajas derivada de respuestas incompletas, manteniendo la integridad del aprendizaje de políticas. Segundo, diseñamos un mecanismo computacionalmente optimizado que permite la optimización independiente de los modelos de política y valor. Al filtrar selectivamente los tokens de entrada y truncados, este mecanismo reduce los cálculos redundantes y acelera el proceso de entrenamiento sin sacrificar el rendimiento de convergencia. Demostramos la efectividad y eficacia de T-PPO en AIME 2024 con un modelo base de 32B. Los resultados experimentales muestran que T-PPO mejora la eficiencia del entrenamiento de LLMs de razonamiento hasta 2.5 veces y supera a sus competidores existentes.

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.