Step Potential Advantage Estimation: Harnessing Intermediate Confidence and Correctness for Efficient Mathematical Reasoning

TL;DR

This paper tackles the sparsity of feedback in Reinforcement Learning with Verifiable Rewards (RLVR) by introducing step-level supervision through a training-free probing mechanism. It defines Step Potential, a unified measure combining intermediate confidence and correctness, to diagnose and steer reasoning progress. The core contribution, Step Potential Advantage Estimation (SPAE), integrates a Saturation Penalty and a Difference Shaping term to perform fine-grained, step-aware credit assignment, significantly reducing redundant post-solution checking and mitigating Right-to-Wrong failures. Extensive experiments across multiple backbones and benchmarks show SPAE improves accuracy while shortening responses, outperforming strong RLVR baselines and token-level methods, with robust generalization to out-of-domain tasks. The approach offers a practical, scalable path to more efficient and reliable long-chain reasoning in large language models.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) elicits long chain-of-thought reasoning in large language models (LLMs), but outcome-based rewards lead to coarse-grained advantage estimation. While existing approaches improve RLVR via token-level entropy or sequence-level length control, they lack a semantically grounded, step-level measure of reasoning progress. As a result, LLMs fail to distinguish necessary deduction from redundant verification: they may continue checking after reaching a correct solution and, in extreme cases, overturn a correct trajectory into an incorrect final answer. To remedy the lack of process supervision, we introduce a training-free probing mechanism that extracts intermediate confidence and correctness and combines them into a Step Potential signal that explicitly estimates the reasoning state at each step. Building on this signal, we propose Step Potential Advantage Estimation (SPAE), a fine-grained credit assignment method that amplifies potential gains, penalizes potential drops, and applies penalty after potential saturates to encourage timely termination. Experiments across multiple benchmarks show SPAE consistently improves accuracy while substantially reducing response length, outperforming strong RL baselines and recent efficient reasoning and token-level advantage estimation methods. The code is available at https://github.com/cii030/SPAE-RL.

Figures (8)

• Figure 1: Given the same user query, the baseline RLVR model DAPO produces a Right-to-Wrong failure due to over-checking, whereas SPAE terminates confidently after reaching the correct solution. The value with colored background appended after each step indicates the corresponding Step Potential.
• Figure 2: Quantifying Over-Checking on AIME2024 (averaged over 16 samples). Top: average solving and checking tokens on correct responses. Bottom: Right-to-Wrong (R2W) rate on incorrect responses.
• Figure 3: Overview of our proposed method. (a) SPAE integrates Step Potential from the probing mechanism into the RLVR optimization loop by combining group-relative outcome advantages with step-aware credit assignment after each rollout. (b) Our training-free probing mechanism estimates Step Potential by inserting a prompt after each reasoning step to compute confidence and correctness from the model’s induced answers; the bottom panel visualizes the resulting Step Potential values as a 2D contour over the (Acc, Conf) space.
• Figure 4: The metric curves of (a) generation entropy during training, (b) test accuracy, and (c) mean response length of DAPO, RF-B and SPAE based on DeepSeek-R1-Distill-Qwen-7B.
• Figure 5: Training efficiency on AIME2024: accuracy vs. cumulative training hours. SPAE reaches higher accuracy under the same wall-clock budget.