Afgeknotte Proximale Beleidsoptimalisatie

Samenvatting

Recentelijk hebben Large Language Models (LLMs) die tijdens de testfase worden geschaald, uitzonderlijke redeneervaardigheden getoond bij wetenschappelijke en professionele taken door lange ketens van gedachten (CoT) te genereren. Als een cruciaal onderdeel voor de ontwikkeling van deze redeneermodellen maakt reinforcement learning (RL), geïllustreerd door Proximal Policy Optimization (PPO) en zijn varianten, het mogelijk dat modellen leren door middel van trial and error. PPO kan echter tijdrovend zijn vanwege zijn inherente on-policy aard, wat verder wordt verergerd door toenemende responslengtes. In dit werk stellen we Truncated Proximal Policy Optimization (T-PPO) voor, een nieuwe uitbreiding van PPO die de trainings efficiëntie verbetert door het beleidsupdateproces en de lengtebeperkte responsgeneratie te stroomlijnen. T-PPO verlicht het probleem van lage hardwarebenutting, een inherent nadeel van volledig gesynchroniseerde lang-generatieprocedures, waarbij resources vaak ongebruikt blijven tijdens de wachtperiodes voor volledige rollouts. Onze bijdragen zijn tweeledig. Ten eerste stellen we Extended Generalized Advantage Estimation (EGAE) voor voor voordeelschatting afgeleid van onvolledige responsen, terwijl de integriteit van beleidsleren behouden blijft. Ten tweede ontwikkelen we een computationeel geoptimaliseerd mechanisme dat onafhankelijke optimalisatie van het beleids- en waardemodel mogelijk maakt. Door selectief prompts en afgekapte tokens te filteren, vermindert dit mechanisme overbodige berekeningen en versnelt het het trainingsproces zonder in te leveren op convergentieprestaties. We demonstreren de effectiviteit en efficiëntie van T-PPO op AIME 2024 met een 32B basismodel. De experimentele resultaten tonen aan dat T-PPO de trainings efficiëntie van redenerende LLMs met tot 2,5x verbetert en zijn bestaande concurrenten overtreft.

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.