Source-linked AI summary
Stabilizing Efficient Reasoning with Step-Level Advantage Selection
Han Wang, Xiaodong Yu, Jialian Wu, Jiang Liu, Ximeng Sun, Mohit Bansal, Zicheng Liu
TL;DR
Long, verbose reasoning traces impose computational costs, while the isolated effect of short-context post-training on compression and stability had not been systematically examined. The paper introduces Step-level Advantage Selection (SAS) to filter unreliable reasoning steps, achieving a stronger accuracy–efficiency trade-off across mathematical and general reasoning benchmarks.
Problem
Long reasoning traces increase inference latency and reduce practical efficiency, while short-context post-training had been conflated with explicit length-control objectives.
Method
Step-level Advantage Selection assigns zero advantage to low-confidence steps in correct rollouts and high-confidence steps in verifier-failed rollouts.
Results
SAS achieves the best accuracy–efficiency trade-off across mathematical benchmarks and maintains or improves accuracy with substantially shorter outputs across general reasoning tasks.
Takeaways & Limitations
Reasoning efficiency depends not only on reward design but also on credit-assignment granularity, with step-level selection helping preserve performance during aggressive compression.
Takeaways & Limitations
Experiments use a single base model and a fixed short-context post-training setting, so generalization across models and context lengths remains untested.
Abstract
from arXiv · showhide
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.
1 Introduction
LLMs gain reasoning performance through longer inference-time computation, but this increases latency and motivates efficient reasoning. The paper isolates short-context post-training as an overlooked compression signal and introduces SAS to stabilize the resulting accuracy–efficiency trade-off.
- Longer chain-of-thought reasoning can improve performance but produces verbose outputs that increase inference latency and reduce practical efficiency.
- Length-aware reinforcement-learning methods reduce reasoning length through token budgets, length-aware rewards, or pruning mechanisms.
- Short-context post-training alone compresses reasoning comparably to efficient-reasoning methods, but training accuracy fluctuates and degrades as exploration collapses into brittle updates.
- SAS treats reasoning as discrete evaluable steps and filters unreliable steps instead of applying one uniform advantage to an entire rollout.
- SAS assigns zero advantage to low-confidence steps in correct rollouts and high-confidence intermediate steps in verifier-failed rollouts.
2 Methodology
The methodology isolates how short-context GRPO affects reasoning compression and stability, then introduces SAS to selectively modulate step-level advantages. SAS suppresses unreliable steps in correct rollouts and protects reliable intermediate steps in verifier-failed rollouts.
- Short-Context Post-Training: Short-context post-training is studied with pure GRPO, excluding length-dependent rewards and pruning, to isolate context length as a compression signal.The controlled setup uses a fixed 4K context window and task-correctness rewards only.
- Short-Context Post-Training: Short-context GRPO sharply reduces output length, reaching levels comparable to or shorter than LAPO and ThinkPrune while initially maintaining comparable accuracy.The reduction occurs early and continues steadily during training.
- Short-Context Post-Training: Short-context training destabilizes accuracy as reasoning length declines, with later performance decay linked to truncation-induced noisy credit assignment.Approximately 29% of originally correct 8K rollouts become verifier-failed after truncation to 4K, often losing only the final answer or closing steps.
- Step-level Advantage Selection: SAS operates at the reasoning-step level, assigning zero advantage to low-confidence steps in correct rollouts and high-confidence steps in verifier-failed rollouts.The method modulates advantages without explicit length-aware rewards or architectural changes.
- Step-level Advantage Selection: For correct rollouts, SAS selects the lowest-confidence steps and zeros their advantages, suppressing unreliable updates while preserving signals from more reliable steps.Reasoning steps are contiguous text segments separated by double-newline delimiters, and confidence is computed from token-level log probabilities.
- Step-level Advantage Selection: For verifier-rejected rollouts, SAS zeros the highest-confidence steps, leaving other negative advantages unchanged to shield reliable intermediate reasoning from penalization.Under group-relative normalization, zero lies above negative failed-rollout advantages while remaining below positive advantages in correct rollouts.
3 Experiments
The experiments evaluate SAS and baselines on five mathematical and three general reasoning benchmarks, measuring correctness, output length, and the accuracy–efficiency trade-off. SAS achieves the strongest reported efficiency results while short-context GRPO and length-aware baselines expose compression–accuracy trade-offs.
- Experimental Setup: The evaluation covers five mathematical datasets and three out-of-domain general reasoning benchmarks, using sampled responses to measure accuracy, output length, and AES.The mathematical benchmarks are AIME2024, AIME2025, AMC, MATH, and OlympiadBench; the general benchmarks are GPQA-Diamond, LSAT, and MMLU.
- Mathematical Results: SAS achieves the best mathematical accuracy–efficiency trade-off, improving average Pass@1 accuracy by more than 2 points over the base model and reducing output length by approximately 1,700 tokens.It attains the highest AES of 0.46 and maintains or exceeds GRPO-4K accuracy across all mathematical benchmarks while generating shorter traces.
- Mathematical Results: Short-context GRPO-4K compresses mathematical reasoning and slightly improves average Pass@1 over the base model, but its training stability is limited.Among length-aware baselines, L1-Max compresses most aggressively at a significant accuracy cost, whereas ThinkPrune and LAPO preserve accuracy with only moderate length reduction.
- General Reasoning Results: On general reasoning benchmarks, GRPO-4K reduces output length but degrades accuracy, while SAS maintains or improves accuracy with substantially shorter outputs.SAS achieves the best overall efficiency score among stable methods, extending the reported benefits beyond mathematical reasoning.
4 Analysis
The ablations show that SAS’s accuracy–efficiency gains depend on verifier-failed rollouts, confidence-based selection, and step-level rather than token-level credit assignment. SAS also preserves higher policy entropy and incurs moderate computational overhead.
- Policy Stability: SAS maintains higher policy entropy than pure GRPO-4K, which exhibits rapid entropy collapse during training.The ablation describes the higher entropy as indicating more robust exploration.
- Verifier-Failed Rollouts: Restricting SAS to correct rollouts lowers average AES from 0.46 to 0.43 and accuracy from 54.54 to 53.90.The authors attribute this degradation to penalization of correct intermediate reasoning in verifier-failed rollouts.
- Confidence-Aware Selection: Random step selection reduces AES to 0.38 and produces longer traces, showing that SAS benefits from confidence-aware filtering rather than sparsification alone.The comparison supports selectively filtering low-confidence reasoning steps.
- Selection Granularity: Token-level selection achieves lower average accuracy and AES than step-level SAS, with AES of 0.39 versus 0.46 and longer outputs.The result favors semantically meaningful reasoning steps over fine-grained token-level selection.
- Selection Ratio: The best selection ratio is r = 0.3, achieving the highest Pass@1 accuracy of 54.54 across five math reasoning datasets.SAS outperforms the base model across all tested ratios, while extreme r = 0.9 remains competitive with AES above 0.36.
- Computational Overhead: SAS increases average training time per step from 279.08 to 327.15 seconds, an approximately 17% increase without extra forward passes, auxiliary models, or rollouts.The added cost comes mainly from step segmentation and step-level advantage computation; memory usage remains unchanged.
5 Related Work
Related work uses test-time scaling and length-aware reinforcement learning to balance reasoning performance against inference cost. SAS instead frames efficient reasoning as a credit-assignment problem using confidence signals to select reasoning steps.
- Test-time Scaling: Test-time scaling improves difficult reasoning tasks by allocating more inference computation, often through longer or multiple reasoning traces.These methods can produce traces several times longer than short-chain-of-thought models.
- Test-time Scaling: Extended reasoning traces can contain redundant verification or repetition, creating an overthinking problem that increases inference cost and latency.The cited discussion connects longer traces with reduced practical efficiency.
- Efficient Reasoning: Efficient-reasoning methods reduce reasoning overhead through approaches such as length-aware rewards, token budgets, or pruning while preserving task performance.SAS differs by pursuing compression through selective advantage assignment without explicit length rewards.
- Confidence-based RL: Confidence- and entropy-based reinforcement learning methods use these signals to shape rewards or update selected tokens, whereas SAS uses confidence only to select reasoning steps.SAS neither modifies the reward nor introduces entropy regularization.
6 Discussion and Conclusion
The paper argues that short-context post-training compresses reasoning but can destabilize rollout-level credit assignment when truncation or verifier failures produce noisy signals. SAS addresses this by assigning credit selectively at the reasoning-step level.
- Discussion and Conclusion: Short-context post-training alone substantially compresses reasoning in long-context models, but truncated or verifier-failed rollouts can undermine training stability.The discussion identifies credit-assignment granularity as a factor in reasoning efficiency.
- Discussion and Conclusion: SAS suppresses low-confidence steps in correct rollouts and shields high-confidence intermediate steps in verifier-failed rollouts from undue penalization.The method uses a single zero-valued advantage operation with asymmetric effects.
- Discussion and Conclusion: The paper concludes that SAS achieves stable reasoning compression with minimal modifications to standard reinforcement-learning pipelines.Its reported advantage is a stronger accuracy–efficiency trade-off than existing length-aware methods.
Limitations
The evaluation is limited to one base model and a fixed short-context post-training setting, while broader generalization and theoretical understanding remain open.
- Scope: Experiments use a single base model, so generalization across model sizes and different pretraining or post-training paradigms remains untested.The limitation explicitly identifies model and training-paradigm scope.
- Scope: All experiments use a fixed short-context post-training setting, leaving SAS behavior across varying training context lengths insufficiently studied.The authors also leave a deeper theoretical account of advantage selection for future work.
A Accuracy-Efficiency Score (AES)
The AES score measures the trade-off between accuracy and computational efficiency by comparing a tuned model with its corresponding base model, with stronger penalties for accuracy degradation.
- AES compares a tuned model against its corresponding base model using output-length and accuracy changes.It defines ∆L as the reduction in output length and ∆Acc as the accuracy difference.
- When accuracy does not decrease, AES combines length reduction with an accuracy reward.The score uses α · ∆L + β · ∆Acc when ∆Acc ≥ 0.
- When accuracy decreases, AES combines length reduction with a penalty proportional to the accuracy loss.The score uses α · ∆L − γ · |∆Acc| when ∆Acc < 0.
- The reported weights are α = 1, β = 3, and γ = 5, making accuracy degradation more heavily penalized than accuracy improvement is rewarded.The condition γ > β emphasizes the penalty for declining accuracy.
B Step Segmentation
SAS segments reasoning using complete steps marked by double-newline delimiters, a format aligned with the model family's training data and adaptable to other structured formats.
- Double-newline delimiters reflect the consistent reasoning-step format used in the model family's supervised and post-training data.The cited data span approximately 800K supervised examples and subsequent post-training data.
- SAS assigns one advantage collectively to all tokens in each complete reasoning step, including its trailing delimiter.This preserves the step formatting pattern during training.
- SAS can use alternative step boundaries for model families with different structured reasoning formats without changing its core algorithm.The method is therefore not intrinsically tied to the double-newline delimiter.
C Licenses
The cited datasets and models are distributed under Apache-2.0 or MIT licenses.
- DeepScaleR-Preview-Dataset and DeepScaleR-1.5B-Preview are released under the MIT license.
- AIME24, AIME25, and AMC are released under the Apache-2.0 license.
- OlympiadBench is released under the MIT license.
- LAPO-I is released under the Apache-2.0 license.
- ThinkPrune-4k is released under the Apache-2.0 license.