SHAPE: Stage-aware Hierarchical Advantage via Potential Estimation for LLM Reasoning

Zhengyang Ai; Zikang Shan; Xiaodong Ai; Jingxian Tang; Hangkai Hu; Pinyan Lu

cs.LGApr 8, 2026

SHAPE: Stage-aware Hierarchical Advantage via Potential Estimation for LLM Reasoning

Authors: Zhengyang Ai, Zikang Shan, Xiaodong Ai, Jingxian Tang, Hangkai Hu, Pinyan Lu

Organizations: Huawei Taylor Lab · Center for Data Science, Peking University · Shanghai University of Finance and Economics

Abstract

Process supervision has emerged as a promising approach for enhancing LLM reasoning, yet existing methods fail to distinguish meaningful progress from mere verbosity, leading to limited reasoning capabilities and unresolved token inefficiency. To address this, we propose Stage-aware Hierarchical Advantage via Potential Estimation (SHAPE), a framework that formalizes reasoning as a trajectory through a state space of empirical solvability. SHAPE introduces a hierarchical credit assignment mechanism: at the segment level, it employs a stage-aware advantage function to prioritize efficient breakthroughs in low-potential states; at the token level, it utilizes entropy-driven redistribution to sharpen execution signals. Extensive experiments in math reasoning across three base models and five benchmarks demonstrate that SHAPE achieves an average accuracy gain of 3% with 30% reduced token consumption.

Figures & tables

Figure 1: Illustration of optimal reasoning path. Segment $s_{2}$ ranks lowest due to insufficient Potential Gain ( $\Delta=0.1$ ). While $s_{1}$ and $s_{3}$ share identical gain ( $\Delta=0.3$ ), $s_{1}$ achieves the highest rank by jointly satisfying Stage Awareness (breaking through from a harder state $\Phi=0$ ) and Efficiency (shorter length $L_{1}<L_{3}$ ).

Figure 2: Overview of the SHAPE framework. The pipeline consists of three steps: (A) decomposing reasoning to estimate state potentials (§ 4.1 ); (B) computing stage-aware segment-level advantages (§ 4.2 ); and (C) redistributing credit to sharpen token-level learning signals (§ 4.3 ).

	AIME 24		AIME 25		AMC 23		MATH500		Minerva		Overall
Method	Acc	Tokens	Acc	Tokens	Acc	Tokens	Acc	Tokens	Acc	Tokens	Acc	Tokens
DeepSeek-R1-Distill-Qwen-1.5B
GRPO	34.7	8772	27.5	8109	79.1	5091	84.8	3354	34.2	5228	52.1	6111
MRT	33.1	6577	28.6	6085	79.3	4058	85.0	2734	33.6	3705	51.9	4632
SHAPE	37.1 +2.4	6164 -29.7%	31.8 +4.3	5425 -33.1%	81.5 +2.4	3612 -29.1%	87.8 +3.0	2415 -28.0%	35.5 +1.3	3207 -38.7%	54.7 +2.6	4165 -31.8%
DeepScaleR-1.5B-Preview

Table 1: Main results on mathematical reasoning benchmarks. We report Pass@1 accuracy (%) and average generated token counts across five datasets. Blue indicates improvement (Acc

\uparrow

or Tokens

\downarrow

) of SHAPE over GRPO. Compared to GRPO and MRT baselines, SHAPE consistently establishes a new Pareto frontier across all three base models, achieving superior accuracy while significantly reducing token consumption.

Panel of Figure 3 — Figure 3: Performance of DS-R1-Distill-Qwen-1.5B.

Component Analysis
Method	AIME 24		AIME 25
Method	Acc	Tokens	Acc	Tokens
GRPO	34.7	8772	27.5	8109
SHAPE (Ours)	37.1	6164	31.8	5425
w/o EBS	36.8 -0.3	6380 +3.5%	31.6 -0.2	5590 +3.0%
w/o TCR	36.2 -0.9	6080 -1.4%	29.8 -2.0	5250 -3.2%

Table 2: Ablation study. Subscripts indicate gaps relative to SHAPE ( blue : improvement; red : degradation).

Figure 4: Marginal effect of potential gains. The steeper regression slope for the Low Start group confirms that improvements in adverse states are more critical for final success.

Figure 5: Distribution of potential gain contributions. SHAPE shifts the focus towards the Low Start regime (

44.4\%

vs. MRT’s

40.6\%

), validating its capability to rectify reasoning paths from poor initial states.

