Gekürzte Proximale Policy-Optimierung

Zusammenfassung

Kürzlich haben testzeit-skalierte Large Language Models (LLMs) außergewöhnliche Fähigkeiten im logischen Denken bei wissenschaftlichen und professionellen Aufgaben durch die Generierung langer Gedankenketten (Chains-of-Thought, CoT) gezeigt. Als entscheidende Komponente für die Entwicklung dieser Denkmodelle ermöglicht das Reinforcement Learning (RL), verkörpert durch Proximal Policy Optimization (PPO) und dessen Varianten, den Modellen das Lernen durch Versuch und Irrtum. Allerdings kann PPO aufgrund seines inhärenten On-Policy-Charakters zeitaufwendig sein, was durch zunehmende Antwortlängen weiter verschärft wird. In dieser Arbeit schlagen wir Truncated Proximal Policy Optimization (T-PPO) vor, eine neuartige Erweiterung von PPO, die die Trainings effizienz durch die Straffung von Policy-Updates und die Längenbegrenzung der Antwortgenerierung verbessert. T-PPO mildert das Problem der geringen Hardware-Auslastung, einen inhärenten Nachteil vollständig synchronisierter Langgenerierungsverfahren, bei denen Ressourcen oft während der Wartezeiten auf vollständige Rollouts inaktiv sind. Unsere Beiträge sind zweifach. Erstens schlagen wir Extended Generalized Advantage Estimation (EGAE) für die Vorteilsschätzung aus unvollständigen Antworten vor, während die Integrität des Policy-Lernens gewahrt bleibt. Zweitens entwickeln wir einen rechnerisch optimierten Mechanismus, der die unabhängige Optimierung der Policy- und Wertmodelle ermöglicht. Durch selektives Filtern von Prompt- und abgeschnittenen Tokens reduziert dieser Mechanismus redundante Berechnungen und beschleunigt den Trainingsprozess, ohne die Konvergenzleistung zu beeinträchtigen. Wir demonstrieren die Wirksamkeit und Effizienz von T-PPO auf AIME 2024 mit einem 32B-Basismodell. Die experimentellen Ergebnisse zeigen, dass T-PPO die Trainings effizienz von logisch denkenden LLMs um bis zu 2,5x verbessert und seine bestehenden Konkurrenten übertrifft.

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.