ESPO: Optimización de Políticas Proximales con Parada Temprana

Resumen

Cuando un modelo de lenguaje de gran tamaño entrenado con aprendizaje por refuerzo comete un paso de razonamiento incorrecto al inicio de una trayectoria, los algoritmos estándar lo obligan a seguir generando hasta alcanzar el horizonte máximo, desperdiciando cómputo en tokens que nunca recibirán una recompensa positiva y contaminando las estimaciones de ventaja con ruido posterior al fallo. Proponemos ESPO (Optimización de Políticas Proximales con Parada Anticipada), que detecta el fallo de la trayectoria sobre la marcha y finaliza los despliegues de forma temprana. En cada paso de generación, ESPO calcula un arrepentimiento sustituto utilizando únicamente los logits ya computados durante el muestreo, y se detiene cuando el arrepentimiento acumulado suavizado supera significativamente sus valores estimados. Las trayectorias truncadas se tratan como estados de fallo absorbentes con una recompensa terminal, concentrando errores negativos de diferencia temporal (TD) cerca del paso de fallo detectado, sin necesidad de ningún modelo de recompensa adicional ni anotación humana. En DeepSeek-R1-Distill-Qwen-7B entrenado para razonamiento matemático, ESPO supera a PPO en AIME 2024 (46,28 % frente a 45,25 %), AMC 2023 (85,83 % frente a 82,94 %) y MATH-500 (87,42 % frente a 85,43 %), al tiempo que ahorra más del 20 % de tokens de despliegue de forma acumulativa.

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.