ESPO: Otimização de Política Proximal com Parada Antecipada

Resumo

Quando um modelo de linguagem de grande escala submetido a aprendizagem por reforço comete uma etapa de raciocínio incorreta no início de uma trajetória, os algoritmos padrão o obrigam a continuar gerando até o horizonte máximo, gastando computação em tokens que nunca recebem recompensa positiva e poluindo as estimativas de vantagem com ruído pós-falha. Propomos ESPO (Proximal Policy Optimization com Parada Antecipada), que detecta falhas na trajetória em tempo real e encerra as rolagens (rollouts) precocemente. A cada etapa de geração, o ESPO calcula um arrependimento substituto usando apenas os logits já computados durante a amostragem e interrompe a geração quando o arrependimento cumulativo suavizado excede significativamente seus valores estimados. Trajetórias truncadas são tratadas como estados de falha absorventes com uma recompensa terminal, concentrando erros negativos de diferença temporal (TD) próximos à etapa de falha detectada, sem exigir qualquer modelo de recompensa adicional ou anotação humana. No DeepSeek-R1-Distill-Qwen-7B treinado para raciocínio matemático, o ESPO supera o PPO no AIME 2024 (46,28% vs. 45,25%), no AMC 2023 (85,83% vs. 82,94%) e no MATH-500 (87,42% vs. 85,43%), economizando mais de 20% dos tokens de rolagem cumulativamente.

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.