ESPO : Optimisation de politique proximale avec arrêt précoce

Résumé

Lorsqu'un grand modèle de langage soumis à un apprentissage par renforcement commet une étape de raisonnement erronée tôt dans une trajectoire, les algorithmes standards le forcent à continuer de générer jusqu'à l'horizon maximal, dépensant des ressources de calcul sur des jetons qui ne reçoivent jamais de récompense positive et polluant les estimations d'avantage avec du bruit post-échec. Nous proposons ESPO (Early-Stopping Proximal Policy Optimization), qui détecte l'échec de la trajectoire à la volée et termine les rollouts de manière anticipée. À chaque étape de génération, ESPO calcule un regret substitut en utilisant uniquement les logits déjà calculés lors de l'échantillonnage, et se termine lorsque le regret cumulatif lissé dépasse significativement ses valeurs estimées. Les trajectoires tronquées sont traitées comme des états d'échec absorbants avec une récompense terminale, concentrant les erreurs de différence temporelle (TD) négatives près de l'étape d'échec détectée, sans aucun modèle de récompense supplémentaire ni annotation humaine. Sur DeepSeek-R1-Distill-Qwen-7B entraîné pour le raisonnement mathématique, ESPO surpasse PPO sur AIME~2024 (46,28 % contre 45,25 %), AMC~2023 (85,83 % contre 82,94 %) et MATH-500 (87,42 % contre 85,43 %), tout en économisant plus de 20 % de jetons de rollout de manière cumulative.

English

When a large language model under reinforcement learning commits a wrong reasoning step early in a trajectory, standard algorithms force it to keep generating until the maximum horizon, spending compute on tokens that never receive positive reward and polluting advantage estimates with post-failure noise. We propose ESPO (Early-Stopping Proximal Policy Optimization), which detects trajectory failure on-the-fly and terminates rollouts early. At each generation step, ESPO computes a surrogate regret using only the logits already computed during sampling, and terminates when the smoothed cumulative regret significantly exceeds its estimated values. Truncated trajectories are treated as absorbing failure states with a terminal reward, concentrating negative temporal-difference (TD) errors near the detected failure step without any additional reward model or human annotation. On DeepSeek-R1-Distill-Qwen-7B trained for mathematical reasoning, ESPO surpasses PPO on AIME~2024 (46.28% vs. 45.25%), AMC~2023 (85.83% vs. 82.94%), and MATH-500 (87.42% vs. 85.43%), while saving more than 20% rollout tokens cumulatively.