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

TL;DR

This paper introduces SWAP, a fine-grained reinforcement learning framework that reduces reasoning chain length in large models by penalizing less important steps, leading to shorter, more accurate reasoning chains.

Contribution

SWAP is a novel step-level length penalty method that dynamically allocates penalties based on step importance, improving reasoning efficiency and accuracy.

Findings

Reduces reasoning length by 64.3% on average.

Improves accuracy by 5.7% relative to the base model.

Demonstrates effective length reduction without sacrificing correctness.

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.