Stepwise Penalization for Length-Efficient Chain-of-Thought Reasoning

Xintong Li; Sha Li; Rongmei Lin; Hongye Jin; Linwei Li; Hejie Cui; Sarah Zhang; Chia-Yuan Chang

cs.CLFeb 27, 2026

Stepwise Penalization for Length-Efficient Chain-of-Thought Reasoning

Authors: Xintong Li, Sha Li, Rongmei Lin, Hongye Jin, Linwei Li, Hejie Cui, Sarah Zhang, Chia-Yuan Chang,

Organizations: University of California, San Diego · Amazon

Abstract

Large reasoning models improve with more test-time computation, but often overthink, producing unnecessarily long chains-of-thought that raise cost without improving accuracy. Prior reinforcement learning approaches typically rely on a single outcome reward with trajectory-level length penalties, which cannot distinguish essential from redundant reasoning steps and therefore yield blunt compression. Although recent work incorporates step-level signals, such as offline pruning, supervised data construction, or verifier-based intermediate rewards, reasoning length is rarely treated as an explicit step-level optimization objective during RL. We propose Step-wise Adaptive Penalization (SWAP), a fine-grained framework that allocates length reduction across steps based on intrinsic contribution. We estimate step importance from the model's on-policy log-probability improvement toward the correct answer, then treat excess length as a penalty mass redistributed to penalize low-importance steps more heavily while preserving high-importance reasoning. We optimize with a unified outcome-process advantage within group-relative policy optimization. Extensive experiments demonstrate that SWAP reduces reasoning length by 64.3% on average while improving accuracy by 5.7% relative to the base model.

Figures & tables

Figure 1: Step-level redundancy in reasoning. We segment the model’s responses into steps by token length and use the change in the probability of generating the ground truth answer to measure each step’s progress. (a) Distribution of step-wise information gain: most steps contribute little or no progress toward the correct answer, and high-gain steps are rare (4% with gain $>$ 0.5). (b) Longer reasoning trajectories contain a lower fraction of high-gain steps. (c) In SwAP , in addition to the outcome reward $r_{o}$ , each step also gets assigned a progress reward $r_{s}$ . When the response is correct but exceeds the reference length, the penalty (dashed grey area) gets distributed to each step based on its progress (darker shades represent more progress).

Method	MATH-500		AMC23		AIME24		AIME25		OlympiadBench		Avg.
Method	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$
DeepSeek-Distill-Qwen-1.5B
Original	83.4	4777	70.0	7722	26.7	12571	23.5	12247	35.1	11557	47.7	9775
ThinkPrune	83.1	2332	72.5	3791	24.7	7469	20.4	7024	37.8	5634	47.7	5250
LC-R1	81.9	2249	68.2	3607	23.3	7061	18.7	6647	33.3	5199	45.1	4953
AdaptThink	83.9	1863	70.0	3628	31.3	7236	22.3	7195	38.1	4900	49.1	4964

Table 1: Performance and reasoning length on mathematical reasoning benchmarks. Avg. denotes the average pass@1 across all datasets.

Method	MATH-500		AMC23		AIME24		AIME25		OlympiadBench		Avg.
Method	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$	Acc $\uparrow$	Len $\downarrow$
Original	83.4	4777	70.0	7722	26.7	12571	23.5	12247	35.1	11557	47.7	9775
Outcome-Only	79.5	1534	70.6	2953	29.8	6287	21.8	6048	41.3	4433	48.6	4251
Step-Only	79.9	1526	67.8	2915	27.1	6019	20.8	5872	31.5	3689	45.4	4004
No-Penalty	85.7	2594	77.1	4169	32.5	8027	24.3	7602	42.7	6237	52.5	5726
Static-Penalty	83.7	1276	72.8	1899	25.2	3264	20.6	2702	43.0	2110	49.1	2250

Table 2: Performance and reasoning length of different model variants. All results are based on the DeepSeek-Distill-Qwen-1.5B model. AVG denotes the average pass@1 across all datasets.

Figure 2: Comparison of Pass@1 accuracy vs. token usage across math benchmarks. SwAP establishes the Pareto frontier for best performance under every token budget.

Figure 3: Impact of step advantage weight

\theta

on model performance and response length. A moderate value of

\theta\in[0.2,0.4]

achieves the best balance. Hard datasets such as AIME 2024/2025 are more sensitive to a large

\theta

value.

Figure 4: Training dynamics of different model ablations. SwAP achieves the best validation performance, with a final average response length similar to Outcome-Only (length penalty assigned at the trajectory level).

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Comparison of Pass@1 accuracy vs. step length across four math benchmarks.

Figure 6: Distribution of step-wise information gain for the original model (red) and SwAP (blue), computed from 50k randomly sampled reasoning steps. After training, SwAP exhibits a clear shift toward positive information gain, indicating that low-importance reasoning steps are suppressed while informative steps are preserved.

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 370 papers

Stabilizing Efficient Reasoning with Step-Level Advantage Selection
Apr 27, 2026Han Wang, Xiaodong Yu, Jialian Wu +4RL for Language Model Reasoning Efficient Language Model Reasoning
UNC Chapel Hill · Advanced Micro Devices, Inc.
To Think or Not to Think: Allocating Reasoning Where It Helps
Aug 30, 2026Zhengdong He, Yunfan Zhou, Jianguo Yao +2Efficient Language Model Reasoning RL for Language Model Reasoning
Shanghai Key Laboratory of Scalable Computing and Systems · School of Computer Science, Shanghai Jiao Tong University
SmartThinker: Progressive Chain-of-Thought Length Calibration for Efficient Large Language Model Reasoning
Mar 9, 2026Chenzhi Hu, Qinzhe Hu, Yuhang Xu +6Efficient Language Model Reasoning Large Reasoning Models
Shanghai Jiao Tong University · Beihang University

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

Stabilizing Efficient Reasoning with Step-Level Advantage Selection

Han Wang, Xiaodong Yu, Jialian Wu, Jiang Liu+3

UNC Chapel Hill · Advanced Micro Devices, Inc.

RL for Language Model Reasoning Efficient Language Model Reasoning Efficient Language Model Inference Fine-Grained Credit Assignment Large Reasoning Models

Open paper AskPDF arXiv

Large language models (LLMs) achieve strong reasoning performance by allocating substantial computation at inference time, often generating long and verbose reasoning traces. While recent work on efficient reasoning reduces this overhead through length-based rewards or pruning, many approaches are post-trained under a much shorter context window than base-model training, a factor whose effect has not been systematically isolated. We first show that short-context post-training alone, using standard GRPO without any length-aware objective, already induces substantial reasoning compression-but at the cost of increasingly unstable training dynamics and accuracy degradation. To address this, we propose Step-level Advantage Selection (SAS), which operates at the reasoning-step level and assigns a zero advantage to low-confidence steps in correct rollouts and to high-confidence steps in verifier-failed rollouts, where failures often arise from truncation or verifier issues rather than incorrect reasoning. Across diverse mathematical and general reasoning benchmarks, SAS improves average Pass@1 accuracy by 0.86 points over the strongest length-aware baseline while reducing average reasoning length by 16.3%, yielding a better accuracy-efficiency trade-off.

Abstract

Efficient Language Model Reasoning

Explore similar work

Stabilizing Efficient Reasoning with Step-Level Advantage Selection

To Think or Not to Think: Allocating Reasoning Where It Helps

SmartThinker: Progressive Chain-of-Thought Length Calibration for Efficient Large Language Model Reasoning