Figure 6: Potential drop rate during training. The curves depict the proportion of adjacent segment transitions where potential decreases (

\Phi(s_{k+1})<\Phi(s_{k})

Figure 7: Granularity trade-off. Accuracy (lines) saturates after

K=8

, while training cost (bars) scales linearly.

K=8

represents the optimal balance between performance gains and computational overhead.

Figure 8: Response length vs. Problem difficulty. SHAPE exhibits a steeper scaling slope and lower variance compared to GRPO, indicating precise resource allocation based on problem hardness.

Panel of Figure 9 — Figure 9: Token length distributions. SHAPE (top) shows a smooth long-tail distribution, while GRPO (bottom) exhibits anomalous spikes near the 32k context limit, indicating degenerate behavior on hard problems. MRT (middle) reduces but does not eliminate such spikes.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

$\gamma$	$\gamma\cdot\Phi_{k+1}$	$\Phi_{k}$	$F$
1.0	0.875	0.625	$+$ 0.250
0.9	0.788	0.625	$+$ 0.163
0.8	0.700	0.625	$+$ 0.075
0.7	0.613	0.625	$\mathbf{-}$ 0.013
0.6	0.525	0.625	$-$ 0.100

Appendix

Table 3: Shaping reward

F=\gamma\cdot\Phi_{k+1}-\Phi_{k}

for a positive reasoning step (

\Phi_{k}\!=\!5/8,\ \Phi_{k+1}\!=\!7/8

Hyperparameter	Value
Global Batch Size	128
Mini-Batch Size	32
Gradient Accumulation Steps	4
Learning Rate	$1\times 10^{-6}$
Warmup Ratio	0.1
Total Epochs	1

Appendix

Table 4: Detailed training hyperparameters.

Benchmark	# Probs	Samples ( $N$ )	Total Resps.
AIME 2024	30	32	960
AIME 2025	30	32	960
AMC 2023	40	32	1,280
MinervaMATH	272	4	1,088
MATH-500	500	2	1,000

Appendix

Table 5: Evaluation dataset specifications and sampling configurations.

DS-Qwen-1.5B
Method	GPQA	LiveCodeBench
GRPO	36.6	19.1
MRT	35.9	18.4
SHAPE	38.4	22.3
DeepScaleR-1.5B
GRPO	40.6	23.4

Appendix

Table 6: Zero-shot performance on OOD benchmarks.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 926 papers

Shorthand for Thought: Compressing LLM Reasoning via Entropy-Guided Supertokens
Apr 29, 2026Zhenyu Zhao, Sander Land, Daniel M. Bikel +1LLM Compression LLM Interpretability
Writer, Inc.
Efficient Reasoning via Constrained Optimization in Latent Space
Sep 28, 2026Zhinan Hou, Xingchen Li, Keyou YouOverthinking in Language Models Efficient Language Model Reasoning
Tsinghua University
SHAPE of Thinking: Mathematical Heuristics for Analyzing and Improving LLM Reasoning
Jun 28, 2026Jonghyun Song, Sangjun Song, Minjae Oh +3LLM Mathematical Reasoning Reasoning Trace Analysis
Graduate School of Data Science, Seoul National University, Seoul, Korea · Department of Mathematics Education, Seoul National University, Seoul, Korea

Apr 29, 2026J/K move · Enter open · S save

Shorthand for Thought: Compressing LLM Reasoning via Entropy-Guided Supertokens

Zhenyu Zhao, Sander Land, Daniel M. Bikel, Waseem Alshikh

Writer, Inc.

LLM Compression LLM Interpretability Efficient Language Model Reasoning LLM Inference Acceleration Structured Reasoning CoT Compression+2

Open paper AskPDF arXiv

Reasoning in Large Language Models incurs significant inference-time compute, yet the token-level information structure of reasoning traces remains underexplored. We observe that reasoning tokens split into two functional types: low-entropy structural tokens (recurring phrases that scaffold the reasoning process) and higher-entropy organic tokens (problem-specific content that drives toward a solution). This asymmetry motivates a simple, model-agnostic compression pipeline: apply cross-word BPE merges on a model's own reasoning traces to derive supertokens that capture frequent structural patterns, then teach the model to adopt them via supervised fine-tuning. Across three model families and five mathematical reasoning benchmarks, our approach shortens reasoning traces by 8.1% on average with no statistically significant accuracy loss on any model--benchmark pair. Beyond compression, supertokens act as interpretable reasoning-move annotations (backtracking, verification, strategy shifts), exposing the model's high-level strategy at a glance. Analyzing transitions between structural categories reveals systematic differences between correct and incorrect traces: correct traces show productive recovery (backtracking followed by strategy shifts and verification), while incorrect traces are dominated by confusion cycles (repeated hedging and unresolved contradictions). These diagnostic signals suggest applications in reward shaping and early stopping for RL-based reasoning training.

Abstract

LLM Reasoning

Explore similar work

Shorthand for Thought: Compressing LLM Reasoning via Entropy-Guided Supertokens

Efficient Reasoning via Constrained Optimization in Latent Space

SHAPE of Thinking: Mathematical Heuristics for Analyzing and Improving LLM Reasoning