Otimização de Política Proximal Truncada

Resumo

Recentemente, modelos de linguagem de grande escala (LLMs) com ajuste em tempo de teste demonstraram capacidades excepcionais de raciocínio em tarefas científicas e profissionais, gerando longas cadeias de pensamento (CoT). Como um componente crucial para o desenvolvimento desses modelos de raciocínio, o aprendizado por reforço (RL), exemplificado pelo Proximal Policy Optimization (PPO) e suas variantes, permite que os modelos aprendam por tentativa e erro. No entanto, o PPO pode ser demorado devido à sua natureza inerentemente on-policy, que é ainda mais exacerbada pelo aumento do comprimento das respostas. Neste trabalho, propomos o Truncated Proximal Policy Optimization (T-PPO), uma extensão inovadora do PPO que melhora a eficiência do treinamento ao simplificar a atualização da política e a geração de respostas com comprimento restrito. O T-PPO mitiga o problema de baixa utilização de hardware, uma desvantagem inerente dos procedimentos de geração longa totalmente sincronizados, onde os recursos frequentemente ficam ociosos durante os períodos de espera para rollouts completos. Nossas contribuições são duplas. Primeiro, propomos a Extended Generalized Advantage Estimation (EGAE) para a estimativa de vantagem derivada de respostas incompletas, mantendo a integridade do aprendizado da política. Segundo, elaboramos um mecanismo computacionalmente otimizado que permite a otimização independente dos modelos de política e valor. Ao filtrar seletivamente tokens de prompt e truncados, esse mecanismo reduz cálculos redundantes e acelera o processo de treinamento sem sacrificar o desempenho de convergência. Demonstramos a eficácia e eficiência do T-PPO no AIME 2024 com um modelo base de 32B. Os resultados experimentais mostram que o T-PPO melhora a eficiência do treinamento de LLMs de raciocínio em até 2,5x e supera seus concorrentes 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.