Ottimizzazione Prossimale della Politica Troncata

Abstract

Recentemente, i modelli di linguaggio di grandi dimensioni (LLMs) con scalatura al momento del test hanno dimostrato capacità di ragionamento eccezionali in compiti scientifici e professionali generando lunghe catene di pensiero (CoT). Come componente cruciale per lo sviluppo di questi modelli di ragionamento, l'apprendimento per rinforzo (RL), esemplificato dall'ottimizzazione delle politiche prossimali (PPO) e dalle sue varianti, consente ai modelli di apprendere attraverso tentativi ed errori. Tuttavia, PPO può essere dispendioso in termini di tempo a causa della sua natura intrinsecamente on-policy, che è ulteriormente aggravata dall'aumento della lunghezza delle risposte. In questo lavoro, proponiamo l'ottimizzazione delle politiche prossimali troncate (T-PPO), una nuova estensione di PPO che migliora l'efficienza dell'addestramento semplificando l'aggiornamento della politica e la generazione di risposte con lunghezza limitata. T-PPO mitiga il problema della bassa utilizzazione dell'hardware, un inconveniente intrinseco delle procedure di generazione lunga completamente sincronizzate, in cui le risorse spesso rimangono inattive durante i periodi di attesa per il completamento dei rollouts. I nostri contributi sono duplici. In primo luogo, proponiamo la stima del vantaggio generalizzata estesa (EGAE) per la stima del vantaggio derivata da risposte incomplete, mantenendo l'integrità dell'apprendimento della politica. In secondo luogo, abbiamo ideato un meccanismo ottimizzato dal punto di vista computazionale che consente l'ottimizzazione indipendente dei modelli di politica e di valore. Filtrando selettivamente i token dei prompt e quelli troncati, questo meccanismo riduce i calcoli ridondanti e accelera il processo di addestramento senza sacrificare le prestazioni di convergenza. Dimostriamo l'efficacia e l'efficienza di T-PPO su AIME 2024 con un modello base da 32B. I risultati sperimentali mostrano che T-PPO migliora l'efficienza dell'addestramento degli LLMs di ragionamento fino a 2,5 volte e supera i suoi concorrenti esistenti.

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.