Optimisation Proximale de la Politique Tronquée

Résumé

Récemment, les modèles de langage à grande échelle (LLMs) ajustés en temps de test ont démontré des capacités de raisonnement exceptionnelles pour des tâches scientifiques et professionnelles en générant de longues chaînes de pensée (CoT). En tant que composant crucial pour le développement de ces modèles de raisonnement, l'apprentissage par renforcement (RL), illustré par l'Optimisation de Politique Proximale (PPO) et ses variantes, permet aux modèles d'apprendre par essais et erreurs. Cependant, PPO peut être chronophage en raison de sa nature intrinsèquement on-policy, ce qui est encore exacerbé par l'augmentation de la longueur des réponses. Dans ce travail, nous proposons l'Optimisation de Politique Proximale Tronquée (T-PPO), une extension novatrice de PPO qui améliore l'efficacité de l'entraînement en rationalisant la mise à jour de la politique et la génération de réponses de longueur limitée. T-PPO atténue le problème de la faible utilisation du matériel, un inconvénient inhérent aux procédures de génération longue entièrement synchronisées, où les ressources restent souvent inactives pendant les périodes d'attente des déploiements complets. Nos contributions sont doubles. Premièrement, nous proposons l'Estimation d'Avantage Généralisée Étendue (EGAE) pour l'estimation de l'avantage dérivée de réponses incomplètes tout en préservant l'intégrité de l'apprentissage de la politique. Deuxièmement, nous concevons un mécanisme optimisé sur le plan computationnel qui permet l'optimisation indépendante des modèles de politique et de valeur. En filtrant de manière sélective les tokens d'invite et tronqués, ce mécanisme réduit les calculs redondants et accélère le processus d'entraînement sans sacrifier les performances de convergence. Nous démontrons l'efficacité et l'efficience de T-PPO sur AIME 2024 avec un modèle de base de 32B. Les résultats expérimentaux montrent que T-PPO améliore l'efficacité de l'entraînement des LLMs de raisonnement jusqu'à 2,5 fois et surpasse ses concurrents existants.

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.