Stepwise Penalization for Length-Efficient Chain-of-Thought Reasoning
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
| Method | MATH-500 | AMC23 | AIME24 | AIME25 | OlympiadBench | Avg. | ||||||
| Acc | Len | Acc | Len | Acc | Len | Acc | Len | Acc | Len | Acc | Len | |
| 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 |
| Method | MATH-500 | AMC23 | AIME24 | AIME25 | OlympiadBench | Avg. | ||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Acc | Len | Acc | Len | Acc | Len | Acc | Len | Acc | Len | Acc | Len | |
| 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 |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
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